adding query expansion model to vector store config

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

Author: franciscojavierarceo

src/llama_stack/core/datatypes.py

@@ -376,6 +376,14 @@ class VectorStoresConfig(BaseModel):
         default=None,
         description="Default embedding model configuration for vector stores.",
     )
+    default_query_expansion_model: QualifiedModel | None = Field(
+        default=None,
+        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:",
+        description="Prompt template for query expansion. Use {query} as placeholder for the original query.",
+    )
 
 
 class SafetyConfig(BaseModel):

src/llama_stack/core/routers/vector_io.py

@@ -99,6 +99,12 @@ 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 = {}
+        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/providers/utils/memory/vector_store.py

@@ -17,6 +17,7 @@
 from numpy.typing import NDArray
 from pydantic import BaseModel
 
+from llama_stack.core.datatypes import VectorStoresConfig
 from llama_stack.log import get_logger
 from llama_stack.models.llama.llama3.tokenizer import Tokenizer
 from llama_stack.providers.utils.inference.prompt_adapter import (
@@ -267,6 +268,7 @@ class VectorStoreWithIndex:
     vector_store: VectorStore
     index: EmbeddingIndex
     inference_api: Api.inference
+    vector_stores_config: VectorStoresConfig | None = None
 
     async def insert_chunks(
         self,
@@ -301,6 +303,11 @@ async def query_chunks(
     ) -> QueryChunksResponse:
         if params is None:
             params = {}
+
+        # Extract configuration if provided by router
+        if "vector_stores_config" in params:
+            self.vector_stores_config = params["vector_stores_config"]
+
         k = params.get("max_chunks", 3)
         mode = params.get("mode")
         score_threshold = params.get("score_threshold", 0.0)
@@ -350,51 +357,40 @@ async def _rewrite_query_for_search(self, query: str) -> str:
         :param query: The original user query
         :returns: The rewritten query optimized for vector search
         """
-        # Get available models and find a suitable chat model
+        # Check if query expansion model is configured
+        if not self.vector_stores_config or not self.vector_stores_config.default_query_expansion_model:
+            raise ValueError("No default_query_expansion_model configured for query rewriting")
+
+        # Use the configured model
+        expansion_model = self.vector_stores_config.default_query_expansion_model
+        chat_model = f"{expansion_model.provider_id}/{expansion_model.model_id}"
+
+        # Validate that the model is available and is an LLM
         try:
             models_response = await self.inference_api.routing_table.list_models()
         except Exception as e:
-            raise RuntimeError(f"Failed to list available models for query rewriting: {e}") from e
-
-        chat_model = None
-        # Look for an LLM model (for chat completion)
-        # Prefer local or non-cloud providers to avoid credential issues
-        llm_models = [m for m in models_response.data if m.model_type == ModelType.llm]
-
-        # Filter out models that are known to be embedding models (misclassified as LLM)
-        embedding_model_patterns = ["minilm", "embed", "embedding", "nomic-embed"]
-        llm_models = [
-            m for m in llm_models if not any(pattern in m.identifier.lower() for pattern in embedding_model_patterns)
-        ]
-
-        # Priority order: ollama (local), then OpenAI, then others
-        provider_priority = ["ollama", "openai", "gemini", "bedrock"]
-
-        for provider in provider_priority:
-            for model in llm_models:
-                model_id = model.identifier.lower()
-                if provider == "ollama" and "ollama/" in model_id:
-                    chat_model = model.identifier
-                    break
-                elif provider in model_id:
-                    chat_model = model.identifier
-                    break
-            if chat_model:
+            raise RuntimeError(f"Failed to list available models for validation: {e}") from e
+
+        model_found = False
+        for model in models_response.data:
+            if model.identifier == chat_model:
+                if model.model_type != ModelType.llm:
+                    raise ValueError(
+                        f"Configured query expansion model '{chat_model}' is not an LLM model "
+                        f"(found type: {model.model_type}). Query rewriting requires an LLM model."
+                    )
+                model_found = True
                 break
 
-        # Fallback: use first available LLM model if no preferred provider found
-        if not chat_model and llm_models:
-            chat_model = llm_models[0].identifier
-
-        # If no suitable model found, raise an error
-        if not chat_model:
-            raise ValueError("No LLM model available for query rewriting")
-
-        rewrite_prompt = f"""Expand this query with relevant synonyms and related terms. Return only the improved query, no explanations:
-
-{query}
+        if not model_found:
+            available_llm_models = [m.identifier for m in models_response.data if m.model_type == ModelType.llm]
+            raise ValueError(
+                f"Configured query expansion model '{chat_model}' is not available. "
+                f"Available LLM models: {available_llm_models}"
+            )
 
-Improved query:"""
+        # Use the configured prompt (has a default value)
+        rewrite_prompt = self.vector_stores_config.query_expansion_prompt.format(query=query)
 
         chat_request = OpenAIChatCompletionRequestWithExtraBody(
             model=chat_model,