From 4dfd0335b3cf7a38aecee8ecd44bcf1e78c1de6f Mon Sep 17 00:00:00 2001
From: James Braza <jamesbraza@gmail.com>
Date: Fri, 15 Nov 2024 10:56:05 -0800
Subject: [PATCH] Fixed LDP runner's `TRUNCATED` not calling `gen_answer`, and
 documented `AgentStatus` (#690)

---
 paperqa/agents/main.py   | 17 +++++------------
 paperqa/agents/models.py |  8 ++++----
 2 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/paperqa/agents/main.py b/paperqa/agents/main.py
index 61f697e3..24022fae 100644
--- a/paperqa/agents/main.py
+++ b/paperqa/agents/main.py
@@ -155,13 +155,15 @@ async def _run_with_timeout_failure(
             f"Agent timeout after {query.settings.agent.timeout}-sec, just answering."
         )
         status = AgentStatus.TRUNCATED
+    except Exception:
+        logger.exception("Trajectory failed.")
+        status = AgentStatus.FAIL
+    if status == AgentStatus.TRUNCATED:
+        # Fail over after truncation (too many steps, timeout): just answer
         generate_answer_tool = next(
             filter(lambda x: x.info.name == GenerateAnswer.TOOL_FN_NAME, env.tools)
         )
         await generate_answer_tool._tool_fn(question=query.query, state=env.state)
-    except Exception:
-        logger.exception("Trajectory failed.")
-        status = AgentStatus.FAIL
     return env.state.session, status
 
 
@@ -263,15 +265,6 @@ async def rollout() -> AgentStatus:
                     f"Agent didn't finish within {max_timesteps} timesteps, just"
                     " answering."
                 )
-                generate_answer_tool = next(
-                    filter(
-                        lambda x: x.info.name == GenerateAnswer.TOOL_FN_NAME,
-                        env.tools,
-                    )
-                )
-                await generate_answer_tool._tool_fn(
-                    question=query.query, state=env.state
-                )
                 return AgentStatus.TRUNCATED
             agent_state.messages += obs
             for attempt in Retrying(
diff --git a/paperqa/agents/models.py b/paperqa/agents/models.py
index 4bad05d4..dffe5f1a 100644
--- a/paperqa/agents/models.py
+++ b/paperqa/agents/models.py
@@ -34,15 +34,15 @@ def __getstate__(self) -> object: ...
     def __setstate__(self, state: object) -> None: ...
 
 
-class AgentStatus(StrEnum):
-    # FAIL - no answer could be generated
+class AgentStatus(StrEnum):  # TODO: rename to AnswerStatus or RolloutStatus
+    # FAIL - during the trajectory encountered an unhandled exception
     FAIL = "fail"
     # SUCCESS - answer was generated
     SUCCESS = "success"
     # TRUNCATED - agent didn't finish naturally (e.g. timeout, too many actions),
-    # so we prematurely answered
+    # so we just generated an answer after the unnatural finish
     TRUNCATED = "truncated"
-    # UNSURE - the agent was unsure, but an answer is present
+    # UNSURE - the gen_answer did not succeed, but an answer is present
     UNSURE = "unsure"