appbuilder：对client的trace功能实现

appbuilder/__init__.py

@@ -139,6 +139,8 @@ 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
+
 __all__ = [
     'logger',
 
@@ -208,5 +210,7 @@ def get_default_header():
     "assistant",
     "StreamRunContext",
     "AssistantEventHandler",
-    "AssistantStreamManager"
+    "AssistantStreamManager",
+
+    "AppBuilderTracer"
 ]

appbuilder/core/_session.py

@@ -15,6 +15,7 @@
 import requests
 import json
 from appbuilder.utils.logger_util import logger
+from appbuilder.utils.trace.tracer_wrapper import session_post
 
 class InnerSession(requests.sessions.Session):
 
@@ -48,3 +49,19 @@ def send(self, request, **kwargs):
         """
         logger.debug("Curl Command:\n" + self.build_curl(request) + "\n")
         return super(InnerSession, self).send(request, **kwargs)
+
+    @session_post
+    def post(self, url, data=None, json=None, **kwargs):
+        return super().post(url = url, data=data, json=json, **kwargs)
+
+    @session_post
+    def delete(self, url, **kwargs):
+        return super().delete(url = url, **kwargs)
+
+    @session_post
+    def get(self, url, **kwargs):
+        return super().get(url = url, **kwargs)
+
+    @session_post
+    def put(self, url, data=None, **kwargs):
+        return super().put(url = url, data=data, **kwargs)

appbuilder/core/console/appbuilder_client/appbuilder_client.py

@@ -23,8 +23,10 @@
 from appbuilder.core._client import HTTPClient
 from appbuilder.utils.func_utils import deprecated
 from appbuilder.utils.logger_util import logger
+from appbuilder.utils.trace.tracer_wrapper import client_run_trace,client_tool_trace
 
 
+@client_tool_trace
 def get_app_list(limit: int = 10, after: str = "", before: str = "", secret_key: Optional[str] = None, gateway_v2: Optional[str] = None) -> list[data_class.AppOverview]:
     """
     该接口查询用户下状态为已发布的应用列表
@@ -97,6 +99,7 @@ def __init__(self, app_id: str, **kwargs):
                              " to get a valid app_id after your application is published.")
         self.app_id = app_id
 
+    @client_tool_trace
     def create_conversation(self) -> str:
         r"""创建会话并返回会话ID，会话ID在服务端用于上下文管理、绑定会话文档等，如需开始新的会话，请创建并使用新的会话ID
                 参数:
@@ -114,6 +117,7 @@ def create_conversation(self) -> str:
         resp = data_class.CreateConversationResponse(**data)
         return resp.conversation_id
 
+    @client_tool_trace
     def upload_local_file(self, conversation_id, local_file_path: str) -> str:
         r"""上传文件并将文件与会话ID进行绑定，后续可使用该文件ID进行对话，目前仅支持上传xlsx、jsonl、pdf、png等文件格式
             该接口用于在对话中上传文件供大模型处理，文件的有效期为7天并且不超过对话的有效期。一次只能上传一个文件。
@@ -143,6 +147,7 @@ def upload_local_file(self, conversation_id, local_file_path: str) -> str:
         resp = data_class.FileUploadResponse(**data)
         return resp.id
 
