From 96575017c10d6d4f2d93dff13a3c7904fcea9a93 Mon Sep 17 00:00:00 2001 From: Omm Prakash Date: Sat, 3 Oct 2026 23:01:50 +0530 Subject: [PATCH] feat(llm): add query transformation and BaseRAGQA alias --- .../pathway/xpacks/llm/question_answering.py | 86 ++++++++++++++++--- python/pathway/xpacks/llm/tests/test_rag.py | 81 ++++++++++++++++- 2 files changed, 151 insertions(+), 16 deletions(-) diff --git a/python/pathway/xpacks/llm/question_answering.py b/python/pathway/xpacks/llm/question_answering.py index 92be39a2a..7e66d581c 100644 --- a/python/pathway/xpacks/llm/question_answering.py +++ b/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,6 +958,7 @@ 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, @@ -912,6 +966,7 @@ def __init__( 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 diff --git a/python/pathway/xpacks/llm/tests/test_rag.py b/python/pathway/xpacks/llm/tests/test_rag.py index b93206a0c..0e7b7517b 100644 --- a/python/pathway/xpacks/llm/tests/test_rag.py +++ b/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