traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/boundary/features.py

@@ -0,0 +1,137 @@
+"""Per-gap featurisation for the boundary classifier (shared train/serve).
+
+A *gap* is the transition point after event ``t``; it is featurised from the
+canonical fields of ``t`` and its successor ``t+1`` (so the classifier sees the
+transition), plus the causal segmentation-detector outputs at ``t``. This is the
+single featuriser used by both leave-session-out training (via the research
+label loader) and runtime/e2e inference, so there is no train/serve skew.
+
+All features are **causal**: ``t+1`` has arrived by the time the gap after ``t``
+is scored, and the segmentation features come from the online BOCPD / trailing
+majority-vote state. The acausal ``position = seq / n`` feature used by the
+earliest boundary baselines is intentionally **not** emitted here — ``n`` (total
+session length) is unknowable mid-stream.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from traceforge.phase.features import (
+    MAX_TEXT_CHARS,
+    _extract_text,
+    _numeric_dict,
+    _symbolic_dict,
+)
+from traceforge.phase.segmentation import SegmentationParams, session_segmentation_features
+
+#: Canonical fields whose change across the gap is an explicit transition signal.
+_CHANGE_FIELDS = ("mechanism", "effect", "tool_name")
+
+
+@dataclass
+class GapExample:
+    """One gap (the transition after ``after_event_id``), ready for vectorisation.
+
+    Field names mirror :class:`traceforge.phase.features.EventExample` so the
+    shared ``merged_symbolic`` / design-matrix machinery applies unchanged.
+    """
+
+    session_id: str
+    source: str
+    after_event_id: str
+    symbolic: dict[str, float]
+    numeric: dict[str, float]
+    text: str
+    seg: dict[str, float] = field(default_factory=dict)
+    nbr: dict[str, float] = field(default_factory=dict)
+    label: str = ""
+
+    @property
+    def event_id(self) -> str:
+        """Alias so predictions can be keyed identically to the phase path."""
+        return self.after_event_id
+
+
+def _gap_symbolic(cur: dict, nxt: dict | None) -> dict[str, float]:
+    """Symbolic features of the current event plus ``next_``-prefixed successor
+    features and explicit change indicators across the gap."""
+
+    sym = _symbolic_dict(cur)
+    if nxt is not None:
+        for k, v in _symbolic_dict(nxt).items():
+            sym[f"next_{k}"] = v
+        for fld in _CHANGE_FIELDS:
+            if cur.get(fld) != nxt.get(fld):
+                sym[f"changed_{fld}"] = 1.0
+    return sym
+
+
+def _gap_text(cur: dict, nxt: dict | None) -> str:
+    text = _extract_text(cur)
+    if nxt is not None:
+        text = (text + " \n " + _extract_text(nxt))[: MAX_TEXT_CHARS * 2]
+    return text
+
+
+def build_gap_example(
+    session_id: str,
+    source: str,
+    cur: dict,
+    nxt: dict | None,
+    seg: dict[str, float] | None = None,
+) -> GapExample:
+    """Build one :class:`GapExample` from a current feature row, its successor
+    (or ``None`` for the final gap), and the causal seg features computed at
+    ``cur``. Shared by the batch featuriser and the live streaming inferencer so
+    there is no train/serve skew."""
+
+    return GapExample(
+        session_id=session_id,
+        source=source,
+        after_event_id=cur["event_id"],
+        symbolic=_gap_symbolic(cur, nxt),
+        numeric=_numeric_dict(cur),  # causal: no acausal position=seq/n
+        text=_gap_text(cur, nxt),
+        seg=seg or {},
+    )
+
+
+def featurize_session_gaps(
+    session_id: str,
+    source: str,
+    events: dict[str, dict],
+    seg_params: SegmentationParams | None = None,
+) -> list[GapExample]:
+    """Featurise every gap of one session causally, in ``seq`` order.
+
+    ``events`` maps ``event_id -> feature-row dict`` (each from
+    :func:`traceforge.phase.event_rows.event_to_feature_row`). Returns one
+    :class:`GapExample` per event (the gap *after* it); the final event's gap has
+    no successor. Labels are left empty — the research loader overlays them; the
+    inference path uses the examples directly.
+    """
+
+    if not events:
+        return []
+    seq = {eid: int(ev.get("seq") or 0) for eid, ev in events.items()}
+    ordered = sorted(events, key=lambda e: seq.get(e, 0))
+    next_of = {ordered[i]: ordered[i + 1] for i in range(len(ordered) - 1)}
+
+    seg_feats: dict[str, dict[str, float]] = {}
+    if seg_params is not None:
+        seg_feats = session_segmentation_features(
+            [events[e].get("phase_signals") for e in ordered],
+            ordered,
+            seg_params,
+        )
+
+    out: list[GapExample] = []
+    for aid in ordered:
+        cur = events[aid]
+        nxt = events.get(next_of.get(aid)) if next_of.get(aid) else None
+        out.append(build_gap_example(session_id, source, cur, nxt, seg_feats.get(aid, {})))
+    return out
+
+
+__all__ = ["GapExample", "build_gap_example", "featurize_session_gaps"]

traceforge/boundary/inference.py

@@ -0,0 +1,267 @@
+"""Persist and apply the trained per-gap boundary classifier.
+
+The shipped contract is ``combined-seg``: canonical symbolic features of the gap
+(current + successor event, with change indicators), a frozen model2vec text
+embedding of both events, and the causal classical-segmentation detector outputs
+(online BOCPD posterior, trailing multi-scale majority vote). It is a 3-class
+single-label problem (``noise`` / ``activity-boundary`` / ``step-boundary``).
+
+This module owns the bundle type so the persisted joblib loads in production
+(``scikit-learn`` / ``scipy`` / ``joblib`` are core deps, imported lazily). It
+mirrors :mod:`traceforge.phase.inference`.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from .decode import DecodeParams
+
+from traceforge.phase.features import (
+    MODEL2VEC_DIM,
+    MODEL2VEC_NAME,
+    embed_texts,
+    feature_set_blocks,
+    merged_symbolic,
+)
+from traceforge.phase.segmentation import SegmentationParams
+
+from .features import GapExample, featurize_session_gaps
+
+#: Bumped whenever the bundle layout or feature contract changes.
+SCHEMA_VERSION = 1
+#: Single-label class vocabulary (order is the model's class order).
+BOUNDARY_CLASSES: tuple[str, ...] = ("noise", "activity-boundary", "step-boundary")
+#: The shipped (causal) feature contract.
+DEFAULT_FEATURE_SET = "combined-seg"
+#: Environment override for the model location.
+MODEL_PATH_ENV = "TRACEFORGE_BOUNDARY_MODEL"
+#: Packaged model location (populated at release time from the trained bundle).
+PACKAGED_MODEL_PATH = Path(__file__).resolve().parent / "data" / "boundary-model.joblib"
+
+
+def resolve_model_path(path: str | Path | None = None) -> Path:
+    """Resolve the bundle location: explicit arg → ``$TRACEFORGE_BOUNDARY_MODEL``
+    → the packaged default."""
+
+    if path is not None:
+        return Path(path)
+    env = os.environ.get(MODEL_PATH_ENV)
+    if env:
+        return Path(env)
+    return PACKAGED_MODEL_PATH
+
+
+@dataclass
+class BoundaryModel:
+    """A self-contained, fitted boundary classifier ready for inference."""
+
+    feature_set: str
+    classes: tuple[str, ...]
+    vectorizer: object  # sklearn DictVectorizer
+    scaler: object  # sklearn StandardScaler
+    estimator: object
+    seg_params: SegmentationParams | None
+    drop_prefixes: tuple[str, ...]
+    model2vec_name: str
+    model2vec_dim: int
+    n_train_examples: int
+    n_train_sessions: int
+    schema_version: int = SCHEMA_VERSION
+    #: Learned causal decode levers (per-class threshold + refractory min-gap).
+    #: ``None`` => legacy bundle; callers fall back to argmax.
+    decode_params: "DecodeParams | None" = None
+
+
+def _filter(d: dict[str, float], drop_prefixes: tuple[str, ...]) -> dict[str, float]:
+    if not drop_prefixes:
+        return d
+    return {k: v for k, v in d.items() if not k.startswith(drop_prefixes)}
+
+
+def _design_matrix(
+    feature_set: str,
+    vectorizer,
+    examples,
+    embeddings: np.ndarray,
+    drop_prefixes: tuple[str, ...],
+    *,
+    fit: bool,
+) -> np.ndarray:
+    """Assemble the (pre-scaling) feature matrix: symbolic block (DictVectorizer)
+    then the frozen model2vec block, horizontally stacked."""
+
+    from scipy import sparse
+
+    use_symbolic, use_embedding, use_seg, nbr_prefixes = feature_set_blocks(feature_set)
+    blocks: list[np.ndarray] = []
+    if use_symbolic:
+        dicts = [
+            _filter(merged_symbolic(ex, use_seg, nbr_prefixes), drop_prefixes) for ex in examples
+        ]
+        m = vectorizer.fit_transform(dicts) if fit else vectorizer.transform(dicts)
+        blocks.append(m.toarray() if sparse.issparse(m) else np.asarray(m))
+    if use_embedding:
+        blocks.append(np.asarray(embeddings))
+    return np.hstack(blocks).astype(np.float64)
+
+
+def fit_boundary_model(
+    examples,
+    feature_set: str,
+    estimator_factory,
+    seg_params: SegmentationParams | None,
+    *,
+    drop_prefixes: tuple[str, ...] = (),
+) -> BoundaryModel:
+    """Fit a boundary classifier on every labelled gap and return the bundle.
+
+    ``examples`` must already carry ``seg`` features for ``feature_set`` and a
+    non-empty ``label`` in :data:`BOUNDARY_CLASSES`.
+    """
+
+    from sklearn.feature_extraction import DictVectorizer
+    from sklearn.preprocessing import StandardScaler
+
+    embeddings = embed_texts([e.text for e in examples])
+    vectorizer = DictVectorizer(sparse=True)
+    x = _design_matrix(feature_set, vectorizer, examples, embeddings, drop_prefixes, fit=True)
+    scaler = StandardScaler().fit(x)
+    y = np.array([e.label for e in examples], dtype=object)
+    estimator = estimator_factory()
+    estimator.fit(scaler.transform(x), y)
+    classes = tuple(str(c) for c in estimator.classes_)
+    return BoundaryModel(
+        feature_set=feature_set,
+        classes=classes,
+        vectorizer=vectorizer,
+        scaler=scaler,
+        estimator=estimator,
+        seg_params=seg_params,
+        drop_prefixes=tuple(drop_prefixes),
+        model2vec_name=MODEL2VEC_NAME,
+        model2vec_dim=MODEL2VEC_DIM,
+        n_train_examples=len(examples),
+        n_train_sessions=len({e.session_id for e in examples}),
+    )
+
+
+def predict_scores(model: BoundaryModel, examples) -> np.ndarray:
+    """Per-class probabilities for each gap, aligned to ``model.classes``."""
+
+    if not examples:
+        return np.zeros((0, len(model.classes)), dtype=np.float64)
+    embeddings = embed_texts([e.text for e in examples])
+    x = _design_matrix(
+        model.feature_set, model.vectorizer, examples, embeddings, model.drop_prefixes, fit=False
+    )
+    xs = model.scaler.transform(x)
+    est = model.estimator
+    if hasattr(est, "predict_proba"):
+        return np.asarray(est.predict_proba(xs), dtype=np.float64)
+    return np.asarray(est.decision_function(xs), dtype=np.float64)
+
+
+def predict_examples(model: BoundaryModel, examples) -> list[dict]:
+    """Label each gap. Returns one dict per example: ``event_id`` (the gap's
+    ``after_event_id``), the single best ``label`` (argmax) and raw ``scores``."""
+
+    scores = predict_scores(model, examples)
+    out: list[dict] = []
+    for ex, row in zip(examples, scores):
+        top = model.classes[int(np.argmax(row))]
+        out.append(
+            {
+                "event_id": ex.event_id,
+                "label": top,
+                "scores": {c: float(s) for c, s in zip(model.classes, row)},
+            }
+        )
+    return out
+
+
+def decode_examples(model: BoundaryModel, examples) -> list[dict]:
+    """Label each gap with the learned causal decoder (per-class threshold +
+    refractory min-gap) instead of ``argmax``.
+
+    ``examples`` must be in session ``seq`` order (as produced by
+    :func:`featurize_session_gaps`) — the refractory is positional. Falls back to
+    :func:`predict_examples` (argmax) when the bundle carries no decode params.
+    """
+
+    if model.decode_params is None:
+        return predict_examples(model, examples)
+    from .decode import decode_scores
+
+    scores = predict_scores(model, examples)
+    labels = decode_scores(model.decode_params, model.classes, scores)
+    return [
+        {
+            "event_id": ex.event_id,
+            "label": lab,
+            "scores": {c: float(s) for c, s in zip(model.classes, row)},
+        }
+        for ex, row, lab in zip(examples, scores, labels)
+    ]
+
+
+def predict_session(
+    model: BoundaryModel, session_id: str, source: str, events: dict[str, dict]
+) -> list[dict]:
+    """Featurise a whole session's gaps causally and label each one. Uses the
+    learned decoder when the bundle carries decode params, else argmax."""
+
+    examples = featurize_session_gaps(session_id, source, events, model.seg_params)
+    if not examples:
+        return []
+    if model.decode_params is not None:
+        return decode_examples(model, examples)
+    return predict_examples(model, examples)
+
+
+def save(model: BoundaryModel, path: str | Path) -> Path:
+    import joblib
+
+    p = Path(path)
+    p.parent.mkdir(parents=True, exist_ok=True)
+    joblib.dump(model, p)
+    return p
+
+
+def load(path: str | Path | None = None) -> BoundaryModel:
+    import joblib
+
+    p = resolve_model_path(path)
+    if not p.exists():
+        raise FileNotFoundError(
+            f"boundary model not found at {p}. Set ${MODEL_PATH_ENV} or place the "
+            "trained bundle at the packaged path. Train it with "
+            "`python -m scripts.persist_boundary_model` in research/."
+        )
+    return joblib.load(p)
+
+
+__all__ = [
+    "BOUNDARY_CLASSES",
+    "DEFAULT_FEATURE_SET",
+    "MODEL_PATH_ENV",
+    "PACKAGED_MODEL_PATH",
+    "SCHEMA_VERSION",
+    "BoundaryModel",
+    "GapExample",
+    "featurize_session_gaps",
+    "fit_boundary_model",
+    "decode_examples",
+    "load",
+    "predict_examples",
+    "predict_scores",
+    "predict_session",
+    "resolve_model_path",
+    "save",
+]

