agentx-python 0.8.27__tar.gz → 0.8.28__tar.gz

{agentx_python-0.8.27 → agentx_python-0.8.28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentx-python
-Version: 0.8.27
+Version: 0.8.28
 Summary: Official Python SDK for AgentX (https://www.agentx.so/)
 Home-page: https://github.com/AgentX-ai/AgentX-python
 Author: Robin Wang and AgentX Team
@@ -260,6 +260,12 @@ extra:
 | LlamaIndex            | `pip install "agentx-python[llamaindex]"`    | `AgentXLlamaIndexHandler` |
 | AutoGen               | `pip install "agentx-python[autogen]"`       | `AgentXAutoGenObserver`   |
 
+Raw-client patches (`patch_openai_client`, `patch_nim_client`, `patch_anthropic_client`) trace
+streaming calls too: the returned stream is a transparent proxy that assembles the reply from the
+chunks you consume, with latency measured to the last chunk and the time to first token in the
+trace metadata. OpenAI-compatible endpoints only send token usage on streams when you pass
+`stream_options={"include_usage": True}`.
+
 > **Warning: pick ONE instrumentation layer per LLM call.** Do not combine
 > `AgentXCallbackHandler` (or any framework integration) with a patched provider client
 > (`patch_openai_client`, `patch_anthropic_client`, `patch_genai_client`) on the same code

{agentx_python-0.8.27 → agentx_python-0.8.28}/README.md

@@ -193,6 +193,12 @@ extra:
 | LlamaIndex            | `pip install "agentx-python[llamaindex]"`    | `AgentXLlamaIndexHandler` |
 | AutoGen               | `pip install "agentx-python[autogen]"`       | `AgentXAutoGenObserver`   |
 
+Raw-client patches (`patch_openai_client`, `patch_nim_client`, `patch_anthropic_client`) trace
+streaming calls too: the returned stream is a transparent proxy that assembles the reply from the
+chunks you consume, with latency measured to the last chunk and the time to first token in the
+trace metadata. OpenAI-compatible endpoints only send token usage on streams when you pass
+`stream_options={"include_usage": True}`.
+
 > **Warning: pick ONE instrumentation layer per LLM call.** Do not combine
 > `AgentXCallbackHandler` (or any framework integration) with a patched provider client
 > (`patch_openai_client`, `patch_anthropic_client`, `patch_genai_client`) on the same code

{agentx_python-0.8.27 → agentx_python-0.8.28}/agentx/integrations/_traced_call.py

@@ -14,9 +14,16 @@ from __future__ import annotations
 import asyncio
 import inspect
 import json
+import logging
+import threading
 import time
 from typing import Any, Callable, Dict, Optional
 
+logger = logging.getLogger(__name__)
+
+# Sentinel for finish_llm_call's `active_span`: "not passed" is distinct from "passed None".
+_UNSET: Any = object()
+
 from agentx.tracing.tracer import Tracer, _safe_serialize
 
 
@@ -109,6 +116,7 @@ def finish_llm_call(
     cache_write_tokens: Optional[int] = None,
     tool_definitions: Optional[list] = None,
     call_metadata: Optional[Dict[str, Any]] = None,
+    active_span: Any = _UNSET,
 ) -> None:
     """
     Close out one raw-client LLM call - shared by the ``on_finish``/exit
@@ -129,7 +137,11 @@ def finish_llm_call(
     if tool_definitions:
         metadata = {**(metadata or {}), "tools": tool_definitions}
 
-    active_span = tracer.current_span
+    # The parent is the span that was active when the CALL was made. Streaming patches pass it
+    # explicitly: a stream finalizes later (exhaustion, close, or garbage collection), by which
+    # time a different span may be active, and the call must not be grafted onto it.
+    if active_span is _UNSET:
+        active_span = tracer.current_span
     if active_span is not None:
         # The definitions describe the whole call's toolbox - attach them to the enclosing
         # span's metadata (first capture wins) so the ROOT trace carries them for the
@@ -168,6 +180,9 @@ def finish_llm_call(
     )
     span.__enter__()
     span._start = start_t
+    # The call ended at end_t (a stream's last chunk), not at whatever later moment this
+    # runs - __exit__ honors the override instead of measuring to time.time().
+    span._end_override = end_t
     span.input = input_repr
     span.output = output
     if error:
@@ -217,7 +232,13 @@ class TracedStream:
 
     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.
+    ``break`` leaves no half-driven generator behind. It is a proxy, not a
+    subclass: ``isinstance(stream, openai.Stream)`` is False and ``repr()``
+    shows the proxy - branch on ``stream=True`` in your own code, not on type.
+
+    Tracing never breaks the caller: a failure while building or sending the
+    trace is logged and swallowed, and the stream's own iteration/close
+    semantics are untouched.
     """
 
     def __init__(
@@ -230,6 +251,8 @@ class TracedStream:
         self._accumulator = accumulator
         self._on_finish = on_finish
         self._done = False
+        # A watchdog close() racing the reader's StopIteration must not finalize twice.
+        self._done_lock = threading.Lock()
         self._first_chunk_t: Optional[float] = None
         self._last_chunk_t: Optional[float] = None
         self._sync_iter: Any = None
@@ -250,9 +273,10 @@ class TracedStream:
             pass
 
     def _finish(self, error: Optional[str]) -> None:
-        if self._done:
-            return
-        self._done = True
+        with self._done_lock:
+            if self._done:
+                return
+            self._done = True
         try:
             result = self._accumulator.result()
         except Exception:
@@ -264,7 +288,12 @@ class TracedStream:
         # 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)
+        try:
+            self._on_finish(result, error)
+        except Exception:
+            # Building or sending the trace failed. The caller's stream ended normally and must
+            # see it end normally - tracing is never allowed to raise into inference code.
+            logger.debug("Streamed call could not be traced", exc_info=True)
 
     @property
     def first_chunk_at(self) -> Optional[float]:
@@ -283,9 +312,14 @@ class TracedStream:
         except StopIteration:
             self._finish(None)
             raise
-        except BaseException as exc:
+        except Exception as exc:
             self._finish(str(exc))
             raise
+        except BaseException:
+            # KeyboardInterrupt / GeneratorExit: a cancellation, not the provider failing -
+            # record what streamed so far without inventing an error message.
+            self._finish(None)
+            raise
         self._observe(chunk)
         return chunk
 
@@ -302,9 +336,12 @@ class TracedStream:
         except StopAsyncIteration:
             self._finish(None)
             raise
-        except BaseException as exc:
+        except Exception as exc:
             self._finish(str(exc))
             raise
+        except BaseException:
+            self._finish(None)
+            raise
         self._observe(chunk)
         return chunk
 
@@ -338,7 +375,17 @@ class TracedStream:
         close = getattr(self._stream, "close", None)
         try:
             if close is not None:
-                close()
+                result = close()
+                if inspect.isawaitable(result):
+                    # openai's AsyncStream spells its close `async def close()`. A sync close()
+                    # on it (an easy slip inside async code) would drop the coroutine and leak
+                    # the connection; run it on the loop when there is one, else at least don't
+                    # leave an un-awaited coroutine behind.
+                    try:
+                        asyncio.get_running_loop().create_task(result)
+                    except RuntimeError:
+                        result.close()  # type: ignore[union-attr]
+                        logger.warning("close() called on an async stream outside an event loop - use aclose()")
         finally:
             self._finish(None)
 
@@ -355,6 +402,10 @@ class TracedStream:
             self._finish(None)
 
     def __getattr__(self, item: str) -> Any:
+        # Only public attributes delegate. Private names must resolve on the proxy itself, or a
+        # half-constructed instance (no _stream yet) would recurse forever looking for it.
+        if item.startswith("_"):
+            raise AttributeError(item)
         return getattr(self._stream, item)
 
     def __del__(self) -> None:

{agentx_python-0.8.27 → agentx_python-0.8.28}/agentx/integrations/anthropic.py

@@ -232,13 +232,21 @@ def _patch_create(
         input_repr = _safe_serialize(input_messages)
 
         if kwargs.get("stream"):
+            # Parent fixed at call time - see openai.py's patched_create for why.
+            parent = tracer.current_span
+
             def on_stream_finish(collected: Dict[str, Any], error: Optional[str]) -> None:
+                ttft = collected.get("time_to_first_token_ms")
+                call_metadata: Dict[str, Any] = {"streaming": True}
+                if ttft is not None:
+                    call_metadata["timeToFirstTokenMs"] = ttft
                 finish_llm_call(
                     tracer,
                     name=name,
                     framework="anthropic",
                     metadata=metadata,
-                    call_metadata={"streaming": True, "timeToFirstTokenMs": collected.get("time_to_first_token_ms")},
+                    call_metadata=call_metadata,
+                    active_span=parent,
                     session_id=session_id,
                     start_t=start_t,
                     end_t=collected.get("end_t") or time.time(),
@@ -318,8 +326,9 @@ def _patch_stream(
         # only shows up in whether `with`/`async with` and
         # `get_final_message()` are used, handled inside `_TracedStream`.
         start_t = time.time()
+        parent = tracer.current_span
         ctx = original_stream(*args, **kwargs)
-        input_repr = _safe_serialize(_prepend_system(kwargs.get("messages"), kwargs.get("system")))
+        input_repr = _safe_serialize(_prepend_system(kwargs.get("messages") or (args[0] if args else None), kwargs.get("system")))
         model = kwargs.get("model")
         tool_definitions = capture_tool_definitions(kwargs.get("tools"))
 
@@ -365,6 +374,29 @@ def _patch_stream(
             """
 
             _inner: Any = None
+            _sent: bool = False
+
+            # What streamed so far, WITHOUT draining the rest of the response: the SDK's
+            # get_final_message() calls until_done(), which would turn an early `break` into a
+            # blocking read of every remaining token. The snapshot is the final message once the
+            # stream was consumed, and honestly partial when the caller stopped early.
+            def _snapshot(self_inner):
+                inner = self_inner._inner
+                if inner is None:
+                    return None
+                try:
+                    return getattr(inner, "current_message_snapshot", None)
+                except Exception:
+                    return None
+
+            def _send_once(self_inner, end_t: float, error: Optional[str], snapshot: Any) -> None:
+                if self_inner._sent:
+                    return
+                self_inner._sent = True
+                try:
+                    build_and_send(end_t, error, snapshot)
+                except Exception:
+                    pass  # tracing never raises into the caller
 
             def __enter__(self_inner):
                 self_inner._inner = ctx.__enter__()
@@ -373,15 +405,12 @@ def _patch_stream(
             def __exit__(self_inner, exc_type, exc_val, tb):
                 end_t = time.time()
                 error = str(exc_val) if exc_val else None
-                final_message = None
-                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
+                # Snapshot BEFORE the manager closes the stream (the earlier bug read it after).
+                snapshot = self_inner._snapshot()
+                try:
+                    return ctx.__exit__(exc_type, exc_val, tb)
+                finally:
+                    self_inner._send_once(end_t, error, snapshot)
 
             async def __aenter__(self_inner):
                 self_inner._inner = await ctx.__aenter__()
@@ -390,16 +419,19 @@ def _patch_stream(
             async def __aexit__(self_inner, exc_type, exc_val, tb):
                 end_t = time.time()
                 error = str(exc_val) if exc_val else None
-                final_message = None
-                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
+                snapshot = self_inner._snapshot()
+                try:
+                    return await ctx.__aexit__(exc_type, exc_val, tb)
+                finally:
+                    self_inner._send_once(end_t, error, snapshot)
+
+            def __del__(self_inner):
+                # A helper stream that was entered but never exited still records what it saw.
+                try:
+                    if self_inner._inner is not None:
+                        self_inner._send_once(time.time(), None, self_inner._snapshot())
+                except Exception:
+                    pass
 
             def __iter__(self_inner):
                 return iter(ctx)

{agentx_python-0.8.27 → agentx_python-0.8.28}/agentx/integrations/nvidia_nim.py

@@ -27,8 +27,11 @@ Works with both ``openai.OpenAI`` and ``openai.AsyncOpenAI`` clients. Token
 usage comes straight off the response's OpenAI-shaped ``usage`` block; NIM
 reports no prompt-cache fields, so cache token counts stay unset.
 
-Streaming calls (``stream=True``) are passed through untouched and are not
-currently traced - same posture as ``patch_openai_client``, see its docstring.
+Streaming calls (``stream=True``) are traced too, exactly as
+``patch_openai_client`` traces them: the stream is wrapped in a transparent
+proxy that assembles the reply from the consumed chunks (token usage when the
+endpoint sends it on the final chunk, e.g. with
+``stream_options={"include_usage": True}``).
 
 Requires: ``pip install "agentx-python[nvidia-nim]"`` (installs the ``openai``
 client package; there is no separate NIM SDK dependency).
@@ -56,8 +59,9 @@ def patch_nim_client(
     call with ``framework="nvidia-nim"``.
 
     The original method is still called and its return value passed through
-    unchanged. Sync and async clients both work; ``stream=True`` calls pass
-    through untraced. Patching is idempotent - and because it shares the guard
+    unchanged. Sync and async clients both work; ``stream=True`` calls are
+    traced through the same stream proxy as ``patch_openai_client``. Patching
+    is idempotent - and because it shares the guard
     with ``patch_openai_client``, whichever of the two patched a given client
     first wins (patch each client with the integration that matches where its
     ``base_url`` actually points).

{agentx_python-0.8.27 → agentx_python-0.8.28}/agentx/integrations/openai.py

@@ -30,7 +30,14 @@ from __future__ import annotations
 import time
 from typing import Any, Dict, Optional, Tuple
 
+import logging
+
 from agentx.tracing.tracer import Tracer, _safe_serialize
+
+logger = logging.getLogger(__name__)
+# Warn once per process, not per call: a streamed OpenAI call carries no usage unless the caller
+# asked for it, and a silent zero would under-report every streaming app's spend.
+_warned_stream_usage = False
 from agentx.integrations._traced_call import (
     StreamAccumulator,
     capture_tool_definitions,
@@ -203,14 +210,30 @@ def _patch_chat_completions_create(
         tool_definitions = capture_tool_definitions(kwargs.get("tools"))
 
         if kwargs.get("stream"):
+            # The parent is fixed at call time: the stream finalizes later, possibly inside an
+            # unrelated span (or none), and must not attach to whatever is active then.
+            parent = tracer.current_span
+
             def on_stream_finish(collected: Dict[str, Any], error: Optional[str]) -> None:
+                global _warned_stream_usage
                 ttft = collected.get("time_to_first_token_ms")
+                call_metadata: Dict[str, Any] = {"streaming": True}
+                if ttft is not None:
+                    call_metadata["timeToFirstTokenMs"] = ttft
+                if error is None and collected.get("input_tokens") is None and not _warned_stream_usage:
+                    _warned_stream_usage = True
+                    logger.warning(
+                        "Streamed %s call carried no token usage - pass stream_options={\"include_usage\": True} "
+                        "so traces (and cost) reflect streamed traffic.",
+                        framework,
+                    )
                 finish_llm_call(
                     tracer,
                     name=name,
                     framework=framework,
                     metadata=metadata,
-                    call_metadata={"streaming": True, "timeToFirstTokenMs": ttft},
+                    call_metadata=call_metadata,
+                    active_span=parent,
                     session_id=session_id,
                     start_t=start_t,
                     end_t=collected.get("end_t") or time.time(),

{agentx_python-0.8.27 → agentx_python-0.8.28}/agentx/monitor/alert_rules.py

@@ -17,6 +17,22 @@ _ALIASES = {
 }
 
 
+def _validate(
+    metric: Optional[str] = None, operator: Optional[str] = None, channels: Optional[List[Dict[str, str]]] = None
+) -> None:
+    """Local checks for the fields the engine would otherwise 400 on - ``None`` means "not given"
+    (an update that leaves the field alone)."""
+    if metric is not None and metric not in ALERT_METRICS:
+        raise ValueError(f"metric must be one of {ALERT_METRICS}, got {metric!r}")
+    if operator is not None and operator not in ("gt", "lt"):
+        raise ValueError(f"operator must be 'gt' or 'lt', got {operator!r}")
+    if channels is not None:
+        for channel in channels:
+            kind = channel.get("kind") if isinstance(channel, dict) else None
+            if kind not in ALERT_CHANNEL_KINDS:
+                raise ValueError(f"channel kind must be one of {ALERT_CHANNEL_KINDS}, got {kind!r}")
+
+
 class AlertRule(dict):
     """Wire object for one KPI alert rule (dict subclass so unknown fields round-trip)."""
 
@@ -48,12 +64,12 @@ class AlertEvent(dict):
 
     @property
     def kind(self) -> str:
-        return str(self.get("kind"))
+        return str(self.get("kind") or "")
 
     @property
     def delivered(self) -> bool:
         deliveries = self.get("deliveries") or []
-        return bool(deliveries) and all(bool(d.get("ok")) for d in deliveries)
+        return bool(deliveries) and all(isinstance(d, dict) and bool(d.get("ok")) for d in deliveries)
 
 
 def slack(url: str) -> Dict[str, str]:
@@ -138,13 +154,7 @@ class AlertRulesClient:
         """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}")
+        _validate(metric=metric, operator=operator, channels=channels)
         payload: Dict[str, Any] = {
             "name": name,
             "metric": metric,
@@ -165,7 +175,10 @@ class AlertRulesClient:
 
     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."""
+        operator, threshold, window, or agent resets the rule's firing state, and a rule that
+        was firing sends its channels a final ``resolved`` notification first. The same local
+        checks as ``create`` apply to whichever of ``metric``, ``operator``, ``channels`` are
+        given."""
         payload: Dict[str, Any] = {}
         for key, value in fields.items():
             wire_key = _ALIASES.get(key, key)
@@ -175,12 +188,13 @@ class AlertRulesClient:
                     "silently ignore this (see AlertRulesClient.create for the field names)."
                 )
             payload[wire_key] = value
+        _validate(metric=payload.get("metric"), operator=payload.get("operator"), channels=payload.get("channels"))
         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."""
+        """Deletes the rule and its history. A rule that is firing sends its channels a final
+        ``resolved`` notification (closing the PagerDuty incident it opened) before it goes."""
         self._request("DELETE", f"/agent-monitoring/alert-rules/{rule_id}", retry=False)
 
     def events(self, rule_id: str, limit: int = 50) -> List[AlertEvent]:

{agentx_python-0.8.27 → agentx_python-0.8.28}/agentx/tracing/tracer.py

@@ -165,6 +165,9 @@ class _TraceSpan:
         self.tool_calls: list = []
 
         self._start: Optional[float] = None
+        # Set by callers that know when the work actually ended (a streamed LLM call's last
+        # chunk) so __exit__ does not measure to "now" - see finish_llm_call's root path.
+        self._end_override: Optional[float] = None
         self._error: Optional[str] = None
 
         self._captured_model: Optional[str] = None
@@ -212,7 +215,8 @@ class _TraceSpan:
 
     def __exit__(self, exc_type, exc_val, tb):
         self._tracer._pop_active_span(self)
-        latency_ms = int((time.time() - self._start) * 1000) if self._start else None
+        ended_at = self._end_override if self._end_override is not None else time.time()
+        latency_ms = int((ended_at - self._start) * 1000) if self._start else None
         if exc_val is not None and self._error is None:
             self._error = str(exc_val)
 

{agentx_python-0.8.27 → agentx_python-0.8.28}/agentx/version.py

@@ -1,7 +1,7 @@
-VERSION = "0.8.27"
+VERSION = "0.8.28"
 
 # The AgentX-trace-eval release this SDK version is tested against - what `agentx-trace-eval`
 # installs and converges to (see agentx/cli.py). Bump together with VERSION when releasing, so
 # every published SDK names a known-good engine+dashboard pair. Users can override with
 # AGENTX_TRACE_EVAL_VERSION=<tag|latest>.
-ENGINE_VERSION = "v0.3.30"
+ENGINE_VERSION = "v0.3.31"

{agentx_python-0.8.27 → agentx_python-0.8.28}/agentx_python.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentx-python
-Version: 0.8.27
+Version: 0.8.28
 Summary: Official Python SDK for AgentX (https://www.agentx.so/)
 Home-page: https://github.com/AgentX-ai/AgentX-python
 Author: Robin Wang and AgentX Team
@@ -260,6 +260,12 @@ extra:
 | LlamaIndex            | `pip install "agentx-python[llamaindex]"`    | `AgentXLlamaIndexHandler` |
 | AutoGen               | `pip install "agentx-python[autogen]"`       | `AgentXAutoGenObserver`   |
 
+Raw-client patches (`patch_openai_client`, `patch_nim_client`, `patch_anthropic_client`) trace
+streaming calls too: the returned stream is a transparent proxy that assembles the reply from the
+chunks you consume, with latency measured to the last chunk and the time to first token in the
+trace metadata. OpenAI-compatible endpoints only send token usage on streams when you pass
+`stream_options={"include_usage": True}`.
+
 > **Warning: pick ONE instrumentation layer per LLM call.** Do not combine
 > `AgentXCallbackHandler` (or any framework integration) with a patched provider client
 > (`patch_openai_client`, `patch_anthropic_client`, `patch_genai_client`) on the same code

{agentx_python-0.8.27 → agentx_python-0.8.28}/tests/test_alert_rules.py

@@ -89,6 +89,18 @@ def test_update_maps_snake_case_and_refuses_unknown_keys():
         AlertRulesClient(fake).update("a1", sample_rate=0.5)  # type: ignore[arg-type]
 
 
+def test_update_applies_the_same_local_checks_as_create():
+    fake = FakeMonitorClient([])
+    client = AlertRulesClient(fake)  # type: ignore[arg-type]
+    with pytest.raises(ValueError, match="metric"):
+        client.update("a1", metric="vibes")
+    with pytest.raises(ValueError, match="operator"):
+        client.update("a1", operator="ge")
+    with pytest.raises(ValueError, match="channel kind"):
+        client.update("a1", channels=[{"kind": "sms", "target": "1"}])
+    assert fake.calls == []
+
+
 def test_events_test_preview_and_sweep_paths():
     fake = FakeMonitorClient(
         [

{agentx_python-0.8.27 → agentx_python-0.8.28}/tests/test_integrations.py

@@ -278,10 +278,17 @@ def test_anthropic_stream_helper_records_the_final_message():
             self.closed = False
             self.text_stream = iter(["streamed ", "reply"])
 
-        def get_final_message(self):
+        @property
+        def current_message_snapshot(self):
             assert not self.closed, "must be read before the manager closes the stream"
             return final
 
+        def get_final_message(self):
+            # The SDK's get_final_message() drains the rest of the response (until_done()):
+            # an early `break` would block until the model finished. The wrapper must never
+            # call it.
+            raise AssertionError("get_final_message() drains the stream - use the snapshot")
+
     class FakeManager:
         def __init__(self):
             self.stream = FakeMessageStream()
@@ -314,6 +321,122 @@ def test_anthropic_stream_helper_records_the_final_message():
     assert kwargs["output_tokens"] == 4
 
 
+def test_stream_root_trace_latency_ends_at_the_last_chunk_not_at_finalization():
+    # The proxy hands finish_llm_call the last-chunk time; the root span must honor it instead
+    # of measuring to whatever later moment the caller dropped the stream.
+    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")
+    result = client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True)
+    for _ in result:
+        pass
+    time.sleep(0.15)  # the caller holds the exhausted stream a while before closing it
+    result.close()
+    _, kwargs = tracer._send.call_args
+    assert kwargs["latency_ms"] < 100
+
+
+def test_stream_tracing_failure_never_escapes_into_the_callers_loop():
+    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()
+    tracer._send = MagicMock(side_effect=RuntimeError("ingest exploded"))
+    patch_openai_client(client, tracer, name="gpt-agent")
+    chunks = list(client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True))
+    assert len(chunks) == 4  # the loop ended normally despite the tracer raising
+
+
+def test_stream_parent_is_the_span_active_at_call_time_not_at_finalization():
+    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()
+    tracer._dispatch = MagicMock(return_value=None)
+    patch_openai_client(client, tracer, name="gpt-agent")
+
+    # Created with NO active span, consumed, then finalized while an unrelated span is active.
+    stream = client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True)
+    for _ in stream:
+        pass
+    with tracer.trace("unrelated-task") as unrelated:
+        stream.close()
+        assert unrelated._child_span_count == 0, "the stream must not graft onto the unrelated span"
+    # It became its own root trace instead (two _send calls: the stream's root + unrelated-task).
+    assert tracer._send.call_count == 2
+    names = [c.kwargs["name"] for c in tracer._send.call_args_list]
+    assert "gpt-agent" in names and "unrelated-task" in names
+
+
+def test_stream_proxy_private_attributes_never_delegate():
+    from agentx.integrations._traced_call import TracedStream
+
+    proxy = TracedStream.__new__(TracedStream)  # half-constructed: no _stream yet
+    with pytest.raises(AttributeError):
+        _ = proxy._done
+    del proxy  # __del__ on the half-built object must not recurse or raise
+
+
+def test_stream_keyboard_interrupt_records_partial_output_without_an_error():
+    from agentx.integrations.openai import patch_openai_client
+
+    class InterruptingStream(_FakeStream):
+        def __iter__(self):
+            yield _chunk(content="Hel")
+            raise KeyboardInterrupt()
+
+    class FakeCompletions:
+        def create(self, **kwargs):
+            return InterruptingStream([])
+
+    client = _fake_openai_client(FakeCompletions())
+    tracer = make_tracer()
+    patch_openai_client(client, tracer, name="gpt-agent")
+    with pytest.raises(KeyboardInterrupt):
+        list(client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True))
+    _, kwargs = tracer._send.call_args
+    assert kwargs["output"] == "Hel"
+    assert kwargs["error"] is None
+
+
+def test_stream_without_a_first_chunk_omits_time_to_first_token():
+    from agentx.integrations.openai import patch_openai_client
+
+    class DeadStream(_FakeStream):
+        def __iter__(self):
+            raise RuntimeError("connection refused before the first chunk")
+            yield  # noqa: unreachable - makes this a generator like the real Stream
+
+    class FakeCompletions:
+        def create(self, **kwargs):
+            return DeadStream([])
+
+    client = _fake_openai_client(FakeCompletions())
+    tracer = make_tracer()
+    patch_openai_client(client, tracer, name="gpt-agent")
+    with pytest.raises(RuntimeError):
+        list(client.chat.completions.create(model="gpt-4o-mini", messages=[], stream=True))
+    _, kwargs = tracer._send.call_args
+    assert kwargs["metadata"]["streaming"] is True
+    assert "timeToFirstTokenMs" not in kwargs["metadata"]
+
+
 def test_abandoned_stream_still_records_what_it_saw_when_collected():
     import gc
     from agentx.integrations.openai import patch_openai_client

{agentx_python-0.8.27 → agentx_python-0.8.28}/agentx/monitor/__init__.py

@@ -19,14 +19,14 @@ from agentx.monitor.sessions import MonitorSessionClient
 from agentx.monitor.signals import MonitorSignalClient
 
 __all__ = [
-    "AlertEvent",
-    "AlertRule",
-    "AlertRulesClient",
     "AgentXImprovementGroupsError",
     "AgentXJudgeScorersError",
     "AgentXMonitorError",
     "AgentXScorerGroupsError",
     "AgentXScorersError",
+    "AlertEvent",
+    "AlertRule",
+    "AlertRulesClient",
     "ImprovementGroupsClient",
     "JudgeScorer",
     "JudgeScorerBuilder",