dm-aioaiagent 0.5.0__tar.gz → 0.5.2__tar.gz

{dm_aioaiagent-0.5.0 → dm_aioaiagent-0.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dm-aioaiagent
-Version: 0.5.0
+Version: 0.5.2
 Summary: This is my custom aioaiagent client
 Home-page: https://pypi.org/project/dm-aioaiagent
 Author: dimka4621

{dm_aioaiagent-0.5.0 → dm_aioaiagent-0.5.2}/dm_aioaiagent/ai_agent.py

@@ -28,6 +28,8 @@ class DMAIAgent:
         # tools
         tools: list[BaseTool] = None,
         parallel_tool_calls: bool = None,
+        before_tool_call_callback: BeforeToolCallCallback = None,
+        after_tool_call_callback: AfterToolCallCallback = None,
         # memory
         is_memory_enabled: bool = True,
         save_tools_responses_in_memory: bool = True,
@@ -55,6 +57,8 @@ class DMAIAgent:
         self._tools = tools or []
         self._is_tools_exists = bool(tools)
         self._parallel_tool_calls = parallel_tool_calls
+        self._before_tool_call_callback = before_tool_call_callback
+        self._after_tool_call_callback = after_tool_call_callback
         # memory
         self._memory_messages = []
         self._is_memory_enabled = bool(is_memory_enabled)
@@ -170,13 +174,25 @@ class DMAIAgent:
             tool_args = tool_call["args"]
 
             def tool_callback(tool_id=tool_id, tool_name=tool_name, tool_args=tool_args) -> None:
-                self._logger.debug("Invoke tool", tool_id=tool_id, tool_name=tool_name, tool_args=tool_args)
                 if tool_name in self._tool_map:
                     try:
+                        if self._before_tool_call_callback:
+                            self._before_tool_call_callback(tool_name, tool_args)
+                    except Exception as e:
+                        self._logger.error(str(e), callback_type="before_tool_call")
+
+                    try:
+                        self._logger.debug("Invoke tool", tool_id=tool_id, tool_name=tool_name, tool_args=tool_args)
                         tool_response = self._tool_map[tool_name].run(tool_args)
                     except Exception as e:
-                        self._logger.error(e, tool_id=tool_id)
+                        self._logger.error(str(e), tool_id=tool_id)
                         tool_response = "Tool executed with an error!"
+
+                    try:
+                        if self._after_tool_call_callback:
+                            self._after_tool_call_callback(tool_name, tool_args, tool_response)
+                    except Exception as e:
+                        self._logger.error(str(e), callback_type="after_tool_call")
                 else:
                     tool_response = f"Tool not found!"
                 self._logger.debug(f"Tool response:\n{tool_response}", tool_id=tool_id)

{dm_aioaiagent-0.5.0 → dm_aioaiagent-0.5.2}/dm_aioaiagent/async_ai_agent.py

@@ -8,8 +8,19 @@ from .types import *
 
 
 class DMAioAIAgent(DMAIAgent):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(
+        self,
+        *args,
+        before_tool_call_callback: AsyncBeforeToolCallCallback = None,
+        after_tool_call_callback: AsyncAfterToolCallCallback = None,
+        **kwargs
+    ):
+        super().__init__(
+            *args,
+            before_tool_call_callback=before_tool_call_callback,
+            after_tool_call_callback=after_tool_call_callback,
+            **kwargs
+        )
 
     async def run(self, query: str, *args, **kwargs) -> Union[str, TypedDict, BaseModel]:
         new_messages = await self.run_messages(messages=[TextMessage(role="user", content=query)], *args, **kwargs)
@@ -79,13 +90,25 @@ class DMAioAIAgent(DMAIAgent):
             tool_args = tool_call["args"]
 
             async def tool_callback(tool_id=tool_id, tool_name=tool_name, tool_args=tool_args) -> None:
-                self._logger.debug("Invoke tool", tool_id=tool_id, tool_name=tool_name, tool_args=tool_args)
                 if tool_name in self._tool_map:
                     try:
+                        if self._before_tool_call_callback:
+                            await self._before_tool_call_callback(tool_name, tool_args)
+                    except Exception as e:
+                        self._logger.error(str(e), callback_type="before_tool_call")
+
+                    try:
+                        self._logger.debug("Invoke tool", tool_id=tool_id, tool_name=tool_name, tool_args=tool_args)
                         tool_response = await self._tool_map[tool_name].arun(tool_args)
                     except Exception as e:
-                        self._logger.error(e, tool_id=tool_id)
+                        self._logger.error(str(e), tool_id=tool_id)
                         tool_response = "Tool executed with an error!"
+
+                    try:
+                        if self._after_tool_call_callback:
+                            await self._after_tool_call_callback(tool_name, tool_args, tool_response)
+                    except Exception as e:
+                        self._logger.error(str(e), callback_type="after_tool_call")
                 else:
                     tool_response = f"Tool '{tool_name}' not found!"
                 self._logger.debug(f"Tool response:\n{tool_response}", tool_id=tool_id)

{dm_aioaiagent-0.5.0 → dm_aioaiagent-0.5.2}/dm_aioaiagent/types.py

@@ -1,10 +1,13 @@
-from typing import Literal, Union, Type
-from typing_extensions import TypedDict
+from typing import Literal, Union, Type, TypedDict, Callable, Coroutine, Any
 from pydantic import BaseModel
 from langchain_core.messages import BaseMessage
 
 OutputSchemaType = Union[Type[TypedDict], Type[BaseModel], None]
 
+BeforeToolCallCallback = Callable[[str, dict], None]
+AfterToolCallCallback = Callable[[str, dict, str], None]
+AsyncBeforeToolCallCallback = Callable[[str, dict], Coroutine[Any, Any, None]]
+AsyncAfterToolCallCallback = Callable[[str, dict, str], Coroutine[Any, Any, None]]
 
 class ImageMessageTextItem(TypedDict):
     type: Literal['text']

{dm_aioaiagent-0.5.0 → dm_aioaiagent-0.5.2}/dm_aioaiagent.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dm-aioaiagent
-Version: 0.5.0
+Version: 0.5.2
 Summary: This is my custom aioaiagent client
 Home-page: https://pypi.org/project/dm-aioaiagent
 Author: dimka4621

{dm_aioaiagent-0.5.0 → dm_aioaiagent-0.5.2}/setup.py

@@ -8,7 +8,7 @@ def readme():
 
 setup(
     name='dm-aioaiagent',
-    version='v0.5.0',
+    version='v0.5.2',
     author='dimka4621',
     author_email='mismartconfig@gmail.com',
     description='This is my custom aioaiagent client',