feat: Add Milvus vector store integration for RAG tool and memory

contributing/samples/milvus_rag_agent/README.md

@@ -0,0 +1,86 @@
+# 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]"
+```
+
+## Supported Milvus Backends
+
+The integration works with all three Milvus deployment modes by changing
+the `uri` (and optionally `token`) setting:
+
+| Backend | `uri` | `token` |
+|---------|-------|---------|
+| **Milvus Lite** (local, no server needed) | `./milvus.db` | — |
+| **Milvus Server** (self-hosted) | `http://localhost:19530` | — |
+| **Zilliz Cloud** (fully managed) | `https://in01-xxx.serverless.gcp-us-west1.cloud.zilliz.com` | `your-api-key` |
+
+> For Milvus Lite, install the extra package: `pip install "pymilvus[milvus_lite]"`.
+
+## Setup
+
+### Environment Variables
+
+```bash
+export GOOGLE_API_KEY="your-google-api-key"
+
+# Pick one of the following:
+export MILVUS_URI="./milvus.db"                # Milvus Lite
+export MILVUS_URI="http://localhost:19530"      # Milvus Server
+export MILVUS_URI="https://in01-xxx.cloud.zilliz.com"  # Zilliz Cloud
+
+export MILVUS_COLLECTION="knowledge_base"
+# Required for Zilliz Cloud only:
+# export MILVUS_TOKEN="your-api-key"
+```
+
+### 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,9 @@ 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_MEMORY_SERVICE = "MILVUS_MEMORY_SERVICE"
+  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 +117,15 @@ class FeatureConfig:
     FeatureName.JSON_SCHEMA_FOR_FUNC_DECL: FeatureConfig(
         FeatureStage.WIP, default_on=False
     ),
+    FeatureName.MILVUS_MEMORY_SERVICE: FeatureConfig(
+        FeatureStage.EXPERIMENTAL, default_on=True
+    ),
+    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/memory/__init__.py

@@ -35,3 +35,14 @@
       ' VertexAiRagMemoryService please install it. If not, you can ignore this'
       ' warning.'
   )
+
+try:
+  from .milvus_memory_service import MilvusMemoryService
+
+  __all__.append('MilvusMemoryService')
+except ImportError:
+  logger.debug(
+      'pymilvus is not installed. If you want to use the'
+      ' MilvusMemoryService please install it with:'
+      ' pip install "google-adk[milvus]". If not, you can ignore this warning.'
+  )