+    @client_run_trace
     def run(self, conversation_id: str,
             query: str,
             file_ids: list = [],

appbuilder/core/console/dataset/dataset.py

@@ -6,6 +6,7 @@
 import os
 from appbuilder.utils.func_utils import deprecated
 from appbuilder.utils.logger_util import logger
+from appbuilder.utils.trace.tracer_wrapper import client_tool_trace
 
 
 class Dataset:
@@ -50,6 +51,7 @@ def http_client(self):
         return self._http_client
 
     @classmethod
+    @client_tool_trace
     def create_dataset(cls, dataset_name: str):
         """
         创建知识库
@@ -70,6 +72,7 @@ def create_dataset(cls, dataset_name: str):
         return Dataset(dataset_id=response["id"], dataset_name=response["name"])
 
     @deprecated
+    @client_tool_trace
     def add_documents(self, file_path_list: List[str], is_custom_process_rule: bool = False,
                       custom_process_rule: Dict = None, is_enhanced: bool = False) -> AddDocumentsResponse:
         """
@@ -133,6 +136,7 @@ def _upload_document(self, file_path: str):
         return res
 
     @deprecated
+    @client_tool_trace
     def delete_documents(self, document_ids: List[str]):
         """
         删除知识库中的文档
@@ -159,6 +163,7 @@ def _delete_document(self, document_id):
         self.http_client.check_console_response(response)
 
     @deprecated
+    @client_tool_trace
     def get_documents(self, page: int, limit: int, keyword: str = "") -> DocumentListResponse:
         """
         获取知识库中的文档列表

appbuilder/core/console/knowledge_base/knowledge_base.py

@@ -21,6 +21,7 @@
 from appbuilder.core._client import HTTPClient
 from appbuilder.core.console.knowledge_base import data_class
 from appbuilder.core.component import Message, Component
+from appbuilder.utils.trace.tracer_wrapper import client_tool_trace
 
 
 class KnowledgeBase(Component):
@@ -31,6 +32,7 @@ def __init__(self, knowledge_id: Optional[str] = None, knowledge_name: Optional[
         self.knowledge_name = knowledge_name
 
     @classmethod
+    @client_tool_trace
     def create_knowledge(cls, knowledge_name: str) -> 'KnowledgeBase':
         payload = json.dumps({"name": knowledge_name})
         http_client = HTTPClient()
@@ -44,6 +46,7 @@ def create_knowledge(cls, knowledge_name: str) -> 'KnowledgeBase':
         response = response.json()["result"]
         return KnowledgeBase(knowledge_id=response["id"], knowledge_name=response["name"])
 
+    @client_tool_trace
     def upload_file(self, file_path: str) -> data_class.KnowledgeBaseUploadFileResponse:
         if not os.path.exists(file_path):
             raise FileNotFoundError("File {} does not exist".format(file_path))
@@ -69,6 +72,7 @@ def upload_file(self, file_path: str) -> data_class.KnowledgeBaseUploadFileRespo
 
         return resp
 
+    @client_tool_trace
     def add_document(self,
                      content_type: str,
                      file_ids: list[str] = [],
@@ -105,6 +109,7 @@ def add_document(self,
         resp = data_class.KnowledgeBaseAddDocumentResponse(**data)
         return resp
 
+    @client_tool_trace
     def delete_document(self, document_id: str, knowledge_base_id: Optional[str] = None):
         if self.knowledge_id == None and knowledge_base_id == None:
             raise ValueError(
@@ -131,6 +136,7 @@ def delete_document(self, document_id: str, knowledge_base_id: Optional[str] = N
         resp = data_class.KnowledgeBaseDeleteDocumentResponse(**data)
         return resp
 
+    @client_tool_trace
     def get_documents_list(self, limit: int = 10, after: Optional[str] = "", before: Optional[str] = "", knowledge_base_id: Optional[str] = None):
         if self.knowledge_id == None and knowledge_base_id == None:
             raise ValueError(

appbuilder/tests/test_appbuilder_trace.py

@@ -0,0 +1,86 @@
+# Copyright (c) 2024 Baidu, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import requests
+import appbuilder
+import os
+
+from unittest.mock import patch,MagicMock
+
+from appbuilder.utils.trace.tracer import AppBuilderTracer, AppbuilderInstrumentor
+from appbuilder.utils.trace.phoenix_wrapper import runtime_main,stop_phoenix
+from appbuilder.core.console.appbuilder_client import get_app_list
+
+@unittest.skipUnless(os.getenv("TEST_CASE", "UNKNOWN") == "CPU_SERIAL", "")
+class TestAppBuilderTrace(unittest.TestCase):
+    def setUp(self):
+        self.app_id = "2a19f6dd-de02-46d9-841d-ef5c52b00466"
+
+    def test_appbuilder_client_trace(self):
+
+        tracer=AppBuilderTracer(
+            enable_phoenix = True,
+            enable_console = True,
+            )
+
+        tracer.start_trace()
+
+        builder = appbuilder.AppBuilderClient(self.app_id)
+        conversation_id = builder.create_conversation()
+
+        # test stream = True
+        msg = builder.run(conversation_id=conversation_id, query="人参有什么用？",stream=True)
+        for m in msg.content:
+            print(m)
+
+        # test stream = False
+        builder.run(conversation_id=conversation_id, query="人参有什么用？")
+
+        # test get_app_list
+        get_app_list()
+
+
+        tracer.end_trace()
+
+
+    def test_client_trace_function(self):
+        from appbuilder.utils.trace._function import _client_tool_trace_output,_client_tool_trace_output_deep_iterate
+        class Test:
+            test = 'test'
+        _client_tool_trace_output(Test,None)
+
+        _client_tool_trace_output_deep_iterate({},None)
+
+    def test_trace_tracer(self):
+        tracer=AppbuilderInstrumentor()
+        tracer.instrumentation_dependencies()
+        tracer._instrument()
+
+
+    def test_appbuilder_phoenix_run(self):
+
+        runtime_main()
+        with self.assertRaises(TypeError):
+            stop_phoenix()
+
+
+if __name__ == '__main__':
+    unittest.main()
+
+
+
+
+
+

appbuilder/utils/trace/README.md

@@ -0,0 +1,52 @@
+# Appbuilder-SDK TRACE 
+
+## 使用方法
+
+### 启动Phoenix可视化软件
+
+**注意:没有可视化需求的可忽略此步骤**
+
+**如果未安装phoenix，需要先安装phoenix软件(这里提供清华园下载方式)** 
+
+```bash
+python3 -m pip install arize-phoenix==4.5.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
+```
+
+**启动phoenix服务**
+
+```bash
+appbuilder_trace_server
+```
+
+- 使用ctrl+c停止phoenix服务
+
+### 启动Appbuilder-SDK TRACE
+
+```python
+from appbuilder.utils.trace.tracer import AppBuilderTracer
+tracer=AppBuilderTracer(
+    enable_phoenix = True,
+    enable_console = False,
+    )
+```
+
+- enable_phoenix:是否启动phoenix可视化跟踪
+- enable_console:是否将trace信息反馈到控制台
+
+```python
+# 启动trace 
+tracer.start_trace()
+
+os.environ["APPBUILDER_TOKEN"] = "YOUR_APPBUILDER_TOKEN"
+app_id = "YOUR_APPBUILDER_ID"
+
+builder = appbuilder.AppBuilderClient(app_id)
+conversation_id = builder.create_conversation()
+msg = builder.run(conversation_id=conversation_id, query="你可以做什么？",stream=True)
+
+for m in msg.content:
+    print(m)
+
+# 结束trace
+tracer.end_trace()
+```

appbuilder/utils/trace/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2024 Baidu, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.