mainwave 0.1.0rc1__py3-none-any.whl

mainwave/__init__.py

@@ -0,0 +1,133 @@
+"""mainwave — Python SDK for the Mainwave AI observability platform.
+
+Public surface (stable):
+    init, instrument, register_framework
+    run, arun, start_run, span, flush, aflush, Run
+    langchain.CallbackHandler  (lazy via `from mainwave import langchain`)
+    testing  (test helpers)
+"""
+
+import logging as _logging
+
+from . import _config, _instrumentors
+from . import testing as testing  # noqa: F401  (re-export module for users)
+from ._config import init
+from ._errors import MainwaveConfigError, MainwaveError, MainwaveReinitConflict
+from ._flush import aflush, flush
+from ._runtime import Run, arun, run, span, start_run
+
+__version__ = _config._SDK_VERSION
+
+_log = _logging.getLogger("mainwave")
+
+
+def _active_injector():
+    """The attribute-injector processor that stamps framework names, or None if absent."""
+    from ._processors import MainwaveAttrInjector
+
+    procs = _config._state.mainwave_owned_processors
+    if procs and isinstance(procs[0], MainwaveAttrInjector):
+        return procs[0]
+    return None
+
+
+def _stamp_frameworks(injector, names: list[str]) -> None:
+    """Record framework names (+ detected versions) on the injector, deduped."""
+    if injector is None:
+        return
+    injector.record_frameworks(names)
+
+
+def register_framework(name: str) -> None:
+    """Stamp a framework name on spans without wiring an OpenInference instrumentor.
+
+    Use when the instrumentor for a framework is unavailable or broken but you
+    still want the framework to appear in telemetry (e.g. crewai>=1.x while
+    openinference-instrumentation-crewai is incompatible).
+    """
+    if not _config.is_initialized() or _config.is_disabled():
+        return
+    _stamp_frameworks(_active_injector(), [name])
+
+
+def instrument(name: str | list[str] | None = None) -> list[str]:
+    """Wire OpenInference instrumentors after `init()`.
+
+    Use after a late import (lazy import, Click subcommand, FastAPI lifespan)
+    where the target library wasn't loaded at init() time.
+
+    - `name=None`: re-walk the full registry, instrumenting any newly-importable
+      packages. Returns the names that were newly wired this call.
+    - `name="openai"`: instrument exactly that platform. Raises KeyError if
+      the name isn't registered.
+    - `name=["crewai", "openai"]`: instrument each named platform (handy when a
+      framework needs more than one, e.g. CrewAI's structure instrumentor plus
+      its LLM-provider instrumentor). Raises KeyError before wiring anything if
+      any name is unregistered. Returns the subset that actually wired.
+
+    Returns the names that wired — a name whose package is missing or whose
+    installed version conflicts with the instrumentor is silently dropped from
+    the result (with a WARN from the wiring layer), not raised.
+
+    No-op when the SDK is disabled or not yet initialized (logs at WARN
+    if not initialized).
+    """
+    if not _config.is_initialized():
+        _log.warning("mainwave.instrument() called before init(); ignoring")
+        return []
+    if _config.is_disabled():
+        return []
+    provider = _config.get_provider()
+    if provider is None:
+        return []
+
+    injector = _active_injector()
+
+    if name is None:
+        newly_wired = _instrumentors.wire_instrumentors(provider)
+        _stamp_frameworks(injector, newly_wired)
+        if newly_wired:
+            _log.info("mainwave: instrumented %s", ", ".join(newly_wired))
+        return newly_wired
+
+    raw = [name] if isinstance(name, str) else list(name)
+    # Accept hyphen or underscore spellings ("llama-index" == "llama_index"):
+    # the install extra is hyphenated, so users naturally type hyphens.
+    names = [_instrumentors.normalize_name(n) for n in raw]
+    # Validate the whole list up front so a typo in one name doesn't leave the
+    # crew half-instrumented (KeyError after some have already wired).
+    unknown = [n for n in names if n not in _instrumentors.known_names()]
+    if unknown:
+        raise KeyError(unknown[0] if len(unknown) == 1 else unknown)
+
+    wired = [n for n in names if _instrumentors.wire_one(n, provider)]
+    _stamp_frameworks(injector, wired)
+    if wired:
+        _log.info("mainwave: instrumented %s", ", ".join(wired))
+    not_wired = [n for n in names if n not in wired]
+    if not_wired:
+        _log.warning(
+            "mainwave: requested instrumentor(s) not wired (package not importable "
+            "or version conflict): %s",
+            ", ".join(not_wired),
+        )
+    return wired
+
+
+__all__ = [
+    "init",
+    "instrument",
+    "register_framework",
+    "run",
+    "arun",
+    "start_run",
+    "span",
+    "flush",
+    "aflush",
+    "Run",
+    "testing",
+    "MainwaveError",
+    "MainwaveConfigError",
+    "MainwaveReinitConflict",
+    "__version__",
+]

