[krisnaparahita]

Fixes #67

Summary

• Add an optional query_transformer_prompt constructor argument with no behavior change when omitted.
• Rewrite the user query through the configured LLM before document retrieval.
• Keep the original user prompt for final answer generation, avoiding the unsafe behavior identified in feat: add query_transformer parameter to BaseRAGQuestionAnswerer #209.
• Accept either a callable or Pathway UDF as the transformer prompt.

Validation

• Added four tests covering configuration, default behavior, rewritten-query retrieval, and preservation of the original answer prompt.
• Full test_rag.py suite passes: 13 tests.

[zxqfd555]

Thanks for the PR! The design looks good: the prompt goes through the LLM, the rewritten query is used only for retrieval, and the answer is built from the original question. One blocker, though.

It fails with the real chat wrappers. OpenAIChat, LiteLLMChat and CohereChat are annotated as -> str | None, so search_query becomes Optional(STR) while RetrieveQuerySchema.query requires STR. answer_query then fails while building the graph:

AssertionError: type of column query does not match - its type is Optional(STR) ... and STR in RetrieveQuerySchema
AssertionError: argument retrieval_queries has incorrect schema
    Line: pw_ai_results = pw_ai_queries + self.indexer.retrieve_query(

The tests don't catch it because _QueryRewriteMockChat.__wrapped__ is annotated as -> str. Changing it to -> str | None reproduces the failure.

Wrapping the LLM result in pw.coalesce(..., pw.this.prompt) after await_futures() fixes the type and also falls back to the original query when the LLM returns None.

[CLAassistant]

All committers have signed the CLA.

[krisnaparahita]

Thanks for the PR! The design looks good: the prompt goes through the LLM, the rewritten query is used only for retrieval, and the answer is built from the original question. One blocker, though.

It fails with the real chat wrappers. OpenAIChat, LiteLLMChat and CohereChat are annotated as -> str | None, so search_query becomes Optional(STR) while RetrieveQuerySchema.query requires STR. answer_query then fails while building the graph:

AssertionError: type of column query does not match - its type is Optional(STR) ... and STR in RetrieveQuerySchema
AssertionError: argument retrieval_queries has incorrect schema
    Line: pw_ai_results = pw_ai_queries + self.indexer.retrieve_query(

The tests don't catch it because _QueryRewriteMockChat.__wrapped__ is annotated as -> str. Changing it to -> str | None reproduces the failure.

Wrapping the LLM result in pw.coalesce(..., pw.this.prompt) after await_futures() fixes the type and also falls back to the original query when the LLM returns None.

Hi @zxqfd555 Thanks for the detailed feedback, that pinpointed it exactly. Fixed by wrapping the rewrite result in pw.coalesce(pw.this.search_query, pw.this.prompt) right after await_futures(), so the column type is STR (not Optional(STR)) and a None rewrite falls back to the original query, matching RetrieveQuerySchema.

 Also fixed the test gap: changed _QueryRewriteMockChat.__wrapped__'s annotation to -> str | None and added a test case for the None-fallback path, so this can't silently regress again. I confirmed by reverting the fix locally first, it reproduces your exact AssertionError on retrieval_queries's schema, then passes once reapplied.

 Pushed to the branch. Let me know if you'd like anything else adjusted.

[zxqfd555]

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.

[zxqfd555]

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]

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.