monkeybot 2.1.1__py3-none-any.whl

monkeybot/observability/__init__.py

@@ -0,0 +1,220 @@
+"""OpenTelemetry tracing for monkeybot (optional ``[observability]`` extra).
+
+Master switch: ``MONKEYBOT_OTEL_ENABLED`` must be truthy to opt in. When unset,
+tracing stays disabled even if OTel packages are installed (explicit rollout).
+"""
+
+from __future__ import annotations
+
+import importlib
+import importlib.metadata
+import logging
+import os
+from typing import Any
+
+from monkeybot.observability import _state
+from monkeybot.observability.instrumentation import (
+    ObservingProvider,
+    add_tool_hook_span_event,
+    get_current_trace_id_hex_optional,
+    instrument_fastapi_app,
+)
+from monkeybot.observability.propagation import extract_traceparent, inject_traceparent
+from monkeybot.observability.spans import (
+    set_llm_io,
+    set_llm_usage,
+    set_run_output,
+    set_summarize_turns,
+    set_turn_io,
+    set_turn_prompt_tokens,
+    span_llm,
+    span_run,
+    span_subagent,
+    span_summarize,
+    span_tool,
+    span_turn,
+    truncate,
+)
+
+logger = logging.getLogger(__name__)
+
+_TRUTHY = frozenset({"1", "true", "yes", "on"})
+
+
+def _env_truthy(name: str) -> bool:
+    raw = os.environ.get(name)
+    if raw is None:
+        return False
+    return raw.strip().lower() in _TRUTHY
+
+
+def _env_sdk_disabled() -> bool:
+    return os.environ.get("OTEL_SDK_DISABLED", "").strip().lower() in _TRUTHY
+
+
+def is_observability_enabled() -> bool:
+    return _state.is_observability_enabled()
+
+
+def _create_span_processor(exporter: str) -> Any:
+    """Build a span processor for ``otlp`` or ``console``. Tests may monkeypatch this."""
+    from opentelemetry.sdk.trace.export import BatchSpanProcessor
+
+    if exporter == "otlp":
+        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+
+        return BatchSpanProcessor(OTLPSpanExporter())
+    if exporter == "console":
+        from opentelemetry.sdk.trace.export import ConsoleSpanExporter
+
+        return BatchSpanProcessor(ConsoleSpanExporter())
+    return None
+
+
+def init_observability() -> bool:
+    if _state._initialized:
+        return _state._enabled
+
+    _state._initialized = True
+    _state._enabled = False
+
+    if not _env_truthy("MONKEYBOT_OTEL_ENABLED"):
+        logger.info("observability disabled (MONKEYBOT_OTEL_ENABLED is not set or false)")
+        return False
+
+    if _env_sdk_disabled():
+        logger.info("observability disabled (OTEL_SDK_DISABLED=true)")
+        return False
+
+    try:
+        trace_api = importlib.import_module("opentelemetry.trace")
+        resources = importlib.import_module("opentelemetry.sdk.resources")
+        trace_sdk = importlib.import_module("opentelemetry.sdk.trace")
+        w3c = importlib.import_module("opentelemetry.trace.propagation.tracecontext")
+        propagate = importlib.import_module("opentelemetry.propagate")
+    except ImportError:
+        logger.warning("observability disabled (OpenTelemetry packages not installed)")
+        return False
+
+    exporter = os.environ.get("OTEL_TRACES_EXPORTER", "").strip().lower()
+    if exporter not in {"otlp", "console"}:
+        logger.info(
+            "observability disabled (OTEL_TRACES_EXPORTER must be otlp or console, got %r)",
+            exporter or "<unset>",
+        )
+        return False
+
+    service_name = os.environ.get("OTEL_SERVICE_NAME", "monkeybot")
+    attrs: dict[str, str] = {"service.name": service_name}
+    try:
+        version = importlib.metadata.version("monkeybot")
+    except importlib.metadata.PackageNotFoundError:
+        version = None
+    if version:
+        attrs["service.version"] = version
+
+    try:
+        processor = _create_span_processor(exporter)
+        if processor is None:
+            return False
+        resource = resources.Resource.create(attrs)
+        provider = trace_sdk.TracerProvider(resource=resource)
+        provider.add_span_processor(processor)
+        trace_api.set_tracer_provider(provider)
+        propagate.set_global_textmap(w3c.TraceContextTextMapPropagator())
+        _state._enabled = True
+        logger.info("observability enabled (exporter=%s, service=%s)", exporter, service_name)
+        return True
+    except Exception as exc:
+        logger.warning("observability init failed: %s", exc)
+        return False
+
+
+def shutdown_observability() -> None:
+    if not _state._initialized:
+        return
+    try:
+        if _state._enabled:
+            trace_api = importlib.import_module("opentelemetry.trace")
+            provider = trace_api.get_tracer_provider()
+            shutdown = getattr(provider, "shutdown", None)
+            if callable(shutdown):
+                try:
+                    shutdown(timeout_millis=5000)
+                except TypeError:
+                    shutdown()
+    except Exception as exc:
+        logger.warning("observability shutdown failed: %s", exc)
+    finally:
+        _state._initialized = False
+        _state._enabled = False
+
+
+class _NoOpTracer:
+    def start_span(self, *_a: Any, **_k: Any) -> _NoOpSpan:
+        return _NoOpSpan()
+
+    def start_as_current_span(self, *_a: Any, **_k: Any) -> _NoOpSpan:
+        return _NoOpSpan()
+
+
+class _NoOpSpan:
+    def __enter__(self) -> _NoOpSpan:
+        return self
+
+    def __exit__(self, *_a: Any) -> None:
+        return None
+
+    def set_attribute(self, *_a: Any, **_k: Any) -> None:
+        return None
+
+    def record_exception(self, *_a: Any, **_k: Any) -> None:
+        return None
+
+    def set_status(self, *_a: Any, **_k: Any) -> None:
+        return None
+
+    def end(self) -> None:
+        return None
+
+    def is_recording(self) -> bool:
+        return False
+
+
+def get_tracer(name: str = "monkeybot") -> Any:
+    if not _state._enabled:
+        try:
+            from opentelemetry.trace import NoOpTracer
+
+            return NoOpTracer()
+        except ImportError:
+            return _NoOpTracer()
+    trace_api = importlib.import_module("opentelemetry.trace")
+    return trace_api.get_tracer(name)
+
+
+__all__ = [
+    "ObservingProvider",
+    "add_tool_hook_span_event",
+    "extract_traceparent",
+    "get_current_trace_id_hex_optional",
+    "instrument_fastapi_app",
+    "get_tracer",
+    "init_observability",
+    "inject_traceparent",
+    "is_observability_enabled",
+    "set_llm_io",
+    "set_llm_usage",
+    "set_run_output",
+    "set_summarize_turns",
+    "set_turn_io",
+    "set_turn_prompt_tokens",
+    "shutdown_observability",
+    "span_llm",
+    "span_run",
+    "span_subagent",
+    "span_summarize",
+    "span_tool",
+    "span_turn",
+    "truncate",
+]