mainwave/_anthropic_thinking.py

@@ -0,0 +1,164 @@
+"""Recover Anthropic extended-thinking content the OpenInference instrumentor drops.
+
+`openinference-instrumentation-anthropic` enumerates a response's content blocks
+into OpenInference message attributes, but does not currently emit the `thinking`
+and `redacted_thinking` blocks — so an agent that calls Claude directly (or via
+`AnthropicBedrock`) emits no chain-of-thought, unlike the native LangChain
+handler (`mainwave.langchain`).
+
+This module monkeypatches the two extractor seams in the installed instrumentor
+so they additionally emit the model's reasoning as OpenInference reasoning
+content parts, using the same `message_content.type = "reasoning"` shape — so
+direct-Claude runs capture reasoning with no opt-in.
+
+The patch is coupled to instrumentor internals, so it is defensive: if the seams
+are missing (upstream renamed them) it logs and skips rather than crashing
+instrumentation. `tests/test_anthropic_thinking.py` asserts capture works against
+those seams, so a silent upstream shape change fails CI loudly."""
+
+from __future__ import annotations
+
+import logging
+import os
+from collections.abc import Iterator
+from typing import Any
+
+log = logging.getLogger("mainwave")
+
+# A `redacted_thinking` block's payload is encrypted; we surface that the model
+# reasoned but redacted it, never the opaque `data`.
+_REDACTED_MARKER = "[redacted thinking]"
+
+# Mirrors the langchain handler's content-capture knobs (own env var so the two
+# paths tune independently).
+_CONTENT_MAX_ENV = "MAINWAVE_ANTHROPIC_CONTENT_MAX"
+_DEFAULT_CONTENT_MAX = 8192
+
+# Sentinel marking a wrapped callable so re-wiring (instrument() called again)
+# never double-wraps and double-emits.
+_PATCHED = "_mainwave_thinking_patched"
+
+
+def _max_len() -> int:
+    try:
+        return max(0, int(os.environ.get(_CONTENT_MAX_ENV, _DEFAULT_CONTENT_MAX)))
+    except ValueError:
+        return _DEFAULT_CONTENT_MAX
+
+
+def _truncate(value: str) -> str:
+    limit = _max_len()
+    return value if len(value) <= limit else value[:limit] + "…[truncated]"
+
+
+def _block_attr(block: Any, key: str) -> Any:
+    """Read a content-block field whether the block is a pydantic object
+    (the SDK response shape) or a plain dict (defensive / test shape)."""
+    if isinstance(block, dict):
+        return block.get(key)
+    return getattr(block, key, None)
+
+
+def _reasoning_parts(content: Any) -> list[dict[str, str]]:
+    """Reasoning parts from a response's `content` list, in order, each
+    `{text, signature}`. `thinking` carries its text + Anthropic signature;
+    `redacted_thinking` carries a marker (never the encrypted payload). Anything
+    else (text, tool_use, …) is skipped. Returns [] when there is no reasoning."""
+    if not isinstance(content, (list, tuple)):
+        return []
+    parts: list[dict[str, str]] = []
+    for block in content:
+        block_type = _block_attr(block, "type")
+        if block_type == "thinking":
+            text = _block_attr(block, "thinking")
+            if not isinstance(text, str) or not text:
+                continue
+            sig = _block_attr(block, "signature")
+            parts.append({"text": text, "signature": sig if isinstance(sig, str) else ""})
+        elif block_type == "redacted_thinking":
+            parts.append({"text": _REDACTED_MARKER, "signature": ""})
+    return parts
+
+
+def reasoning_attributes(content: Any) -> Iterator[tuple[str, Any]]:
+    """OpenInference reasoning content-part attributes for a response's content.
+
+    Emits `llm.output_messages.0.message.contents.{j}.message_content.{type,text}`
+    (plus `.signature` when present), per the OpenInference reasoning content-part
+    shape. The j index runs over reasoning parts only; the instrumentor's own text
+    branch emits the flat `message.content`, so the two never collide."""
+    for j, part in enumerate(_reasoning_parts(content)):
+        prefix = f"llm.output_messages.0.message.contents.{j}.message_content"
+        yield f"{prefix}.type", "reasoning"
+        yield f"{prefix}.text", _truncate(part["text"])
+        if part["signature"]:
+            yield f"{prefix}.signature", part["signature"]
+
+
+def patch_anthropic_thinking() -> bool:
+    """Patch the installed Anthropic instrumentor to emit reasoning content parts.
+
+    Wraps the non-streaming output extractor and the streaming snapshot extractor
+    so both append reasoning parts after the upstream attributes. Idempotent and
+    defensive: returns False (and logs) if the package or a seam is missing, or if
+    already patched — never raises into the instrumentation path."""
+    try:
+        from openinference.instrumentation.anthropic import _stream, _wrappers
+    except ImportError:
+        return False
+
+    ok = True
+    ok &= _wrap_output_messages(_wrappers)
+    ok &= _wrap_stream_extractor(_stream)
+    return ok
+
+
+def _wrap_output_messages(wrappers: Any) -> bool:
+    """Wrap `_wrappers._get_output_messages` (the non-streaming seam). Both the
+    sync and async message wrappers resolve it as a module global at call time,
+    so one patch covers both."""
+    orig = getattr(wrappers, "_get_output_messages", None)
+    if orig is None:
+        log.error(
+            "mainwave: openinference anthropic instrumentor has no "
+            "_get_output_messages; cannot capture thinking. Update the patch."
+        )
+        return False
+    if getattr(orig, _PATCHED, False):
+        return True
+
+    def patched(response: Any) -> Iterator[tuple[str, Any]]:
+        yield from orig(response)
+        yield from reasoning_attributes(getattr(response, "content", None))
+
+    setattr(patched, _PATCHED, True)
+    patched._mainwave_orig = orig  # type: ignore[attr-defined]
+    wrappers._get_output_messages = patched
+    return True
+
+
+def _wrap_stream_extractor(stream: Any) -> bool:
+    """Wrap `_stream._MessageExtractor.get_attributes` (the streaming seam),
+    which only handles text/tool_use blocks on the accumulated snapshot."""
+    extractor = getattr(stream, "_MessageExtractor", None)
+    orig = getattr(extractor, "get_attributes", None)
+    if extractor is None or orig is None:
+        log.error(
+            "mainwave: openinference anthropic instrumentor has no "
+            "_MessageExtractor.get_attributes; cannot capture streamed thinking. "
+            "Update the patch."
+        )
+        return False
+    if getattr(orig, _PATCHED, False):
+        return True
+
+    def patched(self: Any) -> Iterator[tuple[str, Any]]:
+        yield from orig(self)
+        snapshot = getattr(self, "_snapshot", None)
+        if snapshot is not None:
+            yield from reasoning_attributes(getattr(snapshot, "content", None))
+
+    setattr(patched, _PATCHED, True)
+    patched._mainwave_orig = orig  # type: ignore[attr-defined]
+    extractor.get_attributes = patched
+    return True

