Skip to content

docs: Update RAG snippets to use latest RAG module from AgentPlatform SDK #14369

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sararob wants to merge 6 commits into
base: main
Choose a base branch
from
+323 −253
Open

docs: Update RAG snippets to use latest RAG module from AgentPlatform SDK #14369

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f791c2c
docs: update RAG quickstart python snippet
sararob Jul 1, 2026
4af7bc9
update arg name
sararob Jul 1, 2026
4b41808
update additional rag samples
sararob Jul 1, 2026
bb9cbc7
update remaining samples to use latest rag module
sararob Jul 1, 2026
a8dfedd
update additional rag samples
sararob Jul 1, 2026
7bdf192
Fix sample typos
sararob Jul 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 16 additions & 14 deletions generative_ai/rag/create_corpus_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@

from typing import Optional

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Import agentplatform and types globally so they can be used in the function signature's return type annotation without raising a NameError at import time.

Suggested change
from typing import Optional
from typing import Optional
import agentplatform
from agentplatform import types


from vertexai.preview.rag import RagCorpus

PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")

Expand All @@ -26,30 +25,33 @@ def create_corpus(
) -> RagCorpus:

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Update the return type annotation to use types.RagCorpus instead of the undefined RagCorpus.

Suggested change
) -> RagCorpus:
) -> types.RagCorpus:

# [START generativeaionvertexai_rag_create_corpus]

from vertexai import rag
import vertexai
import agentplatform
from agentplatform import types

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

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

# Configure backend_config
backend_config = rag.RagVectorDbConfig(
rag_embedding_model_config=rag.RagEmbeddingModelConfig(
vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
publisher_model="publishers/google/models/text-embedding-005"
# Configure project-level config
backend_config = types.RagVectorDbConfig(
rag_embedding_model_config=types.RagEmbeddingModelConfig(
vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
endpoint="publishers/google/models/text-embedding-005"
)
)
)

corpus = rag.create_corpus(
display_name=display_name,
description=description,
backend_config=backend_config,
# Create a corpus
corpus = client.rag.create_corpus(
rag_corpus=types.RagCorpus(
display_name=display_name,
description=description,
rag_vector_db_config=backend_config,
)
)
print(corpus)
# Example response:
Expand Down
35 changes: 19 additions & 16 deletions generative_ai/rag/create_corpus_feature_store_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,6 @@

from typing import Optional

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Import agentplatform and types globally so they can be used in the function signature's return type annotation without raising a NameError at import time.

Suggested change
from typing import Optional
from typing import Optional
import agentplatform
from agentplatform import types


from vertexai.preview.rag import RagCorpus

PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")


Expand All @@ -27,31 +25,36 @@ def create_corpus_feature_store(
) -> RagCorpus:

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Update the return type annotation to use types.RagCorpus instead of the undefined RagCorpus.

Suggested change
) -> RagCorpus:
) -> types.RagCorpus:

# [START generativeaionvertexai_rag_create_corpus_feature_store]

from vertexai.preview import rag
import vertexai
import agentplatform
from agentplatform import types

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

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

# Configure embedding model (Optional)
embedding_model_config = rag.EmbeddingModelConfig(
publisher_model="publishers/google/models/text-embedding-004"
backend_config = types.RagVectorDbConfig(
rag_embedding_model_config=types.RagEmbeddingModelConfig(
vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
endpoint="publishers/google/models/text-embedding-005"
),
),
vertex_feature_store=types.RagDbConfigVertexFeatureStore(
feature_view_resource_name=feature_view_name
)
)

# Configure Vector DB
vector_db = rag.VertexFeatureStore(resource_name=feature_view_name)

corpus = rag.create_corpus(
display_name=display_name,
description=description,
embedding_model_config=embedding_model_config,
vector_db=vector_db,
corpus = client.rag.create_corpus(
rag_corpus=types.RagCorpus(
display_name=display_name,
description=description,
rag_vector_db_config=backend_config,
)
)
print(corpus)
# Example response:
Expand Down
38 changes: 18 additions & 20 deletions generative_ai/rag/create_corpus_pinecone_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,6 @@

from typing import Optional

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Import agentplatform and types globally so they can be used in the function signature's return type annotation without raising a NameError at import time.

Suggested change
from typing import Optional
from typing import Optional
import agentplatform
from agentplatform import types


from vertexai.preview.rag import RagCorpus

PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")


Expand All @@ -28,39 +26,39 @@ def create_corpus_pinecone(
) -> RagCorpus:

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Update the return type annotation to use types.RagCorpus instead of the undefined RagCorpus.

Suggested change
) -> RagCorpus:
) -> types.RagCorpus:

