Callbacks Refactor [base] (#3256)

Co-authored-by: Nuno Campos <nuno@boringbits.io>
Co-authored-by: Davis Chase <130488702+dev2049@users.noreply.github.com>
Co-authored-by: Zander Chase <130414180+vowelparrot@users.noreply.github.com>
Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>

tests/unit_tests/callbacks/fake_callback_handler.py

@@ -1,10 +1,9 @@
 """A fake callback handler for testing purposes."""
-from typing import Any, Dict, List, Union
+from typing import Any
 
 from pydantic import BaseModel
 
 from langchain.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler
-from langchain.schema import AgentAction, AgentFinish, LLMResult
 
 
 class BaseFakeCallbackHandler(BaseModel):
@@ -17,12 +16,72 @@ class BaseFakeCallbackHandler(BaseModel):
     ignore_llm_: bool = False
     ignore_chain_: bool = False
     ignore_agent_: bool = False
-    always_verbose_: bool = False
 
-    @property
-    def always_verbose(self) -> bool:
-        """Whether to call verbose callbacks even if verbose is False."""
-        return self.always_verbose_
+    # add finer-grained counters for easier debugging of failing tests
+    chain_starts: int = 0
+    chain_ends: int = 0
+    llm_starts: int = 0
+    llm_ends: int = 0
+    llm_streams: int = 0
+    tool_starts: int = 0
+    tool_ends: int = 0
+    agent_actions: int = 0
+    agent_ends: int = 0
+
+
+class BaseFakeCallbackHandlerMixin(BaseFakeCallbackHandler):
+    """Base fake callback handler mixin for testing."""
+
+    def on_llm_start_common(self) -> None:
+        self.llm_starts += 1
+        self.starts += 1
+
+    def on_llm_end_common(self) -> None:
+        self.llm_ends += 1
+        self.ends += 1
+
+    def on_llm_error_common(self) -> None:
+        self.errors += 1
+
+    def on_llm_new_token_common(self) -> None:
+        self.llm_streams += 1
+
+    def on_chain_start_common(self) -> None:
+        self.chain_starts += 1
+        self.starts += 1
+
+    def on_chain_end_common(self) -> None:
+        self.chain_ends += 1
+        self.ends += 1
+
+    def on_chain_error_common(self) -> None:
+        self.errors += 1
+
+    def on_tool_start_common(self) -> None:
+        self.tool_starts += 1
+        self.starts += 1
+
+    def on_tool_end_common(self) -> None:
+        self.tool_ends += 1
+        self.ends += 1
+
+    def on_tool_error_common(self) -> None:
+        self.errors += 1
+
+    def on_agent_action_common(self) -> None:
+        self.agent_actions += 1
+        self.starts += 1
+
+    def on_agent_finish_common(self) -> None:
+        self.agent_ends += 1
+        self.ends += 1
+
+    def on_text_common(self) -> None:
+        self.text += 1
+
+
+class FakeCallbackHandler(BaseCallbackHandler, BaseFakeCallbackHandlerMixin):
+    """Fake callback handler for testing."""
 
     @property
     def ignore_llm(self) -> bool:
@@ -39,164 +98,209 @@ class BaseFakeCallbackHandler(BaseModel):
         """Whether to ignore agent callbacks."""
         return self.ignore_agent_
 
-    # add finer-grained counters for easier debugging of failing tests
-    chain_starts: int = 0
-    chain_ends: int = 0
-    llm_starts: int = 0
-    llm_ends: int = 0
-    llm_streams: int = 0
-    tool_starts: int = 0
-    tool_ends: int = 0
-    agent_ends: int = 0
-
-
-class FakeCallbackHandler(BaseFakeCallbackHandler, BaseCallbackHandler):
-    """Fake callback handler for testing."""
-
     def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
-    ) -> None:
-        """Run when LLM starts running."""
-        self.llm_starts += 1
-        self.starts += 1
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_llm_start_common()
 
