Google ADK with Milvus
Google Agent Development Kit (ADK) 可帮助开发者使用工具、会话、Runner 和内存服务构建 AI Agent。Milvus 是一款开源向量数据库，专为 Embedding 相似性搜索和 AI 内存工作负载而构建。
本教程将使用 adk-milvus，在两个常见场景中连接 ADK 与 Milvus：为知识库提供检索工具集，以及通过跨会话内存服务为不同用户保存专属的 AI Agent 内存。Notebook 默认使用 Milvus Lite，因此无需单独部署 Milvus Server，即可在本地或 Google Colab 中运行。
Prerequisites
安装 ADK Milvus 集成及 Milvus 依赖项。
%%capture
! pip install --upgrade adk-milvus google-genai pymilvus milvus-lite
如果你使用 Google Colab，可能需要 重启运行时 才能启用刚刚安装的依赖项（点击屏幕顶部的“Runtime”菜单，然后从下拉菜单中选择“Restart session”）。
本 Notebook 使用 Gemini 生成 Embedding，并完成 Agent 的最终轮次。运行前，请准备 GEMINI_API_KEY 或 GOOGLE_API_KEY 环境变量。以下示例使用 gemini-embedding-001 生成真实的 Embedding，并使用 gemini-2.5-flash 作为 ADK Agent。
设置本地 Milvus 工作区
创建一个临时工作区，定义 Milvus Lite 数据库文件，并为本演示准备 Gemini Embedding 函数。
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}")
工作区：/tmp/google_adk_milvus_demo__btzq981 Embedding 模型：gemini-embedding-001 Embedding 维度：3072
关于集成所使用的
MilvusClient参数：
- 将
uri设置为本地文件（例如./milvus.db）是最便捷的方式，因为系统会自动使用 Milvus Lite 将所有数据存储在该文件中。- 如果数据量较大，可以使用 Docker 或 Kubernetes 部署性能更高的 Milvus Server。在此部署方式下，请将服务器 URI（例如
http://localhost:19530）用作uri。- 如果要使用 Milvus 的全托管云服务 Zilliz Cloud，请调整
uri和token，它们分别对应 Zilliz Cloud 中的 Public Endpoint 和 API key。
使用 Milvus 构建 ADK 检索工具集
MilvusVectorStore 将经过 Embedding 的文本存储在 Milvus 中，而 MilvusToolset 则将该存储公开为名为 milvus_similarity_search 的 ADK 检索工具。我们将为一个小型知识库创建索引，其中既包含相关文档，也包含无关的干扰文档。
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
接下来，从 ADK 工具集中获取工具，并直接运行 Milvus 检索工具。在使用 LLM 之前直接运行该工具，可以验证由 Milvus 支持的检索链路；在完整的 ADK 应用中，Agent 可以在模型调用期间使用同一个工具。
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.
由于底层存储由 Milvus 支持，你还可以使用元数据 Filter 缩小检索范围。以下查询检索生产环境中的 Milvus 运维信息，并将结果限定为来自 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.
你可以使用 MilvusClient 检查同一个 Milvus Lite 数据库。这可以确认 ADK 集成写入的是普通 Milvus 数据行，其中包含 ID、内容、来源元数据和 Embedding。
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.
Store ADK memory in Milvus
检索工具适合用于共享知识库。AI Agent 的 memory 则有所不同：它应限定于特定的 app 和 user，并且能够跨 session 持久保留。MilvusMemoryService 实现了 ADK 的 memory service 接口，并在底层使用 Milvus 作为向量存储。
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
搜索某个 user 的 memory。该服务会自动按 app_name 和 user_id 过滤，因此其他 user 的事件不会泄露到结果集中。
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
Attach Milvus tools to an ADK agent
前面的单元格直接执行了检索工具，以便在调用模型之前验证基于 Milvus 的工具。你也可以将同一个工具列表附加到 ADK Agent。下一个单元格通过 ADK Runner 执行一次实时 Gemini 交互，并展示模型在回答前调用 milvus_similarity_search。
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.
Conclusion
本 Notebook 展示了 Milvus 如何为 ADK 的两个重要功能提供底层支持：用于共享知识的检索工具，以及用于保存用户级跨 Session 上下文的记忆服务。此外，本 Notebook 还实际运行了一轮 ADK Runner，由 Gemini 在回答问题前调用 Milvus 检索工具。借助 Milvus Lite，你可以轻松地在 Notebook 中构建这一集成的原型；使用 Milvus Server 或 Zilliz Cloud 时，同样的配置结构可以支持更大规模的团队和生产环境中的 AI Agent 工作负载。
其核心思路是：ADK 让 AI Agent 接口保持简洁，而 Milvus 则在底层负责持久化向量检索、元数据过滤和可扩展的记忆存储。