Skip to content

feat: allow setting query transformers in BaseRAGQuestionAnswerer #274

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
krisnaparahita wants to merge 3 commits into
base: main
Choose a base branch
from
+213 −1
Open

feat: allow setting query transformers in BaseRAGQuestionAnswerer #274

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
96a13ee
feat: allow setting query transformer in BaseRAGQuestionAnswerer
krisnaparahita Sep 4, 2026
a079385
fix: fall back to original query when RAG rewrite returns None
krisnaparahita Oct 6, 2026
b67afef
fix: await async RAG rewrite prompts and verify retrieved documents
krisnaparahita Oct 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 48 additions & 1 deletion python/pathway/xpacks/llm/question_answering.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand All @@ -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()
Expand Down Expand Up @@ -639,11 +666,31 @@ 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.await_futures()
pw_ai_queries += pw_ai_queries.select(
search_query=self.llm(
llms.prompt_chat_single_qa(pw.this.query_transformer_prompt),

@zxqfd555 zxqfd555 Oct 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for fixing the nullable LLM result! There is another case that fails while building the graph: a query_transformer_prompt UDF configured with fully_async_executor().

The prompt UDF produces Future[str], but this call to prompt_chat_single_qa consumes it before the existing await_futures() call. Here is a standalone reproduction, using a mock LLM and local documents (no API keys or external requests):

import pathway as pw
from pathway.xpacks.llm.llms import BaseChat
from pathway.xpacks.llm.question_answering import BaseRAGQuestionAnswerer
from pathway.xpacks.llm.vector_store import VectorStoreServer


class MockChat(BaseChat):
    def _accepts_call_arg(self, arg_name: str) -> bool:
        return False

    async def __wrapped__(
        self, messages: list[dict] | pw.Json, model: str | None = None
    ) -> str | None:
        return "rewritten query"


@pw.udf
def embed(text: str) -> list[float]:
    return [1.0, 0.0, 0.0]


@pw.udf(executor=pw.udfs.fully_async_executor())
async def rewrite_prompt(query: str) -> str:
    return f"Rewrite this question for retrieval: {query}"


documents = pw.debug.table_from_rows(
    pw.schema_from_types(data=bytes, _metadata=dict),
    [(b"A local example document", {})],
)
rag = BaseRAGQuestionAnswerer(
    llm=MockChat(),
    indexer=VectorStoreServer(documents, embedder=embed),
    query_transformer_prompt=rewrite_prompt,
    search_topk=1,
)
queries = pw.debug.table_from_rows(
    rag.AnswerQuerySchema,
    [("What is in the document?", None, None, False)],
)
rag.answer_query(queries)

This raises:

TypeError: Cannot perform pathway.apply when column of type
_asyncio.Future[str] is involved. Consider applying `await_futures()`
to the table used here.

Switching only the prompt executor to pw.udfs.async_executor() makes the graph build successfully.

Could you add pw_ai_queries = pw_ai_queries.await_futures() immediately after computing the query_transformer_prompt column, before constructing the chat message? The existing wait after the LLM call is still needed. I verified locally that this additional wait fixes the reproduction.

Please also add a regression test using a fully asynchronous prompt UDF. The constructor accepts pw.UDF without an executor restriction, so this configuration should work too.

@krisnaparahita krisnaparahita Oct 14, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for both finds, that's a real gap.

On the test assertions: added checks on the retrieved document itself. With rewrite bar, the retrieved doc now matches bar's expected content while the answer stays built from the original foo and with a None rewrite, it matches foo's. I confirmed the old tests would've passed even with the feature disabled, and reverting either fix now fails its corresponding test.

On the async UDF bug: reproduced your script exactly, it failed with the same TypeError. Added await_futures() right after the query_transformer_prompt column is computed, keeping the existing wait after the LLM rewrite. Added a regression test using a fully async prompt UDF, following your repro's structure.

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(
Expand Down
165 changes: 165 additions & 0 deletions python/pathway/xpacks/llm/tests/test_rag.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 [
Expand Down Expand Up @@ -99,6 +114,156 @@ 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,

@zxqfd555 zxqfd555 Oct 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could we also assert which document was retrieved here? The test checks the search_query column and the final answer, but neither assertion proves that the rewritten query actually reaches retrieval.

I verified this by changing query=pw.this.search_query back to query=pw.this.prompt in answer_query: all 14 tests in test_rag.py still pass.

With the existing documents and embeddings, please assert that:

  • When the original question is foo and the rewrite is bar, the retrieved document text is bar, while the final answer still uses the original question foo.
  • When the rewrite returns None, the retrieved document text is foo.

For example, include answer_output.docs[0]["text"] in the output assertions, or check the context documents in the response with return_context_docs=True. I checked these cases separately and the current implementation retrieves the correct documents; this request is to make the test catch a regression in that behavior.

retrieved_text=pw.unwrap(pw.this.docs[0]["text"].as_str()),
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 | retrieved_text | result
{rewritten_query or 'foo'} | {rewritten_query or 'foo'} | gpt3.5,foo
"""
),
)


def test_base_rag_fully_async_query_transformer_prompt():
@pw.udf(executor=pw.udfs.fully_async_executor())
async def query_transformer_prompt(query: str) -> str:
return f"rewrite:{query}"

rag = BaseRAGQuestionAnswerer(
_QueryRewriteMockChat(),
build_vector_store(fake_embeddings_model),
prompt_template=_prompt_template,
query_transformer_prompt=query_transformer_prompt,
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)

assert_table_equality(
answer_output.select(
search_query=pw.this.search_query,
result=pw.apply_with_type(
lambda x: x.value, str, pw.this.result["response"]
),
),
pw.debug.table_from_markdown(
"""
search_query | result
foo | gpt3.5,foo
"""
),
)


def test_rag_app_set_prompt():
prompt_template = "Answer the question. Context: {context}\nQuestion: {query}"

Expand Down
Loading