Browse Source

Update toolkits: add register_toolkit decorator

StrayWarrior 4 days ago
parent
commit
e2c543fcc5

+ 2 - 0
pqai_agent/toolkit/image_describer.py

@@ -6,11 +6,13 @@ from pqai_agent.chat_service import VOLCENGINE_MODEL_DOUBAO_1_5_VISION_PRO
 from pqai_agent.logging_service import logger
 from pqai_agent.toolkit.base import BaseToolkit
 from pqai_agent.toolkit.function_tool import FunctionTool
+from pqai_agent.toolkit.tool_registry import register_toolkit
 
 # 不同实例间复用cache,但不是很好的实践
 _image_describer_caches = {}
 _cache_mutex = threading.Lock()
 
+@register_toolkit
 class ImageDescriber(BaseToolkit):
     def __init__(self, cache_dir: str = None):
         self.model = VOLCENGINE_MODEL_DOUBAO_1_5_VISION_PRO

+ 2 - 0
pqai_agent/toolkit/message_notifier.py

@@ -3,8 +3,10 @@ from typing import List, Dict
 from pqai_agent.logging_service import logger
 from pqai_agent.toolkit.base import BaseToolkit
 from pqai_agent.toolkit.function_tool import FunctionTool
+from pqai_agent.toolkit.tool_registry import register_toolkit
 
 
+@register_toolkit
 class MessageNotifier(BaseToolkit):
     def __init__(self):
         super().__init__()

+ 3 - 0
pqai_agent/toolkit/pq_video_searcher.py

@@ -3,7 +3,10 @@ import requests
 
 from pqai_agent.toolkit.base import BaseToolkit
 from pqai_agent.toolkit.function_tool import FunctionTool
+from pqai_agent.toolkit.tool_registry import register_toolkit
 
+
+@register_toolkit
 class PQVideoSearcher(BaseToolkit):
     API_URL = "https://vlogapi.piaoquantv.com/longvideoapi/search/userandvideo/list"
     def search_pq_video(self, keywords: List[str]) -> List[Dict]:

+ 2 - 0
pqai_agent/toolkit/search_toolkit.py

@@ -4,8 +4,10 @@ import requests
 
 from pqai_agent.toolkit.base import BaseToolkit
 from pqai_agent.toolkit.function_tool import FunctionTool
+from pqai_agent.toolkit.tool_registry import register_toolkit
 
 
+@register_toolkit
 class SearchToolkit(BaseToolkit):
     r"""A class representing a toolkit for web search.
     """