agentevals-cli 0.5.2__py3-none-any.whl

agentevals/streaming/incremental_processor.py

@@ -0,0 +1,337 @@
+"""Incremental span processor for extracting conversation elements in real-time.
+
+Processes OTLP spans as they arrive and extracts:
+- User input from call_llm spans
+- Tool calls from execute_tool spans
+- Agent responses from call_llm spans
+- Token usage information
+
+This enables real-time display of agent execution progress without waiting for
+session completion.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from ..extraction import (
+    extract_agent_response_from_attrs,
+    extract_token_usage_from_attrs,
+    extract_tool_call_from_attrs,
+    extract_tool_result_from_attrs,
+    extract_user_text_from_attrs,
+    flatten_otlp_attributes,
+    parse_tool_response_content,
+)
+from ..trace_attrs import (
+    ADK_INVOCATION_ID,
+    ADK_SCOPE_VALUE,
+    OTEL_GENAI_REQUEST_MODEL,
+    OTEL_GENAI_TOOL_NAME,
+    OTEL_SCOPE,
+)
+from ..utils.genai_messages import parse_json_attr
+
+logger = logging.getLogger(__name__)
+
+
+def _normalize_ts(raw_ts) -> float:
+    """Normalize a nanosecond timestamp (string or int) to seconds."""
+    try:
+        ns = int(raw_ts)
+    except (TypeError, ValueError):
+        return 0.0
+    if ns > 1e15:
+        return ns / 1e9
+    return float(ns)
+
+
+class IncrementalInvocationExtractor:
+    """Extracts conversation elements from spans and logs as they arrive."""
+
+    def __init__(self):
+        self.seen_user_input = set()
+        self.seen_tool_calls = {}
+        self.seen_agent_response = set()
+        self.llm_spans_by_invocation = {}
+        self.token_totals = {}
+        self.current_invocation_id = None
+        self.seen_message_contents = set()  # Track message contents to avoid duplicates
+        self.tool_names_by_id: dict[str, str] = {}  # tool_call_id -> tool_name
+
+    def process_span(self, span: dict) -> list[dict]:
+        """Process a single OTLP span and return conversation updates to broadcast.
+
+        Args:
+            span: OTLP JSON span dictionary
+
+        Returns:
+            List of update events to broadcast via SSE
+        """
+        updates = []
+        operation_name = span.get("name", "")
+
+        attributes = flatten_otlp_attributes(span.get("attributes", []))
+
+        is_adk = attributes.get(OTEL_SCOPE) == ADK_SCOPE_VALUE
+        is_genai_llm = bool(attributes.get(OTEL_GENAI_REQUEST_MODEL))
+        is_genai_tool = bool(attributes.get(OTEL_GENAI_TOOL_NAME))
+
+        if not (is_adk or is_genai_llm or is_genai_tool):
+            return updates
+
+        invocation_id = self._get_invocation_id(span, attributes)
+        if not invocation_id:
+            return updates
+
+        self.current_invocation_id = invocation_id
+
+        is_llm = operation_name.startswith("call_llm") or is_genai_llm
+        if is_llm:
+            if invocation_id not in self.llm_spans_by_invocation:
+                self.llm_spans_by_invocation[invocation_id] = []
+            self.llm_spans_by_invocation[invocation_id].append(span)
+
+            if invocation_id not in self.seen_user_input:
+                user_text = extract_user_text_from_attrs(attributes)
+                if user_text:
+                    message_key = f"user:{user_text.strip()}"
+                    if message_key not in self.seen_message_contents:
+                        logger.debug(f"Extracted user input for invocation {invocation_id}")
+                        updates.append(
+                            {
+                                "type": "user_input",
+                                "invocationId": invocation_id,
+                                "text": user_text,
+                                "timestamp": int(span.get("startTimeUnixNano", 0)) / 1e9,
+                            }
+                        )
+                        self.seen_message_contents.add(message_key)
+                self.seen_user_input.add(invocation_id)
+
+            agent_text = extract_agent_response_from_attrs(attributes)
+            if agent_text and invocation_id not in self.seen_agent_response:
+                message_key = f"agent:{agent_text.strip()}"
+                if message_key not in self.seen_message_contents:
+                    logger.debug(f"Extracted agent response for invocation {invocation_id}")
+                    updates.append(
+                        {
+                            "type": "agent_response",
+                            "invocationId": invocation_id,
+                            "text": agent_text,
+                            "timestamp": int(span.get("endTimeUnixNano", 0)) / 1e9,
+                        }
+                    )
+                    self.seen_message_contents.add(message_key)
+                self.seen_agent_response.add(invocation_id)
+
+            in_toks, out_toks, model = extract_token_usage_from_attrs(attributes)
+            if in_toks or out_toks:
+                if invocation_id not in self.token_totals:
+                    self.token_totals[invocation_id] = {
+                        "inputTokens": 0,
+                        "outputTokens": 0,
+                        "model": model,
+                    }
+
+                self.token_totals[invocation_id]["inputTokens"] += in_toks
+                self.token_totals[invocation_id]["outputTokens"] += out_toks
+
+                logger.debug("Token update for %s: +%d input, +%d output", invocation_id, in_toks, out_toks)
+
+                updates.append(
+                    {
+                        "type": "token_update",
+                        "invocationId": invocation_id,
+                        "inputTokens": in_toks,
+                        "outputTokens": out_toks,
+                        "model": model,
+                    }
+                )
+
+        elif operation_name.startswith("execute_tool") or is_genai_tool:
+            span_id = span.get("spanId", "")
+            tool_call = extract_tool_call_from_attrs(attributes, operation_name, span_id=span_id)
+            if tool_call:
+                call_id = tool_call["id"]
+                if invocation_id not in self.seen_tool_calls:
+                    self.seen_tool_calls[invocation_id] = set()
+
+                self.tool_names_by_id[call_id] = tool_call["name"]
+
+                if call_id not in self.seen_tool_calls[invocation_id]:
+                    updates.append(
+                        {
+                            "type": "tool_call",
+                            "invocationId": invocation_id,
+                            "toolCall": tool_call,
+                            "timestamp": int(span.get("startTimeUnixNano", 0)) / 1e9,
+                        }
+                    )
+                    self.seen_tool_calls[invocation_id].add(call_id)
+
+                    tool_result = extract_tool_result_from_attrs(attributes)
+                    if tool_result:
+                        updates.append(
+                            {
+                                "type": "tool_result",
+                                "invocationId": invocation_id,
+                                "toolCallId": call_id,
+                                "toolName": tool_call["name"],
+                                "response": tool_result["response"],
+                                "isError": tool_result["isError"],
+                                "timestamp": int(span.get("endTimeUnixNano", 0)) / 1e9,
+                            }
+                        )
+
+        return updates
+
+    def process_log(self, log_event: dict) -> list[dict]:
+        """Process a GenAI log event and extract conversation updates.
+
+        Args:
+            log_event: Log event dict with event_name, body, attributes
+
+        Returns:
+            List of update events to broadcast via SSE
+        """
+        updates = []
+        event_name = log_event.get("event_name", "")
+        body = log_event.get("body", {})
+
+        invocation_id = log_event.get("span_id")
+        if not invocation_id:
+            invocation_id = self.current_invocation_id
+        if not invocation_id:
+            return updates
+
+        # Extract user messages (gen_ai.user.message)
+        if event_name == "gen_ai.user.message":
+            if isinstance(body, dict) and "content" in body:
+                user_text = body["content"]
+                message_key = f"user:{user_text.strip() if isinstance(user_text, str) else user_text}"
+                if user_text and message_key not in self.seen_message_contents:
+                    logger.debug(f"Extracted user input from log for invocation {invocation_id}")
+                    updates.append(
+                        {
+                            "type": "user_input",
+                            "invocationId": invocation_id,
+                            "text": user_text,
+                            "timestamp": _normalize_ts(log_event.get("timestamp", 0)),
+                        }
+                    )
+                    self.seen_message_contents.add(message_key)
+                    self.seen_user_input.add(invocation_id)
+
+        # Extract assistant messages (gen_ai.assistant.message or gen_ai.choice)
+        elif event_name in ("gen_ai.assistant.message", "gen_ai.choice"):
+            agent_text = None
+
+            if isinstance(body, dict):
+                # Check for direct content
+                if "content" in body:
+                    agent_text = body["content"]
+                # Check for message.content (gen_ai.choice format)
+                elif "message" in body and isinstance(body["message"], dict):
+                    if "content" in body["message"]:
+                        agent_text = body["message"]["content"]
+
+            if agent_text:
+                message_key = f"agent:{agent_text.strip() if isinstance(agent_text, str) else agent_text}"
+                if message_key not in self.seen_message_contents:
+                    logger.debug(f"Extracted agent response from log for invocation {invocation_id}")
+                    updates.append(
+                        {
+                            "type": "agent_response",
+                            "invocationId": invocation_id,
+                            "text": agent_text,
+                            "timestamp": _normalize_ts(log_event.get("timestamp", 0)),
+                        }
+                    )
+                    self.seen_message_contents.add(message_key)
+                    self.seen_agent_response.add(invocation_id)
+
+            # Extract tool calls from assistant message
+            if isinstance(body, dict):
+                tool_calls = None
+                if "tool_calls" in body:
+                    tool_calls = body["tool_calls"]
+                elif "message" in body and isinstance(body["message"], dict) and "tool_calls" in body["message"]:
+                    tool_calls = body["message"]["tool_calls"]
+
+                if tool_calls and isinstance(tool_calls, list):
+                    for tc in tool_calls:
+                        if isinstance(tc, dict):
+                            tool_id = tc.get("id", "unknown")
+                            tool_key = f"tool:{tool_id}"
+
+                            tc_name = (
+                                tc.get("function", {}).get("name", "unknown")
+                                if "function" in tc
+                                else tc.get("name", "unknown")
+                            )
+                            self.tool_names_by_id[tool_id] = tc_name
+
+                            if tool_key not in self.seen_message_contents:
+                                tool_call = {
+                                    "id": tool_id,
+                                    "name": tc_name,
+                                    "args": {},
+                                }
+
+                                if "function" in tc and "arguments" in tc["function"]:
+                                    parsed = parse_json_attr(
+                                        tc["function"]["arguments"], "tool_call.function.arguments"
+                                    )
+                                    if isinstance(parsed, dict):
+                                        tool_call["args"] = parsed
+
+                                logger.debug(f"Extracted tool call from log for invocation {invocation_id}")
+                                updates.append(
+                                    {
+                                        "type": "tool_call",
+                                        "invocationId": invocation_id,
+                                        "toolCall": tool_call,
+                                        "timestamp": _normalize_ts(log_event.get("timestamp", 0)),
+                                    }
+                                )
+                                self.seen_message_contents.add(tool_key)
+
+                                if invocation_id not in self.seen_tool_calls:
+                                    self.seen_tool_calls[invocation_id] = set()
+                                self.seen_tool_calls[invocation_id].add(tool_id)
+
+        # Extract tool results from gen_ai.tool.message logs
+        elif event_name == "gen_ai.tool.message":
+            if isinstance(body, dict):
+                tool_id = body.get("id", "unknown")
+                tool_name = body.get("name") or self.tool_names_by_id.get(tool_id, "unknown")
+                content = body.get("content")
+                if content is not None:
+                    response = parse_tool_response_content(content)
+                    result_key = f"tool_result:{tool_id}"
+                    if result_key not in self.seen_message_contents:
+                        is_error = bool(response.get("isError", False))
+                        updates.append(
+                            {
+                                "type": "tool_result",
+                                "invocationId": invocation_id,
+                                "toolCallId": tool_id,
+                                "toolName": tool_name,
+                                "response": response,
+                                "isError": is_error,
+                                "timestamp": _normalize_ts(log_event.get("timestamp", 0)),
+                            }
+                        )
+                        self.seen_message_contents.add(result_key)
+
+        return updates
+
+    def _get_invocation_id(self, span: dict, attributes: dict) -> str | None:
+        invocation_id = attributes.get(ADK_INVOCATION_ID)
+        if invocation_id:
+            return invocation_id
+        parent_span_id = span.get("parentSpanId")
+        if parent_span_id:
+            return parent_span_id
+        return span.get("spanId")

agentevals/streaming/processor.py

@@ -0,0 +1,285 @@
+"""OpenTelemetry SpanProcessor and LogRecordProcessor for streaming to agentevals dev server."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from typing import Any
+
+try:
+    import websockets
+    from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
+    from opentelemetry.trace import SpanKind
+except ImportError:
+    websockets = None
+    ReadableSpan = None
+    SpanProcessor = None
+    SpanKind = None
+
+logger = logging.getLogger(__name__)
+
+
+class AgentEvalsStreamingProcessor:
+    """OTel span processor that streams spans to agentevals dev server via WebSocket."""
+
+    def __init__(self, ws_url: str, session_id: str, trace_id: str):
+        if websockets is None or SpanProcessor is None:
+            raise ImportError(
+                "websockets and opentelemetry-sdk required for streaming. "
+                "Install with: pip install websockets opentelemetry-sdk"
+            )
+
+        self.ws_url = ws_url
+        self.session_id = session_id
+        self.trace_id = trace_id
+        self.websocket: Any | None = None
+        self.loop: asyncio.AbstractEventLoop | None = None
+        self._connected = False
+        self._span_buffer: list[dict] = []
+        self._failed_spans: list[dict] = []
+        self._pending_sends: set[asyncio.Task] = set()
+
+    async def connect(self, eval_set_id: str | None = None, metadata: dict | None = None):
+        try:
+            self.websocket = await websockets.connect(self.ws_url)
+            self.loop = asyncio.get_running_loop()
+
+            await self.websocket.send(
+                json.dumps(
+                    {
+                        "type": "session_start",
+                        "session_id": self.session_id,
+                        "trace_id": self.trace_id,
+                        "eval_set_id": eval_set_id,
+                        "metadata": metadata or {},
+                    }
+                )
+            )
+
+            self._connected = True
+            logger.info("Connected to agentevals dev server: %s", self.session_id)
+
+        except Exception as exc:
+            logger.error("Failed to connect to agentevals server: %s", exc)
+            self._connected = False
+
+    def on_start(self, span: ReadableSpan, parent_context=None) -> None:
+        pass
+
+    def on_end(self, span: ReadableSpan) -> None:
+        if not self._connected or not self.websocket or not self.loop:
+            logger.debug(f"Skipping span {span.name}: not connected")
+            return
+
+        try:
+            otlp_span = self._span_to_otlp(span)
+            self._span_buffer.append(otlp_span)
+
+            future = asyncio.run_coroutine_threadsafe(self._send_span(otlp_span), self.loop)
+            self._pending_sends.add(future)
+
+            def handle_send_complete(fut):
+                self._pending_sends.discard(fut)
+                try:
+                    fut.result()
+                    logger.debug(f"Sent span: {span.name}")
+                except Exception as exc:
+                    logger.error(f"Failed to send span {span.name}: {exc}")
+                    self._failed_spans.append(otlp_span)
+
+            future.add_done_callback(handle_send_complete)
+
+        except Exception as exc:
+            logger.warning("Failed to convert span: %s", exc)
+
+    def shutdown(self) -> None:
+        pass
+
+    async def shutdown_async(self) -> None:
+        if self.websocket and self._connected:
+            try:
+                await self._send_session_end()
+            except Exception as exc:
+                logger.warning("Failed to shutdown cleanly: %s", exc)
+
+    async def _send_span(self, otlp_span: dict) -> None:
+        if not self.websocket:
+            raise ConnectionError("WebSocket not connected")
+
+        message = {
+            "type": "span",
+            "session_id": self.session_id,
+            "span": otlp_span,
+        }
+        await self.websocket.send(json.dumps(message))
+        logger.debug(f"Sent span: {otlp_span.get('name')}")
+
+    async def _send_session_end(self) -> None:
+        try:
+            if not self.websocket:
+                return
+
+            if self._pending_sends:
+                logger.info("Waiting for %d pending span sends to complete...", len(self._pending_sends))
+                for future in list(self._pending_sends):
+                    try:
+                        await asyncio.wrap_future(future)
+                    except Exception as exc:
+                        logger.warning("Pending send failed during shutdown: %s", exc)
+                logger.info("All pending sends completed")
+
+            if self._failed_spans:
+                logger.info("Retrying %d failed spans at shutdown", len(self._failed_spans))
+                for otlp_span in self._failed_spans:
+                    try:
+                        await self._send_span(otlp_span)
+                    except Exception as exc:
+                        logger.error("Failed to send span even at shutdown: %s", exc)
+
+            self._failed_spans.clear()
+            self._span_buffer.clear()
+            self._pending_sends.clear()
+
+            await self.websocket.send(json.dumps({"type": "session_end", "session_id": self.session_id}))
+
+            await self.websocket.close()
+            self._connected = False
+        except Exception as exc:
+            logger.error("Failed to send session_end: %s", exc)
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+    def _span_to_otlp(self, span: ReadableSpan) -> dict:
+        scope_name = span.instrumentation_scope.name if span.instrumentation_scope else ""
+        scope_version = span.instrumentation_scope.version if span.instrumentation_scope else ""
+
+        attributes = []
+        if scope_name:
+            attributes.append({"key": "otel.scope.name", "value": {"stringValue": scope_name}})
+        if scope_version:
+            attributes.append({"key": "otel.scope.version", "value": {"stringValue": scope_version}})
+
+        if span.attributes:
+            for key, value in span.attributes.items():
+                attributes.append(self._to_otlp_attribute(key, value))
+
+        self._promote_genai_event_attributes(span, attributes)
+
+        parent_span_id = None
+        if span.parent and hasattr(span.parent, "span_id"):
+            parent_span_id = format(span.parent.span_id, "016x")
+
+        return {
+            "traceId": format(span.context.trace_id, "032x"),
+            "spanId": format(span.context.span_id, "016x"),
+            "parentSpanId": parent_span_id,
+            "name": span.name,
+            "kind": span.kind.value if span.kind else 1,
+            "startTimeUnixNano": str(span.start_time),
+            "endTimeUnixNano": str(span.end_time),
+            "attributes": attributes,
+            "status": {"code": span.status.status_code.value} if span.status else {},
+        }
+
+    def _promote_genai_event_attributes(self, span: ReadableSpan, attributes: list[dict]) -> None:
+        """Promote OTel GenAI message attributes from span events to span attributes.
+
+        OTel GenAI semantic convention frameworks may store message content
+        (gen_ai.input.messages, gen_ai.output.messages) in span events rather
+        than span attributes. This promotes those event attributes so downstream
+        processors can access them uniformly from span attributes alone.
+        """
+        if not hasattr(span, "events") or not span.events:
+            return
+
+        from ..trace_attrs import OTEL_GENAI_INPUT_MESSAGES, OTEL_GENAI_OUTPUT_MESSAGES
+
+        _genai_event_keys = {OTEL_GENAI_INPUT_MESSAGES, OTEL_GENAI_OUTPUT_MESSAGES}
+        existing_keys = {a["key"] for a in attributes}
+
+        for event in span.events:
+            if not event.attributes:
+                continue
+            for key, value in event.attributes.items():
+                if key in _genai_event_keys and key not in existing_keys:
+                    attributes.append(self._to_otlp_attribute(key, value))
+                    existing_keys.add(key)
+
+    def _to_otlp_attribute(self, key: str, value: Any) -> dict:
+        if isinstance(value, bool):
+            return {"key": key, "value": {"boolValue": value}}
+        elif isinstance(value, int):
+            return {"key": key, "value": {"intValue": value}}
+        elif isinstance(value, float):
+            return {"key": key, "value": {"doubleValue": value}}
+        else:
+            return {"key": key, "value": {"stringValue": str(value)}}
+
+
+class AgentEvalsLogStreamingProcessor:
+    """OTel log processor that streams GenAI logs to agentevals dev server via WebSocket.
+
+    This processor shares the same WebSocket connection and session as AgentEvalsStreamingProcessor.
+    It extracts input/output messages from GenAI semantic convention logs and streams them.
+    """
+
+    def __init__(self, span_processor: AgentEvalsStreamingProcessor):
+        """Initialize with a reference to the span processor for shared connection."""
+        self.span_processor = span_processor
+
+    def on_emit(self, log_data):
+        """Called when a log record is emitted."""
+        log_record = log_data.log_record
+
+        # Only process GenAI message logs
+        if not log_record.event_name or not log_record.event_name.startswith("gen_ai."):
+            return
+
+        logger.info(f"Log emitted: event={log_record.event_name}")
+
+        if not self.span_processor._connected or not self.span_processor.websocket or not self.span_processor.loop:
+            return
+
+        try:
+            log_json = {
+                "event_name": log_record.event_name,
+                "timestamp": str(log_record.timestamp) if log_record.timestamp else None,
+                "body": log_record.body,
+                "attributes": {},
+            }
+
+            if log_record.attributes:
+                for key, value in log_record.attributes.items():
+                    log_json["attributes"][key] = value
+
+            future = asyncio.run_coroutine_threadsafe(self._send_log(log_json), self.span_processor.loop)
+
+            def handle_send_complete(fut):
+                try:
+                    fut.result()
+                except Exception as exc:
+                    logger.error(f"Failed to send log: {exc}")
+
+            future.add_done_callback(handle_send_complete)
+
+        except Exception as exc:
+            logger.warning("Failed to process log: %s", exc)
+
+    async def _send_log(self, log_json: dict) -> None:
+        if not self.span_processor.websocket:
+            raise ConnectionError("WebSocket not connected")
+
+        message = {
+            "type": "log",
+            "session_id": self.span_processor.session_id,
+            "log": log_json,
+        }
+        await self.span_processor.websocket.send(json.dumps(message))
+
+    def shutdown(self):
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000):
+        return True

agentevals/streaming/session.py

@@ -0,0 +1,36 @@
+"""Trace session tracking for live streaming."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import UTC, datetime
+
+MAX_SPANS_PER_SESSION = 10000
+MAX_LOGS_PER_SESSION = 5000
+
+
+@dataclass
+class TraceSession:
+    """Represents an active trace session from a streaming agent."""
+
+    session_id: str
+    trace_id: str
+    eval_set_id: str | None
+    spans: list[dict] = field(default_factory=list)
+    logs: list[dict] = field(default_factory=list)
+    started_at: datetime = field(default_factory=lambda: datetime.now(UTC))
+    is_complete: bool = False
+    completed_at: datetime | None = None
+    metadata: dict = field(default_factory=dict)
+    source: str = "websocket"
+    has_root_span: bool = False
+    trace_ids: set[str] = field(default_factory=set)
+    invocations: list[dict] = field(default_factory=list)
+
+    def can_accept_span(self) -> bool:
+        """Check if session can accept another span without exceeding limits."""
+        return len(self.spans) < MAX_SPANS_PER_SESSION
+
+    def can_accept_log(self) -> bool:
+        """Check if session can accept another log without exceeding limits."""
+        return len(self.logs) < MAX_LOGS_PER_SESSION