traceforge/boundary/inferencer.py

@@ -0,0 +1,146 @@
+"""Runtime boundary inference over live :class:`SessionEvent` objects.
+
+Bridges the production event type to the trained per-gap boundary classifier and
+runs it **live, as events stream** — no waiting for the session to end. A *gap*
+is the transition after event ``t``; it is featurised from ``t`` and its
+successor ``t+1`` plus the causal segmentation state at ``t``. Because the gap
+needs the successor, the boundary it decodes is stamped on that successor — the
+event that *opens* the new activity/step — via ``metadata.boundary``. So the
+moment event ``t+1`` arrives, the gap after ``t`` is fully known and stamped with
+zero look-ahead and zero buffering.
+
+The causal decoder (per-class threshold + refractory min-gap) runs incrementally
+through :class:`traceforge.boundary.decode.StreamingBoundaryDecoder`, and the
+segmentation features come from :class:`IncrementalSegmentation`, so the streamed
+labels are identical to the batch :func:`predict_session` path (guarded by tests)
+while holding only O(1) state per session.
+
+Mirrors :mod:`traceforge.phase.inferencer`. The model is the only boundary
+producer; a missing bundle raises rather than silently degrading.
+"""
+
+from __future__ import annotations
+
+from traceforge.phase.event_rows import event_to_feature_row
+from traceforge.types import SessionEvent
+
+from .features import build_gap_example
+
+
+class BoundaryInferencer:
+    """Loads a trained boundary bundle once and applies it to live events.
+
+    Construct with an explicit ``model``/``model_path`` or rely on
+    :func:`traceforge.boundary.inference.resolve_model_path` (env var / packaged
+    default). The model loads lazily on first use.
+    """
+
+    def __init__(self, model=None, model_path=None) -> None:
+        self._model = model
+        self._model_path = model_path
+
+    @property
+    def model(self):
+        if self._model is None:
+            from .inference import load
+
+            self._model = load(self._model_path)
+        return self._model
+
+    def apply(self, event: SessionEvent, boundary: str | None) -> SessionEvent:
+        """Return a copy of ``event`` with ``metadata.boundary`` set.
+
+        ``boundary`` is an opening label (``"activity-boundary"`` /
+        ``"step-boundary"``) or ``None`` to leave the event continuing the
+        current segment. Returns the event unchanged when ``boundary`` is
+        ``None`` or the event carries no metadata.
+        """
+
+        if boundary is None or event.metadata is None:
+            return event
+        new_md = event.metadata.model_copy(update={"boundary": boundary})
+        return event.model_copy(update={"metadata": new_md})
+
+    def new_stream(self, session_id: str, source: str = "") -> "SessionBoundaryStream":
+        """Open a live per-event boundary stream for one session.
+
+        Feed events through :meth:`SessionBoundaryStream.push` in arrival order;
+        each call returns the event stamped with the opening boundary the batch
+        decoder would assign to the gap *into* it — computed online from bounded
+        incremental state.
+        """
+
+        return SessionBoundaryStream(self, session_id, source)
+
+
+class SessionBoundaryStream:
+    """Live per-event boundary stamper for a single session.
+
+    Drives :class:`~traceforge.phase.segmentation.IncrementalSegmentation` so each
+    gap's causal seg features match the batch path exactly, scores the gap after
+    the previous event against the just-arrived successor, decodes it with the
+    refractory :class:`~traceforge.boundary.decode.StreamingBoundaryDecoder`, and
+    stamps the resulting opening label on the successor. The very first event has
+    no incoming gap, so it is emitted unstamped (it opens the root segment).
+    """
+
+    def __init__(self, inferencer: "BoundaryInferencer", session_id: str, source: str) -> None:
+        from traceforge.phase.segmentation import IncrementalSegmentation
+
+        from .decode import StreamingBoundaryDecoder
+
+        self._inf = inferencer
+        self._session_id = session_id
+        self._source = source
+        model = inferencer.model
+        self._seg = (
+            IncrementalSegmentation(model.seg_params) if model.seg_params is not None else None
+        )
+        self._decoder = (
+            StreamingBoundaryDecoder(model.decode_params)
+            if model.decode_params is not None
+            else None
+        )
+        self._priority = model.decode_params.priority if model.decode_params is not None else ()
+        self._seq = 0
+        self._prev_row: dict | None = None
+        self._prev_seg: dict[str, float] = {}
+
+    def push(self, event: SessionEvent) -> SessionEvent:
+        """Ingest one event; return it stamped with its opening boundary (if any)."""
+
+        row = event_to_feature_row(event, self._seq)
+        self._seq += 1
+        seg_cur = self._seg.push(row.get("phase_signals")) if self._seg is not None else {}
+
+        boundary: str | None = None
+        if self._prev_row is not None:
+            boundary = self._decode_gap(self._prev_row, row, self._prev_seg)
+
+        self._prev_row = row
+        self._prev_seg = seg_cur
+        return self._inf.apply(event, boundary)
+
+    def _decode_gap(self, cur: dict, nxt: dict, seg: dict[str, float]) -> str | None:
+        """Score the gap (cur -> nxt) and advance the refractory decoder.
+
+        Falls back to ``argmax`` (no refractory) when the bundle carries no
+        decode params, mirroring the batch path. Returns an opening label or
+        ``None`` (noise / suppressed)."""
+
+        from .inference import predict_scores
+
+        model = self._inf.model
+        gap = build_gap_example(self._session_id, self._source, cur, nxt, seg)
+        scores = predict_scores(model, [gap])[0]
+        score_by_class = {c: float(s) for c, s in zip(model.classes, scores)}
+
+        if self._decoder is None:
+            top = max(model.classes, key=lambda c: score_by_class[c])
+            return top if top != "noise" else None
+
+        label = self._decoder.push({c: score_by_class.get(c, 0.0) for c in self._priority})
+        return label if label != "noise" else None
+
+
+__all__ = ["BoundaryInferencer", "SessionBoundaryStream"]

