refactor to only configuration of model at build time

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

Author: franciscojavierarceo

src/llama_stack/core/constants.py

@@ -0,0 +1,8 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# Default prompt template for query expansion in vector search
+DEFAULT_QUERY_EXPANSION_PROMPT = "Expand this query with relevant synonyms and related terms. Return only the improved query, no explanations:\n\n{query}\n\nImproved query:"

src/llama_stack/core/datatypes.py

@@ -12,6 +12,7 @@
 from pydantic import BaseModel, Field, field_validator, model_validator
 
 from llama_stack.core.access_control.datatypes import AccessRule
+from llama_stack.core.constants import DEFAULT_QUERY_EXPANSION_PROMPT
 from llama_stack.core.storage.datatypes import (
     KVStoreReference,
     StorageBackendType,
@@ -381,9 +382,17 @@ class VectorStoresConfig(BaseModel):
         description="Default LLM model for query expansion/rewriting in vector search.",
     )
     query_expansion_prompt: str = Field(
-        default="Expand this query with relevant synonyms and related terms. Return only the improved query, no explanations:\n\n{query}\n\nImproved query:",
+        default=DEFAULT_QUERY_EXPANSION_PROMPT,
         description="Prompt template for query expansion. Use {query} as placeholder for the original query.",
     )
+    query_expansion_max_tokens: int = Field(
+        default=100,
+        description="Maximum number of tokens for query expansion responses.",
+    )
+    query_expansion_temperature: float = Field(
+        default=0.3,
+        description="Temperature for query expansion model (0.0 = deterministic, 1.0 = creative).",
+    )
 
 
 class SafetyConfig(BaseModel):

src/llama_stack/core/resolver.py

@@ -374,13 +374,6 @@ 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"
 
@@ -401,11 +394,6 @@ 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/stack.py

@@ -34,6 +34,7 @@
 from llama_stack.core.store.registry import create_dist_registry
 from llama_stack.core.utils.dynamic import instantiate_class_type
 from llama_stack.log import get_logger
+from llama_stack.providers.utils.memory.vector_store import set_default_query_expansion_config
 from llama_stack_api import (
     Agents,
     Api,
@@ -144,35 +145,62 @@ async def validate_vector_stores_config(vector_stores_config: VectorStoresConfig
     if vector_stores_config is None:
         return
 
+    # Validate default embedding model
     default_embedding_model = vector_stores_config.default_embedding_model
-    if default_embedding_model is None:
-        return
+    if default_embedding_model is not None:
+        provider_id = default_embedding_model.provider_id
+        model_id = default_embedding_model.model_id
+        default_model_id = f"{provider_id}/{model_id}"
 
-    provider_id = default_embedding_model.provider_id
-    model_id = default_embedding_model.model_id
-    default_model_id = f"{provider_id}/{model_id}"
+        if Api.models not in impls:
+            raise ValueError(
+                f"Models API is not available but vector_stores config requires model '{default_model_id}'"
+            )
 
-    if Api.models not in impls:
-        raise ValueError(f"Models API is not available but vector_stores config requires model '{default_model_id}'")
+        models_impl = impls[Api.models]
+        response = await models_impl.list_models()
+        models_list = {m.identifier: m for m in response.data if m.model_type == "embedding"}
 
-    models_impl = impls[Api.models]
-    response = await models_impl.list_models()
-    models_list = {m.identifier: m for m in response.data if m.model_type == "embedding"}
+        default_model = models_list.get(default_model_id)
+        if default_model is None:
+            raise ValueError(
+                f"Embedding model '{default_model_id}' not found. Available embedding models: {models_list}"
+            )
 
-    default_model = models_list.get(default_model_id)
-    if default_model is None:
-        raise ValueError(f"Embedding model '{default_model_id}' not found. Available embedding models: {models_list}")
+        embedding_dimension = default_model.metadata.get("embedding_dimension")
+        if embedding_dimension is None:
+            raise ValueError(f"Embedding model '{default_model_id}' is missing 'embedding_dimension' in metadata")
 
-    embedding_dimension = default_model.metadata.get("embedding_dimension")
-    if embedding_dimension is None:
-        raise ValueError(f"Embedding model '{default_model_id}' is missing 'embedding_dimension' in metadata")
+        try:
+            int(embedding_dimension)
+        except ValueError as err:
+            raise ValueError(f"Embedding dimension '{embedding_dimension}' cannot be converted to an integer") from err
 
-    try:
-        int(embedding_dimension)
-    except ValueError as err:
-        raise ValueError(f"Embedding dimension '{embedding_dimension}' cannot be converted to an integer") from err
+        logger.debug(f"Validated default embedding model: {default_model_id} (dimension: {embedding_dimension})")
 
-    logger.debug(f"Validated default embedding model: {default_model_id} (dimension: {embedding_dimension})")
+    # Validate default query expansion model
+    default_query_expansion_model = vector_stores_config.default_query_expansion_model
+    if default_query_expansion_model is not None:
+        provider_id = default_query_expansion_model.provider_id
+        model_id = default_query_expansion_model.model_id
+        query_model_id = f"{provider_id}/{model_id}"
+
+        if Api.models not in impls:
+            raise ValueError(
+                f"Models API is not available but vector_stores config requires query expansion model '{query_model_id}'"
+            )
+
+        models_impl = impls[Api.models]
+        response = await models_impl.list_models()
+        llm_models_list = {m.identifier: m for m in response.data if m.model_type == "llm"}
+
+        query_expansion_model = llm_models_list.get(query_model_id)
+        if query_expansion_model is None:
+            raise ValueError(
+                f"Query expansion model '{query_model_id}' not found. Available LLM models: {list(llm_models_list.keys())}"
+            )
+
+        logger.debug(f"Validated default query expansion model: {query_model_id}")
 
 
 async def validate_safety_config(safety_config: SafetyConfig | None, impls: dict[Api, Any]):
@@ -437,6 +465,10 @@ async def initialize(self):
         await refresh_registry_once(impls)
         await validate_vector_stores_config(self.run_config.vector_stores, impls)
         await validate_safety_config(self.run_config.safety, impls)
+
+        # Set global query expansion configuration
+        set_default_query_expansion_config(self.run_config.vector_stores)
+
         self.impls = impls
 
     def create_registry_refresh_task(self):

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

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

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

@@ -14,7 +14,6 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from llama_stack.core.datatypes import VectorStoresConfig
 from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
@@ -190,12 +189,10 @@ def __init__(
         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:
@@ -211,7 +208,6 @@ 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
 
@@ -250,7 +246,6 @@ 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]:
@@ -284,7 +279,6 @@ 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,18 +6,15 @@
 
 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], vector_stores_config: VectorStoresConfig | None = None
-):
+async def get_provider_impl(config: SQLiteVectorIOConfig, deps: dict[Api, Any]):
     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), vector_stores_config)