# [START generativeaionvertexai_rag_create_corpus_pinecone]

from vertexai import rag
import vertexai
import agentplatform
from agentplatform import types

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

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

# Configure embedding model (Optional)
embedding_model_config = rag.RagEmbeddingModelConfig(
vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
publisher_model="publishers/google/models/text-embedding-005"
embedding_model_config = types.RagEmbeddingModelConfig(
vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
endpoint="publishers/google/models/text-embedding-005"
)
)

# Configure Vector DB
vector_db = rag.Pinecone(
index_name=pinecone_index_name,
api_key=pinecone_api_key_secret_manager_version,
vector_db = types.RagVectorDbConfig(
pinecone=types.RagVectorDbConfigPinecone(
index_name=pinecone_index_name,
),
rag_embedding_model_config=embedding_model_config,
)

corpus = rag.create_corpus(
display_name=display_name,
description=description,
backend_config=rag.RagVectorDbConfig(
rag_embedding_model_config=embedding_model_config,
vector_db=vector_db,
),
corpus = client.rag.create_corpus(
rag_corpus=types.RagCorpus(
display_name=display_name,
description=description,
rag_vector_db_config=vector_db,
)
)
print(corpus)
# Example response:
Expand Down
37 changes: 19 additions & 18 deletions generative_ai/rag/create_corpus_vector_search_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,6 @@

from typing import Optional

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Import agentplatform and types globally so they can be used in the function signature's return type annotation without raising a NameError at import time.

Suggested change
from typing import Optional
from typing import Optional
import agentplatform
from agentplatform import types


from vertexai.preview.rag import RagCorpus

PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")


Expand All @@ -28,8 +26,8 @@ def create_corpus_vector_search(
) -> RagCorpus:

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Update the return type annotation to use types.RagCorpus instead of the undefined RagCorpus.

Suggested change
) -> RagCorpus:
) -> types.RagCorpus:

# [START generativeaionvertexai_rag_create_corpus_vector_search]

