Callback Updates

agents/abstract.py

@@ -10,10 +10,13 @@
 from asyncio import Task, create_task
 from collections.abc import Awaitable, Callable, Mapping, Sequence
 from dataclasses import dataclass, field
+from enum import StrEnum
 from typing import (
     TYPE_CHECKING,
     Any,
     Literal,
+    Protocol,
+    Self,
     Union,
 )
 
@@ -256,7 +259,7 @@ class _Agent(Observable, metaclass=abc.ABCMeta):
     BASE_PROMPT: str = ""
     SYSTEM_PROMPT: str = ""
     oai_kwargs: dict[str, Any]
-    CALLBACKS: list[Callable[..., Any]]
+    CALLBACKS: list["Callback[Self]"]
     callback_output: list[Any]
     tool_res_payload: list[dict[str, Any]]
     provider: _Provider[Any]
@@ -268,7 +271,7 @@ def __init__(
         model_name: str | None = None,
         provider: _Provider[Any] | None = None,
         tools: Sequence[Any] | None = None,
-        callbacks: Sequence[Callable[..., Any]] | None = None,
+        callbacks: Sequence["Callback[Self]"] | None = None,
         oai_kwargs: dict[str, Any] | None = None,
         **fmt_kwargs: Any,
     ) -> None:
@@ -370,19 +373,29 @@ def register_provider(self, provider: ProviderT) -> None:
         raise NotImplementedError()
 
 
-class _Callback[CallingAgentT: _Agent](metaclass=abc.ABCMeta):
+class CallbackCondition(StrEnum):
+    """
+    Conditions under which a callback is dispatched.
+    """
+
+    ALWAYS = "always"
+    ON_SUCCESS = "on_success"
+    ON_ERROR = "on_error"
+
+
+class Callback[CallingAgentT: _Agent](Protocol):
     """
     A Callback virtual class
     """
 
-    @abc.abstractmethod
-    async def __call__(self, cls: CallingAgentT, answer: Any, scratchpad: str) -> None:
+    condition: CallbackCondition = CallbackCondition.ALWAYS
+
+    async def __call__(self, cls: CallingAgentT, exc: Exception | None) -> None:
         """
         Primary method called by agent during callback process
 
         :param Agent cls: Instantitated class of calling agent for possible modification
-        :param answer: The final response of the calling Agent
-        :param str scratchpad: The full interaction history of the calling Agent
+        :param exc: An exception raised during runtime of the agent, if present
 
         """
-        raise NotImplementedError()
+        ...

agents/agent/base.py

@@ -1,13 +1,20 @@
 import asyncio
 import logging
-from collections.abc import Awaitable, Callable, Sequence
-from typing import Any, cast
+from collections.abc import Sequence
+from typing import Any, Self, cast
 
 import openai
 from openai.types.chat.chat_completion import Choice
 from pydantic import BaseModel
 
-from ..abstract import Message, _Agent, _Provider, _StoppingCondition
+from ..abstract import (
+    Callback,
+    CallbackCondition,
+    Message,
+    _Agent,
+    _Provider,
+    _StoppingCondition,
+)
 from ..json_tool_gen import ResolvedTool, Tool, ToolDefinition
 from ..providers import AzureOpenAIProvider
 from ..stopping_conditions import StopOnDataModel
@@ -30,7 +37,7 @@ class Agent(_Agent):
     :param dict oai_kwargs: OpenAI arguments passed as-is to API (temperature, top_p, etc.)
     :param list TOOLS: Runtime collection of tools available to the agent. Tools
         are declared with the constructor or ``@agent_callable`` decorators.
-    :param list CALLBACKS: List of callbacks to evaluate at completion. Should be a list of callables with a signature `fun(self, answer, scratchpad)`
+    :param list CALLBACKS: Callbacks with a condition and an async signature `callback(agent, exc)`.
     :param _StoppingCondition stopping_condition: The StoppingCondition handler class which will be called after each step to determine if the task is completed.
 
 
