Skip to content

Commit 8b6869e

Browse files
sararobsararob
authored
docs: Update RAG snippets to use latest RAG module from AgentPlatform SDK (#14369)
* docs: update RAG quickstart python snippet * update arg name * update additional rag samples * update remaining samples to use latest rag module * update additional rag samples * Fix sample typos * Update imports and return types in RAG samples * Remove unused imports * Fix linter errors * Fix linter errors * Fix more linter errors * Fix sample test runner * fix import order * fix import error * Fix typos
1 parent ce8d461 commit 8b6869e

19 files changed

Lines changed: 353 additions & 293 deletions

‎generative_ai/rag/create_corpus_example.py‎

Lines changed: 18 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -15,41 +15,44 @@
1515

1616
from typing import Optional
1717

18-
from vertexai.preview.rag import RagCorpus
18+
from agentplatform import types
1919

2020
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
2121

2222

2323
def create_corpus(
2424
display_name: Optional[str] = None,
2525
description: Optional[str] = None,
26-
) -> RagCorpus:
26+
) -> types.RagCorpus:
2727
# [START generativeaionvertexai_rag_create_corpus]
2828

29-
from vertexai import rag
30-
import vertexai
29+
import agentplatform
30+
from agentplatform import types
3131

3232
# TODO(developer): Update and un-comment below lines
3333
# PROJECT_ID = "your-project-id"
3434
# display_name = "test_corpus"
3535
# description = "Corpus Description"
3636

37-
# Initialize Vertex AI API once per session
38-
vertexai.init(project=PROJECT_ID, location="us-central1")
37+
# Initialize Agent Platform client once per session
38+
client = agentplatform.Client(project=PROJECT_ID, location="us-central1")
3939

