ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/aworld/trace/server/__init__.py

@@ -0,0 +1,66 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+import threading
+import os
+from fastapi.staticfiles import StaticFiles
+from .routes import setup_routes
+from aworld.logs.util import logger
+from aworld.utils.import_package import import_package
+
+GLOBAL_TRACE_SERVER = None
+
+
+class TraceServer:
+    def __init__(self, storage, port: int = 7079):
+        self._storage = storage
+        self._port = port
+        self._thread = None
+        self.app = None
+        self._started = False
+
+    def start(self):
+        self._thread = threading.Thread(target=self._start_app, daemon=True)
+        self._thread.start()
+        self._started = True
+
+    def join(self):
+        if self._thread:
+            self._thread.join()
+        else:
+            raise Exception("Trace server not started.")
+
+    def get_storage(self):
+        return self._storage
+
+    def is_started(self):
+        return self._started
+
+    def _start_app(self):
+        import_package('uvicorn')  # noqa
+        import uvicorn
+        app = setup_routes(self._storage)
+
+        webui_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../cmd/web/webui")
+        static_path = os.path.join(webui_path, "public")
+        app.mount("/static", StaticFiles(directory=static_path), name="static")
+
+        self.app = app
+        # app.run(port=self._port)
+        uvicorn.run(app, host="0.0.0.0", port=self._port, loop="asyncio")
+
+
+def set_trace_server(storage, port: int = 7079, start_server=False):
+    global GLOBAL_TRACE_SERVER
+    if GLOBAL_TRACE_SERVER is None:
+        GLOBAL_TRACE_SERVER = TraceServer(storage, port)
+    if GLOBAL_TRACE_SERVER.is_started():
+        setup_routes(storage)
+        return
+    if start_server:
+        GLOBAL_TRACE_SERVER.start()
+
+
+def get_trace_server():
+    if GLOBAL_TRACE_SERVER is None:
+        logger.warning("No trace server has been set.")
+    return GLOBAL_TRACE_SERVER

chapter9/gaia-experience/AWorld/aworld/trace/server/routes.py

@@ -0,0 +1,54 @@
+from aworld.trace.opentelemetry.memory_storage import TraceStorage
+from aworld.utils.import_package import import_package
+from aworld.trace.server.util import build_trace_tree
+
+import_package('fastapi')  # noqa
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse, RedirectResponse
+
+app = FastAPI()
+current_storage = None
+routes_setup = False
+
+
+def setup_routes(storage: TraceStorage):
+    global current_storage
+    current_storage = storage
+
+    global routes_setup
+
+    if routes_setup:
+        return app
+
+    @app.get("/")
+    async def root():
+        return RedirectResponse("/static/trace_ui.html")
+
+    @app.get('/api/trace/list')
+    async def traces():
+        trace_data = []
+        for trace_id in current_storage.get_all_traces():
+            spans = current_storage.get_all_spans(trace_id)
+            spans_sorted = sorted(spans, key=lambda x: x.start_time)
+            trace_tree = build_trace_tree(spans_sorted)
+            trace_data.append({
+                'trace_id': trace_id,
+                'root_span': trace_tree,
+            })
+            response = {
+                "data": trace_data
+            }
+        return JSONResponse(content=response)
+
+    @app.get('/api/traces/{trace_id}')
+    async def get_trace(trace_id):
+        spans = current_storage.get_all_spans(trace_id)
+        spans_sorted = sorted(spans, key=lambda x: x.start_time)
+        trace_tree = build_trace_tree(spans_sorted)
+        return JSONResponse(content={
+            'trace_id': trace_id,
+            'root_span': trace_tree,
+        })
+
+    routes_setup = True
+    return app

chapter9/gaia-experience/AWorld/aworld/trace/server/util.py

