Google ADK com Milvus

O Google Agent Development Kit (ADK) ajuda os programadores a criar agentes com ferramentas, sessões, executores e serviços de memória. O Milvus é uma base de dados vetorial de código aberto concebida para integrar pesquisas de similaridade e cargas de trabalho de memória de IA.

Neste tutorial, iremos utilizar adk-milvus para ligar o ADK ao Milvus em dois contextos comuns: um conjunto de ferramentas de recuperação numa base de conhecimento e um serviço de memória entre sessões para a memória do agente específica do utilizador. O notebook utiliza o Milvus Lite por predefinição, pelo que pode ser executado localmente ou no Google Colab sem necessidade de um servidor Milvus separado.

Pré-requisitos

Instale a integração do ADK com o Milvus e as dependências do Milvus.

%%capture
! pip install --upgrade adk-milvus google-genai pymilvus milvus-lite

Se estiver a utilizar o Google Colab, para ativar as dependências que acabou de instalar, poderá ser necessário reiniciar o ambiente de execução (clique no menu «Runtime» na parte superior do ecrã e selecione «Restart session» no menu suspenso).

Este notebook utiliza o Gemini tanto para as representações como para a vez final do agente. Prepare uma variável de ambiente GEMINI_API_KEY ou GOOGLE_API_KEY antes de o executar. Os exemplos abaixo utilizam gemini-embedding-001 para gerar representações reais e gemini-2.5-flash para o agente ADK.

Configurar um espaço de trabalho Milvus local

Crie um espaço de trabalho temporário, defina os ficheiros da base de dados do Milvus Lite e prepare uma função de embedding do Gemini para a demonstração.

import os
import tempfile
import warnings
from pathlib import Path
from typing import Sequence

from adk_milvus import (
    MilvusMemoryService,
    MilvusMemoryServiceConfig,
    MilvusToolset,
    MilvusVectorStore,
    MilvusVectorStoreSettings,
)
from google.adk.agents import Agent
from google.adk.events.event import Event
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.genai import Client, types
from pymilvus import MilvusClient

work_dir = Path(tempfile.mkdtemp(prefix="google_adk_milvus_demo_"))
rag_db_path = work_dir / "adk_rag.db"
memory_db_path = work_dir / "adk_memory.db"

GOOGLE_EMBEDDING_MODEL = "gemini-embedding-001"
google_api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
if not google_api_key:
    raise RuntimeError(
        "Set GEMINI_API_KEY or GOOGLE_API_KEY before running this notebook."
    )

embedding_client = Client(api_key=google_api_key)


def google_embedding(texts: Sequence[str]) -> list[list[float]]:
    response = embedding_client.models.embed_content(
        model=GOOGLE_EMBEDDING_MODEL,
        contents=list(texts),
    )
    return [list(embedding.values) for embedding in response.embeddings]


EMBEDDING_DIMENSION = len(google_embedding(["Milvus vector database"])[0])


print(f"Workspace: {work_dir}")
print(f"Embedding model: {GOOGLE_EMBEDDING_MODEL}")
print(f"Embedding dimension: {EMBEDDING_DIMENSION}")

Workspace: /tmp/google_adk_milvus_demo__btzq981
Embedding model: gemini-embedding-001
Embedding dimension: 3072

Quanto ao argumento de MilvusClient utilizado pela integração:

• Definir o ` uri ` como um ficheiro local, por exemplo, `./milvus.db`, é o método mais conveniente, uma vez que utiliza automaticamente o Milvus Lite para armazenar todos os dados neste ficheiro.
• Se tiver uma grande quantidade de dados, pode configurar um servidor Milvus com melhor desempenho no Docker ou no Kubernetes. Nesta configuração, utilize o URI do servidor, por exemplo,http://localhost:19530, como o seu uri.
• Se pretender utilizar o Zilliz Cloud, o serviço na nuvem totalmente gerido para o Milvus, ajuste o uri e o token, que correspondem ao Endpoint Público e à Chave API no Zilliz Cloud.

