feat(llm): allow setting query transformation behavior in BaseRAGQuestionAnswerer / BaseRAGQA (#67)

python/pathway/xpacks/llm/question_answering.py

@@ -2,8 +2,9 @@
 import json
 import logging
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 from warnings import warn
 
 import requests
@@ -45,9 +46,9 @@ class BaseContextProcessor(ABC):
     def maybe_unwrap_docs(self, docs: pw.Json | list[pw.Json] | list[Doc]):
         if isinstance(docs, pw.Json):
             doc_ls: list[Doc] = docs.as_list()
-        elif isinstance(docs, list) and all([isinstance(dc, dict) for dc in docs]):
+        elif isinstance(docs, list) and all(isinstance(dc, dict) for dc in docs):
             doc_ls = docs  # type: ignore
-        elif all([isinstance(doc, pw.Json) for doc in docs]):
+        elif all(isinstance(doc, pw.Json) for doc in docs):
             doc_ls = [doc.as_dict() for doc in docs]  # type: ignore
         else:
             raise ValueError(
@@ -463,6 +464,11 @@ 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_transform: Query transformation behavior to apply before retrieval.
+            Can be ``"rewrite"`` for query rewriting (using :py:func:`~pathway.xpacks.llm.prompts.prompt_query_rewrite`),
+            ``"hyde"`` for Hypothetical Document Embeddings (using :py:func:`~pathway.xpacks.llm.prompts.prompt_query_rewrite_hyde`),
+            a custom prompt callable/pw.UDF, or ``None`` to skip transformation and search with the raw query.
+            Defaults to ``None``.
 
 
     Example:
@@ -508,7 +514,7 @@ class BaseRAGQuestionAnswerer(SummaryQuestionAnswerer):
     ... )
     >>> app = QASummaryRestServer(app_host, app_port, rag)  # doctest: +SKIP
     >>> app.run()  # doctest: +SKIP
-    """  # noqa: E501
+    """
 
     def __init__(
         self,
@@ -524,11 +530,36 @@ def __init__(
         search_topk: int = 6,
         reranker: pw.UDF | None = None,
         rerank_topk: int | None = None,
+        query_transform: str | Callable | pw.UDF | None = None,
     ) -> None:
 
         self.llm = llm
         self.indexer = indexer
         self.reranker = reranker
+        self.query_transform = query_transform
+
+        if query_transform is None:
+            self.query_transform_udf = None
+        elif isinstance(query_transform, str):
+            qt_normalized = query_transform.strip().lower()
+            if qt_normalized == "rewrite":
+                self.query_transform_udf = prompts.prompt_query_rewrite
+            elif qt_normalized == "hyde":
+                self.query_transform_udf = prompts.prompt_query_rewrite_hyde
+            else:
+                raise ValueError(
+                    f"Unknown query_transform: '{query_transform}'. "
+                    "Expected 'rewrite', 'hyde', a callable/pw.UDF, or None."
+                )
+        elif isinstance(query_transform, pw.UDF):
+            self.query_transform_udf = query_transform
+        elif callable(query_transform):
+            self.query_transform_udf = pw.udf(query_transform)
+        else:
+            raise ValueError(
+                f"Unknown query_transform: '{query_transform}'. "
+                "Expected 'rewrite', 'hyde', a callable/pw.UDF, or None."
+            )
 
         if default_llm_name is None:
             default_llm_name = llm.model
@@ -639,16 +670,38 @@ 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."""
 
-        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,
-                k=self.search_topk,
+        if self.query_transform_udf is not None:
+            transformed_queries = pw_ai_queries.select(
+                pw.this.filters,
+                search_query=self.llm(
+                    llms.prompt_chat_single_qa(
+                        self.query_transform_udf(pw.this.prompt)
+                    ),
+                    model=pw.this.model,
+                ),
+            ).await_futures()
+
+            pw_ai_results = pw_ai_queries + self.indexer.retrieve_query(
+                transformed_queries.select(
+                    metadata_filter=pw.this.filters,
+                    filepath_globpattern=pw.cast(str | None, None),
+                    query=pw.this.search_query,
+                    k=self.search_topk,
+                )
+            ).select(
+                docs=pw.this.result,
+            )
+        else:
+            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,
+                    k=self.search_topk,
+                )
+            ).select(
+                docs=pw.this.result,
             )
-        ).select(
-            docs=pw.this.result,
-        )
 
         if self.reranker is not None:
             pw_ai_results = self._apply_reranking(pw_ai_results)
@@ -886,7 +939,7 @@ class AdaptiveRAGQuestionAnswerer(BaseRAGQuestionAnswerer):
     ... )
     >>> app.build_server(host=app_host, port=app_port)  # doctest: +SKIP
     >>> app.run_server()  # doctest: +SKIP
-    """  # noqa: E501
+    """
 
     def __init__(
         self,
@@ -905,13 +958,15 @@ def __init__(
         n_starting_documents: int = 2,
         factor: int = 2,
         max_iterations: int = 4,
+        query_transform: str | Callable | pw.UDF | None = None,
     ) -> None:
         super().__init__(
             llm,
             indexer,
             default_llm_name=default_llm_name,
             summarize_template=summarize_template,
             context_processor=context_processor,
+            query_transform=query_transform,
         )
         self.prompt_template = prompt_template
         self.n_starting_documents = n_starting_documents
@@ -1251,3 +1306,6 @@ def pw_list_documents(self, *args, **kwargs):
             DeprecationWarning,
         )
         return self.list_documents(*args, **kwargs)
+
+
+BaseRAGQA = BaseRAGQuestionAnswerer

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

@@ -6,9 +6,9 @@
 
 import pathway as pw
 from pathway.tests.utils import assert_table_equality
-from pathway.xpacks.llm import llms
+from pathway.xpacks.llm import llms, prompts
 from pathway.xpacks.llm._utils import _unwrap_udf
-from pathway.xpacks.llm.question_answering import BaseRAGQuestionAnswerer
+from pathway.xpacks.llm.question_answering import BaseRAGQA, BaseRAGQuestionAnswerer
 from pathway.xpacks.llm.vector_store import VectorStoreServer
 
 from .mocks import IdentityMockChat
@@ -182,3 +182,80 @@ def test_rag_client_accepts_all_timeout_forms():
 
     with pytest.raises(ValueError, match="'timeout' must be positive"):
         RAGClient(url="http://localhost:8080", timeout=-1)
+
+
+def test_base_rag_qa_alias():
+    assert BaseRAGQA is BaseRAGQuestionAnswerer
+
+
+def test_rag_query_transform_initialization():
+    @pw.udf
+    def fake_embeddings_model(x: str) -> list[float]:
+        return [1.0, 1.0, 0.0]
+
+    class FakeChatModel(llms.BaseChat):
+        async def __wrapped__(self, *args, **kwargs) -> str:
+            return "Text"
+
+        def _accepts_call_arg(self, arg_name: str) -> bool:
+            return True
+
+    chat = FakeChatModel()
+    vector_server = build_vector_store(fake_embeddings_model)
+
+    # 1. Default is None
+    rag_default = BaseRAGQA(llm=chat, indexer=vector_server)
+    assert rag_default.query_transform is None
+    assert rag_default.query_transform_udf is None
+
+    # 2. Rewrite
+    rag_rewrite = BaseRAGQA(llm=chat, indexer=vector_server, query_transform="rewrite")
+    assert rag_rewrite.query_transform == "rewrite"
+    assert rag_rewrite.query_transform_udf == prompts.prompt_query_rewrite
+
+    # 3. HyDE
+    rag_hyde = BaseRAGQA(llm=chat, indexer=vector_server, query_transform="hyde")
+    assert rag_hyde.query_transform == "hyde"
+    assert rag_hyde.query_transform_udf == prompts.prompt_query_rewrite_hyde
+
+    # 4. Custom callable / UDF
+    def custom_transform(query: str) -> str:
+        return f"custom: {query}"
+
+    rag_custom = BaseRAGQA(llm=chat, indexer=vector_server, query_transform=custom_transform)
+    assert isinstance(rag_custom.query_transform_udf, pw.UDF)
+
+    # 5. Invalid string raises ValueError
+    with pytest.raises(ValueError, match="Unknown query_transform"):
+        BaseRAGQA(llm=chat, indexer=vector_server, query_transform="unknown_mode")
+
+
+def test_rag_execution_with_query_transform():
+    schema = pw.schema_from_types(data=bytes, _metadata=dict)
+    input_data = pw.debug.table_from_rows(
+        schema=schema, rows=[("foo", {}), ("bar", {}), ("baz", {})]
+    )
+
+    vector_server = VectorStoreServer(
+        input_data,
+        embedder=fake_embeddings_model,
+    )
+
+    rag = BaseRAGQA(
+        IdentityMockChat(),
+        vector_server,
+        prompt_template=_prompt_template,
+        summarize_template=_summarize_template,
+        search_topk=1,
+        query_transform="rewrite",
+    )
+
+    answer_queries = pw.debug.table_from_rows(
+        schema=rag.AnswerQuerySchema,
+        rows=[
+            ("foo", None, "gpt3.5", False),
+        ],
+    )
+
+    answer_output = rag.answer_query(answer_queries)
+    assert answer_output is not None