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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentx-python
-Version: 0.8.26
+Version: 0.8.27
 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

agentx_python-0.8.27/agentx/integrations/_traced_call.py

@@ -0,0 +1,395 @@
+"""
+Shared helper for tracing a raw SDK client call that may be sync or async.
+
+Anthropic's and Google GenAI's raw clients expose the same method names for
+both their sync and async client variants (``client.messages.create``,
+``client.models.generate_content``) - the only way to tell them apart is to
+call the method and check whether the result is awaitable.
+``inspect.iscoroutinefunction`` is unreliable for this: it returns ``False``
+even for ``AsyncAnthropic().messages.create``, since these SDKs don't
+implement the async variant as a plain top-level ``async def``.
+"""
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+import time
+from typing import Any, Callable, Dict, Optional
+
+from agentx.tracing.tracer import Tracer, _safe_serialize
+
+
+def call_and_trace(
+    original: Callable[..., Any],
+    args: tuple,
+    kwargs: dict,
+    on_finish: Callable[[Optional[Any], Optional[str]], None],
+) -> Any:
+    """
+    Call ``original(*args, **kwargs)``.
+
+    If the result is awaitable (async client), return a coroutine that awaits
+    it and calls ``on_finish(response, error)`` only after the real await
+    completes, so timing/output/tokens reflect the actual call rather than
+    the moment the coroutine object was constructed. If the result is a
+    normal value (sync client), call ``on_finish`` immediately.
+
+    Either way, the original call's own return value / exception behavior is
+    unchanged for the caller - this only affects when/how the trace is built.
+    """
+    try:
+        result = original(*args, **kwargs)
+    except Exception as exc:
+        on_finish(None, str(exc))
+        raise
+
+    if asyncio.iscoroutine(result) or inspect.isawaitable(result):
+        return _await_and_finish(result, on_finish)
+
+    on_finish(result, None)
+    return result
+
+
+async def _await_and_finish(
+    awaitable: Any,
+    on_finish: Callable[[Optional[Any], Optional[str]], None],
+) -> Any:
+    try:
+        response = await awaitable
+    except Exception as exc:
+        on_finish(None, str(exc))
+        raise
+    on_finish(response, None)
+    return response
+
+
+# The engine's unregistered-tool surfacing (Tools & MCPs -> Unregistered) reads a trace's
+# metadata "tools" key to show the REAL definition instead of one inferred from observed
+# arguments (see AgentX-trace-eval's toolSchemas.ts draftFromMetadata). Raw-client patches see
+# the request's tools=[...] right in kwargs, so capture it - capped so a huge toolbox never
+# blows up the trace's metadata budget.
+_MAX_TOOL_DEFINITIONS = 20
+_MAX_TOOL_DEFINITIONS_BYTES = 12_000
+
+
+def capture_tool_definitions(tools: Any) -> Optional[list]:
+    """Return a metadata-ready copy of a request's ``tools=[...]`` list, or None."""
+    if not isinstance(tools, list) or not tools:
+        return None
+    # A plain JSON round-trip preserves nested schema objects exactly (default=str catches the
+    # odd SDK object inside); _safe_serialize would repr-stringify nested dicts, turning a
+    # parameters schema into an unusable string.
+    try:
+        text = json.dumps(tools[:_MAX_TOOL_DEFINITIONS], default=str)
+        if len(text) > _MAX_TOOL_DEFINITIONS_BYTES:
+            return None
+        serialized = json.loads(text)
+    except (TypeError, ValueError):
+        return None
+    return serialized if isinstance(serialized, list) else None
+
+
+def finish_llm_call(
+    tracer: Tracer,
+    *,
+    name: str,
+    framework: str,
+    metadata: Optional[Dict[str, Any]],
+    session_id: Optional[str],
+    start_t: float,
+    end_t: float,
+    input_repr: Any,
+    output: Optional[str],
+    model: Optional[str],
+    input_tokens: Optional[int],
+    output_tokens: Optional[int],
+    error: Optional[str],
+    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
+    callbacks of every integration that patches a raw provider client
+    (``anthropic.py``, ``google_genai.py``, ``openai.py``, ``litellm.py``) rather than a
+    framework-level callback/plugin system.
+
+    If the call happened inside a ``with tracer.trace(...)`` block, it becomes that span's own
+    real child span (via _record_llm_call) instead of an independent trace - the same "part of a
+    multi-call agentic loop" behavior ``anthropic.py`` already had; folded in here so every
+    raw-client integration gets it instead of each having to remember to check
+    ``tracer.current_span`` itself. Otherwise it becomes its own real root span, opened/closed
+    directly here (not via ``tracer._send()``) so it still gets a real span_id/session_id and the
+    call's exact timing rather than wall-clock "now".
+    """
+    latency_ms = int((end_t - start_t) * 1000)
+
+    if tool_definitions:
+        metadata = {**(metadata or {}), "tools": tool_definitions}
+
+    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
+        # unregistered-tool listing, same as the standalone-trace path below.
+        if tool_definitions and not (active_span._metadata or {}).get("tools"):
+            active_span._metadata = {**(active_span._metadata or {}), "tools": tool_definitions}
+        if error is not None:
+            active_span.set_error(error)
+        active_span._record_llm_call(
+            duration_ms=latency_ms,
+            start_time=start_t,
+            end_time=end_t,
+            input=input_repr,
+            output=output,
+            model=model,
+            # Stamp the provider literal on a span opened without one (adoption keeps an
+            # explicit framework= or a framework integration's label winning over this).
+            framework=framework,
+            input_tokens=input_tokens,
+            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 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
+    span.input = input_repr
+    span.output = output
+    if error:
+        span.set_error(error)
+    if input_tokens:
+        span._input_tokens = input_tokens
+    if output_tokens:
+        span._output_tokens = output_tokens
+    if cache_read_tokens:
+        span._cache_read_tokens = cache_read_tokens
+    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)

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

@@ -13,6 +13,12 @@ Usage::
 
 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 @@ import time
 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 _patch_create(
 
         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 _patch_stream(
             )
 
         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