Crie um conjunto de ferramentas de recuperação ADK com o Milvus

MilvusVectorStore armazena texto incorporado no Milvus, enquanto MilvusToolset expõe esse armazenamento como uma ferramenta de recuperação ADK denominada milvus_similarity_search. Iremos indexar uma pequena base de conhecimento com documentos relevantes e elementos de distração não relacionados.

RAG_COLLECTION = "google_adk_milvus_rag"

knowledge_docs = [
    {
        "id": "adk-toolset-doc",
        "source": "adk-toolset",
        "topic": "retrieval",
        "content": (
            "MilvusToolset exposes milvus_similarity_search as an ADK retrieval "
            "tool so agents can search product docs, runbooks, and other RAG content."
        ),
    },
    {
        "id": "adk-memory-doc",
        "source": "adk-memory",
        "topic": "memory",
        "content": (
            "MilvusMemoryService implements ADK BaseMemoryService and stores "
            "cross-session user memory with app_name and user_id scope."
        ),
    },
    {
        "id": "zilliz-cloud-doc",
        "source": "zilliz-cloud",
        "topic": "production",
        "content": (
            "Zilliz Cloud provides managed Milvus for production vector search, "
            "with cloud operations, backup planning, and deployment controls."
        ),
    },
    {
        "id": "milvus-lite-doc",
        "source": "milvus-lite",
        "topic": "local-development",
        "content": (
            "Milvus Lite stores vectors in a local database file and is useful "
            "for offline ADK prototypes before moving to a server or cloud deployment."
        ),
    },
    {
        "id": "latency-runbook-doc",
        "source": "operations-runbook",
        "topic": "operations",
        "content": (
            "The production runbook tracks vector search latency, index readiness, "
            "and restore steps for Milvus-backed applications."
        ),
    },
    {
        "id": "recipe-doc",
        "source": "team-recipe",
        "topic": "distractor",
        "content": "A pasta recipe uses tomato sauce, fresh basil, and slow cooking notes.",
    },
    {
        "id": "travel-doc",
        "source": "travel-plan",
        "topic": "distractor",
        "content": "The travel plan compares hotel options, train tickets, and city walks.",
    },
    {
        "id": "payroll-doc",
        "source": "payroll-note",
        "topic": "distractor",
        "content": "The payroll note explains invoice timing and monthly expense categories.",
    },
]

vector_store = MilvusVectorStore(
    embedding_function=google_embedding,
    settings=MilvusVectorStoreSettings(
        uri=str(rag_db_path),
        collection_name=RAG_COLLECTION,
        dimension=EMBEDDING_DIMENSION,
        search_top_k=4,
        consistency_level="Strong",
    ),
)

insert_result = await vector_store.add_texts_async(
    [doc["content"] for doc in knowledge_docs],
    metadatas=[
        {"source": doc["source"], "topic": doc["topic"]} for doc in knowledge_docs
    ],
    ids=[doc["id"] for doc in knowledge_docs],
)

print(insert_result)
print("Indexed sources:", ", ".join(doc["source"] for doc in knowledge_docs))

{'status': 'SUCCESS', 'inserted_count': 8}
Indexed sources: adk-toolset, adk-memory, zilliz-cloud, milvus-lite, operations-runbook, team-recipe, travel-plan, payroll-note

Agora, solicite as ferramentas ao conjunto de ferramentas ADK e execute diretamente a ferramenta de recuperação do Milvus. A execução direta da ferramenta verifica o caminho de recuperação suportado pelo Milvus antes de envolvermos o LLM; numa aplicação ADK completa, o agente pode chamar a mesma ferramenta durante um turno do modelo.

toolset = MilvusToolset(vector_store=vector_store)
tools = await toolset.get_tools_with_prefix()
print("ADK tools:", [tool.name for tool in tools])

