feat: allow setting query transformers in BaseRAGQuestionAnswerer

[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]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
0 out of 2 committers have signed the CLA.

❌ krisnaparahita
❌ vdrace-design
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[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.