Search private data stores with Agent Search

Answer a signing-key question from internal engineering runbooks. Agent Search indexes four bundled documents, and ADK’s VertexAiSearchTool adds retrieval settings to the Gemini request. The TeamSearch subclass requires one of the fixture teams before a request can run. You compare the same question with a docs-team filter and a platform-team filter, then inspect the retrieved sources and answer citations.

When you need this. When an assistant needs answers from internal documents, such as runbooks, design records or incident reports, a managed search index can retrieve relevant passages at answer time. Agent Search fits when you want Google to run the parsing, chunking and ranking. An explicit retrieval function gives the application a place to inspect or redact results before sending them to the model.

The index belongs to your project. Retrieved document text is made available to Gemini; project ownership does not keep that text out of the model request.

Create the data store

Download the complete lab ZIP, extract it, and open a terminal in adk-agent-search/. The commands below use that directory and its pinned dependencies. The [gcp] extra of google-adk installs google-cloud-discoveryengine, the client library for Agent Search. Sign in with gcloud auth application-default login and set GOOGLE_CLOUD_PROJECT.

python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-lock.txt
gcloud auth application-default login
export GOOGLE_CLOUD_PROJECT=your-project-id

"""Create, check or delete the lesson's Agent Search data store: four private engineering documents.

python datastore.py create   # creates the data store and adds private/*.txt with a team field
python datastore.py status   # lists each document and whether it is indexed
python datastore.py delete   # deletes the data store, waits, and confirms it is gone
"""
import json
import os
import sys
import warnings
from pathlib import Path

warnings.filterwarnings("ignore")

from google.api_core import exceptions  # noqa: E402
from google.cloud import discoveryengine_v1 as de  # noqa: E402

HERE = Path(__file__).resolve().parent
PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
STORE_ID = "aiml-lesson-agent-search"  # data stores take no labels; the ID carries the lesson tag
COLLECTION = f"projects/{PROJECT}/locations/global/collections/default_collection"
STORE = f"{COLLECTION}/dataStores/{STORE_ID}"
BRANCH = f"{STORE}/branches/default_branch"
# User ADC has no quota project; Discovery Engine requires one, so name it on each client.
OPTIONS = {"quota_project_id": PROJECT}


def create() -> None:
    stores = de.DataStoreServiceClient(client_options=OPTIONS)
    operation = stores.create_data_store(
        parent=COLLECTION, data_store_id=STORE_ID,
        data_store=de.DataStore(
            display_name="aiml lesson: search-private-data-stores-with-agent-search",
            industry_vertical=de.IndustryVertical.GENERIC,
            solution_types=[de.SolutionType.SOLUTION_TYPE_SEARCH],
            content_config=de.DataStore.ContentConfig.CONTENT_REQUIRED))
    operation.result(timeout=600)
    documents = de.DocumentServiceClient(client_options=OPTIONS)
    for path in sorted((HERE / "private").glob("*.txt")):
        text = path.read_text(encoding="utf-8")
        team = next(line.split(":", 1)[1].strip() for line in text.splitlines() if line.startswith("Team:"))
        documents.create_document(parent=BRANCH, document_id=path.stem, document=de.Document(
            json_data=json.dumps({"title": text.splitlines()[0], "team": team}),
            content=de.Document.Content(raw_bytes=text.encode(), mime_type="text/plain")))


def status() -> list[tuple[str, str]]:
    documents = de.DocumentServiceClient(client_options=OPTIONS)
    return [(d.id, "indexed" if d.index_time else "not indexed yet")
            for d in documents.list_documents(parent=BRANCH)]


def delete() -> None:
    stores = de.DataStoreServiceClient(client_options=OPTIONS)
    try:
        stores.delete_data_store(name=STORE).result(timeout=600)
    except exceptions.GoogleAPICallError as error:
        # google-cloud-discoveryengine 0.13.12 raised "Long-running operation had neither response
        # nor error set" for a delete that succeeded, so check the store itself before trusting it.
        print(f"delete operation: {error}")
    try:
        stores.get_data_store(name=STORE)
    except exceptions.NotFound:
        return
    raise SystemExit(f"{STORE_ID} still exists")


