adding config to providers so that it can properly be used

Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>

Author: franciscojavierarceo

src/llama_stack/core/resolver.py

@@ -374,6 +374,13 @@ async def instantiate_provider(
         method = "get_adapter_impl"
         args = [config, deps]
 
+        # Add vector_stores_config for vector_io providers
+        if (
+            "vector_stores_config" in inspect.signature(getattr(module, method)).parameters
+            and provider_spec.api == Api.vector_io
+        ):
+            args.append(run_config.vector_stores)
+
     elif isinstance(provider_spec, AutoRoutedProviderSpec):
         method = "get_auto_router_impl"
 
@@ -394,6 +401,11 @@ async def instantiate_provider(
             args.append(policy)
         if "telemetry_enabled" in inspect.signature(getattr(module, method)).parameters and run_config.telemetry:
             args.append(run_config.telemetry.enabled)
+        if (
+            "vector_stores_config" in inspect.signature(getattr(module, method)).parameters
+            and provider_spec.api == Api.vector_io
+        ):
+            args.append(run_config.vector_stores)
 
     fn = getattr(module, method)
     impl = await fn(*args)

src/llama_stack/core/routers/vector_io.py

@@ -103,6 +103,13 @@ async def query_chunks(
         # Ensure params dict exists and add vector_stores_config for query rewriting
         if params is None:
             params = {}
+
+        logger.debug(f"Router vector_stores_config: {self.vector_stores_config}")
+        if self.vector_stores_config and hasattr(self.vector_stores_config, "default_query_expansion_model"):
+            logger.debug(
+                f"Router default_query_expansion_model: {self.vector_stores_config.default_query_expansion_model}"
+            )
+
         params["vector_stores_config"] = self.vector_stores_config
 
         return await provider.query_chunks(vector_store_id, query, params)

src/llama_stack/providers/inline/vector_io/faiss/__init__.py

@@ -6,16 +6,19 @@
 
 from typing import Any
 
+from llama_stack.core.datatypes import VectorStoresConfig
 from llama_stack_api import Api
 
 from .config import FaissVectorIOConfig
 
 
-async def get_provider_impl(config: FaissVectorIOConfig, deps: dict[Api, Any]):
+async def get_provider_impl(
+    config: FaissVectorIOConfig, deps: dict[Api, Any], vector_stores_config: VectorStoresConfig | None = None
+):
     from .faiss import FaissVectorIOAdapter
 
     assert isinstance(config, FaissVectorIOConfig), f"Unexpected config type: {type(config)}"
 
-    impl = FaissVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
+    impl = FaissVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files), vector_stores_config)
     await impl.initialize()
     return impl

src/llama_stack/providers/inline/vector_io/faiss/faiss.py

@@ -14,8 +14,9 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from llama_stack.core.storage.kvstore import kvstore_impl
+from llama_stack.core.datatypes import VectorStoresConfig
 from llama_stack.log import get_logger
+from llama_stack.providers.utils.kvstore import kvstore_impl
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
 from llama_stack_api import (
@@ -31,7 +32,7 @@
     VectorStoreNotFoundError,
     VectorStoresProtocolPrivate,
 )
-from llama_stack_api.internal.kvstore import KVStore
+from llama_stack.providers.utils.kvstore.api import KVStore
 
 from .config import FaissVectorIOConfig
 
@@ -184,10 +185,17 @@ async def query_hybrid(
 
 
 class FaissVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtocolPrivate):
-    def __init__(self, config: FaissVectorIOConfig, inference_api: Inference, files_api: Files | None) -> None:
+    def __init__(
+        self,
+        config: FaissVectorIOConfig,
+        inference_api: Inference,
+        files_api: Files | None,
+        vector_stores_config: VectorStoresConfig | None = None,
+    ) -> None:
         super().__init__(files_api=files_api, kvstore=None)
         self.config = config
         self.inference_api = inference_api
+        self.vector_stores_config = vector_stores_config
         self.cache: dict[str, VectorStoreWithIndex] = {}
 
     async def initialize(self) -> None:
@@ -203,6 +211,7 @@ async def initialize(self) -> None:
                 vector_store,
                 await FaissIndex.create(vector_store.embedding_dimension, self.kvstore, vector_store.identifier),
                 self.inference_api,
+                self.vector_stores_config,
             )
             self.cache[vector_store.identifier] = index
 