+    impl = SQLiteVecVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
     await impl.initialize()
     return impl

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

@@ -14,7 +14,6 @@
 import sqlite_vec  # type: ignore[import-untyped]
 from numpy.typing import NDArray
 
-from llama_stack.core.datatypes import VectorStoresConfig
 from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin
@@ -391,12 +390,10 @@ def __init__(
         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
 
@@ -411,9 +408,7 @@ 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.vector_stores_config
-            )
+            self.cache[vector_store.identifier] = VectorStoreWithIndex(vector_store, index, self.inference_api)
 
         # Load existing OpenAI vector stores into the in-memory cache
         await self.initialize_openai_vector_stores()
@@ -437,9 +432,7 @@ 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.vector_stores_config
-        )
+        self.cache[vector_store.identifier] = VectorStoreWithIndex(vector_store, index, self.inference_api)
 
     async def _get_and_cache_vector_store_index(self, vector_store_id: str) -> VectorStoreWithIndex | None:
         if vector_store_id in self.cache:
@@ -464,7 +457,6 @@ 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,17 +4,14 @@
 # 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], vector_stores_config: VectorStoresConfig | None = None
-):
+async def get_adapter_impl(config: ChromaVectorIOConfig, deps: dict[Api, ProviderSpec]):
     from .chroma import ChromaVectorIOAdapter
 
-    impl = ChromaVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files), vector_stores_config)
+    impl = ChromaVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
     await impl.initialize()
     return impl

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

@@ -11,7 +11,6 @@
 import chromadb
 from numpy.typing import NDArray
 
-from llama_stack.core.datatypes import VectorStoresConfig
 from llama_stack.core.storage.kvstore import kvstore_impl
 from llama_stack.log import get_logger
 from llama_stack.providers.inline.vector_io.chroma import ChromaVectorIOConfig as InlineChromaVectorIOConfig
@@ -126,13 +125,11 @@ 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
@@ -165,7 +162,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, self.vector_stores_config
+            vector_store, ChromaIndex(self.client, collection), self.inference_api
         )
 
     async def unregister_vector_store(self, vector_store_id: str) -> None:
@@ -210,9 +207,7 @@ 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, self.vector_stores_config
-        )
+        index = VectorStoreWithIndex(vector_store, ChromaIndex(self.client, collection), self.inference_api)
         self.cache[vector_store_id] = index
         return index