feat: Add Milvus vector store integration for RAG

Add Milvus as a vector database option for knowledge base RAG workflows.
This follows the existing Spanner integration pattern with BaseToolset,
providing data ingestion, similarity search, and a ready-to-use toolset
for LLM agents.

- MilvusVectorStoreSettings / MilvusToolSettings (Pydantic config)
- MilvusVectorStore: setup, add_contents, search, close
- MilvusTool: FunctionTool subclass injecting vector_store param
- MilvusToolset: BaseToolset providing similarity_search tool
- Registered MILVUS_TOOLSET / MILVUS_VECTOR_STORE as experimental features
- Added pymilvus>=2.5.0 to test, extensions, and new milvus extras
- 30 unit tests (all passing)
- Sample milvus_rag_agent under contributing/samples/

Author: zc277584121

contributing/samples/milvus_rag_agent/README.md

@@ -0,0 +1,67 @@
+# Milvus RAG Agent Sample
+
+This sample demonstrates how to build a knowledge base agent using
+[Milvus](https://milvus.io/) as the vector database for
+retrieval-augmented generation (RAG) with Google ADK.
+
+## Prerequisites
+
+1. A running Milvus instance, or use Milvus Lite (local file path).
+2. A Google GenAI API key for embedding generation.
+3. Install dependencies:
+
+```bash
+pip install "google-adk[milvus]"
+```
+
+## Setup
+
+### Environment Variables
+
+```bash
+export GOOGLE_API_KEY="your-google-api-key"
+export MILVUS_URI="http://localhost:19530"  # or ./milvus.db for Milvus Lite
+export MILVUS_COLLECTION="knowledge_base"
+# Optional: export MILVUS_TOKEN="your-token"  # for Zilliz Cloud
+```
+
+### Data Ingestion
+
+Before running the agent, you need to populate the Milvus collection
+with your knowledge base data:
+
+```python
+from google.adk.tools.milvus.milvus_vector_store import MilvusVectorStore
+from google.adk.tools.milvus.settings import MilvusToolSettings
+from google.adk.tools.milvus.settings import MilvusVectorStoreSettings
+from google.genai import Client
+
+# Define your embedding function (example using Google GenAI).
+genai_client = Client()
+
+def embedding_fn(texts):
+    resp = genai_client.models.embed_content(
+        model="text-embedding-004", contents=texts)
+    return [list(e.values) for e in resp.embeddings]
+
+settings = MilvusToolSettings(
+    vector_store_settings=MilvusVectorStoreSettings(
+        uri="http://localhost:19530",
+        collection_name="knowledge_base",
+        dimension=768,
+    ),
+)
+
+store = MilvusVectorStore(settings=settings, embedding_fn=embedding_fn)
+store.setup()
+store.add_contents([
+    "Your document text here...",
+    "Another document...",
+])
+```
+
+## Run
+
+```bash
+adk run contributing/samples/milvus_rag_agent
+```

contributing/samples/milvus_rag_agent/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from . import agent

contributing/samples/milvus_rag_agent/agent.py

@@ -0,0 +1,85 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Sample Milvus RAG Agent.
+
+This sample demonstrates how to build a knowledge base agent using Milvus
+as the vector database for retrieval-augmented generation (RAG).
+
+Prerequisites:
+  1. A running Milvus instance (or use Milvus Lite with a local file path).
+  2. A Google GenAI API key set in the GOOGLE_API_KEY environment variable.
+  3. Install dependencies: pip install "google-adk[milvus]"
+
+Usage:
+  adk run contributing/samples/milvus_rag_agent
+"""
+
+import os
+
+from dotenv import load_dotenv
+from google.adk.agents.llm_agent import LlmAgent
+from google.adk.tools.milvus.milvus_toolset import MilvusToolset
+from google.adk.tools.milvus.settings import MilvusToolSettings
+from google.adk.tools.milvus.settings import MilvusVectorStoreSettings
+from google.genai import Client
+
+load_dotenv()
+
+# --- Embedding function using Google GenAI ---
+genai_client = Client()
+
+
+def embedding_fn(texts: list[str]) -> list[list[float]]:
+  response = genai_client.models.embed_content(
+      model="text-embedding-004",
+      contents=texts,
+  )
+  return [list(e.values) for e in response.embeddings]
+
+
+# --- Milvus vector store settings ---
+# Replace these with your own Milvus connection and collection settings.
+vector_store_settings = MilvusVectorStoreSettings(
+    # Use a remote Milvus instance or a local Milvus Lite file path.
+    uri=os.environ.get("MILVUS_URI", "http://localhost:19530"),
+    token=os.environ.get("MILVUS_TOKEN", None),
+    collection_name=os.environ.get("MILVUS_COLLECTION", "knowledge_base"),
+    dimension=768,
+    metric_type="COSINE",
+    top_k=5,
+)
+
+# --- Milvus toolset ---
+milvus_toolset = MilvusToolset(
+    milvus_tool_settings=MilvusToolSettings(
+        vector_store_settings=vector_store_settings,
+    ),
+    embedding_fn=embedding_fn,
+    tool_filter=["similarity_search"],
+)
+
+# --- Agent definition ---
+root_agent = LlmAgent(
+    model="gemini-2.5-flash",
+    name="milvus_knowledge_agent",
+    description="Agent that answers questions using a Milvus knowledge base.",
+    instruction="""
+    You are a helpful assistant with access to a knowledge base.
+    1. Always use the `similarity_search` tool to find relevant information.
+    2. Present the search results naturally in your response.
+    3. If no results are found, say you don't know.
+    """,
+    tools=[milvus_toolset],
+)

pyproject.toml

@@ -127,6 +127,7 @@ test = [
   "litellm>=1.75.5, <1.80.17",                    # For LiteLLM tests
   "llama-index-readers-file>=0.4.0",              # For retrieval tests
   "openai>=1.100.2",                              # For LiteLLM
+  "pymilvus>=2.5.0",                              # For Milvus vector store tests
   "opentelemetry-instrumentation-google-genai>=0.3b0, <1.0.0",
   "pypika>=0.50.0",                               # For crewai->chromadb dependency
   "pytest-asyncio>=0.25.0",
@@ -159,11 +160,16 @@ extensions = [
   "litellm>=1.75.5, <1.80.17",                  # For LiteLlm class. Currently has OpenAI limitations. TODO: once LiteLlm fix it
   "llama-index-readers-file>=0.4.0",            # For retrieval using LlamaIndex.
   "llama-index-embeddings-google-genai>=0.3.0", # For files retrieval using LlamaIndex.
+  "pymilvus>=2.5.0",                            # For Milvus vector store integration.
   "lxml>=5.3.0",                                # For load_web_page tool.
   "pypika>=0.50.0",                             # For crewai->chromadb dependency
   "toolbox-adk>=0.5.7, <0.6.0",                 # For tools.toolbox_toolset.ToolboxToolset
 ]
 
+milvus = [
+  "pymilvus>=2.5.0",
+]
+
 otel-gcp = ["opentelemetry-instrumentation-google-genai>=0.6b0, <1.0.0"]
 
 toolbox = ["toolbox-adk>=0.5.7, <0.6.0"]

src/google/adk/features/_feature_registry.py

@@ -38,6 +38,8 @@ class FeatureName(str, Enum):
   GOOGLE_CREDENTIALS_CONFIG = "GOOGLE_CREDENTIALS_CONFIG"
   GOOGLE_TOOL = "GOOGLE_TOOL"
   JSON_SCHEMA_FOR_FUNC_DECL = "JSON_SCHEMA_FOR_FUNC_DECL"
+  MILVUS_TOOLSET = "MILVUS_TOOLSET"
+  MILVUS_VECTOR_STORE = "MILVUS_VECTOR_STORE"
   PROGRESSIVE_SSE_STREAMING = "PROGRESSIVE_SSE_STREAMING"
   PUBSUB_TOOL_CONFIG = "PUBSUB_TOOL_CONFIG"
   PUBSUB_TOOLSET = "PUBSUB_TOOLSET"
@@ -114,6 +116,12 @@ class FeatureConfig:
     FeatureName.JSON_SCHEMA_FOR_FUNC_DECL: FeatureConfig(
         FeatureStage.WIP, default_on=False
     ),
+    FeatureName.MILVUS_TOOLSET: FeatureConfig(
+        FeatureStage.EXPERIMENTAL, default_on=True
+    ),
+    FeatureName.MILVUS_VECTOR_STORE: FeatureConfig(
+        FeatureStage.EXPERIMENTAL, default_on=True
+    ),
     FeatureName.PROGRESSIVE_SSE_STREAMING: FeatureConfig(
         FeatureStage.EXPERIMENTAL, default_on=True
     ),

src/google/adk/tools/milvus/__init__.py

@@ -0,0 +1,29 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Milvus Tools (Experimental).
+
+Milvus Tools provide vector similarity search capabilities using Milvus
+as the vector database backend. This module offers:
+
+1. A MilvusToolset for easy integration with ADK agents.
+2. A MilvusVectorStore utility for collection management and data ingestion.
+3. Support for any third-party embedding function (Google GenAI, OpenAI, etc.).
+"""
+
+from .milvus_toolset import MilvusToolset
+
+__all__ = [
+    "MilvusToolset",
+]

src/google/adk/tools/milvus/milvus_tool.py

@@ -0,0 +1,55 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""MilvusTool wraps a function and injects the MilvusVectorStore instance."""
+
+from __future__ import annotations
+
+import inspect
+from typing import Any
+from typing import Callable
+
+from typing_extensions import override
+
+from ..function_tool import FunctionTool
+from ..tool_context import ToolContext
+from .milvus_vector_store import MilvusVectorStore
+
+
+class MilvusTool(FunctionTool):
+  """A FunctionTool that injects MilvusVectorStore into the wrapped function.
+
+  The ``vector_store`` parameter is hidden from the LLM function
+  declaration and automatically injected at runtime.
+  """
+
+  def __init__(
+      self,
+      func: Callable[..., Any],
+      *,
+      vector_store: MilvusVectorStore,
+  ):
+    super().__init__(func=func)
+    self._ignore_params.append("vector_store")
+    self._vector_store = vector_store
+
+  @override
+  async def run_async(
+      self, *, args: dict[str, Any], tool_context: ToolContext
+  ) -> Any:
+    args_to_call = args.copy()
+    signature = inspect.signature(self.func)
+    if "vector_store" in signature.parameters:
+      args_to_call["vector_store"] = self._vector_store
+    return await super().run_async(args=args_to_call, tool_context=tool_context)