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

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'],
+                    })