jatin.blog ~ $
$ cat ai-engineering/episode-segmentation.md

Episode Segmentation and Salience Scoring

How episode boundaries and salience scores shape long-term memory retrieval.

Jatin Bansal@blog:~/ai-engineering$ open episode-segmentation

An episodic store needs boundaries that keep one task coherent without burying it inside an entire session. A conversation that moves from authentication debugging to deployment and then retry logging should produce separate retrievable episodes. Salience scoring then controls how strongly each episode influences later recall.

Segmentation signals

The five practical signals, ordered from cheapest to most-sophisticated. Production systems usually hybridize three or four of them; the goal of knowing them individually is to know which signal is failing when the segmentor mis-fires.

1; Fixed-window timeout (the baseline). Every N turns (typically 10-20) or every M minutes of clock-time, emit a boundary. Cheapest possible policy; zero model calls; trivially implementable. The failure mode is the mid-task split. A single coherent sub-task that runs longer than N turns gets cut in half, and the second half loses the context of the first half. Worth shipping as the floor of any segmentation system; it guarantees segments don’t grow unboundedly even when other signals fail.

2; Semantic-shift threshold. Embed each turn (or each user+assistant exchange) and compute the cosine distance to a running representation of the current episode; typically the mean embedding of the last K turns. When the distance exceeds a threshold (commonly 0.3-0.5 with Sentence-Transformers-style embeddings), emit a boundary. Costs one embedding call per turn (~$0.0001 at current rates) and a vector arithmetic op. The failure mode is the gradual drift. A conversation that slowly transitions from one topic to another over five turns never triggers the threshold on any single turn, and the boundary lands in the wrong place or is missed entirely. Mitigation: combine with the fixed-window floor to bound the worst case, and decay the running mean toward the most recent turns to make the threshold more reactive to sustained shifts.

3; Prediction-error boundary. Maintain a one-line running summary of the current episode (kept in working memory). Before each new turn, the model is implicitly predicting “what comes next in this conversation”; when the new turn diverges from that prediction above a threshold, emit a boundary. Concretely: every K turns, run a small-model prompt; “Here’s the current episode summary: <summary>. Here’s the next turn: <turn>. Does this turn continue the episode or start a new one?” The model emits continue|new|sub-episode plus an updated summary. Costs a small-model call every K turns (K=3-5 in practice). This is the closest direct port of the cognitive-neuroscience prediction-error mechanism, and it catches the gradual-drift case that pure semantic-shift misses; by the time five turns have drifted from the summary, the summary is the wrong anchor and the prediction-error flag fires explicitly.

4; Structural turn-type marker. Tag each turn with its type (user-question, user-correction, assistant-answer, tool-call, tool-error, tool-success, system-event) and apply rules: a tool-error followed by a user-question opens a new sub-episode; a system-event (new file uploaded, new tool granted) is always a boundary; a user-correction typically modifies the current episode rather than starting a new one. Costs nothing at runtime if the turn types are emitted by the orchestration layer (which a well-built agent harness already does); costs a small-classifier call per turn otherwise. The failure mode is type-fidelity; if the turn-type tagger is wrong, the structural signal misleads the segmentor; mitigation is to treat structural markers as hints to the higher-level segmentor rather than authoritative boundaries.

5; Agent-emitted marker. The agent itself emits explicit <episode_boundary> or <new_subtask> tokens during its own reasoning, and the segmentor reads them as authoritative. Costs essentially nothing (no separate model call; the agent emits the marker as part of its normal output) and has the highest fidelity because the agent is the closest possible observer of its own task structure. The failure mode is agent forgetting to emit. The agent gets absorbed in a sub-task and never marks the boundary; mitigation is to combine with signals 1-4 as fallbacks, so a missed agent marker is caught by the running prediction-error or semantic-shift detector. This is analogous to the explicit boundaries exposed by Anthropic’s context editing and compaction tooling; letting the agent self-segment is the highest-quality signal when it works, but never the only signal.

