Retrying if ToolSelector fails to select a tool (#373)

.pre-commit-config.yaml

@@ -76,7 +76,7 @@ repos:
         additional_dependencies:
           - aiohttp
           - fhaviary[llm]>=0.6 # Match pyproject.toml
-          - ldp>=0.4 # Match pyproject.toml
+          - ldp>=0.6 # Match pyproject.toml
           - html2text
           - httpx
           - numpy

paperqa/agents/main.py

@@ -5,7 +5,7 @@
 from pydoc import locate
 from typing import TYPE_CHECKING, Any, cast
 
-from aviary.message import Message
+from aviary.message import MalformedMessageError, Message
 from aviary.tools import (
     Tool,
     ToolCall,
@@ -14,6 +14,12 @@
     ToolSelector,
 )
 from pydantic import BaseModel, Field, TypeAdapter
+from tenacity import (
+    Retrying,
+    before_sleep_log,
+    retry_if_exception_type,
+    stop_after_attempt,
+)
 
 try:
     from ldp.agent import (
@@ -317,7 +323,14 @@ async def run_aviary_agent(
 
             while not done:
                 agent_state.messages += obs
-                action = await agent(agent_state.messages, tools)
+                for attempt in Retrying(
+                    stop=stop_after_attempt(5),
+                    retry=retry_if_exception_type(MalformedMessageError),
+                    before_sleep=before_sleep_log(logger, logging.WARNING),
+                    reraise=True,
+                ):
+                    with attempt:  # Retrying if ToolSelector fails to select a tool
+                        action = await agent(agent_state.messages, tools)
                 agent_state.messages = [*agent_state.messages, action]
                 if on_agent_action_callback:
                     await on_agent_action_callback(action, agent_state)

pyproject.toml

@@ -50,7 +50,7 @@ requires-python = ">=3.11"
 
 [project.optional-dependencies]
 ldp = [
-    "ldp>=0.4",
+    "ldp>=0.6",  # For LiteLLM 1.42.1 pinning's UnsupportedParamsError fix
 ]
 typing = [
     "types-PyYAML",