Cut database egress 189x by deferring the prompt snapshot
The free-tier 5 GB/month network transfer allowance ran out, which blocks connections outright. The database is only ~55 MB, so 5 GB meant the whole thing was being pulled roughly 90 times over. Cause: actions is 39 MB of that 55 MB -- 541 rows at ~74 KB each, almost entirely context_snapshot, which stores the whole assembled prompt for a turn. Every adventure load and every turn fetched all of it in order to read two small things out of it: the world-change chips under an AI message (Action.world_changes) and the emit block re-attached when replaying history to the model (_history_text). The Insights viewer is the only consumer that wants the whole snapshot, and it asks for one action at a time. Lifts that slice into its own small actions.world_delta column (migration 36) and marks context_snapshot, state_before and world_state_before deferred, so they load only when something touches the attribute -- Insights, undo and retry, all single-action paths. The backfill runs server-side, dialect-specific (json_extract on SQLite, #> on Postgres), because pulling 39 MB of snapshots into Python to rewrite a slice of each would defeat the purpose. Measured at production shape (541 actions, 72 KB snapshots), one adventure load goes from 38.46 MB to 0.20 MB. The traffic that consumed 5 GB would now be about 27 MB. Deliberately not included: limiting the history query to recent actions, and removing the redundant db.refresh(adventure) calls. Both were sized against the old numbers; against a 0.20 MB load they would take ~27 MB a month down to ~10 MB, which is not worth the complexity. tests/test_egress.py hooks before_cursor_execute and asserts the emitted SQL never names the deferred columns during a bulk load, so this cannot regress silently. 123 tests pass. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UeQVy5bEjLhfgWNc27Efet
This commit is contained in:
co-authored by
Claude Opus 5
parent
7c538c8235
commit
f1bd099ec8
@@ -62,13 +62,15 @@ def _history_text(action: models.Action) -> str:
|
|||||||
delta). The block is stripped before storage/UI, so without this every past
|
delta). The block is stripped before storage/UI, so without this every past
|
||||||
AI turn would look like one that emitted nothing — biasing the model, by
|
AI turn would look like one that emitted nothing — biasing the model, by
|
||||||
imitation, to stop emitting too. Player turns and blockless turns are
|
imitation, to stop emitting too. Player turns and blockless turns are
|
||||||
returned unchanged."""
|
returned unchanged.
|
||||||
|
|
||||||
|
Reads `world_delta`, not `context_snapshot`: this runs for every action in
|
||||||
|
the replayed history, and the snapshot is deferred precisely so a turn
|
||||||
|
never drags the prompt archive out of the database."""
|
||||||
text = action.text
|
text = action.text
|
||||||
snap = action.context_snapshot if isinstance(action.context_snapshot, dict) else None
|
wd = action.world_delta if isinstance(action.world_delta, dict) else None
|
||||||
if snap:
|
if wd:
|
||||||
ws = snap.get("world_state")
|
block = worldstate.render_delta_block(wd.get("delta") or {})
|
||||||
if isinstance(ws, dict):
|
|
||||||
block = worldstate.render_delta_block(ws.get("delta") or {})
|
|
||||||
if block:
|
if block:
|
||||||
text = f"{text}\n{block}"
|
text = f"{text}\n{block}"
|
||||||
return text
|
return text
|
||||||
|
|||||||
@@ -110,10 +110,49 @@ MIGRATIONS: list[tuple[int, str]] = [
|
|||||||
# predates this column.
|
# predates this column.
|
||||||
(34, "ALTER TABLE actions ADD COLUMN variants JSON"),
|
(34, "ALTER TABLE actions ADD COLUMN variants JSON"),
|
||||||
(35, "ALTER TABLE actions ADD COLUMN variant_index INTEGER NOT NULL DEFAULT 0"),
|
(35, "ALTER TABLE actions ADD COLUMN variant_index INTEGER NOT NULL DEFAULT 0"),
|
||||||
|
# Egress: context_snapshot holds the whole assembled prompt (~74 KB/row) and
|
||||||
|
# was being loaded in bulk for two tiny things — the world-change chips and
|
||||||
|
# the emit block replayed into history. Lift just that slice into its own
|
||||||
|
# column so the snapshot can be deferred. Backfilled by _backfill_world_delta.
|
||||||
|
(36, "ALTER TABLE actions ADD COLUMN world_delta JSON"),
|
||||||
]
|
]
|
||||||
|
|
||||||
LATEST_VERSION = max((v for v, _ in MIGRATIONS), default=1)
|
LATEST_VERSION = max((v for v, _ in MIGRATIONS), default=1)
|
||||||
|
|
||||||
|
# Migrations that need a data pass after their DDL, keyed by version.
|
||||||
|
WORLD_DELTA_VERSION = 36
|
||||||
|
|
||||||
|
|
||||||
|
def _backfill_world_delta(conn) -> None:
|
||||||
|
"""Populate actions.world_delta from the existing context_snapshot.
|
||||||
|
|
||||||
|
Runs entirely server-side: the snapshots are the reason this change exists,
|
||||||
|
so pulling ~40 MB of them into Python to rewrite a slice would defeat the
|
||||||
|
point. Dialect-specific because SQLite and Postgres spell JSON access
|
||||||
|
differently, and both have to work (SQLite locally and in tests).
|
||||||
|
"""
|
||||||
|
if conn.dialect.name == "sqlite":
|
||||||
|
sql = """
|
||||||
|
UPDATE actions SET world_delta = json_object(
|
||||||
|
'delta', json_extract(context_snapshot, '$.world_state.delta'),
|
||||||
|
'applied', json_extract(context_snapshot, '$.world_state.report.applied')
|
||||||
|
)
|
||||||
|
WHERE world_delta IS NULL
|
||||||
|
AND context_snapshot IS NOT NULL
|
||||||
|
AND json_extract(context_snapshot, '$.world_state') IS NOT NULL
|
||||||
|
"""
|
||||||
|
else:
|
||||||
|
sql = """
|
||||||
|
UPDATE actions SET world_delta = jsonb_build_object(
|
||||||
|
'delta', context_snapshot::jsonb #> '{world_state,delta}',
|
||||||
|
'applied', context_snapshot::jsonb #> '{world_state,report,applied}'
|
||||||
|
)
|
||||||
|
WHERE world_delta IS NULL
|
||||||
|
AND context_snapshot IS NOT NULL
|
||||||
|
AND jsonb_exists(context_snapshot::jsonb, 'world_state')
|
||||||
|
"""
|
||||||
|
conn.execute(text(sql))
|
||||||
|
|
||||||
|
|
||||||
def _get_version(conn) -> int:
|
def _get_version(conn) -> int:
|
||||||
if conn.dialect.name == "sqlite":
|
if conn.dialect.name == "sqlite":
|
||||||
@@ -153,6 +192,8 @@ def bootstrap(engine: Engine) -> None:
|
|||||||
for version, sql in MIGRATIONS:
|
for version, sql in MIGRATIONS:
|
||||||
if version > current:
|
if version > current:
|
||||||
conn.execute(text(sql))
|
conn.execute(text(sql))
|
||||||
|
if version == WORLD_DELTA_VERSION:
|
||||||
|
_backfill_world_delta(conn)
|
||||||
current = version
|
current = version
|
||||||
_set_version(conn, current)
|
_set_version(conn, current)
|
||||||
_encrypt_plaintext_api_keys(conn)
|
_encrypt_plaintext_api_keys(conn)
|
||||||
|
|||||||
+31
-11
@@ -205,13 +205,30 @@ class Action(Base):
|
|||||||
text: Mapped[str] = mapped_column(Text, default="")
|
text: Mapped[str] = mapped_column(Text, default="")
|
||||||
# Reasoning-model "thinking" that preceded the text (AI actions only).
|
# Reasoning-model "thinking" that preceded the text (AI actions only).
|
||||||
reasoning: Mapped[str | None] = mapped_column(Text, nullable=True)
|
reasoning: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||||
context_snapshot: Mapped[dict | None] = mapped_column(JSON, nullable=True)
|
# The full assembled prompt for this turn, for the Insights viewer. By far
|
||||||
|
# the biggest column in the database (~74 KB/row in production), and needed
|
||||||
|
# by exactly one endpoint, one action at a time — so it is deferred: never
|
||||||
|
# loaded unless something actually touches the attribute. Bulk readers must
|
||||||
|
# NOT touch it; that is what `world_delta` below exists for.
|
||||||
|
context_snapshot: Mapped[dict | None] = mapped_column(
|
||||||
|
JSON, nullable=True, deferred=True
|
||||||
|
)
|
||||||
|
# The small slice of the snapshot that IS needed in bulk: this turn's RPG
|
||||||
|
# state changes, for the inline chips under an AI message (world_changes)
|
||||||
|
# and for re-attaching the emit block when replaying history to the model.
|
||||||
|
# Mirrors the active variant, same as text/reasoning/context_snapshot.
|
||||||
|
world_delta: Mapped[dict | None] = mapped_column(JSON, nullable=True)
|
||||||
# Copy of Adventure.script_state as it was immediately BEFORE this action's
|
# Copy of Adventure.script_state as it was immediately BEFORE this action's
|
||||||
# script hooks ran, so undo/retry can roll the shared scoreboard back.
|
# script hooks ran, so undo/retry can roll the shared scoreboard back.
|
||||||
# NULL for actions created before this column existed.
|
# NULL for actions created before this column existed. Deferred: only ever
|
||||||
state_before: Mapped[dict | None] = mapped_column(JSON, nullable=True)
|
# read for the one action being undone or retried.
|
||||||
|
state_before: Mapped[dict | None] = mapped_column(
|
||||||
|
JSON, nullable=True, deferred=True
|
||||||
|
)
|
||||||
# Phase 12: same idea for the RPG world_state, so undo/retry rolls it back too.
|
# Phase 12: same idea for the RPG world_state, so undo/retry rolls it back too.
|
||||||
world_state_before: Mapped[dict | None] = mapped_column(JSON, nullable=True)
|
world_state_before: Mapped[dict | None] = mapped_column(
|
||||||
|
JSON, nullable=True, deferred=True
|
||||||
|
)
|
||||||
# Retry history (AI actions): every attempt made for this turn, oldest
|
# Retry history (AI actions): every attempt made for this turn, oldest
|
||||||
# first, INCLUDING the active one. NULL/empty means never retried — the row
|
# first, INCLUDING the active one. NULL/empty means never retried — the row
|
||||||
# is its own only version. `variant_index` says which entry `text`,
|
# is its own only version. `variant_index` says which entry `text`,
|
||||||
@@ -225,14 +242,17 @@ class Action(Base):
|
|||||||
|
|
||||||
@property
|
@property
|
||||||
def world_changes(self) -> list[dict]:
|
def world_changes(self) -> list[dict]:
|
||||||
"""Compact per-turn RPG state changes (Phase 12), derived from the
|
"""Compact per-turn RPG state changes (Phase 12), for the inline summary
|
||||||
stored snapshot, for the inline summary under an AI message. Labels are
|
under an AI message. Labels are path-based (no schema needed):
|
||||||
path-based (no schema needed): `npc.gwen.trust` -> "gwen trust"."""
|
`npc.gwen.trust` -> "gwen trust".
|
||||||
cs = self.context_snapshot if isinstance(self.context_snapshot, dict) else None
|
|
||||||
ws = cs.get("world_state") if cs else None
|
Reads `world_delta`, never `context_snapshot` — this runs for every
|
||||||
if not isinstance(ws, dict):
|
action in a list response, and touching the deferred snapshot here
|
||||||
|
would drag the whole prompt archive out of the database."""
|
||||||
|
wd = self.world_delta if isinstance(self.world_delta, dict) else None
|
||||||
|
if wd is None:
|
||||||
return []
|
return []
|
||||||
applied = (ws.get("report") or {}).get("applied") or []
|
applied = wd.get("applied") or []
|
||||||
out: list[dict] = []
|
out: list[dict] = []
|
||||||
for entry in applied:
|
for entry in applied:
|
||||||
parts = str(entry.get("path", "")).split(".")
|
parts = str(entry.get("path", "")).split(".")
|
||||||
|
|||||||
@@ -333,6 +333,22 @@ def snapshot_world_state(adventure: models.Adventure) -> dict:
|
|||||||
VARIANT_SNAPSHOT_KEYS = ("world_state", "script", "raw_output")
|
VARIANT_SNAPSHOT_KEYS = ("world_state", "script", "raw_output")
|
||||||
|
|
||||||
|
|
||||||
|
def world_delta_of(snapshot: dict | None) -> dict | None:
|
||||||
|
"""The bulk-read slice of a context snapshot, for Action.world_delta.
|
||||||
|
|
||||||
|
context_snapshot is deferred (it holds the whole assembled prompt), so the
|
||||||
|
two things that ARE needed for every action — the world-change chips and
|
||||||
|
the emit block replayed into history — get their own small column. Keep
|
||||||
|
this in step with the snapshot wherever one is written."""
|
||||||
|
ws = (snapshot or {}).get("world_state")
|
||||||
|
if not isinstance(ws, dict):
|
||||||
|
return None
|
||||||
|
return {
|
||||||
|
"delta": ws.get("delta") or {},
|
||||||
|
"applied": (ws.get("report") or {}).get("applied") or [],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
def variant_of(action: models.Action, adventure: models.Adventure) -> dict:
|
def variant_of(action: models.Action, adventure: models.Adventure) -> dict:
|
||||||
"""Freeze an action's *current* content as a variant entry.
|
"""Freeze an action's *current* content as a variant entry.
|
||||||
|
|
||||||
@@ -365,6 +381,7 @@ def apply_variant(action: models.Action, adventure: models.Adventure, index: int
|
|||||||
else:
|
else:
|
||||||
snapshot.pop(key, None)
|
snapshot.pop(key, None)
|
||||||
action.context_snapshot = snapshot
|
action.context_snapshot = snapshot
|
||||||
|
action.world_delta = world_delta_of(snapshot)
|
||||||
action.variant_index = index
|
action.variant_index = index
|
||||||
if isinstance(entry.get("script_state"), dict):
|
if isinstance(entry.get("script_state"), dict):
|
||||||
adventure.script_state = copy.deepcopy(entry["script_state"])
|
adventure.script_state = copy.deepcopy(entry["script_state"])
|
||||||
@@ -637,6 +654,7 @@ async def _generate_turn(
|
|||||||
text=text,
|
text=text,
|
||||||
reasoning=reasoning,
|
reasoning=reasoning,
|
||||||
context_snapshot=snapshot,
|
context_snapshot=snapshot,
|
||||||
|
world_delta=world_delta_of(snapshot),
|
||||||
state_before=state_before,
|
state_before=state_before,
|
||||||
world_state_before=world_state_before,
|
world_state_before=world_state_before,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -0,0 +1,176 @@
|
|||||||
|
"""Guards on how much the database is asked for.
|
||||||
|
|
||||||
|
context_snapshot holds the entire assembled prompt for a turn (~74 KB/row in
|
||||||
|
production, 94% of the database). It used to be pulled for every action on
|
||||||
|
every adventure load and every turn, to read two tiny things out of it. These
|
||||||
|
tests fail if that regresses.
|
||||||
|
|
||||||
|
python -m pytest tests/test_egress.py -v
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
import tempfile
|
||||||
|
|
||||||
|
_tmp = tempfile.NamedTemporaryFile(suffix=".db", delete=False)
|
||||||
|
_tmp.close()
|
||||||
|
os.environ["AIDND_DB_PATH"] = _tmp.name
|
||||||
|
os.environ.pop("AIDND_DATABASE_URL", None)
|
||||||
|
os.environ.pop("DATABASE_URL", None)
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from fastapi import Depends
|
||||||
|
from fastapi.testclient import TestClient
|
||||||
|
from sqlalchemy import event, text
|
||||||
|
|
||||||
|
from app import auth, limits, migrations, models
|
||||||
|
from app.database import Base, SessionLocal, engine, get_db
|
||||||
|
from app.main import app
|
||||||
|
|
||||||
|
# A stand-in for the real thing: the assembled prompt, which is what makes the
|
||||||
|
# column enormous, plus the small world_state slice the UI actually needs.
|
||||||
|
BIG_SNAPSHOT = {
|
||||||
|
"system": "x" * 20_000,
|
||||||
|
"story": "y" * 40_000,
|
||||||
|
"world_state": {
|
||||||
|
"delta": {"player.hp": -15},
|
||||||
|
"report": {"applied": [{"path": "player.hp", "old": 100, "new": 85}]},
|
||||||
|
"state": {"player": {"hp": 85}},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture()
|
||||||
|
def sql_log():
|
||||||
|
"""Every statement the ORM sends, for asserting on what was selected."""
|
||||||
|
statements: list[str] = []
|
||||||
|
|
||||||
|
def record(conn, cursor, statement, parameters, context, executemany):
|
||||||
|
statements.append(statement)
|
||||||
|
|
||||||
|
event.listen(engine, "before_cursor_execute", record)
|
||||||
|
try:
|
||||||
|
yield statements
|
||||||
|
finally:
|
||||||
|
event.remove(engine, "before_cursor_execute", record)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture()
|
||||||
|
def client(monkeypatch):
|
||||||
|
Base.metadata.create_all(bind=engine)
|
||||||
|
setup = SessionLocal()
|
||||||
|
user = models.User(is_guest=False, email="egress@example.com")
|
||||||
|
setup.add(user)
|
||||||
|
setup.flush()
|
||||||
|
setup.add(models.Settings(user_id=user.id, api_key="enc:dummy", model="test-model"))
|
||||||
|
adventure = models.Adventure(user_id=user.id, title="Cave", script_state={})
|
||||||
|
setup.add(adventure)
|
||||||
|
setup.flush()
|
||||||
|
for i in range(12):
|
||||||
|
setup.add(models.Action(
|
||||||
|
adventure_id=adventure.id, index=i,
|
||||||
|
type="ai" if i % 2 else "do", text=f"Action {i}.",
|
||||||
|
context_snapshot=BIG_SNAPSHOT,
|
||||||
|
world_delta={"delta": {"player.hp": -15},
|
||||||
|
"applied": [{"path": "player.hp", "old": 100, "new": 85}]},
|
||||||
|
))
|
||||||
|
setup.commit()
|
||||||
|
adv_id, user_id = adventure.id, user.id
|
||||||
|
setup.close()
|
||||||
|
|
||||||
|
monkeypatch.setattr(limits, "rate_limit", lambda *a, **k: None)
|
||||||
|
monkeypatch.setattr(limits, "check_row_cap", lambda *a, **k: None)
|
||||||
|
|
||||||
|
def _current_user(db=Depends(get_db)):
|
||||||
|
return db.get(models.User, user_id)
|
||||||
|
|
||||||
|
app.dependency_overrides[auth.get_current_user] = _current_user
|
||||||
|
c = TestClient(app)
|
||||||
|
c.adv_id = adv_id
|
||||||
|
try:
|
||||||
|
yield c
|
||||||
|
finally:
|
||||||
|
app.dependency_overrides.clear()
|
||||||
|
Base.metadata.drop_all(bind=engine)
|
||||||
|
|
||||||
|
|
||||||
|
def action_selects(statements: list[str]) -> list[str]:
|
||||||
|
return [s for s in statements if "FROM actions" in s and s.lstrip().upper().startswith("SELECT")]
|
||||||
|
|
||||||
|
|
||||||
|
# ------------------------------------------------------- the deferred columns
|
||||||
|
|
||||||
|
def test_loading_an_adventure_does_not_fetch_context_snapshot(client, sql_log):
|
||||||
|
r = client.get(f"/api/adventures/{client.adv_id}")
|
||||||
|
assert r.status_code == 200, r.text
|
||||||
|
assert len(r.json()["actions"]) == 12
|
||||||
|
|
||||||
|
selects = action_selects(sql_log)
|
||||||
|
assert selects, "expected at least one SELECT against actions"
|
||||||
|
offenders = [s for s in selects if "context_snapshot" in s]
|
||||||
|
assert offenders == [], f"context_snapshot was fetched in bulk:\n{offenders[0][:400]}"
|
||||||
|
|
||||||
|
|
||||||
|
def test_state_before_and_world_state_before_are_not_fetched_in_bulk(client, sql_log):
|
||||||
|
"""Both are rollback snapshots, only ever needed for the single action
|
||||||
|
being undone or retried."""
|
||||||
|
client.get(f"/api/adventures/{client.adv_id}")
|
||||||
|
selects = action_selects(sql_log)
|
||||||
|
for column in ("state_before", "world_state_before"):
|
||||||
|
offenders = [s for s in selects if column in s]
|
||||||
|
assert offenders == [], f"{column} was fetched in bulk"
|
||||||
|
|
||||||
|
|
||||||
|
def test_world_changes_still_works_without_the_snapshot(client):
|
||||||
|
"""The chips under an AI message must survive the snapshot being deferred."""
|
||||||
|
r = client.get(f"/api/adventures/{client.adv_id}")
|
||||||
|
ai = [a for a in r.json()["actions"] if a["type"] == "ai"]
|
||||||
|
assert ai, "fixture should have AI actions"
|
||||||
|
assert ai[0]["world_changes"] == [
|
||||||
|
{"kind": "stat", "label": "hp", "delta": -15, "value": 85}
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def test_snapshot_is_still_reachable_on_demand(client):
|
||||||
|
"""Deferred means lazy, not gone — Insights still gets the full thing."""
|
||||||
|
r = client.get(f"/api/adventures/{client.adv_id}")
|
||||||
|
action_id = r.json()["actions"][0]["id"]
|
||||||
|
r = client.get(f"/api/adventures/{client.adv_id}/actions/{action_id}/context")
|
||||||
|
assert r.status_code == 200, r.text
|
||||||
|
assert r.json()["system"] == "x" * 20_000
|
||||||
|
|
||||||
|
|
||||||
|
# ------------------------------------------------------------------ backfill
|
||||||
|
|
||||||
|
def test_backfill_populates_world_delta_from_existing_snapshots(client):
|
||||||
|
"""Migration 36 lifts the slice out server-side, without reading the
|
||||||
|
snapshots into Python."""
|
||||||
|
db = SessionLocal()
|
||||||
|
try:
|
||||||
|
db.execute(text("UPDATE actions SET world_delta = NULL"))
|
||||||
|
db.commit()
|
||||||
|
assert db.query(models.Action).filter(models.Action.world_delta.isnot(None)).count() == 0
|
||||||
|
|
||||||
|
with engine.begin() as conn:
|
||||||
|
migrations._backfill_world_delta(conn)
|
||||||
|
|
||||||
|
db.expire_all()
|
||||||
|
actions = db.query(models.Action).all()
|
||||||
|
assert all(a.world_delta is not None for a in actions)
|
||||||
|
assert actions[0].world_delta["delta"] == {"player.hp": -15}
|
||||||
|
assert actions[0].world_delta["applied"] == [
|
||||||
|
{"path": "player.hp", "old": 100, "new": 85}
|
||||||
|
]
|
||||||
|
finally:
|
||||||
|
db.close()
|
||||||
|
|
||||||
|
|
||||||
|
def test_backfill_leaves_actions_without_world_state_alone(client):
|
||||||
|
db = SessionLocal()
|
||||||
|
try:
|
||||||
|
db.execute(text("UPDATE actions SET world_delta = NULL, context_snapshot = '{\"story\": \"s\"}'"))
|
||||||
|
db.commit()
|
||||||
|
with engine.begin() as conn:
|
||||||
|
migrations._backfill_world_delta(conn)
|
||||||
|
db.expire_all()
|
||||||
|
assert all(a.world_delta is None for a in db.query(models.Action).all())
|
||||||
|
finally:
|
||||||
|
db.close()
|
||||||
Reference in New Issue
Block a user