Appbuilder-sdk的Trace功能【实现Components&Assistants API的Trace】 (baidubce#407)

* trace-for-assistant-components * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update * update --------- Co-authored-by: yinjiaqi <[email protected]>
C9luster · Jul 18, 2024 · fa905b5 · fa905b5
1 parent 1f4f85c
commit fa905b5
Show file tree

Hide file tree

Showing 66 changed files with 1,363 additions and 78 deletions.
diff --git a/appbuilder/__init__.py b/appbuilder/__init__.py
@@ -139,7 +139,7 @@ def get_default_header():
 from appbuilder.core.assistant.threads.runs import AssistantEventHandler
 from appbuilder.core.assistant.threads.runs import AssistantStreamManager
 
-from appbuilder.utils.trace.tracer import AppBuilderTracer
+from appbuilder.utils.trace.tracer import AppBuilderTracer, AppbuilderInstrumentor
 
 __all__ = [
     'logger',
@@ -212,5 +212,6 @@ def get_default_header():
     "AssistantEventHandler",
     "AssistantStreamManager",
 
-    "AppBuilderTracer"
+    "AppBuilderTracer",
+    "AppbuilderInstrumentor"
 ]
diff --git a/appbuilder/core/assistant/assistants/assistants.py b/appbuilder/core/assistant/assistants/assistants.py
@@ -18,6 +18,7 @@
 from appbuilder.utils.collector import AssistantKeys
 from appbuilder.core._client import AssistantHTTPClient
 from appbuilder.core.assistant.assistants.files import Files
+from appbuilder.utils.trace.tracer_wrapper import assistent_tool_trace
 
 
 class Assistants(object):
@@ -28,7 +29,7 @@ def __init__(self):
     def files(self):
         return Files()
 
-
+    @assistent_tool_trace
     def create(self,
                name: str,
                description: str,
@@ -91,6 +92,7 @@ def create(self,
         Collector().add_to_collection(AssistantKeys.ASSISTANT, resp, resp.id)
         return resp
 
+    @assistent_tool_trace
     def update(self,
                assistant_id: str,
                model: Optional[str],
@@ -157,6 +159,7 @@ def update(self,
         resp = assistant_type.AssistantUpdateResponse(**data)
         return resp
 
+    @assistent_tool_trace
     def list(self,
              limit: Optional[int] = 20,
              order: Optional[str] = "desc",
@@ -203,6 +206,7 @@ def list(self,
         resp = assistant_type.AssistantListResponse(**data)
         return resp
 
+    @assistent_tool_trace
     def query(self,
               assistant_id: Optional[str]) -> assistant_type.AssistantQueryResponse:
         """
@@ -240,6 +244,7 @@ def query(self,
         resp = assistant_type.AssistantQueryResponse(**data)
         return resp
 
+    @assistent_tool_trace
     def delete(self,
                assistant_id: Optional[str]) -> assistant_type.AssistantDeleteResponse:
         """
@@ -277,7 +282,7 @@ def delete(self,
         resp = assistant_type.AssistantDeleteResponse(**data)
         return resp
 
-
+    @assistent_tool_trace
     def mount_files(self,
             assistant_id: Optional[str],
             file_id: Optional[str]
@@ -328,6 +333,7 @@ def mount_files(self,
         resp = assistant_type.AssistantFilesResponse(**data)
         return resp
 
+    @assistent_tool_trace
     def mounted_files_list(self,
                    assistant_id: Optional[str],  
                    limit: Optional[int] = 20,
@@ -374,6 +380,7 @@ def mounted_files_list(self,
         resp = assistant_type.AssistantMountedFilesListResponse(**data)
         return resp
 
+    @assistent_tool_trace
     def unmount_files(self,
                      assistant_id: Optional[str],
                      file_id: Optional[str]

diff --git a/appbuilder/core/assistant/assistants/files.py b/appbuilder/core/assistant/assistants/files.py
@@ -22,11 +22,13 @@
 from appbuilder.core._client import AssistantHTTPClient
 
 from appbuilder.core._exception import AppBuilderServerException,HTTPConnectionException
+from appbuilder.utils.trace.tracer_wrapper import assistent_tool_trace
 
 class Files(object):
     def __init__(self):
         self._http_client = AssistantHTTPClient()
 
+    @assistent_tool_trace
     def create(self, file_path: str, purpose: str = "assistant") -> assistant_type.AssistantFilesCreateResponse:
         """
         上传文件到助理存储中。
@@ -71,7 +73,7 @@ def create(self, file_path: str, purpose: str = "assistant") -> assistant_type.A
         resp = assistant_type.AssistantFilesCreateResponse(**data)
         return resp
 
-        
+    @assistent_tool_trace
     def list(self) -> assistant_type.AssistantFilesListResponse:
         """
         列出存储中的文件列表
@@ -105,7 +107,7 @@ def list(self) -> assistant_type.AssistantFilesListResponse:
         return resp
 
 
-
+    @assistent_tool_trace
     def query(self,
               file_id: str,
             ) -> assistant_type.AssistantFilesQueryResponse:
@@ -150,6 +152,7 @@ def query(self,
             raise ValueError('file_id {} is not exist'.format(file_id))
         return resp 
 
+    @assistent_tool_trace
     def delete(self,
                file_id: str,
             ) -> assistant_type.AssistantFilesDeleteResponse:
@@ -180,6 +183,7 @@ def delete(self,
         resp = assistant_type.AssistantFilesDeleteResponse(**data)
         return resp
 
+    @assistent_tool_trace
     def download(self,
                  file_id:str,
                  file_path:str="", # 要求若文件路径不为空，需要以/结尾，默认下载到当前文件夹
@@ -245,9 +249,7 @@ def download(self,
         except Exception as e:
             raise Exception("出现错误,错误信息{}".format(e))
 
-
-
-
+    @assistent_tool_trace
     def content(self,
                 file_id:str,
                 timeout:Optional[int]=None):

diff --git a/appbuilder/core/assistant/threads/messages/messages.py b/appbuilder/core/assistant/threads/messages/messages.py
@@ -16,12 +16,14 @@
 from appbuilder.core.assistant.type import thread_type
 from appbuilder.core._client import AssistantHTTPClient
 from typing import Optional
+from appbuilder.utils.trace.tracer_wrapper import assistent_tool_trace
 
 
 class Messages(object):
     def __init__(self):
         self._http_client = AssistantHTTPClient()
 
+    @assistent_tool_trace
     def create(self, 
                thread_id: str,
                content: str,
@@ -67,6 +69,7 @@ def create(self,
         response = thread_type.AssistantMessageCreateResponse(**data)
         return response
 
+    @assistent_tool_trace
     def update(self,
                thread_id: str,
                message_id: str,
@@ -107,6 +110,7 @@ def update(self,
         response = thread_type.AssistantMessageUpdateResponse(**data)
         return response  
 
+    @assistent_tool_trace
     def list(self,
             thread_id: str,
             limit: int = 20,
@@ -152,7 +156,7 @@ def list(self,
         response = thread_type.AssistantMessageListResponse(**data)
         return response 
 
-
+    @assistent_tool_trace
     def query(self,
             thread_id:str,
             message_id:str) -> thread_type.AssistantMessageQueryResponse:
@@ -189,7 +193,7 @@ def query(self,
         response = thread_type.AssistantMessageQueryResponse(**data)
         return response  
 
-
+    @assistent_tool_trace
     def files(self,
               thread_id:str,
               message_id:str,

diff --git a/appbuilder/core/assistant/threads/runs/runs.py b/appbuilder/core/assistant/threads/runs/runs.py
@@ -22,6 +22,7 @@
 from appbuilder.core.assistant.type import public_type
 from appbuilder.core._client import AssistantHTTPClient
 from appbuilder.utils.sse_util import SSEClient
+from appbuilder.utils.trace.tracer_wrapper import assistent_tool_trace, assistant_run_trace, assistent_stream_run_trace, assistent_stream_run_with_handler_trace
 
 
 
@@ -32,7 +33,8 @@ def __init__(self) -> None:
     @property
     def steps(self) -> Steps:
         return Steps()
-
+
+    @assistant_run_trace
     def run(self,
             assistant_id: str,
             thread_id: Optional[str] = "",
@@ -234,9 +236,9 @@ def _stream(self,
             stream=True,
             timeout=None
         )
-
         return response
 
+    @assistent_stream_run_trace
     def stream_run(self,
                    assistant_id: str,
                    thread_id: Optional[str] = "",
@@ -302,6 +304,7 @@ def stream_run(self,
         sse_client = SSEClient(response)
         return self._iterate_events(sse_client.events())
 
+    @assistent_stream_run_with_handler_trace
     def stream_run_with_handler(self,
                    assistant_id: str,
                    thread_id: Optional[str] = "",
@@ -375,6 +378,8 @@ def _iterate_events(self, events):
 
             yield result
 
+
+    @assistent_tool_trace
     def submit_tool_outputs(self,
                             run_id: str,
                             thread_id: str,
@@ -416,6 +421,7 @@ def submit_tool_outputs(self,
         resp = thread_type.RunResult(**data)
         return resp
 
+    @assistent_tool_trace
     def cancel(self, run_id: str, thread_id: str) -> thread_type.RunResult:
         """
         取消指定线程的运行
@@ -451,6 +457,7 @@ def cancel(self, run_id: str, thread_id: str) -> thread_type.RunResult:
         resp = thread_type.RunResult(**data)
         return resp
 
+    @assistent_tool_trace
     def list(self, thread_id: str, limit: int = 20,
              order: str = 'desc', after: str = "", before: str = "") -> thread_type.RunListResponse:
         """
@@ -495,6 +502,7 @@ def list(self, thread_id: str, limit: int = 20,
         resp = thread_type.RunListResponse(**data)
         return resp
 
+    @assistent_tool_trace
     def query(self, thread_id: str, run_id: str) -> thread_type.RunResult:
         """
         根据thread_id和run_id，查询run的详情

diff --git a/appbuilder/core/assistant/threads/runs/steps.py b/appbuilder/core/assistant/threads/runs/steps.py
@@ -18,11 +18,13 @@
 from appbuilder.core.assistant.type import assistant_type
 from appbuilder.core._client import AssistantHTTPClient
 from appbuilder.utils.sse_util import SSEClient
+from appbuilder.utils.trace.tracer_wrapper import assistent_tool_trace
 
 class Steps():
     def __init__(self) -> None:
         self._http_client = AssistantHTTPClient()
 
+    @assistent_tool_trace
     def list(self, thread_id: str, run_id: str, limit: int = 20,
                    order: str = 'desc', after: str = "", before: str = "") -> thread_type.RunStepListResponse:
         """
@@ -62,6 +64,7 @@ def list(self, thread_id: str, run_id: str, limit: int = 20,
         resp = thread_type.RunStepListResponse(**data)
         return resp
 
+    @assistent_tool_trace
     def query(self, thread_id: str, run_id: str, step_id: str) -> thread_type.RunStepResult:
         """
         根据thread_id，run_id和step_id，查询对应step的信息

diff --git a/appbuilder/core/assistant/threads/threads.py b/appbuilder/core/assistant/threads/threads.py
@@ -18,7 +18,7 @@
 from appbuilder.core.assistant.threads.messages import Messages
 from appbuilder.core.assistant.threads.runs import Runs
 from appbuilder.core._client import AssistantHTTPClient
-
+from appbuilder.utils.trace.tracer_wrapper import assistent_tool_trace
 class Threads():
     def __init__(self) -> None:
         self._http_client = AssistantHTTPClient()
@@ -31,6 +31,7 @@ def messages(self) -> Messages:
     def runs(self) -> Runs:
         return Runs()
 
+    @assistent_tool_trace
     def create(self, messages: Optional[list[thread_type.AssistantMessage]] = []) -> thread_type.ThreadCreateResponse:
         """
         创建一个新的对话线程。
@@ -69,6 +70,7 @@ def create(self, messages: Optional[list[thread_type.AssistantMessage]] = []) ->
         response = thread_type.ThreadCreateResponse(**data)
         return response
 
+    @assistent_tool_trace
     def query(self,
               thread_id:str)->thread_type.ThreadQueryResponse:
         """
@@ -100,6 +102,7 @@ def query(self,
         response = thread_type.ThreadQueryResponse(**data)
         return response
 
+    @assistent_tool_trace
     def delete(self,
                thread_id:str)->thread_type.ThreadDeleteResponse:
         """
@@ -131,6 +134,7 @@ def delete(self,
         response = thread_type.ThreadDeleteResponse(**data)
         return response
 
+    @assistent_tool_trace
     def update(self,
                thread_id:str ,
                metadata:Optional[dict] ={} )->thread_type.ThreadUpdateResponse:

diff --git a/appbuilder/core/components/animal_recognize/component.py b/appbuilder/core/components/animal_recognize/component.py
@@ -23,6 +23,7 @@
 from appbuilder.core._client import HTTPClient
 from appbuilder.core._exception import AppBuilderServerException
 from typing import Generator, Union
+from appbuilder.utils.trace.tracer_wrapper import components_run_trace, components_run_stream_trace
 
 TOP_NUM = 1
 BAIKE_NUM = 0
@@ -81,6 +82,7 @@ class AnimalRecognition(Component):
     ]
 
     @HTTPClient.check_param
+    @components_run_trace
     def run(self, message: Message, timeout: float = None, retry: int = 0) -> Message:
         r""" 动物识别
 
@@ -141,6 +143,7 @@ def _recognize(
         animalRes.request_id = request_id
         return animalRes
 
+    @components_run_stream_trace
     def tool_eval(
         self,
         name: str,

diff --git a/appbuilder/core/components/asr/component.py b/appbuilder/core/components/asr/component.py
@@ -28,6 +28,7 @@
 from appbuilder.core._client import HTTPClient
 from appbuilder.core.components.asr.model import ShortSpeechRecognitionRequest, ShortSpeechRecognitionResponse, \
     ASRInMsg, ASROutMsg
+from appbuilder.utils.trace.tracer_wrapper import components_run_trace, components_run_stream_trace
 
 DEFAULT_AUDIO_MAX_DURATION = 55 * 1000  # 55s
 # 参考短语音极速版API(https://ai.baidu.com/ai-doc/SPEECH/Jlbxdezuf)
@@ -94,6 +95,7 @@ class ASR(Component):
     ]
 
     @HTTPClient.check_param
+    @components_run_trace
     def run(self, message: Message, audio_format: str = "pcm", rate: int = 16000,
             timeout: float = None, retry: int = 0, **kwargs) -> Message:
         """
@@ -176,6 +178,7 @@ def _check_service_error(request_id: str, data: dict):
                     service_err_message=data["err_msg"]
                 )
 
+    @components_run_stream_trace
     def tool_eval(self, name: str, streaming: bool, **kwargs):
         """
         asr for function call

diff --git a/appbuilder/core/components/dish_recognize/component.py b/appbuilder/core/components/dish_recognize/component.py
@@ -25,6 +25,7 @@
 from appbuilder.core._client import HTTPClient
 from appbuilder.core._exception import AppBuilderServerException
 from appbuilder.core.components.dish_recognize.model import *
+from appbuilder.utils.trace.tracer_wrapper import components_run_trace, components_run_stream_trace
 
 
 class DishRecognition(Component):
@@ -48,6 +49,7 @@ class DishRecognition(Component):
     """
 
     @HTTPClient.check_param
+    @components_run_trace
     def run(self, message: Message, timeout: float = None, retry: int = 0) -> Message:
         """
         根据输入图片进行菜品识别。