monkeybot/observability/_state.py

@@ -0,0 +1,10 @@
+"""Process-local observability flags (no OpenTelemetry imports)."""
+
+from __future__ import annotations
+
+_initialized = False
+_enabled = False
+
+
+def is_observability_enabled() -> bool:
+    return _enabled

monkeybot/observability/instrumentation.py

@@ -0,0 +1,153 @@
+"""Extended instrumentation: provider wrapper, FastAPI, hook span events, trace id helpers."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import AsyncIterator, Sequence
+from typing import TYPE_CHECKING, Any, Literal, cast
+
+from monkeybot.core.context import TurnContext
+from monkeybot.core.llm.provider import (
+    Message,
+    Provider,
+    ProviderEvent,
+    ToolDef,
+    provider_count_input_tokens,
+    provider_stream,
+)
+from monkeybot.observability._state import is_observability_enabled
+from monkeybot.observability.spans import span_llm
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from fastapi import FastAPI
+
+HookPhase = Literal["pre_tool", "post_tool"]
+
+
+class ObservingProvider:
+    """Wraps a :class:`~monkeybot.core.llm.provider.Provider` for direct (non-loop) ``stream`` calls."""
+
+    def __init__(self, inner: Provider, *, turn_ctx: TurnContext | None = None) -> None:
+        self._inner = inner
+        self._turn_ctx = turn_ctx
+
+    @property
+    def name(self) -> str:
+        return self._inner.name
+
+    @property
+    def supports_streaming(self) -> bool:
+        return self._inner.supports_streaming
+
+    async def count_input_tokens(
+        self,
+        messages: Sequence[Message],
+        tools: Sequence[ToolDef],
+        *,
+        model: str,
+        thinking_budget: int | None = None,
+        vertex_google_search: bool = False,
+    ) -> int:
+        return await provider_count_input_tokens(
+            self._inner,
+            messages,
+            tools,
+            model=model,
+            thinking_budget=thinking_budget,
+            vertex_google_search=vertex_google_search,
+        )
+
+    def stream(
+        self,
+        messages: Sequence[Message],
+        tools: Sequence[ToolDef],
+        *,
+        model: str,
+        thinking_budget: int | None = None,
+        vertex_google_search: bool = False,
+    ) -> AsyncIterator[ProviderEvent]:
+        inner_stream = provider_stream(
+            self._inner,
+            messages,
+            tools,
+            model=model,
+            thinking_budget=thinking_budget,
+            vertex_google_search=vertex_google_search,
+        )
+        if not is_observability_enabled():
+            return inner_stream
+
+        async def _observed() -> AsyncIterator[ProviderEvent]:
+            ctx = self._correlation_ctx(model)
+            async with span_llm(
+                ctx=ctx,
+                model=model,
+                vertex_google_search=vertex_google_search,
+            ):
+                async for evt in inner_stream:
+                    yield evt
+
+        return _observed()
+
+    def _correlation_ctx(self, model: str) -> TurnContext:
+        if self._turn_ctx is not None:
+            return self._turn_ctx
+        return TurnContext(
+            thread_id="direct",
+            request_id="direct",
+            agent_md="",
+            memory_index=[],
+            skills=[],
+            tools=[],
+            user_id=None,
+            parent_run_id=None,
+            model=model,
+        )
+
+
+def instrument_fastapi_app(app: object) -> None:
+    """Instrument a FastAPI app with OTel HTTP server spans (idempotent, defensive)."""
+    if not is_observability_enabled():
+        return
+    if getattr(app, "_monkeybot_otel_fastapi_instrumented", False):
+        return
+    try:
+        from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+    except ImportError:
+        logger.warning(
+            "FastAPI OpenTelemetry instrumentation unavailable (install [observability] extra)"
+        )
+        return
+    FastAPIInstrumentor().instrument_app(cast("FastAPI", app))
+    cast(Any, app)._monkeybot_otel_fastapi_instrumented = True
+
+
+def get_current_trace_id_hex_optional() -> str | None:
+    """Return W3C trace id (32 lower-hex chars) for the current span, or ``None``."""
+    if not is_observability_enabled():
+        return None
+    from opentelemetry import trace
+    from opentelemetry.trace import format_trace_id
+
+    span = trace.get_current_span()
+    ctx = span.get_span_context()
+    if not ctx.is_valid:
+        return None
+    return format_trace_id(ctx.trace_id)
+
+
+def add_tool_hook_span_event(*, phase: HookPhase, tool_name: str) -> None:
+    """Record a hook lifecycle event on the active span (``monkeybot.tool`` when in loop)."""
+    if not is_observability_enabled():
+        return
+    from opentelemetry import trace
+
+    event_name = (
+        "monkeybot.hook.pre_tool" if phase == "pre_tool" else "monkeybot.hook.post_tool"
+    )
+    span = trace.get_current_span()
+    if not span.is_recording():
+        return
+    span.add_event(event_name, {"tool.name": tool_name})

monkeybot/observability/propagation.py

@@ -0,0 +1,65 @@
+"""W3C trace context propagation for subagent subprocesses."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from monkeybot.observability._state import is_observability_enabled
+
+if TYPE_CHECKING:
+    from opentelemetry.context import Context
+
+logger = logging.getLogger(__name__)
+
+
+def inject_traceparent(carrier: dict[str, str]) -> None:
+    """Inject W3C traceparent (and tracestate if present) from the current context.
+
+    No-op when observability is disabled, OTel is unavailable, or the current span
+    context is invalid. Never raises.
+    """
+    if not is_observability_enabled():
+        return
+    try:
+        from opentelemetry import propagate, trace
+    except ImportError:
+        return
+    try:
+        span_ctx = trace.get_current_span().get_span_context()
+        if not span_ctx.is_valid:
+            return
+        propagate.inject(carrier)
+    except Exception:
+        return
+
+
+def extract_traceparent(carrier: dict[str, str]) -> Context | None:
+    """Extract an OTel Context from carrier keys ``traceparent`` / ``tracestate``.
+
+    Returns None when the value is missing, malformed, or yields a non-valid span context.
+    Logs a WARNING once per bad value (include no user content). Never raises.
+    """
+    raw = carrier.get("traceparent")
+    if raw is None:
+        return None
+    if not isinstance(raw, str):
+        logger.warning("malformed traceparent in carrier: expected string")
+        return None
+    traceparent = raw.strip()
+    if not traceparent:
+        return None
+    try:
+        from opentelemetry import propagate, trace
+    except ImportError:
+        return None
+    try:
+        ctx = propagate.extract(carrier)
+        span_ctx = trace.get_current_span(ctx).get_span_context()
+        if not span_ctx.is_valid:
+            logger.warning("malformed traceparent: invalid span context after extract")
+            return None
+        return ctx
+    except Exception:
+        logger.warning("malformed traceparent: propagation extract failed")
+        return None