feat: Add trace query tools. (#14)

1 year ago · 43e10c64fd
parent 292d19fbfd
commit 43e10c64fd
8 changed files with 401 additions and 1 deletions
--- a/api/core/tools/builtin_tool/providers/apo_select/tools/fault_log.py
+++ b/api/core/tools/builtin_tool/providers/apo_select/tools/fault_log.py
@ -21,12 +21,14 @@ class FaultLogTool(BuiltinTool):
        service = tool_parameters.get("service")
        start_time = tool_parameters.get("startTime")
        end_time = tool_parameters.get("endTime")
+        pod = tool_parameters.get("pod")
        params = {
          'service': [service],
          'startTime': start_time,
          'endTime': end_time,
          'pageNum': 1,
-          'pageSize': 10
+          'pageSize': 10,
+          'pod': pod,
          }
        url = dify_config.APO_BACKEND_URL + "/api/log/fault/pagelist" 
        resp = requests.post(url, json=params)
--- a/api/core/tools/builtin_tool/providers/apo_select/tools/fault_log.yaml
+++ b/api/core/tools/builtin_tool/providers/apo_select/tools/fault_log.yaml
@ -29,6 +29,19 @@ parameters:
      pt_BR: service name
    llm_description: service name
    form: llm
+  - name: pod
+    type: string
+    required: false
+    label:
+      en_US: pod
+      zh_Hans: pod
+      pt_BR: pod
+    human_description:
+      en_US: pod name
+      zh_Hans: 查询的日志pod名
+      pt_BR: pod name
+    llm_description: pod name
+    form: llm
  - name: startTime
    type: number
    required: true
--- a/api/core/tools/builtin_tool/providers/apo_select/tools/query_span_consumption.py
+++ b/api/core/tools/builtin_tool/providers/apo_select/tools/query_span_consumption.py
@ -0,0 +1,43 @@
+import json
+from collections.abc import Generator
+from typing import Any, Optional
+
+import requests
+
+from configs import dify_config
+from core.tools.builtin_tool.tool import BuiltinTool
+from core.tools.entities.tool_entities import ToolInvokeMessage
+
+
+class QuerySpanTool(BuiltinTool):
+    def _invoke(
+        self,
+        user_id: str,
+        tool_parameters: dict[str, Any],
+        conversation_id: Optional[str] = None,
+        app_id: Optional[str] = None,
+        message_id: Optional[str] = None,
+    ) -> Generator[ToolInvokeMessage, None, None]:
+        start_time = tool_parameters.get("startTime")
+        end_time = tool_parameters.get("endTime")
+        node_name = tool_parameters.get("nodeName")
+        pid = tool_parameters.get("pid")
+
+        params = {
+            'startTime': start_time,
+            'endTime': end_time,
+            'nodeName': node_name,
+            'pid': pid
+        }
+        
+        content_url = dify_config.APO_BACKEND_URL + "/api/trace/onoffcpu"
+        response = requests.get(
+            content_url,
+            params=params,
+        ).json()
+        list = json.dumps({
+            'type': 'span',
+            'display': True,
+            'data': response,
+        })
+        yield self.create_text_message(list)
--- a/api/core/tools/builtin_tool/providers/apo_select/tools/query_span_consumption.yaml
+++ b/api/core/tools/builtin_tool/providers/apo_select/tools/query_span_consumption.yaml
@ -0,0 +1,62 @@
+identity:
+  name: 查询span消耗
+  author: APO
+  label:
+    en_US: Query span consumption.
+    zh_Hans: 查询span消耗
+description:
+  human:
+    en_US: Query span consumption.
+    zh_Hans: 查询span消耗
+  llm: Query span consumption.
+display:
+  type: span
+  title: span
+  unit: span
+parameters:
+  - name: pid
+    type: string
+    required: true
+    label:
+      en_US: Process ID
+      zh_Hans: 进程ID
+    human_description:
+      en_US: Process ID
+      zh_Hans: 进程ID
+    llm_description: Process ID
+    form: llm
+  - name: nodeName
+    type: string
+    required: true
+    label:
+      en_US: Node Name
+      zh_Hans: 主机名
+    human_description:
+      en_US: Node name
+      zh_Hans: 主机名
+    llm_description: Node name
+    form: llm
+
+  - name: startTime
+    type: number
+    required: true
+    label:
+      en_US: Start Time
+      zh_Hans: 开始时间
+    human_description:
+      en_US: Start timestamp in microseconds
+      zh_Hans: 查询开始时间（微秒）
+    llm_description: Microsecond timestamp for start time
+    form: llm
+
+  - name: endTime
+    type: number
+    required: true
+    label:
+      en_US: End Time
+      zh_Hans: 结束时间
+    human_description:
+      en_US: End timestamp in microseconds
+      zh_Hans: 查询结束时间（微秒）
+    llm_description: Microsecond timestamp for end time
+    form: llm
--- a/api/core/tools/builtin_tool/providers/apo_select/tools/query_trace.py
+++ b/api/core/tools/builtin_tool/providers/apo_select/tools/query_trace.py
@ -0,0 +1,67 @@
+import json
+from collections.abc import Generator
+from typing import Any, Optional
+
+import requests
+
+from configs import dify_config
+from core.tools.builtin_tool.tool import BuiltinTool
+from core.tools.entities.tool_entities import ToolInvokeMessage
+
+
+class QueryTraceTool(BuiltinTool):
+    def _invoke(
+        self,
+        user_id: str,
+        tool_parameters: dict[str, Any],
+        conversation_id: Optional[str] = None,
+        app_id: Optional[str] = None,
+        message_id: Optional[str] = None,
+    ) -> Generator[ToolInvokeMessage, None, None]:
+        start_time = tool_parameters.get("startTime")
+        end_time = tool_parameters.get("endTime")
+        service = tool_parameters.get("service")
+        namespace = tool_parameters.get("namespace")
+        endpoint = tool_parameters.get("endpoint")
+        instance = tool_parameters.get("instance")
+        node_name = tool_parameters.get("nodeName")
+        container_id = tool_parameters.get("containerId")
+        pid = tool_parameters.get("pid")
+        trace_id = tool_parameters.get("traceId")
+        page_num = tool_parameters.get("pageNum")
+        page_size = tool_parameters.get("pageSize")
+
+        params = {}
+        if service:
+            params["service"] = [service]
+        if start_time:
+            params["startTime"] = start_time
+        if end_time:
+            params["endTime"] = end_time
+        if namespace:
+            params["namespace"] = [namespace]
+        if endpoint:
+            params["endpoint"] = endpoint
+        if instance:
+            params["instance"] = instance
+        if node_name:
+            params["nodeName"] = node_name
+        if container_id:
+            params["containerId"] = container_id
+        if pid:
+            params["pid"] = pid
+        if trace_id:
+            params["traceId"] = trace_id
+        if page_num:
+            params["pageNum"] = page_num
+        if page_size:
+            params["pageSize"] = page_size
+        
+        content_url = dify_config.APO_BACKEND_URL + "/api/trace/pagelist"
+        content = requests.post(content_url, json=params).json()
+        list = json.dumps({
+            'type': 'trace',
+            'display': True,
+            'data': content,
+        })
+        yield self.create_text_message(list)
--- a/api/core/tools/builtin_tool/providers/apo_select/tools/query_trace.yaml
+++ b/api/core/tools/builtin_tool/providers/apo_select/tools/query_trace.yaml
@ -0,0 +1,148 @@
+identity:
+  name: 查询链路数据
+  author: APO
+  label:
+    en_US: Query trance data
+    zh_Hans: 查询链路数据
+description:
+  human:
+    en_US: Query trace data
+    zh_Hans: 查询链路数据
+  llm: Query trace data
+display:
+  type: trace
+  title: trace
+  unit: trace
+parameters:
+  - name: service
+    type: string
+    required: false
+    label:
+      en_US: service
+      zh_Hans: service
+    human_description:
+      en_US: service name
+      zh_Hans: 服务名
+    llm_description: service name
+    form: llm
+  - name: namespace
+    type: string
+    required: false
+    label:
+      en_US: namespace
+      zh_Hans: namespace
+    human_description:
+      en_US: namespace
+      zh_Hans: 命名空间
+    llm_description: namespace
+    form: llm
+  - name: endpoint
+    type: string
+    required: false
+    label:
+      en_US: endpoint
+      zh_Hans: endpoint
+    human_description:
+      en_US: endpoint name
+      zh_Hans: 服务端点
+    llm_description: endpoint name
+    form: llm
+  - name: instance
+    type: string
+    required: false
+    label:
+      en_US: instance
+      zh_Hans: instance
+    human_description:
+      en_US: instance name
+      zh_Hans: 实例名
+    llm_description: instance name
+    form: llm
+  - name: nodeName
+    type: string
+    required: false
+    label:
+      en_US: nodeName
+      zh_Hans: 主机名
+    human_description:
+      en_US: node name
+      zh_Hans: 服务端点
+    llm_description: node name
+    form: llm
+  - name: containerId
+    type: string
+    required: false
+    label:
+      en_US: containerId
+      zh_Hans: 容器id
+    human_description:
+      en_US: containerId
+      zh_Hans: 服务端点
+    llm_description: node name
+    form: llm
+  - name: pid
+    type: string
+    required: false
+    label:
+      en_US: pid
+      zh_Hans: 进程id
+    human_description:
+      en_US: process id
+      zh_Hans: 进程id
+    llm_description: process id
+    form: llm
+  - name: traceId
+    type: string
+    required: false
+    label:
+      en_US: traceId
+      zh_Hans: traceId
+    human_description:
+      en_US: trace id
+      zh_Hans: trace id
+    llm_description: trace id
+    form: llm
+  - name: pageNum
+    type: number
+    required: false
+    label:
+      en_US: pageNum
+      zh_Hans: 页码
+    human_description:
+      en_US: page number
+      zh_Hans: 页码
+    llm_description: page number
+    form: llm
+  - name: pageSize
+    type: number
+    required: false
+    label:
+      en_US: pageSize
+      zh_Hans: 每页条数
+    human_description:
+      en_US: page size
+      zh_Hans: 每页条数
+    llm_description: page size
+    form: llm
+  - name: startTime
+    type: number
+    required: true
+    label:
+      en_US: startTime
+      zh_Hans: startTime
+    human_description:
+      en_US: Data query start time(Microsecond)
+      zh_Hans: 开始时间 (微秒)
+    llm_description: Data query start time(Microsecond)
+    form: llm
+  - name: endTime
+    type: number
+    required: true
+    label:
+      en_US: endTime
+      zh_Hans: endTime
+    human_description:
+      en_US: Data query end time(Microsecond)
+      zh_Hans: 结束时间 (微秒)
+    llm_description: Data query end time(Microsecond)
+    form: llm
--- a/api/core/tools/builtin_tool/providers/apo_select/tools/query_trace_info.py
+++ b/api/core/tools/builtin_tool/providers/apo_select/tools/query_trace_info.py
@ -0,0 +1,38 @@
+import json
+from collections.abc import Generator
+from typing import Any, Optional
+
+import requests
+
+from configs import dify_config
+from core.tools.builtin_tool.tool import BuiltinTool
+from core.tools.entities.tool_entities import ToolInvokeMessage
+
+
+class QueryTraceInfoTool(BuiltinTool):
+    def _invoke(
+        self,
+        user_id: str,
+        tool_parameters: dict[str, Any],
+        conversation_id: Optional[str] = None,
+        app_id: Optional[str] = None,
+        message_id: Optional[str] = None,
+    ) -> Generator[ToolInvokeMessage, None, None]:
+        trace_id = tool_parameters.get("traceId")
+
+        params = {
+            'traceId': trace_id,
+        }
+        
+        content_url = dify_config.APO_BACKEND_URL + "/api/trace/info"
+        response = requests.get(
+            content_url,
+            params=params 
+        ).json()
+
+        result = json.dumps({
+            'type': 'trace',
+            'display': True,
+            'data': response,
+        })
+        yield self.create_text_message(result)
--- a/api/core/tools/builtin_tool/providers/apo_select/tools/query_trace_info.yaml
+++ b/api/core/tools/builtin_tool/providers/apo_select/tools/query_trace_info.yaml
@ -0,0 +1,27 @@
+identity:
+  name: 查询trace信息
+  author: APO
+  label:
+    en_US: Query trace information.
+    zh_Hans: 查询trace信息
+description:
+  human:
+    en_US: Query trace information.
+    zh_Hans: 查询trace信息
+  llm: Query trace information.
+display:
+  type: trace
+  title: trace
+  unit: trace
+parameters:
+  - name: traceId
+    type: string
+    required: true
+    label:
+      en_US: Trace ID
+      zh_Hans: 链路ID
+    human_description:
+      en_US: Trace ID
+      zh_Hans: 链路ID
+    llm_description: Trace ID
+    form: llm