diff --git a/docs/tools/built-in-tools.md b/docs/tools/built-in-tools.md index b2957480..f213660e 100644 --- a/docs/tools/built-in-tools.md +++ b/docs/tools/built-in-tools.md @@ -58,6 +58,19 @@ like calculations, data manipulation, or running small scripts. --8<-- "examples/java/snippets/src/main/java/tools/CodeExecutionAgentApp.java:full_code" ``` +### Vertex AI RAG Engine + +The `vertex_ai_rag_retrieval` tool allows the agent to perform private data retrieval using Vertex +AI RAG Engine. + +When you use grounding with Vertex AI RAG Engine, you need to prepare a RAG corpus before hand. +Please refer to the [RAG ADK agent sample](https://github.com/google/adk-samples/blob/main/python/agents/RAG/rag/shared_libraries/prepare_corpus_and_data.py) or [Vertex AI RAG Engine page](https://cloud.google.com/vertex-ai/generative-ai/docs/rag-engine/rag-quickstart) for setting it up. + +=== "Python" + + ```py + --8<-- "examples/python/snippets/tools/built-in-tools/vertexai_rag_engine.py" + ``` ### Vertex AI Search diff --git a/examples/python/snippets/tools/built-in-tools/vertexai_rag_engine.py b/examples/python/snippets/tools/built-in-tools/vertexai_rag_engine.py new file mode 100644 index 00000000..0eed3654 --- /dev/null +++ b/examples/python/snippets/tools/built-in-tools/vertexai_rag_engine.py @@ -0,0 +1,50 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +from google.adk.agents import Agent +from google.adk.tools.retrieval.vertex_ai_rag_retrieval import VertexAiRagRetrieval +from vertexai.preview import rag + +from dotenv import load_dotenv +from .prompts import return_instructions_root + +load_dotenv() + +ask_vertex_retrieval = VertexAiRagRetrieval( + name='retrieve_rag_documentation', + description=( + 'Use this tool to retrieve documentation and reference materials for the question from the RAG corpus,' + ), + rag_resources=[ + rag.RagResource( + # please fill in your own rag corpus + # here is a sample rag corpus for testing purpose + # e.g. projects/123/locations/us-central1/ragCorpora/456 + rag_corpus=os.environ.get("RAG_CORPUS") + ) + ], + similarity_top_k=10, + vector_distance_threshold=0.6, +) + +root_agent = Agent( + model='gemini-2.0-flash-001', + name='ask_rag_agent', + instruction=return_instructions_root(), + tools=[ + ask_vertex_retrieval, + ] +) \ No newline at end of file