traceforge/classify/__init__.py

@@ -0,0 +1,95 @@
+"""Classification package — multi-dimensional tool and shell analysis."""
+
+from traceforge.classify.cmd import classify_cmd_command
+from traceforge.classify.config import (
+    ClassificationEngine,
+    ClassifyConfig,
+    get_default_engine,
+    load_config,
+    reset_default_engine,
+    set_default_engine,
+)
+from traceforge.classify.core import (
+    Action,
+    Capability,
+    Classification,
+    Effect,
+    Mechanism,
+    PhaseSegment,
+    Role,
+    Scope,
+    Structure,
+    aggregate_effect,
+)
+from traceforge.classify.coding import (
+    CodingAction,
+    CodingMechanism,
+    CodingRole,
+    CodingScope,
+    ShellDialect,
+    ShellStructure,
+)
+from traceforge.classify.mcp import (
+    McpServerProfile,
+    McpToolOverride,
+    classify_mcp_tool,
+    extract_mcp_namespace,
+)
+from traceforge.classify.powershell import classify_powershell_command
+from traceforge.classify.registry import DimensionRegistry, get_default_registry
+from traceforge.classify.risk import Confidence, RiskAssessment, assess_risk, assess_tool_risk
+from traceforge.classify.shell import (
+    _extract_commands_from_ast,
+    classify_shell,
+)
+from traceforge.classify.tools import classify_tool, normalize_tool_name
+from traceforge.classify.workflow import Phase, Visibility
+
+__all__ = [
+    # Core types
+    "Action",
+    "Capability",
+    "Classification",
+    "Effect",
+    "Mechanism",
+    "Phase",
+    "PhaseSegment",
+    "Role",
+    "Scope",
+    "ShellDialect",
+    "ShellStructure",
+    "Structure",
+    "Visibility",
+    "aggregate_effect",
+    # Coding domain
+    "CodingAction",
+    "CodingMechanism",
+    "CodingRole",
+    "CodingScope",
+    # Config
+    "ClassificationEngine",
+    "ClassifyConfig",
+    "get_default_engine",
+    "load_config",
+    "reset_default_engine",
+    "set_default_engine",
+    # Registry
+    "DimensionRegistry",
+    "get_default_registry",
+    # Classifiers
+    "_extract_commands_from_ast",
+    "classify_shell",
+    "classify_cmd_command",
+    "classify_powershell_command",
+    "classify_tool",
+    "classify_mcp_tool",
+    "extract_mcp_namespace",
+    "McpServerProfile",
+    "McpToolOverride",
+    "normalize_tool_name",
+    # Risk scoring
+    "Confidence",
+    "RiskAssessment",
+    "assess_risk",
+    "assess_tool_risk",
+]