Milvus를 활용한 Google ADK
Google Agent Development Kit(ADK)는 개발자가 도구, 세션, 실행기 및 메모리 서비스를 활용해 에이전트를 구축할 수 있도록 지원합니다. Milvus는 유사도 검색 및 AI 메모리 워크로드를 임베딩하기 위해 구축된 오픈소스 벡터 데이터베이스입니다.
이 튜토리얼에서는 adk-milvus ADK와 Milvus를 두 가지 일반적인 시나리오, 즉 지식 기반에 대한 검색 툴셋과 사용자별 에이전트 메모리를 위한 세션 간 메모리 서비스에 연결해 보겠습니다. 이 노트북은 기본적으로 Milvus Lite를 사용하므로 별도의 Milvus 서버 없이도 로컬 또는 Google Colab에서 실행할 수 있습니다.
필수 조건
ADK Milvus 통합 및 Milvus 종속성을 설치하십시오.
%%capture
! pip install --upgrade adk-milvus google-genai pymilvus milvus-lite
Google Colab을 사용하는 경우, 방금 설치한 종속성을 활성화하려면 런타임을 다시 시작해야 할 수 있습니다(화면 상단의 “Runtime” 메뉴를 클릭하고 드롭다운 메뉴에서 “Restart session”을 선택하세요).
이 노트북은 임베딩과 최종 에이전트 턴 모두에 Gemini를 사용합니다. 실행하기 전에 GEMINI_API_KEY 또는 GOOGLE_API_KEY 환경 변수를 설정하십시오. 아래 예제에서는 gemini-embedding-001 를 사용하여 실제 임베딩을 생성하고, gemini-2.5-flash 를 ADK 에이전트에 사용합니다.
로컬 Milvus 작업 공간 설정
임시 워크스페이스를 생성하고, Milvus Lite 데이터베이스 파일을 정의하며, 데모용 Gemini 임베딩 함수를 준비합니다.
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
통합에 사용되는
MilvusClient의 인자에 대해서는 다음과 같습니다.
uri를 로컬 파일(예:./milvus.db)로 설정하는 것이 가장 편리한 방법이며, 이 경우 Milvus Lite가 자동으로 작동하여 이 파일에 모든 데이터를 저장합니다.- 대규모 데이터를 다루는 경우, Docker나 Kubernetes에 더 높은 성능을 발휘하는 Milvus 서버를 구축할 수 있습니다. 이 설정에서는
uri대신 서버 URI(예:http://localhost:19530)를 사용해야 합니다.- Milvus 전용 완전 관리형 클라우드 서비스인 Zilliz Cloud를 사용하려면, Zilliz Cloud의 ‘Public Endpoint ’ 및 ‘API 키 ’에 해당하는
uri와token를 조정하십시오.
Milvus를 사용하여 ADK 검색 도구 세트 구축하기
MilvusVectorStore 는 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 앱에서는 에이전트가 모델 실행 단계에서 동일한 도구를 호출할 수 있습니다.
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를 기반으로 하므로, 메타데이터 필터를 사용하여 검색 범위를 좁힐 수도 있습니다. 다음 쿼리는 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 통합이 ID, 콘텐츠, 소스 메타데이터 및 임베딩을 포함하는 일반적인 Milvus 행을 기록했음을 확인할 수 있습니다.
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.
Milvus에 ADK 메모리 저장
검색 도구는 공유 지식 기반에 유용합니다. 에이전트 메모리는 이와 다릅니다. 특정 앱과 사용자로 범위가 한정되어야 하며, 세션 간에도 유지되어야 합니다. MilvusMemoryService 는 Milvus를 기본 벡터 스토어로 사용하면서 ADK의 메모리 서비스 인터페이스를 구현합니다.
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
특정 사용자의 메모리를 검색합니다. 이 서비스는 자동으로 app_name 및 user_id 을 기준으로 필터링하므로, 다른 사용자의 이벤트가 결과 집합에 포함되지 않습니다.
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
ADK 에이전트에 Milvus 도구 연결
이전 셀에서는 검색 도구를 직접 실행하여, 모델을 활용하기 전에 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.
결론
이 노트북은 Milvus가 공유 지식을 위한 검색 도구와 사용자 범위 내의 세션 간 컨텍스트를 위한 메모리 서비스라는 두 가지 중요한 ADK 인터페이스 뒤에서 어떻게 작동할 수 있는지 보여주었습니다. 또한 Gemini가 응답하기 전에 Milvus 검색 도구를 호출하는 실시간 ADK Runner 턴을 실행했습니다. Milvus Lite를 사용하면 노트북에서 동일한 통합을 쉽게 프로토타입으로 구현할 수 있으며, Milvus 서버나 Zilliz Cloud를 사용하면 동일한 구성 구조를 통해 더 큰 규모의 팀과 프로덕션 에이전트 워크로드를 지원할 수 있습니다.
핵심 개념은 ADK가 에이전트 인터페이스를 간결하게 유지하는 동안, Milvus가 내부적으로 내구성 있는 벡터 검색, 메타데이터 필터링 및 확장 가능한 메모리 저장을 처리한다는 점입니다.
