traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/sdk/__init__.py

@@ -0,0 +1,59 @@
+"""Traceforge SDK — pipeline setup and gating.
+
+Usage:
+    from traceforge.sdk import Pipeline, GatePolicy, Verdict, ToolCallRequest
+
+    def my_preflight(request: ToolCallRequest, ctx: GateContext) -> Verdict:
+        if request.risk_score > 60:
+            return Verdict.deny(f"score {request.risk_score} exceeds threshold")
+        return Verdict.allow()
+
+    policy = GatePolicy().preflight(my_preflight)
+
+    pipeline = Pipeline.create(policy=policy)
+    pipeline.gate_crewai()          # CrewAI hooks
+    pipeline.gate_langchain(tool)   # LangChain tool wrap
+    pipeline.gate_maf()             # MAF middleware
+
+The preflight callback returns a Verdict. Traceforge enforces it using each framework's
+native blocking mechanism. The postflight callback receives the tool output for audit.
+"""
+
+from __future__ import annotations
+
+from traceforge.sdk.gate_policy import GatePolicy  # noqa: E402
+from traceforge.sdk.gate_types import (  # noqa: E402
+    GateContext,
+    PostflightAction,
+    PostflightVerdict,
+    ToolCallRequest,
+    ToolCallResult,
+)
+from traceforge.sdk.verdict import (  # noqa: E402
+    Decision,
+    PostflightGate,
+    PreflightGate,
+    Verdict,
+)
+from traceforge.trace import EventTrace, TraceStage  # noqa: E402
+
+# Imported last so every SDK submodule symbol above is already defined before the
+# facade (which pulls in the governance engine + core EventPipeline) resolves,
+# keeping this package free of import-cycle ordering hazards.
+from traceforge.sdk.pipeline import Pipeline  # noqa: E402
+
+__all__ = [
+    "Pipeline",
+    "EventTrace",
+    "TraceStage",
+    "Verdict",
+    "Decision",
+    "PreflightGate",
+    "PostflightGate",
+    "GatePolicy",
+    "GateContext",
+    "ToolCallRequest",
+    "ToolCallResult",
+    "PostflightVerdict",
+    "PostflightAction",
+]

traceforge/sdk/gate_policy.py

