feat: allow setting query transformers in BaseRAGQuestionAnswerer

python/pathway/xpacks/llm/question_answering.py

@@ -129,6 +129,22 @@ def _get_RAG_prompt_udf(prompt_template: str | Callable[[str, str], str] | pw.UD
     return verified_template.as_udf()
 
 
+def _get_query_transformer_prompt_udf(
+    query_transformer_prompt: Callable[[str], str] | pw.UDF | None,
+) -> pw.UDF | None:
+    if query_transformer_prompt is None:
+        return None
+    elif isinstance(query_transformer_prompt, pw.UDF):
+        return query_transformer_prompt
+    elif callable(query_transformer_prompt):
+        return pw.udf(query_transformer_prompt)
+    else:
+        raise ValueError(
+            "Query transformer prompt must be type of one of the following: "
+            "Callable[[str], str] | ~pw.UDF | None"
+        )
+
+
 def _get_context_processor_udf(
     context_processor: (
         BaseContextProcessor | Callable[[list[dict] | list[Doc]], str] | pw.UDF
@@ -463,6 +479,13 @@ class BaseRAGQuestionAnswerer(SummaryQuestionAnswerer):
         search_topk: Top k parameter for the retrieval. Adjusts number of chunks in the context.
         rerank_topk: Number of top-scoring documents to retain after reranking, when a reranker is provided.
             If ``None``, reranking is disabled. Defaults to ``None``.
+        query_transformer_prompt: Prompt for transforming the user query before retrieval.
+            Must be a callable or ``pw.UDF`` that accepts the user query and returns a prompt
+            to send to ``llm`` for query rewriting. The transformed query is only used for
+            document retrieval; the answer is still generated from the original user query.
+            See ``pathway.xpacks.llm.prompts.prompt_query_rewrite`` and
+            ``pathway.xpacks.llm.prompts.prompt_query_rewrite_hyde`` for ready-to-use options.
+            Defaults to ``None``, which skips query transformation.
 
 
     Example:
@@ -524,6 +547,7 @@ def __init__(
         search_topk: int = 6,
         reranker: pw.UDF | None = None,
         rerank_topk: int | None = None,
+        query_transformer_prompt: Callable[[str], str] | pw.UDF | None = None,
     ) -> None:
 
         self.llm = llm
@@ -536,6 +560,9 @@ def __init__(
         self._init_schemas(default_llm_name)
 
         self.prompt_udf = _get_RAG_prompt_udf(prompt_template)
+        self.query_transformer_prompt_udf = _get_query_transformer_prompt_udf(
+            query_transformer_prompt
+        )
 
         if isinstance(context_processor, BaseContextProcessor):
             self.docs_to_context_transformer = context_processor.as_udf()
@@ -639,11 +666,30 @@ def add_score_to_doc(doc: pw.Json, score: float) -> dict:
     def answer_query(self, pw_ai_queries: pw.Table) -> pw.Table:
         """Answer a question based on the available information."""
 
+        if self.query_transformer_prompt_udf is not None:
+            pw_ai_queries += pw_ai_queries.select(
+                query_transformer_prompt=self.query_transformer_prompt_udf(
+                    pw.this.prompt
+                )
+            )
+            pw_ai_queries += pw_ai_queries.select(
+                search_query=self.llm(
+                    llms.prompt_chat_single_qa(pw.this.query_transformer_prompt),
+                    model=pw.this.model,
+                )
+            )
+            pw_ai_queries = pw_ai_queries.await_futures()
+            pw_ai_queries = pw_ai_queries.with_columns(
+                search_query=pw.coalesce(pw.this.search_query, pw.this.prompt)
+            )
+        else:
+            pw_ai_queries += pw_ai_queries.select(search_query=pw.this.prompt)
+
         pw_ai_results = pw_ai_queries + self.indexer.retrieve_query(
             pw_ai_queries.select(
                 metadata_filter=pw.this.filters,
                 filepath_globpattern=pw.cast(str | None, None),
-                query=pw.this.prompt,
+                query=pw.this.search_query,
                 k=self.search_topk,
             )
         ).select(

python/pathway/xpacks/llm/tests/test_rag.py

@@ -15,6 +15,21 @@
 from .utils import build_vector_store, create_rag_app
 
 
+class _QueryRewriteMockChat(llms.BaseChat):
+    """Rewrites the rewrite-prompt to a fixed search query; passes other prompts through."""
+
+    def _accepts_call_arg(self, arg_name: str) -> bool:
+        return False
+
+    async def __wrapped__(
+        self, messages: list[dict] | pw.Json, model: str
+    ) -> str | None:
+        content = messages[0]["content"].as_str()
+        if content.startswith("rewrite:"):
+            return content.removeprefix("rewrite:") or None
+        return model + "," + content
+
+
 @pw.udf
 def fake_embeddings_model(x: str) -> list[float]:
     return [
@@ -99,6 +114,120 @@ def test_base_rag():
     )
 
 
+def test_rag_app_set_query_transformer_prompt():
+    def query_transformer_prompt(query: str) -> str:
+        return f"rewrite:{query}"
+
+    rag_app = create_rag_app(query_transformer_prompt=query_transformer_prompt)
+
+    assert isinstance(rag_app.query_transformer_prompt_udf, pw.UDF)
+
+    assert _unwrap_udf(rag_app.query_transformer_prompt_udf)("foo") == "rewrite:foo"
+
+
+def test_rag_app_no_query_transformer_by_default():
+    rag_app = create_rag_app()
+
+    assert rag_app.query_transformer_prompt_udf is None
+
+
+def test_base_rag_uses_original_prompt_for_answer_without_transformer():
+    schema = pw.schema_from_types(data=bytes, _metadata=dict)
+    input = pw.debug.table_from_rows(
+        schema=schema, rows=[("foo", {}), ("bar", {}), ("baz", {})]
+    )
+
+    vector_server = VectorStoreServer(
+        input,
+        embedder=fake_embeddings_model,
+    )
+
+    rag = BaseRAGQuestionAnswerer(
+        IdentityMockChat(),
+        vector_server,
+        prompt_template=_prompt_template,
+        summarize_template=_summarize_template,
+        search_topk=1,
+    )
+
+    answer_queries = pw.debug.table_from_rows(
+        schema=rag.AnswerQuerySchema,
+        rows=[
+            ("foo", None, "gpt3.5", False),
+        ],
+    )
+
+    answer_output = rag.answer_query(answer_queries)
+
+    casted_table = answer_output.select(
+        result=pw.apply_with_type(lambda x: x.value, str, pw.this.result["response"])
+    )
+
+    assert_table_equality(
+        casted_table,
+        pw.debug.table_from_markdown(
+            """
+            result
+            gpt3.5,foo
+            """
+        ),
+    )
+
+
+@pytest.mark.parametrize("rewritten_query", ["bar", None])
+def test_base_rag_query_transformer_used_only_for_retrieval(
+    rewritten_query: str | None,
+):
+    @pw.udf
+    def query_transformer_prompt(query: str) -> str:
+        return f"rewrite:{rewritten_query or ''}"
+
+    schema = pw.schema_from_types(data=bytes, _metadata=dict)
+    input = pw.debug.table_from_rows(
+        schema=schema, rows=[("foo", {}), ("bar", {}), ("baz", {})]
+    )
+
+    vector_server = VectorStoreServer(
+        input,
+        embedder=fake_embeddings_model,
+    )
+
+    rag = BaseRAGQuestionAnswerer(
+        _QueryRewriteMockChat(),
+        vector_server,
+        prompt_template=_prompt_template,
+        query_transformer_prompt=query_transformer_prompt,
+        summarize_template=_summarize_template,
+        search_topk=1,
+    )
+
+    answer_queries = pw.debug.table_from_rows(
+        schema=rag.AnswerQuerySchema,
+        rows=[
+            ("foo", None, "gpt3.5", False),
+        ],
+    )
+
+    answer_output = rag.answer_query(answer_queries)
+
+    casted_table = answer_output.select(
+        search_query=pw.this.search_query,
+        result=pw.apply_with_type(lambda x: x.value, str, pw.this.result["response"]),
+    )
+
+    # Retrieval uses the rewritten query, falling back to the original prompt
+    # when the LLM returns None. The final LLM call still uses the original prompt.
+    assert_table_equality(
+        casted_table,
+        pw.debug.table_from_markdown(
+            f"""
+            search_query | result
+            {rewritten_query or 'foo'} | gpt3.5,foo
+            """
+        ),
+    )
+
+
 def test_rag_app_set_prompt():
     prompt_template = "Answer the question. Context: {context}\nQuestion: {query}"