perf: batch embeddings in background thread to fix SSE timeouts

Root cause: 306 embedding calls at 61ms each blocked the request
thread for ~19s before forwarding to Anthropic.

- Batch all admitted objects into single embed_batch() call
- Run in background thread (non-blocking)
- store_object accepts pre-computed embeddings
- Goal detection uses turn heuristic instead of blocking embed
This commit is contained in:
Joey Yakimowich-Payne 2026-03-15 19:29:23 -06:00
commit 4f50359d01
2 changed files with 54 additions and 36 deletions

View file

@ -1983,6 +1983,7 @@ def create_app(
new_objects = segmented[len(session._segmented_objects) :] new_objects = segmented[len(session._segmented_objects) :]
admitted_count = 0 admitted_count = 0
rejected_count = 0 rejected_count = 0
admitted_objects = []
for seg_obj in new_objects: for seg_obj in new_objects:
# Phase 4d: admission control — score and gate # Phase 4d: admission control — score and gate
has_dup = False has_dup = False
@ -2005,23 +2006,45 @@ def create_app(
rejected_count += 1 rejected_count += 1
continue continue
admitted_count += 1 admitted_count += 1
with Timer(session.benchmark.latency["embedding"]): admitted_objects.append(seg_obj)
# Batch embed + store in a single background thread.
# One embed_batch() call replaces N individual embed() calls.
if admitted_objects:
_batch = list(admitted_objects) # snapshot
def _store_batch(s=session, objs=_batch):
try:
embedder = s.object_store._embedder
embeddings: list[list[float]] = []
if embedder is not None:
texts = [o.content for o in objs]
with Timer(s.benchmark.latency["embedding"]):
embeddings = embedder.embed_batch(texts)
for i, so in enumerate(objs):
emb = embeddings[i] if i < len(embeddings) else []
_run_async( _run_async(
session.object_store.store_object( s.object_store.store_object(
session_id=session.id, session_id=s.id,
content=seg_obj.content, content=so.content,
object_type=seg_obj.object_type, object_type=so.object_type,
source_tool=seg_obj.source_tool, source_tool=so.source_tool,
source_key=seg_obj.source_key, source_key=so.source_key,
stub=seg_obj.stub, stub=so.stub,
tags=seg_obj.tags, tags=so.tags,
key_entities=seg_obj.key_entities, key_entities=so.key_entities,
turn=seg_obj.turn_start, turn=so.turn_start,
embedding=emb,
) )
) )
except Exception:
pass # non-critical — object store is for search
_threading.Thread(target=_store_batch, daemon=True).start()
for _so in admitted_objects:
session.benchmark.segmentation.record_object( session.benchmark.segmentation.record_object(
seg_obj.object_type, _so.object_type,
max(1, len(seg_obj.content) // 4), max(1, len(_so.content) // 4),
) )
session._segmented_objects = segmented session._segmented_objects = segmented
parts = [ parts = [
@ -2104,21 +2127,13 @@ def create_app(
) )
break break
if user_text and session.object_store._embedder is not None: # Skip embedding-based goal detection in the hot path — it takes
from mnemosyne.object_store import _cosine_similarity # ~60ms per call and blocks the request thread. Use turn-based
# heuristic instead: classify goal every 10 turns.
embedder = session.object_store._embedder
current_embedding = embedder.embed(user_text)
# Check for topic shift
goal_changed = False goal_changed = False
if session._last_user_embedding is not None: turn = session.token_state.get("turn", 0)
sim = _cosine_similarity(current_embedding, session._last_user_embedding) if turn <= 1 or turn % 10 == 0:
goal_changed = sim < 0.5 goal_changed = True
else:
goal_changed = True # First message — classify
session._last_user_embedding = current_embedding
if goal_changed and helper_llm is not None: if goal_changed and helper_llm is not None:
session.benchmark.goals.record_topic_shift() session.benchmark.goals.record_topic_shift()

View file

@ -409,10 +409,12 @@ class ObjectStore:
tags: list[str] | None = None, tags: list[str] | None = None,
key_entities: list[str] | None = None, key_entities: list[str] | None = None,
turn: int | None = None, turn: int | None = None,
embedding: list[float] | None = None,
) -> StoredObject: ) -> StoredObject:
"""Create and store a new semantic object. """Create and store a new semantic object.
Auto-generates: ID, timestamps, token estimates, embedding, stub. Auto-generates: ID, timestamps, token estimates, embedding, stub.
Pass ``embedding`` to skip the per-object embed call (for batch use).
""" """
now = datetime.now(timezone.utc).isoformat() now = datetime.now(timezone.utc).isoformat()
object_id = uuid.uuid4().hex[:16] object_id = uuid.uuid4().hex[:16]
@ -422,8 +424,9 @@ class ObjectStore:
preview = content[:80].replace("\n", " ") preview = content[:80].replace("\n", " ")
stub = f"{object_type}: {preview}..." stub = f"{object_type}: {preview}..."
# Compute embedding # Use pre-computed embedding or compute one
embedding: list[float] = [] if embedding is None:
embedding = []
if self._embedder is not None: if self._embedder is not None:
embedding = self._embedder.embed(content) embedding = self._embedder.embed(content)