Rank the memory bank without reading the memory bank

Retrieval walked adventure.memories, so every turn loaded every row of the
bank with its vector attached -- 3.1 MB, 96% of everything a turn read. It
now asks SQL which memories are in play (an id and a flag per row), ranks
against vectors held in process, and fetches text only for the five it picks.

Two more callers were doing the same thing and the production SQL could not
see them: _evict_over_capacity walked the bank to count it, and _embed_pending
walked it to find the rows with no vector. Both are counts and filters the
database can do without sending anything back.

    one turn      3,258.7 kB -> 723.4 kB cold, 122.3 kB warm
    run_post_turn 3,139.1 kB -> 0.7 kB
    Insights      3,223.7 kB -> 117.9 kB
    Memories drawer  ~3.1 MB -> 23.6 kB

A played turn is turn plus post-turn work: 6.4 MB down to 123 kB.

The cache needs no invalidation callbacks, which is what makes it safe. A
vector can only change through set_vector, which drops that one entry;
anything that removes a memory from play leaves the catalogue query, and
entries missing from the catalogue are dropped on the next read. So eviction,
deletion and pruning have nothing to remember to call.

memories.embedded joins the blob, for the same reason actions.variant_count
sits beside actions.variants: with the vector deferred, every "is this
embedded?" check would otherwise be a 6 KB lazy load, once per row.