if __name__ == "__main__":
    command = sys.argv[1] if sys.argv[1:] else "status"
    if command == "create":
        create()
        print(f"created {STORE_ID}")
    elif command == "status":
        print(status())
    elif command == "delete":
        delete()
        print(f"deleted {STORE_ID}")

The data store lives in the global location, in the project’s default collection, and uses IndustryVertical.GENERIC, SolutionType.SOLUTION_TYPE_SEARCH and ContentConfig.CONTENT_REQUIRED: a general document-search store that requires document content. Each file in private/ becomes one document with the text as content and a small JSON record: the title and the team named in the file. That field is what the filter in step 2 uses. Data stores take no labels, so the ID aiml-lesson-agent-search marks it as lesson data.

With user credentials, the recorded API calls failed until a quota project was set: “The discoveryengine.googleapis.com API requires a quota project, which is not set by default.” client_options={"quota_project_id": ...} names the project on each client without changing your gcloud configuration. create also returns before the documents are searchable. status printed “not indexed yet” for several minutes, so check index readiness before routing questions to it.

Create the lesson store, then check whether its documents are indexed:

python datastore.py create
python datastore.py status

A store that is not indexed yet. The recorded search returns no chunks and the model says it found no document. Wait for document index times before routing traffic; successful store creation does not establish search readiness.

Add the search tool to the agent

"""The docs assistant answers from the team's private engineering docs through Agent Search."""
import asyncio
import os
import sys

from google.adk import Agent, Runner
from google.adk.agents.readonly_context import ReadonlyContext
from google.adk.sessions import InMemorySessionService
from google.adk.tools import VertexAiSearchTool
from google.genai import types

PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
# VertexAiSearchTool is a Gemini built-in tool, so the model must run on Agent Platform, not the Gemini API.
os.environ.setdefault("GOOGLE_GENAI_USE_ENTERPRISE", "true")
os.environ.setdefault("GOOGLE_CLOUD_PROJECT", PROJECT)
os.environ.setdefault("GOOGLE_CLOUD_LOCATION", "global")  # gemini-3.8-flash is served from the global endpoint
MODEL = os.environ.get("DOCS_MODEL", "gemini-3.8-flash")
STORE = f"projects/{PROJECT}/locations/global/collections/default_collection/dataStores/aiml-lesson-agent-search"


class TeamSearch(VertexAiSearchTool):
    """Build a filter from a validated application-supplied team; this demo has no login."""

    def __init__(self, **kwargs):
        if kwargs.get("bypass_multi_tools_limit"):
            raise ValueError("TeamSearch requires its dynamic team filter; bypass mode discards it")
        super().__init__(**kwargs)

    def _build_vertex_ai_search_config(self, readonly_context: ReadonlyContext) -> types.VertexAISearch:
        team = readonly_context.state.get("team")
        if not isinstance(team, str) or team not in {"docs", "platform"}:
            raise ValueError("Search requires a known team from application state")
        config = super()._build_vertex_ai_search_config(readonly_context)
        config.filter = f'team: ANY("{team}")'
        return config


def build_agent(model=None, search: VertexAiSearchTool | None = None) -> Agent:
    return Agent(
        name="private_docs",
        model=model or MODEL,
        instruction="""You answer engineers' questions from the team's private engineering documents:
decision records, runbooks, postmortems and style guides. Answer only from what the search returns,
and name the document you used. If the search returns nothing relevant, say you found no document.""",
        tools=[search or TeamSearch(data_store_id=STORE, max_results=3)],
        generate_content_config=types.GenerateContentConfig(temperature=0.2),
    )


async def ask(text: str, agent: Agent | None = None, team: str | None = None) -> list:
    sessions = InMemorySessionService()
    runner = Runner(app_name="docs", agent=agent or build_agent(), session_service=sessions)
    try:
        session = await sessions.create_session(app_name="docs", user_id="dev", state={"team": team} if team else None)
        message = types.Content(role="user", parts=[types.Part(text=text)])
        return [e async for e in runner.run_async(user_id="dev", session_id=session.id, new_message=message)]
    finally:
        await runner.close()