The defensible production stack runs all five: a structural-marker fast path (signal 4) catches the easy cases at zero cost; an agent-emitted marker (signal 5) catches the cases the agent self-identifies; a periodic prediction-error check (signal 3) every 3-5 turns catches the gradual drifts; a continuous semantic-shift detector (signal 2) catches sharp shifts the others miss; a fixed-window timeout (signal 1) bounds the worst case. Any single signal alone has a known failure mode; the layered combination handles each other’s blind spots.

Salience scoring

The salience score is the second output of the segmentation pass: once you have an episode, how much should it weigh on future recall? The canonical formulation, from Park et al.’s Generative Agents paper, is the 1-10 poignancy prompt with explicit endpoint anchoring. The exact text from the paper’s §A.1:

“On the scale of 1 to 10, where 1 is purely mundane (e.g., brushing teeth, making bed) and 10 is extremely poignant (e.g., a break up, college acceptance), rate the likely poignancy of the following piece of memory. Memory: <episode>. Rating: <fill in>.”

The anchoring is the whole architecture. Without it, LLM-rated importance scores cluster at 6-8 regardless of content, the rerank’s importance term becomes near-constant, and the entire α·recency + β·importance + γ·similarity formula collapses to α·recency + γ·similarity; losing one-third of its information. With anchoring, the model has explicit endpoints to calibrate against; “is this more like brushing teeth or more like college acceptance?”, and the resulting distribution is wide enough that the rerank can actually use it. The anchoring works because the model has read enough text to know what a 1 looks like and what a 10 looks like; pinning the scale lets the rest of the range calibrate itself against the named endpoints.

A useful salience scorer needs the following properties.

Pattern 1; Anchor both endpoints with concrete examples. “1 = mundane, 10 = significant” is too abstract; “1 = brushing teeth, 10 = getting divorced” pins the scale. Use endpoints that are vivid enough that the model can’t read them as the same thing. For an agent in a specific domain, replace the Generative Agents’ personal-life endpoints with domain-appropriate ones; “1 = a one-line clarifying question, 10 = a critical production incident root-cause” for an SRE agent; “1 = boilerplate getter method, 10 = the core algorithm of the system” for a coding agent. Domain-specific anchors produce sharper scores than generic ones.

Pattern 2; Make the score deterministic. Run the scorer with temperature=0 and a tightly-constrained output schema (Literal[1, 2, ..., 10] or JSON-schema-constrained). Free-form generation will produce “7-8, depending on context” answers that the downstream pipeline has to parse; structured output forces a single integer and the rerank can rely on it.

Pattern 3; Score at the right unit. Score the episode, not the turn. Per-turn scoring is the most common production mistake; it inflates the rerank’s bias toward whichever turns happened to score high and makes the salience term roughly equivalent to a noisy reweighting of similarity. Per-episode scoring is what the Generative Agents paper actually does, and it’s the unit at which “how much did this matter?” is a well-posed question. The segmentor emits boundaries first; the salience scorer runs once per emitted segment.

A non-obvious refinement from the MemoryBank paper (Zhong, Guo et al., AAAI 2024) and the Generative Agents follow-up work: the salience score is updated at retrieval time. Each successful retrieval of an episode increments a recall-count or resets a recency-decay term; exactly the LRU/LFU hybrid the hierarchical memory piece named. The write-time score is the starting weight; the read pattern adjusts it. An episode written with importance=3 that gets retrieved 50 times has demonstrated more salience than the original score reflects; the maintenance pass should reflect that. Conversely, an importance=8 episode that’s never retrieved is a sign the write-time score was too high or the topic became irrelevant; either way it should decay.

Segment-and-score pipeline in Python

The smallest interesting build: a streaming segmentor that combines signals 1, 2, and 3 (fixed-window, semantic-shift, prediction-error) and a Generative-Agents-style salience scorer. Writes to Chroma for the episodic store. Uses the Anthropic SDK for both the prediction-error check and the salience score, Sentence-Transformers for the embedding-based semantic-shift detector. Install: pip install anthropic chromadb sentence-transformers.

python
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
import os
import json
import time
import uuid
from dataclasses import dataclass, field
from typing import Literal

