spanforge 1.0.0__py3-none-any.whl

spanforge/integrations/langchain.py

@@ -0,0 +1,349 @@
+"""spanforge.integrations.langchain — LangChain callback handler.
+
+Provides :class:`LLMSchemaCallbackHandler`, a LangChain-compatible
+callback handler that records LLM and tool-call activity as SpanForge events.
+
+Usage::
+
+    from spanforge.integrations.langchain import LLMSchemaCallbackHandler
+
+    handler = LLMSchemaCallbackHandler(source="my-app", org_id="org-1")
+    chain = SomeChain(callbacks=[handler])
+    chain.run("What is 2+2?")
+
+    for event in handler.events:
+        print(event.to_json())
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any
+
+from spanforge.event import Event
+from spanforge.ulid import generate as gen_ulid
+
+if TYPE_CHECKING:
+    from uuid import UUID
+
+__all__ = [
+    "LLMSchemaCallbackHandler",
+    "is_patched",
+    "unpatch",
+]
+
+# ---------------------------------------------------------------------------
+# Module resolver
+# ---------------------------------------------------------------------------
+
+
+def _require_langchain() -> Any:
+    """Return the LangChain callbacks module from whichever package is installed.
+
+    Tries ``langchain_core.callbacks`` first (preferred, modern API), then
+    falls back to the legacy ``langchain.callbacks`` module.
+
+    Raises:
+        ImportError: If neither ``langchain_core`` nor ``langchain`` is installed.
+    """
+    # Try modern langchain_core first.  Import the parent module first so that
+    # a None sentinel in sys.modules propagates as ImportError correctly.
+    try:
+        import sys
+
+        import langchain_core
+        import langchain_core.callbacks
+
+        return sys.modules["langchain_core.callbacks"]
+    except ImportError:
+        pass
+    # Fall back to legacy langchain package.
+    try:
+        import sys
+
+        import langchain
+        import langchain.callbacks
+
+        return sys.modules["langchain.callbacks"]
+    except ImportError:
+        pass
+    raise ImportError(
+        "LangChain package is required for the spanforge LangChain integration.\n"
+        "Install it with: pip install 'spanforge[langchain]'"
+    )
+
+
+# ---------------------------------------------------------------------------
+# Callback handler
+# ---------------------------------------------------------------------------
+
+
+class LLMSchemaCallbackHandler:
+    """LangChain callback handler that emits SpanForge events.
+
+    Compatible with both ``langchain_core`` and legacy ``langchain`` SDK
+    versions.  Events are accumulated in :attr:`events` and optionally
+    forwarded to an async exporter.
+
+    Args:
+        source:   Value for ``Event.source`` (e.g. ``"my-llm-app@1.0.0"``).
+        org_id:   Optional organisation identifier.
+        exporter: Optional async exporter; must have an ``export(event)``
+                  coroutine method.  When the event loop is running, export is
+                  scheduled as a task; otherwise the call is silently skipped.
+
+    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,
+    ) -> None:
+        self._source = source
+        self._org_id = org_id
+        self._exporter = exporter
+        self.events: list[Event] = []
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _make_event(self, event_type: str, payload: dict[str, Any]) -> Event:
+        """Create a SpanForge event 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))
+            except RuntimeError:
+                pass  # no running event loop
+
+        return event
+
+    # ------------------------------------------------------------------
+    # LangChain callback interface
+    # ------------------------------------------------------------------
+
+    def on_llm_start(
+        self,
+        serialized: dict[str, Any],
+        prompts: list[str],
+        *,
+        run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when an LLM invocation begins.
+
+        Args:
+            serialized: Serialised LLM config dict (contains ``id`` list).
+            prompts:    List of prompt strings being sent to the LLM.
+            run_id:     LangChain run identifier.
+            **kwargs:   Additional keyword arguments (ignored).
+        """
+        llm_name = ""
+        if serialized and "id" in serialized and serialized["id"]:
+            llm_name = str(serialized["id"][-1])
+        self._make_event(
+            "llm.trace.span.started",
+            {
+                "llm_name": llm_name,
+                "prompt_count": len(prompts),
+                "run_id": str(run_id) if run_id is not None else None,
+            },
+        )
+
+    def on_llm_end(
+        self,
+        response: Any,
+        *,
+        run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when an LLM invocation completes.
+
+        Args:
+            response: LangChain ``LLMResult`` object with ``llm_output`` attribute.
+            run_id:   LangChain run identifier.
+            **kwargs: Additional keyword arguments (ignored).
+        """
+        llm_output = getattr(response, "llm_output", None)
+        prompt_tokens: int | None = None
+        completion_tokens: int | None = None
+        total_tokens: int | None = None
+
+        if isinstance(llm_output, dict):
+            token_usage = llm_output.get("token_usage") or {}
+            if isinstance(token_usage, dict):
+                prompt_tokens = token_usage.get("prompt_tokens")
+                completion_tokens = token_usage.get("completion_tokens")
+                total_tokens = token_usage.get("total_tokens")
+
+        self._make_event(
+            "llm.trace.span.completed",
+            {
+                "run_id": str(run_id) if run_id is not None else None,
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+            },
+        )
+
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when an LLM invocation raises an error.
+
+        Args:
+            error:    The exception that was raised.
+            run_id:   LangChain run identifier.
+            **kwargs: Additional keyword arguments (ignored).
+        """
+        self._make_event(
+            "llm.trace.span.error",
+            {
+                "run_id": str(run_id) if run_id is not None else None,
+                "error": str(error),
+                "error_type": type(error).__name__,
+            },
+        )
+
+    def on_tool_start(
+        self,
+        serialized: dict[str, Any],
+        input_str: str,
+        *,
+        run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when a tool invocation begins.
+
+        Args:
+            serialized: Serialised tool config dict (usually has ``"name"`` key).
+            input_str:  String input passed to the tool.
+            run_id:     LangChain run identifier.
+            **kwargs:   Additional keyword arguments (ignored).
+        """
+        tool_name = serialized.get("name", "") if serialized else ""
+        self._make_event(
+            "llm.trace.tool_call.started",
+            {
+                "tool_name": tool_name,
+                "input": input_str,
+                "run_id": str(run_id) if run_id is not None else None,
+            },
+        )
+
+    def on_tool_end(
+        self,
+        output: str,
+        *,
+        run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when a tool invocation completes.
+
+        Args:
+            output:   String output returned by the tool.
+            run_id:   LangChain run identifier.
+            **kwargs: Additional keyword arguments (ignored).
+        """
+        self._make_event(
+            "llm.trace.tool_call.completed",
+            {
+                "run_id": str(run_id) if run_id is not None else None,
+                "output": str(output)[:1024] if output else None,
+            },
+        )
+
+    def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when a tool invocation raises an error.
+
+        Args:
+            error:    The exception that was raised.
+            run_id:   LangChain run identifier.
+            **kwargs: Additional keyword arguments (ignored).
+        """
+        self._make_event(
+            "llm.trace.tool_call.error",
+            {
+                "run_id": str(run_id) if run_id is not None else None,
+                "error": str(error),
+                "error_type": type(error).__name__,
+            },
+        )
+
+    def clear_events(self) -> None:
+        """Remove all accumulated events from :attr:`events`."""
+        self.events.clear()
+
+    # ------------------------------------------------------------------
+    # dunder
+    # ------------------------------------------------------------------
+
+    def __repr__(self) -> str:
+        return f"LLMSchemaCallbackHandler(source={self._source!r}, events={len(self.events)})"
+
+
+# ---------------------------------------------------------------------------
+# Module-level patch / unpatch API (for API consistency with other integrations)
+# ---------------------------------------------------------------------------
+
+
+def is_patched() -> bool:
+    """Return ``True`` if the LangChain integration module is importable.
+
+    LangChain uses a callback-handler pattern rather than monkey-patching, so
+    there is no global state to check.  This function simply verifies that the
+    LangChain package is available and the handler class is ready to use.
+
+    Returns:
+        ``True`` when LangChain is installed and the handler can be created.
+    """
+    try:
+        _require_langchain()
+    except ImportError:
+        return False
+    else:
+        return True
+
+
+def unpatch() -> None:
+    """No-op for the LangChain integration.
+
+    LangChain uses an explicit callback handler (not monkey-patching), so
+    there is nothing to undo globally.  Remove the handler from any
+    ``CallbackManager`` or chain you attached it to::
+
+        chain = SomeChain(callbacks=[])  # re-create without the handler
+    """

spanforge/integrations/langgraph.py

@@ -0,0 +1,306 @@
+"""spanforge.integrations.langgraph - LangGraph governance demo integration.
+
+Implements a lightweight handler that records graph and node execution while
+optionally invoking SpanForge runtime-governance services.  This is intended to
+cover the GA demo path rather than deep framework auto-patching.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+from spanforge.event import Event
+from spanforge.ulid import generate as gen_ulid
+
+__all__ = [
+    "LangGraphGovernanceHandler",
+    "LangGraphNodeResult",
+    "LangGraphRunRecord",
+    "is_available",
+]
+
+
+def _utc_now() -> str:
+    return datetime.now(tz=timezone.utc).isoformat(timespec="microseconds").replace("+00:00", "Z")
+
+
+def _require_langgraph() -> Any:
+    try:
+        import langgraph
+    except ImportError as exc:
+        raise ImportError(
+            "LangGraph is required for this integration.\n"
+            "Install it with: pip install 'spanforge[langgraph]'"
+        ) from exc
+    return langgraph
+
+
+def is_available() -> bool:
+    """Return ``True`` when LangGraph is importable."""
+    try:
+        _require_langgraph()
+    except ImportError:
+        return False
+    return True
+
+
+@dataclass
+class LangGraphNodeResult:
+    """Recorded result for one node in a LangGraph run."""
+
+    node_id: str
+    trace_id: str
+    node_name: str
+    node_type: str
+    started_at: str
+    completed_at: str | None = None
+    status: str = "started"
+    scope_result: Any | None = None
+    rbac_result: Any | None = None
+    lineage_result: Any | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        data = {
+            "node_id": self.node_id,
+            "trace_id": self.trace_id,
+            "node_name": self.node_name,
+            "node_type": self.node_type,
+            "started_at": self.started_at,
+            "status": self.status,
+            "metadata": dict(self.metadata),
+        }
+        if self.completed_at is not None:
+            data["completed_at"] = self.completed_at
+        if self.scope_result is not None:
+            data["scope_result"] = _to_dict(self.scope_result)
+        if self.rbac_result is not None:
+            data["rbac_result"] = _to_dict(self.rbac_result)
+        if self.lineage_result is not None:
+            data["lineage_result"] = _to_dict(self.lineage_result)
+        return data
+
+
+@dataclass
+class LangGraphRunRecord:
+    """High-level trace record for one governed LangGraph run."""
+
+    run_id: str
+    trace_id: str
+    graph_name: str
+    environment: str
+    started_at: str
+    completed_at: str | None = None
+    status: str = "started"
+    agent_id: str | None = None
+    actor_id: str | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+    nodes: list[LangGraphNodeResult] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        data = {
+            "run_id": self.run_id,
+            "trace_id": self.trace_id,
+            "graph_name": self.graph_name,
+            "environment": self.environment,
+            "started_at": self.started_at,
+            "status": self.status,
+            "metadata": dict(self.metadata),
+            "nodes": [node.to_dict() for node in self.nodes],
+        }
+        if self.completed_at is not None:
+            data["completed_at"] = self.completed_at
+        if self.agent_id is not None:
+            data["agent_id"] = self.agent_id
+        if self.actor_id is not None:
+            data["actor_id"] = self.actor_id
+        return data
+
+
+def _to_dict(value: Any) -> Any:
+    if hasattr(value, "to_dict"):
+        return value.to_dict()
+    return value
+
+
+class LangGraphGovernanceHandler:
+    """Record LangGraph execution and invoke runtime-governance controls."""
+
+    def __init__(
+        self,
+        *,
+        source: str = "spanforge.langgraph@1.0.0",
+        environment: str = "prod",
+        policy_client: Any | None = None,
+        scope_client: Any | None = None,
+        rbac_client: Any | None = None,
+        lineage_client: Any | None = None,
+    ) -> None:
+        self._source = source
+        self._environment = environment
+        self._policy_client = policy_client
+        self._scope_client = scope_client
+        self._rbac_client = rbac_client
+        self._lineage_client = lineage_client
+        self.events: list[Event] = []
+        self.runs: dict[str, LangGraphRunRecord] = {}
+        self.nodes: dict[str, LangGraphNodeResult] = {}
+
+    def on_graph_start(
+        self,
+        *,
+        trace_id: str,
+        graph_name: str,
+        agent_id: str | None = None,
+        actor_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> LangGraphRunRecord:
+        """Start recording one LangGraph run."""
+        run = LangGraphRunRecord(
+            run_id=gen_ulid(),
+            trace_id=trace_id,
+            graph_name=graph_name,
+            environment=self._environment,
+            started_at=_utc_now(),
+            agent_id=agent_id,
+            actor_id=actor_id,
+            metadata=metadata or {},
+        )
+        self.runs[trace_id] = run
+        self._emit_event(
+            "llm.langgraph.run.started",
+            trace_id=trace_id,
+            payload=run.to_dict(),
+        )
+        return run
+
+    def on_node_start(
+        self,
+        *,
+        trace_id: str,
+        node_name: str,
+        node_type: str = "chain",
+        agent_id: str | None = None,
+        actor_id: str | None = None,
+        resource: str | None = None,
+        action_name: str | None = None,
+        capability: str | None = None,
+        required_roles: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> LangGraphNodeResult:
+        """Record node start and optionally run scope/RBAC checks."""
+        started_at = _utc_now()
+        result = LangGraphNodeResult(
+            node_id=gen_ulid(),
+            trace_id=trace_id,
+            node_name=node_name,
+            node_type=node_type,
+            started_at=started_at,
+            metadata=metadata or {},
+        )
+        if self._scope_client is not None and agent_id and resource and action_name:
+            result.scope_result = self._scope_client.evaluate_with_policy(
+                environment=self._environment,
+                trace_id=trace_id,
+                agent_id=agent_id,
+                resource=resource,
+                action_name=action_name,
+                checked_at=started_at,
+                capability=capability,
+                policy_client=self._policy_client,
+            )
+        if self._rbac_client is not None and actor_id and resource and action_name:
+            result.rbac_result = self._rbac_client.authorize_with_policy(
+                environment=self._environment,
+                trace_id=trace_id,
+                actor_id=actor_id,
+                resource=resource,
+                action_name=action_name,
+                checked_at=started_at,
+                required_roles=required_roles,
+                policy_client=self._policy_client,
+            )
+        self.nodes[result.node_id] = result
+        run = self.runs.get(trace_id)
+        if run is not None:
+            run.nodes.append(result)
+        self._emit_event(
+            "llm.langgraph.node.started",
+            trace_id=trace_id,
+            payload=result.to_dict(),
+        )
+        return result
+
+    def on_node_end(
+        self,
+        *,
+        trace_id: str,
+        node_id: str,
+        decision_id: str,
+        subject_type: str = "langgraph_node",
+        output_refs: list[str] | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> LangGraphNodeResult:
+        """Complete one node and optionally capture lineage."""
+        node = self.nodes[node_id]
+        completed_at = _utc_now()
+        node.completed_at = completed_at
+        node.status = "completed"
+        if self._lineage_client is not None:
+            node.lineage_result = self._lineage_client.record_with_policy(
+                environment=self._environment,
+                trace_id=trace_id,
+                decision_id=decision_id,
+                subject_type=subject_type,
+                subject_id=node.node_id,
+                operation=node.node_name,
+                recorded_at=completed_at,
+                policy_client=self._policy_client,
+                output_refs=output_refs,
+                metadata=metadata,
+            )
+        self._emit_event(
+            "llm.langgraph.node.completed",
+            trace_id=trace_id,
+            payload=node.to_dict(),
+        )
+        return node
+
+    def on_node_error(self, *, trace_id: str, node_id: str, error: BaseException) -> LangGraphNodeResult:
+        """Record node failure."""
+        node = self.nodes[node_id]
+        node.completed_at = _utc_now()
+        node.status = "error"
+        node.metadata["error"] = str(error)
+        node.metadata["error_type"] = type(error).__name__
+        self._emit_event(
+            "llm.langgraph.node.error",
+            trace_id=trace_id,
+            payload=node.to_dict(),
+        )
+        return node
+
+    def on_graph_end(self, *, trace_id: str, status: str = "completed") -> LangGraphRunRecord:
+        """Complete one LangGraph run."""
+        run = self.runs[trace_id]
+        run.completed_at = _utc_now()
+        run.status = status
+        self._emit_event(
+            "llm.langgraph.run.completed",
+            trace_id=trace_id,
+            payload=run.to_dict(),
+        )
+        return run
+
+    def _emit_event(self, event_type: str, *, trace_id: str, payload: dict[str, Any]) -> Event:
+        event = Event(
+            event_type=event_type,
+            source=self._source,
+            payload=payload,
+            trace_id=trace_id,
+            event_id=gen_ulid(),
+        )
+        self.events.append(event)
+        return event