retrieval_result = await tools[0].run_async(
    args={"query": "Which ADK tool should retrieve Milvus product docs for an agent?"},
    tool_context=None,
)

for rank, row in enumerate(retrieval_result["rows"], start=1):
    metadata = row.get("metadata") or {}
    print(f"#{rank} | source={row['source']} | topic={metadata.get('topic')}")
    print(row["content"])
    print()

assert retrieval_result["rows"], "The retrieval tool should return matching rows."
assert retrieval_result["rows"][0]["source"] == "adk-toolset"

ADK tools: ['milvus_similarity_search']


#1 | source=adk-toolset | topic=retrieval
MilvusToolset exposes milvus_similarity_search as an ADK retrieval tool so agents can search product docs, runbooks, and other RAG content.

#2 | source=milvus-lite | topic=local-development
Milvus Lite stores vectors in a local database file and is useful for offline ADK prototypes before moving to a server or cloud deployment.

#3 | source=adk-memory | topic=memory
MilvusMemoryService implements ADK BaseMemoryService and stores cross-session user memory with app_name and user_id scope.

#4 | source=operations-runbook | topic=operations
The production runbook tracks vector search latency, index readiness, and restore steps for Milvus-backed applications.

Como o repositório é suportado pelo Milvus, também é possível utilizar filtros de metadados para uma pesquisa mais restrita. A consulta seguinte procura operações de produção do Milvus e restringe os resultados à fonte Zilliz Cloud.

filtered_result = await vector_store.similarity_search_async(
    "managed cloud production Milvus operations",
    top_k=3,
    filter_expr='source == "zilliz-cloud"',
)

for rank, row in enumerate(filtered_result["rows"], start=1):
    print(f"#{rank} | source={row['source']}")
    print(row["content"])

assert filtered_result["rows"]
assert all(row["source"] == "zilliz-cloud" for row in filtered_result["rows"])

#1 | source=zilliz-cloud
Zilliz Cloud provides managed Milvus for production vector search, with cloud operations, backup planning, and deployment controls.

Podemos inspecionar a mesma base de dados do Milvus Lite com MilvusClient. Isto confirma que a integração do ADK escreveu linhas normais do Milvus contendo IDs, conteúdo, metadados de origem e embeddings.

inspection_client = MilvusClient(uri=str(rag_db_path))
stats = inspection_client.get_collection_stats(RAG_COLLECTION)
sample_rows = inspection_client.query(
    collection_name=RAG_COLLECTION,
    filter='source in ["adk-toolset", "team-recipe"]',
    output_fields=["id", "source", "content"],
    limit=4,
)
inspection_client.close()

print("Collection stats:", stats)
print("Sample rows:")
for row in sample_rows:
    print(f"- {row['id']} | {row['source']} | {row['content'][:90]}")

assert stats["row_count"] == len(knowledge_docs)

Collection stats: {'row_count': 8}
Sample rows:
- adk-toolset-doc | adk-toolset | MilvusToolset exposes milvus_similarity_search as an ADK retrieval tool so agents can sear
- recipe-doc | team-recipe | A pasta recipe uses tomato sauce, fresh basil, and slow cooking notes.

Armazenar a memória do ADK no Milvus

As ferramentas de pesquisa são úteis para bases de conhecimento partilhadas. A memória do agente é diferente: deve estar limitada a uma aplicação e a um utilizador específicos e deve ser mantida entre sessões. MilvusMemoryService implementa a interface do serviço de memória do ADK, utilizando o Milvus como armazenamento vetorial subjacente.

MEMORY_COLLECTION = "google_adk_milvus_memory"
APP_NAME = "google-adk-milvus-demo"

memory_service = MilvusMemoryService(
    embedding_function=google_embedding,
    config=MilvusMemoryServiceConfig(
        uri=str(memory_db_path),
        collection_name=MEMORY_COLLECTION,
        dimension=EMBEDDING_DIMENSION,
        search_top_k=2,
        consistency_level="Strong",
    ),
)

