diff --git a/agentx/integrations/_traced_call.py b/agentx/integrations/_traced_call.py index bdd5ade..11c704b 100644 --- a/agentx/integrations/_traced_call.py +++ b/agentx/integrations/_traced_call.py @@ -14,6 +14,7 @@ import asyncio import inspect import json +import time from typing import Any, Callable, Dict, Optional from agentx.tracing.tracer import Tracer, _safe_serialize @@ -107,6 +108,7 @@ def finish_llm_call( cache_read_tokens: Optional[int] = None, cache_write_tokens: Optional[int] = None, tool_definitions: Optional[list] = None, + call_metadata: Optional[Dict[str, Any]] = None, ) -> None: """ Close out one raw-client LLM call - shared by the ``on_finish``/exit @@ -150,13 +152,19 @@ def finish_llm_call( output_tokens=output_tokens, cache_read_tokens=cache_read_tokens, cache_write_tokens=cache_write_tokens, + metadata=call_metadata, ) return # A patched provider call outside any active span becomes its own root trace - it is a bare # model call, so stamp it "llm" rather than leaving the kind unset. span = tracer.trace( - name, metadata=metadata, framework=framework, model=model, session_id=session_id, span_kind="llm" + name, + metadata={**(metadata or {}), **(call_metadata or {})} if (metadata or call_metadata) else None, + framework=framework, + model=model, + session_id=session_id, + span_kind="llm", ) span.__enter__() span._start = start_t @@ -173,3 +181,215 @@ def finish_llm_call( if cache_write_tokens: span._cache_write_tokens = cache_write_tokens span.__exit__(None, None, None) + + +# --------------------------------------------------------------------------- +# Streaming: wrap a provider's chunk stream so the trace is built from what +# was actually streamed, without touching the caller's consumption of it. +# --------------------------------------------------------------------------- + +class StreamAccumulator: + """ + What a streaming patch feeds each chunk into. Subclasses collect the + provider-specific pieces (text deltas, tool-call deltas, the usage block + that only arrives on the final chunk) and hand back the finished picture + in ``result()``. + """ + + def feed(self, chunk: Any) -> None: # pragma: no cover - interface + raise NotImplementedError + + def result(self) -> Dict[str, Any]: # pragma: no cover - interface + raise NotImplementedError + + +class TracedStream: + """ + Transparent proxy over a provider ``Stream``/``AsyncStream``: iterates the + real object, feeds every chunk to the accumulator, and calls ``on_finish`` + exactly once when the stream is exhausted, raises, is closed (``close()``, + ``with``/``async with`` exit), or is dropped part-way and garbage + collected - so an abandoned stream still records what it streamed. + + Latency is measured to the LAST chunk (the response as the caller saw it), + and the time to the FIRST chunk is reported separately as + ``time_to_first_token_ms`` - the two numbers a streaming call is judged by. + + Attribute access falls through to the wrapped stream (``.response``, + provider helpers), and ``__iter__``/``__aiter__`` return ``self`` so early + ``break`` leaves no half-driven generator behind. + """ + + def __init__( + self, + stream: Any, + accumulator: StreamAccumulator, + on_finish: Callable[[Dict[str, Any], Optional[str]], None], + ) -> None: + self._stream = stream + self._accumulator = accumulator + self._on_finish = on_finish + self._done = False + self._first_chunk_t: Optional[float] = None + self._last_chunk_t: Optional[float] = None + self._sync_iter: Any = None + self._async_iter: Any = None + + # -- bookkeeping --------------------------------------------------------- + + def _observe(self, chunk: Any) -> None: + now = time.time() + if self._first_chunk_t is None: + self._first_chunk_t = now + self._last_chunk_t = now + try: + self._accumulator.feed(chunk) + except Exception: + # A malformed chunk must never break the caller's stream; it just + # goes uncounted in the trace. + pass + + def _finish(self, error: Optional[str]) -> None: + if self._done: + return + self._done = True + try: + result = self._accumulator.result() + except Exception: + result = {} + start_t = result.pop("_start_t", None) + result["time_to_first_token_ms"] = ( + int((self._first_chunk_t - start_t) * 1000) if self._first_chunk_t is not None and start_t is not None else None + ) + # The response "ended" at its last chunk, not at whatever later moment the caller closed + # or dropped the stream - that is the latency the user experienced. + result["end_t"] = self._last_chunk_t if self._last_chunk_t is not None else time.time() + self._on_finish(result, error) + + @property + def first_chunk_at(self) -> Optional[float]: + return self._first_chunk_t + + # -- sync iteration ------------------------------------------------------ + + def __iter__(self) -> "TracedStream": + return self + + def __next__(self) -> Any: + if self._sync_iter is None: + self._sync_iter = iter(self._stream) + try: + chunk = next(self._sync_iter) + except StopIteration: + self._finish(None) + raise + except BaseException as exc: + self._finish(str(exc)) + raise + self._observe(chunk) + return chunk + + # -- async iteration ----------------------------------------------------- + + def __aiter__(self) -> "TracedStream": + return self + + async def __anext__(self) -> Any: + if self._async_iter is None: + self._async_iter = self._stream.__aiter__() + try: + chunk = await self._async_iter.__anext__() + except StopAsyncIteration: + self._finish(None) + raise + except BaseException as exc: + self._finish(str(exc)) + raise + self._observe(chunk) + return chunk + + # -- context managers / close -------------------------------------------- + + def __enter__(self) -> "TracedStream": + enter = getattr(self._stream, "__enter__", None) + if enter is not None: + enter() + return self + + def __exit__(self, exc_type, exc_val, tb) -> Any: + exit_ = getattr(self._stream, "__exit__", None) + result = exit_(exc_type, exc_val, tb) if exit_ is not None else None + self._finish(str(exc_val) if exc_val else None) + return result + + async def __aenter__(self) -> "TracedStream": + enter = getattr(self._stream, "__aenter__", None) + if enter is not None: + await enter() + return self + + async def __aexit__(self, exc_type, exc_val, tb) -> Any: + exit_ = getattr(self._stream, "__aexit__", None) + result = await exit_(exc_type, exc_val, tb) if exit_ is not None else None + self._finish(str(exc_val) if exc_val else None) + return result + + def close(self) -> None: + close = getattr(self._stream, "close", None) + try: + if close is not None: + close() + finally: + self._finish(None) + + async def aclose(self) -> None: + # openai's AsyncStream spells its close as `async def close()`; httpx-style streams + # spell it `aclose()`. Await whichever one answers with an awaitable. + close = getattr(self._stream, "aclose", None) or getattr(self._stream, "close", None) + try: + if close is not None: + result = close() + if inspect.isawaitable(result): + await result + finally: + self._finish(None) + + def __getattr__(self, item: str) -> Any: + return getattr(self._stream, item) + + def __del__(self) -> None: + # Best effort only: a stream the caller stopped reading and dropped still records the + # chunks it did see. Never raises - a destructor exception is unactionable noise. + try: + self._finish(None) + except Exception: + pass + + +def trace_stream( + result: Any, + accumulator: StreamAccumulator, + on_finish: Callable[[Dict[str, Any], Optional[str]], None], +) -> Any: + """ + Wrap the value a patched ``create(..., stream=True)`` returned. A sync + client hands back the stream object directly; an async client hands back + a coroutine that resolves to it, so the wrapping is deferred until the + real stream exists - the caller's ``await`` is unchanged either way. + """ + if asyncio.iscoroutine(result) or inspect.isawaitable(result): + return _await_and_wrap(result, accumulator, on_finish) + return TracedStream(result, accumulator, on_finish) + + +async def _await_and_wrap( + awaitable: Any, + accumulator: StreamAccumulator, + on_finish: Callable[[Dict[str, Any], Optional[str]], None], +) -> Any: + try: + stream = await awaitable + except Exception as exc: + on_finish({}, str(exc)) + raise + return TracedStream(stream, accumulator, on_finish) diff --git a/agentx/integrations/anthropic.py b/agentx/integrations/anthropic.py index 002cd2f..fd3e02b 100644 --- a/agentx/integrations/anthropic.py +++ b/agentx/integrations/anthropic.py @@ -13,6 +13,12 @@ Works with both ``anthropic.Anthropic`` and ``anthropic.AsyncAnthropic`` clients. +Both streaming shapes are traced: the ``client.messages.stream(...)`` helper +(a context manager with ``get_final_message()``) and the raw +``messages.create(..., stream=True)`` event stream, which is wrapped in a +transparent proxy that assembles the reply, tool-use blocks, and token usage +from the events as the caller consumes them. + Requires: ``pip install "agentx-python[anthropic]"`` """ from __future__ import annotations @@ -22,7 +28,13 @@ from typing import Any, Dict, Optional, Tuple from agentx.tracing.tracer import Tracer, _safe_serialize -from agentx.integrations._traced_call import capture_tool_definitions, call_and_trace, finish_llm_call +from agentx.integrations._traced_call import ( + StreamAccumulator, + capture_tool_definitions, + call_and_trace, + finish_llm_call, + trace_stream, +) def _extract_output_text(response: Any) -> Optional[str]: @@ -92,6 +104,85 @@ def _extract_usage_tokens( return input_tokens, output_tokens, cache_read, cache_creation +class _MessageEventStreamAccumulator(StreamAccumulator): + """ + Rebuild a ``Message`` from the raw ``create(stream=True)`` event sequence: + ``message_start`` carries the input-side usage, ``content_block_start`` opens + a text or tool_use block, ``content_block_delta`` appends ``text_delta`` / + ``input_json_delta`` fragments to it, ``message_delta`` carries the + output-token count. Token accounting mirrors ``_extract_usage_tokens``. + """ + + def __init__(self, start_t: float) -> None: + self._start_t = start_t + self._blocks: Dict[int, Dict[str, Any]] = {} + self._input_tokens: Optional[int] = None + self._output_tokens: Optional[int] = None + self._cache_read: Optional[int] = None + self._cache_write: Optional[int] = None + self._model: Optional[str] = None + + def feed(self, event: Any) -> None: + event_type = getattr(event, "type", None) + if event_type == "message_start": + message = getattr(event, "message", None) + self._model = getattr(message, "model", None) or self._model + input_tokens, output_tokens, cache_read, cache_write = _extract_usage_tokens(getattr(message, "usage", None)) + self._input_tokens = input_tokens + self._cache_read = cache_read + self._cache_write = cache_write + if output_tokens: + self._output_tokens = output_tokens + elif event_type == "content_block_start": + index = getattr(event, "index", 0) or 0 + block = getattr(event, "content_block", None) + self._blocks[index] = { + "type": getattr(block, "type", None), + "name": getattr(block, "name", None), + "text": [getattr(block, "text", None) or ""] if getattr(block, "type", None) == "text" else [], + "json": [], + } + elif event_type == "content_block_delta": + index = getattr(event, "index", 0) or 0 + delta = getattr(event, "delta", None) + entry = self._blocks.setdefault(index, {"type": None, "name": None, "text": [], "json": []}) + delta_type = getattr(delta, "type", None) + if delta_type == "text_delta": + entry["type"] = entry["type"] or "text" + entry["text"].append(getattr(delta, "text", None) or "") + elif delta_type == "input_json_delta": + entry["type"] = entry["type"] or "tool_use" + entry["json"].append(getattr(delta, "partial_json", None) or "") + elif event_type == "message_delta": + usage = getattr(event, "usage", None) + output_tokens = getattr(usage, "output_tokens", None) if usage is not None else None + if output_tokens is not None: + self._output_tokens = output_tokens + + def result(self) -> Dict[str, Any]: + texts = [] + tool_calls = [] + for _, block in sorted(self._blocks.items()): + if block["type"] == "text": + text = "".join(block["text"]) + if text: + texts.append(text) + elif block["type"] == "tool_use": + tool_calls.append(f"{block['name'] or 'unknown'}({''.join(block['json'])})") + output: Optional[str] = "\n".join(texts) if texts else None + if output is None and tool_calls: + output = "[tool call] " + ", ".join(tool_calls) + return { + "_start_t": self._start_t, + "output": output, + "model": self._model, + "input_tokens": self._input_tokens, + "output_tokens": self._output_tokens, + "cache_read_tokens": self._cache_read, + "cache_write_tokens": self._cache_write, + } + + def patch_anthropic_client( client: Any, tracer: Tracer, @@ -140,6 +231,35 @@ def patched_create(*args, **kwargs): input_repr = _safe_serialize(input_messages) + if kwargs.get("stream"): + def on_stream_finish(collected: Dict[str, Any], error: Optional[str]) -> None: + finish_llm_call( + tracer, + name=name, + framework="anthropic", + metadata=metadata, + call_metadata={"streaming": True, "timeToFirstTokenMs": collected.get("time_to_first_token_ms")}, + session_id=session_id, + start_t=start_t, + end_t=collected.get("end_t") or time.time(), + input_repr=input_repr, + output=collected.get("output"), + model=collected.get("model") or model, + input_tokens=collected.get("input_tokens"), + output_tokens=collected.get("output_tokens"), + cache_read_tokens=collected.get("cache_read_tokens"), + cache_write_tokens=collected.get("cache_write_tokens"), + error=error, + tool_definitions=tool_definitions, + ) + + try: + result = original(*args, **kwargs) + except Exception as exc: + on_stream_finish({}, str(exc)) + raise + return trace_stream(result, _MessageEventStreamAccumulator(start_t), on_stream_finish) + def on_finish(response: Optional[Any], error: Optional[str]) -> None: end_t = time.time() output = None @@ -237,36 +357,47 @@ def build_and_send(end_t: float, error: Optional[str], final_message: Optional[A ) class _TracedStream: - """Thin wrapper that records timing when the stream context exits.""" + """ + Thin wrapper that records the final message when the stream context exits. + ``ctx`` is the SDK's stream *manager*; the ``MessageStream`` it yields on enter is + what carries ``get_final_message()``, and it must be read BEFORE the manager's exit + closes it - reading it off the manager after close silently yielded no output. + """ + + _inner: Any = None def __enter__(self_inner): - return ctx.__enter__() + self_inner._inner = ctx.__enter__() + return self_inner._inner def __exit__(self_inner, exc_type, exc_val, tb): - result = ctx.__exit__(exc_type, exc_val, tb) end_t = time.time() error = str(exc_val) if exc_val else None final_message = None - try: - final_message = ctx.get_final_message() - except Exception: - pass + if error is None and self_inner._inner is not None: + try: + final_message = self_inner._inner.get_final_message() + except Exception: + pass + result = ctx.__exit__(exc_type, exc_val, tb) build_and_send(end_t, error, final_message) return result async def __aenter__(self_inner): - return await ctx.__aenter__() + self_inner._inner = await ctx.__aenter__() + return self_inner._inner async def __aexit__(self_inner, exc_type, exc_val, tb): - result = await ctx.__aexit__(exc_type, exc_val, tb) end_t = time.time() error = str(exc_val) if exc_val else None final_message = None - try: - raw = ctx.get_final_message() - final_message = await raw if inspect.isawaitable(raw) else raw - except Exception: - pass + if error is None and self_inner._inner is not None: + try: + raw = self_inner._inner.get_final_message() + final_message = await raw if inspect.isawaitable(raw) else raw + except Exception: + pass + result = await ctx.__aexit__(exc_type, exc_val, tb) build_and_send(end_t, error, final_message) return result diff --git a/agentx/integrations/openai.py b/agentx/integrations/openai.py index 74ad839..3f6ac0f 100644 --- a/agentx/integrations/openai.py +++ b/agentx/integrations/openai.py @@ -17,8 +17,11 @@ Works with both ``openai.OpenAI`` and ``openai.AsyncOpenAI`` clients. -Streaming calls (``stream=True``) are passed through untouched and are not -currently traced - see ``patch_openai_client``'s docstring. +Streaming calls (``stream=True``) are traced too: the returned stream is +wrapped in a transparent proxy that assembles the reply from the chunks as the +caller consumes them, so the trace carries the full text, tool calls, and +(with ``stream_options={"include_usage": True}``) token usage, plus the time +to first token. Requires: ``pip install "agentx-python[openai]"`` """ @@ -28,7 +31,13 @@ from typing import Any, Dict, Optional, Tuple from agentx.tracing.tracer import Tracer, _safe_serialize -from agentx.integrations._traced_call import capture_tool_definitions, call_and_trace, finish_llm_call +from agentx.integrations._traced_call import ( + StreamAccumulator, + capture_tool_definitions, + call_and_trace, + finish_llm_call, + trace_stream, +) def _extract_output_text(response: Any) -> Optional[str]: @@ -75,6 +84,68 @@ def _extract_usage_tokens(usage: Any) -> Tuple[Optional[int], Optional[int], Opt return getattr(usage, "prompt_tokens", None), getattr(usage, "completion_tokens", None), cached_tokens +class _ChatCompletionStreamAccumulator(StreamAccumulator): + """ + Rebuild a ``ChatCompletion``-shaped result from ``ChatCompletionChunk``s: + text deltas concatenate per choice, tool-call deltas merge by index (name + arrives once, arguments arrive as fragments), and the ``usage`` block - + present only on the final chunk, and only when the caller asked for it + with ``stream_options={"include_usage": True}`` - is kept when it appears. + """ + + def __init__(self, start_t: float) -> None: + self._start_t = start_t + self._texts: Dict[int, list] = {} + self._tool_calls: Dict[int, Dict[str, Any]] = {} + self._usage: Any = None + self._model: Optional[str] = None + + def feed(self, chunk: Any) -> None: + usage = getattr(chunk, "usage", None) + if usage is not None: + self._usage = usage + model = getattr(chunk, "model", None) + if model and not self._model: + self._model = model + for choice in getattr(chunk, "choices", None) or []: + index = getattr(choice, "index", 0) or 0 + delta = getattr(choice, "delta", None) + if delta is None: + continue + content = getattr(delta, "content", None) + if content: + self._texts.setdefault(index, []).append(content) + for tc in getattr(delta, "tool_calls", None) or []: + key = getattr(tc, "index", 0) or 0 + entry = self._tool_calls.setdefault(key, {"name": None, "arguments": []}) + fn = getattr(tc, "function", None) + fn_name = getattr(fn, "name", None) if fn is not None else None + fn_args = getattr(fn, "arguments", None) if fn is not None else None + if fn_name: + entry["name"] = fn_name + if fn_args: + entry["arguments"].append(fn_args) + + def result(self) -> Dict[str, Any]: + texts = ["".join(parts) for _, parts in sorted(self._texts.items())] + output: Optional[str] = "\n".join(t for t in texts if t) or None + if output is None and self._tool_calls: + described = [ + f"{entry['name'] or 'unknown'}({''.join(entry['arguments'])})" + for _, entry in sorted(self._tool_calls.items()) + ] + output = "[tool call] " + ", ".join(described) + input_tokens, output_tokens, cache_read_tokens = _extract_usage_tokens(self._usage) + return { + "_start_t": self._start_t, + "output": output, + "model": self._model, + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cache_read_tokens": cache_read_tokens, + } + + def patch_openai_client( client: Any, tracer: Tracer, @@ -92,12 +163,14 @@ def patch_openai_client( client's ``create()`` returns a coroutine, which is detected and awaited before the trace is built. - Calls made with ``stream=True`` are passed through untouched and are not - traced by this function: safely wrapping a (sync or async) chunk - iterator without disrupting the caller's own consumption of it needs - different handling than a single request/response call, so it's left - unpatched rather than risking a partially-consumed or double-consumed - stream for the caller. + Calls made with ``stream=True`` return a transparent proxy over the + provider's stream (see ``_traced_call.TracedStream``): iteration, ``with``, + ``close()`` and attribute access all pass through to the real stream, and + the trace is built from the chunks the caller actually consumed - text and + tool calls assembled from the deltas, token usage from the final chunk + when ``stream_options={"include_usage": True}`` was requested (OpenAI omits + usage from streams otherwise), latency to the last chunk, and the time to + first token in the trace metadata. """ chat = getattr(client, "chat", None) completions = getattr(chat, "completions", None) if chat is not None else None @@ -123,17 +196,41 @@ def _patch_chat_completions_create( return # already patched def patched_create(*args, **kwargs): - if kwargs.get("stream"): - # Not traced - see patch_openai_client's docstring. Passed - # through completely untouched, sync or async. - return original(*args, **kwargs) - start_t = time.time() input_messages = kwargs.get("messages") or (args[0] if args else None) model = kwargs.get("model") input_repr = _safe_serialize(input_messages) tool_definitions = capture_tool_definitions(kwargs.get("tools")) + if kwargs.get("stream"): + def on_stream_finish(collected: Dict[str, Any], error: Optional[str]) -> None: + ttft = collected.get("time_to_first_token_ms") + finish_llm_call( + tracer, + name=name, + framework=framework, + metadata=metadata, + call_metadata={"streaming": True, "timeToFirstTokenMs": ttft}, + session_id=session_id, + start_t=start_t, + end_t=collected.get("end_t") or time.time(), + input_repr=input_repr, + output=collected.get("output"), + model=collected.get("model") or model, + input_tokens=collected.get("input_tokens"), + output_tokens=collected.get("output_tokens"), + cache_read_tokens=collected.get("cache_read_tokens"), + error=error, + tool_definitions=tool_definitions, + ) + + try: + result = original(*args, **kwargs) + except Exception as exc: + on_stream_finish({}, str(exc)) + raise + return trace_stream(result, _ChatCompletionStreamAccumulator(start_t), on_stream_finish) + def on_finish(response: Optional[Any], error: Optional[str]) -> None: end_t = time.time() output = None diff --git a/agentx/monitor/__init__.py b/agentx/monitor/__init__.py index 8f7d2ad..03c4bc0 100644 --- a/agentx/monitor/__init__.py +++ b/agentx/monitor/__init__.py @@ -12,12 +12,16 @@ from agentx.monitor.profile import MonitorProfileClient from agentx.monitor.review_queue import ReviewQueueClient, ReviewQueueItem from agentx.monitor.rules import MonitorRule, MonitorRulesClient +from agentx.monitor.alert_rules import AlertEvent, AlertRule, AlertRulesClient from agentx.monitor.scorers import AgentXScorersError, ScorersClient from agentx.monitor.scorer_groups import AgentXScorerGroupsError, ScorerGroup, ScorerGroupsClient from agentx.monitor.sessions import MonitorSessionClient from agentx.monitor.signals import MonitorSignalClient __all__ = [ + "AlertEvent", + "AlertRule", + "AlertRulesClient", "AgentXImprovementGroupsError", "AgentXJudgeScorersError", "AgentXMonitorError", diff --git a/agentx/monitor/alert_rules.py b/agentx/monitor/alert_rules.py new file mode 100644 index 0000000..d2d3c82 --- /dev/null +++ b/agentx/monitor/alert_rules.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional, TYPE_CHECKING + +if TYPE_CHECKING: + from agentx.monitor.client import MonitorClient + +ALERT_METRICS = ("failureRate", "toolFailureRate", "p95LatencyMs", "estimatedCostUsd", "judgeFailures", "traceCount") +ALERT_CHANNEL_KINDS = ("slack", "teams", "pagerduty", "email", "webhook") + +# snake_case kwargs -> wire keys. The engine reads camelCase only (the wire convention); an +# unknown snake_case key would be silently ignored, so update() refuses it instead. +_ALIASES = { + "window_minutes": "windowMinutes", + "agent_id": "agentId", + "cooldown_minutes": "cooldownMinutes", +} + + +class AlertRule(dict): + """Wire object for one KPI alert rule (dict subclass so unknown fields round-trip).""" + + @property + def id(self) -> str: + return self["_id"] + + @property + def enabled(self) -> bool: + return bool(self.get("enabled")) + + @property + def state(self) -> str: + """``"ok"`` or ``"firing"``.""" + return str(self.get("state") or "ok") + + @property + def last_value(self) -> Optional[float]: + return self.get("lastValue") + + @property + def fired_count(self) -> int: + return int(self.get("firedCount") or 0) + + +class AlertEvent(dict): + """One row of a rule's notification history: ``kind`` is ``triggered`` / ``repeat`` / + ``resolved`` / ``test``; ``deliveries`` lists each channel's outcome.""" + + @property + def kind(self) -> str: + return str(self.get("kind")) + + @property + def delivered(self) -> bool: + deliveries = self.get("deliveries") or [] + return bool(deliveries) and all(bool(d.get("ok")) for d in deliveries) + + +def slack(url: str) -> Dict[str, str]: + """A Slack incoming-webhook channel.""" + return {"kind": "slack", "target": url} + + +def teams(url: str) -> Dict[str, str]: + """A Microsoft Teams incoming-webhook (or Workflows) channel.""" + return {"kind": "teams", "target": url} + + +def pagerduty(routing_key: str) -> Dict[str, str]: + """A PagerDuty Events API v2 integration - ``routing_key`` is the integration key.""" + return {"kind": "pagerduty", "target": routing_key} + + +def email(address: str) -> Dict[str, str]: + """An email recipient (needs a mailer configured on the engine).""" + return {"kind": "email", "target": address} + + +def webhook(url: str) -> Dict[str, str]: + """A generic JSON webhook receiving the full structured notification.""" + return {"kind": "webhook", "target": url} + + +class AlertRulesClient: + """Surfaced as ``client.monitor.alert_rules``: KPI alert rules. + + An alert rule watches an AGGREGATE over a sliding window - one of + ``failureRate``, ``toolFailureRate``, ``p95LatencyMs``, ``estimatedCostUsd``, + ``judgeFailures``, ``traceCount`` - and pages typed channels (Slack, Teams, + PagerDuty, email, generic webhook) when it crosses a threshold. The engine + evaluates every enabled rule once a minute with an Alertmanager-style + lifecycle: one ``triggered`` notification when the rule starts breaching, a + ``repeat`` every ``cooldown_minutes`` while it keeps breaching, and a + ``resolved`` notification when it recovers. This is distinct from a scorer's + per-verdict alert threshold and from an automation rule's per-trace routing. + + Example:: + + from agentx.monitor.alert_rules import slack, pagerduty + + rule = client.monitor.alert_rules.create( + "Failure rate above 10%", + metric="failureRate", operator="gt", threshold=0.10, window_minutes=15, + severity="high", + channels=[slack("https://hooks.slack.com/services/..."), pagerduty("R0123...")], + ) + client.monitor.alert_rules.test(rule.id) # sends a TEST page to every channel + """ + + def __init__(self, client: "MonitorClient"): + self._client = client + + def _request(self, method: str, path: str, **kwargs: Any) -> Any: + return self._client._request(method, path, base=self._client._api_root(), **kwargs) + + def list(self) -> List[AlertRule]: + data = self._request("GET", "/agent-monitoring/alert-rules") + return [AlertRule(r) for r in data.get("rules", [])] + + def get(self, rule_id: str) -> AlertRule: + data = self._request("GET", f"/agent-monitoring/alert-rules/{rule_id}") + return AlertRule(data.get("rule", data)) + + def create( + self, + name: str, + *, + metric: str, + operator: str, + threshold: float, + window_minutes: int, + channels: List[Dict[str, str]], + agent_id: Optional[str] = None, + severity: str = "high", + cooldown_minutes: int = 60, + enabled: bool = True, + ) -> AlertRule: + """Create a rule. ``operator`` is ``"gt"`` (above) or ``"lt"`` (below); rates are + fractions (``0.10`` = 10%), latency is milliseconds, cost is USD. ``channels`` takes the + dicts the module-level helpers build (``slack(url)``, ``pagerduty(key)``, ...).""" + if metric not in ALERT_METRICS: + raise ValueError(f"metric must be one of {ALERT_METRICS}, got {metric!r}") + if operator not in ("gt", "lt"): + raise ValueError(f"operator must be 'gt' or 'lt', got {operator!r}") + for channel in channels: + if channel.get("kind") not in ALERT_CHANNEL_KINDS: + raise ValueError(f"channel kind must be one of {ALERT_CHANNEL_KINDS}, got {channel.get('kind')!r}") + payload: Dict[str, Any] = { + "name": name, + "metric": metric, + "operator": operator, + "threshold": threshold, + "windowMinutes": window_minutes, + "channels": channels, + "severity": severity, + "cooldownMinutes": cooldown_minutes, + "enabled": enabled, + } + if agent_id is not None: + payload["agentId"] = agent_id + # Server-side write: a timeout retry would create a duplicate rule that pages twice on + # every incident - no transport retry (same posture as rules.create). + data = self._request("POST", "/agent-monitoring/alert-rules", json=payload, retry=False) + return AlertRule(data.get("rule", data)) + + def update(self, rule_id: str, **fields: Any) -> AlertRule: + """Sparse update; snake_case kwargs are mapped to the wire. Changing the metric, + operator, threshold, window, or agent resets the rule's firing state.""" + payload: Dict[str, Any] = {} + for key, value in fields.items(): + wire_key = _ALIASES.get(key, key) + if "_" in wire_key: + raise ValueError( + f"Unknown alert rule field {key!r} - the engine reads camelCase keys and would " + "silently ignore this (see AlertRulesClient.create for the field names)." + ) + payload[wire_key] = value + data = self._request("PUT", f"/agent-monitoring/alert-rules/{rule_id}", json=payload) + return AlertRule(data.get("rule", data)) + + def delete(self, rule_id: str) -> None: + """Deletes the rule and its history. A PagerDuty incident the rule opened is not + resolved by this - close it in PagerDuty.""" + self._request("DELETE", f"/agent-monitoring/alert-rules/{rule_id}", retry=False) + + def events(self, rule_id: str, limit: int = 50) -> List[AlertEvent]: + """The rule's notification history, newest first, with per-channel delivery results.""" + data = self._request("GET", f"/agent-monitoring/alert-rules/{rule_id}/events", params={"limit": limit}) + return [AlertEvent(e) for e in data.get("events", [])] + + def test(self, rule_id: str) -> AlertEvent: + """Send a TEST notification to the rule's channels with the metric's live value and + return the recorded event - ``event.delivered`` says whether every channel accepted it. + Never changes the rule's firing state.""" + data = self._request("POST", f"/agent-monitoring/alert-rules/{rule_id}/test", json={}, retry=False) + return AlertEvent(data.get("event", data)) + + def preview(self, metric: str, window_minutes: int, agent_id: Optional[str] = None) -> Dict[str, Any]: + """What ``metric`` reads right now over the last ``window_minutes`` - the same + computation the sweep runs. Returns ``{"value": float | None, "valueLabel": str, ...}``; + ``value`` is ``None`` when the window has no data for a rate metric.""" + payload: Dict[str, Any] = {"metric": metric, "windowMinutes": window_minutes} + if agent_id is not None: + payload["agentId"] = agent_id + return self._request("POST", "/agent-monitoring/alert-rules/preview", json=payload) + + def run_sweep(self) -> Dict[str, Any]: + """Evaluate this project's rules now instead of waiting for the next minute tick.""" + return self._request("POST", "/agent-monitoring/alert-rules/sweep/run", json={}, retry=False) diff --git a/agentx/monitor/client.py b/agentx/monitor/client.py index b422406..8af4152 100644 --- a/agentx/monitor/client.py +++ b/agentx/monitor/client.py @@ -139,6 +139,11 @@ def __init__( # Automation rules: route matching traffic into review / a dataset / a webhook. self.rules = MonitorRulesClient(self) + from agentx.monitor.alert_rules import AlertRulesClient + + # KPI alert rules: threshold pages (Slack/Teams/PagerDuty/email/webhook) on failure + # rate, p95 latency, spend, judge failures, and traffic volume over a window. + self.alert_rules = AlertRulesClient(self) from agentx.monitor.scorers import ScorersClient # Scorers-catalog administration as code: template enable/disable, code/external scorer # CRUD and dry runs - full parity with the dashboard's Scorers page (P1.3). diff --git a/agentx/tracing/tracer.py b/agentx/tracing/tracer.py index 4e7b1b4..a5a1d32 100644 --- a/agentx/tracing/tracer.py +++ b/agentx/tracing/tracer.py @@ -297,6 +297,7 @@ def _record_llm_call( output_tokens: Optional[int] = None, cache_read_tokens: Optional[int] = None, cache_write_tokens: Optional[int] = None, + metadata: Optional[Dict[str, Any]] = None, ) -> None: """Record one LLM-call child span (e.g. one patched Anthropic call) under this span - name left unset so _merge_child_run auto-numbers it "LLM Call N". ``framework`` lets the @@ -315,6 +316,9 @@ def _record_llm_call( "outputTokenSize": output_tokens, "cacheReadTokenSize": cache_read_tokens, "cacheWriteTokenSize": cache_write_tokens, + # Per-call facts that belong on the child row (a streamed call's time to first + # token), not on the parent trace's metadata. + "metadata": metadata, }], input=input, output=output, @@ -479,6 +483,7 @@ def _merge_child_run( output_tokens=step.get("outputTokenSize"), cache_read_tokens=step.get("cacheReadTokenSize"), cache_write_tokens=step.get("cacheWriteTokenSize"), + metadata=step.get("metadata") or None, # Stated, so a step named anything other than "LLM Call N" still classifies - # the backend's name regex was the only thing holding this together. Steps # may state their own kind (crewai.py's task steps carry "agent"); the diff --git a/tests/test_alert_rules.py b/tests/test_alert_rules.py new file mode 100644 index 0000000..51c2fef --- /dev/null +++ b/tests/test_alert_rules.py @@ -0,0 +1,128 @@ +"""Unit tests for client.monitor.alert_rules - wire-level, no engine required. The engine-side +contract (metric evaluation, the firing/resolved lifecycle, channel delivery) is pinned by the +engine's alertRules integration suite; this pins the SDK's half: camelCase wire keys, the +no-retry posture on writes, argument validation, and the channel helpers.""" + +from typing import Any, Dict, List + +import pytest + +from agentx.monitor.alert_rules import AlertEvent, AlertRule, AlertRulesClient, email, pagerduty, slack, teams, webhook + + +class FakeMonitorClient: + def __init__(self, responses: List[Any]): + self.calls: List[Dict[str, Any]] = [] + self._responses = responses + + def _api_root(self) -> str: + return "http://engine:4700/api/v1" + + def _request(self, method: str, path: str, base: str = "", **kwargs: Any) -> Any: + self.calls.append({"method": method, "path": path, "base": base, **kwargs}) + return self._responses.pop(0) if self._responses else {} + + +def test_channel_helpers_build_wire_dicts(): + assert slack("https://hooks.slack.com/x") == {"kind": "slack", "target": "https://hooks.slack.com/x"} + assert teams("https://t.example/hook") == {"kind": "teams", "target": "https://t.example/hook"} + assert pagerduty("R0123456789abcdef0123456789abcdef") == {"kind": "pagerduty", "target": "R0123456789abcdef0123456789abcdef"} + assert email("oncall@example.com") == {"kind": "email", "target": "oncall@example.com"} + assert webhook("https://ops.example/agentx") == {"kind": "webhook", "target": "https://ops.example/agentx"} + + +def test_create_sends_camelcase_wire_without_retry(): + fake = FakeMonitorClient([{"rule": {"_id": "a1", "state": "ok", "enabled": True, "firedCount": 0}}]) + rule = AlertRulesClient(fake).create( # type: ignore[arg-type] + "Failure rate above 10%", + metric="failureRate", + operator="gt", + threshold=0.1, + window_minutes=15, + channels=[slack("https://hooks.slack.com/x")], + agent_id="agent-7", + severity="critical", + cooldown_minutes=30, + ) + call = fake.calls[0] + assert call["method"] == "POST" + assert call["path"] == "/agent-monitoring/alert-rules" + assert call["base"] == "http://engine:4700/api/v1" + assert call["retry"] is False + assert call["json"] == { + "name": "Failure rate above 10%", + "metric": "failureRate", + "operator": "gt", + "threshold": 0.1, + "windowMinutes": 15, + "channels": [{"kind": "slack", "target": "https://hooks.slack.com/x"}], + "severity": "critical", + "cooldownMinutes": 30, + "enabled": True, + "agentId": "agent-7", + } + assert isinstance(rule, AlertRule) + assert rule.id == "a1" + assert rule.state == "ok" + assert rule.fired_count == 0 + + +def test_create_validates_metric_operator_and_channel_kind_locally(): + client = AlertRulesClient(FakeMonitorClient([])) # type: ignore[arg-type] + common = dict(operator="gt", threshold=1, window_minutes=5, channels=[slack("https://h/x")]) + with pytest.raises(ValueError, match="metric"): + client.create("x", metric="vibes", **common) # type: ignore[arg-type] + with pytest.raises(ValueError, match="operator"): + client.create("x", metric="traceCount", operator="ge", threshold=1, window_minutes=5, channels=[slack("https://h/x")]) + with pytest.raises(ValueError, match="channel kind"): + client.create("x", metric="traceCount", operator="lt", threshold=1, window_minutes=5, channels=[{"kind": "sms", "target": "1"}]) + + +def test_update_maps_snake_case_and_refuses_unknown_keys(): + fake = FakeMonitorClient([{"rule": {"_id": "a1", "state": "ok"}}]) + AlertRulesClient(fake).update("a1", window_minutes=60, cooldown_minutes=15, enabled=False) # type: ignore[arg-type] + call = fake.calls[0] + assert call["method"] == "PUT" + assert call["path"] == "/agent-monitoring/alert-rules/a1" + assert call["json"] == {"windowMinutes": 60, "cooldownMinutes": 15, "enabled": False} + with pytest.raises(ValueError, match="Unknown alert rule field"): + AlertRulesClient(fake).update("a1", sample_rate=0.5) # type: ignore[arg-type] + + +def test_events_test_preview_and_sweep_paths(): + fake = FakeMonitorClient( + [ + {"events": [{"kind": "triggered", "deliveries": [{"kind": "slack", "ok": True}, {"kind": "email", "ok": False}]}]}, + {"event": {"kind": "test", "deliveries": [{"kind": "slack", "ok": True}]}, "delivered": True}, + {"metric": "p95LatencyMs", "value": 812, "valueLabel": "812 ms"}, + {"evaluated": 2, "results": []}, + ] + ) + client = AlertRulesClient(fake) # type: ignore[arg-type] + + history = client.events("a1", limit=10) + assert fake.calls[0]["path"] == "/agent-monitoring/alert-rules/a1/events" + assert fake.calls[0]["params"] == {"limit": 10} + assert isinstance(history[0], AlertEvent) + assert history[0].kind == "triggered" + assert history[0].delivered is False # one channel failed + + sent = client.test("a1") + assert fake.calls[1]["method"] == "POST" + assert fake.calls[1]["path"] == "/agent-monitoring/alert-rules/a1/test" + assert fake.calls[1]["retry"] is False + assert sent.delivered is True + + preview = client.preview("p95LatencyMs", 15, agent_id="agent-7") + assert fake.calls[2]["json"] == {"metric": "p95LatencyMs", "windowMinutes": 15, "agentId": "agent-7"} + assert preview["value"] == 812 + + client.run_sweep() + assert fake.calls[3]["path"] == "/agent-monitoring/alert-rules/sweep/run" + assert fake.calls[3]["retry"] is False + + +def test_delete_never_retries(): + fake = FakeMonitorClient([{}]) + AlertRulesClient(fake).delete("a1") # type: ignore[arg-type] + assert fake.calls[0] == {"method": "DELETE", "path": "/agent-monitoring/alert-rules/a1", "base": "http://engine:4700/api/v1", "retry": False} diff --git a/tests/test_integrations.py b/tests/test_integrations.py index e3b1057..4fed5aa 100644 --- a/tests/test_integrations.py +++ b/tests/test_integrations.py @@ -189,6 +189,154 @@ async def create(self, **kwargs): assert kwargs["latency_ms"] >= 5 +def _anthropic_events(with_tool=False): + ns = types.SimpleNamespace + events = [ + ns(type="message_start", message=ns(model="claude-x", usage=_FakeAnthropicUsage(input_tokens=20, output_tokens=1, cache_read_input_tokens=8))), + ns(type="content_block_start", index=0, content_block=ns(type="text", text="")), + ns(type="content_block_delta", index=0, delta=ns(type="text_delta", text="Hel")), + ns(type="content_block_delta", index=0, delta=ns(type="text_delta", text="lo")), + ns(type="content_block_stop", index=0), + ] + if with_tool: + events += [ + ns(type="content_block_start", index=1, content_block=ns(type="tool_use", name="lookup_order")), + ns(type="content_block_delta", index=1, delta=ns(type="input_json_delta", partial_json='{"id":')), + ns(type="content_block_delta", index=1, delta=ns(type="input_json_delta", partial_json=' "A1"}')), + ns(type="content_block_stop", index=1), + ] + events += [ + ns(type="message_delta", delta=ns(stop_reason="end_turn"), usage=ns(output_tokens=7)), + ns(type="message_stop"), + ] + return events + + +def test_anthropic_raw_create_stream_is_traced_from_events(): + from agentx.integrations.anthropic import patch_anthropic_client + + class FakeMessages: + def create(self, **kwargs): + assert kwargs.get("stream") is True + return _FakeStream(_anthropic_events(with_tool=True)) + + client = types.SimpleNamespace(messages=FakeMessages()) + tracer = make_tracer() + patch_anthropic_client(client, tracer, name="claude-agent") + + events = list(client.messages.create(model="claude-x", system="be brief", messages=[{"role": "user", "content": "hi"}], stream=True)) + assert len(events) == 11 + + tracer._send.assert_called_once() + _, kwargs = tracer._send.call_args + # Text wins over the tool description when both exist, matching _extract_output_text. + assert kwargs["output"] == "Hello" + assert kwargs["model"] == "claude-x" + # input total folds the cached subset in (20 + 8), cache_read reported alongside. + assert kwargs["input_tokens"] == 28 + assert kwargs["cache_read_tokens"] == 8 + assert kwargs["output_tokens"] == 7 + assert kwargs["framework"] == "anthropic" + assert kwargs["metadata"]["streaming"] is True + # The system kwarg still lands in the traced input for streams. + assert "be brief" in str(kwargs["input"]) + + +def test_anthropic_raw_stream_tool_only_reply_is_described(): + from agentx.integrations.anthropic import patch_anthropic_client + ns = types.SimpleNamespace + events = [ + ns(type="message_start", message=ns(model="claude-x", usage=_FakeAnthropicUsage())), + ns(type="content_block_start", index=0, content_block=ns(type="tool_use", name="lookup_order")), + ns(type="content_block_delta", index=0, delta=ns(type="input_json_delta", partial_json='{"id": "A1"}')), + ns(type="message_delta", delta=ns(stop_reason="tool_use"), usage=ns(output_tokens=3)), + ns(type="message_stop"), + ] + + class FakeMessages: + def create(self, **kwargs): + return _FakeStream(events) + + client = types.SimpleNamespace(messages=FakeMessages()) + tracer = make_tracer() + patch_anthropic_client(client, tracer, name="claude-agent") + list(client.messages.create(model="claude-x", messages=[], stream=True)) + _, kwargs = tracer._send.call_args + assert kwargs["output"] == '[tool call] lookup_order({"id": "A1"})' + + +def test_anthropic_stream_helper_records_the_final_message(): + # Regression: the helper wrapper used to call get_final_message() on the stream MANAGER + # (which has no such method) after exit had closed it, so every .stream() trace silently + # carried no output and no tokens. + from agentx.integrations.anthropic import patch_anthropic_client + + final = _FakeAnthropicMessage("streamed reply", usage=_FakeAnthropicUsage(input_tokens=11, output_tokens=4)) + + class FakeMessageStream: + def __init__(self): + self.closed = False + self.text_stream = iter(["streamed ", "reply"]) + + def get_final_message(self): + assert not self.closed, "must be read before the manager closes the stream" + return final + + class FakeManager: + def __init__(self): + self.stream = FakeMessageStream() + + def __enter__(self): + return self.stream + + def __exit__(self, *exc): + self.stream.closed = True + return False + + class FakeMessages: + def create(self, **kwargs): + raise AssertionError("not used") + + def stream(self, **kwargs): + return FakeManager() + + client = types.SimpleNamespace(messages=FakeMessages()) + tracer = make_tracer() + patch_anthropic_client(client, tracer, name="claude-agent") + + with client.messages.stream(model="claude-x", messages=[{"role": "user", "content": "hi"}]) as s: + assert "".join(s.text_stream) == "streamed reply" + + tracer._send.assert_called_once() + _, kwargs = tracer._send.call_args + assert kwargs["output"] == "streamed reply" + assert kwargs["input_tokens"] == 11 + assert kwargs["output_tokens"] == 4 + + +def test_abandoned_stream_still_records_what_it_saw_when_collected(): + import gc + from agentx.integrations.openai import patch_openai_client + + class FakeCompletions: + def create(self, **kwargs): + return _FakeStream(_stream_chunks()) + + client = _fake_openai_client(FakeCompletions()) + tracer = make_tracer() + patch_openai_client(client, tracer, name="gpt-agent") + + result = client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True) + next(result) + next(result) + tracer._send.assert_not_called() + del result + gc.collect() + tracer._send.assert_called_once() + _, kwargs = tracer._send.call_args + assert kwargs["output"] == "Hello" + + def test_anthropic_async_client_propagates_and_records_errors(): from agentx.integrations.anthropic import patch_anthropic_client @@ -472,26 +620,254 @@ async def create(self, **kwargs): assert kwargs["latency_ms"] >= 5 -def test_openai_streaming_calls_are_passed_through_untraced(): +# --------------------------------------------------------------------------- +# 6a. Streaming: the wrapped stream is transparent to the caller and the trace +# is assembled from the chunks actually consumed. +# --------------------------------------------------------------------------- + +def _chunk(content=None, tool_calls=None, usage=None, model="gpt-4o-mini", with_choice=True): + delta = types.SimpleNamespace(content=content, tool_calls=tool_calls) + choices = [types.SimpleNamespace(index=0, delta=delta)] if with_choice else [] + return types.SimpleNamespace(choices=choices, usage=usage, model=model) + + +class _FakeStream: + """Mimics openai.Stream: iterable, context manager, closeable, with an attribute to delegate.""" + + def __init__(self, chunks, fail_after=None): + self._chunks = list(chunks) + self._fail_after = fail_after + self.closed = False + self.response = "raw-http-response" + + def __iter__(self): + for i, chunk in enumerate(self._chunks): + if self._fail_after is not None and i == self._fail_after: + raise RuntimeError("connection reset mid-stream") + yield chunk + + def __enter__(self): + return self + + def __exit__(self, *exc): + self.close() + return False + + def close(self): + self.closed = True + + +class _FakeAsyncStream: + def __init__(self, chunks): + self._chunks = list(chunks) + + def __aiter__(self): + return self._gen() + + async def _gen(self): + for chunk in self._chunks: + await asyncio.sleep(0) + yield chunk + + +def _stream_chunks(): + return [ + _chunk(content="Hel"), + _chunk(content="lo"), + _chunk(content=" there"), + # Final usage chunk (stream_options={"include_usage": True}): no choices, usage present. + _chunk(with_choice=False, usage=_FakeOpenAIUsage(prompt_tokens=9, completion_tokens=4)), + ] + + +def test_openai_stream_is_traced_from_consumed_chunks(): from agentx.integrations.openai import patch_openai_client - sentinel_stream = object() + stream = _FakeStream(_stream_chunks()) class FakeCompletions: def create(self, **kwargs): assert kwargs.get("stream") is True - return sentinel_stream + return stream client = _fake_openai_client(FakeCompletions()) tracer = make_tracer() patch_openai_client(client, tracer, name="gpt-agent") result = client.chat.completions.create( - model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], stream=True + model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], stream=True, stream_options={"include_usage": True} ) + # Transparent: the caller sees every chunk unchanged, attributes delegate to the real stream, + # and nothing is sent until the stream is exhausted. + seen = [] + for chunk in result: + seen.append(chunk) + tracer._send.assert_not_called() + assert len(seen) == 4 + assert result.response == "raw-http-response" - assert result is sentinel_stream - tracer._send.assert_not_called() + tracer._send.assert_called_once() + _, kwargs = tracer._send.call_args + assert kwargs["output"] == "Hello there" + assert kwargs["input_tokens"] == 9 + assert kwargs["output_tokens"] == 4 + assert kwargs["model"] == "gpt-4o-mini" + assert kwargs["framework"] == "openai" + assert kwargs["metadata"]["streaming"] is True + assert isinstance(kwargs["metadata"]["timeToFirstTokenMs"], int) + assert kwargs["error"] is None + + +def test_openai_stream_assembles_tool_call_deltas(): + from agentx.integrations.openai import patch_openai_client + + def tc(index, name=None, arguments=None): + return types.SimpleNamespace(index=index, function=types.SimpleNamespace(name=name, arguments=arguments)) + + stream = _FakeStream([ + _chunk(tool_calls=[tc(0, name="lookup_order")]), + _chunk(tool_calls=[tc(0, arguments='{"order_id":')]), + _chunk(tool_calls=[tc(0, arguments=' "A1"}')]), + ]) + + class FakeCompletions: + def create(self, **kwargs): + return stream + + client = _fake_openai_client(FakeCompletions()) + tracer = make_tracer() + patch_openai_client(client, tracer, name="gpt-agent") + list(client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True)) + + _, kwargs = tracer._send.call_args + assert kwargs["output"] == '[tool call] lookup_order({"order_id": "A1"})' + # No usage chunk requested: token counts stay unset, never a fabricated 0. + assert kwargs["input_tokens"] is None + + +def test_openai_stream_records_partial_output_when_caller_stops_early(): + from agentx.integrations.openai import patch_openai_client + + stream = _FakeStream(_stream_chunks()) + + class FakeCompletions: + def create(self, **kwargs): + return stream + + client = _fake_openai_client(FakeCompletions()) + tracer = make_tracer() + patch_openai_client(client, tracer, name="gpt-agent") + + with client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True) as result: + for chunk in result: + if chunk.choices and chunk.choices[0].delta.content == "lo": + break + # `with` exit closes the real stream and finalizes ONCE with what was seen. + assert stream.closed is True + tracer._send.assert_called_once() + _, kwargs = tracer._send.call_args + assert kwargs["output"] == "Hello" + # A later explicit close() must not send a second trace. + result.close() + tracer._send.assert_called_once() + + +def test_openai_stream_records_a_mid_stream_error(): + from agentx.integrations.openai import patch_openai_client + + stream = _FakeStream(_stream_chunks(), fail_after=2) + + class FakeCompletions: + def create(self, **kwargs): + return stream + + client = _fake_openai_client(FakeCompletions()) + tracer = make_tracer() + patch_openai_client(client, tracer, name="gpt-agent") + + with pytest.raises(RuntimeError, match="connection reset"): + list(client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True)) + + tracer._send.assert_called_once() + _, kwargs = tracer._send.call_args + assert kwargs["error"] == "connection reset mid-stream" + assert kwargs["output"] == "Hello" + + +def test_openai_async_stream_is_traced(): + from agentx.integrations.openai import patch_openai_client + + class FakeAsyncCompletions: + async def create(self, **kwargs): + await asyncio.sleep(0.005) + return _FakeAsyncStream(_stream_chunks()) + + client = _fake_openai_client(FakeAsyncCompletions()) + tracer = make_tracer() + patch_openai_client(client, tracer, name="gpt-agent") + + async def consume(): + stream = await client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True) + parts = [] + async for chunk in stream: + if chunk.choices and chunk.choices[0].delta.content: + parts.append(chunk.choices[0].delta.content) + return parts + + assert asyncio.run(consume()) == ["Hel", "lo", " there"] + tracer._send.assert_called_once() + _, kwargs = tracer._send.call_args + assert kwargs["output"] == "Hello there" + assert kwargs["output_tokens"] == 4 + assert kwargs["latency_ms"] >= 5 + + +def test_openai_stream_inside_active_span_becomes_child_llm_call(): + from agentx.integrations.openai import patch_openai_client + + stream = _FakeStream(_stream_chunks()) + + class FakeCompletions: + def create(self, **kwargs): + return stream + + client = _fake_openai_client(FakeCompletions()) + tracer = make_tracer() + tracer._dispatch = MagicMock(return_value=None) # child rows bypass _send and go here + patch_openai_client(client, tracer, name="gpt-agent") + + with tracer.trace("agent-loop") as span: + list(client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True)) + # Folded into the enclosing span: no independent trace was sent mid-span. + assert tracer._send.call_count == 0 + assert span._model == "gpt-4o-mini" or span._captured_model == "gpt-4o-mini" + tracer._send.assert_called_once() + # The child LLM row carries the streaming markers; the parent's metadata does not. + child_rows = [c.args[0] for c in tracer._dispatch.call_args_list if c.args and c.args[0].get("parent_span_id")] + assert len(child_rows) == 1 + assert child_rows[0]["metadata"]["streaming"] is True + assert isinstance(child_rows[0]["metadata"]["timeToFirstTokenMs"], int) + _, parent_kwargs = tracer._send.call_args + assert not (parent_kwargs.get("metadata") or {}).get("streaming") + + +def test_nim_stream_is_traced_with_nim_framework(): + from agentx.integrations.nvidia_nim import patch_nim_client + + stream = _FakeStream(_stream_chunks()) + + class FakeCompletions: + def create(self, **kwargs): + return stream + + client = _fake_openai_client(FakeCompletions()) + tracer = make_tracer() + patch_nim_client(client, tracer, name="nim-agent") + list(client.chat.completions.create(model="meta/llama-3.1-8b-instruct", messages=[], stream=True)) + + _, kwargs = tracer._send.call_args + assert kwargs["framework"] == "nvidia-nim" + assert kwargs["output"] == "Hello there" def test_openai_sync_client_records_errors(): @@ -591,28 +967,6 @@ async def create(self, **kwargs): assert kwargs["input_tokens"] == 12 -def test_nim_streaming_calls_are_passed_through_untraced(): - from agentx.integrations.nvidia_nim import patch_nim_client - - sentinel_stream = object() - - class FakeCompletions: - def create(self, **kwargs): - assert kwargs.get("stream") is True - return sentinel_stream - - client = _fake_openai_client(FakeCompletions()) - tracer = make_tracer() - patch_nim_client(client, tracer, name="nim-agent") - - result = client.chat.completions.create( - model="meta/llama-3.1-8b-instruct", messages=[{"role": "user", "content": "hi"}], stream=True - ) - - assert result is sentinel_stream - tracer._send.assert_not_called() - - def test_nim_sync_client_records_errors(): from agentx.integrations.nvidia_nim import patch_nim_client