zexl 0.1.0__py3-none-any.whl

zexl/__init__.py

@@ -0,0 +1,103 @@
+"""Zexl SDK: records metadata about model calls. Prompt and response text never leave the process.
+
+import zexl.auto                      # reads ZEXL_API_KEY and starts capturing
+
+zexl.set_customer("acct_123")         # optional
+with zexl.run("research"):            # optional: group calls into one run explicitly
+    ...
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import Any
+
+from . import _config
+from ._context import feature, run, set_customer
+from ._export import Transport
+
+__version__ = "0.1.0"
+__all__ = ["__version__", "feature", "flush", "init", "run", "set_customer", "shutdown", "status"]
+
+log = logging.getLogger("zexl")
+_lock = threading.Lock()
+_runtime: Any = None
+_patches: Any = None
+
+
+def init(
+    api_key: str | None = None,
+    *,
+    endpoint: str | None = None,
+    environment: str | None = None,
+    _transport: Transport | None = None,
+) -> bool:
+    """Starts capture. Safe to call more than once. Returns whether capture is on.
+
+    Never raises: with a missing or invalid key the SDK logs one warning and stays off.
+    """
+    global _runtime, _patches
+    with _lock:
+        if _runtime is not None:
+            return True
+        if _config.disabled_by_env():
+            return False
+        try:
+            config = _config.load(api_key, endpoint, environment)
+        except _config.ConfigError as e:
+            log.warning("zexl: disabled (%s)", e)
+            return False
+        try:
+            from . import _hooks
+            from ._core import Runtime
+            from ._export import Exporter
+
+            runtime = Runtime(config, Exporter(config, __version__, _transport))
+            _patches = _hooks.install(runtime)
+            _runtime = runtime
+        except Exception:  # noqa: BLE001 - starting Zexl must never break the application
+            log.warning("zexl: disabled (could not start)", exc_info=True)
+            return False
+        return True
+
+
+def flush(timeout: float = 5.0) -> bool:
+    """Sends queued records now. Call before a short-lived process (serverless) exits."""
+    rt = _runtime
+    if rt is None:
+        return True
+    try:
+        return bool(rt.exporter.flush(timeout))
+    except Exception:  # noqa: BLE001
+        return False
+
+
+def shutdown() -> None:
+    """Flushes, removes the hooks and stops the SDK."""
+    global _runtime, _patches
+    with _lock:
+        rt, patches = _runtime, _patches
+        _runtime = _patches = None
+    try:
+        if patches is not None:
+            patches.undo()
+        if rt is not None:
+            rt.exporter.shutdown()
+    except Exception:  # noqa: BLE001, S110
+        pass
+
+
+def status() -> dict[str, Any]:
+    """What the SDK is doing: whether it is on, which hooks are active, and delivery counts."""
+    rt = _runtime
+    if rt is None:
+        return {"enabled": False}
+    return {
+        "enabled": True,
+        "endpoint": rt.config.endpoint,
+        "environment": rt.config.environment,
+        "hooks": dict(rt.hooks),
+        "sent": rt.exporter.sent,
+        "dropped": rt.exporter.dropped,
+    }

zexl/_canon.py

@@ -0,0 +1,193 @@
+"""Canonical form of chat content.
+
+The same content must produce the same canonical value whether it arrives as a provider SDK
+object, a dumped dict, or a hand-built dict, and whichever provider shape it uses. Rules (each
+one was needed by a case in the M1 hash-stability spike):
+
+- A message is a role plus an ordered list of parts: text, tool call, tool result, other.
+- Ids are left out, so id formats cannot break a match.
+- Tool arguments are parsed and re-serialised with sorted keys.
+- A string and a single text part are the same content; adjacent text parts merge.
+- Reasoning blocks are ignored, because callers may or may not send them back.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+Part = list[str]
+Message = dict[str, Any]
+
+_TEXT_TYPES = ("text", "input_text", "output_text")
+
+
+def dumps(value: Any) -> str:
+    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False, default=str)
+
+
+def plain(obj: Any) -> Any:
+    """Pydantic models, typed dicts and plain dicts all become plain JSON-like data."""
+    if isinstance(obj, (str, int, float, bool)) or obj is None:
+        return obj
+    if isinstance(obj, dict):
+        return {k: plain(v) for k, v in obj.items() if v is not None}
+    if isinstance(obj, (list, tuple)):
+        return [plain(v) for v in obj]
+    dump = getattr(obj, "model_dump", None)
+    if dump is not None:
+        try:
+            return dump(mode="json", exclude_none=True, warnings=False)
+        except Exception:  # noqa: BLE001, S110 - fall through to a best-effort dict
+            pass
+    to_dict = getattr(obj, "__dict__", None)
+    if isinstance(to_dict, dict):
+        return {k: plain(v) for k, v in to_dict.items() if v is not None and not k.startswith("_")}
+    return str(obj)
+
+
+def canon_args(value: Any) -> str:
+    """Tool arguments as canonical JSON, whether given as a JSON string or as data."""
+    if isinstance(value, str):
+        try:
+            value = json.loads(value)
+        except ValueError:
+            return str(value)
+    return dumps(value)
+
+
+def text_of(content: Any) -> str:
+    if content is None:
+        return ""
+    if isinstance(content, str):
+        return content
+    if not isinstance(content, list):
+        return dumps(content)
+    out: list[str] = []
+    for part in content:
+        if isinstance(part, str):
+            out.append(part)
+        elif isinstance(part, dict) and part.get("type") in _TEXT_TYPES:
+            out.append(str(part.get("text", "")))
+        else:
+            out.append(dumps(part))
+    return "".join(out)
+
+
+def parts_of(content: Any) -> list[Part]:
+    """Message content (already plain) as ordered parts."""
+    if content is None:
+        return []
+    if isinstance(content, str):
+        return [["text", content]] if content else []
+    if not isinstance(content, list):
+        return [["other", dumps(content)]]
+    parts: list[Part] = []
+    for p in content:
+        if isinstance(p, str):
+            if p:
+                _add_text(parts, p)
+            continue
+        if not isinstance(p, dict):
+            parts.append(["other", dumps(p)])
+            continue
+        t = p.get("type")
+        if t in _TEXT_TYPES:
+            if p.get("text"):
+                _add_text(parts, str(p["text"]))
+        elif t == "tool_use":
+            parts.append(["tool_call", str(p.get("name", "")), canon_args(p.get("input", {}))])
+        elif t == "tool_result":
+            parts.append(["tool_result", text_of(p.get("content"))])
+        elif t in ("thinking", "redacted_thinking"):
+            continue
+        else:
+            rest = {k: v for k, v in p.items() if k not in ("id", "cache_control")}
+            parts.append(["other", dumps(rest)])
+    return parts
+
+
+def _add_text(parts: list[Part], text: str) -> None:
+    if parts and parts[-1][0] == "text":
+        parts[-1][1] += text
+    else:
+        parts.append(["text", text])
+
+
+def canon_message(msg: Any) -> Message:
+    """One chat message (OpenAI or Anthropic shape, object or dict)."""
+    m = plain(msg)
+    if not isinstance(m, dict):
+        return {"role": "", "parts": [["other", dumps(m)]]}
+    role = str(m.get("role", ""))
+    if role == "tool":  # OpenAI tool result message
+        return {"role": role, "parts": [["tool_result", text_of(m.get("content"))]]}
+    parts = parts_of(m.get("content"))
+    for tc in m.get("tool_calls") or []:  # OpenAI assistant tool calls
+        fn = (tc.get("function") if isinstance(tc, dict) else None) or {}
+        parts.append(["tool_call", str(fn.get("name", "")), canon_args(fn.get("arguments", "{}"))])
+    return {"role": role, "parts": parts}
+
+
+def canon_messages(messages: Any) -> list[Message]:
+    if isinstance(messages, str):
+        return [{"role": "user", "parts": [["text", messages]]}]
+    out = [canon_message(m) for m in messages or []]
+    # An assistant message with no content carries nothing and is dropped by some callers.
+    return [m for m in out if m["parts"] or m["role"] != "assistant"]
+
+
+def canon_response_items(items: Any) -> list[Message]:
+    """OpenAI Responses `input` or `output` (a string or an item list)."""
+    if items is None:
+        return []
+    if isinstance(items, str):
+        return [{"role": "user", "parts": [["text", items]]}] if items else []
+    out: list[Message] = []
+    for item in plain(items):
+        if not isinstance(item, dict):
+            out.append({"role": "other", "parts": [["other", dumps(item)]]})
+            continue
+        t = item.get("type", "message")
+        if t == "message":
+            out.append({"role": str(item.get("role", "")), "parts": parts_of(item.get("content"))})
+        elif t == "function_call":
+            call = ["tool_call", str(item.get("name", "")), canon_args(item.get("arguments", "{}"))]
+            out.append({"role": "assistant", "parts": [call]})
+        elif t == "function_call_output":
+            out.append({"role": "tool", "parts": [["tool_result", text_of(item.get("output"))]]})
+        elif t == "reasoning":
+            continue
+        else:
+            rest = {k: v for k, v in item.items() if k not in ("id", "status")}
+            out.append({"role": "other", "parts": [["other", dumps(rest)]]})
+    return out
+
+
+def canon_system(system: Any) -> list[Part]:
+    return parts_of(plain(system))
+
+
+def canon_tools(tools: Any) -> list[Any]:
+    plain_tools = [plain(t) for t in tools or []]
+    return sorted(plain_tools, key=dumps)
+
+
+def text_chars(messages: list[Message]) -> int:
+    return sum(len(p[-1]) for m in messages for p in m["parts"])
+
+
+def message_text(message: Message) -> str:
+    return "".join(p[1] for p in message["parts"] if p[0] == "text")
+
+
+def tool_calls(message: Message) -> list[tuple[str, str]]:
+    return [(p[1], p[2]) for p in message["parts"] if p[0] == "tool_call"]
+
+
+def split_at_last_assistant(messages: list[Message]) -> int:
+    """Index just past the last assistant message; 0 when there is none."""
+    for i in range(len(messages) - 1, -1, -1):
+        if messages[i]["role"] == "assistant":
+            return i + 1
+    return 0

zexl/_config.py

@@ -0,0 +1,76 @@
+"""Configuration and the two-part API key."""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import hmac
+import os
+import re
+from dataclasses import dataclass
+
+DEFAULT_ENDPOINT = "https://zexl-intake.fly.dev"
+_KEY_RE = re.compile(r"^(zx_[A-Za-z0-9]+_[A-Za-z0-9]+)\.([A-Za-z0-9_\-]{22,})$")
+_TRUE = {"1", "true", "yes", "on"}
+
+
+class ConfigError(ValueError):
+    """The SDK cannot start with the given settings. Never raised into customer calls."""
+
+
+@dataclass(frozen=True)
+class Config:
+    token: str
+    """The part of the key that is sent to Zexl for authentication."""
+    salt: bytes
+    """The part of the key that never leaves this process. Keys every hash."""
+    endpoint: str
+    environment: str
+    openai_stream_usage: bool
+
+    @property
+    def salt_check(self) -> str:
+        """Short value that lets the server notice two keys with different salts. Reveals nothing
+        useful about the salt: 32 bits of an HMAC keyed with 128+ bits of entropy."""
+        return hmac.new(self.salt, b"zexl-salt-check", hashlib.sha256).hexdigest()[:8]
+
+
+def parse_key(key: str) -> tuple[str, bytes]:
+    m = _KEY_RE.match(key.strip())
+    if not m:
+        if key.strip().startswith("zx_") and "." not in key:
+            raise ConfigError(
+                "ZEXL_API_KEY has no salt part. Use the full key shown at creation "
+                "(zx_<id>_<secret>.<salt>); the part after the dot stays on this machine."
+            )
+        raise ConfigError("ZEXL_API_KEY is not a valid Zexl key")
+    token, salt_text = m.group(1), m.group(2)
+    try:
+        salt = base64.urlsafe_b64decode(salt_text + "=" * (-len(salt_text) % 4))
+    except ValueError:
+        salt = salt_text.encode()
+    if len(salt) < 16:
+        raise ConfigError("ZEXL_API_KEY salt part is too short")
+    return token, salt
+
+
+def load(
+    api_key: str | None = None,
+    endpoint: str | None = None,
+    environment: str | None = None,
+) -> Config:
+    key = api_key or os.environ.get("ZEXL_API_KEY", "")
+    if not key:
+        raise ConfigError("ZEXL_API_KEY is not set")
+    token, salt = parse_key(key)
+    return Config(
+        token=token,
+        salt=salt,
+        endpoint=(endpoint or os.environ.get("ZEXL_ENDPOINT") or DEFAULT_ENDPOINT).rstrip("/"),
+        environment=(environment or os.environ.get("ZEXL_ENV") or "production")[:64],
+        openai_stream_usage=os.environ.get("ZEXL_OPENAI_STREAM_USAGE", "").lower() in _TRUE,
+    )
+
+
+def disabled_by_env() -> bool:
+    return os.environ.get("ZEXL_DISABLED", "").lower() in _TRUE

zexl/_context.py

@@ -0,0 +1,210 @@
+"""Run, feature and customer labels carried in context variables."""
+
+from __future__ import annotations
+
+import contextvars
+import functools
+import inspect
+import os
+import sys
+import sysconfig
+import threading
+from collections import OrderedDict
+from collections.abc import Callable
+from types import CodeType, TracebackType
+from typing import Any, TypeVar
+
+from ._stitch import new_run_id
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+_LABEL_MAX = 128
+
+
+class _RunState:
+    __slots__ = ("run_id", "_step", "_lock")
+
+    def __init__(self) -> None:
+        self.run_id = new_run_id()
+        self._step = 0
+        self._lock = threading.Lock()
+
+    def next_step(self) -> int:
+        with self._lock:
+            self._step += 1
+            return self._step
+
+
+_run: contextvars.ContextVar[_RunState | None] = contextvars.ContextVar("zexl_run", default=None)
+_feature: contextvars.ContextVar[str | None] = contextvars.ContextVar("zexl_feature", default=None)
+_customer: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "zexl_customer", default=None
+)
+
+
+def clean_label(value: Any) -> str:
+    text = " ".join(str(value).split())
+    return text[:_LABEL_MAX]
+
+
+class _Scope:
+    """Context manager and decorator that sets context variables for its body."""
+
+    def __init__(self, new_run: bool, feature: str | None) -> None:
+        self._new_run = new_run
+        self._feature = clean_label(feature) if feature else None
+        self._tokens: list[tuple[contextvars.ContextVar[Any], contextvars.Token[Any]]] = []
+
+    def __enter__(self) -> _Scope:
+        if self._new_run:
+            self._tokens.append((_run, _run.set(_RunState())))
+        if self._feature:
+            self._tokens.append((_feature, _feature.set(self._feature)))
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        while self._tokens:
+            var, token = self._tokens.pop()
+            try:
+                var.reset(token)
+            except ValueError:
+                # Exited in a different context (for example a generator moved between tasks).
+                var.set(None)
+
+    async def __aenter__(self) -> _Scope:
+        return self.__enter__()
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        self.__exit__(exc_type, exc, tb)
+
+    def __call__(self, fn: F) -> F:
+        new_run, feature = self._new_run, self._feature or clean_label(fn.__qualname__)
+        if inspect.iscoroutinefunction(fn):
+
+            @functools.wraps(fn)
+            async def awrapper(*a: Any, **kw: Any) -> Any:
+                with _Scope(new_run, feature):
+                    return await fn(*a, **kw)
+
+            return awrapper  # type: ignore[return-value]
+
+        @functools.wraps(fn)
+        def wrapper(*a: Any, **kw: Any) -> Any:
+            with _Scope(new_run, feature):
+                return fn(*a, **kw)
+
+        return wrapper  # type: ignore[return-value]
+
+
+def run(name: str | None = None, *, feature: str | None = None) -> _Scope:
+    """Marks one unit of agent work. Calls inside share a run id and count steps from 1.
+
+    Use as `with zexl.run("research"):` or as a decorator. `name` is used as the feature label
+    unless `feature` is given.
+    """
+    return _Scope(True, feature or name)
+
+
+def feature(name: str) -> _Scope:
+    """Sets the feature label for calls inside, without starting a run."""
+    return _Scope(False, name)
+
+
+def set_customer(customer_id: str | None) -> None:
+    """Attributes later calls in this context to a customer. Pass an id, not personal data."""
+    _customer.set(clean_label(customer_id) if customer_id else None)
+
+
+def current_customer() -> str | None:
+    return _customer.get()
+
+
+def explicit_run() -> _RunState | None:
+    return _run.get()
+
+
+def explicit_feature() -> str | None:
+    return _feature.get()
+
+
+# --------------------------------------------------------------------------- ambient trace
+
+_trace_steps: OrderedDict[str, int] = OrderedDict()
+_trace_lock = threading.Lock()
+
+
+def ambient_trace() -> tuple[str, int, str | None] | None:
+    """(trace id, step, http route) of the application's own active OpenTelemetry span, if any."""
+    otel = sys.modules.get("opentelemetry.trace")
+    if otel is None:
+        return None
+    try:
+        span = otel.get_current_span()
+        ctx = span.get_span_context()
+        if not ctx.is_valid:
+            return None
+        trace_id = format(ctx.trace_id, "032x")
+        attrs = getattr(span, "attributes", None) or {}
+        route = attrs.get("http.route")
+    except Exception:  # noqa: BLE001
+        return None
+    with _trace_lock:
+        step = _trace_steps.get(trace_id, 0) + 1
+        _trace_steps[trace_id] = step
+        _trace_steps.move_to_end(trace_id)
+        while len(_trace_steps) > 4096:
+            _trace_steps.popitem(last=False)
+    return trace_id, step, clean_label(route) if route else None
+
+
+# --------------------------------------------------------------------------- caller frame
+
+_PKG_DIR = os.path.dirname(os.path.abspath(__file__))
+_STDLIB = sysconfig.get_paths().get("stdlib", "\0")
+_frame_cache: dict[CodeType, str | None] = {}
+
+
+def _frame_label(code: CodeType) -> str | None:
+    filename = code.co_filename
+    if (
+        filename.startswith(("<", _PKG_DIR, _STDLIB))
+        or "site-packages" in filename
+        or "dist-packages" in filename
+    ):
+        return None
+    module = os.path.splitext(os.path.basename(filename))[0]
+    name = getattr(code, "co_qualname", code.co_name)
+    if name == "<module>":
+        return clean_label(module)
+    return clean_label(f"{module}.{name}".replace(".<locals>", ""))
+
+
+def caller_feature() -> str | None:
+    """`module.function` of the nearest frame that is application code."""
+    try:
+        frame = sys._getframe(2)
+    except ValueError:
+        return None
+    depth = 0
+    while frame is not None and depth < 80:
+        code = frame.f_code
+        if code not in _frame_cache:
+            if len(_frame_cache) > 2048:
+                _frame_cache.clear()
+            _frame_cache[code] = _frame_label(code)
+        label = _frame_cache[code]
+        if label:
+            return label
+        frame = frame.f_back  # type: ignore[assignment]
+        depth += 1
+    return None