def sources(event) -> list[str]:
    """The documents Agent Search returned for this answer, from the grounding metadata."""
    chunks = (event.grounding_metadata.grounding_chunks or []) if event.grounding_metadata else []
    return [c.retrieved_context.title or c.retrieved_context.uri for c in chunks if c.retrieved_context]


def show(events: list) -> str:
    lines = []
    for event in events:
        if event.grounding_metadata:
            queries = event.grounding_metadata.retrieval_queries or []
            lines.append(f"search {queries} -> {sources(event)}")
        if event.is_final_response() and event.author != "user" and event.content:
            lines.append("answer " + "".join(p.text or "" for p in event.content.parts if not p.thought).strip())
    return "\n".join(lines)


if __name__ == "__main__":
    team = os.environ.get("DOCS_TEAM")
    print(show(asyncio.run(ask(" ".join(sys.argv[1:]), team=team))))

VertexAiSearchTool is a Gemini built-in tool. It has no function the model calls: before each request, ADK adds a retrieval entry to the request’s tools, and Gemini runs the search on Google’s side. That is why the model must run on Agent Platform (GOOGLE_GENAI_USE_ENTERPRISE=true) rather than the Gemini API, and why it must be a Gemini model. gemini-3.8-flash is served from the global endpoint; in us-central1 the same model ID returned 404.

TeamSearch overrides _build_vertex_ai_search_config, the hook ADK calls for every request with the invocation’s context. It reads team from session state and returns a VertexAISearch config with the filter team: ANY("platform"). This demo reads the team from DOCS_TEAM and has no login. The subclass accepts only docs or platform, rejects missing or unknown values, and refuses bypass_multi_tools_limit because that mode discards the dynamic hook. A deployed service must derive team membership from authenticated application state. The filter narrows ranking results; it is an application control. For access control that holds whatever the client sends, Agent Search supports document-level access control on the data store (acl_enabled), tied to the end user’s identity.

A non-Gemini model. process_llm_request raises ValueError before sending a request. Use a Gemini model on Agent Platform for this built-in retrieval tool.

bypass_multi_tools_limit with other tools. The installed base class replaces itself with a function tool, dropping the subclass’s dynamic hook. TeamSearch rejects this mode at construction; keep this filtered retrieval tool on its own agent.

Missing or invalid team. TeamSearch raises before the model request rather than searching unfiltered. Supply an allowed team from authenticated application state; this demo’s DOCS_TEAM is not a login.

Ask questions from the private documents

DOCS_TEAM=docs python agent.py "How do I rotate the release signing key?"
DOCS_TEAM=platform python agent.py "How do I rotate the release signing key?"

The signing-key runbook belongs to the platform team. Compare the sources returned for the two commands. With the docs-team filter, the agent should report that it could not find the procedure. With the platform-team filter, check that the answer cites the runbook and follows its four steps. The team selector is a fixture input, not authentication.

A question the store cannot answer. The historical model run broadens its queries and then reports no match. Keep the instruction to say so, and log empty final grounding_chunks instead of treating an ordinary answer event as retrieval success.

A filter that excludes the document. The docs-team recording returns no runbook while the platform team finds it. Test each team’s expected documents; this result exercises a search filter, not user authorization.

An empty final grounding_chunks list can mean the index is not ready, the filter excluded the document or the store has no answer. Check index status and the selected team before changing the question. The offline tests preserve an earlier unfiltered empty-index example separately; it does not exercise the current team-validation path.

How Gemini applies the team filter

The tool prepares retrieval configuration rather than executing a search function in the agent process. Runner still manages the session and emits the model’s events. The saved fixture shows the platform team’s retrieval configuration. Its original documents and model replies retain the earlier shipctl scenario; current document inputs use docuctl, and the saved answers are not new docuctl results:

[
  {
    "vertex_ai_search": {
      "datastore": "projects/vertex-ai-experminent/locations/global/collections/default_collection/dataStores/aiml-lesson-agent-search",
      "filter": "team: ANY(\"platform\")",
      "max_results": 3
    }
  }
]

