langchain-agentx-python 0.1__py3-none-any.whl

langchain_agentx/observability/trace/collector.py

@@ -0,0 +1,560 @@
+"""
+TraceCollector：统一的 Trace 数据汇聚点。
+
+职责：
+1. 维护内存中的 span 树（按 session_id 隔离）
+2. 提供幂等的 session 管理
+3. 提供查询 API（替代 PerformanceMonitor）
+
+设计要点：
+- 不依赖 ContextVar，所有方法显式传递 session_id
+- start_session() 幂等：已存在则跳过
+- 线程安全：使用 Lock 保护 _sessions 字典
+"""
+
+import time
+import uuid
+import logging
+from threading import Lock
+from typing import Dict, List, Optional, Any
+
+from .models import (
+    TraceSession,
+    TraceSpan,
+    SessionBuffer,
+    SpanType,
+    SpanStatus,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class TraceCollector:
+    """Trace 数据收集器"""
+    
+    def __init__(self, store: Optional[Any] = None):
+        """
+        初始化 TraceCollector。
+        
+        Args:
+            store: SqliteTraceStore 实例（可选，Phase 2 实现）
+        """
+        self._store = store
+        self._sessions: Dict[str, SessionBuffer] = {}
+        self._lock = Lock()
+    
+    # ========= 采集 API（幂等） =========
+    
+    def has_session(self, session_id: str) -> bool:
+        """
+        检查 session 是否已启动。
+        
+        Args:
+            session_id: session 唯一标识
+            
+        Returns:
+            True 如果 session 已存在
+        """
+        with self._lock:
+            return session_id in self._sessions
+
+    def get_session_buffer(self, session_id: str) -> Optional[SessionBuffer]:
+        """返回 session 的内存缓冲区，供评估系统读取。"""
+        with self._lock:
+            return self._sessions.get(session_id)
+    
+    def start_session(
+        self,
+        session_id: str,
+        graph_name: str = "unknown",
+        user_query: Optional[str] = None,
+        **kwargs
+    ) -> None:
+        """
+        启动一个 trace session（幂等：已存在则跳过）。
+        
+        Args:
+            session_id: session 唯一标识（通常是 run_id）
+            graph_name: Agent 名称
+            user_query: 用户查询（可选）
+            **kwargs: 其他元数据
+        """
+        with self._lock:
+            if session_id in self._sessions:
+                return  # 已存在，跳过
+
+            conversation_session_id = str(kwargs.pop("conversation_session_id", "") or session_id)
+            parent_session_id = kwargs.pop("parent_session_id", None)
+            parent_tool_call_id = kwargs.pop("parent_tool_call_id", None)
+            workflow_id = kwargs.pop("workflow_id", None)
+            unit_type = str(kwargs.pop("unit_type", "main_loop") or "main_loop")
+            container_type = str(kwargs.pop("container_type", "agent_session") or "agent_session")
+            origin = str(kwargs.pop("origin", "user") or "user")
+            visibility = str(kwargs.pop("visibility", "default") or "default")
+            session = TraceSession(
+                session_id=session_id,
+                conversation_session_id=conversation_session_id,
+                parent_session_id=parent_session_id,
+                parent_tool_call_id=parent_tool_call_id,
+                workflow_id=workflow_id,
+                unit_type=unit_type,
+                container_type=container_type,
+                origin=origin,
+                visibility=visibility,
+                graph_name=graph_name,
+                user_query=user_query,
+                start_time=time.time(),
+                metadata=kwargs
+            )
+
+            self._sessions[session_id] = SessionBuffer(session=session)
+    
+    def end_session(
+        self,
+        session_id: str,
+        finish_reason: Optional[str] = None,
+        **kwargs
+    ) -> None:
+        """
+        结束一个 session。
+        
+        Args:
+            session_id: session 唯一标识
+            finish_reason: 退出原因
+            **kwargs: 其他元数据
+        """
+        should_flush = False
+        with self._lock:
+            if session_id not in self._sessions:
+                return
+            
+            buffer = self._sessions[session_id]
+            buffer.session.end_time = time.time()
+            buffer.session.status = str(kwargs.get("status", "completed") or "completed")
+            buffer.session.finish_reason = finish_reason
+            if "exit_code" in kwargs:
+                buffer.session.exit_code = kwargs.get("exit_code")
+            if "terminal_reason" in kwargs:
+                buffer.session.terminal_reason = kwargs.get("terminal_reason")
+            
+            # 计算聚合指标
+            buffer.session.loop_count = sum(
+                1 for span in buffer.spans if span.span_type == SpanType.LOOP_STEP
+            )
+            llm_spans = [span for span in buffer.spans if span.span_type == SpanType.LLM_CALL]
+            buffer.session.total_input_tokens = sum(
+                self._as_int((span.metadata or {}).get("input_tokens")) for span in llm_spans
+            )
+            buffer.session.total_output_tokens = sum(
+                self._as_int((span.metadata or {}).get("output_tokens")) for span in llm_spans
+            )
+            buffer.session.total_latency_ms = (
+                (buffer.session.end_time - buffer.session.start_time) * 1000
+                if buffer.session.end_time
+                else 0
+            )
+            should_flush = True
+
+        # 在会话结束时自动持久化，避免调用方漏调 flush() 导致数据仅停留在内存中。
+        if should_flush:
+            self.flush(session_id)
+
+    def emit_event(
+        self,
+        *,
+        session_id: str,
+        event_type: str,
+        payload: Optional[Dict[str, Any]] = None,
+        parent_span_id: Optional[str] = None,
+    ) -> str:
+        """统一事件采集入口（用于 loop/task 决策类事件）。"""
+        payload = payload or {}
+        span_type_map = {
+            "loop.decision": SpanType.LOOP_DECISION.value,
+            "task.reserve": SpanType.TASK_EVENT.value,
+            "task.ack": SpanType.TASK_EVENT.value,
+            "hook.execution": SpanType.HOOK_EXEC.value,
+        }
+        span_type = span_type_map.get(event_type, SpanType.TASK_EVENT.value)
+        span_id = self.span_start(
+            span_type=span_type,
+            name=event_type,
+            session_id=session_id,
+            parent_span_id=parent_span_id,
+            metadata=payload,
+        )
+        self.span_end(span_id=span_id, session_id=session_id)
+        return span_id
+    
+    def span_start(
+        self,
+        span_type: str,
+        name: str,
+        session_id: str,
+        parent_span_id: Optional[str] = None,
+        input_data: Optional[Dict[str, Any]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        """
+        开始一个 span。
+        
+        Args:
+            span_type: span 类型（agent_run / loop_step / llm_call / tool_call）
+            name: span 名称
+            session_id: 所属 session
+            parent_span_id: 父 span ID（可选）
+            input_data: 输入数据
+            metadata: 元数据
+            
+        Returns:
+            span_id: 新创建的 span ID
+            
+        Raises:
+            RuntimeError: 如果 session 未启动
+        """
+        with self._lock:
+            if session_id not in self._sessions:
+                raise RuntimeError(
+                    f"trace session is not started: {session_id}. "
+                    f"Call start_session() first or ensure _ensure_session() is called."
+                )
+            
+            buffer = self._sessions[session_id]
+            metadata = metadata or {}
+            metadata = dict(metadata)
+            metadata["seq_no"] = buffer.next_seq_no
+            buffer.next_seq_no += 1
+
+            # 自动父 span 推断：若调用方未显式传 parent_span_id，则挂到“当前最近打开且未关闭”的 span。
+            # 依赖 buffer.span_stack 维护的 FILO 语义；显式传参优先，不覆盖调用方意图。
+            if parent_span_id is None and buffer.span_stack:
+                parent_span_id = buffer.span_stack[-1]
+            
+            span_id = str(uuid.uuid4())
+            span = TraceSpan(
+                span_id=span_id,
+                session_id=session_id,
+                parent_span_id=parent_span_id,
+                span_type=SpanType(span_type) if isinstance(span_type, str) else span_type,
+                name=name,
+                start_time=time.time(),
+                input_data=input_data,
+                metadata=metadata
+            )
+            
+            buffer.add_span(span)
+            buffer.span_stack.append(span_id)
+            
+            return span_id
+    
+    def span_end(
+        self,
+        span_id: str,
+        session_id: str,
+        output_data: Optional[Dict[str, Any]] = None,
+        error: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """
+        结束一个 span。
+        
+        Args:
+            span_id: span 唯一标识
+            session_id: 所属 session
+            output_data: 输出数据
+            error: 错误信息（如果失败）
+            metadata: 额外元数据
+        """
+        span_to_persist: TraceSpan | None = None
+        with self._lock:
+            if session_id not in self._sessions:
+                return
+            
+            buffer = self._sessions[session_id]
+            span = buffer.get_span(span_id)
+            
+            if not span:
+                return
+            
+            span.end_time = time.time()
+            span.latency_ms = (span.end_time - span.start_time) * 1000
+            span.status = SpanStatus.FAILED if error else SpanStatus.COMPLETED
+            span.output_data = output_data
+            span.error = error
+            
+            if metadata:
+                span.metadata.update(metadata)
+            
+            # 从栈中移除
+            if span_id in buffer.span_stack:
+                buffer.span_stack.remove(span_id)
+            span_to_persist = span
+
+        # 双写：span 结束时立即落盘，降低异常中断时数据丢失概率。
+        if span_to_persist is not None and self._store is not None:
+            self._store.save_spans([span_to_persist])
+    
+    def flush(self, session_id: str) -> None:
+        """
+        将 session 聚合信息写入存储并清理内存。
+        
+        Args:
+            session_id: session 唯一标识
+        """
+        with self._lock:
+            if session_id not in self._sessions:
+                return
+            
+            buffer = self._sessions[session_id]
+            
+            # 写入 SQLite：span 已在 span_end 双写，不再重复落盘。
+            if self._store:
+                self._store.save_session(buffer.session)
+                logger.debug("Flushed session %s to SQLite", session_id)
+            else:
+                logger.debug("Flushing session %s (no store)", session_id)
+            logger.debug(
+                "Trace session summary: graph=%s spans=%s loop_count=%s total_latency_ms=%.2f",
+                buffer.session.graph_name,
+                len(buffer.spans),
+                buffer.session.loop_count,
+                buffer.session.total_latency_ms,
+            )
+            
+            # 清理内存
+            del self._sessions[session_id]
+    
+    # ========= 查询 API（Phase 2 实现） =========
+    
+    def get_session(self, session_id: str) -> Optional[TraceSession]:
+        """获取 session"""
+        # 先查内存
+        with self._lock:
+            if session_id in self._sessions:
+                return self._sessions[session_id].session
+        
+        # 再查 SQLite
+        if self._store:
+            return self._store.get_session(session_id)
+        
+        return None
+    
+    def query_sessions(self, **filters) -> List[TraceSession]:
+        """查询 session 列表"""
+        if self._store:
+            return self._store.query_sessions(**filters)
+        return []
+    
+    def get_span_tree(self, session_id: str) -> List[TraceSpan]:
+        """获取 session 的所有 span"""
+        # 先查内存
+        with self._lock:
+            if session_id in self._sessions:
+                return self._sessions[session_id].spans.copy()
+        
+        # 再查 SQLite
+        if self._store:
+            return self._store.get_span_tree(session_id)
+        
+        return []
+
+    def get_session_metrics(self, session_id: str) -> Dict[str, Any]:
+        """返回单个 session 的聚合指标。"""
+        spans = self.get_span_tree(session_id)
+        session = self.get_session(session_id)
+        return self._build_session_metrics(session_id=session_id, session=session, spans=spans)
+
+    def get_sessions_analysis(self, session_ids: List[str]) -> Dict[str, Dict[str, Any]]:
+        """
+        批量计算多个 session 的 metrics/warnings/degradation，避免重复查库。
+        返回字典按传入 session_ids 顺序插入，便于上层直接按输入顺序消费。
+        返回结构：
+        {
+          session_id: {
+            "metrics": ...,
+            "warnings": ...,
+            "degradation": ...,
+          }
+        }
+        """
+        if not session_ids:
+            return {}
+        sessions: Dict[str, TraceSession] = {}
+        if self._store and hasattr(self._store, "get_sessions_by_ids"):
+            sessions = self._store.get_sessions_by_ids(session_ids)
+        else:
+            sessions = {
+                session_id: session
+                for session_id in session_ids
+                for session in [self.get_session(session_id)]
+                if session is not None
+            }
+        spans_map: Dict[str, List[TraceSpan]] = {}
+        if self._store and hasattr(self._store, "get_span_trees"):
+            spans_map = self._store.get_span_trees(session_ids)
+        else:
+            spans_map = {sid: self.get_span_tree(sid) for sid in session_ids}
+
+        result: Dict[str, Dict[str, Any]] = {}
+        for sid in session_ids:
+            spans = spans_map.get(sid, [])
+            session = sessions.get(sid)
+            metrics = self._build_session_metrics(session_id=sid, session=session, spans=spans)
+            warnings = self._build_performance_warnings(spans=spans)
+            degradation = self._build_degradation_stats(spans=spans)
+            result[sid] = {
+                "metrics": metrics,
+                "warnings": warnings,
+                "degradation": degradation,
+            }
+        return result
+
+    def _build_session_metrics(
+        self,
+        *,
+        session_id: str,
+        session: Optional[TraceSession],
+        spans: List[TraceSpan],
+    ) -> Dict[str, Any]:
+        ordered = self._ordered_spans(spans)
+        tool_spans = [span for span in ordered if span.span_type == SpanType.TOOL_CALL]
+        llm_spans = [span for span in ordered if span.span_type == SpanType.LLM_CALL]
+        tool_call_chain = [span.name for span in tool_spans]
+        llm_call_chain = [span.name for span in llm_spans]
+        complete_call_chain = [f"{span.span_type.value}:{span.name}" for span in ordered]
+        total_llm_latency_ms = sum((span.latency_ms or 0.0) for span in llm_spans)
+        total_input_tokens = sum(
+            self._as_int((span.metadata or {}).get("input_tokens")) for span in llm_spans
+        )
+        total_output_tokens = sum(
+            self._as_int((span.metadata or {}).get("output_tokens")) for span in llm_spans
+        )
+        return {
+            "session_id": session_id,
+            "total_tool_calls": len(tool_spans),
+            "total_latency_ms": session.total_latency_ms if session else 0.0,
+            "tool_call_chain": tool_call_chain,
+            "unique_tools": sorted(set(tool_call_chain)),
+            "context_tokens": total_input_tokens,
+            "total_input_tokens": total_input_tokens,
+            "total_output_tokens": total_output_tokens,
+            "total_llm_calls": len(llm_spans),
+            "llm_call_chain": llm_call_chain,
+            "total_llm_latency_ms": total_llm_latency_ms,
+            "complete_call_chain": complete_call_chain,
+        }
+
+    def get_performance_warnings(self, session_id: str) -> List[str]:
+        """返回单个 session 的性能告警。"""
+        spans = self.get_span_tree(session_id)
+        return self._build_performance_warnings(spans=spans)
+
+    def _build_performance_warnings(self, *, spans: List[TraceSpan]) -> List[str]:
+        ordered = self._ordered_spans(spans)
+        warnings: List[str] = []
+        tool_spans = [span for span in ordered if span.span_type == SpanType.TOOL_CALL]
+        llm_spans = [span for span in ordered if span.span_type == SpanType.LLM_CALL]
+        if len(tool_spans) > 10:
+            warnings.append(f"工具调用次数偏高: {len(tool_spans)} (>10)")
+        if tool_spans:
+            avg_tool_latency_ms = sum((span.latency_ms or 0.0) for span in tool_spans) / len(tool_spans)
+            if avg_tool_latency_ms > 2000:
+                warnings.append(f"工具平均耗时偏高: {avg_tool_latency_ms:.0f}ms (>2000ms)")
+        if llm_spans:
+            avg_llm_latency_ms = sum((span.latency_ms or 0.0) for span in llm_spans) / len(llm_spans)
+            if avg_llm_latency_ms > 5000:
+                warnings.append(f"LLM 平均耗时偏高: {avg_llm_latency_ms:.0f}ms (>5000ms)")
+        return warnings
+
+    def get_degradation_stats(self, session_id: str) -> Dict[str, Any]:
+        """返回工具调用退化统计。"""
+        spans = self.get_span_tree(session_id)
+        return self._build_degradation_stats(spans=spans)
+
+    def _build_degradation_stats(self, *, spans: List[TraceSpan]) -> Dict[str, Any]:
+        llm_spans = [span for span in spans if span.span_type == SpanType.LLM_CALL]
+        degradation_hits: List[TraceSpan] = []
+        tool_name_distribution: Dict[str, int] = {}
+        for span in llm_spans:
+            meta = span.metadata or {}
+            degradation = meta.get("degradation")
+            detected = False
+            tool_names: List[str] = []
+            if isinstance(degradation, dict):
+                detected = bool(degradation.get("detected"))
+                raw_names = degradation.get("tool_names")
+                if isinstance(raw_names, list):
+                    tool_names = [str(item) for item in raw_names if str(item)]
+            elif isinstance(degradation, bool):
+                detected = degradation
+            if detected:
+                degradation_hits.append(span)
+                for tool_name in tool_names:
+                    tool_name_distribution[tool_name] = tool_name_distribution.get(tool_name, 0) + 1
+        total_llm_calls = len(llm_spans)
+        degradation_count = len(degradation_hits)
+        total_input_on_degradation = sum(
+            self._as_int((span.metadata or {}).get("input_tokens")) for span in degradation_hits
+        )
+        return {
+            "total_llm_calls": total_llm_calls,
+            "degradation_count": degradation_count,
+            "degradation_rate_percent": (
+                (degradation_count / total_llm_calls) * 100 if total_llm_calls else 0.0
+            ),
+            "tool_name_distribution": tool_name_distribution,
+            "avg_prompt_length_on_degradation": (
+                total_input_on_degradation / degradation_count if degradation_count else 0.0
+            ),
+        }
+
+    def generate_report(self, session_id: str, mode: str = "debug") -> str:
+        """生成 session 报告文本。"""
+        mode_normalized = str(mode).lower()
+        if mode_normalized == "production":
+            return ""
+        metrics = self.get_session_metrics(session_id)
+        warnings = self.get_performance_warnings(session_id)
+        degradation = self.get_degradation_stats(session_id)
+        lines = [
+            f"Trace Report [{session_id}]",
+            f"- Tool calls: {metrics['total_tool_calls']}",
+            f"- LLM calls: {metrics['total_llm_calls']}",
+            f"- Total latency: {metrics['total_latency_ms']:.1f}ms",
+            f"- LLM latency: {metrics['total_llm_latency_ms']:.1f}ms",
+            f"- Tokens: in={metrics['total_input_tokens']} out={metrics['total_output_tokens']}",
+            (
+                f"- Degradation: {degradation['degradation_count']}/{degradation['total_llm_calls']}"
+                f" ({degradation['degradation_rate_percent']:.1f}%)"
+            ),
+        ]
+        if warnings:
+            lines.append("- Warnings:")
+            lines.extend([f"  - {item}" for item in warnings])
+        else:
+            lines.append("- Warnings: none")
+        if mode_normalized == "benchmark":
+            lines.append(f"- Tool chain: {', '.join(metrics['tool_call_chain']) or '-'}")
+            lines.append(f"- LLM chain: {', '.join(metrics['llm_call_chain']) or '-'}")
+            if degradation["tool_name_distribution"]:
+                lines.append(f"- Degradation tools: {degradation['tool_name_distribution']}")
+        return "\n".join(lines)
+
+    @staticmethod
+    def _as_int(value: Any) -> int:
+        try:
+            return int(value)
+        except (TypeError, ValueError):
+            return 0
+
+    @staticmethod
+    def _ordered_spans(spans: List[TraceSpan]) -> List[TraceSpan]:
+        def seq_no(span: TraceSpan) -> int:
+            raw = (span.metadata or {}).get("seq_no")
+            try:
+                return int(raw)
+            except (TypeError, ValueError):
+                return 10**9
+
+        return sorted(spans, key=lambda span: (seq_no(span), span.start_time or 0.0))