@@ -241,6 +250,7 @@ async def register_vector_store(self, vector_store: VectorStore) -> None:
             vector_store=vector_store,
             index=await FaissIndex.create(vector_store.embedding_dimension, self.kvstore, vector_store.identifier),
             inference_api=self.inference_api,
+            vector_stores_config=self.vector_stores_config,
         )
 
     async def list_vector_stores(self) -> list[VectorStore]:
@@ -274,6 +284,7 @@ async def _get_and_cache_vector_store_index(self, vector_store_id: str) -> Vecto
             vector_store=vector_store,
             index=await FaissIndex.create(vector_store.embedding_dimension, self.kvstore, vector_store.identifier),
             inference_api=self.inference_api,
+            vector_stores_config=self.vector_stores_config,
         )
         self.cache[vector_store_id] = index
         return index

src/llama_stack/providers/inline/vector_io/sqlite_vec/__init__.py

@@ -6,15 +6,18 @@
 
 from typing import Any
 
+from llama_stack.core.datatypes import VectorStoresConfig
 from llama_stack_api import Api
 
 from .config import SQLiteVectorIOConfig
 
 
-async def get_provider_impl(config: SQLiteVectorIOConfig, deps: dict[Api, Any]):
+async def get_provider_impl(
+    config: SQLiteVectorIOConfig, deps: dict[Api, Any], vector_stores_config: VectorStoresConfig | None = None
+):
     from .sqlite_vec import SQLiteVecVectorIOAdapter
 
     assert isinstance(config, SQLiteVectorIOConfig), f"Unexpected config type: {type(config)}"
-    impl = SQLiteVecVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
+    impl = SQLiteVecVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files), vector_stores_config)
     await impl.initialize()
     return impl

src/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py

@@ -14,8 +14,9 @@
 import sqlite_vec  # type: ignore[import-untyped]
 from numpy.typing import NDArray
 
-from llama_stack.core.storage.kvstore import kvstore_impl
+from llama_stack.core.datatypes import VectorStoresConfig
 from llama_stack.log import get_logger
+from llama_stack.providers.utils.kvstore import kvstore_impl
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import (
     RERANKER_TYPE_RRF,
@@ -34,7 +35,7 @@
     VectorStoreNotFoundError,
     VectorStoresProtocolPrivate,
 )
-from llama_stack_api.internal.kvstore import KVStore
+from llama_stack.providers.utils.kvstore.api import KVStore
 
 logger = get_logger(name=__name__, category="vector_io")
 
@@ -385,10 +386,17 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresPro
     and creates a cache of VectorStoreWithIndex instances (each wrapping a SQLiteVecIndex).
     """
 
-    def __init__(self, config, inference_api: Inference, files_api: Files | None) -> None:
+    def __init__(
+        self,
+        config,
+        inference_api: Inference,
+        files_api: Files | None,
+        vector_stores_config: VectorStoresConfig | None = None,
+    ) -> None:
         super().__init__(files_api=files_api, kvstore=None)
         self.config = config
         self.inference_api = inference_api
+        self.vector_stores_config = vector_stores_config
         self.cache: dict[str, VectorStoreWithIndex] = {}
         self.vector_store_table = None
 
@@ -403,7 +411,9 @@ async def initialize(self) -> None:
             index = await SQLiteVecIndex.create(
                 vector_store.embedding_dimension, self.config.db_path, vector_store.identifier
             )
-            self.cache[vector_store.identifier] = VectorStoreWithIndex(vector_store, index, self.inference_api)
+            self.cache[vector_store.identifier] = VectorStoreWithIndex(
+                vector_store, index, self.inference_api, self.vector_stores_config
+            )
 
         # Load existing OpenAI vector stores into the in-memory cache
         await self.initialize_openai_vector_stores()
@@ -427,7 +437,9 @@ async def register_vector_store(self, vector_store: VectorStore) -> None:
         index = await SQLiteVecIndex.create(
             vector_store.embedding_dimension, self.config.db_path, vector_store.identifier
         )
-        self.cache[vector_store.identifier] = VectorStoreWithIndex(vector_store, index, self.inference_api)
+        self.cache[vector_store.identifier] = VectorStoreWithIndex(
+            vector_store, index, self.inference_api, self.vector_stores_config
+        )
 
     async def _get_and_cache_vector_store_index(self, vector_store_id: str) -> VectorStoreWithIndex | None:
         if vector_store_id in self.cache:
@@ -452,6 +464,7 @@ async def _get_and_cache_vector_store_index(self, vector_store_id: str) -> Vecto
                 kvstore=self.kvstore,
             ),
             inference_api=self.inference_api,
+            vector_stores_config=self.vector_stores_config,
         )
         self.cache[vector_store_id] = index
         return index

src/llama_stack/providers/remote/vector_io/chroma/__init__.py

@@ -4,14 +4,17 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+from llama_stack.core.datatypes import VectorStoresConfig
 from llama_stack_api import Api, ProviderSpec
 
 from .config import ChromaVectorIOConfig
 
 
-async def get_adapter_impl(config: ChromaVectorIOConfig, deps: dict[Api, ProviderSpec]):
+async def get_adapter_impl(
+    config: ChromaVectorIOConfig, deps: dict[Api, ProviderSpec], vector_stores_config: VectorStoresConfig | None = None
+):
     from .chroma import ChromaVectorIOAdapter
 
-    impl = ChromaVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
+    impl = ChromaVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files), vector_stores_config)
     await impl.initialize()
     return impl

src/llama_stack/providers/remote/vector_io/chroma/chroma.py

@@ -11,8 +11,9 @@
 import chromadb
 from numpy.typing import NDArray
 
-from llama_stack.core.storage.kvstore import kvstore_impl
+from llama_stack.core.datatypes import VectorStoresConfig
 from llama_stack.log import get_logger
+from llama_stack.providers.utils.kvstore import kvstore_impl
 from llama_stack.providers.inline.vector_io.chroma import ChromaVectorIOConfig as InlineChromaVectorIOConfig
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import ChunkForDeletion, EmbeddingIndex, VectorStoreWithIndex
@@ -26,7 +27,7 @@
     VectorStore,
     VectorStoresProtocolPrivate,
 )
-from llama_stack_api.internal.kvstore import KVStore
+from llama_stack.providers.utils.kvstore.api import KVStore
 
 from .config import ChromaVectorIOConfig as RemoteChromaVectorIOConfig
 
@@ -125,11 +126,13 @@ def __init__(
         config: RemoteChromaVectorIOConfig | InlineChromaVectorIOConfig,
         inference_api: Inference,
         files_api: Files | None,
+        vector_stores_config: VectorStoresConfig | None = None,
     ) -> None:
         super().__init__(files_api=files_api, kvstore=None)
         log.info(f"Initializing ChromaVectorIOAdapter with url: {config}")
         self.config = config
         self.inference_api = inference_api
+        self.vector_stores_config = vector_stores_config
         self.client = None
         self.cache = {}
         self.vector_store_table = None
@@ -162,7 +165,7 @@ async def register_vector_store(self, vector_store: VectorStore) -> None:
             )
         )
         self.cache[vector_store.identifier] = VectorStoreWithIndex(
-            vector_store, ChromaIndex(self.client, collection), self.inference_api
+            vector_store, ChromaIndex(self.client, collection), self.inference_api, self.vector_stores_config
         )
 
     async def unregister_vector_store(self, vector_store_id: str) -> None:
@@ -207,7 +210,9 @@ async def _get_and_cache_vector_store_index(self, vector_store_id: str) -> Vecto
         collection = await maybe_await(self.client.get_collection(vector_store_id))
         if not collection:
             raise ValueError(f"Vector DB {vector_store_id} not found in Chroma")
-        index = VectorStoreWithIndex(vector_store, ChromaIndex(self.client, collection), self.inference_api)
+        index = VectorStoreWithIndex(
+            vector_store, ChromaIndex(self.client, collection), self.inference_api, self.vector_stores_config
+        )
         self.cache[vector_store_id] = index
         return index
 

src/llama_stack/providers/remote/vector_io/milvus/__init__.py

@@ -4,15 +4,18 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+from llama_stack.core.datatypes import VectorStoresConfig
 from llama_stack_api import Api, ProviderSpec
 
 from .config import MilvusVectorIOConfig
 
 
-async def get_adapter_impl(config: MilvusVectorIOConfig, deps: dict[Api, ProviderSpec]):
+async def get_adapter_impl(
+    config: MilvusVectorIOConfig, deps: dict[Api, ProviderSpec], vector_stores_config: VectorStoresConfig | None = None
+):
     from .milvus import MilvusVectorIOAdapter
 
     assert isinstance(config, MilvusVectorIOConfig), f"Unexpected config type: {type(config)}"
-    impl = MilvusVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
+    impl = MilvusVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files), vector_stores_config)
     await impl.initialize()
     return impl

src/llama_stack/providers/remote/vector_io/milvus/milvus.py

@@ -11,8 +11,9 @@
 from numpy.typing import NDArray
 from pymilvus import AnnSearchRequest, DataType, Function, FunctionType, MilvusClient, RRFRanker, WeightedRanker
 
-from llama_stack.core.storage.kvstore import kvstore_impl
+from llama_stack.core.datatypes import VectorStoresConfig
 from llama_stack.log import get_logger
+from llama_stack.providers.utils.kvstore import kvstore_impl
 from llama_stack.providers.inline.vector_io.milvus import MilvusVectorIOConfig as InlineMilvusVectorIOConfig
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
 from llama_stack.providers.utils.memory.vector_store import (
@@ -33,7 +34,7 @@
     VectorStoreNotFoundError,
     VectorStoresProtocolPrivate,
 )
-from llama_stack_api.internal.kvstore import KVStore
+from llama_stack.providers.utils.kvstore.api import KVStore
 
 from .config import MilvusVectorIOConfig as RemoteMilvusVectorIOConfig
 
@@ -272,12 +273,14 @@ def __init__(
         config: RemoteMilvusVectorIOConfig | InlineMilvusVectorIOConfig,
         inference_api: Inference,
         files_api: Files | None,
+        vector_stores_config: VectorStoresConfig | None = None,
     ) -> None:
         super().__init__(files_api=files_api, kvstore=None)
         self.config = config
         self.cache = {}
         self.client = None
         self.inference_api = inference_api
+        self.vector_stores_config = vector_stores_config
         self.vector_store_table = None
         self.metadata_collection_name = "openai_vector_stores_metadata"
 
@@ -298,6 +301,7 @@ async def initialize(self) -> None:
                     kvstore=self.kvstore,
                 ),
                 inference_api=self.inference_api,
+                vector_stores_config=self.vector_stores_config,
             )
             self.cache[vector_store.identifier] = index
         if isinstance(self.config, RemoteMilvusVectorIOConfig):
@@ -325,6 +329,7 @@ async def register_vector_store(self, vector_store: VectorStore) -> None:
             vector_store=vector_store,
             index=MilvusIndex(self.client, vector_store.identifier, consistency_level=consistency_level),
             inference_api=self.inference_api,
+            vector_stores_config=self.vector_stores_config,
         )
 
         self.cache[vector_store.identifier] = index
@@ -347,6 +352,7 @@ async def _get_and_cache_vector_store_index(self, vector_store_id: str) -> Vecto
             vector_store=vector_store,
             index=MilvusIndex(client=self.client, collection_name=vector_store.identifier, kvstore=self.kvstore),
             inference_api=self.inference_api,
+            vector_stores_config=self.vector_stores_config,
         )
         self.cache[vector_store_id] = index
         return index