Gemini decides the search queries, sends them to the data store with the filter, and reads the top max_results results. The answer comes back with grounding_metadata, which ADK copies onto the event. grounding_chunks holds the retrieved text and the document it came from; grounding_supports maps spans of the answer to those chunks, which lets a UI connect each citation to the text it supports.

The trade-off against a function tool is control. You do not see the search request or its raw results, only what Gemini reports in the metadata, and you cannot post-process results before the model reads them. The provider runs retrieval; the application still maintains its retrieval configuration and access controls. If you need to rank, filter or redact results yourself, call the Agent Search API from a function tool instead.

Test the agent offline

test_agent.py replays the recorded model turns, including the grounding metadata, through the real agent and retrieval configuration. No cloud search executes in these tests. ADK still builds every request, so the tests check the retrieval config and the filter for each team.

python -m unittest -v test_agent.py

The two team-filtered runs replay through the corrected subclass. Earlier unfiltered runs replay through the base tool only to preserve their historical output; they do not test current team validation.

"""Replay the recorded Gemini turns through the real agent and VertexAiSearchTool, offline.

Run: python -m unittest -v test_agent.py
The model is replaced by its recorded turns, including the grounding metadata Agent Search returned,
so nothing reaches Google Cloud. ADK still builds every request, so the tests see the retrieval config.
"""
import asyncio
import json
import logging
import os
import unittest
import warnings
from pathlib import Path
from types import SimpleNamespace

warnings.filterwarnings("ignore")
logging.disable(logging.CRITICAL)

from google.adk.agents.readonly_context import ReadonlyContext  # noqa: E402
from google.adk.models.base_llm import BaseLlm  # noqa: E402
from google.adk.models.llm_request import LlmRequest  # noqa: E402
from google.adk.models.llm_response import LlmResponse  # noqa: E402
from google.adk.tools.discovery_engine_search_tool import DiscoveryEngineSearchTool  # noqa: E402
from google.genai import types  # noqa: E402

# Fixture project only: cloud clients and model calls are replaced in these tests.
os.environ.setdefault("GOOGLE_CLOUD_PROJECT", "vertex-ai-experminent")
import agent  # noqa: E402

RECORD = json.loads(Path(__file__).with_name("recorded-runs.json").read_text(encoding="utf-8"))


def setUpModule():
    warnings.simplefilter("ignore")


class ReplayLlm(BaseLlm):
    turns: list
    requests: list = []

    async def generate_content_async(self, llm_request, stream=False):
        self.requests.append(llm_request)
        turn = self.turns.pop(0)
        yield LlmResponse(
            content=types.Content.model_validate(turn["content"]) if turn["content"] else None,
            grounding_metadata=types.GroundingMetadata.model_validate(turn["grounding_metadata"])
            if turn["grounding_metadata"] else None)


def replay(run: dict) -> tuple[list, ReplayLlm]:
    model = ReplayLlm(model=RECORD["model"], turns=list(run["model_turns"]), requests=[])
    # Some earlier recordings deliberately searched without a team. Keep them historical;
    # current TeamSearch rejects missing teams instead of silently doing this.
    search = agent.VertexAiSearchTool(data_store_id=agent.STORE, max_results=3) if run["team"] is None else None
    events = asyncio.run(agent.ask("replayed question", agent.build_agent(model, search=search), team=run["team"]))
    return events, model


def retrieval(model: ReplayLlm) -> list[dict]:
    return [t.retrieval.model_dump(mode="json", exclude_none=True)
            for t in model.requests[0].config.tools if t.retrieval]


