renaming to query_rewrite, consolidating, and cleaning up validation

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

Author: franciscojavierarceo

src/llama_stack/core/datatypes.py

@@ -366,6 +366,27 @@ class QualifiedModel(BaseModel):
     model_id: str
 
 
+class RewriteQueryParams(BaseModel):
+    """Parameters for query rewriting/expansion."""
+
+    model: QualifiedModel | None = Field(
+        default=None,
+        description="LLM model for query rewriting/expansion in vector search.",
+    )
+    prompt: str = Field(
+        default=DEFAULT_QUERY_EXPANSION_PROMPT,
+        description="Prompt template for query rewriting. Use {query} as placeholder for the original query.",
+    )
+    max_tokens: int = Field(
+        default=100,
+        description="Maximum number of tokens for query expansion responses.",
+    )
+    temperature: float = Field(
+        default=0.3,
+        description="Temperature for query expansion model (0.0 = deterministic, 1.0 = creative).",
+    )
+
+
 class VectorStoresConfig(BaseModel):
     """Configuration for vector stores in the stack."""
 
@@ -377,21 +398,9 @@ class VectorStoresConfig(BaseModel):
         default=None,
         description="Default embedding model configuration for vector stores.",
     )
-    default_query_expansion_model: QualifiedModel | None = Field(
+    rewrite_query_params: RewriteQueryParams | None = Field(
         default=None,
-        description="Default LLM model for query expansion/rewriting in vector search.",
-    )
-    query_expansion_prompt: str = Field(
-        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).",
+        description="Parameters for query rewriting/expansion. None disables query rewriting.",
     )
 
 

src/llama_stack/core/stack.py

@@ -14,7 +14,7 @@
 import yaml
 
 from llama_stack.core.conversations.conversations import ConversationServiceConfig, ConversationServiceImpl
-from llama_stack.core.datatypes import Provider, SafetyConfig, StackRunConfig, VectorStoresConfig
+from llama_stack.core.datatypes import Provider, QualifiedModel, SafetyConfig, StackRunConfig, VectorStoresConfig
 from llama_stack.core.distribution import get_provider_registry
 from llama_stack.core.inspect import DistributionInspectConfig, DistributionInspectImpl
 from llama_stack.core.prompts.prompts import PromptServiceConfig, PromptServiceImpl
@@ -145,61 +145,67 @@ async def validate_vector_stores_config(vector_stores_config: VectorStoresConfig
         return
 
     # Validate default embedding model
-    default_embedding_model = vector_stores_config.default_embedding_model
-    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}"
+    if vector_stores_config.default_embedding_model is not None:
+        await _validate_embedding_model(vector_stores_config.default_embedding_model, impls)
 
-        if Api.models not in impls:
-            raise ValueError(
-                f"Models API is not available but vector_stores config requires model '{default_model_id}'"
-            )
+    # Validate default rewrite query model
+    if vector_stores_config.rewrite_query_params and vector_stores_config.rewrite_query_params.model:
+        await _validate_rewrite_query_model(vector_stores_config.rewrite_query_params.model, impls)
 
-        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}"
-            )
+async def _validate_embedding_model(embedding_model: QualifiedModel, impls: dict[Api, Any]) -> None:
+    """Validate that an embedding model exists and has required metadata."""
+    provider_id = embedding_model.provider_id
+    model_id = embedding_model.model_id
+    model_identifier = f"{provider_id}/{model_id}"
 
-        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")
+    if Api.models not in impls:
+        raise ValueError(f"Models API is not available but vector_stores config requires model '{model_identifier}'")
 
-        try:
-            int(embedding_dimension)
-        except ValueError as err:
-            raise ValueError(f"Embedding dimension '{embedding_dimension}' cannot be converted to an integer") from err
+    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"}
 
-        logger.debug(f"Validated default embedding model: {default_model_id} (dimension: {embedding_dimension})")
+    model = models_list.get(model_identifier)
+    if model is None:
+        raise ValueError(
+            f"Embedding model '{model_identifier}' not found. Available embedding models: {list(models_list.keys())}"
+        )
 
-    # 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}"
+    embedding_dimension = model.metadata.get("embedding_dimension")
+    if embedding_dimension is None:
+        raise ValueError(f"Embedding model '{model_identifier}' is missing 'embedding_dimension' in metadata")
 
-        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}'"
-            )
+    try:
+        int(embedding_dimension)
+    except ValueError as err:
+        raise ValueError(f"Embedding dimension '{embedding_dimension}' cannot be converted to an integer") from err
 
