Skip to content

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

Open
sararob wants to merge 6 commits into
GoogleCloudPlatform:mainfrom
sararob:update-rag-snippet
Open

docs: Update RAG snippets to use latest RAG module from AgentPlatform SDK#14369
sararob wants to merge 6 commits into
GoogleCloudPlatform:mainfrom
sararob:update-rag-snippet

Conversation

@sararob

@sararob sararob commented Jul 1, 2026

Copy link
Copy Markdown

Description

Fixes #

Checklist

Testing

  • I have tested this change on a live environment and verified it works as intended.

Compliance & Style


Post-Approval Actions

  • Please merge this PR for me once it is approved

@sararob sararob requested review from a team as code owners July 1, 2026 20:52
@product-auto-label product-auto-label Bot added the samples Issues that are directly related to samples. label Jul 1, 2026

@gemini-code-assist gemini-code-assist Bot left a comment

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.

Code Review

This pull request migrates various Vertex AI RAG example scripts from the legacy vertexai SDK to the new agentplatform and google.genai SDKs. While the migration updates client initialization, corpus creation, and retrieval mechanisms, the review highlights several critical issues. Most notably, multiple files will raise NameErrors because types used in function return annotations (such as RagCorpus and GenerationResponse) are either not imported globally or need to be updated to their new namespaces (types.RagCorpus and genai_types.GenerateContentResponse). Additionally, there are bugs involving incorrect configuration nesting for Weaviate, passing uninstantiated classes as arguments, using the wrong client for content generation, and referencing undefined variables in return statements.

@@ -15,7 +15,6 @@

from typing import Optional

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

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

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:

@@ -15,8 +15,6 @@

from typing import Optional

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

@@ -27,31 +25,41 @@ def create_corpus_feature_store(
) -> RagCorpus:

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:

@@ -15,8 +15,6 @@

from typing import Optional

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

Comment on lines 41 to 43
rag_resources=[
rag.RagResource(
rag_corpus=corpus_name,
# Optional: supply IDs from `rag.list_files()`.
# rag_file_ids=["rag-file-1", "rag-file-2", ...],
)
genai_types.VertexRagStoreRagResource
],

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

Instantiate genai_types.VertexRagStoreRagResource with rag_corpus=corpus_name instead of passing the class itself.

                rag_resources=[
                    genai_types.VertexRagStoreRagResource(
                        rag_corpus=corpus_name,
                    )
                ],

Comment on lines +57 to +66
rag_model = GenerativeModel(
model_name="gemini-2.0-flash-001", tools=[rag_retrieval_tool]
)
response = rag_model.generate_content("Why is the sky blue?")
response = client.models.generate_content(
model="gemini-2.5-pro",
contents="Why is the sky blue?",
config=genai_types.GenerateContentConfig(
tools=[rag_retrieval_tool]
)
)

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

Remove the unused and undefined GenerativeModel instantiation, and use genai_client instead of client to call models.generate_content.

    response = genai_client.models.generate_content(
        model="gemini-2.5-pro",
        contents="Why is the sky blue?",
        config=genai_types.GenerateContentConfig(
            tools=[rag_retrieval_tool]
        )
    )


result = await response.result()
print(f"Imported {result.imported_rag_files_count} files.")
print(f"Imported {response.imported_rag_files_count} files.")

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

The variable result returned on line 72 is undefined because result = await response.result() was removed. Please update line 72 to return response instead of result.

Comment on lines +36 to +37
files_response = client.rag.list_files(name=corpus_name)
for file in files_response.rag_files:

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

The variable files returned on line 47 is undefined because files = rag.list_files(...) was replaced by files_response = client.rag.list_files(...). Please update line 47 to return files_response instead of files.

@@ -27,75 +29,85 @@ def quickstart(
paths: List[str],
) -> Tuple[rag.RagCorpus, GenerationResponse]:

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 and genai_types.GenerateContentResponse instead of the undefined rag.RagCorpus and GenerationResponse.

Suggested change
) -> Tuple[rag.RagCorpus, GenerationResponse]:
) -> Tuple[types.RagCorpus, genai_types.GenerateContentResponse]:

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

samples Issues that are directly related to samples.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant