diff --git a/python/pathway/xpacks/llm/question_answering.py b/python/pathway/xpacks/llm/question_answering.py index 92be39a2a..20a23795a 100644 --- a/python/pathway/xpacks/llm/question_answering.py +++ b/python/pathway/xpacks/llm/question_answering.py @@ -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 @@ -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: @@ -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 @@ -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() @@ -639,11 +666,30 @@ 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.select( + search_query=self.llm( + llms.prompt_chat_single_qa(pw.this.query_transformer_prompt), + 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( diff --git a/python/pathway/xpacks/llm/tests/test_rag.py b/python/pathway/xpacks/llm/tests/test_rag.py index b93206a0c..585fe38bc 100644 --- a/python/pathway/xpacks/llm/tests/test_rag.py +++ b/python/pathway/xpacks/llm/tests/test_rag.py @@ -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 [ @@ -99,6 +114,120 @@ 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, + 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 | result + {rewritten_query or 'foo'} | gpt3.5,foo + """ + ), + ) + + def test_rag_app_set_prompt(): prompt_template = "Answer the question. Context: {context}\nQuestion: {query}"