@@ -0,0 +1,241 @@
+import uuid
+from aworld.logs.util import logger
+from aworld.trace.opentelemetry.memory_storage import SpanModel
+from aworld.trace.constants import RunType, SPAN_NAME_PREFIX_EVENT_AGENT
+from aworld.trace.instrumentation import semconv
+
+
+def build_trace_tree(spans: list[SpanModel]):
+    spans_dict = {span.span_id: span.dict() for span in spans}
+    for span in list(spans_dict.values()):
+        parent_id = span['parent_id'] if span['parent_id'] else None
+        if parent_id:
+            parent_span = spans_dict.get(parent_id)
+            if not parent_span:
+                logger.warning(f"span[{parent_id}] not be exported")
+                parent_span = {
+                    'span_id': parent_id,
+                    'trace_id': span['trace_id'],
+                    'name': 'Pengding-Span',
+                    'start_time': span['start_time'],
+                    'end_time': span['end_time'],
+                    'duration_ms': span['duration_ms'],
+                    'attributes': {},
+                    'status': {},
+                    'parent_id': None,
+                    'run_type': 'OTHER'
+                }
+                spans_dict[parent_id] = parent_span
+            if 'children' not in parent_span:
+                parent_span['children'] = []
+            parent_span['children'].append(span)
+
+    root_spans = [span for span in spans_dict.values()
+                  if span['parent_id'] is None]
+    return root_spans
+
+
+def _get_agent_show_name(span: dict):
+    agent_name_prefix = SPAN_NAME_PREFIX_EVENT_AGENT
+    name = span.get("name")
+    if name and name.startswith(agent_name_prefix):
+        name = name[len(agent_name_prefix):]
+    if name and '---' in name:
+        name = name.split('---', 1)[0]
+    return name
+
+
+def _remove_span_detail(root_spans: list):
+    keys_to_keep = {'span_id', 'show_name', 'task_group_id', 'event_id'}
+    for span in root_spans:
+        keys_to_remove = [key for key in span.keys() if key not in keys_to_keep]
+        for key in keys_to_remove:
+            span.pop(key, None)
+        if 'children' in span:
+            _remove_span_detail(span['children'])
+
+
+def _get_top_task_nodes(spans_dict):
+    task_nodes = [
+        span for span in spans_dict.values()
+        if span.get('name', '').startswith('task.')
+    ]
+    top_task_nodes = []
+    for task_node in task_nodes:
+        parent_id = task_node.get('parent_id')
+        is_top = True
+
+        while parent_id:
+            parent_span = spans_dict.get(parent_id)
+            if not parent_span:
+                break
+
+            if parent_span.get('name', '').startswith('task.'):
+                is_top = False
+                break
+
+            parent_id = parent_span.get('parent_id')
+
+        if is_top:
+            top_task_nodes.append(task_node)
+
+    return top_task_nodes
+
+
+def _get_root_nodes(edges):
+    sources = set()
+    targets = set()
+    for edge in edges:
+        sources.add(edge['source'])
+        targets.add(edge['target'])
+    return sources - targets
+
+
+def _build_graph(root_spans: list):
+    nodes = []
+    edges = []
+
+    group_id_counter = 0
+
+    def __process_group_span(parent_spans, group_id, group_spans):
+        nonlocal group_id_counter
+        group_id_counter += 1
+        # add group node
+        group_node = {
+            'span_id': f'group_{group_id_counter}',
+            'group_id': group_id,
+            'show_name': 'Task Group'
+        }
+        nodes.append(group_node)
+
+        # add edges from parent_spans to group node
+        for parent_span in parent_spans:
+            edges.append({
+                'source': parent_span['span_id'],
+                'target': group_node['span_id']
+            })
+
+        # add edges from group node to children spans
+        last_spans = []
+        for child in group_spans:
+            edges.append({
+                'source': group_node['span_id'],
+                'target': child['span_id']
+            })
+            last_spans.extend(__process_span(child))
+        return last_spans
+
+    def __process_span(span):
+        nonlocal group_id_counter
+        nodes.append(span)
+        if 'children' in span:
+            groups = {}
+            for child in span['children']:
+                group_id = child.get('task_group_id', id(child))
+                if group_id not in groups:
+                    groups[group_id] = []
+                groups[group_id].append(child)
+
+            last_spans = [span]  # The leaf nodes of the current subtree
+            for group_id, group_spans in groups.items():
+                if len(group_spans) > 1:
+                    parent_spans = last_spans
+                    last_spans = __process_group_span(parent_spans, group_id, group_spans)
+                else:
+                    child_span = group_spans[0]
+                    # add edges from last_spans to child
+                    for prev_node in last_spans:
+                        edges.append({
+                            'source': prev_node['span_id'],
+                            'target': child_span['span_id']
+                        })
+                    last_spans = __process_span(child_span)
+        return last_spans
+
+    for span in root_spans:
+        __process_span(span)
+
+    return {
+        'nodes': nodes,
+        'edges': edges
+    }
+
+
+def get_agent_flow(trace_id):
+    from aworld.trace.server import get_trace_server
+
+    storage = get_trace_server().get_storage()
+    spans = storage.get_all_spans(trace_id)
+    spans_dict = {span.span_id: span.dict() for span in spans}
+    children_spans = []
+    top_task_nodes = _get_top_task_nodes(spans_dict)
+
+    filtered_spans = {}
+    for span_id, span in spans_dict.items():
+        if span.get('is_event', False) and span.get('run_type') == RunType.AGNET.value:
+            span['show_name'] = _get_agent_show_name(span)
+            span['event_id'] = span.get('attributes', {}).get('event.id')
+            filtered_spans[span_id] = span
+
+    sub_task_spans = []
+    for span in list(filtered_spans.values()):
+        skip_this_span = False
+        parent_id = span['parent_id'] if span['parent_id'] else None
+
+        while parent_id and parent_id not in filtered_spans:
+            parent_span = spans_dict.get(parent_id)
+            if parent_span and parent_span.get('run_type') == RunType.TASK.value:
+                # if str(parent_span['attributes'].get(semconv.TASK_IS_SUB_TASK)).lower() == 'true':
+                #     sub_task_spans.append(span)
+                #     skip_this_span = True
+                #     break
+                # else:
+                # print(f"parent_span_name: {parent_span['name']}")
+                span['task_group_id'] = parent_span['attributes'].get(semconv.TASK_GROUP_ID)
+            parent_id = parent_span['parent_id'] if parent_span and parent_span['parent_id'] else None
+
+        if skip_this_span:
+            continue
+        if parent_id:
+            parent_span = filtered_spans.get(parent_id)
+            if not parent_span:
+                continue
+
+            if 'children' not in parent_span:
+                parent_span['children'] = []
+            parent_span['children'].append(span)
+            children_spans.append(span)
+
+    filtered_span_list = [span for span in filtered_spans.values() if span not in sub_task_spans]
+    root_spans = [span for span in filtered_span_list
+                  if span not in children_spans]
+
+    data = _build_graph(root_spans)
+    _remove_span_detail(data["nodes"])
+
+    # add query start node
+    _add_query_node(data, top_task_nodes)
+    return data
+
+def _add_query_node(data, top_task_nodes):
+    top_task_node = top_task_nodes[0] if top_task_nodes else None
+    if top_task_node:
+        start_node_span_id = f'{uuid.uuid4().hex[:16]}'
+        data['nodes'].append({
+            'span_id': start_node_span_id,
+            'show_name': top_task_node['attributes'].get(semconv.TASK_INPUT),
+        })
+        root_nodes = _get_root_nodes(data['edges'])
+        if root_nodes:
+            for root_node in root_nodes:
+                data['edges'].append({
+                    'source': start_node_span_id,
+                    'target': root_node,
+                })
+        else:
+            for root_node in data['nodes']:
+                if root_node['span_id'] != start_node_span_id:
+                    data['edges'].append({
+                        'source': start_node_span_id,
+                        'target': root_node['span_id'],
+                    })