import chromadb
import numpy as np
from anthropic import Anthropic
from sentence_transformers import SentenceTransformer

client = Anthropic()
embed_model = SentenceTransformer("all-MiniLM-L6-v2")
chroma = chromadb.PersistentClient(path="./episode_store")
episodes = chroma.get_or_create_collection("episodes")

# --------- tunables ---------
MAX_SEGMENT_TURNS = 15           # signal 1 — fixed-window floor
SEMANTIC_SHIFT_THRESHOLD = 0.45  # signal 2 — cosine distance to running mean
PREDICTION_CHECK_EVERY = 4       # signal 3 — small-model boundary check cadence
SCORER_MODEL = "claude-haiku-4-5"
SUMMARIZER_MODEL = "claude-haiku-4-5"


@dataclass
class Turn:
    actor: Literal["user", "assistant", "tool"]
    text: str
    ts: float


@dataclass
class Segment:
    user: str
    turns: list[Turn] = field(default_factory=list)
    embeds: list[np.ndarray] = field(default_factory=list)
    summary: str = ""  # one-line running summary
    opened_at: float = field(default_factory=time.time)


# --------- signal 2: semantic-shift detector ---------
def semantic_shift(seg: Segment, turn_embed: np.ndarray) -> float:
    """Cosine distance of new turn to running mean of segment."""
    if not seg.embeds:
        return 0.0
    running = np.mean(seg.embeds[-5:], axis=0)
    cos_sim = np.dot(running, turn_embed) / (
        np.linalg.norm(running) * np.linalg.norm(turn_embed) + 1e-9
    )
    return 1.0 - cos_sim


# --------- signal 3: prediction-error check via small model ---------
def prediction_error_check(seg: Segment, new_turn: Turn) -> tuple[bool, str]:
    """Returns (is_boundary, updated_summary)."""
    if not seg.summary:
        # First check: build the running summary, no boundary yet.
        return False, _summarize_segment(seg)

    prompt = (
        f"Current episode summary: {seg.summary}\n\n"
        f"New turn ({new_turn.actor}): {new_turn.text}\n\n"
        "Does this turn continue the current episode or start a new one?\n"
        'Return JSON: {"decision":"continue"|"new","updated_summary":"<one line>"}'
    )
    resp = client.messages.create(
        model=SCORER_MODEL,
        max_tokens=200,
        temperature=0,
        messages=[{"role": "user", "content": prompt}],
    )
    raw = resp.content[0].text.strip()
    try:
        parsed = json.loads(raw[raw.index("{") : raw.rindex("}") + 1])
    except (ValueError, json.JSONDecodeError):
        return False, seg.summary
    return parsed["decision"] == "new", parsed["updated_summary"]


def _summarize_segment(seg: Segment) -> str:
    if not seg.turns:
        return ""
    transcript = "\n".join(f"{t.actor}: {t.text}" for t in seg.turns[-6:])
    resp = client.messages.create(
        model=SUMMARIZER_MODEL,
        max_tokens=80,
        temperature=0,
        messages=[
            {
                "role": "user",
                "content": (
                    "Summarize the topic of this dialogue in one short line "
                    f"(<=15 words):\n{transcript}"
                ),
            }
        ],
    )
    return resp.content[0].text.strip()


# --------- salience scorer (anchored 1-10) ---------
def score_salience(seg: Segment, domain: str = "general") -> int:
    transcript = "\n".join(f"{t.actor}: {t.text}" for t in seg.turns)
    if domain == "coding":
        anchors = (
            "1 = trivial syntactic clarification or boilerplate; "
            "10 = the core algorithm of the system or a critical production fix"
        )
    else:
        anchors = (
            "1 = purely mundane (brushing teeth, making bed); "
            "10 = extremely poignant (a break up, college acceptance)"
        )
    prompt = (
        f"On the scale of 1 to 10, where {anchors}, rate the likely poignancy "
        f"of the following episode for future recall.\n\nEpisode:\n{transcript}\n\n"
        'Return JSON: {"score": <integer 1-10>}'
    )
    resp = client.messages.create(
        model=SCORER_MODEL,
        max_tokens=80,
        temperature=0,
        messages=[{"role": "user", "content": prompt}],
    )
    raw = resp.content[0].text.strip()
    try:
        return int(json.loads(raw[raw.index("{") : raw.rindex("}") + 1])["score"])
    except (ValueError, json.JSONDecodeError, KeyError):
        return 5


