feat: Actualize query rewrite in search API

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

Author: franciscojavierarceo

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

@@ -611,6 +611,7 @@ async def openai_search_vector_store(
                 "max_chunks": max_num_results * CHUNK_MULTIPLIER,
                 "score_threshold": score_threshold,
                 "mode": search_mode,
+                "rewrite_query": rewrite_query,
             }
             # TODO: Add support for ranking_options.ranker
 

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

@@ -34,6 +34,11 @@
     RAGDocument,
     VectorStore,
 )
+from llama_stack_api.inference import (
+    OpenAIChatCompletionRequestWithExtraBody,
+    OpenAIUserMessageParam,
+)
+from llama_stack_api.models import ModelType
 
 log = get_logger(name=__name__, category="providers::utils")
 
@@ -318,6 +323,11 @@ async def query_chunks(
                 reranker_params = {"impact_factor": k_value}
 
         query_string = interleaved_content_as_str(query)
+
+        # Apply query rewriting if enabled
+        if params.get("rewrite_query", False):
+            query_string = await self._rewrite_query_for_search(query_string)
+
         if mode == "keyword":
             return await self.index.query_keyword(query_string, k, score_threshold)
 
@@ -333,3 +343,78 @@ async def query_chunks(
             )
         else:
             return await self.index.query_vector(query_vector, k, score_threshold)
+
+    async def _rewrite_query_for_search(self, query: str) -> str:
+        """Rewrite the user query to improve vector search performance.
+
+        :param query: The original user query
+        :returns: The rewritten query optimized for vector search
+        """
+        # Get available models and find a suitable chat model
+        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:
+                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}
+
+Improved query:"""
+
+        chat_request = OpenAIChatCompletionRequestWithExtraBody(
+            model=chat_model,
+            messages=[
+                OpenAIUserMessageParam(
+                    role="user",
+                    content=rewrite_prompt,
+                )
+            ],
+            max_tokens=100,
+        )
+
+        try:
+            response = await self.inference_api.openai_chat_completion(chat_request)
+        except Exception as e:
+            raise RuntimeError(f"Failed to generate rewritten query: {e}") from e
+
+        if response.choices and len(response.choices) > 0:
+            rewritten_query = response.choices[0].message.content.strip()
+            log.info(f"Query rewritten: '{query}' → '{rewritten_query}'")
+            return rewritten_query
+        else:
+            raise RuntimeError("No response received from LLM model for query rewriting")