class AgentSearchTests(unittest.TestCase):
    def test_the_tool_is_a_retrieval_config_not_a_function(self):
        events, model = replay(RECORD["runs"]["postmortem"])
        self.assertEqual(retrieval(model), RECORD["runs"]["postmortem"]["retrieval"])
        self.assertFalse(any(t.function_declarations for t in model.requests[0].config.tools))
        self.assertEqual(agent.sources([e for e in events if e.grounding_metadata][-1]),
                         RECORD["runs"]["postmortem"]["sources"])

    def test_the_team_filter_comes_from_session_state(self):
        for key, team in (("team-docs", "docs"), ("team-platform", "platform")):
            _, model = replay(RECORD["runs"][key])
            self.assertEqual(retrieval(model)[0]["vertex_ai_search"]["filter"], f'team: ANY("{team}")')
        self.assertEqual(RECORD["runs"]["team-docs"]["sources"], [])
        self.assertEqual(RECORD["runs"]["team-platform"]["sources"], ["Runbook: rotate the release signing key"])

    def test_an_unindexed_store_returns_no_sources(self):
        run = RECORD["not_indexed"]
        self.assertEqual(run["sources"], [])
        self.assertIn("found no document", run["transcript"])

    def test_a_question_the_store_cannot_answer_gets_no_sources(self):
        run = RECORD["runs"]["not-found"]
        events, _ = replay(run)
        self.assertEqual(agent.show(events), run["transcript"])
        self.assertIn("no document", run["transcript"].lower())

    def test_a_non_gemini_model_is_refused(self):
        request = LlmRequest(model="ollama_chat/gemma4:e2b")
        tool = agent.TeamSearch(data_store_id=agent.STORE)
        with self.assertRaisesRegex(ValueError, "Vertex AI search tool is not supported for model"):
            asyncio.run(tool.process_llm_request(tool_context=None, llm_request=request))

    def test_bypass_multi_tools_limit_drops_the_team_filter(self):
        def list_pages() -> list[str]:
            """List the public docs pages."""
            return []

        # Demonstrate the installed base-class behavior; TeamSearch refuses this mode.
        search = agent.VertexAiSearchTool(data_store_id=agent.STORE, max_results=3, bypass_multi_tools_limit=True)
        mixed = agent.build_agent(search=search).clone(update={"tools": [search, list_pages]})
        tools = asyncio.run(mixed.canonical_tools(None))
        replaced = [t for t in tools if isinstance(t, DiscoveryEngineSearchTool)]
        self.assertEqual(len(replaced), 1)
        self.assertNotIsInstance(replaced[0], agent.TeamSearch)
        self.assertIsNone(replaced[0]._filter)  # the per-team filter from TeamSearch is gone


class TeamFilterSafetyTests(unittest.TestCase):
    def test_missing_unknown_and_filter_syntax_teams_are_refused(self):
        tool = agent.TeamSearch(data_store_id=agent.STORE)
        for team in (None, "", "unknown", 'docs") OR team: ANY("platform', ["docs"]):
            with self.subTest(team=team), self.assertRaisesRegex(ValueError, "known team"):
                tool._build_vertex_ai_search_config(SimpleNamespace(state={"team": team}))

    def test_team_search_refuses_a_mode_that_discards_its_dynamic_filter(self):
        with self.assertRaisesRegex(ValueError, "dynamic team filter"):
            agent.TeamSearch(data_store_id=agent.STORE, bypass_multi_tools_limit=True)


if __name__ == "__main__":
    unittest.main()

Cost and cleanup

Agent Search charges for index storage and queries; model requests have separate charges. The activity indexes four small documents, but a model can search more than once per question. Check the pricing pages below and inspect the selected project before loading your own data.

Delete the data store when you are done:

python datastore.py delete

The delete is a long-running operation. delete() checks the store afterward and succeeds only when get_data_store returns NotFound, including when the operation wait reports an error. Inspect the project and lesson store ID before running it; deleting a store removes its indexed documents.

Check your understanding

What’s next

• Expose the docs assistant over A2A and consume it: make this assistant available to another agent.

References

Official API documentation and the installed ADK interfaces were checked October 22, 2026. Exact package versions are pinned in the lab lockfile.

• ADK: Agent Search
• ADK: grounding with search
• Agent Search: create a data store
• Agent Search: filter search results
• Agent Search: pricing
• Authentication: user credentials and quota projects

Previously on the blog: Google ADK + Vertex AI Search Tool (VertexAiSearchTool) covers the ADK 1.x version of this tool, from before Vertex AI Search was renamed Agent Search.