# --------- close-segment and persist ---------
def close_and_persist(seg: Segment, domain: str = "general") -> str:
    if not seg.turns:
        return ""
    salience = score_salience(seg, domain)
    transcript = "\n".join(f"{t.actor}: {t.text}" for t in seg.turns)
    seg_embed = embed_model.encode(seg.summary or transcript[:500]).tolist()
    seg_id = str(uuid.uuid4())
    episodes.add(
        ids=[seg_id],
        embeddings=[seg_embed],
        documents=[transcript],
        metadatas=[
            {
                "user": seg.user,
                "summary": seg.summary,
                "salience": salience / 10.0,
                "opened_at": seg.opened_at,
                "closed_at": time.time(),
                "n_turns": len(seg.turns),
                "last_read": time.time(),
            }
        ],
    )
    return seg_id


# --------- streaming segmentor ---------
class Segmentor:
    def __init__(self, user: str, domain: str = "general") -> None:
        self.user = user
        self.domain = domain
        self.current = Segment(user=user)
        self.turn_count = 0

    def observe(self, turn: Turn) -> str | None:
        """Process one turn. Returns segment_id if a boundary closed a segment."""
        self.current.turns.append(turn)
        emb = embed_model.encode(turn.text)
        self.current.embeds.append(emb)
        self.turn_count += 1

        boundary = False
        # signal 1 — fixed-window floor
        if len(self.current.turns) >= MAX_SEGMENT_TURNS:
            boundary = True
        # signal 2 — semantic shift
        elif semantic_shift(self.current, emb) > SEMANTIC_SHIFT_THRESHOLD:
            boundary = True
        # signal 3 — prediction-error check (cheaper cadence)
        elif self.turn_count % PREDICTION_CHECK_EVERY == 0:
            is_new, updated_summary = prediction_error_check(self.current, turn)
            self.current.summary = updated_summary
            if is_new:
                boundary = True

        if not boundary:
            return None

        # The current turn opens the *new* segment; close the prior one.
        closing = self.current
        closing.turns = closing.turns[:-1]
        closing.embeds = closing.embeds[:-1]
        seg_id = close_and_persist(closing, self.domain)
        self.current = Segment(user=self.user)
        self.current.turns.append(turn)
        self.current.embeds.append(emb)
        return seg_id

    def flush(self) -> str | None:
        """Close any open segment (e.g., at session-end)."""
        if not self.current.turns:
            return None
        seg_id = close_and_persist(self.current, self.domain)
        self.current = Segment(user=self.user)
        return seg_id


# --------- minimal usage ---------
if __name__ == "__main__":
    seg = Segmentor(user="alice", domain="coding")
    convo = [
        ("user", "I'm getting a 401 on the auth callback. Stack trace?"),
        ("assistant", "Looks like the JWT signing key got rotated."),
        ("user", "Where do I check that?"),
        ("assistant", "config/auth.yml, look for jwt_secret_rotation."),
        ("user", "Yep. The rotation timestamp is from yesterday."),
        # Sharp topic pivot here
        ("user", "Different question — how do I bump the prod deploy version?"),
        ("assistant", "Edit infra/prod/version.tf and run terraform apply."),
        ("user", "Got it, applying now."),
    ]
    for actor, text in convo:
        sid = seg.observe(Turn(actor=actor, text=text, ts=time.time()))
        if sid:
            print(f"Closed segment: {sid}")
    final = seg.flush()
    if final:
        print(f"Final segment: {final}")

Segments are scoped per user so retrieval can enforce tenant isolation. Prediction-error checks run every few turns because model calls are expensive; fixed-window and semantic-shift checks cover the intervening turns. Salience is normalized to [0, 1] at write time to match the downstream reranker.