@@ -0,0 +1,63 @@
+"""GatePolicy — composable, testable gate registration.
+
+Gates are registered via a GatePolicy which owns the ordered chain
+of preflight/postflight gates. The pipeline uses this instead of
+per-attach kwargs.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from traceforge.sdk.verdict import PostflightGate, PreflightGate
+
+
+@dataclass
+class GatePolicy:
+    """Ordered collection of gates. Testable without a pipeline.
+
+    Usage:
+        policy = (
+            GatePolicy()
+            .preflight(block_destructive_shell)
+            .preflight(rate_limit_gate)
+            .postflight(redact_secrets)
+        )
+
+        pipeline = Pipeline.create(policy=policy)
+
+    Gates run in registration order (top-down).
+    """
+
+    _preflight: list["PreflightGate"] = field(default_factory=list)
+    _postflight: list["PostflightGate"] = field(default_factory=list)
+
+    def preflight(self, gate: "PreflightGate") -> "GatePolicy":
+        """Register a preflight gate. Runs in registration order."""
+        self._preflight.append(gate)
+        return self
+
+    def postflight(self, gate: "PostflightGate") -> "GatePolicy":
+        """Register a postflight gate. Runs in registration order."""
+        self._postflight.append(gate)
+        return self
+
+    @property
+    def preflight_gates(self) -> tuple["PreflightGate", ...]:
+        """Ordered preflight gates (registration order)."""
+        return tuple(self._preflight)
+
+    @property
+    def postflight_gates(self) -> tuple["PostflightGate", ...]:
+        """Ordered postflight gates (registration order)."""
+        return tuple(self._postflight)
+
+    @property
+    def has_preflight(self) -> bool:
+        return len(self._preflight) > 0
+
+    @property
+    def has_postflight(self) -> bool:
+        return len(self._postflight) > 0

traceforge/sdk/gate_types.py

@@ -0,0 +1,140 @@
+"""Gate-facing types: view objects, context, and postflight verdict.
+
+Gate authors import from here. These types provide the policy-relevant surface
+with guaranteed non-optional fields — no None checks needed in gate logic.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import StrEnum
+from types import MappingProxyType
+from typing import TYPE_CHECKING
+
+from traceforge._generated import (
+    Action,
+    Capability,
+    Effect,
+    Mechanism,
+    Recommendation,
+    RiskBand,
+    Role,
+    Scope,
+)
+from traceforge.sdk.verdict import Verdict
+from traceforge.trace import EMPTY_MAP
+
+if TYPE_CHECKING:
+    from traceforge.trace import EventTrace
+
+
+# ─── Gate View Objects ────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True, slots=True)
+class ToolCallRequest:
+    """What preflight gate authors receive. Policy-focused, no None checks.
+
+    All classification/assessment fields are guaranteed non-None — the pipeline
+    only constructs this for fully-assessed EventTraces.
+    """
+
+    # What tool is being called
+    tool: str
+    input: MappingProxyType
+    target: str | None
+
+    # How it's classified
+    mechanism: Mechanism
+    effect: Effect
+    capabilities: tuple[Capability, ...]
+    scope: tuple[Scope, ...]
+    role: tuple[Role, ...]
+    action: tuple[Action, ...]
+
+    # How dangerous
+    risk_score: int
+    risk_band: RiskBand
+    suggested_action: Recommendation
+    reason: str
+
+    # Identity
+    session_id: str
+    tool_call_id: str
+
+    # Escape hatch — full EventTrace for raw_event, attributes, etc.
+    event_trace: EventTrace
+
+
+@dataclass(frozen=True, slots=True)
+class ToolCallResult:
+    """What postflight gate authors receive. Includes tool output."""
+
+    # What tool ran
+    tool: str
+    input: MappingProxyType
+    target: str | None
+
+    # What it produced
+    output: MappingProxyType
+    duration_ms: int | None
+    error: str | None
+
+    # Classification + assessment
+    mechanism: Mechanism
+    effect: Effect
+    capabilities: tuple[Capability, ...]
+    risk_score: int
+    risk_band: RiskBand
+    suggested_action: Recommendation
+    reason: str
+
+    # Identity
+    session_id: str
+    tool_call_id: str
+
+    # Escape hatch
+    event_trace: EventTrace
+
+
+# ─── GateContext ──────────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True, slots=True)
+class GateContext:
+    """Session-scoped state passed as second arg to every gate.
+
+    Pipeline builds and maintains this per session. Gate authors can
+    ignore it for stateless policies.
+    """
+
+    session_id: str
+    tool_call_count: int
+    denied_count: int
+    prior_verdicts: tuple[Verdict, ...] = ()
+    prior_tool_call_ids: tuple[str, ...] = ()
+    agent_id: str | None = None
+    user_id: str | None = None
+    policy: MappingProxyType = field(default_factory=lambda: EMPTY_MAP)
+
+
+# ─── Postflight Verdict ───────────────────────────────────────────────────────
+
+
+class PostflightAction(StrEnum):
+    """What to do with tool output after execution."""
+
+    ACCEPT = "accept"
+    REDACT = "redact"
+    SUPPRESS = "suppress"
+    TERMINATE = "terminate"
+    ALERT = "alert"
+
+
+@dataclass(frozen=True, slots=True)
+class PostflightVerdict:
+    """Returned by postflight gates. Controls what happens to tool output."""
+
+    action: PostflightAction = PostflightAction.ACCEPT
+    reason: str = ""
+    redaction_keys: tuple[str, ...] = ()

traceforge/sdk/pipeline.py

