traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/phase/inference.py

@@ -0,0 +1,279 @@
+"""Persist and apply the trained per-event phase classifier.
+
+The shipped contract is ``combined-seg-nbrcentroid``: the best *causal*
+feature-set (trailing-only neighbor centroid distance, online BOCPD, trailing
+majority-vote), leave-session-out F1_macro 0.931 — within 0.0004 of the
+look-ahead model, so inference needs no future context and runs online.
+
+This module owns the bundle type so the persisted joblib is loadable in
+production. ``scikit-learn`` / ``scipy`` / ``joblib`` are core dependencies;
+they are imported lazily only to keep the base package's import time low.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+
+from .features import (
+    MODEL2VEC_DIM,
+    MODEL2VEC_NAME,
+    PHASES,
+    EventExample,
+    NeighborParams,
+    embed_texts,
+    featurize_session_events,
+    feature_set_blocks,
+    merged_symbolic,
+)
+from .segmentation import SegmentationParams
+
+#: Bumped whenever the bundle layout or feature contract changes.
+SCHEMA_VERSION = 1
+#: The shipped (causal) feature contract.
+DEFAULT_FEATURE_SET = "combined-seg-nbrcentroid"
+#: Environment override for the model location.
+MODEL_PATH_ENV = "TRACEFORGE_PHASE_MODEL"
+#: Packaged model location (populated at release time from the trained bundle).
+PACKAGED_MODEL_PATH = Path(__file__).resolve().parent / "data" / "phase-model.joblib"
+
+
+def resolve_model_path(path: str | Path | None = None) -> Path:
+    """Resolve the bundle location: explicit arg → ``$TRACEFORGE_PHASE_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 PhaseModel:
+    """A self-contained, fitted phase classifier ready for inference."""
+
+    feature_set: str
+    classes: tuple[str, ...]
+    vectorizer: object  # sklearn DictVectorizer
+    scaler: object  # sklearn StandardScaler
+    estimator: object
+    seg_params: SegmentationParams | None
+    neighbor_params: NeighborParams | None
+    drop_prefixes: tuple[str, ...]
+    decision_threshold: float
+    model2vec_name: str
+    model2vec_dim: int
+    n_train_examples: int
+    n_train_sessions: int
+    schema_version: int = SCHEMA_VERSION
+
+
+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)
+    first, 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_phase_model(
+    examples,
+    feature_set: str,
+    estimator_factory,
+    seg_params: SegmentationParams | None,
+    neighbor_params: NeighborParams | None,
+    *,
+    drop_prefixes: tuple[str, ...] = (),
+    decision_threshold: float = 0.5,
+) -> PhaseModel:
+    """Fit a phase classifier on every labelled example and return the bundle.
+
+    ``examples`` must already carry ``seg``/``nbr`` features for ``feature_set``.
+    The same params are stored so inference recomputes identical features.
+    """
+
+    from sklearn.feature_extraction import DictVectorizer
+    from sklearn.preprocessing import MultiLabelBinarizer, 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)
+    mlb = MultiLabelBinarizer(classes=list(PHASES))
+    y = mlb.fit_transform([e.phases for e in examples])
+    estimator = estimator_factory()
+    estimator.fit(scaler.transform(x), y)
+    return PhaseModel(
+        feature_set=feature_set,
+        classes=tuple(PHASES),
+        vectorizer=vectorizer,
+        scaler=scaler,
+        estimator=estimator,
+        seg_params=seg_params,
+        neighbor_params=neighbor_params,
+        drop_prefixes=tuple(drop_prefixes),
+        decision_threshold=decision_threshold,
+        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: PhaseModel, examples) -> np.ndarray:
+    """Per-class probabilities for each example, 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: PhaseModel, examples) -> list[dict]:
+    """Stamp a phase onto each example.
+
+    Returns one dict per example: ``event_id``, the single best ``phase``
+    (argmax — the production stamp), the multi-label ``phases`` (every class
+    above ``decision_threshold``; falls back to argmax when none clear it),
+    and the raw per-class ``scores``.
+    """
+
+    scores = predict_scores(model, examples)
+    out: list[dict] = []
+    for ex, row in zip(examples, scores):
+        top = model.classes[int(np.argmax(row))]
+        multi = tuple(c for c, s in zip(model.classes, row) if s >= model.decision_threshold)
+        if not multi:
+            multi = (top,)
+        out.append(
+            {
+                "event_id": ex.event_id,
+                "phase": top,
+                "phases": multi,
+                "scores": {c: float(s) for c, s in zip(model.classes, row)},
+            }
+        )
+    return out
+
+
+def predict_session(
+    model: PhaseModel, session_id: str, source: str, events: dict[str, dict]
+) -> list[dict]:
+    """Featurise a whole session causally and predict, with plumbing inheriting.
+
+    The phase *decision* is restricted to content-bearing events (messages, tool
+    calls, reasoning). Plumbing events (lifecycle/turn/hook markers) carry no
+    intrinsic phase, so they inherit the most recent content-bearing phase
+    (leading plumbing back-fills the first content-bearing phase). A stamp is
+    still returned for **every** event so downstream sinks see a phase on each.
+
+    Featurisation runs over the full event sequence, so the trailing windows /
+    neighbor centroids are unchanged — only which events the model is *asked*
+    about differs. Inherited stamps are flagged with ``"inherited": True``.
+    """
+
+    examples = featurize_session_events(
+        session_id, source, events, model.seg_params, model.neighbor_params
+    )
+    if not examples:
+        return []
+
+    content = [ex for ex in examples if ex.content_bearing]
+    preds = predict_examples(model, content)
+    pred_by_id = {p["event_id"]: p for p in preds}
+    first_phase = preds[0]["phase"] if preds else model.classes[0]
+
+    out: list[dict] = []
+    last: dict | None = None
+    for ex in examples:
+        if ex.content_bearing and ex.event_id in pred_by_id:
+            last = pred_by_id[ex.event_id]
+            out.append(last)
+        else:
+            inherited = last["phase"] if last is not None else first_phase
+            out.append(
+                {
+                    "event_id": ex.event_id,
+                    "phase": inherited,
+                    "phases": (inherited,),
+                    "scores": {},
+                    "inherited": True,
+                }
+            )
+    return out
+
+
+def save(model: PhaseModel, 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) -> PhaseModel:
+    import joblib
+
+    p = resolve_model_path(path)
+    if not p.exists():
+        raise FileNotFoundError(
+            f"phase 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_phase_model` in research/."
+        )
+    return joblib.load(p)
+
+
+__all__ = [
+    "DEFAULT_FEATURE_SET",
+    "MODEL_PATH_ENV",
+    "PACKAGED_MODEL_PATH",
+    "SCHEMA_VERSION",
+    "EventExample",
+    "PhaseModel",
+    "fit_phase_model",
+    "load",
+    "predict_examples",
+    "predict_scores",
+    "predict_session",
+    "resolve_model_path",
+    "save",
+]

traceforge/phase/inferencer.py

@@ -0,0 +1,171 @@
+"""Runtime phase inference over live :class:`SessionEvent` objects.
+
+Bridges the production event type to the trained classifier: projects each
+event through the same :func:`event_to_feature_row` contract the corpus was
+built from, runs the causal session featuriser + bundle, and stamps the
+predicted workflow stage onto ``metadata.phase``.
+
+The trained model is the *only* phase producer — there is no deterministic
+fallback. If the model bundle is missing, inference raises rather than
+silently degrading.
+"""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+from traceforge.classify.workflow import Phase
+from traceforge.types import SessionEvent
+
+from .event_rows import event_to_feature_row
+
+
+class PhaseInferencer:
+    """Loads a trained bundle once and applies it to whole sessions.
+
+    Construct with an explicit ``model``/``model_path`` or rely on
+    :func:`traceforge.phase.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 predict(self, events: Sequence[SessionEvent]) -> list[dict]:
+        """Predict a phase for every event of one session (in given order).
+
+        Events must belong to a single session and be in session sequence
+        order. Returns the per-event prediction dicts from
+        :func:`traceforge.phase.inference.predict_examples`.
+        """
+
+        from .inference import predict_session
+
+        ordered = list(events)
+        if not ordered:
+            return []
+        rows: dict[str, dict] = {}
+        for seq, ev in enumerate(ordered):
+            row = event_to_feature_row(ev, seq)
+            rows[row["event_id"]] = row
+        first_md = ordered[0].metadata
+        source = (first_md.source_framework if first_md else None) or ""
+        return predict_session(self.model, ordered[0].session_id, source, rows)
+
+    def stamp(self, events: Sequence[SessionEvent]) -> list[SessionEvent]:
+        """Return copies of ``events`` with ``metadata.phase`` set by the model."""
+
+        preds = {p["event_id"]: p for p in self.predict(events)}
+        out: list[SessionEvent] = []
+        for ev in events:
+            p = preds.get(ev.id)
+            out.append(self.apply(ev, p["phase"] if p else None))
+        return out
+
+    def apply(self, event: SessionEvent, phase: str | None) -> SessionEvent:
+        """Return a copy of ``event`` with ``metadata.phase`` set to ``phase``.
+
+        Returns the event unchanged when ``phase`` is ``None`` or the event has
+        no metadata. Used by the streaming pipeline to stamp one event at a time.
+        """
+
+        if phase is None or event.metadata is None:
+            return event
+        new_md = event.metadata.model_copy(update={"phase": Phase(phase)})
+        return event.model_copy(update={"metadata": new_md})
+
+    def is_content_bearing(self, event: SessionEvent) -> bool:
+        """Whether ``event`` carries an intrinsic phase the model predicts.
+
+        Plumbing events (lifecycle/turn/hook markers) are *not* content-bearing;
+        in the live stream they inherit the prevailing content phase instead of
+        being classified.
+        """
+
+        from .features import is_content_bearing as _is_content_bearing
+
+        return _is_content_bearing(event_to_feature_row(event, 0))
+
+    def new_stream(self, session_id: str, source: str = "") -> "SessionPhaseStream":
+        """Open a live per-event phase stream for one session.
+
+        Feed events through :meth:`SessionPhaseStream.push` in arrival order;
+        each call returns the events ready to emit, stamped with the phase the
+        batch path would assign — computed online from bounded incremental state
+        (no waiting for the session to end, no moving-window resets).
+        """
+
+        return SessionPhaseStream(self, session_id, source)
+
+
+class SessionPhaseStream:
+    """Live per-event phase stamper for a single session.
+
+    Drives :class:`~traceforge.phase.features.StreamingSessionFeaturizer` so each
+    event's causal features (segmentation BOCPD posterior, trailing centroids,
+    windowed majority/entropy) are carried forward exactly, then classifies and
+    stamps content-bearing events the instant they arrive. Plumbing events
+    inherit the prevailing content phase. Only contiguous *leading* plumbing is
+    held (so it can inherit the first content phase, matching the batch path);
+    everything after the first content event is emitted one event at a time.
+    """
+
+    def __init__(self, inferencer: "PhaseInferencer", session_id: str, source: str) -> None:
+        from .features import StreamingSessionFeaturizer
+
+        self._inf = inferencer
+        model = inferencer.model
+        self._classes = model.classes
+        self._feat = StreamingSessionFeaturizer(
+            session_id, source, model.seg_params, model.neighbor_params
+        )
+        self._seq = 0
+        self._seen_content = False
+        self._leading: list[SessionEvent] = []
+        self._last_phase: str | None = None
+
+    def _predict(self, example) -> str:
+        from .inference import predict_examples
+
+        return predict_examples(self._inf.model, [example])[0]["phase"]
+
+    def push(self, event: SessionEvent) -> list[SessionEvent]:
+        """Ingest one event; return the events now ready to emit (stamped)."""
+
+        row = event_to_feature_row(event, self._seq)
+        self._seq += 1
+        example = self._feat.push(row["event_id"], row)
+
+        if not self._seen_content:
+            self._leading.append(event)
+            if example.content_bearing:
+                self._last_phase = self._predict(example)
+                out = [self._inf.apply(ev, self._last_phase) for ev in self._leading]
+                self._leading = []
+                self._seen_content = True
+                return out
+            return []
+
+        if example.content_bearing:
+            self._last_phase = self._predict(example)
+        return [self._inf.apply(event, self._last_phase)]
+
+    def flush(self) -> list[SessionEvent]:
+        """Emit any held leading plumbing for a session that never produced a
+        content event (matches the batch fallback: inherit the first class)."""
+
+        if self._seen_content or not self._leading:
+            return []
+        phase = self._classes[0]
+        out = [self._inf.apply(ev, phase) for ev in self._leading]
+        self._leading = []
+        return out