From d89c24d48c931f46b31d5c37ffe5dc60a5903e66 Mon Sep 17 00:00:00 2001
From: Peter Gustafsson <peter.gustafsson6@gmail.com>
Date: Fri, 4 Jul 2025 21:59:48 +0200
Subject: [PATCH] fix(milvus): add missing files_api parameter and kvstore
 config

- Fix constructor call missing files_api parameter
- Add kvstore field to MilvusVectorIOConfig
- Resolves #2626
---
 llama_stack/providers/remote/vector_io/milvus/__init__.py | 4 +++-
 llama_stack/providers/remote/vector_io/milvus/config.py   | 2 ++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/llama_stack/providers/remote/vector_io/milvus/__init__.py b/llama_stack/providers/remote/vector_io/milvus/__init__.py
index 92dbfda2e..8f0c35c7e 100644
--- a/llama_stack/providers/remote/vector_io/milvus/__init__.py
+++ b/llama_stack/providers/remote/vector_io/milvus/__init__.py
@@ -14,6 +14,8 @@ async def get_adapter_impl(config: MilvusVectorIOConfig, deps: dict[Api, Provide
 
     assert isinstance(config, MilvusVectorIOConfig), f"Unexpected config type: {type(config)}"
 
-    impl = MilvusVectorIOAdapter(config, deps[Api.inference])
+    files_api = deps.get(Api.files, None)
+
+    impl = MilvusVectorIOAdapter(config, deps[Api.inference], files_api)
     await impl.initialize()
     return impl
diff --git a/llama_stack/providers/remote/vector_io/milvus/config.py b/llama_stack/providers/remote/vector_io/milvus/config.py
index b42233d6d..e57b8dfb6 100644
--- a/llama_stack/providers/remote/vector_io/milvus/config.py
+++ b/llama_stack/providers/remote/vector_io/milvus/config.py
@@ -8,6 +8,7 @@ from typing import Any
 
 from pydantic import BaseModel, ConfigDict, Field
 
+from llama_stack.providers.utils.kvstore.config import KVStoreConfig
 from llama_stack.schema_utils import json_schema_type
 
 
@@ -16,6 +17,7 @@ class MilvusVectorIOConfig(BaseModel):
     uri: str = Field(description="The URI of the Milvus server")
     token: str | None = Field(description="The token of the Milvus server")
     consistency_level: str = Field(description="The consistency level of the Milvus server", default="Strong")
+    kvstore: KVStoreConfig | None = None
 
     # This configuration allows additional fields to be passed through to the underlying Milvus client.
     # See the [Milvus](https://milvus.io/docs/install-overview.md) documentation for more details about Milvus in general.