Usar a pesquisa vetorial da Agent Platform com o mecanismo RAG

import agentplatform
from agentplatform import types

# TODO(developer): Update and un-comment below lines
# PROJECT_ID = "your-project-id"
# vector_search_index_name = "projects/{PROJECT_ID}/locations/{LOCATION}/indexes/{INDEX_ID}"
# vector_search_index_endpoint_name = "projects/{PROJECT_ID}/locations/{LOCATION}/indexEndpoints/{INDEX_ENDPOINT_ID}"
# display_name = "test_corpus"
# description = "Corpus Description"

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

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

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

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:
# RagCorpus(name='projects/1234567890/locations/us-central1/ragCorpora/1234567890',
# display_name='test_corpus', description='Corpus Description', embedding_model_config=...
# ...

import agentplatform
from agentplatform import types

from google.genai import types as genai_types

# 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}"

# Supports Google Cloud Storage and Google Drive Links
# paths = ["https://drive.google.com/file/d/123", "gs://my_bucket/my_files_dir/*"]

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

response = client.rag.import_files(
    name=corpus_name,
    import_config=types.ImportRagFilesConfig(
        gcs_source=genai_types.GcsSource(uris=[paths[1]]),
        google_drive_source=types.GoogleDriveSource(
            resource_ids=[
                types.GoogleDriveSourceResourceId(
                    resource_id=paths[0],
                    resource_type=types.ResourceType.RESOURCE_TYPE_FILE
                )
            ]
        ), # optional
        rag_file_transformation_config=types.RagFileTransformationConfig(
            rag_file_chunking_config=types.RagFileChunkingConfig(
                chunk_size=512,
                chunk_overlap=100,
            )
        ), # optional
        max_embedding_requests_per_min=900, # optional
    )
)

print(f"Imported {response.imported_rag_files_count} files.")
# Example response:
# Imported 2 files.

import agentplatform

from agentplatform import types
from google.genai import types as genai_types

# 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 Agent Platform client once per session
client = agentplatform.Client(project=PROJECT_ID, location="us-east4")

response = client.rag.retrieve_contexts(
    vertex_rag_store=genai_types.VertexRagStore(
        rag_resources=[
            genai_types.VertexRagStoreRagResource(
                rag_corpus=corpus_name,
                # Optional: supply IDs from `rag.list_files()`.
                # rag_file_ids=["rag-file-1", "rag-file-2", ...],
            )
        ],
    ),
    query=types.RagQuery(
        text="Hello World!",
        rag_retrieval_config=genai_types.RagRetrievalConfig(
            top_k=10,
            filter=genai_types.RagRetrievalConfigFilter(
                vector_distance_threshold=0.5
            ),
        ),
    )
)
print(response)
# Example response:
# contexts {
#   contexts {
#     source_uri: "gs://your-bucket-name/file.txt"
#     text: "....
#   ....

from google import genai
from google.genai import types as genai_types

# 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}"

rag_retrieval_tool = genai_types.Tool(
    retrieval=genai_types.Retrieval(
        vertex_rag_store=genai_types.VertexRagStore(
            rag_resources=[
                genai_types.VertexRagStoreRagResource(
                    rag_corpus=corpus_name
                )
            ],
            rag_retrieval_config=genai_types.RagRetrievalConfig(
                top_k=10,
                filter=genai_types.RagRetrievalConfigFilter(
                    vector_distance_threshold=0.5
                ),
            ),
        ),
    )
)

# Create a GenAI SDK client to make a generate_content request
genai_client = genai.Client(enterprise=True, project=PROJECT_ID, location="us-central1")

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]
    )
)
print(response.text)
# Example response:
#   The sky appears blue due to a phenomenon called Rayleigh scattering.
#   Sunlight, which contains all colors of the rainbow, is scattered
#   by the tiny particles in the Earth's atmosphere....
#   ...