mainwave/_config.py

@@ -0,0 +1,302 @@
+"""init() — read config, build TracerProvider (or coexist), wire processors.
+
+Config-equality contract for idempotent re-init lives here too."""
+
+import importlib.metadata
+import logging
+import os
+import threading
+import uuid
+from dataclasses import dataclass
+from typing import Any
+from urllib.parse import urlsplit, urlunsplit
+
+from opentelemetry import trace
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+
+from ._errors import MainwaveConfigError, MainwaveReinitConflict
+
+log = logging.getLogger("mainwave")
+
+_DEFAULT_COLLECTOR = "https://collect.mainwave.ai"
+_DEFAULT_ENV = "production"
+_UCID_NAMESPACE = uuid.UUID("a1b2c3d4-e5f6-7890-abcd-ef1234567890")
+_OTLP_TRACES_PATH = "/v1/traces"
+
+try:
+    _SDK_VERSION = importlib.metadata.version("mainwave")
+except importlib.metadata.PackageNotFoundError:
+    _SDK_VERSION = "0.0.0+unknown"
+
+
+@dataclass(frozen=True)
+class _Config:
+    api_key: str
+    use_case_id: str
+    ucid_source: str  # "explicit" | "derived" | "" (disabled)
+    collector_url: str
+    service_name: str
+    deployment_environment: str
+    disabled: bool
+
+
+@dataclass
+class _State:
+    config: _Config | None = None
+    provider: TracerProvider | None = None
+    coexist_mode: bool = False
+    mainwave_owned_processors: list[Any] | None = None
+    # The MainwaveExporter inside the BatchSpanProcessor, kept so flush() can
+    # drain its trace-end buffers (the BSP's force_flush doesn't reach it).
+    mainwave_exporter: Any = None
+    init_lock: threading.Lock = None  # type: ignore[assignment]
+
+    def __post_init__(self) -> None:
+        if self.init_lock is None:
+            self.init_lock = threading.Lock()
+
+
+_state = _State()
+
+
+def _normalize_collector_url(url: str) -> str:
+    """Lowercase host, strip trailing slash on path, drop default ports."""
+    parts = urlsplit(url)
+    if not parts.scheme or not parts.netloc:
+        raise MainwaveConfigError(f"collector_url is unparseable: {url!r}")
+    host = parts.hostname or ""
+    port = parts.port
+    if (parts.scheme == "http" and port == 80) or (parts.scheme == "https" and port == 443):
+        port = None
+    netloc = host.lower()
+    if port is not None:
+        netloc = f"{netloc}:{port}"
+    if parts.username:
+        userinfo = parts.username
+        if parts.password:
+            userinfo = f"{userinfo}:{parts.password}"
+        netloc = f"{userinfo}@{netloc}"
+    path = parts.path.rstrip("/")
+    return urlunsplit((parts.scheme.lower(), netloc, path, parts.query, parts.fragment))
+
+
+def _coerce_bool(raw: str | None, *, name: str) -> bool:
+    if raw is None:
+        return False
+    if raw.lower() in {"1", "true", "yes", "on"}:
+        return True
+    if raw.lower() in {"0", "false", "no", "off", ""}:
+        return False
+    raise MainwaveConfigError(f"{name}={raw!r} is not a valid boolean")
+
+
+def _build_config(
+    *,
+    api_key: str | None,
+    use_case_id: str | None,
+    collector_url: str | None,
+    service_name: str | None,
+    deployment_environment: str | None,
+    disabled: bool | None,
+) -> _Config:
+    api_key = api_key if api_key is not None else os.environ.get("MAINWAVE_API_KEY")
+    use_case_id = use_case_id if use_case_id is not None else os.environ.get("MAINWAVE_USE_CASE_ID")
+    collector_url = (
+        collector_url
+        if collector_url is not None
+        else os.environ.get("MAINWAVE_COLLECTOR_URL", _DEFAULT_COLLECTOR)
+    )
+    service_name = (
+        service_name
+        if service_name is not None
+        else os.environ.get("OTEL_SERVICE_NAME", "mainwave-app")
+    )
+    deployment_environment = (
+        deployment_environment
+        if deployment_environment is not None
+        else os.environ.get("MAINWAVE_ENVIRONMENT", _DEFAULT_ENV)
+    )
+    if disabled is None:
+        disabled = _coerce_bool(os.environ.get("MAINWAVE_DISABLED"), name="MAINWAVE_DISABLED")
+
+    ucid_source = ""
+    if not disabled:
+        if not api_key:
+            raise MainwaveConfigError("api_key is required (pass api_key= or set MAINWAVE_API_KEY)")
+        if not use_case_id:
+            use_case_id = str(uuid.uuid5(_UCID_NAMESPACE, f"{api_key}:{service_name}"))
+            ucid_source = "derived"
+            log.info(
+                "mainwave: use_case_id not set — derived stable ucid %s from api_key + service_name",  # noqa: E501
+                use_case_id,
+            )
+        else:
+            ucid_source = "explicit"
+
+    normalized_url = _normalize_collector_url(collector_url)
+
+    return _Config(
+        api_key=api_key or "",
+        use_case_id=use_case_id or "",
+        ucid_source=ucid_source,
+        collector_url=normalized_url,
+        service_name=service_name,
+        deployment_environment=deployment_environment,
+        disabled=disabled,
+    )
+
+
+def _build_resource(cfg: _Config) -> Resource:
+    return Resource(
+        attributes={
+            "service.name": cfg.service_name,
+            "deployment.environment": cfg.deployment_environment,
+            "mainwave.ucid": cfg.use_case_id,
+            "mainwave.ucid_source": cfg.ucid_source,
+            "mainwave.sdk.version": _SDK_VERSION,
+            "mainwave.sdk.language": "python",
+        }
+    )
+
+
+def _is_real_provider(p: trace.TracerProvider) -> bool:
+    """True if the global provider is a real SDK TracerProvider (not the API
+    no-op default)."""
+    return isinstance(p, TracerProvider)
+
+
+def init(
+    *,
+    api_key: str | None = None,
+    use_case_id: str | None = None,
+    collector_url: str | None = None,
+    service_name: str | None = None,
+    deployment_environment: str | None = None,
+    disabled: bool | None = None,
+    skip_instrumentors: list[str] | None = None,
+) -> None:
+    """Initialize the SDK. Idempotent on identical config; raises
+    MainwaveReinitConflict on diverging re-init.
+
+    Bad config (missing api_key, unparseable URL, malformed env bool) raises
+    MainwaveConfigError. Once initialized, runtime errors are swallowed
+    per the 'cannot break the host' guarantee.
+
+    `skip_instrumentors` names OpenInference instrumentors to leave unwired even
+    when their package is installed — for a framework you drive through an
+    explicit handler instead (e.g. `["langchain"]` when using
+    `mainwave.langchain.CallbackHandler`), so the two paths don't double-count.
+    It only affects the one-time wiring on first init, so it is not part of the
+    re-init equality contract.
+    """
+    cfg = _build_config(
+        api_key=api_key,
+        use_case_id=use_case_id,
+        collector_url=collector_url,
+        service_name=service_name,
+        deployment_environment=deployment_environment,
+        disabled=disabled,
+    )
+
+    with _state.init_lock:
+        if _state.config is not None:
+            if _state.config == cfg:
+                return
+            raise MainwaveReinitConflict(
+                "mainwave.init() was already called with a different config"
+            )
+
+        if cfg.disabled:
+            _state.config = cfg
+            log.info("mainwave: disabled; no spans will be exported")
+            return
+
+        from ._exporter import build_exporter
+        from ._instrumentors import wire_instrumentors
+        from ._processors import MainwaveAttrInjector
+
+        existing = trace.get_tracer_provider()
+        coexist = _is_real_provider(existing)
+
+        injector = MainwaveAttrInjector(
+            use_case_id=cfg.use_case_id, ucid_source=cfg.ucid_source, coexist=coexist
+        )
+        exporter_processor, mainwave_exporter = build_exporter(cfg)
+
+        if coexist:
+            log.info("mainwave: detected existing TracerProvider; running in coexist mode")
+            existing.add_span_processor(injector)  # type: ignore[attr-defined]
+            existing.add_span_processor(exporter_processor)
+            _state.provider = existing  # type: ignore[assignment]
+            _state.coexist_mode = True
+        else:
+            resource = _build_resource(cfg)
+            provider = TracerProvider(resource=resource)
+            provider.add_span_processor(injector)
+            provider.add_span_processor(exporter_processor)
+            trace.set_tracer_provider(provider)
+            _state.provider = provider
+            _state.coexist_mode = False
+
+        _state.mainwave_owned_processors = [injector, exporter_processor]
+        _state.mainwave_exporter = mainwave_exporter
+        _state.config = cfg
+
+        wired = wire_instrumentors(_state.provider, skip=skip_instrumentors)
+        if wired:
+            injector.record_frameworks(wired)
+            log.info("mainwave: instrumented %s", ", ".join(wired))
+
+        from ._flush import register_atexit
+
+        register_atexit()
+
+
+def get_config() -> _Config | None:
+    return _state.config
+
+
+def get_provider() -> TracerProvider | None:
+    return _state.provider
+
+
+def get_mainwave_exporter() -> Any:
+    """The MainwaveExporter for the active init, or None (uninitialized,
+    disabled, or an in-memory test provider). flush() drains its buffers."""
+    return _state.mainwave_exporter
+
+
+def is_initialized() -> bool:
+    return _state.config is not None
+
+
+def is_disabled() -> bool:
+    return _state.config is not None and _state.config.disabled
+
+
+def is_coexist_mode() -> bool:
+    return _state.coexist_mode
+
+
+def reset_for_testing() -> None:
+    """Tear down state. For test fixtures only."""
+    with _state.init_lock:
+        if _state.provider is not None and not _state.coexist_mode:
+            try:
+                _state.provider.shutdown()
+            except Exception:
+                pass
+        _state.config = None
+        _state.provider = None
+        _state.coexist_mode = False
+        _state.mainwave_owned_processors = None
+        _state.mainwave_exporter = None
+        trace._TRACER_PROVIDER = None  # type: ignore[attr-defined]
+        trace._PROXY_TRACER_PROVIDER._real_tracer_provider = None  # type: ignore[attr-defined]
+        # Clear a possibly-leaked run contextvar: a test that intentionally leaks
+        # a start_run (no .end()) would otherwise leave _current_run set, making
+        # the next test's run()/start_run() look nested and skip its wrapper span.
+        from ._runtime import _current_run
+
+        _current_run.set(None)

mainwave/_errors.py

@@ -0,0 +1,15 @@
+"""Public exception types raised by the SDK."""
+
+
+class MainwaveError(Exception):
+    """Base class for all SDK-raised exceptions."""
+
+
+class MainwaveConfigError(MainwaveError):
+    """Raised at init() when configuration is invalid (missing api key,
+    unparseable collector URL, malformed env value)."""
+
+
+class MainwaveReinitConflict(MainwaveError):
+    """Raised when init() is called a second time with a config that diverges
+    from the first call. Re-init with identical config is a no-op."""