Google ADK avec Milvus

Le Google Agent Development Kit (ADK) aide les développeurs à créer des agents à l'aide d'outils, de sessions, de runners et de services de mémoire. Milvus est une base de données vectorielle open source conçue pour intégrer la recherche par similarité et les charges de travail de mémoire liées à l'IA.

Dans ce tutoriel, nous allons utiliser adk-milvus pour connecter l’ADK à Milvus dans deux cas de figure courants : un ensemble d’outils de recherche sur une base de connaissances et un service de mémoire intersessionnel pour la mémoire d’agent spécifique à l’utilisateur. Le notebook utilise Milvus Lite par défaut, ce qui permet de l’exécuter localement ou dans Google Colab sans avoir besoin d’un serveur Milvus distinct.

Prérequis

Installez l’intégration ADK Milvus et les dépendances Milvus.

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

Si vous utilisez Google Colab, pour activer les dépendances qui viennent d’être installées, vous devrez peut-être redémarrer le runtime (cliquez sur le menu « Runtime » en haut de l’écran, puis sélectionnez « Redémarrer la session » dans le menu déroulant).

Ce notebook utilise Gemini à la fois pour les représentations vectorielles et pour le tour final de l’agent. Préparez une variable d’environnement GEMINI_API_KEY ou GOOGLE_API_KEY avant de l’exécuter. Les exemples ci-dessous utilisent gemini-embedding-001 pour générer des représentations vectorielles réelles et gemini-2.5-flash pour l’agent ADK.

Configurer un espace de travail Milvus local

Créez un espace de travail temporaire, définissez les fichiers de base de données Milvus Lite et préparez une fonction d’embedding Gemini pour la démo.

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

En ce qui concerne l’argument de MilvusClient utilisé par l’intégration :

• Définir l’argument « uri » comme un fichier local, par exemple./milvus.db, est la méthode la plus pratique, car cela utilise automatiquement Milvus Lite pour stocker toutes les données dans ce fichier.
• Si vous disposez d’un volume important de données, vous pouvez configurer un serveur Milvus plus performant sur Docker ou Kubernetes. Dans cette configuration, veuillez utiliser l’URI du serveur, par exemplehttp://localhost:19530, comme uri.
• Si vous souhaitez utiliser Zilliz Cloud, le service cloud entièrement géré pour Milvus, modifiez les paramètres « uri » et « token », qui correspondent respectivement au point de terminaison public et à la clé API dans Zilliz Cloud.

Créer un ensemble d’outils de récupération ADK avec Milvus

MilvusVectorStore stores stocke du texte intégré dans Milvus, tandis que MilvusToolset expose ce stockage sous la forme d’un outil de recherche ADK nommé milvus_similarity_search. Nous allons indexer une petite base de connaissances contenant à la fois des documents pertinents et des éléments de distraction sans rapport.

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

Demandez maintenant des outils à la suite ADK et exécutez directement l’outil de recherche Milvus. L’exécution directe de l’outil permet de vérifier le chemin de recherche basé sur Milvus avant d’impliquer le LLM ; dans une application ADK complète, l’agent peut appeler le même outil lors d’un tour de modèle.

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.

Comme le magasin s’appuie sur Milvus, vous pouvez également utiliser des filtres de métadonnées pour affiner la recherche. La requête suivante recherche les opérations Milvus en production et limite les résultats à la source 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.

Nous pouvons inspecter la même base de données Milvus Lite à l’aide de ` MilvusClient`. Cela confirme que l’intégration ADK a écrit des lignes Milvus standard contenant des identifiants, du contenu, des métadonnées de source et des représentations vectorielles.

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.

Stockage de la mémoire ADK dans Milvus

Les outils de recherche sont utiles pour les bases de connaissances partagées. La mémoire de l’agent est différente : elle doit être limitée à une application et à un utilisateur spécifiques, et elle doit persister d’une session à l’autre. MilvusMemoryService implémente l’interface du service de mémoire de l’ADK tout en utilisant Milvus comme magasin vectoriel sous-jacent.

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

Recherche dans la mémoire d’un utilisateur. Le service effectue automatiquement un filtrage par app_name et user_id, afin que les événements d’un autre utilisateur n’apparaissent pas dans le jeu de résultats.

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

Associer des outils Milvus à un agent ADK

Les cellules précédentes exécutaient directement l’outil de récupération, ce qui permet de vérifier le bon fonctionnement de l’outil s’appuyant sur Milvus avant d’impliquer le modèle. La même liste d’outils peut également être associée à une Agent ADK. La cellule suivante exécute un tour Gemini en direct via ADK Runner et montre que le modèle appelle milvus_similarity_search avant de répondre.

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

Ce notebook a montré comment Milvus peut s’intégrer derrière deux interfaces ADK importantes : les outils de recherche pour les connaissances partagées et les services de mémoire pour le contexte inter-sessions au niveau de l’utilisateur. Il a également exécuté un tour ADK Runner en direct au cours duquel Gemini a appelé l’outil de recherche Milvus avant de répondre. Avec Milvus Lite, cette intégration est facile à prototyper dans un notebook ; avec le serveur Milvus ou Zilliz Cloud, cette même architecture de configuration peut prendre en charge des équipes plus importantes et des charges de travail d’agents en production.

L’idée principale est que l’ADK maintient l’interface de l’agent épurée tandis que Milvus gère en arrière-plan la recherche vectorielle durable, le filtrage des métadonnées et le stockage en mémoire évolutif.