user_1_events = [
    Event(
        id="user-1-event-1",
        invocation_id="inv-user-1-1",
        author="user",
        timestamp=10001,
        content=types.Content(
            parts=[
                types.Part(
                    text=(
                        "Remember that I prefer Milvus Lite for local ADK memory "
                        "prototypes before using a shared server."
                    )
                )
            ]
        ),
    ),
    Event(
        id="user-1-event-2",
        invocation_id="inv-user-1-2",
        author="user",
        timestamp=10002,
        content=types.Content(
            parts=[
                types.Part(
                    text=(
                        "For production, remember that our ADK agent should use "
                        "Zilliz Cloud for managed Milvus vector memory."
                    )
                )
            ]
        ),
    ),
    Event(
        id="user-1-event-3",
        invocation_id="inv-user-1-3",
        author="user",
        timestamp=10003,
        content=types.Content(
            parts=[types.Part(text="I also like cooking noodles on Friday evenings.")]
        ),
    ),
]

user_2_events = [
    Event(
        id="user-2-event-1",
        invocation_id="inv-user-2-1",
        author="user",
        timestamp=20001,
        content=types.Content(
            parts=[
                types.Part(
                    text=(
                        "User two keeps travel planning notes and hotel preferences "
                        "in a separate ADK memory scope."
                    )
                )
            ]
        ),
    )
]

await memory_service.add_events_to_memory(
    app_name=APP_NAME,
    user_id="user-1",
    session_id="session-local-and-cloud",
    events=user_1_events,
)
await memory_service.add_events_to_memory(
    app_name=APP_NAME,
    user_id="user-2",
    session_id="session-other-user",
    events=user_2_events,
)

print("Stored memory events:", len(user_1_events) + len(user_2_events))

Stored memory events: 4

Pesquisar a memória de um utilizador. O serviço filtra automaticamente por app_name e user_id, para que os eventos de outro utilizador não apareçam no conjunto de resultados.

memory_result = await memory_service.search_memory(
    app_name=APP_NAME,
    user_id="user-1",
    query="production Milvus memory preference for my ADK agent",
)

print("User 1 memory search:")
for rank, memory in enumerate(memory_result.memories, start=1):
    print(f"#{rank} | author={memory.author} | timestamp={memory.timestamp}")
    print(memory.content.parts[0].text)
    print()

user_2_result = await memory_service.search_memory(
    app_name=APP_NAME,
    user_id="user-2",
    query="travel planning memory",
)
empty_user_result = await memory_service.search_memory(
    app_name=APP_NAME,
    user_id="user-3",
    query="production Milvus memory preference for my ADK agent",
)
wrong_app_result = await memory_service.search_memory(
    app_name="different-adk-app",
    user_id="user-1",
    query="production Milvus memory preference for my ADK agent",
)

print("User 2 scoped result:")
for memory in user_2_result.memories:
    print(memory.content.parts[0].text)
print("User 3 result count:", len(empty_user_result.memories))
print("Different app result count:", len(wrong_app_result.memories))

user_1_texts = [memory.content.parts[0].text for memory in memory_result.memories]
assert any("Zilliz Cloud" in text for text in user_1_texts)
assert all(
    "Zilliz Cloud" not in memory.content.parts[0].text
    for memory in user_2_result.memories
)
assert empty_user_result.memories == []
assert wrong_app_result.memories == []

User 1 memory search:
#1 | author=user | timestamp=1970-01-01T02:46:42
For production, remember that our ADK agent should use Zilliz Cloud for managed Milvus vector memory.

#2 | author=user | timestamp=1970-01-01T02:46:41
Remember that I prefer Milvus Lite for local ADK memory prototypes before using a shared server.



User 2 scoped result:
User two keeps travel planning notes and hotel preferences in a separate ADK memory scope.
User 3 result count: 0
Different app result count: 0

Associar ferramentas do Milvus a um agente ADK