from vertexai import rag
import vertexai
import agentplatform
from agentplatform import types

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
Expand All @@ -38,28 +36,31 @@ def create_corpus_vector_search(
# display_name = "test_corpus"
# description = "Corpus Description"

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

# Configure embedding model (Optional)
embedding_model_config = rag.RagEmbeddingModelConfig(
vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
publisher_model="publishers/google/models/text-embedding-005"
embedding_model_config = types.RagEmbeddingModelConfig(
vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
endpoint="publishers/google/models/text-embedding-005"
)
)

# Configure Vector DB
vector_db = rag.VertexVectorSearch(
index=vector_search_index_name, index_endpoint=vector_search_index_endpoint_name
vector_db = types.RagVectorDbConfigVertexVectorSearch(
index=vector_search_index_name,
index_endpoint=vector_search_index_endpoint_name
)

corpus = rag.create_corpus(
display_name=display_name,
description=description,
backend_config=rag.RagVectorDbConfig(
rag_embedding_model_config=embedding_model_config,
vector_db=vector_db,
),
corpus = client.rag.create_corpus(
rag_corpus=types.RagCorpus(
display_name=display_name,
description=description,
rag_vector_db_config=types.RagVectorDbConfig(
rag_embedding_model_config=embedding_model_config,
vertex_vector_search=vector_db,
),
)
)
print(corpus)
# Example response:
Expand Down
22 changes: 11 additions & 11 deletions generative_ai/rag/create_corpus_vertex_ai_search_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,6 @@

from typing import Optional

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Import agentplatform and types globally so they can be used in the function signature's return type annotation without raising a NameError at import time.

Suggested change
from typing import Optional
from typing import Optional
import agentplatform
from agentplatform import types


from vertexai import rag

PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")


Expand All @@ -27,27 +25,29 @@ def create_corpus_vertex_ai_search(
) -> rag.RagCorpus:

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Update the return type annotation to use types.RagCorpus instead of the undefined rag.RagCorpus.

Suggested change
) -> rag.RagCorpus:
) -> types.RagCorpus:

# [START generativeaionvertexai_rag_create_corpus_vertex_ai_search]

from vertexai import rag
import vertexai
import agentplatform
from agentplatform import types

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

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

# Configure Search
vertex_ai_search_config = rag.VertexAiSearchConfig(
vertex_ai_search_config = types.VertexAiSearchConfig(
serving_config=f"{vertex_ai_search_engine_name}/servingConfigs/default_search",
)

corpus = rag.create_corpus(
display_name=display_name,
description=description,
vertex_ai_search_config=vertex_ai_search_config,
corpus = client.rag.create_corpus(
rag_corpus=types.RagCorpus(
display_name=display_name,
description=description,
vertex_ai_search_config=vertex_ai_search_config,
),
)
print(corpus)
# Example response:
Expand Down
35 changes: 19 additions & 16 deletions generative_ai/rag/create_corpus_weaviate_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,6 @@

from typing import Optional

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Import agentplatform and types globally so they can be used in the function signature's return type annotation without raising a NameError at import time.

Suggested change
from typing import Optional
from typing import Optional
import agentplatform
from agentplatform import types


from vertexai.preview.rag import RagCorpus

PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")


Expand All @@ -29,8 +27,8 @@ def create_corpus_weaviate(
) -> RagCorpus:

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Update the return type annotation to use types.RagCorpus instead of the undefined RagCorpus.

Suggested change
) -> RagCorpus:
) -> types.RagCorpus:

# [START generativeaionvertexai_rag_create_corpus_weaviate]

from vertexai.preview import rag
import vertexai
import agentplatform
from agentplatform import types

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

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

# Configure embedding model (Optional)
embedding_model_config = rag.EmbeddingModelConfig(
publisher_model="publishers/google/models/text-embedding-004"
embedding_model_config = types.RagEmbeddingModelConfig(
vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
endpoint="publishers/google/models/text-embedding-004"
)
)

# Configure Vector DB
vector_db = rag.Weaviate(
weaviate_http_endpoint=weaviate_http_endpoint,
vector_db = types.RagVectorDbConfigWeaviate(
http_endpoint=weaviate_http_endpoint,
collection_name=weaviate_collection_name,
api_key=weaviate_api_key_secret_manager_version,
)

corpus = rag.create_corpus(
display_name=display_name,
description=description,
embedding_model_config=embedding_model_config,
vector_db=vector_db,
corpus = client.rag.create_corpus(
rag_corpus=types.RagCorpus(
display_name=display_name,
description=description,
rag_embedding_model_config=embedding_model_config,
rag_vector_db_config=types.RagVectorDbConfig(
weaviate=vector_db
),
)
)
Comment on lines +57 to 66

@gemini-code-assist gemini-code-assist Bot Jul 1, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Correct the field names and nesting for types.RagCorpus. The rag_embedding_model_config and weaviate fields should be nested inside types.RagVectorDbConfig under rag_vector_db_config.

Suggested change
corpus = client.rag.create_corpus(
rag_corpus=types.RagCorpus(
display_name=display_name,
description=description,
rag_embedding_model_config=embedding_model_config,
vector_db=types.RagVectorDbConfig(
weaviate=vector_db
),
)
)
corpus = client.rag.create_corpus(
rag_corpus=types.RagCorpus(
display_name=display_name,
description=description,
rag_vector_db_config=types.RagVectorDbConfig(
rag_embedding_model_config=embedding_model_config,
weaviate=vector_db,
),
)
)

print(corpus)
# Example response:
Expand Down
9 changes: 4 additions & 5 deletions generative_ai/rag/delete_corpus_example.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,17 +20,16 @@
def delete_corpus(corpus_name: str) -> None:
# [START generativeaionvertexai_rag_delete_corpus]

from vertexai import rag
import vertexai
import agentplatform

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"

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

rag.delete_corpus(name=corpus_name)
client.rag.delete_corpus(name=corpus_name)
print(f"Corpus {corpus_name} deleted.")
# Example response:
# Successfully deleted the RagCorpus.
Expand Down
Loading
Loading