Upload trace_tool.py

examples/tools/trace/trace_tool.py

@@ -0,0 +1,155 @@
+import aworld.trace as trace
+import aworld.trace.instrumentation.semconv as semconv
+from aworld.trace.server import get_trace_server
+from aworld.trace.server.util import build_trace_tree
+from aworld.core.tool.base import Tool, AgentInput, ToolFactory
+from examples.tools.tool_action import GetTraceAction
+from aworld.tools.utils import build_observation
+from aworld.config.conf import ToolConfig
+from aworld.core.common import Observation, ActionModel, ActionResult
+from typing import Tuple, Dict, Any, List
+from aworld.logs.util import logger
+
+
+@ToolFactory.register(name="trace",
+                      desc="Get the trace of the current execution.",
+                      supported_action=GetTraceAction,
+                      conf_file_name=f'trace_tool.yaml')
+class TraceTool(Tool):
+    def __init__(self,
+                 conf: ToolConfig,
+                 **kwargs) -> None:
+        """
+        Initialize the TraceTool
+        Args:
+            conf: tool config
+            **kwargs: -
+        Return:
+            None
+        """
+        super(TraceTool, self).__init__(conf, **kwargs)
+        self.type = "function"
+        self.get_trace_url = self.conf.get('get_trace_url')
+
+    def reset(self,
+              *,
+              seed: int | None = None,
+              options: Dict[str, str] | None = None) -> Tuple[AgentInput, dict[str, Any]]:
+        """
+        Reset the executor
+        Args:
+            seed: -
+            options: -
+        Returns:
+            AgentInput, dict[str, Any]: -
+        """
+        self._finished = False
+        return build_observation(observer=self.name(),
+                                 ability=GetTraceAction.GET_TRACE.value.name), {}
+
+    def close(self) -> None:
+        """
+        Close the executor
+        Returns:
+            None
+        """
+        self._finished = True
+
+    def do_step(self,
+                actions: List[ActionModel],
+                **kwargs) -> Tuple[Observation, float, bool, bool, dict[str, Any]]:
+        reward = 0
+        fail_error = ""
+        observation = build_observation(observer=self.name(),
+                                        ability=GetTraceAction.GET_TRACE.value.name)
+        results = []
+        try:
+            if not actions:
+                return (observation, reward,
+                        kwargs.get("terminated",
+                                   False), kwargs.get("truncated", False), {
+                            "exception": "actions is empty"
+                        })
+            for action in actions:
+                trace_id = action.params.get("trace_id", "")
+                if not trace_id:
+                    current_span = trace.get_current_span()
+                    if current_span:
+                        trace_id = current_span.get_trace_id()
+                if not trace_id:
+                    logger.warning(f"{action} no trace_id to fetch.")
+                    observation.action_result.append(
+                        ActionResult(is_done=True,
+                                     success=False,
+                                     content="",
+                                     error="no trace_id to fetch",
+                                     keep=False))
+                    continue
+                try:
+                    trace_data = self.fetch_trace_data(trace_id)
+                    logger.info(f"trace_data={trace_data}")
+                    error = ""
+                except Exception as e:
+                    error = str(e)
+                results.append(trace_data)
+                observation.action_result.append(
+                    ActionResult(is_done=True,
+                                 success=False if error else True,
+                                 content=f"{trace_data}",
+                                 error=f"{error}",
+                                 keep=False))
+
+            observation.content = f"{results}"
+            reward = 1
+        except Exception as e:
+            fail_error = str(e)
+        finally:
+            self._finished = True
+
+        info = {"exception": fail_error}
+        info.update(kwargs)
+        return (observation, reward, kwargs.get("terminated", False),
+                kwargs.get("truncated", False), info)
+
+    def fetch_trace_data(self, trace_id=None):
+        '''
+            fetch trace data from trace server.
+            return trace data, like:
+            {
+                'trace_id': trace_id,
+                'root_span': [],
+            }
+        '''
+        try:
+            if trace_id:
+                trace_server = get_trace_server()
+                if not trace_server:
+                    logger.error("No memory trace server has been set.")
+                else:
+                    trace_storage = trace_server.get_storage()
+                    spans = trace_storage.get_all_spans(trace_id)
+                    if spans:
+                        return self.proccess_trace(build_trace_tree(spans))
+            return {"trace_id": trace_id, "root_span": []}
+        except Exception as e:
+            logger.error(f"Error fetching trace data: {e}")
+            return {"trace_id": trace_id, "root_span": []}
+
+    def proccess_trace(self, trace_data):
+        root_spans = trace_data.get("root_span")
+        for span in root_spans:
+            self.choose_attribute(span)
+        return trace_data
+
+    def choose_attribute(self, span):
+        include_attr = [semconv.GEN_AI_USAGE_INPUT_TOKENS,
+                        semconv.GEN_AI_USAGE_OUTPUT_TOKENS, semconv.GEN_AI_USAGE_TOTAL_TOKENS]
+        result_attributes = {}
+        origin_attributes = span.get("attributes") or {}
+        for key, value in origin_attributes.items():
+            if key in include_attr:
+                result_attributes[key] = value
+        span["attributes"] = result_attributes
+        if span.get("children"):
+            for child in span.get("children"):
+                self.choose_attribute(child)