As células anteriores executaram diretamente a ferramenta de recuperação, que verifica a ferramenta suportada pelo Milvus antes de envolver o modelo. A mesma lista de ferramentas também pode ser associada a um « Agent » do ADK. A célula seguinte executa uma iteração do Gemini em tempo real através do ADK Runner e mostra que o modelo chama « milvus_similarity_search » antes de responder.

agent = Agent(
    name="milvus_research_agent",
    model="gemini-2.5-flash",
    instruction=(
        "You are a concise assistant. Use milvus_similarity_search before "
        "answering questions about ADK, Milvus deployment, or vector memory. "
        "Mention source names from retrieved rows when useful."
    ),
    tools=tools,
)

print("Agent:", agent.name)
print("Attached tools:", [tool.name for tool in agent.tools])

model_key_available = bool(os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY"))
if not model_key_available:
    print("Set GEMINI_API_KEY or GOOGLE_API_KEY to run the live LLM turn.")
else:
    session_service = InMemorySessionService()
    llm_user_id = "user-llm"
    llm_session_id = "session-llm"
    await session_service.create_session(
        app_name=APP_NAME,
        user_id=llm_user_id,
        session_id=llm_session_id,
    )
    runner = Runner(
        app_name=APP_NAME,
        agent=agent,
        session_service=session_service,
    )
    prompt = (
        "Use the Milvus retrieval tool to answer: "
        "What does the ADK Milvus integration provide for agents?"
    )

    tool_calls = []
    tool_responses = []
    final_answer = ""
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            message=".*JSON_SCHEMA_FOR_FUNC_DECL.*",
            category=UserWarning,
        )
        async for event in runner.run_async(
            user_id=llm_user_id,
            session_id=llm_session_id,
            new_message=types.Content(
                role="user",
                parts=[types.Part(text=prompt)],
            ),
        ):
            tool_calls.extend(call.name for call in event.get_function_calls())
            tool_responses.extend(
                response.name for response in event.get_function_responses()
            )
            if event.is_final_response() and event.content and event.content.parts:
                final_answer = "".join(part.text or "" for part in event.content.parts)

    print("LLM tool calls:", tool_calls)
    print("LLM tool responses:", tool_responses)
    print("Final answer:")
    print(final_answer)

    assert "milvus_similarity_search" in tool_calls
    assert final_answer

Agent: milvus_research_agent
Attached tools: ['milvus_similarity_search']


LLM tool calls: ['milvus_similarity_search']
LLM tool responses: ['milvus_similarity_search']
Final answer:
The ADK Milvus integration provides agents with the ability to search product documentation, runbooks, and other RAG content through the `milvus_similarity_search` tool, as stated in the "adk-toolset" source. It also offers a `MilvusMemoryService` for storing cross-session user memory, as mentioned in the "adk-memory" source. For development, "milvus-lite" allows for offline prototyping by storing vectors in a local database file, and for production, "zilliz-cloud" provides managed Milvus for vector search with cloud operations and deployment controls.

await toolset.close()
await memory_service.close()
print("Milvus clients closed.")

Milvus clients closed.

Conclusão

Este notebook demonstrou como o Milvus pode estar por trás de duas importantes interfaces do ADK: ferramentas de recuperação para conhecimento partilhado e serviços de memória para contexto inter-sessões no âmbito do utilizador. Também executou uma ronda ao vivo do ADK Runner, na qual o Gemini chamou a ferramenta de recuperação do Milvus antes de responder. Com o Milvus Lite, é fácil criar um protótipo desta integração num notebook; com o servidor Milvus ou a Zilliz Cloud, a mesma configuração pode dar resposta a equipas maiores e a cargas de trabalho de agentes em produção.

A ideia principal é que o ADK mantém a interface do agente simples, enquanto o Milvus trata da pesquisa vetorial duradoura, da filtragem de metadados e do armazenamento de memória escalável em segundo plano.