@@ -41,7 +48,8 @@ class Agent(_Agent):
 
     Callback Use
     ------------
-    Each callback will have access to class object, scratchpad, and final answer, thus the signature must match.
+    Each callback receives the calling agent and its run exception, or None on success.
+    Its condition selects whether it runs on success, on error, or both.
     This is still quite experimental, but the intended usecase is for reflection / refinement applications.
     """
 
@@ -51,7 +59,7 @@ def __init__(
         model_name: str | None = None,
         provider: _Provider[Any] | None = None,
         tools: Sequence[Tool[Any] | ToolDefinition] | None = None,
-        callbacks: Sequence[Callable[..., Any]] | None = None,
+        callbacks: Sequence[Callback[Self]] | None = None,
         oai_kwargs: dict[str, Any] | None = None,
         **fmt_kwargs: Any,
     ) -> None:
@@ -62,7 +70,7 @@ def __init__(
         :param str model_name: Name of model to use (or deployment name for AzureOpenAI) (optional if provider is passed)
         :param Type[_Provider] provider: Instantiated OpenAI instance to use (optional)
         :param tools: Executable tools or legacy definitions backed by agent methods (optional)
-        :param List[Callable] callbacks: List of callbacks to evaluate at end of run (optional)
+        :param callbacks: Callbacks dispatched according to their condition (optional)
         :param dict[str, any] oai_kwargs: Dict of additional OpenAI arguments to pass thru to chat call
         :param fmt_kwargs: Additional named arguments which will be inserted into the :func:`BASE_PROMPT` via fstring
         """
@@ -119,26 +127,36 @@ async def run(self, reset: bool = False, *kwargs) -> None:
         if reset:
             self.reset()
 
-        while not (self.is_terminated or self.is_truncated):
-            logger.debug(f"Running step {self.curr_step}.")
-            await self.step()
+        try:
+            while not (self.is_terminated or self.is_truncated):
+                logger.debug(f"Running step {self.curr_step}.")
+                await self.step()
+        except Exception as err:
+            await self._dispatch_callbacks(err)
+            raise
+        else:
+            await self._dispatch_callbacks(None)
+
+    async def _dispatch_callbacks(self: Self, exc: Exception | None) -> None:
+        outcome = (
+            CallbackCondition.ON_SUCCESS if exc is None else CallbackCondition.ON_ERROR
+        )
 
-        # Evaluate callbacks, if available
         for callback in self.CALLBACKS:
-            await self._handle_callback(
-                callback, self, answer=self.answer, scratchpad=self.scratchpad
-            )
+            if callback.condition in (CallbackCondition.ALWAYS, outcome):
+                await self._handle_callback(callback, exc)
 
     async def _handle_callback(
-        self, func: Callable[..., Awaitable[None]], *args, **kwargs
-    ):
+        self, func: Callback[Self], exc: Exception | None
+    ) -> None:
         """
         Wrapper to handle callback gracefully
         """
         try:
-            await func(*args, **kwargs)
+            await func(self, exc)
         except Exception as err:
             logger.exception("Agent encountered an error during callback evaluation")
+            # TODO: This output needs to be more traceable/indexable than a list
             self.callback_output.append(err)
 
     async def __call__(self, *args: Any, **kwargs: Any) -> Any:
@@ -409,7 +427,7 @@ def __init__(
         stopping_condition: _StoppingCondition | None = None,
         provider: _Provider[Any] | None = None,
         tools: Sequence[Tool[Any] | ToolDefinition] | None = None,
-        callbacks: Sequence[Callable[..., Any]] | None = None,
+        callbacks: Sequence[Callback[Self]] | None = None,
         oai_kwargs: dict[str, Any] | None = None,
         **fmt_kwargs: Any,
     ) -> None:
@@ -424,7 +442,7 @@ def __init__(
         :param str model_name: Name of model to use (or deployment name for AzureOpenAI) (optional if provider is passed)
         :param Type[_Provider] provider: Instantiated OpenAI instance to use (optional)
         :param tools: Executable tools or legacy definitions backed by agent methods (optional)
-        :param List[Callable] callbacks: List of callbacks to evaluate at end of run (optional)
+        :param callbacks: Callbacks dispatched according to their condition (optional)
         :param dict[str, any] oai_kwargs: Dict of additional OpenAI arguments to pass thru to chat call
         :param fmt_kwargs: Additional named arguments which will be inserted into the :func:`BASE_PROMPT` via fstring
 

agents/agent/prediction.py

@@ -4,12 +4,12 @@
 """
 
 import logging
-from collections.abc import Callable, Sequence
-from typing import Any, Literal
+from collections.abc import Sequence
+from typing import Any, Literal, Self
 
 import pydantic
 
-from ..abstract import _Provider, _StoppingCondition
+from ..abstract import Callback, _Provider, _StoppingCondition
 from ..json_tool_gen import Tool, ToolDefinition
 from .base import StructuredOutputAgent
 
@@ -34,7 +34,7 @@ def __init__(
         model_name: str | None = None,
         provider: _Provider | None = None,
         tools: Sequence[Tool[Any] | ToolDefinition] | None = None,
-        callbacks: list[Callable] | None = None,
+        callbacks: Sequence[Callback[Self]] | None = None,
         oai_kwargs: dict[str, Any] | None = None,
         **fmt_kwargs,
     ):
@@ -48,7 +48,7 @@ def __init__(
         :param str model_name: Name of model to use (or deployment name for AzureOpenAI) (optional if provider is passed)
         :param _Provider provider: Instantiated OpenAI instance to use (optional)
         :param tools: Executable tools or legacy definitions backed by agent methods (optional)
-        :param List[Callable] callbacks: List of callbacks to evaluate at end of run (optional)
+        :param callbacks: Callbacks dispatched according to their condition (optional)
         :param dict[str, any] oai_kwargs: Dict of additional OpenAI arguments to pass thru to chat call
         :param fmt_kwargs: Additional named arguments which will be inserted into the :func:`BASE_PROMPT` via fstring
         """

agents/callbacks.py

@@ -4,15 +4,17 @@
 
 from typing import Any
 
-from .abstract import _Callback
+from .abstract import Callback, CallbackCondition
 from .agent import Agent
 
 
-class AgentCallback[CallbackAgentT: Agent](_Callback[Agent]):
+class AgentCallback[CallbackAgentT: Agent](Callback[Agent]):
     """
     Call another agent with the answer and scratchpad of a completed agent
     """
 
+    condition = CallbackCondition.ALWAYS
+
     callback_agent: CallbackAgentT
 
     def __init__(self, agent_class: type[CallbackAgentT], **agent_kwargs: Any) -> None:
@@ -32,13 +34,13 @@ def __init__(self, agent_class: type[CallbackAgentT], **agent_kwargs: Any) -> No
         self.agent_class = agent_class
         self.agent_kwargs = agent_kwargs
 
-    async def __call__(self, cls: Agent, answer: Any, scratchpad: str) -> None:
+    async def __call__(self, cls: Agent, exc: Exception | None) -> None:
         """
         Run new callback agent on calling agent's answer and scratchpad and append output.
         """
 
         self.callback_agent = self.agent_class(
-            **self.agent_kwargs, answer=answer, scratchpad=scratchpad
+            **self.agent_kwargs, answer=cls.answer, scratchpad=cls.scratchpad
         )
         await self.callback_agent.run()
         cls.callback_output.append(self.callback_agent.answer)

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "agents"
-version = "0.4.2"
+version = "0.4.3"
 description = "Yet another langchain-esque package to use language agents and compose agent-systems"
 readme = "README.md"
 requires-python = ">=3.12"