-    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Run when LLM generates a new token."""
-        self.llm_streams += 1
+    def on_llm_new_token(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_llm_new_token_common()
 
-    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        """Run when LLM ends running."""
-        self.llm_ends += 1
-        self.ends += 1
+    def on_llm_end(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_llm_end_common()
 
     def on_llm_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        """Run when LLM errors."""
-        self.errors += 1
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_llm_error_common()
 
     def on_chain_start(
-        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
-    ) -> None:
-        """Run when chain starts running."""
-        self.chain_starts += 1
-        self.starts += 1
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_chain_start_common()
 
-    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
-        """Run when chain ends running."""
-        self.chain_ends += 1
-        self.ends += 1
+    def on_chain_end(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_chain_end_common()
 
     def on_chain_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        """Run when chain errors."""
-        self.errors += 1
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_chain_error_common()
 
     def on_tool_start(
-        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
-    ) -> None:
-        """Run when tool starts running."""
-        self.tool_starts += 1
-        self.starts += 1
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_tool_start_common()
 
-    def on_tool_end(self, output: str, **kwargs: Any) -> None:
-        """Run when tool ends running."""
-        self.tool_ends += 1
-        self.ends += 1
+    def on_tool_end(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_tool_end_common()
 
     def on_tool_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
-    ) -> None:
-        """Run when tool errors."""
-        self.errors += 1
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_tool_error_common()
 
-    def on_text(self, text: str, **kwargs: Any) -> None:
-        """Run when agent is ending."""
-        self.text += 1
+    def on_agent_action(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_agent_action_common()
 
-    def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
-        """Run when agent ends running."""
-        self.agent_ends += 1
-        self.ends += 1
+    def on_agent_finish(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_agent_finish_common()
 
-    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
-        """Run on agent action."""
-        self.tool_starts += 1
-        self.starts += 1
+    def on_text(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        self.on_text_common()
+
+    def __deepcopy__(self, memo: dict) -> "FakeCallbackHandler":
+        return self
 
 
-class FakeAsyncCallbackHandler(BaseFakeCallbackHandler, AsyncCallbackHandler):
+class FakeAsyncCallbackHandler(AsyncCallbackHandler, BaseFakeCallbackHandlerMixin):
     """Fake async callback handler for testing."""
 
+    @property
+    def ignore_llm(self) -> bool:
+        """Whether to ignore LLM callbacks."""
+        return self.ignore_llm_
+
+    @property
+    def ignore_chain(self) -> bool:
+        """Whether to ignore chain callbacks."""
+        return self.ignore_chain_
+
+    @property
+    def ignore_agent(self) -> bool:
+        """Whether to ignore agent callbacks."""
+        return self.ignore_agent_
+
     async def on_llm_start(
-        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+        self,
+        *args: Any,
+        **kwargs: Any,
     ) -> None:
-        """Run when LLM starts running."""
-        self.llm_starts += 1
-        self.starts += 1
+        self.on_llm_start_common()
 
-    async def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
-        """Run when LLM generates a new token."""
-        self.llm_streams += 1
+    async def on_llm_new_token(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        self.on_llm_new_token_common()
 
-    async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
-        """Run when LLM ends running."""
-        self.llm_ends += 1
-        self.ends += 1
+    async def on_llm_end(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        self.on_llm_end_common()
 
     async def on_llm_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+        self,
+        *args: Any,
+        **kwargs: Any,
     ) -> None:
-        """Run when LLM errors."""
-        self.errors += 1
+        self.on_llm_error_common()
 
     async def on_chain_start(
-        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
+        self,
+        *args: Any,
+        **kwargs: Any,
     ) -> None:
-        """Run when chain starts running."""
-        self.chain_starts += 1
-        self.starts += 1
+        self.on_chain_start_common()
 
-    async def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
-        """Run when chain ends running."""
-        self.chain_ends += 1
-        self.ends += 1
+    async def on_chain_end(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        self.on_chain_end_common()
 
     async def on_chain_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+        self,
+        *args: Any,
+        **kwargs: Any,
     ) -> None:
-        """Run when chain errors."""
-        self.errors += 1
+        self.on_chain_error_common()
 
     async def on_tool_start(
-        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
+        self,
+        *args: Any,
+        **kwargs: Any,
     ) -> None:
-        """Run when tool starts running."""
-        self.tool_starts += 1
-        self.starts += 1
+        self.on_tool_start_common()
 
-    async def on_tool_end(self, output: str, **kwargs: Any) -> None:
-        """Run when tool ends running."""
-        self.tool_ends += 1
-        self.ends += 1
+    async def on_tool_end(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        self.on_tool_end_common()
 
     async def on_tool_error(
-        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
+        self,
+        *args: Any,
+        **kwargs: Any,
     ) -> None:
-        """Run when tool errors."""
-        self.errors += 1
+        self.on_tool_error_common()
 
-    async def on_text(self, text: str, **kwargs: Any) -> None:
-        """Run when agent is ending."""
-        self.text += 1
+    async def on_agent_action(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        self.on_agent_action_common()
 
-    async def on_agent_finish(self, finish: AgentFinish, **kwargs: Any) -> None:
-        """Run when agent ends running."""
-        self.agent_ends += 1
-        self.ends += 1
+    async def on_agent_finish(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        self.on_agent_finish_common()
 
-    async def on_agent_action(self, action: AgentAction, **kwargs: Any) -> None:
-        """Run on agent action."""
-        self.tool_starts += 1
-        self.starts += 1
+    async def on_text(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        self.on_text_common()
+
+    def __deepcopy__(self, memo: dict) -> "FakeAsyncCallbackHandler":
+        return self