rsp2k
6dacdafd6e
backfill_embeddings reads stored call errors (including rows from before the outcome column) and events, rebuilds the records and sends them through the live sink's extraction, redaction and truncation, so later records dedupe against the backfilled rows. backfill_feedback does the same for feedback kept in a host's own tables. A CLI wraps it: python -m fastmcp_feedback.instrumentation.backfill. List and tuple attrs are embedded one item per line instead of as a Python repr; dicts as key: value lines sorted by key, since JSONB drops key order.
566 lines
21 KiB
Python
566 lines
21 KiB
Python
"""Tests for backfill_embeddings, backfill_feedback, its CLI, and list/dict
|
|
text flattening.
|
|
|
|
DB tests take the ``db_url`` fixture, so they run on SQLite and, with
|
|
FFB_TEST_DATABASE_URL set, on Postgres with pgvector.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import pytest
|
|
from sqlalchemy import func, select, update
|
|
|
|
from fastmcp_feedback.instrumentation import (
|
|
REDACTED,
|
|
DatabaseSink,
|
|
EmbeddingSink,
|
|
EventRecord,
|
|
Redactor,
|
|
backfill_embeddings,
|
|
backfill_feedback,
|
|
extract_texts,
|
|
)
|
|
from fastmcp_feedback.instrumentation.backfill.cli import main as cli_main
|
|
|
|
from .test_embeddings import DIM, SECRET, FakeEmbedder, call_record, event, rows
|
|
|
|
T0 = datetime(2026, 9, 1, 12, 0, tzinfo=UTC)
|
|
|
|
|
|
@pytest.fixture
|
|
async def db(db_url):
|
|
made = []
|
|
|
|
def factory(name="calls", **kwargs):
|
|
kwargs.setdefault("create_tables", True)
|
|
kwargs.setdefault("embedding_dim", DIM)
|
|
sink = DatabaseSink(db_url(name), **kwargs)
|
|
made.append(sink)
|
|
return sink
|
|
|
|
yield factory
|
|
for sink in made:
|
|
await sink.engine.dispose()
|
|
|
|
|
|
def at(minutes: float) -> datetime:
|
|
return T0 + timedelta(minutes=minutes)
|
|
|
|
|
|
def error_call(minute, tool="render", message="render timed out", **kw):
|
|
record = call_record(tool=tool, message=message, **kw)
|
|
record.started_at = at(minute)
|
|
return record
|
|
|
|
|
|
def ev(minute, kind, key=None, **attrs):
|
|
record = event(kind, key, **attrs)
|
|
record.occurred_at = at(minute)
|
|
return record
|
|
|
|
|
|
def mixed_records():
|
|
"""One of everything backfill should pick up, and some it should not."""
|
|
return [
|
|
error_call(1, message="render timed out after 30s"),
|
|
error_call(2, tool="find", outcome="soft_error", message="no object Cube"),
|
|
call_record(outcome="ok"),
|
|
ev(
|
|
3, "feedback.submitted", "bug-1", title="Render hangs", description="at 99%"
|
|
),
|
|
ev(
|
|
4, "llm.call", model="m", prompt="summarize the log", completion="it failed"
|
|
),
|
|
ev(5, "llm.call", model="m", ok=False, error="rate limited"),
|
|
ev(6, "llm.call", model="m", ok=True),
|
|
ev(7, "job.failed", "job-1", error_excerpt=["line one", "line two"]),
|
|
ev(8, "job.progress", "job-1", percent=50),
|
|
]
|
|
|
|
|
|
def keyed(stored):
|
|
return {(r["source_type"], r["source_id"]): r for r in stored}
|
|
|
|
|
|
async def count_rows(sink):
|
|
t = sink.embeddings
|
|
async with sink.engine.connect() as conn:
|
|
return (await conn.execute(select(func.count()).select_from(t))).scalar_one()
|
|
|
|
|
|
class TestFlattening:
|
|
def test_list_and_tuple_join_lines(self):
|
|
records = [
|
|
event("job.failed", error_excerpt=[" a ", "", None, "b"]),
|
|
event("job.failed", error=("x", 1, 2.5)),
|
|
]
|
|
texts = [t.text for t in extract_texts(records)]
|
|
assert texts == ["a\nb", "x\n1\n2.5"]
|
|
|
|
def test_dict_becomes_sorted_key_lines(self):
|
|
records = [event("job.failed", error={"zeta": "last", "alpha": 1, "mid": None})]
|
|
[item] = extract_texts(records)
|
|
assert item.text == "alpha: 1\nzeta: last"
|
|
|
|
def test_nested(self):
|
|
records = [
|
|
event(
|
|
"job.failed",
|
|
error_excerpt=[["a", "b"], {"k": ["c", "d"]}, ("e",)],
|
|
)
|
|
]
|
|
[item] = extract_texts(records)
|
|
assert item.text == "a\nb\nk: c\nd\ne"
|
|
assert "[" not in item.text and "{" not in item.text
|
|
|
|
def test_flattening_applies_to_feedback_and_llm(self):
|
|
records = [
|
|
event("feedback.submitted", "f1", title=["Two", "parts"], description=None),
|
|
event("llm.call", prompt=["p1", "p2"], completion={"text": "done"}),
|
|
]
|
|
by_type = {t.source_type: t.text for t in extract_texts(records)}
|
|
assert by_type["feedback"] == "Two\nparts"
|
|
assert by_type["llm"] == "prompt: p1\np2\ncompletion: text: done"
|
|
|
|
def test_empty_containers_are_no_text(self):
|
|
records = [event("job.failed", error_excerpt=[], error={}, message=[[""]])]
|
|
assert extract_texts(records) == []
|
|
|
|
def test_bounded(self):
|
|
deep: list = ["bottom"]
|
|
for _ in range(100):
|
|
deep = [deep]
|
|
huge = ["x" * 1000] * 10_000
|
|
records = [event("job.failed", error=deep), event("job.failed", message=huge)]
|
|
texts = extract_texts(records)
|
|
assert len(texts) == 1 # the deep one is cut off before any text
|
|
assert len(texts[0].text) < 70_000
|
|
|
|
def test_other_values_use_str(self):
|
|
[item] = extract_texts([event("job.failed", error=42)])
|
|
assert item.text == "42"
|
|
|
|
|
|
class TestBackfill:
|
|
async def test_everything_once(self, db):
|
|
store = db()
|
|
await store.write(mixed_records())
|
|
emb = EmbeddingSink(FakeEmbedder(), store)
|
|
counts = await backfill_embeddings(emb)
|
|
assert counts == {
|
|
"feedback": 1,
|
|
"call_error": 2,
|
|
"event": 2,
|
|
"llm": 1,
|
|
"skipped": 0,
|
|
"errors": 0,
|
|
}
|
|
stored = keyed(await rows(store))
|
|
texts = {k[0]: [] for k in stored}
|
|
for (source_type, _), row in stored.items():
|
|
texts[source_type].append(row["text"])
|
|
assert sorted(texts["call_error"]) == [
|
|
"find: TimeoutError: no object Cube",
|
|
"render: TimeoutError: render timed out after 30s",
|
|
]
|
|
assert texts["feedback"] == ["Render hangs\nat 99%"]
|
|
assert texts["llm"] == ["prompt: summarize the log\ncompletion: it failed"]
|
|
assert sorted(texts["event"]) == ["line one\nline two", "rate limited"]
|
|
assert ("feedback", "bug-1") in stored
|
|
|
|
async def test_pre_outcome_error_rows(self, db):
|
|
store = db()
|
|
old = error_call(1, message="boom from before outcome existed")
|
|
old_ok = call_record(outcome="ok")
|
|
await store.write([old, old_ok])
|
|
t = store.table
|
|
async with store.engine.begin() as conn:
|
|
await conn.execute(update(t).values(outcome=None))
|
|
emb = EmbeddingSink(FakeEmbedder(), store)
|
|
counts = await backfill_embeddings(emb)
|
|
assert counts["call_error"] == 1
|
|
[row] = await rows(store)
|
|
assert row["source_id"] == old.id
|
|
assert row["text"] == "render: TimeoutError: boom from before outcome existed"
|
|
|
|
async def test_same_text_hash_and_id_as_live_sink(self, db):
|
|
records = mixed_records() + [
|
|
ev(9, "job.failed", error={"zeta": 1, "alpha": {"b": 2, "a": [1, 2]}}),
|
|
]
|
|
backfilled, live = db("backfilled"), db("live")
|
|
await backfilled.write(records)
|
|
await backfill_embeddings(EmbeddingSink(FakeEmbedder(), backfilled))
|
|
live_emb = EmbeddingSink(FakeEmbedder(), live)
|
|
await live_emb.write(records)
|
|
await live_emb.flush()
|
|
assert live_emb.errors == 0
|
|
|
|
def shape(stored):
|
|
return sorted(
|
|
(r["source_type"], r["source_id"], r["text"], r["text_hash"])
|
|
for r in stored
|
|
)
|
|
|
|
assert shape(await rows(backfilled)) == shape(await rows(live))
|
|
assert len(await rows(live)) == 7
|
|
|
|
async def test_live_records_dedupe_against_backfill(self, db):
|
|
records = mixed_records() + [
|
|
ev(9, "job.failed", error={"zeta": 1, "alpha": {"b": 2, "a": [1, 2]}}),
|
|
]
|
|
store = db()
|
|
await store.write(records)
|
|
await backfill_embeddings(EmbeddingSink(FakeEmbedder(), store))
|
|
before = await count_rows(store)
|
|
embedder = FakeEmbedder()
|
|
live = EmbeddingSink(embedder, store)
|
|
# The records as they come back from the database, as a live sink
|
|
# in another process would have seen them before storing.
|
|
await live.write(records)
|
|
await live.flush()
|
|
assert embedder.calls == []
|
|
assert live.embedded == 0 and live.skipped == 7
|
|
assert await count_rows(store) == before
|
|
|
|
async def test_second_run_embeds_nothing(self, db):
|
|
store = db()
|
|
await store.write(mixed_records())
|
|
embedder = FakeEmbedder()
|
|
emb = EmbeddingSink(embedder, store)
|
|
await backfill_embeddings(emb)
|
|
calls = len(embedder.calls)
|
|
again = await backfill_embeddings(emb)
|
|
assert again == {
|
|
"feedback": 0,
|
|
"call_error": 0,
|
|
"event": 0,
|
|
"llm": 0,
|
|
"skipped": 6,
|
|
"errors": 0,
|
|
}
|
|
assert len(embedder.calls) == calls
|
|
|
|
async def test_changed_text_replaces(self, db):
|
|
store = db()
|
|
record = error_call(1, message="first")
|
|
await store.write([record])
|
|
emb = EmbeddingSink(FakeEmbedder(), store)
|
|
await backfill_embeddings(emb)
|
|
async with store.engine.begin() as conn:
|
|
await conn.execute(
|
|
update(store.table).values(error_message="second"),
|
|
)
|
|
counts = await backfill_embeddings(emb)
|
|
assert counts["call_error"] == 1
|
|
[row] = await rows(store)
|
|
assert row["text"] == "render: TimeoutError: second"
|
|
|
|
async def test_sources(self, db):
|
|
store = db()
|
|
await store.write(mixed_records())
|
|
emb = EmbeddingSink(FakeEmbedder(), store)
|
|
counts = await backfill_embeddings(emb, sources=["feedback", "llm"])
|
|
assert counts == {"feedback": 1, "llm": 1, "skipped": 0, "errors": 0}
|
|
assert {r["source_type"] for r in await rows(store)} == {"feedback", "llm"}
|
|
|
|
async def test_defaults_to_sink_sources(self, db):
|
|
store = db()
|
|
await store.write(mixed_records())
|
|
emb = EmbeddingSink(FakeEmbedder(), store, sources=["call_error"])
|
|
counts = await backfill_embeddings(emb)
|
|
assert counts == {"call_error": 2, "skipped": 0, "errors": 0}
|
|
|
|
async def test_since_until(self, db):
|
|
store = db()
|
|
await store.write([error_call(m, message=f"error {m}") for m in range(10)])
|
|
emb = EmbeddingSink(FakeEmbedder(), store)
|
|
counts = await backfill_embeddings(emb, since=at(3), until=at(6))
|
|
assert counts["call_error"] == 3
|
|
texts = sorted(r["text"] for r in await rows(store))
|
|
assert texts == [f"render: TimeoutError: error {m}" for m in (3, 4, 5)]
|
|
# A naive datetime is UTC.
|
|
naive = (T0 + timedelta(minutes=8)).replace(tzinfo=None)
|
|
counts = await backfill_embeddings(emb, since=naive)
|
|
assert counts["call_error"] == 2
|
|
|
|
async def test_limit_takes_oldest_and_resumes(self, db):
|
|
store = db()
|
|
await store.write([error_call(m, message=f"error {m}") for m in range(10)])
|
|
emb = EmbeddingSink(FakeEmbedder(), store)
|
|
counts = await backfill_embeddings(emb, limit=4, batch_size=3)
|
|
assert counts["call_error"] == 4
|
|
texts = sorted(r["text"] for r in await rows(store))
|
|
assert texts == [f"render: TimeoutError: error {m}" for m in range(4)]
|
|
counts = await backfill_embeddings(emb, limit=4, batch_size=3)
|
|
assert counts["call_error"] == 4 and counts["skipped"] == 4
|
|
counts = await backfill_embeddings(emb, limit=0)
|
|
assert counts["call_error"] == 0
|
|
|
|
async def test_pages_through_many_rows(self, db):
|
|
store = db()
|
|
# Identical timestamps too, so paging must break ties on id.
|
|
records = [error_call(m // 5, message=f"error {m}") for m in range(53)]
|
|
records += [ev(m // 5, "job.failed", error=f"job error {m}") for m in range(41)]
|
|
await store.write(records)
|
|
embedder = FakeEmbedder()
|
|
emb = EmbeddingSink(embedder, store, batch_size=8)
|
|
counts = await backfill_embeddings(emb, batch_size=7)
|
|
assert counts["call_error"] == 53 and counts["event"] == 41
|
|
assert await count_rows(store) == 94
|
|
assert len(embedder.seen) == 94 # nothing read twice
|
|
assert max(len(c) for c in embedder.calls) <= 8
|
|
|
|
async def test_embedder_errors_are_counted(self, db):
|
|
store = db()
|
|
await store.write([error_call(m, message=f"error {m}") for m in range(5)])
|
|
emb = EmbeddingSink(FakeEmbedder(fail=True), store, batch_size=2)
|
|
counts = await backfill_embeddings(emb)
|
|
assert counts["errors"] == 5 and counts["call_error"] == 0
|
|
assert await count_rows(store) == 0
|
|
emb.embedder.fail = False
|
|
counts = await backfill_embeddings(emb)
|
|
assert counts["call_error"] == 5 and counts["errors"] == 0
|
|
|
|
async def test_wrong_vector_count_is_counted(self, db):
|
|
class Short(FakeEmbedder):
|
|
async def embed(self, texts):
|
|
return (await super().embed(texts))[:-1]
|
|
|
|
store = db()
|
|
await store.write([error_call(1), error_call(2, message="other")])
|
|
counts = await backfill_embeddings(EmbeddingSink(Short(), store))
|
|
assert counts["errors"] == 2
|
|
|
|
async def test_redaction(self, db):
|
|
store = db()
|
|
leaked = error_call(1, message=f"auth failed with {SECRET}")
|
|
custom = ev(2, "job.failed", error="ticket ACME-1234 failed")
|
|
await store.write([leaked, custom])
|
|
embedder = FakeEmbedder()
|
|
emb = EmbeddingSink(
|
|
embedder, store, redactor=Redactor(extra_patterns=[r"ACME-\d+"])
|
|
)
|
|
await backfill_embeddings(emb)
|
|
texts = sorted(r["text"] for r in await rows(store))
|
|
assert SECRET not in "".join(embedder.seen + texts)
|
|
assert "ACME-1234" not in "".join(embedder.seen + texts)
|
|
assert f"render: TimeoutError: auth failed with {REDACTED}" in texts
|
|
assert f"ticket {REDACTED} failed" in texts
|
|
|
|
async def test_truncation(self, db):
|
|
store = db()
|
|
await store.write([error_call(1, message="x" * 500)])
|
|
emb = EmbeddingSink(FakeEmbedder(), store, max_chars=50)
|
|
await backfill_embeddings(emb)
|
|
[row] = await rows(store)
|
|
assert len(row["text"]) == 50
|
|
|
|
async def test_dry_run_writes_nothing(self, db):
|
|
store = db()
|
|
await store.write(mixed_records())
|
|
embedder = FakeEmbedder()
|
|
emb = EmbeddingSink(embedder, store)
|
|
counts = await backfill_embeddings(emb, dry_run=True)
|
|
assert counts["call_error"] == 2 and counts["feedback"] == 1
|
|
assert embedder.calls == []
|
|
assert await rows(store) == []
|
|
|
|
async def test_missing_tables_raise(self, db):
|
|
store = db(create_tables=False)
|
|
emb = EmbeddingSink(FakeEmbedder(), store)
|
|
with pytest.raises(RuntimeError, match="does not exist"):
|
|
await backfill_embeddings(emb)
|
|
with pytest.raises(RuntimeError, match="does not exist"):
|
|
await backfill_embeddings(emb, dry_run=True)
|
|
|
|
async def test_bad_arguments(self, db):
|
|
emb = EmbeddingSink(FakeEmbedder(), db())
|
|
with pytest.raises(ValueError, match="unknown sources"):
|
|
await backfill_embeddings(emb, sources=["nope"])
|
|
with pytest.raises(ValueError):
|
|
await backfill_embeddings(emb, batch_size=0)
|
|
with pytest.raises(ValueError):
|
|
await backfill_embeddings(emb, limit=-1)
|
|
|
|
|
|
class TestBackfillFeedback:
|
|
async def test_matches_feedback_event(self, db):
|
|
own, events = db("own"), db("events")
|
|
items = [
|
|
("bug-1", " Render hangs ", "Stuck at 99%"),
|
|
("bug-2", "No description", None),
|
|
("bug-3", "", ""),
|
|
]
|
|
counts = await backfill_feedback(EmbeddingSink(FakeEmbedder(), own), items)
|
|
assert counts == {"feedback": 2, "skipped": 0, "errors": 0}
|
|
live = EmbeddingSink(FakeEmbedder(), events)
|
|
await live.write(
|
|
[
|
|
EventRecord(
|
|
kind="feedback.submitted",
|
|
occurred_at=datetime.now(UTC),
|
|
key=ref,
|
|
attrs={"title": title, "description": desc},
|
|
)
|
|
for ref, title, desc in items
|
|
]
|
|
)
|
|
await live.flush()
|
|
|
|
def shape(stored):
|
|
return sorted((r["source_id"], r["text"], r["text_hash"]) for r in stored)
|
|
|
|
assert shape(await rows(own)) == shape(await rows(events))
|
|
assert shape(await rows(own))[0][1] == "Render hangs\nStuck at 99%"
|
|
|
|
async def test_idempotent_and_dedupes_with_events(self, db):
|
|
store = db()
|
|
emb = EmbeddingSink(FakeEmbedder(), store)
|
|
await backfill_feedback(emb, [("bug-1", "Render hangs", "at 99%")])
|
|
await store.write(
|
|
[
|
|
ev(
|
|
1,
|
|
"feedback.submitted",
|
|
"bug-1",
|
|
title="Render hangs",
|
|
description="at 99%",
|
|
)
|
|
]
|
|
)
|
|
counts = await backfill_embeddings(emb)
|
|
assert counts["feedback"] == 0 and counts["skipped"] == 1
|
|
counts = await backfill_feedback(
|
|
emb, iter([("bug-1", "Render hangs", "at 99%")]), batch_size=1
|
|
)
|
|
assert counts == {"feedback": 0, "skipped": 1, "errors": 0}
|
|
|
|
async def test_redacted(self, db):
|
|
store = db()
|
|
emb = EmbeddingSink(FakeEmbedder(), store)
|
|
await backfill_feedback(emb, [("bug-1", "Leak", f"my key is {SECRET}")])
|
|
[row] = await rows(store)
|
|
assert SECRET not in row["text"] and REDACTED in row["text"]
|
|
|
|
|
|
class TestCli:
|
|
async def seeded(self, db):
|
|
store = db()
|
|
await store.write(mixed_records())
|
|
await store.ensure_embeddings()
|
|
return store
|
|
|
|
async def test_dry_run_main(self, db, capsys):
|
|
store = await self.seeded(db)
|
|
url = store.engine.url.render_as_string(hide_password=False)
|
|
code = await _in_thread(
|
|
cli_main,
|
|
[
|
|
"--database-url",
|
|
url,
|
|
"--dim",
|
|
str(DIM),
|
|
"--dry-run",
|
|
"--sources",
|
|
"call_error,feedback",
|
|
],
|
|
)
|
|
assert code == 0
|
|
out = json.loads(capsys.readouterr().out.strip().splitlines()[-1])
|
|
assert out == {
|
|
"dry_run": True,
|
|
"feedback": 1,
|
|
"call_error": 2,
|
|
"skipped": 0,
|
|
"errors": 0,
|
|
}
|
|
assert await rows(store) == []
|
|
|
|
async def test_dry_run_subprocess(self, db):
|
|
store = await self.seeded(db)
|
|
url = store.engine.url.render_as_string(hide_password=False)
|
|
env = {k: v for k, v in os.environ.items() if k != "FFB_EMBED_API_KEY"}
|
|
proc = await _in_thread(
|
|
subprocess.run,
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"fastmcp_feedback.instrumentation.backfill",
|
|
"--database-url",
|
|
url,
|
|
"--dim",
|
|
str(DIM),
|
|
"--dry-run",
|
|
],
|
|
capture_output=True,
|
|
text=True,
|
|
env=env,
|
|
timeout=60,
|
|
)
|
|
assert proc.returncode == 0, proc.stderr
|
|
out = json.loads(proc.stdout.strip().splitlines()[-1])
|
|
assert out["call_error"] == 2 and out["llm"] == 1 and out["event"] == 2
|
|
assert "RuntimeWarning" not in proc.stderr
|
|
|
|
async def test_needs_embed_url_without_dry_run(self, capsys):
|
|
with pytest.raises(SystemExit) as exc:
|
|
cli_main(["--database-url", "sqlite+aiosqlite:///x.db"])
|
|
assert exc.value.code == 2
|
|
assert "--embed-base-url" in capsys.readouterr().err
|
|
|
|
async def test_no_api_key_flag(self, capsys):
|
|
with pytest.raises(SystemExit):
|
|
cli_main(["--database-url", "sqlite+aiosqlite:///x.db", "--api-key", "k"])
|
|
|
|
async def test_setup_error_exit_code(self, tmp_path, capsys):
|
|
url = f"sqlite+aiosqlite:///{tmp_path / 'empty.db'}"
|
|
code = await _in_thread(cli_main, ["--database-url", url, "--dry-run"])
|
|
assert code == 2
|
|
assert "does not exist" in capsys.readouterr().err
|
|
|
|
async def test_api_key_from_env(self, db, monkeypatch):
|
|
seen = {}
|
|
|
|
class Recorder(FakeEmbedder):
|
|
def __init__(self, base_url, api_key=None, model="m", dim=DIM, **kw):
|
|
super().__init__(dim=dim, model=model)
|
|
seen["api_key"] = api_key
|
|
seen["base_url"] = base_url
|
|
|
|
store = await self.seeded(db)
|
|
url = store.engine.url.render_as_string(hide_password=False)
|
|
monkeypatch.setattr(
|
|
"fastmcp_feedback.instrumentation.backfill.cli.OpenAIEmbedder", Recorder
|
|
)
|
|
monkeypatch.setenv("MY_KEY", "secret-value")
|
|
code = await _in_thread(
|
|
cli_main,
|
|
[
|
|
"--database-url",
|
|
url,
|
|
"--dim",
|
|
str(DIM),
|
|
"--embed-base-url",
|
|
"https://embed.example/v1",
|
|
"--api-key-env",
|
|
"MY_KEY",
|
|
],
|
|
)
|
|
assert code == 0
|
|
assert seen == {
|
|
"api_key": "secret-value",
|
|
"base_url": "https://embed.example/v1",
|
|
}
|
|
assert await count_rows(store) == 6
|
|
|
|
|
|
async def _in_thread(fn, *args, **kwargs):
|
|
# The CLI runs its own event loop, which cannot start on the test's.
|
|
import asyncio
|
|
|
|
return await asyncio.to_thread(fn, *args, **kwargs)
|