Support for answer iterations in the gen_answer flow (#911)

Author: mskarlin

paperqa/docs.py

@@ -778,6 +778,11 @@ async def aquery(  # noqa: PLR0912
             answer_reasoning = None
         else:
             with set_llm_session_ids(session.id):
+                prior_answer_prompt = ""
+                if prompt_config.answer_iteration_prompt and session.answer:
+                    prior_answer_prompt = prompt_config.answer_iteration_prompt.format(
+                        prior_answer=session.answer
+                    )
                 messages = [
                     Message(role="system", content=prompt_config.system),
                     Message(
@@ -787,6 +792,7 @@ async def aquery(  # noqa: PLR0912
                             answer_length=answer_config.answer_length,
                             question=session.question,
                             example_citation=prompt_config.EXAMPLE_CITATION,
+                            prior_answer_prompt=prior_answer_prompt,
                         ),
                     ),
                 ]

paperqa/prompts.py

@@ -21,6 +21,15 @@
 # 1. Lead to complete tool being called with has_successful_answer=False
 # 2. Can be used for unit testing
 CANNOT_ANSWER_PHRASE = "I cannot answer"
+
+answer_iteration_prompt_template = (
+    "You are iterating on a prior answer, with a potentially different context:\n\n"
+    "{prior_answer}\n\n"
+    "Create a new answer only using keys and data from the included context."
+    " You can not use context keys from the prior answer which are not "
+    "also included in the above context.\n\n"
+)
+
 qa_prompt = (
     "Answer the question below with the context.\n\n"
     "Context (with relevance scores):\n\n{context}\n\n----\n\n"
@@ -37,6 +46,7 @@
     "so there may inaccuracies or ambiguities. If quotes are present and "
     "relevant, use them in the answer. This answer will go directly onto "
     "Wikipedia, so do not add any extraneous information.\n\n"
+    "{prior_answer_prompt}"
     "Answer ({answer_length}):"
 )
 

paperqa/settings.py

@@ -42,6 +42,7 @@
 from paperqa.prompts import (
     CONTEXT_INNER_PROMPT,
     CONTEXT_OUTER_PROMPT,
+    answer_iteration_prompt_template,
     citation_prompt,
     default_system_prompt,
     env_reset_prompt,
@@ -265,6 +266,13 @@ class PromptSettings(BaseModel):
 
     summary: str = summary_prompt
     qa: str = qa_prompt
+    answer_iteration_prompt: str | None = Field(
+        default=answer_iteration_prompt_template,
+        description=(
+            "Prompt to inject existing prior answers into the qa prompt to allow the model to iterate. "
+            "If None, then no prior answers will be injected."
+        ),
+    )
     select: str = select_paper_prompt
     pre: str | None = Field(
         default=None,

tests/test_paperqa.py

@@ -539,7 +539,29 @@ def test_location_awareness(docs_fixture) -> None:
 
 
 def test_query(docs_fixture) -> None:
-    docs_fixture.query("Is XAI usable in chemistry?")
+    settings = Settings(prompts={"answer_iteration_prompt": None})
+    docs_fixture.query("Is XAI usable in chemistry?", settings=settings)
+
+
+def test_query_with_iteration(docs_fixture) -> None:
+    # we store these results to check that the prompts are OK
+    my_results: list[LLMResult] = []
+    # explicitly set the prompt to use QA iterations
+    settings = Settings()
+    llm = settings.get_llm()
+    llm.llm_result_callback = my_results.append
+    prior_answer = "No, it isn't usable in chemistry."
+    question = "Is XAI usable in chemistry?"
+    prior_session = PQASession(question=question, answer=prior_answer)
+    docs_fixture.query(prior_session, llm_model=llm, settings=settings)
+    assert prior_answer in cast(
+        "str", my_results[-1].prompt[1].content  # type: ignore[union-attr, index]
+    ), "prior answer not in prompt"
+    # run without a prior session to check that the flow works correctly
+    docs_fixture.query(question, llm_model=llm, settings=settings)
+    assert settings.prompts.answer_iteration_prompt[:10] not in cast(  # type: ignore[index]
+        "str", my_results[-1].prompt[1].content  # type: ignore[union-attr, index]
+    ), "prior answer prompt should not be inserted"
 
 
 def test_llmresult_callback(docs_fixture: Docs) -> None: