adk-milvus 0.1.0

Milvus memory and retrieval integrations for Google ADK

ADK Milvus

Milvus integrations for Google Agent Development Kit (ADK).

This package provides:

• MilvusMemoryService: a Milvus-backed ADK BaseMemoryService for cross-session memory.
• MilvusVectorStore and MilvusToolset: a Milvus-backed retrieval toolset for RAG-style agent tools.

The same configuration shape supports Milvus Lite, Milvus server, and Zilliz Cloud. New projects default to Milvus Lite, so local development requires no separate database service.

Installation

Install with pip:

pip install adk-milvus

Install with uv:

uv add adk-milvus

Configuration

Use MILVUS_URI and MILVUS_TOKEN for all deployment modes:

# Milvus Lite
export MILVUS_URI="./adk_milvus.db"

# Milvus server
export MILVUS_URI="http://localhost:19530"

# Zilliz Cloud
export MILVUS_URI="https://your-endpoint.api.gcp-us-west1.zillizcloud.com"
export MILVUS_TOKEN="your-token"

MILVUS_TOKEN is only needed for authenticated deployments such as Zilliz Cloud. If you use a non-default Milvus database, set MILVUS_DB_NAME.

Memory Service

MilvusMemoryService accepts any embedding function that returns one vector per input text. It implements ADK's BaseMemoryService, so it can be passed to an ADK Runner or used directly with add_events_to_memory() / search_memory().

from adk_milvus import MilvusMemoryService


memory_service = MilvusMemoryService(
    embedding_function=embedding_function,
    dimension=1536,
)

Run the complete OpenAI embedding example:

export OPENAI_API_KEY="..."
python examples/memory_service_openai.py

The memory service scopes search by app_name and user_id, so one user cannot retrieve another user's memories through this service.

Retrieval Toolset

MilvusVectorStore stores indexed text. MilvusToolset exposes the milvus_similarity_search tool that an ADK agent can call.

from adk_milvus import MilvusToolset
from adk_milvus import MilvusVectorStore
from adk_milvus import MilvusVectorStoreSettings


vector_store = MilvusVectorStore(
    embedding_function=embedding_function,
    settings=MilvusVectorStoreSettings(
        collection_name="adk_rag",
        dimension=1536,
    ),
)

vector_store.add_texts(
    [
        "Milvus Lite is useful for local RAG development.",
        "Zilliz Cloud provides managed Milvus for production workloads.",
    ],
    metadatas=[
        {"source": "milvus-lite"},
        {"source": "zilliz-cloud"},
    ],
)

toolset = MilvusToolset(vector_store=vector_store)

Run the complete OpenAI embedding example:

export OPENAI_API_KEY="..."
python examples/retrieval_toolset_openai.py

The retrieval tool returns:

{
    "status": "SUCCESS",
    "rows": [
        {
            "id": "...",
            "content": "...",
            "source": "...",
            "metadata": {...},
            "distance": 0.12,
        }
    ],
}

End-to-End Usage

The snippets below assume embedding_function returns 1536-dimensional vectors.

Case 1: Cross-session memory

Use MilvusMemoryService when the agent needs to remember user-specific context across sessions:

import asyncio

from adk_milvus import MilvusMemoryService
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 types


async def main():
    app_name = "milvus_memory_app"
    memory_service = MilvusMemoryService(
        embedding_function=embedding_function,
        dimension=1536,
        collection_name="adk_memory",
        consistency_level="Strong",
    )

    runner = Runner(
        app_name=app_name,
        agent=Agent(
            name="memory_agent",
            model="gemini-2.5-flash",
            instruction="Use memory when it is relevant to the user's request.",
        ),
        session_service=InMemorySessionService(),
        memory_service=memory_service,
    )

    try:
        await memory_service.add_events_to_memory(
            app_name=app_name,
            user_id="user-1",
            session_id="session-1",
            events=[
                Event(
                    id="event-1",
                    invocation_id="invocation-1",
                    author="user",
                    timestamp=12345,
                    content=types.Content(
                        parts=[
                            types.Part(
                                text="The user prefers Milvus for vector memory."
                            )
                        ]
                    ),
                )
            ],
        )

        result = await memory_service.search_memory(
            app_name=app_name,
            user_id="user-1",
            query="database preference",
        )
        print(result.memories[0].content.parts[0].text)
    finally:
        await memory_service.close()


asyncio.run(main())

Case 2: Knowledge-base retrieval

Use MilvusToolset when the agent needs a retrieval tool over external documents or product knowledge:

import asyncio

from adk_milvus import MilvusToolset, MilvusVectorStore, MilvusVectorStoreSettings
from google.adk.agents import Agent


async def main():
    vector_store = MilvusVectorStore(
        embedding_function=embedding_function,
        settings=MilvusVectorStoreSettings(
            collection_name="adk_rag",
            dimension=1536,
            consistency_level="Strong",
        ),
    )
    toolset = MilvusToolset(vector_store=vector_store)

    try:
        await vector_store.add_texts_async(
            [
                "Milvus Lite is useful for local RAG development.",
                "Zilliz Cloud provides managed Milvus for production workloads.",
            ],
            metadatas=[
                {"source": "milvus-lite"},
                {"source": "zilliz-cloud"},
            ],
        )

        tools = await toolset.get_tools_with_prefix()
        agent = Agent(
            name="rag_agent",
            model="gemini-2.5-flash",
            instruction="Use retrieval tools before answering knowledge questions.",
            tools=tools,
        )

        result = await tools[0].run_async(
            args={"query": "managed Milvus for production"},
            tool_context=None,
        )
        print(result["rows"][0]["content"])
    finally:
        await toolset.close()


asyncio.run(main())

See examples/memory_service_openai.py and examples/retrieval_toolset_openai.py for runnable scripts that exercise each component.

Contributing

See CONTRIBUTING.md for local development, testing, and release checks.