refactor to only configuration of model at build time

Signed-off-by: Francisco Javier Arceo <[email protected]>

Author: franciscojavierarceo

src/llama_stack/core/datatypes.py

@@ -18,6 +18,7 @@
     StorageConfig,
 )
 from llama_stack.log import LoggingConfig
+from llama_stack.providers.utils.memory.constants import DEFAULT_QUERY_EXPANSION_PROMPT
 from llama_stack_api import (
     Api,
     Benchmark,
@@ -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/routers/vector_io.py

@@ -99,19 +99,6 @@ async def query_chunks(
     ) -> QueryChunksResponse:
         logger.debug(f"VectorIORouter.query_chunks: {vector_store_id}")
         provider = await self.routing_table.get_provider_impl(vector_store_id)
-
-        # 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)
 
     # OpenAI Vector Stores API endpoints

src/llama_stack/core/stack.py

@@ -144,35 +144,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})")
+
+    # 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"}
 
-    logger.debug(f"Validated default embedding model: {default_model_id} (dimension: {embedding_dimension})")
+        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 +464,12 @@ 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 from stack config
+        from llama_stack.providers.utils.memory.query_expansion_config import set_default_query_expansion_config
+
+        set_default_query_expansion_config(self.run_config.vector_stores)
+
         self.impls = impls
 
     def create_registry_refresh_task(self):

src/llama_stack/distributions/ci-tests/run-with-postgres-store.yaml

@@ -296,5 +296,7 @@ vector_stores:
 
 
     Improved query:'
+  query_expansion_max_tokens: 100
+  query_expansion_temperature: 0.3
 safety:
   default_shield_id: llama-guard

src/llama_stack/distributions/ci-tests/run.yaml

@@ -287,5 +287,7 @@ vector_stores:
 
 
     Improved query:'
+  query_expansion_max_tokens: 100
+  query_expansion_temperature: 0.3
 safety:
   default_shield_id: llama-guard

src/llama_stack/distributions/starter-gpu/run-with-postgres-store.yaml

@@ -299,5 +299,7 @@ vector_stores:
 
 
     Improved query:'
+  query_expansion_max_tokens: 100
+  query_expansion_temperature: 0.3
 safety:
   default_shield_id: llama-guard

src/llama_stack/distributions/starter-gpu/run.yaml

@@ -290,5 +290,7 @@ vector_stores:
 
 
     Improved query:'
+  query_expansion_max_tokens: 100
+  query_expansion_temperature: 0.3
 safety:
   default_shield_id: llama-guard

src/llama_stack/distributions/starter/run-with-postgres-store.yaml

@@ -296,5 +296,7 @@ vector_stores:
 
 
     Improved query:'
+  query_expansion_max_tokens: 100
+  query_expansion_temperature: 0.3
 safety:
   default_shield_id: llama-guard

src/llama_stack/distributions/starter/run.yaml

@@ -287,5 +287,7 @@ vector_stores:
 
 
     Improved query:'
+  query_expansion_max_tokens: 100
+  query_expansion_temperature: 0.3
 safety:
   default_shield_id: llama-guard