-        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 embedding model: {model_identifier} (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_rewrite_query_model(rewrite_query_model: QualifiedModel, impls: dict[Api, Any]) -> None:
+    """Validate that a rewrite query model exists and is accessible."""
+    provider_id = rewrite_query_model.provider_id
+    model_id = rewrite_query_model.model_id
+    model_identifier = f"{provider_id}/{model_id}"
+
+    if Api.models not in impls:
+        raise ValueError(
+            f"Models API is not available but vector_stores config requires rewrite query model '{model_identifier}'"
+        )
+
+    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"}
+
+    model = llm_models_list.get(model_identifier)
+    if model is None:
+        raise ValueError(
+            f"Rewrite query model '{model_identifier}' not found. Available LLM models: {list(llm_models_list.keys())}"
+        )
+
+    logger.debug(f"Validated rewrite query model: {model_identifier}")
 
 
 async def validate_safety_config(safety_config: SafetyConfig | None, impls: dict[Api, Any]):
@@ -466,9 +472,9 @@ async def initialize(self):
         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
+        from llama_stack.providers.utils.memory.rewrite_query_config import set_default_rewrite_query_config
 
-        set_default_query_expansion_config(self.run_config.vector_stores)
+        set_default_rewrite_query_config(self.run_config.vector_stores)
 
         self.impls = impls
 

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

@@ -288,15 +288,5 @@ vector_stores:
   default_embedding_model:
     provider_id: sentence-transformers
     model_id: nomic-ai/nomic-embed-text-v1.5
-  query_expansion_prompt: 'Expand this query with relevant synonyms and related terms.
-    Return only the improved query, no explanations:
-
-
-    {query}
-
-
-    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

@@ -279,15 +279,5 @@ vector_stores:
   default_embedding_model:
     provider_id: sentence-transformers
     model_id: nomic-ai/nomic-embed-text-v1.5
-  query_expansion_prompt: 'Expand this query with relevant synonyms and related terms.
-    Return only the improved query, no explanations:
-
-
-    {query}
-
-
-    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

@@ -291,15 +291,5 @@ vector_stores:
   default_embedding_model:
     provider_id: sentence-transformers
     model_id: nomic-ai/nomic-embed-text-v1.5
-  query_expansion_prompt: 'Expand this query with relevant synonyms and related terms.
-    Return only the improved query, no explanations:
-
-
-    {query}
-
-
-    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

@@ -282,15 +282,5 @@ vector_stores:
   default_embedding_model:
     provider_id: sentence-transformers
     model_id: nomic-ai/nomic-embed-text-v1.5
-  query_expansion_prompt: 'Expand this query with relevant synonyms and related terms.
-    Return only the improved query, no explanations:
-
-
-    {query}
-
-
-    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

@@ -288,15 +288,5 @@ vector_stores:
   default_embedding_model:
     provider_id: sentence-transformers
     model_id: nomic-ai/nomic-embed-text-v1.5
-  query_expansion_prompt: 'Expand this query with relevant synonyms and related terms.
-    Return only the improved query, no explanations:
-
-
-    {query}
-
-
-    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

@@ -279,15 +279,5 @@ vector_stores:
   default_embedding_model:
     provider_id: sentence-transformers
     model_id: nomic-ai/nomic-embed-text-v1.5
-  query_expansion_prompt: 'Expand this query with relevant synonyms and related terms.
-    Return only the improved query, no explanations:
-
-
-    {query}
-
-
-    Improved query:'
-  query_expansion_max_tokens: 100
-  query_expansion_temperature: 0.3
 safety:
   default_shield_id: llama-guard

src/llama_stack/providers/utils/memory/query_expansion_config.py

@@ -0,0 +1,38 @@
+# 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.
+
+from llama_stack.core.datatypes import QualifiedModel, VectorStoresConfig
+from llama_stack.providers.utils.memory.constants import DEFAULT_QUERY_EXPANSION_PROMPT
+
+# Global configuration for query rewriting - set during stack startup
+_DEFAULT_REWRITE_QUERY_MODEL: QualifiedModel | None = None
+_DEFAULT_REWRITE_QUERY_MAX_TOKENS: int = 100
+_DEFAULT_REWRITE_QUERY_TEMPERATURE: float = 0.3
+_REWRITE_QUERY_PROMPT_OVERRIDE: str | None = None
+
+
+def set_default_rewrite_query_config(vector_stores_config: VectorStoresConfig | None):
+    """Set the global default query rewriting configuration from stack config."""
+    global \
+        _DEFAULT_REWRITE_QUERY_MODEL, \
+        _REWRITE_QUERY_PROMPT_OVERRIDE, \
+        _DEFAULT_REWRITE_QUERY_MAX_TOKENS, \
+        _DEFAULT_REWRITE_QUERY_TEMPERATURE
+    if vector_stores_config and vector_stores_config.rewrite_query_params:
+        params = vector_stores_config.rewrite_query_params
+        _DEFAULT_REWRITE_QUERY_MODEL = params.model
+        # Only set override if user provided a custom prompt different from default
+        if params.prompt != DEFAULT_QUERY_EXPANSION_PROMPT:
+            _REWRITE_QUERY_PROMPT_OVERRIDE = params.prompt
+        else:
+            _REWRITE_QUERY_PROMPT_OVERRIDE = None
+        _DEFAULT_REWRITE_QUERY_MAX_TOKENS = params.max_tokens
+        _DEFAULT_REWRITE_QUERY_TEMPERATURE = params.temperature
+    else:
+        _DEFAULT_REWRITE_QUERY_MODEL = None
+        _REWRITE_QUERY_PROMPT_OVERRIDE = None
+        _DEFAULT_REWRITE_QUERY_MAX_TOKENS = 100
+        _DEFAULT_REWRITE_QUERY_TEMPERATURE = 0.3