Skip to content

Milvus 与 ADK 集成

在 ADK 中受支持Python

adk-milvus 包将 ADK Python 智能体连接到 Milvus——一个开源向量数据库。你可以通过 MilvusMemoryService 实现持久化的语义记忆,也可以通过 MilvusToolset 暴露一个 milvus_similarity_search 检索工具用于 RAG 工作流。

Milvus 可以使用 Milvus Lite 在本地运行,也可以作为自托管的 Milvus 服务器,或者作为托管的 Zilliz Cloud 部署,它们使用相同的配置字段。

使用场景

  • 智能体的语义记忆:将会话事件持久化到 Milvus 中,并在后续对话中检索相关记忆。
  • 基于私有内容的 RAG:索引文档或片段,让智能体通过工具调用检索相关上下文。
  • 本地到云端开发:使用 Milvus Lite 进行本地开发,然后只需更改 URI 和 token 即可切换到 Milvus 服务器或 Zilliz Cloud。

前提条件

  • Python 3.10 或更高版本
  • ADK for Python 和 adk-milvus
  • 一个能为每段输入文本返回一个向量的嵌入函数
  • 一个 Milvus 部署:
    • Milvus Lite,用于本地开发
    • Milvus 服务器,例如 http://localhost:19530
    • Zilliz Cloud 端点和 token

安装

pip install adk-milvus

这将安装 ADK 运行时依赖、PyMilvus 和 Milvus Lite 支持。

配置

所有部署模式均使用 MILVUS_URIMILVUS_TOKEN 进行配置:

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

# Milvus 服务器
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 仅在需要身份验证的部署(如 Zilliz Cloud)中使用。如果你使用非默认的 Milvus 数据库,请设置 MILVUS_DB_NAME

与智能体配合使用

MilvusMemoryService 插入 Runner,以持久化和搜索跨会话的记忆。

from adk_milvus import MilvusMemoryService
from google.adk.agents import Agent
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.genai import Client

genai_client = Client()

def embedding_function(texts):
    response = genai_client.models.embed_content(
        model="gemini-embedding-001",
        contents=list(texts),
    )
    return [list(embedding.values) for embedding in response.embeddings]

memory_service = MilvusMemoryService(
    embedding_function=embedding_function,
    dimension=3072,
    collection_name="adk_memory",
)

agent = Agent(
    name="memory_agent",
    model="gemini-flash-latest",
    instruction="Use memory to personalize responses when relevant.",
)

runner = Runner(
    app_name="milvus_memory_app",
    agent=agent,
    session_service=InMemorySessionService(),
    memory_service=memory_service,
)

在一次有用的会话之后,将其添加到记忆中以便后续搜索:

session = await runner.session_service.get_session(
    app_name="milvus_memory_app",
    user_id="user-1",
    session_id="session-1",
)
await memory_service.add_session_to_memory(session)

result = await memory_service.search_memory(
    app_name="milvus_memory_app",
    user_id="user-1",
    query="what did the user say about database preferences?",
)
for memory in result.memories:
    print(memory.content.parts[0].text)

使用 MilvusVectorStore 索引文本,然后通过 MilvusToolset 暴露它。

from adk_milvus import MilvusToolset
from adk_milvus import MilvusVectorStore
from adk_milvus import MilvusVectorStoreSettings
from google.adk.agents import Agent
from google.genai import Client

genai_client = Client()

def embedding_function(texts):
    response = genai_client.models.embed_content(
        model="gemini-embedding-001",
        contents=list(texts),
    )
    return [list(embedding.values) for embedding in response.embeddings]

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

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"},
    ],
)

milvus_toolset = MilvusToolset(vector_store=vector_store)
tools = await milvus_toolset.get_tools_with_prefix()

agent = Agent(
    name="rag_agent",
    model="gemini-flash-latest",
    instruction="Use retrieval context when answering questions.",
    tools=tools,
)

可用的工具和操作

RAG 工具集

工具 描述
milvus_similarity_search 在 Milvus 中搜索已索引的文本,并返回包含内容、来源、元数据和距离的匹配行。

记忆服务

方法 描述
add_session_to_memory(session) 从 ADK 会话中持久化包含文本的事件。
search_memory(app_name, user_id, query) 搜索限定在某个 ADK 应用和用户范围内的记忆。

注意事项

  • dimension 必须与嵌入模型的输出维度匹配。
  • MilvusMemoryServiceapp_nameuser_id 限定搜索范围。
  • MilvusVectorStore 会在集合不存在时自动创建,并在复用前验证现有集合的架构。
  • 对于需要更强的读写一致性行为或使用多个 Milvus 数据库的部署,可以配置集合的一致性级别和数据库名称。

相关资源