tokenjam 0.2.0__py3-none-any.whl

tokenjam/sdk/bootstrap.py

@@ -0,0 +1,120 @@
+"""
+Auto-bootstrap: lazily initialise the OCW TracerProvider + IngestPipeline
+the first time @watch() or a provider patch creates a span.
+
+This ensures that SDK users don't need to manually wire up the pipeline.
+"""
+from __future__ import annotations
+
+import atexit
+import logging
+import threading
+
+
+logger = logging.getLogger("tokenjam.sdk")
+
+_lock = threading.Lock()
+_initialised = False
+_provider = None
+_pipeline = None
+
+
+def ensure_initialised() -> None:
+    """
+    Idempotent bootstrap. Safe to call multiple times / from multiple threads.
+    Sets up: config -> DuckDB -> IngestPipeline -> TjSpanExporter -> TracerProvider.
+    """
+    global _initialised, _provider, _pipeline
+    if _initialised:
+        return
+
+    with _lock:
+        if _initialised:
+            return
+
+        try:
+            from tokenjam.core.config import load_config
+            from tokenjam.otel.provider import build_tracer_provider
+
+            config = load_config()
+
+            # Check if tj serve is running — use HTTP exporter if so
+            if _try_http_mode(config):
+                _initialised = True
+                atexit.register(_shutdown)
+                return
+
+            # Direct DuckDB mode
+            from tokenjam.core.db import open_db
+            from tokenjam.core.ingest import build_default_pipeline
+
+            db = open_db(config.storage)
+            pipeline = build_default_pipeline(db, config)
+            _pipeline = pipeline
+            _provider = build_tracer_provider(config, pipeline)
+            _initialised = True
+
+            # Ensure spans are flushed on exit
+            atexit.register(_shutdown)
+
+            logger.debug("OCW: writing spans to local DuckDB (%s)", config.storage.path)
+
+        except Exception as exc:
+            # DuckDB lock error — try HTTP fallback
+            err_msg = str(exc).lower()
+            if "lock" in err_msg or "i/o error" in err_msg:
+                try:
+                    config = load_config()
+                    if _try_http_mode(config):
+                        _initialised = True
+                        atexit.register(_shutdown)
+                        return
+                except Exception:
+                    pass
+            logger.warning("OCW bootstrap failed — spans will not be recorded: %s", exc)
+            _initialised = True  # Don't retry on every call
+
+
+def _try_http_mode(config) -> bool:
+    """Try to connect to tj serve and set up HTTP exporter. Returns True on success."""
+    global _provider
+    import httpx
+    base_url = f"http://{config.api.host}:{config.api.port}"
+    try:
+        resp = httpx.get(f"{base_url}/api/v1/status", timeout=2)
+        if resp.status_code not in (200, 401):
+            return False
+    except (httpx.ConnectError, httpx.TimeoutException):
+        return False
+
+    from tokenjam.sdk.http_exporter import TjHttpExporter
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace.export import BatchSpanProcessor
+    from opentelemetry.sdk.resources import Resource
+    from opentelemetry import trace as trace_api
+    import tokenjam
+
+    endpoint = f"{base_url}/api/v1/spans"
+    exporter = TjHttpExporter(endpoint, config.security.ingest_secret)
+
+    resource = Resource.create({
+        "service.name": "tokenjam",
+        "service.version": tokenjam.__version__,
+    })
+    provider = TracerProvider(resource=resource)
+    provider.add_span_processor(BatchSpanProcessor(exporter))
+    trace_api.set_tracer_provider(provider)
+    _provider = provider
+
+    logger.info("TJ: sending spans to tj serve at %s", base_url)
+    return True
+
+
+def _shutdown() -> None:
+    """Flush pending spans on interpreter exit."""
+    if _provider is not None:
+        try:
+            _provider.force_flush(timeout_millis=5000)
+            _provider.shutdown()
+        except Exception:
+            pass

tokenjam/sdk/http_exporter.py

@@ -0,0 +1,109 @@
+"""
+OTel SpanExporter that POSTs OTLP JSON to tj serve's /api/v1/spans endpoint.
+Used when tj serve is running and holds the DuckDB lock.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Sequence
+
+import httpx
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from opentelemetry.sdk.trace import ReadableSpan
+
+logger = logging.getLogger("tokenjam.sdk")
+
+
+class TjHttpExporter(SpanExporter):
+    """Exports spans via HTTP POST to tj serve."""
+
+    def __init__(self, endpoint: str, ingest_secret: str) -> None:
+        self._endpoint = endpoint
+        self._headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {ingest_secret}",
+        }
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        otlp_spans = [_span_to_otlp(s) for s in spans]
+        payload = {
+            "resourceSpans": [{
+                "resource": {"attributes": [
+                    {"key": "service.name", "value": {"stringValue": "tokenjam"}},
+                ]},
+                "scopeSpans": [{"spans": otlp_spans}],
+            }],
+        }
+        try:
+            resp = httpx.post(
+                self._endpoint, json=payload, headers=self._headers, timeout=5.0,
+            )
+            if resp.status_code < 300:
+                return SpanExportResult.SUCCESS
+            logger.warning("tj serve returned %d on span export", resp.status_code)
+        except (httpx.ConnectError, httpx.TimeoutException) as exc:
+            logger.warning("Failed to export spans to tj serve: %s", exc)
+        return SpanExportResult.FAILURE
+
+    def shutdown(self) -> None:
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+
+def _span_to_otlp(span: ReadableSpan) -> dict:
+    """Convert an OTel ReadableSpan to OTLP JSON dict."""
+    ctx = span.context
+    attrs = []
+    for k, v in (span.attributes or {}).items():
+        attrs.append({"key": k, "value": _to_otlp_value(v)})
+
+    result: dict = {
+        "traceId": format(ctx.trace_id, "032x") if ctx else "",
+        "spanId": format(ctx.span_id, "016x") if ctx else "",
+        "name": span.name,
+        "kind": (span.kind.value + 1) if span.kind is not None else 1,
+        "startTimeUnixNano": str(span.start_time or 0),
+        "endTimeUnixNano": str(span.end_time or 0),
+        "attributes": attrs,
+        "status": {},
+        "events": [],
+    }
+
+    if span.parent and span.parent.span_id:
+        result["parentSpanId"] = format(span.parent.span_id, "016x")
+
+    if span.status:
+        code_map = {0: 0, 1: 1, 2: 2}  # UNSET, OK, ERROR
+        result["status"] = {
+            "code": code_map.get(span.status.status_code.value, 0),
+        }
+        if span.status.description:
+            result["status"]["message"] = span.status.description
+
+    for event in span.events or []:
+        evt_attrs = [
+            {"key": k, "value": _to_otlp_value(v)}
+            for k, v in (event.attributes or {}).items()
+        ]
+        result["events"].append({
+            "name": event.name,
+            "timeUnixNano": str(event.timestamp or 0),
+            "attributes": evt_attrs,
+        })
+
+    return result
+
+
+def _to_otlp_value(v: object) -> dict:
+    """Convert a Python value to an OTLP AttributeValue dict."""
+    if isinstance(v, bool):
+        return {"boolValue": v}
+    if isinstance(v, int):
+        return {"intValue": str(v)}
+    if isinstance(v, float):
+        return {"doubleValue": v}
+    if isinstance(v, (list, tuple)):
+        return {"arrayValue": {"values": [_to_otlp_value(item) for item in v]}}
+    return {"stringValue": str(v)}

tokenjam/sdk/integrations/anthropic.py

@@ -0,0 +1,200 @@
+"""
+Anthropic provider integration.
+
+Wraps anthropic.resources.Messages.create and .stream to automatically
+create OTel spans with token usage and model attributes.
+"""
+from __future__ import annotations
+
+import functools
+import logging
+
+from opentelemetry import trace
+
+from tokenjam.otel.semconv import GenAIAttributes
+
+logger = logging.getLogger(__name__)
+
+
+class AnthropicIntegration:
+    name = "anthropic"
+    installed = False
+
+    def __init__(self) -> None:
+        self._original_create = None
+        self._original_stream = None
+        self._tracer = None
+
+    def install(self, tracer) -> None:
+        """Patch anthropic.resources.Messages.create and .stream."""
+        if self.installed:
+            return
+        self._tracer = tracer
+        try:
+            from anthropic.resources import Messages
+        except ImportError:
+            logger.warning("anthropic package not installed — skipping patch")
+            return
+
+        self._original_create = Messages.create
+        self._original_stream = getattr(Messages, "stream", None)
+
+        integration = self
+
+        @functools.wraps(self._original_create)
+        def patched_create(self_msg, *args, **kwargs):
+            # Skip span if call originates from litellm (avoids double-counting)
+            from tokenjam.sdk.integrations.litellm import _tj_litellm_active
+            if _tj_litellm_active.get(False):
+                return integration._original_create(self_msg, *args, **kwargs)
+            span = integration._tracer.start_span(GenAIAttributes.SPAN_LLM_CALL)
+            span.set_attribute(GenAIAttributes.PROVIDER_NAME, "anthropic")
+            span.set_attribute(
+                GenAIAttributes.REQUEST_MODEL,
+                kwargs.get("model", "unknown"),
+            )
+            # Inherit agent_id from parent span (set by @watch())
+            parent_span = trace.get_current_span()
+            if parent_span and parent_span.is_recording():
+                agent_id = parent_span.attributes.get(GenAIAttributes.AGENT_ID)
+                if agent_id:
+                    span.set_attribute(GenAIAttributes.AGENT_ID, agent_id)
+                conv_id = parent_span.attributes.get(GenAIAttributes.CONVERSATION_ID)
+                if conv_id:
+                    span.set_attribute(GenAIAttributes.CONVERSATION_ID, conv_id)
+            try:
+                response = integration._original_create(self_msg, *args, **kwargs)
+                if hasattr(response, "usage"):
+                    span.set_attribute(
+                        GenAIAttributes.INPUT_TOKENS,
+                        response.usage.input_tokens,
+                    )
+                    span.set_attribute(
+                        GenAIAttributes.OUTPUT_TOKENS,
+                        response.usage.output_tokens,
+                    )
+                    cache_read = getattr(response.usage, "cache_read_input_tokens", None)
+                    if cache_read:
+                        span.set_attribute(GenAIAttributes.CACHE_READ_TOKENS, cache_read)
+                    cache_create = getattr(response.usage, "cache_creation_input_tokens", None)
+                    if cache_create:
+                        span.set_attribute(GenAIAttributes.CACHE_CREATE_TOKENS, cache_create)
+                span.set_status(trace.Status(trace.StatusCode.OK))
+                return response
+            except TypeError as exc:
+                if "api_key" in str(exc) or "auth" in str(exc).lower():
+                    span.set_status(trace.Status(trace.StatusCode.ERROR, str(exc)))
+                    import sys
+                    print(
+                        "\n\033[1;31mError: Anthropic API key not found.\033[0m\n"
+                        "\n"
+                        "  Set it in your environment:\n"
+                        "\n"
+                        "    export ANTHROPIC_API_KEY='sk-ant-...'\n"
+                        "\n"
+                        "  Or pass it directly:\n"
+                        "\n"
+                        "    anthropic.Anthropic(api_key='...')\n",
+                        file=sys.stderr,
+                    )
+                    raise SystemExit(1)
+                span.set_status(trace.Status(trace.StatusCode.ERROR, str(exc)))
+                raise
+            except Exception as exc:
+                span.set_status(trace.Status(trace.StatusCode.ERROR, str(exc)))
+                raise
+            finally:
+                span.end()
+
+        Messages.create = patched_create
+
+        if self._original_stream is not None:
+            @functools.wraps(self._original_stream)
+            def patched_stream(self_msg, *args, **kwargs):
+                from tokenjam.sdk.integrations.litellm import _tj_litellm_active
+                if _tj_litellm_active.get(False):
+                    return integration._original_stream(self_msg, *args, **kwargs)
+                span = integration._tracer.start_span(GenAIAttributes.SPAN_LLM_CALL)
+                span.set_attribute(GenAIAttributes.PROVIDER_NAME, "anthropic")
+                span.set_attribute(
+                    GenAIAttributes.REQUEST_MODEL,
+                    kwargs.get("model", "unknown"),
+                )
+                parent_span = trace.get_current_span()
+                if parent_span and parent_span.is_recording():
+                    agent_id = parent_span.attributes.get(GenAIAttributes.AGENT_ID)
+                    if agent_id:
+                        span.set_attribute(GenAIAttributes.AGENT_ID, agent_id)
+                    conv_id = parent_span.attributes.get(GenAIAttributes.CONVERSATION_ID)
+                    if conv_id:
+                        span.set_attribute(GenAIAttributes.CONVERSATION_ID, conv_id)
+                try:
+                    stream = integration._original_stream(self_msg, *args, **kwargs)
+                    return _StreamWrapper(stream, span)
+                except Exception as exc:
+                    span.set_status(trace.Status(trace.StatusCode.ERROR, str(exc)))
+                    span.end()
+                    raise
+
+            Messages.stream = patched_stream
+
+        self.installed = True
+        logger.debug("Anthropic integration installed")
+
+    def uninstall(self) -> None:
+        if not self.installed:
+            return
+        try:
+            from anthropic.resources import Messages
+            if self._original_create:
+                Messages.create = self._original_create
+            if self._original_stream:
+                Messages.stream = self._original_stream
+        except ImportError:
+            pass
+        self.installed = False
+
+
+class _StreamWrapper:
+    """Wraps an Anthropic stream to capture final usage and end the span."""
+
+    def __init__(self, stream, span):
+        self._stream = stream
+        self._span = span
+
+    def __enter__(self):
+        self._stream.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        result = self._stream.__exit__(exc_type, exc_val, exc_tb)
+        final_message = getattr(self._stream, "get_final_message", lambda: None)()
+        if final_message and hasattr(final_message, "usage"):
+            self._span.set_attribute(
+                GenAIAttributes.INPUT_TOKENS,
+                final_message.usage.input_tokens,
+            )
+            self._span.set_attribute(
+                GenAIAttributes.OUTPUT_TOKENS,
+                final_message.usage.output_tokens,
+            )
+        if exc_type is None:
+            self._span.set_status(trace.Status(trace.StatusCode.OK))
+        else:
+            self._span.set_status(trace.Status(trace.StatusCode.ERROR, str(exc_val)))
+        self._span.end()
+        return result
+
+    def __iter__(self):
+        return iter(self._stream)
+
+    def __next__(self):
+        return next(self._stream)
+
+
+def patch_anthropic() -> None:
+    """Convenience function. Instantiates and installs AnthropicIntegration."""
+    from tokenjam.sdk.bootstrap import ensure_initialised
+    ensure_initialised()
+    integration = AnthropicIntegration()
+    integration.install(trace.get_tracer("tokenjam.sdk"))

tokenjam/sdk/integrations/autogen.py

@@ -0,0 +1,97 @@
+"""
+AutoGen framework integration.
+
+Patches ConversableAgent.generate_reply and ConversableAgent.initiate_chat
+to create OTel spans.
+"""
+from __future__ import annotations
+
+import functools
+import logging
+
+from opentelemetry import trace
+
+from tokenjam.otel.semconv import GenAIAttributes
+
+logger = logging.getLogger(__name__)
+
+
+class AutoGenIntegration:
+    name = "autogen"
+    installed = False
+
+    def __init__(self) -> None:
+        self._original_generate_reply = None
+        self._original_initiate_chat = None
+        self._tracer = None
+
+    def install(self, tracer) -> None:
+        if self.installed:
+            return
+        self._tracer = tracer
+        try:
+            from autogen import ConversableAgent
+        except ImportError:
+            logger.warning("pyautogen not installed — skipping patch")
+            return
+
+        integration = self
+
+        self._original_generate_reply = ConversableAgent.generate_reply
+        @functools.wraps(self._original_generate_reply)
+        def patched_generate_reply(self_agent, *args, **kwargs):
+            span = integration._tracer.start_span("autogen.generate_reply")
+            span.set_attribute(GenAIAttributes.PROVIDER_NAME, "autogen")
+            span.set_attribute("autogen.agent_name", getattr(self_agent, "name", "unknown"))
+            try:
+                result = integration._original_generate_reply(self_agent, *args, **kwargs)
+                span.set_status(trace.Status(trace.StatusCode.OK))
+                return result
+            except Exception as exc:
+                span.set_status(trace.Status(trace.StatusCode.ERROR, str(exc)))
+                raise
+            finally:
+                span.end()
+
+        ConversableAgent.generate_reply = patched_generate_reply
+
+        self._original_initiate_chat = ConversableAgent.initiate_chat
+        @functools.wraps(self._original_initiate_chat)
+        def patched_initiate_chat(self_agent, *args, **kwargs):
+            span = integration._tracer.start_span("autogen.initiate_chat")
+            span.set_attribute(GenAIAttributes.PROVIDER_NAME, "autogen")
+            span.set_attribute("autogen.agent_name", getattr(self_agent, "name", "unknown"))
+            try:
+                result = integration._original_initiate_chat(self_agent, *args, **kwargs)
+                span.set_status(trace.Status(trace.StatusCode.OK))
+                return result
+            except Exception as exc:
+                span.set_status(trace.Status(trace.StatusCode.ERROR, str(exc)))
+                raise
+            finally:
+                span.end()
+
+        ConversableAgent.initiate_chat = patched_initiate_chat
+        self.installed = True
+        logger.debug("AutoGen integration installed")
+
+    def uninstall(self) -> None:
+        if not self.installed:
+            return
+        try:
+            from autogen import ConversableAgent
+            if self._original_generate_reply:
+                ConversableAgent.generate_reply = self._original_generate_reply
+            if self._original_initiate_chat:
+                ConversableAgent.initiate_chat = self._original_initiate_chat
+        except ImportError:
+            pass
+        self.installed = False
+
+
+def patch_autogen() -> None:
+    """Convenience function. Instantiates and installs AutoGenIntegration."""
+    from tokenjam.sdk.bootstrap import ensure_initialised
+    ensure_initialised()
+    integration = AutoGenIntegration()
+    integration.install(trace.get_tracer("tokenjam.sdk"))

tokenjam/sdk/integrations/base.py

@@ -0,0 +1,27 @@
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class Integration(Protocol):
+    """
+    Formal interface for all framework and provider integrations.
+
+    Convenience functions like patch_anthropic() instantiate and install
+    the integration automatically. You can also install manually:
+
+        integration = AnthropicIntegration()
+        integration.install(tracer)
+
+    tj doctor uses `integration.installed` to list active integrations
+    and detect conflicts (two integrations patching the same method).
+    """
+    name: str
+    installed: bool
+
+    def install(self, tracer) -> None:
+        """Register all hooks. Idempotent — safe to call multiple times."""
+        ...
+
+    def uninstall(self) -> None:
+        """Remove all hooks. Called on process shutdown."""
+        ...

tokenjam/sdk/integrations/bedrock.py

@@ -0,0 +1,103 @@
+"""
+AWS Bedrock provider integration.
+
+Wraps boto3 client invoke_model and invoke_agent calls. The Bedrock response
+body is JSON — parsed to extract token counts.
+"""
+from __future__ import annotations
+
+import functools
+import json
+import logging
+
+from opentelemetry import trace
+
+from tokenjam.otel.semconv import GenAIAttributes
+
+logger = logging.getLogger(__name__)
+
+
+class BedrockIntegration:
+    name = "aws.bedrock"
+    installed = False
+
+    def __init__(self) -> None:
+        self._original_invoke_model = None
+        self._tracer = None
+
+    def install(self, tracer) -> None:
+        if self.installed:
+            return
+        self._tracer = tracer
+        try:
+            import botocore.client
+        except ImportError:
+            logger.warning("boto3/botocore not installed — skipping patch")
+            return
+
+        original_api_call = botocore.client.ClientCreator._create_api_method
+        integration = self
+
+        def patched_create_api_method(client_creator, py_operation_name, operation_name, service_model):
+            method = original_api_call(client_creator, py_operation_name, operation_name, service_model)
+            if py_operation_name not in ("invoke_model", "invoke_agent"):
+                return method
+
+            @functools.wraps(method)
+            def wrapped(self_client, *args, **kwargs):
+                span = integration._tracer.start_span(GenAIAttributes.SPAN_LLM_CALL)
+                span.set_attribute(GenAIAttributes.PROVIDER_NAME, "aws.bedrock")
+                model_id = kwargs.get("modelId", "unknown")
+                span.set_attribute(GenAIAttributes.REQUEST_MODEL, model_id)
+                try:
+                    response = method(self_client, *args, **kwargs)
+                    _extract_bedrock_usage(response, span)
+                    span.set_status(trace.Status(trace.StatusCode.OK))
+                    return response
+                except Exception as exc:
+                    span.set_status(trace.Status(trace.StatusCode.ERROR, str(exc)))
+                    raise
+                finally:
+                    span.end()
+
+            return wrapped
+
+        botocore.client.ClientCreator._create_api_method = patched_create_api_method
+        self.installed = True
+        logger.debug("Bedrock integration installed")
+
+    def uninstall(self) -> None:
+        self.installed = False
+
+
+def _extract_bedrock_usage(response: dict, span) -> None:
+    """Extract token counts from a Bedrock response."""
+    body = response.get("body")
+    if body is None:
+        return
+    try:
+        if hasattr(body, "read"):
+            raw = body.read()
+            body_dict = json.loads(raw)
+            # Reset stream for caller
+            import io
+            response["body"] = io.BytesIO(raw)
+        else:
+            body_dict = json.loads(body) if isinstance(body, (str, bytes)) else body
+    except (json.JSONDecodeError, TypeError):
+        return
+
+    usage = body_dict.get("usage", {})
+    if usage:
+        if "input_tokens" in usage:
+            span.set_attribute(GenAIAttributes.INPUT_TOKENS, usage["input_tokens"])
+        if "output_tokens" in usage:
+            span.set_attribute(GenAIAttributes.OUTPUT_TOKENS, usage["output_tokens"])
+
+
+def patch_bedrock() -> None:
+    """Convenience function. Instantiates and installs BedrockIntegration."""
+    from tokenjam.sdk.bootstrap import ensure_initialised
+    ensure_initialised()
+    integration = BedrockIntegration()
+    integration.install(trace.get_tracer("tokenjam.sdk"))