mirror of
https://github.com/google/adk-docs.git
synced 2026-09-14 16:16:59 +08:00
Polong Lin
2ccd53a3d9
* fix: update gemini 2.0 flash to latest model id for Live API * fix: use -flash for most pages instead of the live api model * fix: wording andformatting * fix: resolve merge conflict
91 lines
3.8 KiB
Python
91 lines
3.8 KiB
Python
import asyncio
|
|
|
|
from google.adk.agents import LlmAgent
|
|
from google.adk.runners import Runner
|
|
from google.adk.sessions import InMemorySessionService
|
|
from google.genai import types
|
|
from google.adk.tools import VertexAiSearchTool
|
|
|
|
# Replace with your actual Vertex AI Search Datastore ID
|
|
# Format: projects/<PROJECT_ID>/locations/<LOCATION>/collections/default_collection/dataStores/<DATASTORE_ID>
|
|
# e.g., "projects/12345/locations/us-central1/collections/default_collection/dataStores/my-datastore-123"
|
|
YOUR_DATASTORE_ID = "YOUR_DATASTORE_ID_HERE"
|
|
|
|
# Constants
|
|
APP_NAME_VSEARCH = "vertex_search_app"
|
|
USER_ID_VSEARCH = "user_vsearch_1"
|
|
SESSION_ID_VSEARCH = "session_vsearch_1"
|
|
AGENT_NAME_VSEARCH = "doc_qa_agent"
|
|
GEMINI_2_FLASH = "gemini-2.0-flash"
|
|
|
|
# Tool Instantiation
|
|
# You MUST provide your datastore ID here.
|
|
vertex_search_tool = VertexAiSearchTool(data_store_id=YOUR_DATASTORE_ID)
|
|
|
|
# Agent Definition
|
|
doc_qa_agent = LlmAgent(
|
|
name=AGENT_NAME_VSEARCH,
|
|
model=GEMINI_2_FLASH, # Requires Gemini model
|
|
tools=[vertex_search_tool],
|
|
instruction=f"""You are a helpful assistant that answers questions based on information found in the document store: {YOUR_DATASTORE_ID}.
|
|
Use the search tool to find relevant information before answering.
|
|
If the answer isn't in the documents, say that you couldn't find the information.
|
|
""",
|
|
description="Answers questions using a specific Vertex AI Search datastore.",
|
|
)
|
|
|
|
# Session and Runner Setup
|
|
session_service_vsearch = InMemorySessionService()
|
|
runner_vsearch = Runner(
|
|
agent=doc_qa_agent, app_name=APP_NAME_VSEARCH, session_service=session_service_vsearch
|
|
)
|
|
session_vsearch = session_service_vsearch.create_session(
|
|
app_name=APP_NAME_VSEARCH, user_id=USER_ID_VSEARCH, session_id=SESSION_ID_VSEARCH
|
|
)
|
|
|
|
# Agent Interaction Function
|
|
async def call_vsearch_agent_async(query):
|
|
print("\n--- Running Vertex AI Search Agent ---")
|
|
print(f"Query: {query}")
|
|
if "YOUR_DATASTORE_ID_HERE" in YOUR_DATASTORE_ID:
|
|
print("Skipping execution: Please replace YOUR_DATASTORE_ID_HERE with your actual datastore ID.")
|
|
print("-" * 30)
|
|
return
|
|
|
|
content = types.Content(role='user', parts=[types.Part(text=query)])
|
|
final_response_text = "No response received."
|
|
try:
|
|
async for event in runner_vsearch.run_async(
|
|
user_id=USER_ID_VSEARCH, session_id=SESSION_ID_VSEARCH, new_message=content
|
|
):
|
|
# Like Google Search, results are often embedded in the model's response.
|
|
if event.is_final_response() and event.content and event.content.parts:
|
|
final_response_text = event.content.parts[0].text.strip()
|
|
print(f"Agent Response: {final_response_text}")
|
|
# You can inspect event.grounding_metadata for source citations
|
|
if event.grounding_metadata:
|
|
print(f" (Grounding metadata found with {len(event.grounding_metadata.grounding_attributions)} attributions)")
|
|
|
|
except Exception as e:
|
|
print(f"An error occurred: {e}")
|
|
print("Ensure your datastore ID is correct and the service account has permissions.")
|
|
print("-" * 30)
|
|
|
|
# --- Run Example ---
|
|
async def run_vsearch_example():
|
|
# Replace with a question relevant to YOUR datastore content
|
|
await call_vsearch_agent_async("Summarize the main points about the Q2 strategy document.")
|
|
await call_vsearch_agent_async("What safety procedures are mentioned for lab X?")
|
|
|
|
# Execute the example
|
|
# await run_vsearch_example()
|
|
|
|
# Running locally due to potential colab asyncio issues with multiple awaits
|
|
try:
|
|
asyncio.run(run_vsearch_example())
|
|
except RuntimeError as e:
|
|
if "cannot be called from a running event loop" in str(e):
|
|
print("Skipping execution in running event loop (like Colab/Jupyter). Run locally.")
|
|
else:
|
|
raise e
|