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:
parent
0135298966
commit
4f50359d01
2 changed files with 54 additions and 36 deletions
|
|
@ -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,24 +2006,46 @@ 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)
|
||||||
_run_async(
|
|
||||||
session.object_store.store_object(
|
# Batch embed + store in a single background thread.
|
||||||
session_id=session.id,
|
# One embed_batch() call replaces N individual embed() calls.
|
||||||
content=seg_obj.content,
|
if admitted_objects:
|
||||||
object_type=seg_obj.object_type,
|
_batch = list(admitted_objects) # snapshot
|
||||||
source_tool=seg_obj.source_tool,
|
|
||||||
source_key=seg_obj.source_key,
|
def _store_batch(s=session, objs=_batch):
|
||||||
stub=seg_obj.stub,
|
try:
|
||||||
tags=seg_obj.tags,
|
embedder = s.object_store._embedder
|
||||||
key_entities=seg_obj.key_entities,
|
embeddings: list[list[float]] = []
|
||||||
turn=seg_obj.turn_start,
|
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(
|
||||||
|
s.object_store.store_object(
|
||||||
|
session_id=s.id,
|
||||||
|
content=so.content,
|
||||||
|
object_type=so.object_type,
|
||||||
|
source_tool=so.source_tool,
|
||||||
|
source_key=so.source_key,
|
||||||
|
stub=so.stub,
|
||||||
|
tags=so.tags,
|
||||||
|
key_entities=so.key_entities,
|
||||||
|
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(
|
||||||
|
_so.object_type,
|
||||||
|
max(1, len(_so.content) // 4),
|
||||||
)
|
)
|
||||||
session.benchmark.segmentation.record_object(
|
|
||||||
seg_obj.object_type,
|
|
||||||
max(1, len(seg_obj.content) // 4),
|
|
||||||
)
|
|
||||||
session._segmented_objects = segmented
|
session._segmented_objects = segmented
|
||||||
parts = [
|
parts = [
|
||||||
f"{admitted_count} new objects ({len(segmented)} total)",
|
f"{admitted_count} new objects ({len(segmented)} total)",
|
||||||
|
|
@ -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
|
goal_changed = False
|
||||||
current_embedding = embedder.embed(user_text)
|
turn = session.token_state.get("turn", 0)
|
||||||
|
if turn <= 1 or turn % 10 == 0:
|
||||||
# Check for topic shift
|
goal_changed = True
|
||||||
goal_changed = False
|
|
||||||
if session._last_user_embedding is not None:
|
|
||||||
sim = _cosine_similarity(current_embedding, session._last_user_embedding)
|
|
||||||
goal_changed = sim < 0.5
|
|
||||||
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()
|
||||||
|
|
|
||||||
|
|
@ -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,10 +424,11 @@ 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:
|
||||||
if self._embedder is not None:
|
embedding = []
|
||||||
embedding = self._embedder.embed(content)
|
if self._embedder is not None:
|
||||||
|
embedding = self._embedder.embed(content)
|
||||||
|
|
||||||
obj = StoredObject(
|
obj = StoredObject(
|
||||||
id=object_id,
|
id=object_id,
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue