瀏覽代碼

Update BaseToolkit: add method get_tool

StrayWarrior 1 周之前
父節點
當前提交
c73fca2b11
共有 1 個文件被更改,包括 32 次插入0 次删除
  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.")