40-
# Configure backend_config
41-
backend_config = rag.RagVectorDbConfig(
42-
rag_embedding_model_config=rag.RagEmbeddingModelConfig(
43-
vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
44-
publisher_model="publishers/google/models/text-embedding-005"
40+
# Configure project-level config
41+
backend_config = types.RagVectorDbConfig(
42+
rag_embedding_model_config=types.RagEmbeddingModelConfig(
43+
vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
44+
endpoint="publishers/google/models/text-embedding-005"
4545
)
4646
)
4747
)
4848

49-
corpus = rag.create_corpus(
50-
display_name=display_name,
51-
description=description,
52-
backend_config=backend_config,
49+
# Create a corpus
50+
corpus = client.rag.create_corpus(
51+
rag_corpus=types.RagCorpus(
52+
display_name=display_name,
53+
description=description,
54+
rag_vector_db_config=backend_config,
55+
)
5356
)
5457
print(corpus)
5558
# Example response:

‎generative_ai/rag/create_corpus_feature_store_example.py‎

Lines changed: 21 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
from typing import Optional
1717

18-
from vertexai.preview.rag import RagCorpus
18+
from agentplatform import types
1919

2020
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
2121

@@ -24,34 +24,39 @@ def create_corpus_feature_store(
2424
feature_view_name: str,
2525
display_name: Optional[str] = None,
2626
description: Optional[str] = None,
27-
) -> RagCorpus:
27+
) -> types.RagCorpus:
2828
# [START generativeaionvertexai_rag_create_corpus_feature_store]
2929

30-
from vertexai.preview import rag
31-
import vertexai
30+
import agentplatform
31+
from agentplatform import types
3232

3333
# TODO(developer): Update and un-comment below lines
3434
# PROJECT_ID = "your-project-id"
3535
# feature_view_name = "projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}"
3636
# display_name = "test_corpus"
3737
# description = "Corpus Description"
3838

39-
# Initialize Vertex AI API once per session
40-
vertexai.init(project=PROJECT_ID, location="us-central1")
39+
# Initialize Agent Platform client once per session
40+
client = agentplatform.Client(project=PROJECT_ID, location="us-central1")
4141

4242
# Configure embedding model (Optional)
43-
embedding_model_config = rag.EmbeddingModelConfig(
44-
publisher_model="publishers/google/models/text-embedding-004"
43+
backend_config = types.RagVectorDbConfig(
44+
rag_embedding_model_config=types.RagEmbeddingModelConfig(
45+
vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
46+
endpoint="publishers/google/models/text-embedding-005"
47+
),
48+
),
49+
vertex_feature_store=types.RagVectorDbConfigVertexFeatureStore(
50+
feature_view_resource_name=feature_view_name
51+
)
4552
)
4653

47-
# Configure Vector DB
48-
vector_db = rag.VertexFeatureStore(resource_name=feature_view_name)
49-
50-
corpus = rag.create_corpus(
51-
display_name=display_name,
52-
description=description,
53-
embedding_model_config=embedding_model_config,
54-
vector_db=vector_db,
54+
corpus = client.rag.create_corpus(
55+
rag_corpus=types.RagCorpus(
56+
display_name=display_name,
57+
description=description,
58+
rag_vector_db_config=backend_config,
59+
)
5560
)
5661
print(corpus)
5762
# Example response:

‎generative_ai/rag/create_corpus_pinecone_example.py‎

Lines changed: 21 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -15,52 +15,53 @@
1515

1616
from typing import Optional
1717

18-
from vertexai.preview.rag import RagCorpus
18+
from agentplatform import types
1919

2020
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
2121

2222

23+
2324
def create_corpus_pinecone(
2425
pinecone_index_name: str,
2526
pinecone_api_key_secret_manager_version: str,
2627
display_name: Optional[str] = None,
2728
description: Optional[str] = None,
28-
) -> RagCorpus:
29+
) -> types.RagCorpus:
2930
# [START generativeaionvertexai_rag_create_corpus_pinecone]
3031

31-
from vertexai import rag
32-
import vertexai
32+
import agentplatform
33+
from agentplatform import types
3334

3435
# TODO(developer): Update and un-comment below lines
3536
# PROJECT_ID = "your-project-id"
3637
# pinecone_index_name = "pinecone-index-name"
37-
# pinecone_api_key_secret_manager_version = "projects/{PROJECT_ID}/secrets/{SECRET_NAME}/versions/latest"
3838
# display_name = "test_corpus"
3939
# description = "Corpus Description"
4040

41-
# Initialize Vertex AI API once per session
42-
vertexai.init(project=PROJECT_ID, location="us-central1")
41+
# Initialize Agent Platform client once per session
42+
client = agentplatform.Client(project=PROJECT_ID, location="us-central1")
4343

4444
# Configure embedding model (Optional)
45-
embedding_model_config = rag.RagEmbeddingModelConfig(
46-
vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
47-
publisher_model="publishers/google/models/text-embedding-005"
45+
embedding_model_config = types.RagEmbeddingModelConfig(
46+
vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
47+
endpoint="publishers/google/models/text-embedding-005"
4848
)
4949
)
5050

5151
# Configure Vector DB
52-
vector_db = rag.Pinecone(
53-
index_name=pinecone_index_name,
54-
api_key=pinecone_api_key_secret_manager_version,
52+
vector_db = types.RagVectorDbConfig(
53+
pinecone=types.RagVectorDbConfigPinecone(
54+
index_name=pinecone_index_name,
55+
),
56+
rag_embedding_model_config=embedding_model_config,
5557
)
5658

57-
corpus = rag.create_corpus(
58-
display_name=display_name,
59-
description=description,
60-
backend_config=rag.RagVectorDbConfig(
61-
rag_embedding_model_config=embedding_model_config,
62-
vector_db=vector_db,
63-
),
59+
corpus = client.rag.create_corpus(
60+
rag_corpus=types.RagCorpus(
61+
display_name=display_name,
62+
description=description,
63+
rag_vector_db_config=vector_db,
64+
)
6465
)
6566
print(corpus)
6667
# Example response:

‎generative_ai/rag/create_corpus_vector_search_example.py‎

Lines changed: 21 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -15,21 +15,20 @@
1515

1616
from typing import Optional
1717

18-
from vertexai.preview.rag import RagCorpus
18+
from agentplatform import types
1919

2020
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
2121

22-
2322
def create_corpus_vector_search(
2423
vector_search_index_name: str,
2524
vector_search_index_endpoint_name: str,
2625
display_name: Optional[str] = None,
2726
description: Optional[str] = None,
28-
) -> RagCorpus:
27+
) -> types.RagCorpus:
2928
# [START generativeaionvertexai_rag_create_corpus_vector_search]
3029

31-
from vertexai import rag
32-
import vertexai
30+
import agentplatform
31+
from agentplatform import types
3332

3433
# TODO(developer): Update and un-comment below lines
3534
# PROJECT_ID = "your-project-id"
@@ -38,28 +37,31 @@ def create_corpus_vector_search(
3837
# display_name = "test_corpus"
3938
# description = "Corpus Description"
4039

41-
# Initialize Vertex AI API once per session
42-
vertexai.init(project=PROJECT_ID, location="us-central1")
40+
# Initialize Agent Platform client once per session
41+
client = agentplatform.Client(project=PROJECT_ID, location="us-central1")
4342

4443
# Configure embedding model (Optional)
45-
embedding_model_config = rag.RagEmbeddingModelConfig(
46-
vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
47-
publisher_model="publishers/google/models/text-embedding-005"
44+
embedding_model_config = types.RagEmbeddingModelConfig(
45+
vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
46+
endpoint="publishers/google/models/text-embedding-005"
4847
)
4948
)
5049

5150
# Configure Vector DB
52-
vector_db = rag.VertexVectorSearch(
53-
index=vector_search_index_name, index_endpoint=vector_search_index_endpoint_name
51+
vector_db = types.RagVectorDbConfigVertexVectorSearch(
52+
index=vector_search_index_name,
53+
index_endpoint=vector_search_index_endpoint_name
5454
)
5555

56-
corpus = rag.create_corpus(
57-
display_name=display_name,
58-
description=description,
59-
backend_config=rag.RagVectorDbConfig(
60-
rag_embedding_model_config=embedding_model_config,
61-
vector_db=vector_db,
62-
),
56+
corpus = client.rag.create_corpus(
57+
rag_corpus=types.RagCorpus(
58+
display_name=display_name,
59+
description=description,
60+
rag_vector_db_config=types.RagVectorDbConfig(
61+
rag_embedding_model_config=embedding_model_config,
62+
vertex_vector_search=vector_db,
63+
),
64+
)
6365
)
6466
print(corpus)
6567
# Example response:

‎generative_ai/rag/create_corpus_vertex_ai_search_example.py‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
from typing import Optional
1717

18-
from vertexai import rag
18+
from agentplatform import types
1919

2020
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
2121

@@ -24,30 +24,32 @@ def create_corpus_vertex_ai_search(
2424
vertex_ai_search_engine_name: str,
2525
display_name: Optional[str] = None,
2626
description: Optional[str] = None,
27-
) -> rag.RagCorpus:
27+
) -> types.RagCorpus:
2828
# [START generativeaionvertexai_rag_create_corpus_vertex_ai_search]
2929

30-
from vertexai import rag
31-
import vertexai
30+
import agentplatform
31+
from agentplatform import types
3232

3333
# TODO(developer): Update and un-comment below lines
3434
# PROJECT_ID = "your-project-id"
3535
# vertex_ai_search_engine_name = "projects/{PROJECT_ID}/locations/{LOCATION}/collections/default_collection/engines/{ENGINE_ID}"
3636
# display_name = "test_corpus"
3737
# description = "Corpus Description"
3838

39-
# Initialize Vertex AI API once per session
40-
vertexai.init(project=PROJECT_ID, location="us-central1")
39+
# Initialize Agent Platform client once per session
40+
client = agentplatform.Client(project=PROJECT_ID, location="us-central1")
4141

4242
# Configure Search
43-
vertex_ai_search_config = rag.VertexAiSearchConfig(
43+
vertex_ai_search_config = types.VertexAiSearchConfig(
4444
serving_config=f"{vertex_ai_search_engine_name}/servingConfigs/default_search",
4545
)
4646

47-
corpus = rag.create_corpus(
48-
display_name=display_name,
49-
description=description,
50-
vertex_ai_search_config=vertex_ai_search_config,
47+
corpus = client.rag.create_corpus(
48+
rag_corpus=types.RagCorpus(
49+
display_name=display_name,
50+
description=description,
51+
vertex_ai_search_config=vertex_ai_search_config,
52+
),
5153
)
5254
print(corpus)
5355
# Example response:

‎generative_ai/rag/create_corpus_weaviate_example.py‎

Lines changed: 21 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
from typing import Optional
1717

18-
from vertexai.preview.rag import RagCorpus
18+
from agentplatform import types
1919

2020
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
2121

@@ -26,11 +26,11 @@ def create_corpus_weaviate(
2626
weaviate_api_key_secret_manager_version: str,
2727
display_name: Optional[str] = None,
2828
description: Optional[str] = None,
29-
) -> RagCorpus:
29+
) -> types.RagCorpus:
3030
# [START generativeaionvertexai_rag_create_corpus_weaviate]
3131

32-
from vertexai.preview import rag
33-
import vertexai
32+
import agentplatform
33+
from agentplatform import types
3434

3535
# TODO(developer): Update and un-comment below lines
3636
# PROJECT_ID = "your-project-id"
@@ -40,26 +40,31 @@ def create_corpus_weaviate(
4040
# display_name = "test_corpus"
4141
# description = "Corpus Description"
4242

43-
# Initialize Vertex AI API once per session
44-
vertexai.init(project=PROJECT_ID, location="us-central1")
43+
# Initialize Agent Platform client once per session
44+
client = agentplatform.Client(project=PROJECT_ID, location="us-central1")
4545

4646
# Configure embedding model (Optional)
47-
embedding_model_config = rag.EmbeddingModelConfig(
48-
publisher_model="publishers/google/models/text-embedding-004"
47+
embedding_model_config = types.RagEmbeddingModelConfig(
48+
vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
49+
endpoint="publishers/google/models/text-embedding-004"
50+
)
4951
)
5052

5153
# Configure Vector DB
52-
vector_db = rag.Weaviate(
53-
weaviate_http_endpoint=weaviate_http_endpoint,
54+
vector_db = types.RagVectorDbConfigWeaviate(
55+
http_endpoint=weaviate_http_endpoint,
5456
collection_name=weaviate_collection_name,
55-
api_key=weaviate_api_key_secret_manager_version,
5657
)
5758

58-
corpus = rag.create_corpus(
59-
display_name=display_name,
60-
description=description,
61-
embedding_model_config=embedding_model_config,
62-
vector_db=vector_db,
59+
corpus = client.rag.create_corpus(
60+
rag_corpus=types.RagCorpus(
61+
display_name=display_name,
62+
description=description,
63+
rag_embedding_model_config=embedding_model_config,
64+
rag_vector_db_config=types.RagVectorDbConfig(
65+
weaviate=vector_db
66+
),
67+
)
6368
)
6469
print(corpus)
6570
# Example response:

0 commit comments

Comments
 (0)