diff --git a/src/xrouter_llm/migrations/versions/0004_dedupe_prompt_text.py b/src/xrouter_llm/migrations/versions/0004_dedupe_prompt_text.py new file mode 100644 index 0000000..280bf5f --- /dev/null +++ b/src/xrouter_llm/migrations/versions/0004_dedupe_prompt_text.py @@ -0,0 +1,81 @@ +"""Deduplicate prompt text into a prompts table + +Revision ID: 0004 +Revises: 0003 +Create Date: 2026-07-08 +""" +from __future__ import annotations + +import hashlib + +import sqlalchemy as sa +from alembic import op + +revision: str = "0004" +down_revision: str = "0003" +branch_labels: str | None = None +depends_on: str | None = None + + +def upgrade() -> None: + op.create_table( + "prompts", + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("sha256", sa.String(64), nullable=False), + sa.Column("text", sa.Text(), nullable=False), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("sha256", name="uq_prompts_sha256"), + ) + + op.add_column("calls", sa.Column("prompt_id", sa.Integer(), nullable=True)) + + # Backfill by exact content hash in Python rather than SQL DISTINCT / + # text-equality joins: under case-insensitive collations (MySQL default) + # those would merge prompts differing only in casing, silently rewriting + # the stored text of some calls. Hashing matches the runtime dedup rule. + conn = op.get_bind() + calls = conn.execute(sa.text("SELECT id, prompt FROM calls")).fetchall() + sha_to_id: dict[str, int] = {} + call_updates: list[dict[str, int]] = [] + for call_id, prompt in calls: + sha = hashlib.sha256(prompt.encode("utf-8")).hexdigest() + prompt_id = sha_to_id.get(sha) + if prompt_id is None: + conn.execute( + sa.text("INSERT INTO prompts (sha256, text) VALUES (:sha, :text)"), + {"sha": sha, "text": prompt}, + ) + prompt_id = conn.execute( + sa.text("SELECT id FROM prompts WHERE sha256 = :sha"), {"sha": sha} + ).scalar_one() + sha_to_id[sha] = prompt_id + call_updates.append({"cid": call_id, "pid": prompt_id}) + if call_updates: + conn.execute( + sa.text("UPDATE calls SET prompt_id = :pid WHERE id = :cid"), + call_updates, + ) + + # batch mode: SQLite cannot ALTER to NOT NULL / add FK in place + with op.batch_alter_table("calls") as batch: + batch.alter_column("prompt_id", existing_type=sa.Integer(), nullable=False) + batch.create_foreign_key( + "fk_calls_prompt_id_prompts", "prompts", ["prompt_id"], ["id"] + ) + batch.drop_column("prompt") + op.create_index("ix_calls_prompt_id", "calls", ["prompt_id"]) + + +def downgrade() -> None: + op.drop_index("ix_calls_prompt_id", table_name="calls") + op.add_column("calls", sa.Column("prompt", sa.Text(), nullable=True)) + conn = op.get_bind() + conn.execute(sa.text( + "UPDATE calls SET prompt = " + "(SELECT p.text FROM prompts p WHERE p.id = calls.prompt_id)" + )) + with op.batch_alter_table("calls") as batch: + batch.alter_column("prompt", existing_type=sa.Text(), nullable=False) + batch.drop_constraint("fk_calls_prompt_id_prompts", type_="foreignkey") + batch.drop_column("prompt_id") + op.drop_table("prompts") diff --git a/src/xrouter_llm/server.py b/src/xrouter_llm/server.py index c3522d9..691abc4 100644 --- a/src/xrouter_llm/server.py +++ b/src/xrouter_llm/server.py @@ -200,7 +200,10 @@ def create_app(service: RoutingService) -> FastAPI: .pager { display: flex; align-items: center; gap: 8px; } .pager span { color: #8a93a6; font-size: 13px; } .pager button { padding: 5px 12px; font-size: 14px; font-weight: normal; } - .prompt-cell { max-width: 260px; } + /* overflow-wrap alone can't break a 60-char unspaced token (e.g. a URL) + inside a table cell; break-all keeps it from overflowing into the next + column. CJK text is unaffected (it already breaks per character). */ + .prompt-cell { max-width: 260px; word-break: break-all; } .prompt-short { display: inline; } diff --git a/src/xrouter_llm/store.py b/src/xrouter_llm/store.py index 003510f..e4f8369 100644 --- a/src/xrouter_llm/store.py +++ b/src/xrouter_llm/store.py @@ -2,6 +2,7 @@ from __future__ import annotations +import hashlib from pathlib import Path from typing import Any @@ -10,7 +11,15 @@ from alembic.config import Config as AlembicConfig from sqlalchemy import create_engine from sqlalchemy.engine import Engine -from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker +from sqlalchemy.exc import IntegrityError +from sqlalchemy.orm import ( + DeclarativeBase, + Mapped, + Session, + mapped_column, + relationship, + sessionmaker, +) from sqlalchemy.pool import StaticPool _MIGRATIONS_DIR = Path(__file__).parent / "migrations" @@ -20,13 +29,25 @@ class Base(DeclarativeBase): pass +class PromptRecord(Base): + """Prompt text stored once per distinct prompt, keyed by content hash.""" + + __tablename__ = "prompts" + + id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) + sha256: Mapped[str] = mapped_column(sa.String(64), nullable=False, unique=True) + text: Mapped[str] = mapped_column(sa.Text, nullable=False) + + class CallRecord(Base): __tablename__ = "calls" id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) ts: Mapped[float] = mapped_column(sa.Float, nullable=False) config: Mapped[str] = mapped_column(sa.String(255), nullable=False) - prompt: Mapped[str] = mapped_column(sa.Text, nullable=False) + prompt_id: Mapped[int] = mapped_column( + sa.ForeignKey("prompts.id"), nullable=False + ) task: Mapped[str | None] = mapped_column(sa.String(255), nullable=True) selected: Mapped[Any] = mapped_column(sa.JSON, nullable=False) candidates: Mapped[Any] = mapped_column(sa.JSON, nullable=False) @@ -36,8 +57,11 @@ class CallRecord(Base): feedback: Mapped[Any] = mapped_column(sa.JSON, nullable=True) user_id: Mapped[str | None] = mapped_column(sa.String(255), nullable=True) + prompt_rec: Mapped[PromptRecord] = relationship(lazy="joined") + __table_args__ = ( sa.Index("ix_calls_user_id_id", "user_id", "id"), + sa.Index("ix_calls_prompt_id", "prompt_id"), ) @@ -47,9 +71,14 @@ class CallRecord(Base): _SCHEMA_CHECKPOINTS: list[tuple[str, str]] = [ ("feedback", "0002"), ("user_id", "0003"), + ("prompt_id", "0004"), ] +def prompt_sha256(prompt: str) -> str: + return hashlib.sha256(prompt.encode("utf-8")).hexdigest() + + def run_migrations(engine: Engine) -> None: """Run Alembic migrations using an already-created engine. @@ -105,6 +134,19 @@ def _stamp_legacy_db_if_needed(engine: Engine) -> None: ) +def _enforce_sqlite_fks(engine: Engine) -> Engine: + """SQLite ships with foreign_keys OFF; without it a dangling + calls.prompt_id would be silently accepted instead of raising.""" + + @sa.event.listens_for(engine, "connect") + def _fk_pragma(dbapi_conn, _connection_record) -> None: + cursor = dbapi_conn.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + + return engine + + def make_engine(db_url: str) -> Engine: url = sa.engine.make_url(db_url) if url.drivername.startswith("sqlite"): @@ -114,9 +156,15 @@ def make_engine(db_url: str) -> Engine: Path(expanded).parent.mkdir(parents=True, exist_ok=True) # write expanded path back so SQLite opens the real file, not literal "~" url = url.set(database=expanded) - return create_engine(url, connect_args={"check_same_thread": False}) + return _enforce_sqlite_fks( + create_engine(url, connect_args={"check_same_thread": False}) + ) # in-memory: StaticPool shares one connection so the DB persists across sessions - return create_engine(url, connect_args={"check_same_thread": False}, poolclass=StaticPool) + return _enforce_sqlite_fks( + create_engine( + url, connect_args={"check_same_thread": False}, poolclass=StaticPool + ) + ) return create_engine(db_url, pool_pre_ping=True) @@ -148,12 +196,45 @@ def record( cost: float, latency: float, user_id: str | None = None, + ) -> int: + try: + return self._record_once( + ts=ts, config=config, prompt=prompt, task=task, + selected=selected, candidates=candidates, + expected_quality=expected_quality, cost=cost, + latency=latency, user_id=user_id, + ) + except IntegrityError: + # A concurrent delete() GC'd our prompt row between lookup and + # insert. Reachable on SQLite, where FOR UPDATE is a no-op and + # the FK violation only surfaces at insert; one retry recreates + # the prompt. + return self._record_once( + ts=ts, config=config, prompt=prompt, task=task, + selected=selected, candidates=candidates, + expected_quality=expected_quality, cost=cost, + latency=latency, user_id=user_id, + ) + + def _record_once( + self, + *, + ts: float, + config: str, + prompt: str, + task: str | None, + selected: list[str], + candidates: list[dict[str, Any]], + expected_quality: float, + cost: float, + latency: float, + user_id: str | None, ) -> int: with self._Session() as session: row = CallRecord( ts=ts, config=config, - prompt=prompt, + prompt_id=_get_or_create_prompt(session, prompt), task=task, selected=selected, candidates=candidates, @@ -196,7 +277,32 @@ def delete(self, call_id: int, *, owner_user_id: str | None = None) -> bool: row = session.scalars(stmt).first() if row is None: return False + prompt_id = row.prompt_id + # Lock the prompt row BEFORE deleting the call: every delete() + # and record() for one prompt serializes on this single lock + # first, then acquires per-call row locks — a consistent order + # that cannot deadlock. Locking after the call-row delete could: + # T1 would hold the prompt lock while its reference check waits + # on T2's deleted-but-uncommitted call row, and T2 waits on the + # prompt lock. No-op on SQLite, whose single-writer transactions + # already serialize the write paths. + prompt_row = session.get(PromptRecord, prompt_id, with_for_update=True) session.delete(row) + session.flush() + # GC the prompt text once no call references it (the log holds + # user prompts; orphaned text must not outlive its last call). + if prompt_row is not None: + still_referenced = session.scalar( + sa.select(CallRecord.id) + .where(CallRecord.prompt_id == prompt_id) + .limit(1) + # locking read: sees the latest committed refs even under + # MySQL REPEATABLE READ, where a plain SELECT would read + # the transaction snapshot and miss a concurrent delete. + .with_for_update() + ) + if still_referenced is None: + session.delete(prompt_row) session.commit() return True @@ -232,12 +338,41 @@ def model_counts(self) -> dict[str, int]: return counts +def _get_or_create_prompt(session: Session, prompt: str) -> int: + """Return the id of the deduplicated prompt row, inserting if new. + + Insert-then-recover (savepoint) rather than check-then-insert alone, so a + concurrent writer inserting the same hash cannot fail this transaction. + """ + sha = prompt_sha256(prompt) + # FOR UPDATE holds the prompt row until this transaction commits, so a + # concurrent delete() cannot GC it before our call row is inserted. It + # also makes the post-conflict re-read see the winner's committed row + # under MySQL REPEATABLE READ. No-op on SQLite (see record()). + lookup = ( + sa.select(PromptRecord.id) + .where(PromptRecord.sha256 == sha) + .with_for_update() + ) + prompt_id = session.scalar(lookup) + if prompt_id is not None: + return prompt_id + try: + with session.begin_nested(): + row = PromptRecord(sha256=sha, text=prompt) + session.add(row) + session.flush() + return row.id + except IntegrityError: + return session.scalar(lookup) + + def _row_to_dict(r: CallRecord) -> dict[str, Any]: return { "id": r.id, "ts": r.ts, "config": r.config, - "prompt": r.prompt, + "prompt": r.prompt_rec.text, "task": r.task, "selected": r.selected, "candidates": r.candidates, diff --git a/tests/conftest.py b/tests/conftest.py index 2d42cae..f64c4bc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,7 +6,13 @@ import pytest import sqlalchemy as sa -from xrouter_llm.store import CallRecord, CallStore, make_engine, normalize_db_url +from xrouter_llm.store import ( + CallRecord, + CallStore, + PromptRecord, + make_engine, + normalize_db_url, +) _SAMPLE_RECORD = dict( ts=1_000_000.0, @@ -52,3 +58,4 @@ def store(db_url): engine = make_engine(normalize_db_url(db_url)) with engine.begin() as conn: conn.execute(sa.delete(CallRecord)) + conn.execute(sa.delete(PromptRecord)) diff --git a/tests/test_store.py b/tests/test_store.py index 06892af..90afdda 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -1,8 +1,9 @@ """CallStore integration tests, parameterized over DB backends via conftest.db_url.""" from __future__ import annotations +import pytest import sqlalchemy as sa -from xrouter_llm.store import Base, CallStore, make_engine +from xrouter_llm.store import Base, CallStore, PromptRecord, make_engine def test_record_and_recent(store) -> None: @@ -74,6 +75,44 @@ def test_model_counts(store) -> None: assert counts["strong"] == 1 +def _record(store, **overrides): + kwargs = dict( + ts=1.0, config="all", prompt="p", task=None, + selected=["m"], candidates=[], expected_quality=0.8, + cost=0.0, latency=0.0, + ) + kwargs.update(overrides) + return store.record(**kwargs) + + +def _prompt_count(store) -> int: + with store._Session() as session: + return session.execute( + sa.select(sa.func.count(PromptRecord.id)) + ).scalar_one() + + +def test_prompt_text_is_deduplicated(store) -> None: + _record(store, ts=1.0, prompt="same prompt") + _record(store, ts=2.0, prompt="same prompt") + _record(store, ts=3.0, prompt="other prompt") + assert store.count() == 3 + assert _prompt_count(store) == 2 + rows = store.recent() + assert [r["prompt"] for r in rows] == ["other prompt", "same prompt", "same prompt"] + + +def test_delete_gcs_orphaned_prompt(store) -> None: + id1 = _record(store, ts=1.0, prompt="shared") + id2 = _record(store, ts=2.0, prompt="shared") + assert store.delete(id1) is True + # still referenced by the second call + assert _prompt_count(store) == 1 + assert store.recent()[0]["prompt"] == "shared" + assert store.delete(id2) is True + assert _prompt_count(store) == 0 + + def _legacy_db(tmp_path): """Return a sqlite:// URL for a pre-0002 DB: calls table without feedback column, no alembic_version. @@ -145,6 +184,95 @@ def test_legacy_db_no_alembic_version(tmp_path) -> None: assert row["user_id"] is None +def test_legacy_db_backfill_dedupes_prompts(tmp_path) -> None: + """Migration 0004 moves existing prompt text into the prompts table, deduplicated.""" + url = _legacy_db(tmp_path) + engine = make_engine(url) + with engine.begin() as conn: + for i, prompt in enumerate(["dup", "dup", "unique"]): + conn.execute( + sa.text( + "INSERT INTO calls (ts, config, prompt, selected, candidates) " + "VALUES (:ts, 'all', :prompt, '[\"m\"]', '[]')" + ), + {"ts": float(i), "prompt": prompt}, + ) + engine.dispose() + + store = CallStore(url) + assert store.count() == 3 + assert _prompt_count(store) == 2 + rows = store.recent() + assert [r["prompt"] for r in rows] == ["unique", "dup", "dup"] + + +def test_legacy_db_backfill_keeps_case_variants(tmp_path) -> None: + """Backfill dedupes by exact hash: prompts differing only in casing stay distinct.""" + url = _legacy_db(tmp_path) + engine = make_engine(url) + with engine.begin() as conn: + for i, prompt in enumerate(["Hello", "hello"]): + conn.execute( + sa.text( + "INSERT INTO calls (ts, config, prompt, selected, candidates) " + "VALUES (:ts, 'all', :prompt, '[\"m\"]', '[]')" + ), + {"ts": float(i), "prompt": prompt}, + ) + engine.dispose() + + store = CallStore(url) + assert _prompt_count(store) == 2 + assert [r["prompt"] for r in store.recent()] == ["hello", "Hello"] + + +def test_concurrent_delete_gcs_shared_prompt(store, db_url) -> None: + """Two transactions deleting the last two calls of one prompt must not both + skip GC and leak the text (requires row-level FOR UPDATE).""" + if db_url.startswith("sqlite"): + pytest.skip("SQLite serializes writers instead of row-locking") + import threading + + from xrouter_llm.store import CallRecord, PromptRecord + + id1 = _record(store, ts=1.0, prompt="gc-race") + id2 = _record(store, ts=2.0, prompt="gc-race") + + # Transaction 1 replicates delete()'s exact steps (prompt lock first, + # then call delete) but pauses before its reference check, forcing + # transaction 2 (a real store.delete) to queue on the prompt lock. + session = store._Session() + row = session.get(CallRecord, id1) + prompt_id = row.prompt_id + session.get(PromptRecord, prompt_id, with_for_update=True) + session.delete(row) + session.flush() + + t2_done = threading.Event() + t2 = threading.Thread(target=lambda: (store.delete(id2), t2_done.set())) + t2.start() + assert not t2_done.wait(0.5), "delete() should block on the prompt row lock" + + # The real reference check runs while T2 is concurrently deleting the + # same prompt's other call — under the old lock order (call row first, + # prompt second) this is the step that deadlocked. T2 queued before + # touching its call row, so the locking read finds id2 intact. + still_referenced = session.scalar( + sa.select(CallRecord.id) + .where(CallRecord.prompt_id == prompt_id) + .limit(1) + .with_for_update() + ) + assert still_referenced == id2 + session.commit() + session.close() + assert t2_done.wait(10), "blocked delete() never finished" + t2.join() + + assert store.count() == 0 + assert _prompt_count(store) == 0 + + def test_legacy_db_empty_alembic_version(tmp_path) -> None: """CallStore recovers from a DB where alembic_version exists but is empty.""" url = _legacy_db_empty_version(tmp_path)