Capacity drops 200 -> 80, on retrieval quality as much as cost -- ranking two
hundred memories to pick five buries the five. Eviction was measured at scale
first: trimming 100 to 80 costs 0.8 kB and reads no vectors.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CYEJKobJ2Re4Dv7qUoSA7
This commit is contained in:
parththakkar106
2026-08-16 21:28:00 +05:30
co-authored by Claude Opus 5
parent c56864877a
commit b7e53ae581
9 changed files with 690 additions and 60 deletions
+157 -26
View File
@@ -20,8 +20,11 @@ on a later turn because the cursors only advance on success.
"""
import asyncio
from array import array
from collections import OrderedDict
from sqlalchemy.orm import Session
from sqlalchemy import func, select, update
from sqlalchemy.orm import Session, object_session
from . import models, vectors
from .context import history, story_actions, truncate_to_last_tokens
@@ -82,12 +85,73 @@ def embedding_provider(settings: models.Settings) -> OpenAICompatibleProvider:
def set_vector(memory: models.Memory, vector: list[float] | None) -> None:
"""Store (or clear) a memory's embedding.
Both columns, always together: `embedding_blob` is what will be read, and
the JSON `embedding` stays correct behind it until the follow-up migration
drops it. Going through one function is what keeps them from drifting.
Every column that describes the vector moves together: `embedding_blob` is
what the ranking reads, `embedded` is the flag everything else reads, and
the JSON `embedding` stays correct behind both until the follow-up
migration drops it. Going through one function is what keeps them in step —
and it is also the only place a stored vector can change, which is what
makes the cache below safe to invalidate here and nowhere else.
"""
memory.embedding = vector
memory.embedding_blob = None if vector is None else vectors.pack(vector)
memory.embedded = vector is not None
cached = _vector_cache.get(memory.adventure_id)
if cached is not None:
cached.pop(memory.id, None)
# ---------- The vector cache ----------
# adventure id -> {memory id: vector}, most-recently-used last.
#
# Turns for one adventure arrive back to back, and the bank barely changes
# between them, so re-reading every vector each turn is the same 600 KB over
# and over. Vectors are held as array("f") — 4 bytes a component, the same
# 6 KB the column holds. A list of Python floats would be eight times that.
#
# Correctness rests on two things. Anything that *changes* a vector goes
# through set_vector, which drops that one entry. Anything that *removes* a
# memory from play — eviction, deletion, pruning, an edit clearing the vector —
# takes it out of the catalogue query below, and entries missing from the
# catalogue are dropped on the next read. So nothing has to remember to call an
# invalidate, which is the failure this design is chosen to avoid.
#
# In-process, so it assumes one worker. That is what the deploy runs; a second
# worker would each keep their own copy and both would still be correct on
# eviction and deletion, but a vector rewritten by one could go stale in the
# other until that memory next leaves the catalogue.
_vector_cache: OrderedDict[int, dict[int, array]] = OrderedDict()
VECTOR_CACHE_ADVENTURES = 8 # ~600 KB each at a 100-memory bank
def forget_cached_vectors(adventure_id: int) -> None:
"""Drop an adventure's cached vectors. Only needed when the adventure
itself goes away — everything else self-corrects (see above)."""
_vector_cache.pop(adventure_id, None)
def _vectors_for(db: Session, adventure_id: int, ids: list[int]) -> dict[int, array]:
"""The vectors for `ids`, reading only the ones not already held."""
cached = _vector_cache.get(adventure_id)
if cached is None:
cached = _vector_cache[adventure_id] = {}
_vector_cache.move_to_end(adventure_id)
while len(_vector_cache) > VECTOR_CACHE_ADVENTURES:
_vector_cache.popitem(last=False)
wanted = set(ids)
for gone in set(cached) - wanted:
del cached[gone]
missing = [memory_id for memory_id in ids if memory_id not in cached]
if missing:
rows = db.execute(
select(models.Memory.id, models.Memory.embedding_blob)
.where(models.Memory.id.in_(missing))
).all()
for memory_id, blob in rows:
if blob:
cached[memory_id] = vectors.unpack(blob)
return cached
def settled_count(adventure: models.Adventure) -> int:
@@ -205,9 +269,22 @@ async def retrieve_memories(
return None
if not settings.embedding_model.strip():
return {"used": [], "error": "No embedding model configured in Settings."}
db = object_session(adventure)
if db is None:
return {"used": [], "error": None}
candidates = [m for m in adventure.memories if not m.forgotten and m.embedding]
if not candidates:
# Which memories are in play, and nothing else about them. This used to
# walk adventure.memories, which loaded every row of the bank *including
# its vector* — ~31 KB a memory, three megabytes a turn, 96% of everything
# a turn read. Two ids and a flag per row is about eight bytes.
catalogue = db.execute(
select(models.Memory.id, models.Memory.pinned).where(
models.Memory.adventure_id == adventure.id,
models.Memory.forgotten.is_(False),
models.Memory.embedded.is_(True),
)
).all()
if not catalogue:
return {"used": [], "error": None}
recent = history.tail(adventure, RETRIEVAL_WINDOW_ACTIONS, exclude_action_id)
@@ -222,29 +299,51 @@ async def retrieve_memories(
except ProviderError as exc:
return {"used": [], "error": str(exc)}
held = _vectors_for(db, adventure.id, [memory_id for memory_id, _ in catalogue])
scored = sorted(
((cosine(query_vec, m.embedding), m) for m in candidates),
key=lambda pair: pair[0],
(
(cosine(query_vec, held[memory_id]), memory_id, pinned)
for memory_id, pinned in catalogue
if memory_id in held
),
key=lambda row: row[0],
reverse=True,
)
# Pinned memories are always used and count toward top_k, so the injected
# set never exceeds the configured budget (unless pinned alone exceed it).
top_k = max(1, settings.memory_top_k)
used = [(score, m) for score, m in scored if m.pinned]
used = [row for row in scored if row[2]]
remaining = max(0, top_k - len(used))
used += [(score, m) for score, m in scored if not m.pinned][:remaining]
used.sort(key=lambda pair: pair[0], reverse=True)
used += [row for row in scored if not row[2]][:remaining]
used.sort(key=lambda row: row[0], reverse=True)
if not used:
return {"used": [], "error": None}
# Only now, for at most top_k rows, is the text worth fetching.
used_ids = [memory_id for _, memory_id, _ in used]
texts = dict(
db.execute(
select(models.Memory.id, models.Memory.text)
.where(models.Memory.id.in_(used_ids))
).all()
)
if update_stats:
now = models.utcnow()
for _, m in used:
m.use_count += 1
m.last_used_at = now
# synchronize_session=False: nothing in this request reads the counters
# back, and matching the UPDATE against loaded objects would mean having
# loaded them, which is the cost this whole path exists to avoid.
db.execute(
update(models.Memory)
.where(models.Memory.id.in_(used_ids))
.values(use_count=models.Memory.use_count + 1, last_used_at=models.utcnow())
.execution_options(synchronize_session=False)
)
return {
"used": [
{"id": m.id, "text": m.text, "similarity": round(score, 4), "pinned": m.pinned}
for score, m in used
{"id": memory_id, "text": texts.get(memory_id, ""),
"similarity": round(score, 4), "pinned": pinned}
for score, memory_id, pinned in used
],
"error": None,
}
@@ -391,8 +490,19 @@ async def _update_story_summary(
async def _embed_pending(
adventure: models.Adventure, settings: models.Settings, db: Session
) -> None:
pending = [m for m in adventure.memories if m.embedding is None and not m.forgotten]
pending = pending[:MAX_EMBED_BATCH]
# A query, not a walk of adventure.memories: this ran every turn and pulled
# the whole bank's vectors to find the handful that had none.
pending = (
db.query(models.Memory)
.filter(
models.Memory.adventure_id == adventure.id,
models.Memory.embedded.is_(False),
models.Memory.forgotten.is_(False),
)
.order_by(models.Memory.id)
.limit(MAX_EMBED_BATCH)
.all()
)
if not pending:
return
try:
@@ -407,14 +517,35 @@ async def _embed_pending(
def _evict_over_capacity(
adventure: models.Adventure, settings: models.Settings, db: Session
) -> None:
active = [m for m in adventure.memories if not m.forgotten]
overflow = len(active) - max(1, settings.memory_bank_capacity)
# Counting and ranking are both things the database does without sending
# anything back. Walking adventure.memories to count them fetched every
# vector in the bank, every turn, whether or not anything was over capacity.
in_this_bank = (models.Memory.adventure_id == adventure.id,
models.Memory.forgotten.is_(False))
active = db.execute(
select(func.count(models.Memory.id)).where(*in_this_bank)
).scalar() or 0
overflow = active - max(1, settings.memory_bank_capacity)
if overflow <= 0:
return
evictable = sorted(
(m for m in active if not m.pinned),
key=lambda m: (m.use_count, m.last_used_at or m.created_at),
doomed = db.execute(
select(models.Memory.id)
.where(*in_this_bank, models.Memory.pinned.is_(False))
.order_by(
models.Memory.use_count,
func.coalesce(models.Memory.last_used_at, models.Memory.created_at),
)
.limit(overflow)
).scalars().all()
if not doomed:
return # every active memory is pinned; capacity yields to the pins
db.execute(
update(models.Memory)
.where(models.Memory.id.in_(doomed))
.values(forgotten=True)
.execution_options(synchronize_session=False)
)
for memory in evictable[:overflow]:
memory.forgotten = True
db.commit()
# The bulk UPDATE went around any loaded objects, so anything still holding
# the collection would see the evicted memories as active.
db.expire(adventure, ["memories"])
+10
View File
@@ -134,6 +134,16 @@ MIGRATIONS: list[tuple[int, str | dict[str, str]]] = [
# in place and keeps being written until a follow-up migration drops it.
(38, {"sqlite": "ALTER TABLE memories ADD COLUMN embedding_blob BLOB",
"default": "ALTER TABLE memories ADD COLUMN embedding_blob BYTEA"}),
# ...and the one-bit answer beside it, so the Memories drawer and the embed
# queue can ask "has this got a vector?" without fetching one. Same shape as
# actions.variant_count beside actions.variants. TRUE/FALSE and a boolean
# DEFAULT are spelled the same on both dialects; 0/1 would not be.
(39, "ALTER TABLE memories ADD COLUMN embedded BOOLEAN NOT NULL DEFAULT false"),
(40, "UPDATE memories SET embedded = true WHERE embedding_blob IS NOT NULL"),
# Memory bank capacity 200 -> 80. Only rows still on the old default move,
# so anyone who picked a value keeps it — same rule as migration 29.
# Adventures already over 80 evict down on their next turn.
(41, "UPDATE settings SET memory_bank_capacity = 80 WHERE memory_bank_capacity = 200"),
]
LATEST_VERSION = max((v for v, _ in MIGRATIONS), default=1)
+14 -9
View File
@@ -158,10 +158,10 @@ class Memory(Base):
id: Mapped[int] = mapped_column(primary_key=True)
adventure_id: Mapped[int] = mapped_column(ForeignKey("adventures.id", ondelete="CASCADE"))
text: Mapped[str] = mapped_column(Text, default="")
# Superseded by embedding_blob and written alongside it, so the two stay in
# step until a follow-up migration drops this one. Still the column the
# ranking path reads; that moves next.
embedding: Mapped[list | None] = mapped_column(JSON, nullable=True)
# Superseded by embedding_blob and still written alongside it, so a
# rollback finds the vectors, until a follow-up migration drops it. Nothing
# reads it.
embedding: Mapped[list | None] = mapped_column(JSON, nullable=True, deferred=True)
# The vector, little-endian float32. Deferred because it is wider than the
# rest of the row put together and exactly one code path wants it: anything
# bulk-loading memories (the Memories drawer, eviction, the embed queue)
@@ -172,6 +172,11 @@ class Memory(Base):
# Action index range this memory summarizes (null for manual memories).
source_start: Mapped[int | None] = mapped_column(Integer, nullable=True)
source_end: Mapped[int | None] = mapped_column(Integer, nullable=True)
# Whether embedding_blob is set. Maintained on write by memorybank
# .set_vector, for the same reason actions.variant_count exists beside
# actions.variants: every reader wants the one-bit answer and none of them
# should have to fetch six kilobytes of vector to get it.
embedded: Mapped[bool] = mapped_column(Boolean, default=False)
pinned: Mapped[bool] = mapped_column(Boolean, default=False)
forgotten: Mapped[bool] = mapped_column(Boolean, default=False) # evicted, kept for UI
use_count: Mapped[int] = mapped_column(Integer, default=0)
@@ -180,10 +185,6 @@ class Memory(Base):
adventure: Mapped[Adventure] = relationship(back_populates="memories")
@property
def embedded(self) -> bool:
return self.embedding is not None
class StoryCard(Base):
"""Owned by either a scenario or an adventure (exactly one set)."""
@@ -378,7 +379,11 @@ class Settings(Base):
# Phase 6: auto-summarization + memory bank
summary_model: Mapped[str] = mapped_column(String(200), default="") # "" = main model
embedding_model: Mapped[str] = mapped_column(String(200), default="") # "" = bank disabled
memory_bank_capacity: Mapped[int] = mapped_column(Integer, default=200)
# Was 200. Lowered on retrieval-quality grounds first: ranking two hundred
# memories to pick five means the five are chosen out of a lot of noise,
# and older memories describe a story the player has moved on from. That it
# also cuts what the bank costs to read is the smaller reason.
memory_bank_capacity: Mapped[int] = mapped_column(Integer, default=80)
memory_top_k: Mapped[int] = mapped_column(Integer, default=5)
@property
+3
View File
@@ -427,6 +427,9 @@ def delete_adventure(
adventure = get_adventure_or_404(adventure_id, db, user)
db.delete(adventure)
db.commit()
# Nothing else would ever ask for this adventure's vectors again, so the
# cache would hold them until the process restarted.
memorybank.forget_cached_vectors(adventure_id)
# ---------- Turn engine ----------
+16 -4
View File
@@ -18,16 +18,28 @@ that the packing and the in-process cache together already make cheap.
import math
import struct
import sys
from array import array
def pack(vector: list[float]) -> bytes:
def pack(vector) -> bytes:
"""A vector as little-endian float32."""
return struct.pack(f"<{len(vector)}f", *vector)
def unpack(blob: bytes) -> list[float]:
"""The inverse of `pack`. Length is implied: four bytes per component."""
return list(struct.unpack(f"<{len(blob) // 4}f", blob))
def unpack(blob: bytes) -> array:
"""The inverse of `pack`. Length is implied: four bytes per component.
Returns an `array("f")` rather than a list, because these are held in
memory between turns: the array is the same 4 bytes a component the column
is, where a list of Python floats is eight times that. It indexes, zips and
lens like a list, which is all the ranking needs.
"""
vector = array("f")
vector.frombytes(blob)
if sys.byteorder != "little":
vector.byteswap()
return vector
def cosine(a: list[float], b: list[float]) -> float: