tokenjam 0.2.0__py3-none-any.whl

tokenjam/sdk/integrations/litellm.py

@@ -0,0 +1,323 @@
+"""
+LiteLLM provider integration.
+
+Wraps litellm.completion and litellm.acompletion to automatically
+create OTel spans with token usage and provider attribution.
+
+LiteLLM provides a unified interface across 100+ LLM providers. A single
+patch_litellm() call gives tj coverage across all of them.
+
+When both patch_litellm() and patch_openai()/patch_anthropic() are active,
+the LiteLLM patch wins — a context variable prevents inner provider patches
+from creating duplicate spans.
+"""
+from __future__ import annotations
+
+import contextvars
+import functools
+import logging
+
+from opentelemetry import trace
+
+from tokenjam.otel.semconv import GenAIAttributes
+
+logger = logging.getLogger(__name__)
+
+# Context variable used to suppress inner provider patches (openai, anthropic)
+# when a call originates from litellm.completion/acompletion.
+_tj_litellm_active: contextvars.ContextVar[bool] = contextvars.ContextVar(
+    "_tj_litellm_active", default=False,
+)
+
+
+def _parse_provider(model: str, response: object) -> str:
+    """Extract provider name from response or model string prefix."""
+    # Prefer the actual provider from LiteLLM's hidden params
+    hidden = getattr(response, "_hidden_params", None)
+    if isinstance(hidden, dict):
+        provider = hidden.get("custom_llm_provider")
+        if provider:
+            return str(provider)
+    # Fallback: infer from model string prefix (e.g. "anthropic/claude-...")
+    if "/" in model:
+        return model.split("/", 1)[0]
+    return "litellm"
+
+
+def _strip_provider_prefix(model: str) -> str:
+    """Strip provider prefix from LiteLLM model string.
+
+    'openai/gpt-4o-mini' -> 'gpt-4o-mini', consistent with patch_openai.
+    """
+    if "/" in model:
+        return model.split("/", 1)[1]
+    return model
+
+
+class LiteLLMIntegration:
+    name = "litellm"
+    installed = False
+
+    def __init__(self) -> None:
+        self._original_completion = None
+        self._original_acompletion = None
+        self._tracer = None
+
+    def install(self, tracer) -> None:
+        """Patch litellm.completion and litellm.acompletion."""
+        if self.installed:
+            return
+        self._tracer = tracer
+        try:
+            import litellm
+        except ImportError:
+            logger.warning("litellm package not installed — skipping patch")
+            return
+
+        self._original_completion = litellm.completion
+        self._original_acompletion = litellm.acompletion
+        integration = self
+
+        @functools.wraps(self._original_completion)
+        def patched_completion(*args, **kwargs):
+            raw_model = str(args[0] if args else kwargs.get("model", "unknown"))
+            model = _strip_provider_prefix(raw_model)
+            is_stream = kwargs.get("stream", False)
+
+            span = integration._tracer.start_span(GenAIAttributes.SPAN_LLM_CALL)
+            span.set_attribute(GenAIAttributes.REQUEST_MODEL, model)
+
+            # Inherit agent_id / conversation_id from parent span
+            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,
+                    )
+
+            token = _tj_litellm_active.set(True)
+            try:
+                response = integration._original_completion(*args, **kwargs)
+                if is_stream:
+                    return _SyncStreamWrapper(response, span, raw_model, token)
+                _record_usage(response, span, raw_model)
+                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:
+                if not is_stream:
+                    span.end()
+                    _tj_litellm_active.reset(token)
+
+        litellm.completion = patched_completion
+
+        @functools.wraps(self._original_acompletion)
+        async def patched_acompletion(*args, **kwargs):
+            raw_model = str(args[0] if args else kwargs.get("model", "unknown"))
+            model = _strip_provider_prefix(raw_model)
+            is_stream = kwargs.get("stream", False)
+
+            span = integration._tracer.start_span(GenAIAttributes.SPAN_LLM_CALL)
+            span.set_attribute(GenAIAttributes.REQUEST_MODEL, model)
+
+            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,
+                    )
+
+            token = _tj_litellm_active.set(True)
+            try:
+                response = await integration._original_acompletion(
+                    *args, **kwargs,
+                )
+                if is_stream:
+                    return _AsyncStreamWrapper(response, span, raw_model, token)
+                _record_usage(response, span, raw_model)
+                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:
+                if not is_stream:
+                    span.end()
+                    _tj_litellm_active.reset(token)
+
+        litellm.acompletion = patched_acompletion
+
+        self.installed = True
+        logger.debug("LiteLLM integration installed")
+
+    def uninstall(self) -> None:
+        if not self.installed:
+            return
+        try:
+            import litellm
+            if self._original_completion:
+                litellm.completion = self._original_completion
+            if self._original_acompletion:
+                litellm.acompletion = self._original_acompletion
+        except ImportError:
+            pass
+        self.installed = False
+
+
+def _record_usage(response: object, span, model: str) -> None:
+    """Extract provider and token usage from a LiteLLM ModelResponse."""
+    provider = _parse_provider(model, response)
+    span.set_attribute(GenAIAttributes.PROVIDER_NAME, provider)
+    # Update model attribute to strip provider prefix for consistent pricing
+    span.set_attribute(GenAIAttributes.REQUEST_MODEL, _strip_provider_prefix(model))
+
+    usage = getattr(response, "usage", None)
+    if usage:
+        prompt_tokens = getattr(usage, "prompt_tokens", None)
+        if prompt_tokens is not None:
+            span.set_attribute(GenAIAttributes.INPUT_TOKENS, prompt_tokens)
+        completion_tokens = getattr(usage, "completion_tokens", None)
+        if completion_tokens is not None:
+            span.set_attribute(
+                GenAIAttributes.OUTPUT_TOKENS, completion_tokens,
+            )
+
+
+class _SyncStreamWrapper:
+    """Wraps a LiteLLM sync stream to capture usage and end the span."""
+
+    def __init__(self, stream, span, model: str, token):
+        self._stream = stream
+        self._span = span
+        self._model = model
+        self._token = token
+        self._usage = None
+        self._last_chunk = None
+
+    def __iter__(self):
+        _ok = False
+        try:
+            for chunk in self._stream:
+                usage = getattr(chunk, "usage", None)
+                if usage:
+                    self._usage = usage
+                self._last_chunk = chunk
+                yield chunk
+            _ok = True
+        except Exception as exc:
+            self._span.set_status(
+                trace.Status(trace.StatusCode.ERROR, str(exc)),
+            )
+            raise
+        finally:
+            provider = _parse_provider(self._model, self._last_chunk)
+            self._span.set_attribute(GenAIAttributes.PROVIDER_NAME, provider)
+            self._span.set_attribute(
+                GenAIAttributes.REQUEST_MODEL,
+                _strip_provider_prefix(self._model),
+            )
+            if self._usage:
+                prompt_tokens = getattr(
+                    self._usage, "prompt_tokens", None,
+                )
+                if prompt_tokens is not None:
+                    self._span.set_attribute(
+                        GenAIAttributes.INPUT_TOKENS, prompt_tokens,
+                    )
+                completion_tokens = getattr(
+                    self._usage, "completion_tokens", None,
+                )
+                if completion_tokens is not None:
+                    self._span.set_attribute(
+                        GenAIAttributes.OUTPUT_TOKENS, completion_tokens,
+                    )
+            if _ok:
+                self._span.set_status(trace.Status(trace.StatusCode.OK))
+            self._span.end()
+            _tj_litellm_active.reset(self._token)
+
+    def __next__(self):
+        return self._stream.__next__()
+
+
+class _AsyncStreamWrapper:
+    """Wraps a LiteLLM async stream to capture usage and end the span."""
+
+    def __init__(self, stream, span, model: str, token):
+        self._stream = stream
+        self._span = span
+        self._model = model
+        self._token = token
+        self._usage = None
+        self._last_chunk = None
+
+    def __aiter__(self):
+        return self._iterate()
+
+    async def _iterate(self):
+        _ok = False
+        try:
+            async for chunk in self._stream:
+                usage = getattr(chunk, "usage", None)
+                if usage:
+                    self._usage = usage
+                self._last_chunk = chunk
+                yield chunk
+            _ok = True
+        except Exception as exc:
+            self._span.set_status(
+                trace.Status(trace.StatusCode.ERROR, str(exc)),
+            )
+            raise
+        finally:
+            provider = _parse_provider(self._model, self._last_chunk)
+            self._span.set_attribute(GenAIAttributes.PROVIDER_NAME, provider)
+            self._span.set_attribute(
+                GenAIAttributes.REQUEST_MODEL,
+                _strip_provider_prefix(self._model),
+            )
+            if self._usage:
+                prompt_tokens = getattr(
+                    self._usage, "prompt_tokens", None,
+                )
+                if prompt_tokens is not None:
+                    self._span.set_attribute(
+                        GenAIAttributes.INPUT_TOKENS, prompt_tokens,
+                    )
+                completion_tokens = getattr(
+                    self._usage, "completion_tokens", None,
+                )
+                if completion_tokens is not None:
+                    self._span.set_attribute(
+                        GenAIAttributes.OUTPUT_TOKENS, completion_tokens,
+                    )
+            if _ok:
+                self._span.set_status(trace.Status(trace.StatusCode.OK))
+            self._span.end()
+            _tj_litellm_active.reset(self._token)
+
+
+def patch_litellm() -> None:
+    """Convenience function. Instantiates and installs LiteLLMIntegration."""
+    from tokenjam.sdk.bootstrap import ensure_initialised
+    ensure_initialised()
+    integration = LiteLLMIntegration()
+    integration.install(trace.get_tracer("tokenjam.sdk"))

tokenjam/sdk/integrations/llamaindex.py

@@ -0,0 +1,52 @@
+"""
+LlamaIndex framework integration.
+
+LlamaIndex has native OTel support. This module is a thin convenience wrapper
+that configures LlamaIndex's built-in OTel instrumentation to point at tj serve.
+
+Does NOT monkey-patch LlamaIndex internals — uses its official instrumentation API.
+
+Requires: pip install opentelemetry-instrumentation-llama-index
+"""
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from tokenjam.core.config import load_config
+
+if TYPE_CHECKING:
+    from tokenjam.core.config import TjConfig
+
+logger = logging.getLogger(__name__)
+
+
+def patch_llamaindex(config: TjConfig | None = None) -> None:
+    """
+    Configure LlamaIndex's built-in OTel support to export to tj serve.
+    """
+    if config is None:
+        config = load_config()
+
+    try:
+        from opentelemetry.instrumentation.llama_index import LlamaIndexInstrumentor
+        from opentelemetry.sdk.trace import TracerProvider
+        from opentelemetry.sdk.trace.export import BatchSpanProcessor
+        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+    except ImportError:
+        logger.warning(
+            "opentelemetry-instrumentation-llama-index not installed — "
+            "run: pip install opentelemetry-instrumentation-llama-index"
+        )
+        return
+
+    endpoint = f"http://{config.api.host}:{config.api.port}/api/v1/spans"
+    headers: dict[str, str] = {}
+    if config.security.ingest_secret:
+        headers["Authorization"] = f"Bearer {config.security.ingest_secret}"
+
+    exporter = OTLPSpanExporter(endpoint=endpoint, headers=headers)
+    provider = TracerProvider()
+    provider.add_span_processor(BatchSpanProcessor(exporter))
+    LlamaIndexInstrumentor().instrument(tracer_provider=provider)
+    logger.debug("LlamaIndex integration installed (via native OTel)")

tokenjam/sdk/integrations/nemoclaw.py

@@ -0,0 +1,139 @@
+"""
+NemoClaw/OpenShell Gateway WebSocket observer.
+
+Connects to the OpenShell Gateway as an observer client, receives sandbox
+events, and translates them into OTel spans fed into the ingest pipeline.
+
+NOTICE: NemoClaw is licensed under Apache License 2.0.
+This integration module acknowledges the upstream Apache 2.0 license.
+See https://github.com/NVIDIA/NemoClaw/blob/main/LICENSE
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from typing import TYPE_CHECKING
+
+
+from tokenjam.core.models import NormalizedSpan, SpanKind, SpanStatus
+from tokenjam.otel.semconv import TjAttributes
+from tokenjam.utils.ids import new_span_id, new_trace_id
+from tokenjam.utils.time_parse import utcnow
+
+if TYPE_CHECKING:
+    from tokenjam.core.config import TjConfig
+
+logger = logging.getLogger(__name__)
+
+SANDBOX_EVENT_MAP = {
+    "network_blocked": "network_blocked",
+    "fs_access_denied": "fs_denied",
+    "syscall_blocked": "syscall_denied",
+    "inference_reroute": "inference_rerouted",
+}
+
+
+class NemoClawGatewayObserver:
+    """
+    Observes a NemoClaw OpenShell sandbox by connecting to the
+    OpenShell Gateway WebSocket as an observer client.
+
+    All inference calls, blocked network requests, filesystem denials,
+    and syscall blocks are translated into OTel spans and fed into the
+    standard tj ingest pipeline.
+
+    Usage:
+        observer = NemoClawGatewayObserver(ingest_pipeline)
+        asyncio.run(observer.connect())  # runs until cancelled
+    """
+
+    def __init__(self, ingest_pipeline, gateway_url: str = "ws://127.0.0.1:18789"):
+        self.pipeline = ingest_pipeline
+        self.gateway_url = gateway_url
+
+    async def connect(self) -> None:
+        """Connect and observe. Reconnects with backoff on disconnect."""
+        try:
+            import websockets
+        except ImportError:
+            logger.error("websockets package not installed — cannot observe NemoClaw")
+            return
+
+        backoff = 1.0
+        while True:
+            try:
+                async with websockets.connect(self.gateway_url) as ws:
+                    logger.info("Connected to NemoClaw gateway at %s", self.gateway_url)
+                    backoff = 1.0
+                    async for message in ws:
+                        try:
+                            event = json.loads(message)
+                            span = self._translate_event(event)
+                            if span:
+                                self.pipeline.process(span)
+                        except (json.JSONDecodeError, KeyError) as exc:
+                            logger.debug("Skipping malformed gateway event: %s", exc)
+            except Exception as exc:
+                logger.warning(
+                    "NemoClaw gateway disconnected: %s — reconnecting in %.0fs",
+                    exc, backoff,
+                )
+                await asyncio.sleep(backoff)
+                backoff = min(backoff * 2, 60.0)
+
+    def _translate_event(self, event: dict) -> NormalizedSpan | None:
+        """Convert an OpenShell gateway event to a NormalizedSpan."""
+        event_type = event.get("type", "")
+        ocw_event = SANDBOX_EVENT_MAP.get(event_type)
+        if not ocw_event:
+            return None
+
+        now = utcnow()
+        attrs: dict = {
+            TjAttributes.SANDBOX_EVENT: ocw_event,
+        }
+
+        if ocw_event == "network_blocked":
+            attrs[TjAttributes.EGRESS_HOST] = event.get("host", "unknown")
+            attrs[TjAttributes.EGRESS_PORT] = event.get("port")
+        elif ocw_event == "fs_denied":
+            attrs[TjAttributes.FILESYSTEM_PATH] = event.get("path", "unknown")
+        elif ocw_event == "syscall_denied":
+            attrs[TjAttributes.SYSCALL_NAME] = event.get("syscall", "unknown")
+
+        return NormalizedSpan(
+            span_id=new_span_id(),
+            trace_id=new_trace_id(),
+            name=f"sandbox.{ocw_event}",
+            kind=SpanKind.INTERNAL,
+            status_code=SpanStatus.ERROR,
+            start_time=now,
+            end_time=now,
+            duration_ms=0.0,
+            agent_id=event.get("agent_id"),
+            attributes=attrs,
+        )
+
+
+def watch_nemoclaw(
+    gateway_url: str = "ws://127.0.0.1:18789",
+    config: TjConfig | None = None,
+) -> NemoClawGatewayObserver:
+    """
+    Convenience function. Creates an observer instance.
+    Call observer.connect() in an asyncio task to start observing.
+
+    Note: You must pass an ingest_pipeline instance. If config is None,
+    a default config is loaded.
+    """
+    from tokenjam.core.config import load_config
+    from tokenjam.sdk.bootstrap import ensure_initialised, _pipeline
+    if config is None:
+        config = load_config()
+    ensure_initialised()
+    logger.info("NemoClaw observer created for %s", gateway_url)
+    return NemoClawGatewayObserver(
+        ingest_pipeline=_pipeline,
+        gateway_url=gateway_url,
+    )

tokenjam/sdk/integrations/openai.py

@@ -0,0 +1,159 @@
+"""
+OpenAI provider integration.
+
+Wraps openai.resources.chat.completions.Completions.create to automatically
+create OTel spans with token usage and model attributes.
+
+Also works for OpenAI-compatible providers (Groq, Together, Fireworks, xAI,
+Azure OpenAI) — pass the provider's base_url and set provider name from it.
+"""
+from __future__ import annotations
+
+import functools
+import logging
+
+from opentelemetry import trace
+
+from tokenjam.otel.semconv import GenAIAttributes
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAIIntegration:
+    name = "openai"
+    installed = False
+
+    def __init__(self, provider_name: str = "openai") -> None:
+        self._original_create = None
+        self._tracer = None
+        self._provider_name = provider_name
+
+    def install(self, tracer) -> None:
+        """Patch openai.resources.chat.completions.Completions.create."""
+        if self.installed:
+            return
+        self._tracer = tracer
+        try:
+            from openai.resources.chat.completions import Completions
+        except ImportError:
+            logger.warning("openai package not installed — skipping patch")
+            return
+
+        self._original_create = Completions.create
+        integration = self
+
+        @functools.wraps(self._original_create)
+        def patched_create(self_comp, *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_comp, *args, **kwargs)
+            span = integration._tracer.start_span(GenAIAttributes.SPAN_LLM_CALL)
+            span.set_attribute(GenAIAttributes.PROVIDER_NAME, integration._provider_name)
+            span.set_attribute(
+                GenAIAttributes.REQUEST_MODEL,
+                kwargs.get("model", "unknown"),
+            )
+            is_stream = kwargs.get("stream", False)
+            try:
+                response = integration._original_create(self_comp, *args, **kwargs)
+                if is_stream:
+                    return _StreamWrapper(response, span)
+                if hasattr(response, "usage") and response.usage:
+                    span.set_attribute(
+                        GenAIAttributes.INPUT_TOKENS,
+                        response.usage.prompt_tokens,
+                    )
+                    span.set_attribute(
+                        GenAIAttributes.OUTPUT_TOKENS,
+                        response.usage.completion_tokens,
+                    )
+                span.set_status(trace.Status(trace.StatusCode.OK))
+                span.end()
+                return response
+            except Exception as exc:
+                span.set_status(trace.Status(trace.StatusCode.ERROR, str(exc)))
+                span.end()
+                raise
+
+        Completions.create = patched_create
+        self.installed = True
+        logger.debug("OpenAI integration installed (provider=%s)", self._provider_name)
+
+    def uninstall(self) -> None:
+        if not self.installed:
+            return
+        try:
+            from openai.resources.chat.completions import Completions
+            if self._original_create:
+                Completions.create = self._original_create
+        except ImportError:
+            pass
+        self.installed = False
+
+
+class _StreamWrapper:
+    """Wraps an OpenAI stream to capture final usage chunk and end the span."""
+
+    def __init__(self, stream, span):
+        self._stream = stream
+        self._span = span
+        self._usage = None
+
+    def __iter__(self):
+        _ok = False
+        try:
+            for chunk in self._stream:
+                if hasattr(chunk, "usage") and chunk.usage:
+                    self._usage = chunk.usage
+                yield chunk
+            _ok = True
+        except Exception as exc:
+            self._span.set_status(trace.Status(trace.StatusCode.ERROR, str(exc)))
+            raise
+        finally:
+            if self._usage:
+                self._span.set_attribute(
+                    GenAIAttributes.INPUT_TOKENS,
+                    self._usage.prompt_tokens,
+                )
+                self._span.set_attribute(
+                    GenAIAttributes.OUTPUT_TOKENS,
+                    self._usage.completion_tokens,
+                )
+            if _ok:
+                self._span.set_status(trace.Status(trace.StatusCode.OK))
+            self._span.end()
+
+    def __next__(self):
+        return self._stream.__next__()
+
+
+def patch_openai(base_url: str | None = None) -> None:
+    """
+    Wraps the OpenAI client.
+    Also works for OpenAI-compatible providers (Groq, Together, Fireworks, xAI,
+    Azure OpenAI) — pass the provider's base_url and set provider name from it.
+    """
+    from tokenjam.sdk.bootstrap import ensure_initialised
+    ensure_initialised()
+    provider = "openai"
+    if base_url:
+        # Infer provider name from base_url domain
+        from urllib.parse import urlparse
+        domain = urlparse(base_url).hostname or ""
+        if "groq" in domain:
+            provider = "groq"
+        elif "together" in domain:
+            provider = "together"
+        elif "fireworks" in domain:
+            provider = "fireworks"
+        elif "xai" in domain or "x.ai" in domain:
+            provider = "xai"
+        elif "azure" in domain:
+            provider = "azure.openai"
+        else:
+            provider = domain.split(".")[0] if domain else "openai-compatible"
+
+    integration = OpenAIIntegration(provider_name=provider)
+    integration.install(trace.get_tracer("tokenjam.sdk"))