Milvus 搭載の Google ADK

Google Agent Development Kit (ADK)は、ツール、セッション、ランナー、メモリサービスを活用して、開発者がエージェントを構築できるよう支援します。Milvusは、類似性検索や AI メモリワークロードの組み込みを目的として構築されたオープンソースのベクトルデータベースです。

このチュートリアルでは、 adk-milvus ADKとMilvusを、知識ベース上の検索ツールセットと、ユーザー固有のエージェントメモリのためのクロスセッションメモリサービスという2つの一般的な場面で連携させます。このノートブックはデフォルトで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 を、ADKエージェントにはgemini-2.5-flash を使用しています。

ローカルの 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のパブリックエンドポイントおよび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が2つの重要なADKインターフェース（共有知識のための検索ツールと、ユーザースコープのクロスセッションコンテキストのためのメモリサービス）の背後で機能する方法を示しました。また、Geminiが回答する前にMilvusの検索ツールを呼び出す、ADK Runnerによるライブターンの実行も実施しました。 Milvus Lite を使用すれば、ノートブック上で同様の統合を簡単にプロトタイプ化できます。また、Milvus サーバーや Zilliz Cloud を使用すれば、同様の構成で大規模なチームや本番環境のエージェントワークロードに対応可能です。

重要なポイントは、ADKがエージェントインターフェースをシンプルに保ちつつ、Milvusがその裏側で永続的なベクトル検索、メタデータフィルタリング、スケーラブルなメモリストレージを処理する点にあります。