Parcourir la source

Update BaseToolkit: add method get_tool

StrayWarrior il y a 1 semaine
Parent
commit
c73fca2b11
1 fichiers modifiés avec 32 ajouts et 0 suppressions
  1. 32 0
      pqai_agent/toolkit/base.py

+ 32 - 0
pqai_agent/toolkit/base.py

@@ -3,6 +3,23 @@ import functools
 import threading
 from pqai_agent.toolkit.function_tool import FunctionTool
 
+
+class ToolServiceError(Exception):
+
+    def __init__(self,
+                 exception: Optional[Exception] = None,
+                 code: Optional[str] = None,
+                 message: Optional[str] = None,
+                 extra: Optional[dict] = None):
+        if exception is not None:
+            super().__init__(exception)
+        else:
+            super().__init__(f'\nError code: {code}. Error message: {message}')
+        self.exception = exception
+        self.code = code
+        self.message = message
+        self.extra = extra
+
 def with_timeout(timeout=None):
     r"""Decorator that adds timeout functionality to functions.
 
@@ -100,3 +117,18 @@ class BaseToolkit:
                 representing the functions in the toolkit.
         """
         raise NotImplementedError("Subclasses must implement this method.")
+
+    def get_tool(self, name: str) -> FunctionTool:
+        r"""Returns a FunctionTool object by name.
+
+        Args:
+            name (str): The name of the tool to retrieve.
+
+        Returns:
+            Optional[FunctionTool]: The FunctionTool object if found, else None.
+        """
+        tools = self.get_tools()
+        for tool in tools:
+            if tool.name == name:
+                return tool
+        raise NotImplementedError("Tool not found in the toolkit.")