@@ -0,0 +1,265 @@
+"""The traceforge Pipeline: observe -> enrich -> classify -> structure -> [govern] -> sinks.
+
+This is the SDK's top-level entry point. It composes traceforge's two halves into one
+object:
+
+* the **observation backbone** (:class:`~traceforge.pipeline.EventPipeline`), which
+  enriches, classifies, and ML-structures (phase / boundary / title) every event before
+  fanning it out to sinks, and
+* the **governance engine** (:class:`~traceforge.governance.pipeline.GovernancePipeline`),
+  which scores risk, tracks budgets / drift, and backs the opt-in gating layer.
+
+Governance is wired in as **one stage** of the backbone: when enabled, each event is
+scored and its :class:`~traceforge.governance.results.SessionMeta` stamped onto
+``event.metadata.governance`` just before the sinks see it. It is not a separate pipeline
+and not a precondition -- structuring runs with or without it.
+
+Gating is a separate, opt-in enforcement layer. The ``gate_*`` helpers install a
+:class:`~traceforge.sdk.GatePolicy` into your agent framework using that framework's native
+blocking mechanism; the policy's preflight callback returns a
+:class:`~traceforge.sdk.Verdict` (ALLOW / DENY). Observation itself never issues verdicts --
+only the gating layer does, and only when you opt in.
+
+Usage -- gating a live agent::
+
+    from traceforge.sdk import Pipeline, GatePolicy, Verdict, ToolCallRequest, GateContext
+
+    def preflight(request: ToolCallRequest, ctx: GateContext) -> Verdict:
+        if request.risk_score > 60:
+            return Verdict.deny(f"score {request.risk_score} exceeds threshold")
+        return Verdict.allow()
+
+    pipeline = Pipeline.create(policy=GatePolicy().preflight(preflight))
+    tool = pipeline.gate_langchain(tool)   # wrap a LangChain tool
+    pipeline.gate_crewai()                 # install CrewAI hooks
+
+Usage -- observing an event stream::
+
+    from traceforge.sdk import Pipeline
+    from traceforge.sinks.jsonl import JsonlSink
+
+    async with Pipeline.create(sinks=[JsonlSink("events.jsonl")]) as pipeline:
+        async for event in adapter.stream(...):
+            await pipeline.push(event)     # enriched, structured, governed, emitted
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from traceforge.enricher import Enricher
+from traceforge.governance.pipeline import GovernancePipeline
+from traceforge.pipeline import EventPipeline
+
+if TYPE_CHECKING:
+    from traceforge.config.models import GovernanceConfig
+    from traceforge.sdk.gate_policy import GatePolicy
+    from traceforge.sinks.base import StorageSink
+    from traceforge.trace import EventTrace
+    from traceforge.types import SessionEvent, TelemetrySpan, UsageRecord
+
+
+class Pipeline:
+    """Observe -> structure backbone with governance as one stage and opt-in gating.
+
+    Construct via :meth:`create` or :meth:`from_config`; do not instantiate directly.
+    """
+
+    def __init__(self, *, backbone: EventPipeline, governance: GovernancePipeline) -> None:
+        self._backbone = backbone
+        self._governance = governance
+
+    # ---- construction -------------------------------------------------------
+
+    @classmethod
+    def create(
+        cls,
+        config: "GovernanceConfig | None" = None,
+        *,
+        policy: "GatePolicy | None" = None,
+        sinks: "list[StorageSink] | None" = None,
+        enable_structure: bool = True,
+        enable_title: bool = False,
+        enricher: "Enricher | None" = None,
+        governance: bool = True,
+    ) -> "Pipeline":
+        """Build a pipeline from config.
+
+        Args:
+            config: GovernanceConfig for the governance engine (in-memory DB and
+                sensible defaults when omitted).
+            policy: Optional GatePolicy enabling the gating layer (the ``gate_*``
+                helpers). Omit for observation-only usage.
+            sinks: Observation destinations for pushed events. Omit for gating-only
+                usage (you never call :meth:`push`).
+            enable_structure: Run phase + boundary ML structuring on pushed events
+                (default True). Models load lazily on first push, so gating-only
+                usage pays nothing.
+            enable_title: Also run session-title inference (default False).
+            enricher: Custom :class:`~traceforge.enricher.Enricher`; defaults to a
+                zero-config ``Enricher()``.
+            governance: Wire the governance engine in as a pipeline stage so pushed
+                events get ``metadata.governance`` stamped (default True). Set False
+                for pure observation; ``gate_*`` / :meth:`score_tool_call` still use
+                the engine.
+        """
+        gov_engine = GovernancePipeline.create(config, policy=policy)
+        return cls._assemble(
+            gov_engine,
+            sinks=sinks,
+            enable_structure=enable_structure,
+            enable_title=enable_title,
+            enricher=enricher,
+            governance=governance,
+        )
+
+    @classmethod
+    def from_config(
+        cls,
+        path=None,
+        *,
+        policy: "GatePolicy | None" = None,
+        sinks: "list[StorageSink] | None" = None,
+        enable_structure: bool = True,
+        enable_title: bool = False,
+        enricher: "Enricher | None" = None,
+        governance: bool = True,
+    ) -> "Pipeline":
+        """Build a pipeline from a ``traceforge.yaml`` file.
+
+        Same knobs as :meth:`create`, but the governance engine is loaded from the
+        given config file (see :meth:`GovernancePipeline.from_config`).
+
+        Sink hydration (SDK/CLI parity, issue #116): when ``sinks`` is **omitted**
+        (left as ``None``), the standard sinks declared in the config's pipelines
+        (``pipelines[].sinks``) are auto-built from their ``SinkConfig`` entries via
+        the shared :func:`traceforge.sinks.factory.build_sinks` — the same factory the
+        ``traceforge watch`` daemon uses. Sinks declared across multiple named
+        pipelines are flattened into one list, since the SDK backbone is a single
+        source/adapter-agnostic event stream and every declared sink is a live
+        destination.
+
+        Programmatic sinks always win: passing ``sinks=[...]`` — **including an
+        explicit empty list ``[]``** — is honored exactly and disables
+        auto-hydration. ``None`` is the sole "omitted" sentinel.
+        """
+        gov_engine = GovernancePipeline.from_config(path, policy=policy)
+
+        if sinks is None:
+            from traceforge.config.loader import load_config_from_path
+            from traceforge.sinks.factory import build_sinks
+
+            config = load_config_from_path(path)
+            sinks = build_sinks(
+                [sink_config for pipeline in config.pipelines for sink_config in pipeline.sinks]
+            )
+
+        return cls._assemble(
+            gov_engine,
+            sinks=sinks,
+            enable_structure=enable_structure,
+            enable_title=enable_title,
+            enricher=enricher,
+            governance=governance,
+        )
+
+    @classmethod
+    def _assemble(
+        cls,
+        gov_engine: GovernancePipeline,
+        *,
+        sinks,
+        enable_structure,
+        enable_title,
+        enricher,
+        governance,
+    ) -> "Pipeline":
+        backbone = EventPipeline(
+            sinks=list(sinks) if sinks is not None else [],
+            enricher=enricher if enricher is not None else Enricher(),
+            enable_phase=enable_structure,
+            enable_boundary=enable_structure,
+            enable_title=enable_title,
+            governance=gov_engine if governance else None,
+        )
+        return cls(backbone=backbone, governance=gov_engine)
+
+    # ---- observation backbone ----------------------------------------------
+
+    async def push(self, event: "SessionEvent") -> None:
+        """Enrich, classify, structure, govern, and emit a single event."""
+        await self._backbone.push(event)
+
+    async def push_span(self, span: "TelemetrySpan") -> None:
+        """Emit a telemetry span to the sinks."""
+        await self._backbone.push_span(span)
+
+    async def push_usage(self, usage: "UsageRecord") -> None:
+        """Emit a usage record to the sinks."""
+        await self._backbone.push_usage(usage)
+
+    async def flush(self) -> None:
+        """Flush held plumbing and pending refinements, then flush all sinks."""
+        await self._backbone.flush()
+
+    async def close(self) -> None:
+        """Flush, then close all sinks."""
+        await self._backbone.close()
+
+    async def __aenter__(self) -> "Pipeline":
+        return self
+
+    async def __aexit__(self, *exc) -> None:
+        await self.close()
+
+    # ---- preflight scoring (governance stage, read-only) -------------------
+
+    def score_tool_call(self, payload: dict) -> "EventTrace":
+        """Score a prospective tool call (dict) without mutating session state."""
+        return self._governance.score_tool_call(payload)
+
+    # ---- gating layer (opt-in; delegates to the governance engine) ---------
+
+    def gate_crewai(self) -> None:
+        """Install CrewAI gating hooks."""
+        return self._governance.gate_crewai()
+
+    def gate_langchain(self, tool):
+        """Wrap a LangChain tool so calls are gated. Returns the wrapped tool."""
+        return self._governance.gate_langchain(tool)
+
+    def gate_langgraph(self, tools):
+        """Wrap LangGraph tools. Returns a gated tool node."""
+        return self._governance.gate_langgraph(tools)
+
+    def gate_semantic_kernel(self, kernel) -> None:
+        """Install a Semantic Kernel gating filter."""
+        return self._governance.gate_semantic_kernel(kernel)
+
+    def gate_maf(self):
+        """Return Microsoft Agent Framework gating middleware."""
+        return self._governance.gate_maf()
+
+    def gate_smolagents(self, agent_cls=None):
+        """Gate smolagents. Returns a gated agent class."""
+        return self._governance.gate_smolagents(agent_cls)
+
+    def gate_pydantic_ai(self, agent) -> None:
+        """Install Pydantic AI gating."""
+        return self._governance.gate_pydantic_ai(agent)
+
+    def gate_openai_agents(self, agent):
+        """Gate an OpenAI Agents SDK agent."""
+        return self._governance.gate_openai_agents(agent)
+
+    # ---- escape hatches ----------------------------------------------------
+
+    @property
+    def governance(self) -> GovernancePipeline:
+        """The underlying governance engine (scoring, budgets, gating internals)."""
+        return self._governance
+
+    @property
+    def backbone(self) -> EventPipeline:
+        """The underlying observation / structuring EventPipeline."""
+        return self._backbone

traceforge/sdk/verdict.py

@@ -0,0 +1,81 @@
+"""Verdict types for tool-call gating decisions."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from traceforge.sdk.gate_types import (
+        GateContext,
+        PostflightVerdict,
+        ToolCallRequest,
+        ToolCallResult,
+    )
+
+
+# ─── Decision & Verdict ───────────────────────────────────────────────────────
+
+
+class Decision(Enum):
+    """The outcome of a gating decision."""
+
+    ALLOW = "allow"
+    DENY = "deny"
+
+
+@dataclass(frozen=True, slots=True)
+class Verdict:
+    """A gating decision returned by a gate callback.
+
+    Args:
+        decision: ALLOW or DENY.
+        reason: Human-readable reason propagated to the LLM on denial.
+    """
+
+    decision: Decision
+    reason: str = ""
+
+    @property
+    def allowed(self) -> bool:
+        return self.decision == Decision.ALLOW
+
+    @property
+    def denied(self) -> bool:
+        return self.decision == Decision.DENY
+
+    @staticmethod
+    def allow() -> Verdict:
+        """Convenience factory for ALLOW."""
+        return Verdict(decision=Decision.ALLOW)
+
+    @staticmethod
+    def deny(reason: str = "") -> Verdict:
+        """Convenience factory for DENY with reason."""
+        return Verdict(decision=Decision.DENY, reason=reason)
+
+
+# ─── Callback Protocols ───────────────────────────────────────────────────────
+
+
+@runtime_checkable
+class PreflightGate(Protocol):
+    """Protocol for preflight gate callbacks.
+
+    Receives a ToolCallRequest (policy-focused view) and GateContext.
+    Must return a Verdict (ALLOW or DENY).
+    """
+
+    def __call__(self, request: "ToolCallRequest", ctx: "GateContext") -> Verdict: ...
+
+
+@runtime_checkable
+class PostflightGate(Protocol):
+    """Protocol for postflight gate callbacks.
+
+    Receives a ToolCallResult (includes tool output) and GateContext.
+    Must return a PostflightVerdict.
+    """
+
+    def __call__(self, result: "ToolCallResult", ctx: "GateContext") -> "PostflightVerdict": ...

traceforge/sinks/__init__.py

@@ -0,0 +1,23 @@
+"""Storage sinks for the traceforge pipeline."""
+
+from traceforge.sinks.base import StorageSink
+from traceforge.sinks.callback import CallbackSink
+from traceforge.sinks.console import ConsoleSink
+from traceforge.sinks.jsonl import JsonlSink
+from traceforge.sinks.otel_exporter import OtelExporterSink
+from traceforge.sinks.parquet import ParquetSink
+from traceforge.sinks.s3 import S3Sink
+from traceforge.sinks.sqlite_output import SqliteOutputSink
+from traceforge.sinks.webhook import WebhookSink
+
+__all__ = [
+    "StorageSink",
+    "CallbackSink",
+    "ConsoleSink",
+    "JsonlSink",
+    "OtelExporterSink",
+    "ParquetSink",
+    "S3Sink",
+    "SqliteOutputSink",
+    "WebhookSink",
+]