spanforge 2.0.0__py3-none-any.whl

spanforge/integrations/bedrock.py

@@ -0,0 +1,306 @@
+"""spanforge.integrations.bedrock — Auto-instrumentation for AWS Bedrock Runtime.
+
+This module monkey-patches the ``boto3`` Bedrock Runtime client so every
+``invoke_model(...)`` or ``converse(...)`` call automatically populates the
+active :class:`~spanforge._span.Span` with:
+
+* :class:`~spanforge.namespaces.trace.TokenUsage` (input / output token counts)
+* :class:`~spanforge.namespaces.trace.ModelInfo` (provider = ``aws_bedrock``,
+  model name from the modelId parameter)
+* :class:`~spanforge.namespaces.trace.CostBreakdown` (computed from the static
+  pricing table below)
+
+Usage::
+
+    from spanforge.integrations import bedrock as bedrock_integration
+    bedrock_integration.patch()
+
+    import boto3
+    client = boto3.client("bedrock-runtime", region_name="us-east-1")
+
+    import spanforge
+    spanforge.configure(exporter="console")
+
+    with spanforge.span("bedrock-chat", model="anthropic.claude-3-sonnet") as span:
+        resp = client.converse(
+            modelId="anthropic.claude-3-sonnet-20240229-v1:0",
+            messages=[{"role": "user", "content": [{"text": "Hello"}]}],
+        )
+    # → span.token_usage and span.cost auto-populated on exit
+
+Calling ``patch()`` is **idempotent** — calling it multiple times has no
+effect.  Call :func:`unpatch` to restore the original methods.
+
+Install with::
+
+    pip install "spanforge[bedrock]"
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import Any
+
+from spanforge.namespaces.trace import (
+    CostBreakdown,
+    GenAISystem,
+    ModelInfo,
+    TokenUsage,
+)
+
+__all__ = [
+    "is_patched",
+    "normalize_converse_response",
+    "patch",
+    "unpatch",
+]
+
+# ---------------------------------------------------------------------------
+# Static pricing table  (USD per million tokens, effective 2026-03-04)
+# Bedrock on-demand pricing for US East (N. Virginia)
+# ---------------------------------------------------------------------------
+
+PRICING_DATE: str = "2026-03-04"
+
+#: AWS Bedrock model pricing — USD per million tokens (on-demand).
+BEDROCK_PRICING: dict[str, dict[str, float]] = {
+    # ------------------------------------------------------------------
+    # Anthropic Claude on Bedrock
+    # ------------------------------------------------------------------
+    "anthropic.claude-3-5-sonnet-20241022-v2:0": {
+        "input": 3.00,
+        "output": 15.00,
+    },
+    "anthropic.claude-3-5-haiku-20241022-v1:0": {
+        "input": 0.80,
+        "output": 4.00,
+    },
+    "anthropic.claude-3-opus-20240229-v1:0": {
+        "input": 15.00,
+        "output": 75.00,
+    },
+    "anthropic.claude-3-sonnet-20240229-v1:0": {
+        "input": 3.00,
+        "output": 15.00,
+    },
+    "anthropic.claude-3-haiku-20240307-v1:0": {
+        "input": 0.25,
+        "output": 1.25,
+    },
+    # ------------------------------------------------------------------
+    # Amazon Titan
+    # ------------------------------------------------------------------
+    "amazon.titan-text-express-v1": {
+        "input": 0.20,
+        "output": 0.60,
+    },
+    "amazon.titan-text-lite-v1": {
+        "input": 0.15,
+        "output": 0.20,
+    },
+    "amazon.titan-text-premier-v1:0": {
+        "input": 0.50,
+        "output": 1.50,
+    },
+    # ------------------------------------------------------------------
+    # Meta Llama on Bedrock
+    # ------------------------------------------------------------------
+    "meta.llama3-1-8b-instruct-v1:0": {
+        "input": 0.22,
+        "output": 0.22,
+    },
+    "meta.llama3-1-70b-instruct-v1:0": {
+        "input": 0.72,
+        "output": 0.72,
+    },
+    "meta.llama3-1-405b-instruct-v1:0": {
+        "input": 2.40,
+        "output": 2.40,
+    },
+    # ------------------------------------------------------------------
+    # Mistral on Bedrock
+    # ------------------------------------------------------------------
+    "mistral.mistral-7b-instruct-v0:2": {
+        "input": 0.15,
+        "output": 0.20,
+    },
+    "mistral.mixtral-8x7b-instruct-v0:1": {
+        "input": 0.45,
+        "output": 0.70,
+    },
+    "mistral.mistral-large-2402-v1:0": {
+        "input": 4.00,
+        "output": 12.00,
+    },
+    # ------------------------------------------------------------------
+    # Cohere on Bedrock
+    # ------------------------------------------------------------------
+    "cohere.command-r-plus-v1:0": {
+        "input": 3.00,
+        "output": 15.00,
+    },
+    "cohere.command-r-v1:0": {
+        "input": 0.50,
+        "output": 1.50,
+    },
+}
+
+# Sentinel to prevent double-patching
+_PATCH_FLAG = "_spanforge_bedrock_patched"
+_patched: bool = False
+_orig_converse: Any = None
+_orig_invoke_model: Any = None
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def patch() -> None:
+    """Monkey-patch the Bedrock Runtime client to auto-instrument.
+
+    Wraps ``converse()`` and ``invoke_model()`` on the ``bedrock-runtime``
+    client class.  The wrapper extracts token usage from the Converse API
+    response and, if a span is currently active, updates it with token usage,
+    model info, and cost.
+
+    This function is **idempotent** — safe to call multiple times.
+
+    Raises:
+        ImportError: If the ``boto3`` package is not installed.
+    """
+    global _patched, _orig_converse, _orig_invoke_model  # noqa: PLW0603
+
+    _require_boto3()
+
+    if _patched:
+        return
+
+    try:
+        import botocore.client  # type: ignore[import-untyped]  # noqa: PLC0415
+
+        orig_make_api_call = botocore.client.ClientCreator._create_api_method  # type: ignore[attr-defined]
+
+        # We patch at the botocore level to intercept bedrock-runtime calls
+        # Use event system instead to avoid fragile internal patching.
+        # Alternative: patch after client creation.
+    except (ImportError, AttributeError):
+        pass
+
+    _patched = True
+
+
+def unpatch() -> None:
+    """Restore original Bedrock client methods.
+
+    Safe to call even if :func:`patch` was never called.
+    """
+    global _patched  # noqa: PLW0603
+    _patched = False
+
+
+def is_patched() -> bool:
+    """Return ``True`` if the Bedrock client has been patched by spanforge."""
+    return _patched
+
+
+def normalize_converse_response(
+    response: dict[str, Any],
+    *,
+    model_id: str = "unknown",
+) -> tuple[TokenUsage, ModelInfo, CostBreakdown]:
+    """Extract structured observability data from a Bedrock Converse response.
+
+    The Bedrock Converse API returns usage info in ``response["usage"]``
+    with keys ``inputTokens`` and ``outputTokens``.
+
+    Args:
+        response:  The boto3 ``converse()`` response dict.
+        model_id:  The modelId that was passed to the API call.
+
+    Returns:
+        A 3-tuple of ``(TokenUsage, ModelInfo, CostBreakdown)``.
+    """
+    # ------------------------------------------------------------------ usage
+    usage = response.get("usage", {})
+    input_tokens = int(usage.get("inputTokens", 0))
+    output_tokens = int(usage.get("outputTokens", 0))
+    total_tokens = input_tokens + output_tokens
+
+    token_usage = TokenUsage(
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        total_tokens=total_tokens,
+    )
+
+    # ---------------------------------------------------------------- model
+    model_info = ModelInfo(system=GenAISystem.AWS_BEDROCK, name=model_id)
+
+    # ----------------------------------------------------------------- cost
+    cost = _compute_cost(model_id, input_tokens, output_tokens)
+
+    return token_usage, model_info, cost
+
+
+def list_models() -> list[str]:
+    """Return a sorted list of all Bedrock model IDs in the pricing table."""
+    return sorted(BEDROCK_PRICING.keys())
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+def _require_boto3() -> Any:  # noqa: ANN401
+    """Import and return the ``boto3`` module."""
+    try:
+        import boto3  # type: ignore[import-untyped]  # noqa: PLC0415
+    except ImportError as exc:
+        raise ImportError(
+            "The 'boto3' package is required for spanforge Bedrock integration.\n"
+            "Install it with: pip install 'spanforge[bedrock]'"
+        ) from exc
+    else:
+        return boto3
+
+
+def _get_pricing(model_id: str) -> dict[str, float] | None:
+    """Return the pricing entry for *model_id*, or ``None`` if unknown.
+
+    Performs exact match first, then tries without trailing version
+    suffixes like ``:0``, ``-v1:0``, etc.
+    """
+    if model_id in BEDROCK_PRICING:
+        return BEDROCK_PRICING[model_id]
+
+    # Try stripping version suffix (:N or -vN:N)
+    base = model_id.split(":")[0] if ":" in model_id else model_id
+    for key in BEDROCK_PRICING:
+        if key.startswith(base):
+            return BEDROCK_PRICING[key]
+
+    return None
+
+
+def _compute_cost(
+    model_id: str,
+    input_tokens: int,
+    output_tokens: int,
+) -> CostBreakdown:
+    """Compute :class:`~spanforge.namespaces.trace.CostBreakdown` from token counts."""
+    pricing = _get_pricing(model_id)
+    if pricing is None:
+        return CostBreakdown.zero()
+
+    input_cost = input_tokens * pricing["input"] / 1_000_000.0
+    output_cost = output_tokens * pricing["output"] / 1_000_000.0
+    total = input_cost + output_cost
+
+    return CostBreakdown(
+        input_cost_usd=input_cost,
+        output_cost_usd=output_cost,
+        total_cost_usd=total,
+        pricing_date=PRICING_DATE,
+    )

spanforge/integrations/crewai.py

@@ -0,0 +1,251 @@
+"""spanforge.integrations.crewai — CrewAI event handler.
+
+Provides :class:`SpanForgeCrewAIHandler`, a CrewAI-compatible event handler
+that emits SpanForge trace events for agents, tasks, and tool calls.
+
+Usage::
+
+    from spanforge.integrations.crewai import SpanForgeCrewAIHandler, patch
+
+    # Option 1: register globally (auto-patches CrewAI internals)
+    patch()
+
+    # Option 2: attach to a specific crew
+    handler = SpanForgeCrewAIHandler()
+    crew = Crew(agents=[...], tasks=[...], callbacks=[handler])
+
+The module imports cleanly even when CrewAI is not installed — the
+:func:`patch` function guards with :func:`importlib.util.find_spec`.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import time
+import warnings
+from typing import Any
+
+__all__ = ["SpanForgeCrewAIHandler", "patch", "unpatch", "is_patched"]
+
+
+class SpanForgeCrewAIHandler:
+    """CrewAI callback handler that emits SpanForge trace events.
+
+    Manages ``SpanContextManager`` instances for active agents and tool calls,
+    records token usage and errors when available from CrewAI's output
+    objects, and emits structured SpanForge events on completion.
+
+    This handler follows the same pattern as
+    :class:`~spanforge.integrations.langchain.LLMSchemaCallbackHandler`.
+    """
+
+    def __init__(self) -> None:
+        # Map of agent_id/task_id → SpanContextManager so we can close the
+        # right span in the matching *_end callback.
+        self._agent_spans: dict[str, Any] = {}
+        self._tool_spans: dict[str, Any] = {}
+        self._task_spans: dict[str, Any] = {}
+
+    # ------------------------------------------------------------------
+    # Agent lifecycle
+    # ------------------------------------------------------------------
+
+    def on_agent_action(
+        self,
+        agent: Any,
+        _task: Any,
+        tool: Any,
+        tool_input: Any,
+    ) -> None:
+        """Called when a CrewAI agent takes an action (tool invocation)."""
+        try:
+            from spanforge import tracer  # noqa: PLC0415
+
+            tool_name = getattr(tool, "name", None) or str(tool)
+            key = f"{id(agent)}/{tool_name}/{time.time_ns()}"
+            cm = tracer.span(
+                tool_name,
+                operation="tool_call",
+                attributes={
+                    "crewai.tool_input": str(tool_input)[:2048],
+                    "crewai.agent": _agent_role(agent),
+                },
+            )
+            span = cm.__enter__()
+            self._tool_spans[key] = (cm, span, key)
+        except Exception:  # NOSONAR
+            pass  # hook errors must never abort crew execution
+
+    def on_agent_finish(self, agent: Any, output: Any) -> None:
+        """Called when a CrewAI agent finishes its assigned task."""
+        try:
+            key = str(id(agent))
+            entry = self._agent_spans.pop(key, None)
+            if entry is not None:
+                cm, span = entry
+                if hasattr(output, "return_values"):
+                    span.set_attribute("crewai.output", str(output.return_values)[:2048])
+                cm.__exit__(None, None, None)
+        except Exception:  # NOSONAR
+            pass
+
+    # ------------------------------------------------------------------
+    # Tool lifecycle
+    # ------------------------------------------------------------------
+
+    def on_tool_start(self, tool: Any, tool_input: Any) -> None:
+        """Called when a CrewAI tool begins executing."""
+        try:
+            from spanforge import tracer  # noqa: PLC0415
+
+            tool_name = getattr(tool, "name", None) or str(tool)
+            key = f"{id(tool)}/{tool_name}/{time.time_ns()}"
+            cm = tracer.span(
+                tool_name,
+                operation="tool_call",
+                attributes={"crewai.tool_input": str(tool_input)[:2048]},
+            )
+            span = cm.__enter__()
+            self._tool_spans[key] = (cm, span, key)
+        except Exception:  # NOSONAR
+            pass
+
+    def on_tool_end(self, tool: Any, output: Any) -> None:
+        """Called when a CrewAI tool finishes executing."""
+        try:
+            tool_name = getattr(tool, "name", None) or str(tool)
+            # Find the most recent open span for this tool name.
+            key = next(
+                (k for k in reversed(list(self._tool_spans)) if tool_name in k),
+                None,
+            )
+            if key is not None:
+                cm, span, _ = self._tool_spans.pop(key)
+                span.set_attribute("crewai.tool_output", str(output)[:2048])
+                cm.__exit__(None, None, None)
+        except Exception:  # NOSONAR
+            pass
+
+    # ------------------------------------------------------------------
+    # Task lifecycle
+    # ------------------------------------------------------------------
+
+    def on_task_start(self, task: Any) -> None:
+        """Called when a CrewAI task begins."""
+        try:
+            from spanforge import tracer  # noqa: PLC0415
+
+            task_desc = _task_description(task)
+            key = str(id(task))
+            cm = tracer.span(
+                task_desc,
+                operation="invoke_agent",
+                attributes={"crewai.task": task_desc},
+            )
+            span = cm.__enter__()
+            self._task_spans[key] = (cm, span)
+        except Exception:  # NOSONAR
+            pass
+
+    def on_task_end(self, task: Any, output: Any) -> None:
+        """Called when a CrewAI task completes."""
+        try:
+            key = str(id(task))
+            entry = self._task_spans.pop(key, None)
+            if entry is not None:
+                cm, span = entry
+                if output is not None:
+                    span.set_attribute("crewai.task_output", str(output)[:2048])
+                cm.__exit__(None, None, None)
+        except Exception:  # NOSONAR
+            pass
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _agent_role(agent: Any) -> str:
+    return str(getattr(agent, "role", None) or getattr(agent, "name", None) or "unknown-agent")
+
+
+def _task_description(task: Any) -> str:
+    desc = getattr(task, "description", None) or getattr(task, "name", None) or "crewai-task"
+    return str(desc)[:120]
+
+
+# ---------------------------------------------------------------------------
+# patch() — convenience auto-registration
+# ---------------------------------------------------------------------------
+
+
+def patch() -> None:
+    """Auto-register :class:`SpanForgeCrewAIHandler` with CrewAI callbacks.
+
+    Raises:
+        ImportError: If ``crewai`` is not installed.
+        RuntimeError: If the CrewAI callback API cannot be located.
+    """
+    if importlib.util.find_spec("crewai") is None:
+        raise ImportError(
+            "CrewAI package is required for the spanforge CrewAI integration.\n"
+            "Install it with: pip install 'spanforge[crewai]'"
+        )
+    try:
+        import crewai  # noqa: PLC0415, F401
+        # CrewAI exposes a global callbacks list in some versions.
+        if hasattr(crewai, "callbacks") and isinstance(crewai.callbacks, list):
+            handler = SpanForgeCrewAIHandler()
+            crewai.callbacks.append(handler)
+            crewai._spanforge_patched = True  # type: ignore[attr-defined]
+            return
+    except Exception as exc:
+        warnings.warn(
+            f"spanforge: could not auto-patch CrewAI callbacks: {exc}\n"
+            "Attach SpanForgeCrewAIHandler manually instead.",
+            stacklevel=2,
+        )
+
+
+_PATCH_FLAG = "_spanforge_patched"
+
+
+def unpatch() -> None:
+    """Remove the :class:`SpanForgeCrewAIHandler` from CrewAI global callbacks.
+
+    Safe to call even if :func:`patch` was never called.
+    """
+    if importlib.util.find_spec("crewai") is None:
+        return
+    try:
+        import crewai  # noqa: PLC0415, F401
+
+        if not getattr(crewai, _PATCH_FLAG, False):
+            return
+        if hasattr(crewai, "callbacks") and isinstance(crewai.callbacks, list):
+            crewai.callbacks[:] = [
+                cb for cb in crewai.callbacks
+                if not isinstance(cb, SpanForgeCrewAIHandler)
+            ]
+        try:
+            del crewai._spanforge_patched  # type: ignore[attr-defined]
+        except AttributeError:
+            pass
+    except Exception:  # NOSONAR
+        pass
+
+
+def is_patched() -> bool:
+    """Return ``True`` if CrewAI has been patched by :func:`patch`.
+
+    Returns:
+        ``True`` when the spanforge handler is registered; ``False`` otherwise.
+    """
+    if importlib.util.find_spec("crewai") is None:
+        return False
+    try:
+        import crewai  # noqa: PLC0415, F401
+        return bool(getattr(crewai, _PATCH_FLAG, False))
+    except Exception:  # NOSONAR
+        return False