feat: Actualize query rewrite in search API

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

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,55 @@ 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)
+        for model in models_response.data:
+            if model.model_type == ModelType.llm:
+                chat_model = model.identifier
+                break
+
+        # If no suitable model found, raise an error
+        if not chat_model:
+            raise ValueError("No LLM model available for query rewriting")
+
+        rewrite_prompt = f"""Rewrite this search query to improve vector search results by expanding it with relevant synonyms and related terms while maintaining the original intent:
+
+{query}
+
+Rewritten 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")