spanforge 2.0.0__py3-none-any.whl

spanforge/integrations/llamaindex.py

@@ -0,0 +1,370 @@
+"""spanforge.integrations.llamaindex — LlamaIndex event handler.
+
+Provides :class:`LLMSchemaEventHandler`, a LlamaIndex-compatible callback
+handler that records LLM, tool-call, and query activity as SpanForge events.
+
+Usage::
+
+    from spanforge.integrations.llamaindex import LLMSchemaEventHandler
+
+    handler = LLMSchemaEventHandler(source="rag-app", org_id="org-2")
+
+    Settings.callback_manager = CallbackManager([handler])
+    # or inject via the query engine constructor
+
+    for event in handler.events:
+        print(event.to_json())
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import Any
+
+from spanforge.event import Event
+from spanforge.ulid import generate as gen_ulid
+
+__all__ = [
+    "LLMSchemaEventHandler",
+    "is_patched",
+    "unpatch",
+]
+
+# ---------------------------------------------------------------------------
+# Module resolver
+# ---------------------------------------------------------------------------
+
+
+def _require_llamaindex() -> Any:  # noqa: ANN401
+    """Return the LlamaIndex callbacks module.
+
+    Tries ``llama_index.core.callbacks`` first (preferred, modern API), then
+    falls back to the legacy ``llama_index.callbacks`` module.
+
+    Raises:
+        ImportError: If neither ``llama_index.core`` nor ``llama_index`` is installed.
+    """
+    # Try modern llama_index.core first.
+    try:
+        import sys  # noqa: PLC0415
+
+        import llama_index.core  # noqa: PLC0415
+        import llama_index.core.callbacks  # noqa: PLC0415
+        return sys.modules["llama_index.core.callbacks"]
+    except ImportError:
+        pass
+    # Fall back to legacy llama_index package.
+    try:
+        import sys  # noqa: PLC0415
+
+        import llama_index  # noqa: PLC0415
+        import llama_index.callbacks  # noqa: PLC0415, F401
+        return sys.modules["llama_index.callbacks"]
+    except ImportError:
+        pass
+    raise ImportError(
+        "LlamaIndex package is required for the spanforge LlamaIndex integration.\n"
+        "Install it with: pip install 'spanforge[llamaindex]'"
+    )
+
+
+# ---------------------------------------------------------------------------
+# Event type mapping
+# ---------------------------------------------------------------------------
+
+#: Maps LlamaIndex CBEventType string → (start_event_type, end_event_type)
+_CB_TYPE_MAP: dict[str, tuple[str, str]] = {
+    "LLM": ("llm.trace.span.started", "llm.trace.span.completed"),
+    "FUNCTION_CALL": ("llm.trace.tool_call.started", "llm.trace.tool_call.completed"),
+    "QUERY": ("llm.trace.query.started", "llm.trace.query.completed"),
+    "RETRIEVE": ("llm.trace.retrieve.started", "llm.trace.retrieve.completed"),
+}
+
+
+# ---------------------------------------------------------------------------
+# Event handler
+# ---------------------------------------------------------------------------
+
+
+class LLMSchemaEventHandler:
+    """LlamaIndex callback handler that emits SpanForge events.
+
+    Compatible with both ``llama_index.core`` and legacy ``llama_index``
+    SDK versions.  Events are accumulated in :attr:`events` and optionally
+    forwarded to an async exporter.
+
+    Args:
+        source:   Value for ``Event.source`` (e.g. ``"rag-app@1.0.0"``).
+        org_id:   Optional organisation identifier.
+        exporter: Optional async exporter with an ``export(event)``
+                  coroutine method.  Export is scheduled via ``create_task``
+                  when the event loop is running.
+
+    Attributes:
+        events: List of all :class:`~spanforge.event.Event` objects emitted by
+                this handler in chronological order.
+    """
+
+    def __init__(
+        self,
+        source: str,
+        *,
+        org_id: str | None = None,
+        exporter: Any | None = None,  # noqa: ANN401
+    ) -> None:
+        self._source = source
+        self._org_id = org_id
+        self._exporter = exporter
+        self.events: list[Event] = []
+        # Tracks start monotonic time by event_id
+        self._start_times: dict[str, float] = {}
+
+    # ------------------------------------------------------------------
+    # Static helpers
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _cb_event_type_str(event_type: Any) -> str:  # noqa: ANN401
+        """Convert a LlamaIndex CBEventType (or string) to a plain string.
+
+        If *event_type* has a ``value`` attribute (i.e. it is an ``Enum``),
+        the ``str`` representation of that value is returned.  Otherwise the
+        object is converted to ``str`` directly.
+
+        Args:
+            event_type: A ``CBEventType`` enum member or plain string.
+
+        Returns:
+            The string representation of the event type.
+        """
+        if hasattr(event_type, "value"):
+            return str(event_type.value)
+        return str(event_type)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _make_event(self, event_type: str, payload: dict[str, Any]) -> Event:
+        """Create a SpanForge event, append it, and optionally schedule async export.
+
+        Args:
+            event_type: Dotted event type string.
+            payload:    Event payload dict.
+
+        Returns:
+            The newly created :class:`~spanforge.event.Event`.
+        """
+        event = Event(
+            event_type=event_type,
+            source=self._source,
+            org_id=self._org_id,
+            payload=payload,
+            event_id=gen_ulid(),
+        )
+        self.events.append(event)
+
+        if self._exporter is not None:
+            try:
+                loop = asyncio.get_running_loop()
+                loop.create_task(self._exporter.export(event))  # noqa: RUF006
+            except RuntimeError:
+                pass  # no event loop configured
+
+        return event
+
+    def _duration_ms(self, event_id: str) -> float | None:
+        """Return elapsed milliseconds since *event_id* was started, or ``None``.
+
+        Args:
+            event_id: The event identifier recorded in :attr:`_start_times`.
+
+        Returns:
+            Duration in milliseconds, or ``None`` if *event_id* was not found.
+        """
+        start = self._start_times.pop(event_id, None)
+        if start is None:
+            return None
+        return (time.monotonic() - start) * 1000.0
+
+    def _start_event_payload(self, et: str, payload: dict[str, Any]) -> dict[str, Any]:
+        """Build the payload dict for a start event based on the event type."""
+        result: dict[str, Any] = {}
+        if et == "LLM":
+            model_dict = payload.get("model_dict") or {}
+            if isinstance(model_dict, dict):
+                result["model"] = model_dict.get("model")
+        elif et == "FUNCTION_CALL":
+            tool_info = payload.get("tool") or {}
+            if isinstance(tool_info, dict):
+                result["tool_name"] = tool_info.get("name")
+        elif et == "QUERY":
+            result["query"] = payload.get("query_str")
+        return result
+
+    def _end_event_payload(self, et: str, payload: dict[str, Any]) -> dict[str, Any]:
+        """Build the payload dict for an end event based on the event type."""
+        result: dict[str, Any] = {}
+        if et == "LLM":
+            response = payload.get("response")
+            raw = getattr(response, "raw", None) if response is not None else None
+            if isinstance(raw, dict):
+                usage = raw.get("usage") or {}
+                if isinstance(usage, dict):
+                    result["prompt_tokens"] = usage.get("prompt_tokens")
+                    result["completion_tokens"] = usage.get("completion_tokens")
+                    result["total_tokens"] = usage.get("total_tokens")
+        elif et == "FUNCTION_CALL":
+            result["output"] = payload.get("output")
+        elif et == "QUERY":
+            result["response"] = str(payload.get("response", ""))[:2048]
+        return result
+
+    # ------------------------------------------------------------------
+    # LlamaIndex callback interface
+    # ------------------------------------------------------------------
+
+    def on_event_start(
+        self,
+        event_type: Any,  # noqa: ANN401
+        *,
+        payload: dict[str, Any] | None = None,
+        event_id: str | None = None,
+        **kwargs: Any,  # noqa: ANN401
+    ) -> str:
+        """Called when a LlamaIndex callback event begins.
+
+        Args:
+            event_type: A ``CBEventType`` enum member or plain string.
+            payload:    Optional dict of event-specific data.
+            event_id:   Optional identifier for this event; generated if absent.
+            **kwargs:   Additional keyword arguments (ignored).
+
+        Returns:
+            The ``event_id`` (passed through or generated).
+        """
+        if event_id is None:
+            event_id = gen_ulid()
+
+        et = self._cb_event_type_str(event_type)
+        type_map = _CB_TYPE_MAP.get(et)
+        if type_map is None:
+            # Unknown event type — record start time but emit nothing.
+            self._start_times[event_id] = time.monotonic()
+            return event_id
+
+        self._start_times[event_id] = time.monotonic()
+
+        start_et, _ = type_map
+        payload = payload or {}
+        event_payload: dict[str, Any] = {"event_id": event_id, **self._start_event_payload(et, payload)}
+
+        self._make_event(start_et, event_payload)
+        return event_id
+
+    def on_event_end(
+        self,
+        event_type: Any,  # noqa: ANN401
+        *,
+        payload: dict[str, Any] | None = None,
+        event_id: str | None = None,
+        **kwargs: Any,  # noqa: ANN401
+    ) -> None:
+        """Called when a LlamaIndex callback event ends.
+
+        Args:
+            event_type: A ``CBEventType`` enum member or plain string.
+            payload:    Optional dict of event-specific data (e.g. response).
+            event_id:   Identifies which started event this ends.
+            **kwargs:   Additional keyword arguments (ignored).
+        """
+        et = self._cb_event_type_str(event_type)
+        type_map = _CB_TYPE_MAP.get(et)
+        if type_map is None:
+            # Unknown event type — consume any pending start time silently.
+            if event_id:
+                self._start_times.pop(event_id, None)
+            return
+
+        _, end_et = type_map
+        duration: float | None = self._duration_ms(event_id) if event_id else None
+        payload = payload or {}
+        event_payload: dict[str, Any] = {
+            "event_id": event_id,
+            "duration_ms": duration,
+            **self._end_event_payload(et, payload),
+        }
+
+        self._make_event(end_et, event_payload)
+
+    def start_trace(self, trace_id: str | None = None, **kwargs: Any) -> None:  # noqa: ANN401
+        """No-op — LlamaIndex trace lifecycle hook.
+
+        Args:
+            trace_id: Ignored.
+            **kwargs: Ignored.
+        """
+        pass  # intentionally empty
+
+    def end_trace(
+        self,
+        trace_id: str | None = None,
+        trace_map: dict[str, Any] | None = None,
+        **kwargs: Any,  # noqa: ANN401
+    ) -> None:
+        """No-op — LlamaIndex trace lifecycle hook.
+
+        Args:
+            trace_id:  Ignored.
+            trace_map: Ignored.
+            **kwargs:  Ignored.
+        """
+        pass  # intentionally empty
+
+    def clear_events(self) -> None:
+        """Remove all accumulated events from :attr:`events`."""
+        self.events.clear()
+
+    # ------------------------------------------------------------------
+    # dunder
+    # ------------------------------------------------------------------
+
+    def __repr__(self) -> str:
+        return (
+            f"LLMSchemaEventHandler("
+            f"source={self._source!r}, "
+            f"events={len(self.events)})"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Module-level patch / unpatch API (for API consistency with other integrations)
+# ---------------------------------------------------------------------------
+
+def is_patched() -> bool:
+    """Return ``True`` if the LlamaIndex integration module is importable.
+
+    LlamaIndex uses a callback-handler pattern rather than monkey-patching, so
+    there is no global state to check.  This function simply verifies that the
+    LlamaIndex package is available.
+
+    Returns:
+        ``True`` when LlamaIndex is installed and the handler can be created.
+    """
+    try:
+        _require_llamaindex()
+        return True
+    except ImportError:
+        return False
+
+
+def unpatch() -> None:
+    """No-op for the LlamaIndex integration.
+
+    LlamaIndex uses an explicit callback handler (not monkey-patching), so
+    there is nothing to undo globally.  Remove the handler from your
+    ``CallbackManager`` to stop receiving events::
+
+        Settings.callback_manager = CallbackManager([])  # remove all handlers
+    """

spanforge/integrations/ollama.py

@@ -0,0 +1,286 @@
+"""spanforge.integrations.ollama — Auto-instrumentation for the Ollama Python SDK.
+
+This module monkey-patches the Ollama client so every
+``ollama.chat(...)`` (or ``client.chat(...)``) call automatically populates the
+active :class:`~spanforge._span.Span` with:
+
+* :class:`~spanforge.namespaces.trace.TokenUsage` (prompt / eval token counts
+  mapped to input / output)
+* :class:`~spanforge.namespaces.trace.ModelInfo` (provider = ``ollama``, name
+  from response)
+* :class:`~spanforge.namespaces.trace.CostBreakdown` (always zero — Ollama is a
+  local runtime with no per-token billing)
+
+Usage::
+
+    from spanforge.integrations import ollama as ollama_integration
+    ollama_integration.patch()
+
+    import ollama
+
+    import spanforge
+    spanforge.configure(exporter="console")
+
+    with spanforge.span("ollama-chat", model="llama3") as span:
+        resp = ollama.chat(
+            model="llama3",
+            messages=[{"role": "user", "content": "Hello"}],
+        )
+    # → span.token_usage auto-populated on exit; cost is always zero
+
+Calling ``patch()`` is **idempotent** — calling it multiple times has no
+effect.  Call :func:`unpatch` to restore the original functions.
+
+Install with::
+
+    pip install "spanforge[ollama]"
+
+.. note::
+   Ollama has no per-token pricing.  :func:`normalize_response` always returns
+   :meth:`~spanforge.namespaces.trace.CostBreakdown.zero` for the cost field.
+"""
+
+from __future__ import annotations
+
+import functools
+from typing import Any
+
+from spanforge.namespaces.trace import (
+    CostBreakdown,
+    GenAISystem,
+    ModelInfo,
+    TokenUsage,
+)
+
+__all__ = [
+    "is_patched",
+    "normalize_response",
+    "patch",
+    "unpatch",
+]
+
+# Sentinel to prevent double-patching.
+_PATCH_FLAG = "_spanforge_patched"
+
+
+# ---------------------------------------------------------------------------
+# Public API
+# ---------------------------------------------------------------------------
+
+
+def patch() -> None:
+    """Monkey-patch the Ollama module to auto-instrument all chat calls.
+
+    Wraps the module-level ``ollama.chat`` function and, when instantiated,
+    ``ollama.Client.chat``.  The wrapper calls :func:`normalize_response` on
+    the result and, if a span is currently active, updates it.
+
+    This function is **idempotent** — safe to call multiple times.
+
+    Raises:
+        ImportError: If the ``ollama`` package is not installed.
+    """
+    ollama_mod = _require_ollama()
+
+    if getattr(ollama_mod, _PATCH_FLAG, False):
+        return  # already patched
+
+    # --- module-level ollama.chat --------------------------------------------
+    _orig_chat = getattr(ollama_mod, "chat", None)
+    if _orig_chat is not None:
+        @functools.wraps(_orig_chat)
+        def _patched_chat(*args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+            response = _orig_chat(*args, **kwargs)
+            _auto_populate_span(response)
+            return response
+
+        ollama_mod.chat = _patched_chat  # type: ignore[attr-defined]
+        ollama_mod._spanforge_orig_chat = _orig_chat  # type: ignore[attr-defined]
+
+    # --- Client.chat ---------------------------------------------------------
+    try:
+        from ollama import Client  # type: ignore[import-untyped]  # noqa: PLC0415
+
+        _orig_client_chat = Client.chat  # type: ignore[attr-defined]
+
+        @functools.wraps(_orig_client_chat)
+        def _patched_client_chat(self: Any, *args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+            response = _orig_client_chat(self, *args, **kwargs)
+            _auto_populate_span(response)
+            return response
+
+        Client.chat = _patched_client_chat  # type: ignore[method-assign]
+        Client._spanforge_orig_chat = _orig_client_chat  # type: ignore[attr-defined]
+    except (ImportError, AttributeError):  # pragma: no cover
+        pass
+
+    # --- AsyncClient.chat ----------------------------------------------------
+    try:
+        from ollama import AsyncClient  # type: ignore[import-untyped]  # noqa: PLC0415
+
+        _orig_async_chat = AsyncClient.chat  # type: ignore[attr-defined]
+
+        @functools.wraps(_orig_async_chat)
+        async def _patched_async_chat(self: Any, *args: Any, **kwargs: Any) -> Any:  # noqa: ANN401
+            response = await _orig_async_chat(self, *args, **kwargs)
+            _auto_populate_span(response)
+            return response
+
+        AsyncClient.chat = _patched_async_chat  # type: ignore[method-assign]
+        AsyncClient._spanforge_orig_chat = _orig_async_chat  # type: ignore[attr-defined]
+    except (ImportError, AttributeError):  # pragma: no cover
+        pass
+
+    ollama_mod._spanforge_patched = True  # type: ignore[attr-defined]
+
+
+def unpatch() -> None:
+    """Restore the original Ollama functions and remove the patch flag.
+
+    Safe to call even if :func:`patch` was never called.
+
+    Raises:
+        ImportError: If the ``ollama`` package is not installed.
+    """
+    ollama_mod = _require_ollama()
+
+    if not getattr(ollama_mod, _PATCH_FLAG, False):
+        return  # nothing to do
+
+    orig_chat = getattr(ollama_mod, "_spanforge_orig_chat", None)
+    if orig_chat is not None:
+        ollama_mod.chat = orig_chat  # type: ignore[attr-defined]
+        del ollama_mod._spanforge_orig_chat  # type: ignore[attr-defined]
+
+    try:
+        from ollama import Client  # type: ignore[import-untyped]  # noqa: PLC0415
+
+        Client.chat = Client._spanforge_orig_chat  # type: ignore[attr-defined,method-assign]
+        del Client._spanforge_orig_chat  # type: ignore[attr-defined]
+    except (ImportError, AttributeError):  # pragma: no cover
+        pass
+
+    try:
+        from ollama import AsyncClient  # type: ignore[import-untyped]  # noqa: PLC0415
+
+        AsyncClient.chat = AsyncClient._spanforge_orig_chat  # type: ignore[attr-defined,method-assign]
+        del AsyncClient._spanforge_orig_chat  # type: ignore[attr-defined]
+    except (ImportError, AttributeError):  # pragma: no cover
+        pass
+
+    try:  # noqa: SIM105
+        del ollama_mod._spanforge_patched  # type: ignore[attr-defined]
+    except AttributeError:  # pragma: no cover
+        pass
+
+
+def is_patched() -> bool:
+    """Return ``True`` if Ollama has been patched by spanforge.
+
+    Returns ``False`` if the ``ollama`` package is not installed.
+    """
+    try:
+        ollama_mod = _require_ollama()
+        return bool(getattr(ollama_mod, _PATCH_FLAG, False))
+    except ImportError:
+        return False
+
+
+def normalize_response(
+    response: Any,  # noqa: ANN401
+) -> tuple[TokenUsage, ModelInfo, CostBreakdown]:
+    """Extract structured observability data from an Ollama chat response.
+
+    Works with both ``ollama.ChatResponse`` objects and any duck-typed mock
+    with the same attribute structure.
+
+    Args:
+        response: An Ollama chat response (or compatible object).
+
+    Returns:
+        A 3-tuple of ``(TokenUsage, ModelInfo, CostBreakdown)``.
+        The ``CostBreakdown`` is **always zero** — Ollama has no billing.
+
+    Field mapping:
+
+    +--------------------------------------------+---------------------------+
+    | Ollama field                               | SpanForge field             |
+    +============================================+===========================+
+    | ``response.model``                         | ``ModelInfo.name``        |
+    | ``response.prompt_eval_count``             | ``TokenUsage.input_tokens``|
+    | ``response.eval_count``                    | ``TokenUsage.output_tokens``|
+    +--------------------------------------------+---------------------------+
+    """
+    # Ollama may return a dict or an object depending on SDK version.
+    if isinstance(response, dict):
+        model_name: str = response.get("model", None) or "unknown"
+        input_tokens = int(response.get("prompt_eval_count", 0) or 0)
+        output_tokens = int(response.get("eval_count", 0) or 0)
+    else:
+        model_name = getattr(response, "model", None) or "unknown"
+        input_tokens = int(getattr(response, "prompt_eval_count", 0) or 0)
+        output_tokens = int(getattr(response, "eval_count", 0) or 0)
+
+    total_tokens = input_tokens + output_tokens
+
+    token_usage = TokenUsage(
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        total_tokens=total_tokens,
+    )
+
+    model_info = ModelInfo(system=GenAISystem.OLLAMA, name=model_name)
+
+    # Ollama is local — no billing data available.
+    cost = CostBreakdown.zero()
+
+    return token_usage, model_info, cost
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+def _require_ollama() -> Any:  # noqa: ANN401
+    """Import and return the ``ollama`` module, raising ``ImportError`` if absent."""
+    try:
+        import ollama  # type: ignore[import-untyped]  # noqa: PLC0415
+    except ImportError as exc:
+        raise ImportError(
+            "The 'ollama' package is required for spanforge Ollama integration.\n"
+            "Install it with: pip install 'spanforge[ollama]'"
+        ) from exc
+    else:
+        return ollama
+
+
+def _auto_populate_span(response: Any) -> None:  # noqa: ANN401
+    """If there is an active span on this thread, populate it from *response*.
+
+    Silently does nothing if:
+
+    * There is no active span.
+    * ``normalize_response`` raises (malformed response).
+    * The span already has ``token_usage`` set (don't overwrite manual data).
+    """
+    try:
+        from spanforge._span import _span_stack  # noqa: PLC0415
+
+        stack = _span_stack()
+        if not stack:
+            return
+        span = stack[-1]
+
+        if span.token_usage is not None:
+            return
+
+        token_usage, model_info, cost = normalize_response(response)
+        span.token_usage = token_usage
+        span.cost = cost
+
+        if span.model is None:
+            span.model = model_info.name
+
+    except Exception:  # noqa: S110  # NOSONAR
+        pass