Compare commits
21
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
28816477ff | ||
|
|
c81dc3bf8c | ||
|
|
4922eaecbd | ||
|
|
36fed0267c | ||
|
|
ae52e7a0f9 | ||
|
|
a69bf989fd | ||
|
|
f4aef293a5 | ||
|
|
146bbbff25 | ||
|
|
34d5130bd3 | ||
|
|
7378afae8e | ||
|
|
9f3edc59d3 | ||
|
|
ac1d751045 | ||
|
|
6dacdafd6e | ||
|
|
2bd3d9476c | ||
|
|
b9652da288 | ||
|
|
ecef7c4cbc | ||
|
|
c655949bbb | ||
|
|
0bc53a8b3c | ||
|
|
3f3beda8e3 | ||
|
|
b3ded1c267 | ||
|
|
ee46066383 |
+2
-2
@@ -6,8 +6,8 @@ FFB_PG_USER=ffb
|
||||
FFB_PG_PASSWORD=change-me
|
||||
FFB_PG_DB=ffb_test
|
||||
# Published on 127.0.0.1 only. 55432 is commonly taken by other local stacks.
|
||||
FFB_PG_PORT=55433
|
||||
FFB_PG_PORT=55447
|
||||
|
||||
# Read by the test suite. Unset, the Postgres test cases skip.
|
||||
# Each test gets its own database created on this server and dropped after.
|
||||
FFB_TEST_DATABASE_URL=postgresql+asyncpg://ffb:change-me@127.0.0.1:55433/ffb_test
|
||||
FFB_TEST_DATABASE_URL=postgresql+asyncpg://ffb:change-me@127.0.0.1:55447/ffb_test
|
||||
|
||||
+174
@@ -5,6 +5,180 @@ All notable changes to this project will be documented in this file.
|
||||
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
|
||||
and this project adheres to [Calendar Versioning](https://calver.org/) (YYYY.MM.DD).
|
||||
|
||||
## [2026.10.01.1] - 2026-10-01
|
||||
|
||||
A result that says it succeeded is no longer flagged for its own error string.
|
||||
|
||||
### Changed
|
||||
- Explicit success flag wins over error keys: when any of the classifier's
|
||||
`flag_keys` (`ok`, `success`) is exactly `True`, the `error_keys` check is
|
||||
skipped and the string is treated as a warning, so
|
||||
`{"ok": True, "error": "deprecated arg"}` is recorded as `ok` again. Only
|
||||
that check is skipped: a failing status (`{"success": True, "status":
|
||||
"failed"}`), any flag key exactly `False` and `isError` still flag the
|
||||
result. Truthy values that are not `True`, such as `1` or `"yes"`, are not
|
||||
an explicit success. Results with no flag key are classified as in
|
||||
2026.10.01.
|
||||
|
||||
## [2026.10.01] - 2026-10-01
|
||||
|
||||
Results that report failure only as an `error` string now count as soft
|
||||
errors, and the docs cover joining events to calls.
|
||||
|
||||
### Changed
|
||||
- `soft_error_classifier(..., error_keys=("error",))`: a dict result whose
|
||||
value at any of `error_keys` is a non-empty string (after stripping) is a
|
||||
soft error, with message `<key>: <string>`, redacted and truncated as
|
||||
before. Checked after the status, flag and `isError` checks, so results
|
||||
those already flag keep their message and are counted once. None, empty or
|
||||
blank strings, numbers, booleans, dicts and lists do not count, so
|
||||
`{"error": None}` and `{"error": 0.02}` stay ok. `error_keys=()` turns it
|
||||
off. `DEFAULT_RESULT_CLASSIFIER` has it on.
|
||||
- **Behavior change:** results such as `{"played": False, "error": "..."}`,
|
||||
`{"error": "..."}` and `{"status": "unknown_target", "error": "..."}` were
|
||||
recorded as `ok` and are now `soft_error`, so `NOT ok` and soft-error
|
||||
counts can rise after upgrading. This includes the package's own
|
||||
`list_feedback`, `get_feedback_statistics` and analytics tools, whose
|
||||
failure results carry an `error` string and no flag key. Pass
|
||||
`result_classifier=soft_error_classifier(error_keys=())` for the old
|
||||
behavior.
|
||||
|
||||
### Docs
|
||||
- `record_event` / `events_for`: a `key` used for joining must be unique for
|
||||
at least the retention period and across restarts; a per-boot counter
|
||||
collides. Use a UUID or a startup token plus a counter.
|
||||
- The join from a call to its events through an enricher-written `extra`
|
||||
field, for SQLite (`json_extract(c.extra, '$.speech_id') = e.key`) and
|
||||
Postgres (`c.extra->>'speech_id' = e.key`).
|
||||
- Schema notes: `ffb_tool_calls` has primary key `id` and no `call_id`
|
||||
column; `ffb_events.call_id` holds a call's `id`. Ids are UUIDs, so order
|
||||
events by `occurred_at`.
|
||||
- The soft-error section documents `error_keys`, its default and the
|
||||
upgrade note.
|
||||
|
||||
## [2026.09.30.1] - 2026-09-30
|
||||
|
||||
Embed what was stored before embedding was switched on, and embed list and
|
||||
dict attrs as text rather than as Python reprs.
|
||||
|
||||
### Added
|
||||
- `await backfill_embeddings(embedding_sink, *, sources=None, since=None,
|
||||
until=None, limit=None, batch_size=500, dry_run=False)`. Reads the
|
||||
`DatabaseSink`'s tables for calls with outcome `error` or `soft_error`
|
||||
(and, for rows from before the `outcome` column, outcome NULL with an
|
||||
`error_message`), `feedback.submitted` events, `llm.call` events with
|
||||
prompt or completion, and other events with any of the sink's
|
||||
`event_text_keys`. Each record goes through the live sink's extraction,
|
||||
redactor and `max_chars`, so text, `text_hash` and `source_id` match and
|
||||
later records dedupe against backfilled rows. Idempotent: a second run
|
||||
embeds nothing. Keyset-paged by timestamp and id, `batch_size` rows per
|
||||
query; embeds in chunks of the sink's `batch_size`. `limit` embeds at most
|
||||
that many texts, oldest first. Returns counts per source type plus
|
||||
`skipped` and `errors`; embed failures are counted, never raised. Raises
|
||||
only for bad arguments or a missing table. `sources` defaults to the
|
||||
sink's.
|
||||
- `await backfill_feedback(embedding_sink, items, *, batch_size=500,
|
||||
dry_run=False)` for feedback in the host's own tables, as
|
||||
`(ref, title, description)` tuples; the same text a `feedback.submitted`
|
||||
event with that ref gives.
|
||||
- `python -m fastmcp_feedback.instrumentation.backfill --database-url ...
|
||||
--embed-base-url ... [--model --dim --prefix --sources --since --until
|
||||
--limit --batch-size --max-chars --event-text-keys --create-tables
|
||||
--dry-run]`. The API key comes from `FFB_EMBED_API_KEY` or the variable
|
||||
named by `--api-key-env`. Prints counts as JSON; exits 1 when any text
|
||||
failed, 2 on a setup error.
|
||||
|
||||
### Changed
|
||||
- **Embedded text of list, tuple and dict values.** Lists and tuples are
|
||||
embedded one cleaned, non-empty item per line instead of as their repr
|
||||
(`"['a', 'b']"`); dicts as `key: value` lines, sorted by key because
|
||||
PostgreSQL JSONB does not keep key order and a backfill must rebuild the
|
||||
same text. Nested values are flattened the same way, to a depth of 32 and
|
||||
about 64k characters before `max_chars` applies. This applies to event
|
||||
text keys, feedback title and description, and LLM prompt and completion.
|
||||
Existing rows embedded from such values keep their old text; a backfill
|
||||
run replaces them, since their text changed.
|
||||
|
||||
## [2026.09.30] - 2026-09-30
|
||||
|
||||
Embed feedback, tool errors, event text and (opt-in) LLM text, and search
|
||||
them by meaning: "has anyone reported this before?"
|
||||
|
||||
### Added
|
||||
- `EmbeddingSink(embedder, database_sink, *, sources=("feedback",
|
||||
"call_error", "event", "llm"), event_text_keys=("error_excerpt", "error",
|
||||
"message", "reason"), max_chars=4000, max_queue=1000, batch_size=32,
|
||||
timeout=30.0, scan_limit=5000)`. `write` extracts text and puts it on the
|
||||
sink's own bounded queue, returning at once, so a slow embedding backend
|
||||
never reaches the dispatcher's `sink_timeout` or a tool call. A background
|
||||
task embeds in batches and upserts through the `DatabaseSink`'s engine. A
|
||||
full queue drops and counts (`dropped`); embedder and database failures are
|
||||
logged and counted (`errors`), never raised. An unchanged text (same
|
||||
`text_hash`) is not embedded again; a changed one replaces its row. Text is
|
||||
redacted with the middleware's `Redactor` (adopted when the sink is passed
|
||||
to `instrument`) and truncated before it leaves the process. `aclose()`
|
||||
drains for up to `timeout`, then drops and counts the rest.
|
||||
- `Embedder` protocol (`model`, `dim`, `async embed(texts)`) and
|
||||
`OpenAIEmbedder(base_url, api_key=None, model="mxbai-embed-large", dim=1024,
|
||||
batch_size=32, timeout=30.0)` for any OpenAI-compatible `/v1/embeddings`
|
||||
endpoint. It uses `httpx`, or `httpx2` under FastMCP 4; no new dependency.
|
||||
A vector of the wrong length raises `ValueError`.
|
||||
- `await mw.similar(text, *, source_types=None, k=10, min_score=None)` and
|
||||
`await mw.similar_feedback(feedback_ref, *, k=10, source_types=None,
|
||||
min_score=None)`. Results carry `source_type`, `source_id`, `text`, `score`
|
||||
(1 - cosine distance, clamped to [0, 1]) and `created_at`. PostgreSQL ranks
|
||||
with `embedding <=> :q` on an HNSW index; other databases compare in Python
|
||||
against the newest `scan_limit` rows. `similar_feedback` reuses the stored
|
||||
vector and leaves the item itself out. Neither raises; failures log and
|
||||
return [].
|
||||
- **Feedback flows as an event.** With `instrumentation=`, `submit_feedback`
|
||||
records `feedback.submitted` keyed by the feedback id, with `type`, `title`,
|
||||
`description` and `submitter`. Hosts with their own feedback tools record
|
||||
the same event with their own ref (`bug-7Q2X`).
|
||||
- `InstrumentationMiddleware(capture_llm_text=False, max_llm_text_chars=None)`
|
||||
and `record_llm_call(..., prompt=None, completion=None)`. Prompt and
|
||||
completion are stored (redacted, cut to `8 * max_error_chars`, 4000 by
|
||||
default) only when `capture_llm_text=True`.
|
||||
- `embeddings` extra: `pgvector>=0.5.0`. Postgres hosts also need `asyncpg`.
|
||||
- `build_metadata(prefix, metadata=None, embedding_dim=1024)` and
|
||||
`DatabaseSink(embedding_dim=1024)`. `embedding_dim=None` in `build_metadata`
|
||||
leaves the embeddings table out.
|
||||
|
||||
### Changed
|
||||
- **New table `{prefix}embeddings`** (`id`, `source_type`, `source_id`,
|
||||
`model`, `dim`, `text_hash`, `text`, `created_at`, `embedding`), unique on
|
||||
(`source_type`, `source_id`, `model`). `embedding` is pgvector `vector(dim)`
|
||||
with an HNSW `vector_cosine_ops` index on PostgreSQL, a JSON float list
|
||||
elsewhere. `build_metadata()` now includes it, so Alembic hosts need a
|
||||
migration; the README has the SQL, starting with
|
||||
`CREATE EXTENSION IF NOT EXISTS vector`. With `create_tables=True` the table
|
||||
is created when an `EmbeddingSink` first writes, not at startup, so hosts
|
||||
that do not embed need neither pgvector nor the extension. On PostgreSQL
|
||||
that step first runs `CREATE EXTENSION IF NOT EXISTS vector`; without the
|
||||
privilege it logs a warning, keeps the sink running, and retries at most
|
||||
once a minute.
|
||||
- `DatabaseSink.prune()` also deletes embeddings older than the cutoff,
|
||||
except those of feedback, and returns the total from all three tables.
|
||||
- `record_llm_call` now drops `prompt` and `completion` keys passed in
|
||||
`attrs` unless `capture_llm_text=True`. Earlier versions stored them.
|
||||
|
||||
### Documentation
|
||||
- README: `DatabaseSink` needs the `instrumentation` extra even with
|
||||
SQLAlchemy installed, since SQLAlchemy 2.1 no longer installs `greenlet`.
|
||||
|
||||
## [2026.09.29.1] - 2026-09-29
|
||||
|
||||
### Fixed
|
||||
- `DatabaseSink(create_tables=True)`: two processes starting at the same
|
||||
moment against a database without the tables both ran `create_all`, and one
|
||||
failed (SQLite: `table ... already exists`; Postgres: a unique violation on
|
||||
`pg_type_typname_nsp_index`). That process dropped its first batch and
|
||||
counted a sink error. `ensure_tables()` now waits 50 to 250 ms after a
|
||||
database error and runs `create_all` once more, which skips what the other
|
||||
process created; the recovery is logged at INFO. A second failure is raised
|
||||
as before. Measured on two-writer cold starts: SQLite 9 of 10 runs lost a
|
||||
batch before, 0 of 10 after; Postgres 16 12 of 12 before, 0 of 12 after.
|
||||
|
||||
## [2026.09.29] - 2026-09-29
|
||||
|
||||
Record things that are not tool calls: background job transitions, LLM
|
||||
|
||||
+15
-1
@@ -113,7 +113,7 @@ make pg-down # stop it; the data volume is kept (pg-reset deletes it)
|
||||
```
|
||||
|
||||
The server listens on `127.0.0.1` only, on `FFB_PG_PORT` from `.env`
|
||||
(55433 by default; change it if that port is taken). `make pg-psql` opens a
|
||||
(55447 by default; change it if that port is taken). `make pg-psql` opens a
|
||||
shell on it. Every Postgres test gets its own freshly created database with
|
||||
the `vector` extension, dropped when the test ends, so tests never see each
|
||||
other's rows. To write a test that runs on both backends, take the `db_url`
|
||||
@@ -127,6 +127,20 @@ needs permission to create and drop databases:
|
||||
FFB_TEST_DATABASE_URL=postgresql+asyncpg://user:pass@host:5432/postgres uv run pytest -m postgres
|
||||
```
|
||||
|
||||
### Embedding tests against a real endpoint
|
||||
|
||||
`tests/unit/test_embeddings.py` uses a deterministic fake embedder. One test
|
||||
embeds through a real OpenAI-compatible endpoint when you point it at one, and
|
||||
skips otherwise:
|
||||
|
||||
```bash
|
||||
FFB_TEST_EMBED_URL=https://api.example.com/v1 FFB_TEST_EMBED_KEY=... \
|
||||
uv run pytest tests/unit/test_embeddings.py -k live
|
||||
```
|
||||
|
||||
`FFB_TEST_EMBED_MODEL` (default `mxbai-embed-large`) and `FFB_TEST_EMBED_DIM`
|
||||
(default 1024) pick the model.
|
||||
|
||||
### Writing Tests
|
||||
|
||||
- Place unit tests in `tests/unit/`
|
||||
|
||||
@@ -59,4 +59,5 @@ test-all: test test-pg
|
||||
|
||||
lint:
|
||||
uv run --no-sync ruff check src/ examples/ tests/_compat.py tests/_databases.py \
|
||||
tests/unit/test_test_databases.py tests/unit/test_instrumentation.py
|
||||
tests/unit/test_test_databases.py tests/unit/test_instrumentation.py \
|
||||
tests/unit/test_embeddings.py tests/unit/test_backfill.py
|
||||
|
||||
@@ -181,6 +181,8 @@ instrument(app, [
|
||||
|
||||
`DatabaseSink` needs the `instrumentation` extra
|
||||
(`uv add "fastmcp-feedback[instrumentation]"`), plus `asyncpg` for Postgres.
|
||||
Install the extra even if you already have SQLAlchemy: SQLAlchemy 2.1 no longer
|
||||
pulls in `greenlet` by itself, and its async engine cannot run without it.
|
||||
|
||||
### Using your own database and migrations
|
||||
|
||||
@@ -222,6 +224,49 @@ The default classifier flags a dict result when `ok` or `success` is exactly
|
||||
`False`, when `status` or `state` is `error`, `failed` or `failure` (any case),
|
||||
or when the tool marked the result as an MCP error (`isError`) without raising.
|
||||
The message comes from the first of `error`, `message`, `detail` and `reason`.
|
||||
|
||||
Failing those, a result whose `error` holds a non-empty string is a soft
|
||||
error too, with that string as the message: `{"played": False, "error":
|
||||
"audio output stalled"}` is recorded as `error: audio output stalled`. Only a
|
||||
string with something other than whitespace counts, so `{"error": None}`,
|
||||
`{"error": ""}` and a statistic such as `{"error": 0.02}` stay ok. A dict or
|
||||
list there (`{"error": {"code": 1}}`) is not flagged either; catch that shape
|
||||
with a flag key, a status or your own classifier. `error_keys` names the keys
|
||||
to check, and `error_keys=()` turns the check off.
|
||||
|
||||
A result that says it succeeded wins over its own error string: when `ok` or
|
||||
`success` (the `flag_keys`) is exactly `True`, the `error` check is skipped and
|
||||
the string is read as a warning, so `{"ok": True, "error": "deprecated arg"}`
|
||||
is recorded as `ok`. Only that check is skipped. In order of precedence:
|
||||
|
||||
1. A failing `status` or `state` flags the result, even next to `ok: True`,
|
||||
since `{"success": True, "status": "failed"}` contradicts itself and the
|
||||
status is the more specific of the two.
|
||||
2. Any flag key exactly `False` flags it, even if another is `True`.
|
||||
3. `isError` flags it.
|
||||
4. A flag key exactly `True` makes it ok. Truthy values that are not `True`,
|
||||
such as `1` or `"yes"`, do not count.
|
||||
5. A non-empty string at an `error_keys` key flags it.
|
||||
|
||||
```python
|
||||
from fastmcp_feedback.instrumentation import DEFAULT_RESULT_CLASSIFIER as classify
|
||||
|
||||
assert classify("t", {"ok": True, "error": "deprecated arg"}) is None
|
||||
assert classify("t", {"ok": 1, "error": "x"}) == "error: x"
|
||||
assert classify("t", {"success": True, "status": "failed"}) == "status=failed"
|
||||
assert classify("t", {"played": False, "error": "stalled"}) == "error: stalled"
|
||||
```
|
||||
|
||||
```python
|
||||
from fastmcp_feedback.instrumentation import instrument, soft_error_classifier
|
||||
|
||||
# Tools that report {"failure_reason": "..."} instead
|
||||
instrument(app, result_classifier=soft_error_classifier(error_keys=("failure_reason",)))
|
||||
|
||||
# As before 2026.10.01: an error string alone is not a soft error
|
||||
instrument(app, result_classifier=soft_error_classifier(error_keys=()))
|
||||
```
|
||||
|
||||
Add your own status values, or pass any `(tool, result) -> str | None`:
|
||||
|
||||
```python
|
||||
@@ -258,6 +303,19 @@ table made earlier by `create_tables=True` needs
|
||||
tables are not altered. Until the column exists, database writes fail and
|
||||
those records are dropped (logged).
|
||||
|
||||
**Upgrading to 2026.10.01:** the `error` key check is on by default, so a
|
||||
result that carries a non-empty `error` string with no failing status or flag
|
||||
(such as `{"status": "unknown_target", "error": "no object named Cube"}`, or
|
||||
`{"error": "..."}` alone) is now a soft error where earlier releases recorded
|
||||
it as `ok`. Counts of `NOT ok` and `outcome = 'soft_error'` can rise after
|
||||
upgrading. Results that also have a failing status or flag keep their old
|
||||
message, such as `success=False: ...`, and are counted once. Pass
|
||||
`soft_error_classifier(error_keys=())` to keep the old behavior.
|
||||
|
||||
**Upgrading to 2026.10.01.1:** a result with `ok` or `success` exactly `True`
|
||||
and an `error` string, which 2026.10.01 recorded as `soft_error`, is `ok` again.
|
||||
Nothing else changes.
|
||||
|
||||
### Identity and enrichment hooks
|
||||
|
||||
```python
|
||||
@@ -441,7 +499,12 @@ background queue, the same drop-and-count policy when the queue is full.
|
||||
- `kind` is a dotted name (`job.submitted`, `llm.call`), a non-empty string of
|
||||
at most 128 characters.
|
||||
- `key` is what you join on, such as a job id. It is stored as a string, at
|
||||
most 255 characters.
|
||||
most 255 characters. A key used to join events to calls or to each other
|
||||
must be unique for at least the retention period, **including across
|
||||
restarts**. A per-boot counter (`speech-1`, `speech-2`, ...) starts over
|
||||
on every restart, and the second boot's `speech-1` joins to the first
|
||||
boot's events. Use a UUID, or a token chosen at startup plus a counter
|
||||
(`speech-<boot token>-<n>`).
|
||||
- `attrs` is any JSON-like dict. It is redacted like call arguments, then made
|
||||
JSON-safe: `NaN` and infinities become the strings `"nan"`/`"inf"` (Postgres
|
||||
JSONB rejects them), datetimes become ISO strings, anything else unknown
|
||||
@@ -466,6 +529,23 @@ for ev in await mw.events_for("job-wing"):
|
||||
steps = await mw.events_for("job-wing", kinds=["step.started", "step.finished"], limit=100)
|
||||
```
|
||||
|
||||
To join a call to its events in SQL, have the enricher write the key into
|
||||
`extra` and match on it. This is how mcspeak reads the phases of one `speak`
|
||||
call, with SQLite:
|
||||
|
||||
```sql
|
||||
SELECT e.occurred_at, e.attrs FROM ffb_events e
|
||||
JOIN ffb_tool_calls c ON json_extract(c.extra, '$.speech_id') = e.key
|
||||
WHERE e.kind = 'speech.phase' AND c.id = ?
|
||||
ORDER BY e.occurred_at;
|
||||
```
|
||||
|
||||
On Postgres the join condition is `c.extra->>'speech_id' = e.key`. Events
|
||||
recorded inside the call also carry its id, so `e.call_id = c.id` joins those
|
||||
without a key. `ffb_tool_calls` has no `call_id` column; its primary key is
|
||||
`id`, and `ffb_events.call_id` holds that value. Order events by
|
||||
`occurred_at`: ids are random UUIDs and do not sort in recording order.
|
||||
|
||||
`events_for` reads the `DatabaseSink` and adds events still waiting in the
|
||||
queue, so an event is visible as soon as `record_event` returns. With more
|
||||
than `limit` (1000) matches it returns the newest `limit`. Without a
|
||||
@@ -506,13 +586,35 @@ on SQLite):
|
||||
Anything in `attrs` is kept alongside, but these names win. `key=` works as
|
||||
for `record_event`, for example a chat turn id.
|
||||
|
||||
The prompt and the answer are not stored unless you opt in, since chat text
|
||||
can hold anything a user typed. With `capture_llm_text=True` on the
|
||||
middleware, `record_llm_call(..., prompt=..., completion=...)` stores them as
|
||||
`prompt` and `completion` attributes, redacted and cut to `max_llm_text_chars`
|
||||
(by default `8 * max_error_chars`, so 4000 characters):
|
||||
|
||||
```python
|
||||
from fastmcp_feedback.instrumentation import instrument
|
||||
|
||||
mw = instrument(app, capture_llm_text=True)
|
||||
mw.record_llm_call(
|
||||
model="claude-opus-4",
|
||||
prompt="Why does the render time out?",
|
||||
completion="The scene has 40 million polygons.",
|
||||
)
|
||||
```
|
||||
|
||||
With the default `capture_llm_text=False` both are discarded, and so are
|
||||
`prompt` and `completion` keys passed in `attrs`.
|
||||
|
||||
#### Storage and sinks
|
||||
|
||||
Events go to a `ffb_events` table: `id`, `occurred_at`, `kind`, `key`,
|
||||
`call_id`, `session_id`, `user_sub`, `caller_kind`, `client_id`,
|
||||
`server_version`, `attrs` (JSON; JSONB on Postgres), indexed on
|
||||
`occurred_at`, `kind`, `key`, `call_id` and `user_sub`. It has no foreign key
|
||||
to `ffb_tool_calls`, since calls are pruned and may be written later.
|
||||
`occurred_at`, `kind`, `key`, `call_id` and `user_sub`. `id` is a UUID
|
||||
string in both tables, so sort by `occurred_at` (events) or `started_at`
|
||||
(calls), never by `id`. `call_id` refers to `ffb_tool_calls.id` but has no
|
||||
foreign key, since calls are pruned and may be written later.
|
||||
|
||||
Sinks get calls and events in the same batches. `DatabaseSink` writes each to
|
||||
its table, calls first in their own transaction, so a database without
|
||||
@@ -534,6 +636,224 @@ includes it). `create_tables=True` creates it on first use, since it creates
|
||||
missing tables. Until the table exists, event writes fail and those events
|
||||
are dropped (logged); tool calls and pruning are unaffected.
|
||||
|
||||
### Embeddings: finding similar reports and errors
|
||||
|
||||
"Has anyone reported this before?" is hard to answer with `LIKE`. An
|
||||
`EmbeddingSink` turns the text in the records into vectors and stores them, so
|
||||
you can search by meaning. It embeds four kinds of text:
|
||||
|
||||
| `source_type` | Text | `source_id` |
|
||||
|---------------|------|-------------|
|
||||
| `feedback` | title and description of a `feedback.submitted` event | the feedback ref |
|
||||
| `call_error` | `"{tool}: {error_type}: {error_message}"` of calls with outcome `error` or `soft_error` | the call id |
|
||||
| `event` | the values of `error_excerpt`, `error`, `message`, `reason` in an event's attrs (`event_text_keys=`) | the event id |
|
||||
| `llm` | `"prompt: ...\ncompletion: ..."` of an `llm.call` event, when `capture_llm_text=True` | the event id |
|
||||
|
||||
An `llm.call` without prompt or completion counts as a plain event, so a
|
||||
failed LLM call's `error` is still embedded. A list or tuple value is
|
||||
embedded one item per line, a dict as `key: value` lines sorted by key (at
|
||||
any depth), rather than as its Python repr.
|
||||
|
||||
It needs a `DatabaseSink` to store into and an embedder. `OpenAIEmbedder`
|
||||
talks to any OpenAI-compatible `/v1/embeddings` endpoint (OpenAI, Ollama,
|
||||
LiteLLM, vLLM):
|
||||
|
||||
```python
|
||||
import os
|
||||
|
||||
from fastmcp_feedback import add_feedback_tools
|
||||
from fastmcp_feedback.instrumentation import (
|
||||
DatabaseSink,
|
||||
EmbeddingSink,
|
||||
OpenAIEmbedder,
|
||||
instrument,
|
||||
)
|
||||
|
||||
embedder = OpenAIEmbedder(
|
||||
os.environ["EMBED_URL"], # e.g. https://api.example.com/v1
|
||||
api_key=os.environ["EMBED_API_KEY"],
|
||||
model="mxbai-embed-large",
|
||||
dim=1024,
|
||||
)
|
||||
db = DatabaseSink("sqlite+aiosqlite:///calls.db", create_tables=True, embedding_dim=1024)
|
||||
mw = instrument(app, [db, EmbeddingSink(embedder, db)])
|
||||
add_feedback_tools(app, database_url="sqlite:///feedback.db", instrumentation=mw)
|
||||
```
|
||||
|
||||
Then search:
|
||||
|
||||
```python
|
||||
for hit in await mw.similar("render timed out", k=5):
|
||||
print(f"{hit['score']:.2f}", hit["source_type"], hit["source_id"], hit["text"][:60])
|
||||
|
||||
# Reports and errors like feedback 42, leaving 42 itself out.
|
||||
related = await mw.similar_feedback("42", source_types=["feedback", "call_error"])
|
||||
```
|
||||
|
||||
Each hit has `source_type`, `source_id`, `text` (what was embedded),
|
||||
`score` and `created_at`. `score` is 1 minus the cosine distance, clamped to
|
||||
[0, 1]: 1 is the same direction, 0 unrelated or opposite. `min_score=` drops
|
||||
weaker hits, and `source_types=` limits the search. Only vectors from the
|
||||
embedder's `model` are compared. `similar_feedback` uses the stored vector of
|
||||
that feedback, so it embeds nothing, and returns [] until the feedback has
|
||||
been embedded. Neither ever raises; on any failure they log and return [].
|
||||
|
||||
On PostgreSQL the search runs in the database (`ORDER BY embedding <=> :q`)
|
||||
on an HNSW index. pgvector applies `source_types` and the model filter after
|
||||
the index scan, so a narrow filter over a large table can return fewer than
|
||||
`k` hits; `SET hnsw.iterative_scan = relaxed_order` (pgvector 0.8) fixes that.
|
||||
Elsewhere it compares in Python against the newest `scan_limit` (5000) rows,
|
||||
and logs once when there are more.
|
||||
|
||||
**It stays off the call path.** `write` only picks out the text and puts it on
|
||||
the sink's own bounded queue (`max_queue`, 1000), so a slow embedding backend
|
||||
never holds up the dispatcher or a tool call. A background task embeds up to
|
||||
`batch_size` (32) texts per request, each request limited to `timeout` (30 s),
|
||||
and upserts the rows. A full queue drops texts and counts them in
|
||||
`sink.dropped`; failures of the embedder or the database are logged and
|
||||
counted in `sink.errors`, never raised. A text whose row already holds the
|
||||
same text (by `text_hash`) is not embedded again; a changed text replaces its
|
||||
row. `await sink.flush()` waits for the queue, and `mw.aclose()` lets it drain
|
||||
for up to `timeout` before stopping.
|
||||
|
||||
**Nothing new leaves the process unredacted.** Text is redacted with the
|
||||
middleware's `Redactor` (the sink adopts it when passed to `instrument`) and
|
||||
cut to `max_chars` (4000) before it is sent to the embedder or stored, on top
|
||||
of the redaction every record already had.
|
||||
|
||||
**Feedback flows as an event.** With `instrumentation=mw`, `submit_feedback`
|
||||
records a `feedback.submitted` event keyed by the feedback id, with `type`,
|
||||
`title`, `description` and `submitter` attributes. If you have your own
|
||||
feedback tools, record the same event with your own ref and it is embedded
|
||||
the same way:
|
||||
|
||||
```python
|
||||
mw.record_event(
|
||||
"feedback.submitted",
|
||||
key="bug-7Q2X",
|
||||
attrs={"type": "bug", "title": "Render hangs", "description": "Stuck at 99%"},
|
||||
)
|
||||
```
|
||||
|
||||
**Storage.** Vectors go to `ffb_embeddings`: `id`, `source_type`,
|
||||
`source_id`, `model`, `dim`, `text_hash` (sha256 of `text`), `text`,
|
||||
`created_at` and `embedding`, unique on (`source_type`, `source_id`,
|
||||
`model`). `embedding` is a pgvector `vector(dim)` on PostgreSQL and a JSON
|
||||
list of floats elsewhere. `DatabaseSink(embedding_dim=...)` sets `dim` and must
|
||||
equal the embedder's `dim`, or `EmbeddingSink` raises `ValueError`. Retention
|
||||
covers it: `prune()` deletes embeddings older than the cutoff, **except those
|
||||
of feedback**, which are kept like the calls linked to feedback.
|
||||
|
||||
On PostgreSQL this needs the `embeddings` extra (`pgvector`) and `asyncpg`:
|
||||
|
||||
```bash
|
||||
uv add "fastmcp-feedback[instrumentation,embeddings]" asyncpg
|
||||
```
|
||||
|
||||
With `create_tables=True`, the table is created when an `EmbeddingSink` first
|
||||
writes, not before, so hosts that do not embed never need pgvector. On
|
||||
PostgreSQL it first runs `CREATE EXTENSION IF NOT EXISTS vector`. Creating an
|
||||
extension needs a superuser or the database owner on most setups; without
|
||||
that it logs a warning, and so does the table creation if the extension is
|
||||
still missing. The sink keeps working and counts the failed writes, and
|
||||
retries the creation at most once a minute.
|
||||
|
||||
**Upgrading:** `ffb_embeddings` is new in 2026.09.30. `build_metadata()`
|
||||
includes it (pass `embedding_dim=None` to leave it out). Hosts on Alembic need
|
||||
a migration; for PostgreSQL, with the default prefix and 1024 dimensions:
|
||||
|
||||
```sql
|
||||
CREATE EXTENSION IF NOT EXISTS vector;
|
||||
|
||||
CREATE TABLE ffb_embeddings (
|
||||
id VARCHAR(36) NOT NULL,
|
||||
source_type VARCHAR(16) NOT NULL,
|
||||
source_id VARCHAR(255) NOT NULL,
|
||||
model VARCHAR(128) NOT NULL,
|
||||
dim INTEGER NOT NULL,
|
||||
text_hash VARCHAR(64) NOT NULL,
|
||||
text TEXT NOT NULL,
|
||||
created_at TIMESTAMP WITH TIME ZONE NOT NULL,
|
||||
embedding VECTOR(1024) NOT NULL,
|
||||
PRIMARY KEY (id),
|
||||
CONSTRAINT ffb_embeddings_source_uq UNIQUE (source_type, source_id, model)
|
||||
);
|
||||
CREATE INDEX ffb_embeddings_hnsw ON ffb_embeddings USING hnsw (embedding vector_cosine_ops);
|
||||
CREATE INDEX ix_ffb_embeddings_created_at ON ffb_embeddings (created_at);
|
||||
CREATE INDEX ix_ffb_embeddings_source_id ON ffb_embeddings (source_id);
|
||||
CREATE INDEX ix_ffb_embeddings_source_type ON ffb_embeddings (source_type);
|
||||
```
|
||||
|
||||
Wrap it in `op.execute(...)` in the migration's `upgrade()`, and
|
||||
`DROP TABLE ffb_embeddings` in `downgrade()`. Autogenerate also works but
|
||||
renders the column as `fastmcp_feedback.instrumentation.storage.EmbeddingVector(dim=1024)`,
|
||||
which ties the migration to this package.
|
||||
|
||||
#### Backfilling existing data
|
||||
|
||||
The sink embeds records as they arrive. Calls, events and feedback stored
|
||||
before it was running can be embedded afterwards with `backfill_embeddings`,
|
||||
which reads them back from the `DatabaseSink`'s tables:
|
||||
|
||||
```python
|
||||
from fastmcp_feedback.instrumentation import backfill_embeddings
|
||||
|
||||
counts = await backfill_embeddings(mw.embedding_sink)
|
||||
# {"feedback": 4, "call_error": 90, "event": 12, "llm": 0, "skipped": 0, "errors": 0}
|
||||
```
|
||||
|
||||
It picks up the same four kinds of text as the live sink: calls with outcome
|
||||
`error` or `soft_error`, `feedback.submitted` events, `llm.call` events with
|
||||
a prompt or completion, and other events with any of the sink's
|
||||
`event_text_keys`. Calls recorded before the `outcome` column existed
|
||||
(2026.09.27.4) have it NULL; those with an `error_message` count as errors.
|
||||
|
||||
Each text goes through the sink's own extraction, redactor and `max_chars`,
|
||||
so it comes out with the same text, `text_hash` and `source_id` the live sink
|
||||
would give that record. That makes it idempotent: a second run embeds nothing
|
||||
and reports the rows as `skipped`, and records arriving later dedupe against
|
||||
what was backfilled. When the middleware has a custom `Redactor`, build the
|
||||
sink with it (`EmbeddingSink(..., redactor=mw.redactor)`), or use
|
||||
`mw.embedding_sink`, which already has it.
|
||||
|
||||
`sources=` narrows it (default: the sink's sources), `since=` and `until=`
|
||||
select rows by timestamp (naive means UTC), and `limit=` embeds at most that
|
||||
many texts, oldest first; running it again carries on where it stopped. Rows
|
||||
are read `batch_size` (500) at a time and embedded in chunks of the sink's
|
||||
`batch_size`, so a large table never sits in memory. A failed chunk is
|
||||
logged and counted in `errors`, never raised; only a missing table or bad
|
||||
arguments raise. `dry_run=True` counts what would be embedded without calling
|
||||
the embedder.
|
||||
|
||||
Feedback kept in your own tables rather than as events goes through
|
||||
`backfill_feedback`, which makes the same text a `feedback.submitted` event
|
||||
with that ref would:
|
||||
|
||||
```python
|
||||
from fastmcp_feedback.instrumentation import backfill_feedback
|
||||
|
||||
rows = [("bug-7Q2X", "Render hangs", "Stuck at 99%"), ("bug-8K1D", "Crash", None)]
|
||||
await backfill_feedback(mw.embedding_sink, rows) # (ref, title, description)
|
||||
```
|
||||
|
||||
From a shell, the same over a database URL. The API key comes from
|
||||
`FFB_EMBED_API_KEY` (or the variable `--api-key-env` names), never a flag:
|
||||
|
||||
```bash
|
||||
export FFB_EMBED_API_KEY=...
|
||||
python -m fastmcp_feedback.instrumentation.backfill \
|
||||
--database-url postgresql+asyncpg://user:pass@db/app \
|
||||
--embed-base-url https://api.example.com/v1 \
|
||||
--model mxbai-embed-large --dim 1024 --prefix ffb_ \
|
||||
--since 2026-09-01 --sources call_error,feedback --dry-run
|
||||
```
|
||||
|
||||
It prints the counts as JSON and exits 1 if any text failed to embed, 2 on
|
||||
a setup error. `--create-tables` creates `ffb_embeddings` (and on PostgreSQL
|
||||
the vector extension) when it is missing; without it a missing table is an
|
||||
error. `--help` lists the rest (`--until`, `--limit`, `--batch-size`,
|
||||
`--max-chars`, `--event-text-keys`).
|
||||
|
||||
### Middleware order and shutdown
|
||||
|
||||
FastMCP runs the first middleware added as the outermost. Call `instrument()`
|
||||
|
||||
+1
-1
@@ -11,7 +11,7 @@ services:
|
||||
POSTGRES_PASSWORD: ${FFB_PG_PASSWORD}
|
||||
POSTGRES_DB: ${FFB_PG_DB}
|
||||
ports:
|
||||
- "127.0.0.1:${FFB_PG_PORT:-55433}:5432"
|
||||
- "127.0.0.1:${FFB_PG_PORT:-55447}:5432"
|
||||
volumes:
|
||||
- ffb-postgres-data:/var/lib/postgresql/data
|
||||
- ./docker/postgres-init:/docker-entrypoint-initdb.d:ro
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
# Postgres on Alpine with pgvector built in. Keeps the musl C library, so an
|
||||
# existing data directory's collation and text indexes are unaffected.
|
||||
# with_llvm=no: the base image's Postgres was built with a specific clang for
|
||||
# JIT bitcode that Alpine may no longer ship; pgvector itself builds with gcc.
|
||||
ARG PG_MAJOR=16
|
||||
FROM postgres:${PG_MAJOR}-alpine
|
||||
ARG PGVECTOR_VERSION=v0.8.6
|
||||
RUN apk add --no-cache --virtual .build-deps git build-base \
|
||||
&& git clone --depth 1 --branch ${PGVECTOR_VERSION} https://github.com/pgvector/pgvector.git /tmp/pgvector \
|
||||
&& make -C /tmp/pgvector OPTFLAGS="" with_llvm=no \
|
||||
&& make -C /tmp/pgvector install with_llvm=no \
|
||||
&& rm -rf /tmp/pgvector \
|
||||
&& apk del .build-deps
|
||||
@@ -0,0 +1,28 @@
|
||||
# Postgres on Alpine with pgvector
|
||||
|
||||
For stacks that already run `postgres:NN-alpine` and want the `vector`
|
||||
extension for fastmcp-feedback embeddings.
|
||||
|
||||
Why not switch to the `pgvector/pgvector` image: it's Debian (glibc), and the
|
||||
Alpine images use musl. The two C libraries sort text differently under
|
||||
locales such as `en_US.utf8`, so reusing an Alpine data directory under a
|
||||
Debian image changes the order of every text index. The indexes need a full
|
||||
`REINDEX` or they silently return wrong results. Building pgvector into the
|
||||
same Alpine image keeps the data directory's collation unchanged.
|
||||
|
||||
```bash
|
||||
docker build --build-arg PG_MAJOR=16 -t postgres-pgvector:16-alpine .
|
||||
docker build --build-arg PG_MAJOR=17 -t postgres-pgvector:17-alpine .
|
||||
```
|
||||
|
||||
Then point the stack's Postgres service at the new image (same major version
|
||||
as the existing data directory) and run `CREATE EXTENSION IF NOT EXISTS
|
||||
vector;` once, or let a migration do it.
|
||||
|
||||
Tested 2026-09-30 with pgvector v0.8.6 on postgres:16-alpine and
|
||||
postgres:17-alpine: the extension loads, and an HNSW `vector_cosine_ops`
|
||||
index is used for `ORDER BY embedding <=> ...` queries.
|
||||
|
||||
`with_llvm=no` skips the LLVM bitcode used for JIT inlining. The base images'
|
||||
Postgres was built with a clang version Alpine no longer ships, so the bitcode
|
||||
step fails; pgvector itself still builds with gcc and works normally.
|
||||
+3
-1
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "fastmcp-feedback"
|
||||
version = "2026.09.29" # YYYYMMDD.patch versioning - update for each release
|
||||
version = "2026.10.01.1" # YYYYMMDD.patch versioning - update for each release
|
||||
description = "Production-ready feedback collection system for FastMCP servers"
|
||||
readme = "README.md"
|
||||
license = "MIT"
|
||||
@@ -52,6 +52,8 @@ dependencies = [
|
||||
postgresql = ["psycopg2-binary>=2.9.0"]
|
||||
# Async database sink for per-call instrumentation. Postgres hosts add asyncpg.
|
||||
instrumentation = ["sqlalchemy[asyncio]>=2.0.43", "aiosqlite>=0.22.1"]
|
||||
# Vector column for EmbeddingSink on Postgres. Postgres hosts add asyncpg.
|
||||
embeddings = ["pgvector>=0.5.0"]
|
||||
mysql = ["pymysql>=1.1.0"]
|
||||
dev = [
|
||||
"sqlalchemy[asyncio]>=2.0.43",
|
||||
|
||||
@@ -44,7 +44,7 @@ Quick Start:
|
||||
import logging
|
||||
|
||||
# Version information - CalVer YYYY.MM.DD
|
||||
__version__ = "2026.09.29"
|
||||
__version__ = "2026.10.01.1"
|
||||
__author__ = "Ryan Malloy"
|
||||
__email__ = "ryan@supported.systems"
|
||||
|
||||
|
||||
@@ -6,14 +6,28 @@ mw = instrument(app) # every tool call as JSON lines on stderr
|
||||
mw = instrument(app, [DatabaseSink(engine)], identity_resolver=who)
|
||||
mw = instrument(app, result_classifier=soft_error_classifier(extra_statuses=("unknown_target",)))
|
||||
mw.record_event("job.finished", key=job_id, attrs={"duration_s": 3600})
|
||||
|
||||
db = DatabaseSink(engine)
|
||||
mw = instrument(app, [db, EmbeddingSink(OpenAIEmbedder(url, key), db)])
|
||||
await mw.similar("render timed out")
|
||||
await backfill_embeddings(mw.embedding_sink) # rows stored before embedding
|
||||
"""
|
||||
|
||||
from .backfill import backfill_embeddings, backfill_feedback
|
||||
from .classify import (
|
||||
DEFAULT_RESULT_CLASSIFIER,
|
||||
ResultClassifier,
|
||||
soft_error_classifier,
|
||||
)
|
||||
from .dispatch import RecordDispatcher
|
||||
from .embeddings import (
|
||||
DEFAULT_EVENT_TEXT_KEYS,
|
||||
Embedder,
|
||||
EmbeddingSink,
|
||||
EmbeddingText,
|
||||
OpenAIEmbedder,
|
||||
extract_texts,
|
||||
)
|
||||
from .events import EventRecord
|
||||
from .middleware import MODE_ENV_VAR, MODES, InstrumentationMiddleware, instrument
|
||||
from .record import ToolCallRecord
|
||||
@@ -24,28 +38,46 @@ from .redaction import (
|
||||
Redactor,
|
||||
)
|
||||
from .sinks import CallbackSink, JsonLinesSink, MemorySink, Sink
|
||||
from .storage import DEFAULT_TABLE_PREFIX, DatabaseSink, build_metadata
|
||||
from .storage import (
|
||||
DEFAULT_EMBEDDING_DIM,
|
||||
DEFAULT_TABLE_PREFIX,
|
||||
EMBEDDING_SOURCES,
|
||||
DatabaseSink,
|
||||
EmbeddingVector,
|
||||
build_metadata,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CallbackSink",
|
||||
"DatabaseSink",
|
||||
"DEFAULT_EMBEDDING_DIM",
|
||||
"DEFAULT_EVENT_TEXT_KEYS",
|
||||
"DEFAULT_RESULT_CLASSIFIER",
|
||||
"DEFAULT_SECRET_KEYS",
|
||||
"DEFAULT_SECRET_PATTERNS",
|
||||
"DEFAULT_TABLE_PREFIX",
|
||||
"EMBEDDING_SOURCES",
|
||||
"Embedder",
|
||||
"EmbeddingSink",
|
||||
"EmbeddingText",
|
||||
"EmbeddingVector",
|
||||
"EventRecord",
|
||||
"InstrumentationMiddleware",
|
||||
"JsonLinesSink",
|
||||
"MemorySink",
|
||||
"MODE_ENV_VAR",
|
||||
"MODES",
|
||||
"OpenAIEmbedder",
|
||||
"RecordDispatcher",
|
||||
"REDACTED",
|
||||
"Redactor",
|
||||
"ResultClassifier",
|
||||
"Sink",
|
||||
"ToolCallRecord",
|
||||
"backfill_embeddings",
|
||||
"backfill_feedback",
|
||||
"build_metadata",
|
||||
"extract_texts",
|
||||
"instrument",
|
||||
"soft_error_classifier",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,338 @@
|
||||
"""Embed records that were stored before an EmbeddingSink was running.
|
||||
|
||||
``backfill_embeddings`` reads tool-call errors and events back out of the
|
||||
DatabaseSink's tables, rebuilds the records, and sends them through the same
|
||||
extraction, redaction and truncation as ``EmbeddingSink.write``. The texts,
|
||||
hashes and source ids therefore match what the live sink would produce, so
|
||||
records arriving later dedupe against backfilled rows, and a second run
|
||||
embeds nothing new.
|
||||
|
||||
Command line: ``python -m fastmcp_feedback.instrumentation.backfill --help``.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from collections.abc import AsyncIterator, Iterable
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import and_, inspect, or_, select, tuple_
|
||||
|
||||
from ..embeddings import EmbeddingSink, extract_texts
|
||||
from ..events import EventRecord
|
||||
from ..record import ToolCallRecord
|
||||
from ..sinks import Record
|
||||
from ..storage import EMBEDDING_SOURCES
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
FEEDBACK_KIND = "feedback.submitted"
|
||||
LLM_KIND = "llm.call"
|
||||
|
||||
_CALL_COLUMNS = (
|
||||
"id",
|
||||
"tool",
|
||||
"started_at",
|
||||
"duration_ms",
|
||||
"mode",
|
||||
"ok",
|
||||
"error_type",
|
||||
"error_message",
|
||||
)
|
||||
|
||||
|
||||
def _utc(value: datetime | None) -> datetime | None:
|
||||
if value is None or value.tzinfo is None:
|
||||
return value if value is None else value.replace(tzinfo=UTC)
|
||||
return value.astimezone(UTC)
|
||||
|
||||
|
||||
class _Run:
|
||||
"""Counts and limits for one backfill call."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sink: EmbeddingSink,
|
||||
sources: Iterable[str],
|
||||
limit: int | None,
|
||||
dry_run: bool,
|
||||
):
|
||||
self.sink = sink
|
||||
self.sources = tuple(sources)
|
||||
self.limit = limit
|
||||
self.dry_run = dry_run
|
||||
self.counts = dict.fromkeys(self.sources, 0)
|
||||
self.skipped = 0
|
||||
self.errors = 0
|
||||
self.sent = 0
|
||||
self.have_table = True
|
||||
|
||||
@property
|
||||
def done(self) -> bool:
|
||||
return self.limit is not None and self.sent >= self.limit
|
||||
|
||||
def result(self) -> dict[str, int]:
|
||||
return {**self.counts, "skipped": self.skipped, "errors": self.errors}
|
||||
|
||||
async def feed(self, records: list[Record]) -> None:
|
||||
sink = self.sink
|
||||
items = [
|
||||
sink._prepare(i)
|
||||
for i in extract_texts(
|
||||
records, sources=self.sources, event_text_keys=sink.event_text_keys
|
||||
)
|
||||
]
|
||||
# The last text for a source wins, as in the live sink.
|
||||
latest = list({(i.source_type, i.source_id): i for i in items}.values())
|
||||
if self.have_table:
|
||||
todo, stored = await sink._unembedded(latest)
|
||||
else:
|
||||
todo, stored = latest, set()
|
||||
self.skipped += len(latest) - len(todo)
|
||||
if self.limit is not None:
|
||||
todo = todo[: max(0, self.limit - self.sent)]
|
||||
self.sent += len(todo)
|
||||
if self.dry_run:
|
||||
for item in todo:
|
||||
self.counts[item.source_type] += 1
|
||||
return
|
||||
for start in range(0, len(todo), sink.batch_size):
|
||||
chunk = todo[start : start + sink.batch_size]
|
||||
try:
|
||||
vectors = await asyncio.wait_for(
|
||||
sink.embedder.embed([i.text for i in chunk]), sink.timeout
|
||||
)
|
||||
await sink._store(chunk, vectors, stored)
|
||||
except Exception:
|
||||
self.errors += len(chunk)
|
||||
logger.warning(
|
||||
"Backfill failed to embed %d texts with %s",
|
||||
len(chunk),
|
||||
sink.embedder.model,
|
||||
exc_info=True,
|
||||
)
|
||||
continue
|
||||
for item in chunk:
|
||||
self.counts[item.source_type] += 1
|
||||
|
||||
|
||||
async def _has_table(engine: Any, name: str) -> bool:
|
||||
async with engine.connect() as conn:
|
||||
return await conn.run_sync(lambda sync: inspect(sync).has_table(name))
|
||||
|
||||
|
||||
async def _columns(engine: Any, name: str) -> set[str]:
|
||||
async with engine.connect() as conn:
|
||||
return await conn.run_sync(
|
||||
lambda sync: {c["name"] for c in inspect(sync).get_columns(name)}
|
||||
)
|
||||
|
||||
|
||||
async def _pages(
|
||||
engine: Any,
|
||||
table: Any,
|
||||
columns: list[Any],
|
||||
ts: Any,
|
||||
where: Any,
|
||||
since: datetime | None,
|
||||
until: datetime | None,
|
||||
page_size: int,
|
||||
) -> AsyncIterator[list[dict[str, Any]]]:
|
||||
"""Rows matching ``where``, oldest first, ``page_size`` at a time.
|
||||
|
||||
Keyset paging on (timestamp, id), so rows written during the run neither
|
||||
shift the pages nor get read twice.
|
||||
"""
|
||||
conditions = [where]
|
||||
if since is not None:
|
||||
conditions.append(ts >= since)
|
||||
if until is not None:
|
||||
conditions.append(ts < until)
|
||||
last: tuple[Any, Any] | None = None
|
||||
while True:
|
||||
query = select(*columns).where(*conditions)
|
||||
if last is not None:
|
||||
query = query.where(tuple_(ts, table.c.id) > tuple_(*last))
|
||||
query = query.order_by(ts, table.c.id).limit(page_size)
|
||||
async with engine.connect() as conn:
|
||||
rows = [dict(r) for r in (await conn.execute(query)).mappings()]
|
||||
if not rows:
|
||||
return
|
||||
yield rows
|
||||
if len(rows) < page_size:
|
||||
return
|
||||
last = (rows[-1][ts.name], rows[-1]["id"])
|
||||
|
||||
|
||||
async def _backfill_calls(run: _Run, since, until, page_size) -> None:
|
||||
db = run.sink.database_sink
|
||||
t = db.table
|
||||
have = await _columns(db.engine, t.name)
|
||||
if "outcome" in have:
|
||||
# Rows from before the outcome column was added have it NULL and
|
||||
# were errors exactly when they carry an error message.
|
||||
where = or_(
|
||||
t.c.outcome.in_(("error", "soft_error")),
|
||||
and_(t.c.outcome.is_(None), t.c.error_message.is_not(None)),
|
||||
)
|
||||
columns = [t.c[c] for c in (*_CALL_COLUMNS, "outcome")]
|
||||
else:
|
||||
where = t.c.error_message.is_not(None)
|
||||
columns = [t.c[c] for c in _CALL_COLUMNS]
|
||||
async for rows in _pages(
|
||||
db.engine, t, columns, t.c.started_at, where, since, until, page_size
|
||||
):
|
||||
records: list[Record] = []
|
||||
for row in rows:
|
||||
row["outcome"] = row.get("outcome") or "error"
|
||||
records.append(ToolCallRecord(**row))
|
||||
await run.feed(records)
|
||||
if run.done:
|
||||
return
|
||||
|
||||
|
||||
async def _backfill_events(run: _Run, since, until, page_size) -> None:
|
||||
db = run.sink.database_sink
|
||||
e = db.events
|
||||
kinds = []
|
||||
if "feedback" in run.sources:
|
||||
kinds.append(e.c.kind == FEEDBACK_KIND)
|
||||
if "llm" in run.sources:
|
||||
kinds.append(e.c.kind == LLM_KIND)
|
||||
if "event" in run.sources:
|
||||
# Includes llm.call: one without prompt or completion is a plain event.
|
||||
kinds.append(e.c.kind != FEEDBACK_KIND)
|
||||
columns = [e.c.id, e.c.kind, e.c.key, e.c.occurred_at, e.c.attrs]
|
||||
async for rows in _pages(
|
||||
db.engine, e, columns, e.c.occurred_at, or_(*kinds), since, until, page_size
|
||||
):
|
||||
records: list[Record] = [
|
||||
EventRecord(**{**row, "attrs": row["attrs"] or {}}) for row in rows
|
||||
]
|
||||
await run.feed(records)
|
||||
if run.done:
|
||||
return
|
||||
|
||||
|
||||
async def _start(
|
||||
sink: EmbeddingSink,
|
||||
sources: Iterable[str] | None,
|
||||
limit: int | None,
|
||||
batch_size: int,
|
||||
dry_run: bool,
|
||||
) -> _Run:
|
||||
chosen = tuple(sink.sources if sources is None else sources)
|
||||
unknown = set(chosen) - set(EMBEDDING_SOURCES)
|
||||
if unknown:
|
||||
raise ValueError(
|
||||
f"unknown sources {sorted(unknown)}; choose from {EMBEDDING_SOURCES}"
|
||||
)
|
||||
if batch_size < 1:
|
||||
raise ValueError(f"batch_size must be at least 1, got {batch_size!r}")
|
||||
if limit is not None and limit < 0:
|
||||
raise ValueError(f"limit must not be negative, got {limit!r}")
|
||||
run = _Run(sink, sorted(chosen, key=EMBEDDING_SOURCES.index), limit, dry_run)
|
||||
db = sink.database_sink
|
||||
if not dry_run:
|
||||
await db.ensure_embeddings()
|
||||
run.have_table = await _has_table(db.engine, db.embeddings.name)
|
||||
if not run.have_table and not dry_run:
|
||||
raise RuntimeError(
|
||||
f"{db.embeddings.name} does not exist; create it with a migration "
|
||||
"or pass create_tables=True to the DatabaseSink"
|
||||
)
|
||||
return run
|
||||
|
||||
|
||||
async def backfill_embeddings(
|
||||
embedding_sink: EmbeddingSink,
|
||||
*,
|
||||
sources: Iterable[str] | None = None,
|
||||
since: datetime | None = None,
|
||||
until: datetime | None = None,
|
||||
limit: int | None = None,
|
||||
batch_size: int = 500,
|
||||
dry_run: bool = False,
|
||||
) -> dict[str, int]:
|
||||
"""Embed stored tool-call errors and events that have no embedding yet.
|
||||
|
||||
Reads the tables of ``embedding_sink.database_sink``:
|
||||
|
||||
- ``call_error``: calls with outcome ``error`` or ``soft_error``, and
|
||||
calls with outcome NULL and an error message (rows written before the
|
||||
outcome column existed).
|
||||
- ``feedback``: ``feedback.submitted`` events.
|
||||
- ``llm``: ``llm.call`` events with a prompt or completion.
|
||||
- ``event``: other events with any of the sink's ``event_text_keys``.
|
||||
|
||||
Each text goes through the sink's extraction, redactor and ``max_chars``,
|
||||
so it matches what the live sink makes of the same record. Texts already
|
||||
stored with the same hash are skipped; a changed text replaces its row.
|
||||
|
||||
Args:
|
||||
sources: which of the above; defaults to the sink's ``sources``.
|
||||
since, until: only rows timestamped in [since, until). Naive
|
||||
datetimes are taken as UTC.
|
||||
limit: embed at most this many texts, oldest first. A later run
|
||||
carries on, since what was embedded is then skipped.
|
||||
batch_size: rows read per query. Texts are embedded in chunks of
|
||||
the sink's ``batch_size``, each limited to its ``timeout``.
|
||||
dry_run: count what would be embedded without calling the embedder
|
||||
or writing anything.
|
||||
|
||||
Returns counts of texts embedded per source type, plus ``skipped``
|
||||
(already embedded) and ``errors`` (texts whose embed or store failed;
|
||||
logged, never raised). Raises ValueError for bad arguments and
|
||||
RuntimeError when a table it needs does not exist.
|
||||
"""
|
||||
run = await _start(embedding_sink, sources, limit, batch_size, dry_run)
|
||||
db = embedding_sink.database_sink
|
||||
since, until = _utc(since), _utc(until)
|
||||
if "call_error" in run.sources:
|
||||
if not await _has_table(db.engine, db.table.name):
|
||||
raise RuntimeError(f"{db.table.name} does not exist")
|
||||
if not run.done:
|
||||
await _backfill_calls(run, since, until, batch_size)
|
||||
if set(run.sources) & {"feedback", "event", "llm"}:
|
||||
if not await _has_table(db.engine, db.events.name):
|
||||
raise RuntimeError(f"{db.events.name} does not exist")
|
||||
if not run.done:
|
||||
await _backfill_events(run, since, until, batch_size)
|
||||
return run.result()
|
||||
|
||||
|
||||
async def backfill_feedback(
|
||||
embedding_sink: EmbeddingSink,
|
||||
items: Iterable[tuple[str, str, str | None]],
|
||||
*,
|
||||
batch_size: int = 500,
|
||||
dry_run: bool = False,
|
||||
) -> dict[str, int]:
|
||||
"""Embed feedback kept in the host's own tables, as (ref, title, description).
|
||||
|
||||
The text and source id are those of a ``feedback.submitted`` event with
|
||||
key ``ref`` and those attrs, so feedback recorded later as an event
|
||||
dedupes against these rows. Returns ``{"feedback", "skipped", "errors"}``
|
||||
like ``backfill_embeddings``.
|
||||
"""
|
||||
run = await _start(embedding_sink, ("feedback",), None, batch_size, dry_run)
|
||||
now = datetime.now(UTC)
|
||||
page: list[Record] = []
|
||||
for ref, title, description in items:
|
||||
page.append(
|
||||
EventRecord(
|
||||
kind=FEEDBACK_KIND,
|
||||
occurred_at=now,
|
||||
key=str(ref),
|
||||
attrs={"title": title, "description": description},
|
||||
)
|
||||
)
|
||||
if len(page) >= batch_size:
|
||||
await run.feed(page)
|
||||
page = []
|
||||
if page:
|
||||
await run.feed(page)
|
||||
return run.result()
|
||||
|
||||
|
||||
__all__ = ["backfill_embeddings", "backfill_feedback"]
|
||||
@@ -0,0 +1,5 @@
|
||||
import sys
|
||||
|
||||
from .cli import main
|
||||
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,143 @@
|
||||
"""python -m fastmcp_feedback.instrumentation.backfill
|
||||
|
||||
Prints the counts as JSON. Exits 1 when any text failed to embed, 2 on a
|
||||
setup error (bad arguments, missing table, unreachable database).
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime
|
||||
|
||||
from ..embeddings import DEFAULT_EVENT_TEXT_KEYS, EmbeddingSink, OpenAIEmbedder
|
||||
from ..storage import DEFAULT_TABLE_PREFIX, EMBEDDING_SOURCES, DatabaseSink
|
||||
from . import backfill_embeddings
|
||||
|
||||
API_KEY_ENV = "FFB_EMBED_API_KEY"
|
||||
|
||||
|
||||
def _sources(value: str) -> tuple[str, ...]:
|
||||
chosen = tuple(s.strip() for s in value.split(",") if s.strip())
|
||||
unknown = set(chosen) - set(EMBEDDING_SOURCES)
|
||||
if not chosen or unknown:
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"choose from {','.join(EMBEDDING_SOURCES)}, comma separated"
|
||||
)
|
||||
return chosen
|
||||
|
||||
|
||||
def _when(value: str) -> datetime:
|
||||
try:
|
||||
return datetime.fromisoformat(value)
|
||||
except ValueError as exc:
|
||||
raise argparse.ArgumentTypeError(f"not an ISO 8601 time: {value!r}") from exc
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="python -m fastmcp_feedback.instrumentation.backfill",
|
||||
description=(
|
||||
"Embed stored tool-call errors and events that have no embedding "
|
||||
f"yet. The embedder's API key is read from ${API_KEY_ENV} (or the "
|
||||
"variable named by --api-key-env)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--database-url",
|
||||
required=True,
|
||||
help="async SQLAlchemy URL, e.g. postgresql+asyncpg://... or sqlite+aiosqlite:///calls.db",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--embed-base-url",
|
||||
help="OpenAI-compatible base URL up to the version, e.g. https://api.example.com/v1 "
|
||||
"(required unless --dry-run)",
|
||||
)
|
||||
parser.add_argument("--api-key-env", default=API_KEY_ENV, metavar="NAME")
|
||||
parser.add_argument("--model", default="mxbai-embed-large")
|
||||
parser.add_argument("--dim", type=int, default=1024)
|
||||
parser.add_argument("--prefix", default=DEFAULT_TABLE_PREFIX)
|
||||
parser.add_argument(
|
||||
"--sources",
|
||||
type=_sources,
|
||||
default=EMBEDDING_SOURCES,
|
||||
help=f"comma separated, default {','.join(EMBEDDING_SOURCES)}",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--event-text-keys",
|
||||
default=",".join(DEFAULT_EVENT_TEXT_KEYS),
|
||||
help="attrs making up an event text, comma separated",
|
||||
)
|
||||
parser.add_argument("--since", type=_when, help="ISO 8601; naive means UTC")
|
||||
parser.add_argument("--until", type=_when, help="ISO 8601, exclusive")
|
||||
parser.add_argument("--limit", type=int, help="embed at most this many texts")
|
||||
parser.add_argument("--batch-size", type=int, default=500, help="rows per query")
|
||||
parser.add_argument("--max-chars", type=int, default=4000)
|
||||
parser.add_argument(
|
||||
"--create-tables",
|
||||
action="store_true",
|
||||
help="create the embeddings table (and vector extension) if missing",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
help="count what would be embedded; call no embedder, write nothing",
|
||||
)
|
||||
parser.add_argument("-v", "--verbose", action="store_true")
|
||||
return parser
|
||||
|
||||
|
||||
async def _run(args: argparse.Namespace) -> dict[str, int]:
|
||||
embedder = OpenAIEmbedder(
|
||||
args.embed_base_url or "http://embedder.invalid/v1",
|
||||
api_key=os.environ.get(args.api_key_env) or None,
|
||||
model=args.model,
|
||||
dim=args.dim,
|
||||
)
|
||||
db = DatabaseSink(
|
||||
args.database_url,
|
||||
prefix=args.prefix,
|
||||
create_tables=args.create_tables and not args.dry_run,
|
||||
embedding_dim=args.dim,
|
||||
)
|
||||
try:
|
||||
sink = EmbeddingSink(
|
||||
embedder,
|
||||
db,
|
||||
sources=args.sources,
|
||||
event_text_keys=[k for k in args.event_text_keys.split(",") if k],
|
||||
max_chars=args.max_chars,
|
||||
)
|
||||
return await backfill_embeddings(
|
||||
sink,
|
||||
since=args.since,
|
||||
until=args.until,
|
||||
limit=args.limit,
|
||||
batch_size=args.batch_size,
|
||||
dry_run=args.dry_run,
|
||||
)
|
||||
finally:
|
||||
await db.aclose()
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
if not args.dry_run and not args.embed_base_url:
|
||||
parser.error("--embed-base-url is required unless --dry-run")
|
||||
# force: importing fastmcp_feedback already configured the root logger
|
||||
# at INFO, which would print every embedding request.
|
||||
logging.basicConfig(
|
||||
level=logging.INFO if args.verbose else logging.WARNING,
|
||||
format="%(levelname)s %(name)s: %(message)s",
|
||||
force=True,
|
||||
)
|
||||
try:
|
||||
counts = asyncio.run(_run(args))
|
||||
except Exception as exc:
|
||||
print(f"backfill failed: {type(exc).__name__}: {exc}", file=sys.stderr)
|
||||
return 2
|
||||
print(json.dumps({"dry_run": args.dry_run, **counts}))
|
||||
return 1 if counts["errors"] else 0
|
||||
@@ -58,6 +58,7 @@ def soft_error_classifier(
|
||||
status_keys: Iterable[str] = ("status", "state"),
|
||||
flag_keys: Iterable[str] = ("ok", "success"),
|
||||
extra_statuses: Iterable[str] = (),
|
||||
error_keys: Iterable[str] = ("error",),
|
||||
) -> ResultClassifier:
|
||||
"""Build a classifier for dict results that describe their own failure.
|
||||
|
||||
@@ -67,13 +68,31 @@ def soft_error_classifier(
|
||||
``isError``. The message is the first of ``error``, ``message``,
|
||||
``detail`` and ``reason`` present, e.g. ``status=failed: disk full``.
|
||||
|
||||
Failing those, it is a soft error when any of ``error_keys`` holds a
|
||||
non-empty string, as in ``{"played": False, "error": "output stalled"}``,
|
||||
with that string as the message (``error: output stalled``). None, empty
|
||||
or blank strings, numbers, booleans, dicts and lists there do not count,
|
||||
so ``{"error": None}`` and ``{"error": 0.02}`` stay ok. ``error_keys=()``
|
||||
turns this off.
|
||||
|
||||
A result that says it succeeded, with any of ``flag_keys`` exactly True,
|
||||
skips only that last check: ``{"ok": True, "error": "deprecated arg"}``
|
||||
is ok, the string being a warning. A failing status, a False flag and
|
||||
``isError`` still flag it. Truthy non-bools such as ``1`` or ``"yes"``
|
||||
are not an explicit success.
|
||||
|
||||
>>> classify = soft_error_classifier(extra_statuses=("unknown_target",))
|
||||
>>> classify("render", {"status": "unknown_target", "error": "no such object"})
|
||||
'status=unknown_target: no such object'
|
||||
>>> classify("speak", {"played": False, "error": "output stalled"})
|
||||
'error: output stalled'
|
||||
>>> print(classify("speak", {"ok": True, "error": "deprecated arg"}))
|
||||
None
|
||||
"""
|
||||
bad = frozenset(s.lower() for s in (*statuses, *extra_statuses))
|
||||
status_keys = tuple(status_keys)
|
||||
flag_keys = tuple(flag_keys)
|
||||
error_keys = tuple(error_keys)
|
||||
|
||||
def classify(tool: str, result: Any) -> str | None:
|
||||
if not isinstance(result, dict):
|
||||
@@ -89,6 +108,12 @@ def soft_error_classifier(
|
||||
if result.get(key) is True:
|
||||
detail = _message(result) or _content_text(result.get("content"))
|
||||
return _describe("isError", detail)
|
||||
if any(result.get(key) is True for key in flag_keys):
|
||||
return None
|
||||
for key in error_keys:
|
||||
value = result.get(key)
|
||||
if isinstance(value, str) and value.strip():
|
||||
return _describe(key, value.strip())
|
||||
return None
|
||||
|
||||
return classify
|
||||
|
||||
@@ -0,0 +1,670 @@
|
||||
"""Embed feedback, tool errors, event text and LLM text for similarity search.
|
||||
|
||||
EmbeddingSink sits next to a DatabaseSink. It picks text out of each batch,
|
||||
queues it on its own bounded queue and returns at once, so a slow embedding
|
||||
backend never holds up the dispatcher. A background task embeds and stores
|
||||
the vectors in ``{prefix}embeddings``.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import hashlib
|
||||
import logging
|
||||
import math
|
||||
import operator
|
||||
import uuid
|
||||
from collections.abc import Iterable, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
from sqlalchemy import Float, and_, literal, not_, select, tuple_
|
||||
|
||||
from .events import EventRecord
|
||||
from .record import ToolCallRecord
|
||||
from .redaction import Redactor
|
||||
from .sinks import Record
|
||||
from .storage import EMBEDDING_SOURCES, EmbeddingVector
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_EVENT_TEXT_KEYS = ("error_excerpt", "error", "message", "reason")
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class Embedder(Protocol):
|
||||
"""Turns texts into vectors of length ``dim``, one per text, in order."""
|
||||
|
||||
model: str
|
||||
dim: int
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
|
||||
def _http_module() -> Any:
|
||||
# FastMCP 2.x and 3.x depend on httpx, FastMCP 4 on its successor httpx2.
|
||||
try:
|
||||
import httpx
|
||||
|
||||
return httpx
|
||||
except ImportError:
|
||||
pass
|
||||
try:
|
||||
import httpx2
|
||||
|
||||
return httpx2
|
||||
except ImportError as exc:
|
||||
raise ImportError("OpenAIEmbedder needs httpx or httpx2") from exc
|
||||
|
||||
|
||||
class OpenAIEmbedder:
|
||||
"""An OpenAI-compatible ``POST {base_url}/embeddings`` endpoint.
|
||||
|
||||
Args:
|
||||
base_url: up to and including the version, e.g.
|
||||
``https://api.example.com/v1``.
|
||||
api_key: sent as ``Authorization: Bearer ...`` when given.
|
||||
model: the ``model`` field of each request.
|
||||
dim: expected vector length. A response with any other length raises
|
||||
ValueError, since it cannot be stored or compared.
|
||||
batch_size: texts per request.
|
||||
timeout: seconds per request.
|
||||
transport: passed to the HTTP client. For tests.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
api_key: str | None = None,
|
||||
model: str = "mxbai-embed-large",
|
||||
dim: int = 1024,
|
||||
batch_size: int = 32,
|
||||
timeout: float = 30.0,
|
||||
transport: Any = None,
|
||||
):
|
||||
if batch_size < 1:
|
||||
raise ValueError(f"batch_size must be at least 1, got {batch_size!r}")
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.dim = dim
|
||||
self.batch_size = batch_size
|
||||
self.timeout = timeout
|
||||
self.transport = transport
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||||
if not texts:
|
||||
return []
|
||||
http = _http_module()
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if self.api_key:
|
||||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||||
vectors: list[list[float]] = []
|
||||
kwargs: dict[str, Any] = {"timeout": self.timeout}
|
||||
if self.transport is not None:
|
||||
kwargs["transport"] = self.transport
|
||||
async with http.AsyncClient(**kwargs) as client:
|
||||
for start in range(0, len(texts), self.batch_size):
|
||||
chunk = texts[start : start + self.batch_size]
|
||||
response = await client.post(
|
||||
f"{self.base_url}/embeddings",
|
||||
json={"model": self.model, "input": chunk},
|
||||
headers=headers,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = sorted(
|
||||
response.json()["data"], key=lambda item: item.get("index", 0)
|
||||
)
|
||||
if len(data) != len(chunk):
|
||||
raise ValueError(
|
||||
f"asked for {len(chunk)} embeddings, got {len(data)}"
|
||||
)
|
||||
for item in data:
|
||||
vector = [float(v) for v in item["embedding"]]
|
||||
if len(vector) != self.dim:
|
||||
raise ValueError(
|
||||
f"{self.model} returned a {len(vector)}-dim vector, "
|
||||
f"expected {self.dim}"
|
||||
)
|
||||
vectors.append(vector)
|
||||
return vectors
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbeddingText:
|
||||
"""One text waiting to be embedded, and where it came from."""
|
||||
|
||||
source_type: str
|
||||
source_id: str
|
||||
text: str
|
||||
|
||||
@property
|
||||
def text_hash(self) -> str:
|
||||
return hashlib.sha256(self.text.encode()).hexdigest()
|
||||
|
||||
|
||||
# Flattening stops adding items past this many characters; max_chars cuts
|
||||
# the result again later. Keeps a huge list from being joined in full.
|
||||
_FLAT_CHARS = 65536
|
||||
_FLAT_DEPTH = 32
|
||||
|
||||
|
||||
def _clean(value: Any, depth: int = 0) -> str | None:
|
||||
"""``value`` as stripped text, or None when there is none.
|
||||
|
||||
Lists and tuples become their items' texts, one per line. Dicts become
|
||||
``key: value`` lines sorted by key: PostgreSQL JSONB does not keep key
|
||||
order, and the text has to come out the same when a backfill rebuilds it
|
||||
from a stored row. Other values go through ``str()``.
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, str):
|
||||
text = value
|
||||
elif isinstance(value, list | tuple | dict):
|
||||
if depth >= _FLAT_DEPTH:
|
||||
return None
|
||||
if isinstance(value, dict):
|
||||
pairs = sorted(value.items(), key=lambda kv: str(kv[0]))
|
||||
parts = ((f"{k}: ", _clean(v, depth + 1)) for k, v in pairs)
|
||||
else:
|
||||
parts = (("", _clean(v, depth + 1)) for v in value)
|
||||
lines: list[str] = []
|
||||
size = 0
|
||||
for label, part in parts:
|
||||
if not part:
|
||||
continue
|
||||
lines.append(label + part)
|
||||
size += len(lines[-1]) + 1
|
||||
if size > _FLAT_CHARS:
|
||||
break
|
||||
text = "\n".join(lines)
|
||||
else:
|
||||
text = str(value)
|
||||
text = text.strip()
|
||||
return text or None
|
||||
|
||||
|
||||
def extract_texts(
|
||||
records: Iterable[Record],
|
||||
*,
|
||||
sources: Iterable[str] = EMBEDDING_SOURCES,
|
||||
event_text_keys: Sequence[str] = DEFAULT_EVENT_TEXT_KEYS,
|
||||
) -> list[EmbeddingText]:
|
||||
"""The texts worth embedding in a batch of records, unredacted and untruncated.
|
||||
|
||||
``call_error``: calls with outcome ``error`` or ``soft_error``, as
|
||||
``"{tool}: {error_type}: {error_message}"``, keyed by the call id.
|
||||
``feedback``: ``feedback.submitted`` events, title and description,
|
||||
keyed by the event's ``key`` (the feedback ref).
|
||||
``llm``: ``llm.call`` events with ``prompt`` or ``completion`` attrs.
|
||||
``event``: any other event with one of ``event_text_keys`` in its attrs,
|
||||
those values joined by newlines, keyed by the event id. An ``llm.call``
|
||||
without prompt or completion counts as a plain event, so a failed LLM
|
||||
call's ``error`` is still embedded.
|
||||
|
||||
List and tuple values become one line per item, dicts ``key: value``
|
||||
lines sorted by key, at any depth.
|
||||
"""
|
||||
wanted = frozenset(sources)
|
||||
out: list[EmbeddingText] = []
|
||||
for record in records:
|
||||
if isinstance(record, ToolCallRecord):
|
||||
if "call_error" in wanted and record.outcome in ("error", "soft_error"):
|
||||
parts = [record.tool, record.error_type, _clean(record.error_message)]
|
||||
out.append(
|
||||
EmbeddingText(
|
||||
"call_error", record.id, ": ".join(p for p in parts if p)
|
||||
)
|
||||
)
|
||||
continue
|
||||
if not isinstance(record, EventRecord):
|
||||
continue
|
||||
attrs = record.attrs or {}
|
||||
if record.kind == "feedback.submitted":
|
||||
if "feedback" in wanted:
|
||||
body = "\n".join(
|
||||
t
|
||||
for t in (
|
||||
_clean(attrs.get("title")),
|
||||
_clean(attrs.get("description")),
|
||||
)
|
||||
if t
|
||||
)
|
||||
if body:
|
||||
out.append(EmbeddingText("feedback", record.key or record.id, body))
|
||||
continue
|
||||
if record.kind == "llm.call":
|
||||
prompt, completion = (
|
||||
_clean(attrs.get("prompt")),
|
||||
_clean(attrs.get("completion")),
|
||||
)
|
||||
if prompt or completion:
|
||||
if "llm" in wanted:
|
||||
lines = []
|
||||
if prompt:
|
||||
lines.append(f"prompt: {prompt}")
|
||||
if completion:
|
||||
lines.append(f"completion: {completion}")
|
||||
out.append(EmbeddingText("llm", record.id, "\n".join(lines)))
|
||||
continue
|
||||
if "event" in wanted:
|
||||
values = [_clean(attrs.get(k)) for k in event_text_keys]
|
||||
body = "\n".join(v for v in values if v)
|
||||
if body:
|
||||
out.append(EmbeddingText("event", record.id, body))
|
||||
return out
|
||||
|
||||
|
||||
def _dot(a: Sequence[float], b: Sequence[float]) -> float:
|
||||
sumprod = getattr(math, "sumprod", None) # 3.12+
|
||||
if sumprod is not None:
|
||||
return sumprod(a, b)
|
||||
return sum(map(operator.mul, a, b))
|
||||
|
||||
|
||||
def cosine_score(a: Sequence[float], b: Sequence[float]) -> float | None:
|
||||
"""1 - cosine distance, clamped to [0, 1]. None for a zero vector."""
|
||||
na, nb = math.sqrt(_dot(a, a)), math.sqrt(_dot(b, b))
|
||||
if not na or not nb:
|
||||
return None
|
||||
return _clamp(_dot(a, b) / (na * nb))
|
||||
|
||||
|
||||
def _clamp(score: float) -> float:
|
||||
return min(1.0, max(0.0, score))
|
||||
|
||||
|
||||
class EmbeddingSink:
|
||||
"""Embed text from records and store the vectors through a DatabaseSink.
|
||||
|
||||
``write`` only queues, so it returns at once whatever the embedder is
|
||||
doing. A background task embeds up to ``batch_size`` texts at a time and
|
||||
upserts them into ``{prefix}embeddings`` on ``database_sink.engine``. A
|
||||
full queue drops the text and counts it in ``dropped``; embedder and
|
||||
database failures are logged and counted in ``errors``, never raised.
|
||||
A text whose (source_type, source_id, model) row already holds the same
|
||||
text is not embedded again; a changed text replaces the row.
|
||||
|
||||
Text is redacted with ``redactor`` (the middleware's, when this sink is
|
||||
passed to ``instrument``) and cut to ``max_chars`` before it is embedded
|
||||
or stored.
|
||||
|
||||
Args:
|
||||
embedder: anything with ``model``, ``dim`` and ``async embed(texts)``.
|
||||
database_sink: the DatabaseSink whose engine and tables to use. Its
|
||||
``embedding_dim`` must equal ``embedder.dim``.
|
||||
sources: which of ``feedback``, ``call_error``, ``event`` and ``llm``
|
||||
to embed.
|
||||
event_text_keys: attrs whose values make up an ``event`` text.
|
||||
max_chars: longest text embedded.
|
||||
max_queue: texts waiting to be embedded before new ones are dropped.
|
||||
batch_size: texts per embed call.
|
||||
timeout: seconds one embed call may take, and how long ``aclose``
|
||||
waits for the queue to drain.
|
||||
scan_limit: without pgvector, ``similar`` compares against the newest
|
||||
this many rows.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
embedder: Embedder,
|
||||
database_sink: Any,
|
||||
*,
|
||||
sources: Iterable[str] = EMBEDDING_SOURCES,
|
||||
event_text_keys: Sequence[str] = DEFAULT_EVENT_TEXT_KEYS,
|
||||
max_chars: int = 4000,
|
||||
max_queue: int = 1000,
|
||||
batch_size: int = 32,
|
||||
timeout: float = 30.0,
|
||||
scan_limit: int = 5000,
|
||||
redactor: Redactor | None = None,
|
||||
):
|
||||
self.sources = frozenset(sources)
|
||||
unknown = self.sources - set(EMBEDDING_SOURCES)
|
||||
if unknown:
|
||||
raise ValueError(
|
||||
f"unknown sources {sorted(unknown)}; choose from {EMBEDDING_SOURCES}"
|
||||
)
|
||||
if embedder.dim != database_sink.embedding_dim:
|
||||
raise ValueError(
|
||||
f"embedder {embedder.model!r} makes {embedder.dim}-dim vectors but "
|
||||
f"the DatabaseSink's embeddings table holds "
|
||||
f"{database_sink.embedding_dim}; pass embedding_dim={embedder.dim} "
|
||||
"to DatabaseSink"
|
||||
)
|
||||
if max_chars < 1 or max_queue < 1 or batch_size < 1 or scan_limit < 1:
|
||||
raise ValueError(
|
||||
"max_chars, max_queue, batch_size and scan_limit must be positive"
|
||||
)
|
||||
self.embedder = embedder
|
||||
self.database_sink = database_sink
|
||||
self.table = database_sink.embeddings
|
||||
self.event_text_keys = tuple(event_text_keys)
|
||||
self.max_chars = max_chars
|
||||
self.max_queue = max_queue
|
||||
self.batch_size = batch_size
|
||||
self.timeout = timeout
|
||||
self.scan_limit = scan_limit
|
||||
self.redactor = redactor
|
||||
self._default_redactor = Redactor()
|
||||
self.dropped = 0
|
||||
self.errors = 0
|
||||
self.embedded = 0
|
||||
self.skipped = 0
|
||||
self._queue: asyncio.Queue[EmbeddingText] | None = None
|
||||
self._task: asyncio.Task | None = None
|
||||
self._loop: asyncio.AbstractEventLoop | None = None
|
||||
self._warned_scan = False
|
||||
|
||||
# Writing
|
||||
|
||||
async def write(self, records: Sequence[Record]) -> None:
|
||||
try:
|
||||
items = extract_texts(
|
||||
records, sources=self.sources, event_text_keys=self.event_text_keys
|
||||
)
|
||||
except Exception:
|
||||
self.errors += 1
|
||||
logger.warning("Failed to extract texts to embed", exc_info=True)
|
||||
return
|
||||
for item in items:
|
||||
self._enqueue(self._prepare(item))
|
||||
|
||||
def _prepare(self, item: EmbeddingText) -> EmbeddingText:
|
||||
redactor = self.redactor or self._default_redactor
|
||||
item.text = redactor.redact_text(item.text)[: self.max_chars]
|
||||
return item
|
||||
|
||||
def _enqueue(self, item: EmbeddingText) -> None:
|
||||
loop = asyncio.get_running_loop()
|
||||
if self._queue is None or self._loop is not loop:
|
||||
self._loop = loop
|
||||
self._queue = asyncio.Queue(maxsize=self.max_queue)
|
||||
self._task = loop.create_task(self._run(self._queue))
|
||||
try:
|
||||
self._queue.put_nowait(item)
|
||||
except asyncio.QueueFull:
|
||||
self.dropped += 1
|
||||
if self.dropped == 1 or self.dropped % 1000 == 0:
|
||||
logger.warning(
|
||||
"Embedding queue full; %d texts dropped so far", self.dropped
|
||||
)
|
||||
|
||||
async def _run(self, queue: asyncio.Queue) -> None:
|
||||
while True:
|
||||
batch = [await queue.get()]
|
||||
while len(batch) < self.batch_size:
|
||||
try:
|
||||
batch.append(queue.get_nowait())
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
try:
|
||||
await self._process(batch)
|
||||
except Exception:
|
||||
self.errors += 1
|
||||
logger.warning(
|
||||
"Failed to embed %d texts with %s",
|
||||
len(batch),
|
||||
self.embedder.model,
|
||||
exc_info=True,
|
||||
)
|
||||
finally:
|
||||
for _ in batch:
|
||||
queue.task_done()
|
||||
|
||||
async def _process(self, batch: list[EmbeddingText]) -> None:
|
||||
# The last text for a source wins within a batch.
|
||||
latest = {(i.source_type, i.source_id): i for i in batch}
|
||||
db = self.database_sink
|
||||
await db._ready()
|
||||
await db.ensure_embeddings()
|
||||
todo, stored = await self._unembedded(list(latest.values()))
|
||||
self.skipped += len(batch) - len(todo)
|
||||
if not todo:
|
||||
return
|
||||
vectors = await asyncio.wait_for(
|
||||
self.embedder.embed([i.text for i in todo]), self.timeout
|
||||
)
|
||||
await self._store(todo, vectors, stored)
|
||||
self.embedded += len(todo)
|
||||
|
||||
async def _unembedded(
|
||||
self, items: list[EmbeddingText]
|
||||
) -> tuple[list[EmbeddingText], set[tuple[str, str]]]:
|
||||
"""The items whose stored row is missing or holds another text, and
|
||||
the keys that have a row (to be replaced). Items must be unique by
|
||||
(source_type, source_id)."""
|
||||
if not items:
|
||||
return [], set()
|
||||
t = self.table
|
||||
keys = [(i.source_type, i.source_id) for i in items]
|
||||
async with self.database_sink.engine.connect() as conn:
|
||||
rows = (
|
||||
await conn.execute(
|
||||
select(t.c.source_type, t.c.source_id, t.c.text_hash).where(
|
||||
t.c.model == self.embedder.model,
|
||||
tuple_(t.c.source_type, t.c.source_id).in_(keys),
|
||||
)
|
||||
)
|
||||
).all()
|
||||
stored = {(r.source_type, r.source_id): r.text_hash for r in rows}
|
||||
todo = [
|
||||
i for i in items if stored.get((i.source_type, i.source_id)) != i.text_hash
|
||||
]
|
||||
return todo, set(stored)
|
||||
|
||||
async def _store(
|
||||
self,
|
||||
todo: list[EmbeddingText],
|
||||
vectors: list[list[float]],
|
||||
replace: set[tuple[str, str]],
|
||||
) -> None:
|
||||
db = self.database_sink
|
||||
if len(vectors) != len(todo):
|
||||
raise ValueError(
|
||||
f"embedder returned {len(vectors)} vectors for {len(todo)}"
|
||||
)
|
||||
for vector in vectors:
|
||||
if len(vector) != db.embedding_dim:
|
||||
raise ValueError(
|
||||
f"embedder returned a {len(vector)}-dim vector, "
|
||||
f"expected {db.embedding_dim}"
|
||||
)
|
||||
now = datetime.now(UTC)
|
||||
values = [
|
||||
{
|
||||
"id": str(uuid.uuid4()),
|
||||
"source_type": i.source_type,
|
||||
"source_id": i.source_id,
|
||||
"model": self.embedder.model,
|
||||
"dim": len(v),
|
||||
"text_hash": i.text_hash,
|
||||
"text": i.text,
|
||||
"created_at": now,
|
||||
"embedding": [float(x) for x in v],
|
||||
}
|
||||
for i, v in zip(todo, vectors, strict=True)
|
||||
]
|
||||
await self._upsert(values, replace=replace)
|
||||
|
||||
async def _upsert(self, values: list[dict[str, Any]], replace: set) -> None:
|
||||
t = self.table
|
||||
engine = self.database_sink.engine
|
||||
dialect = engine.dialect.name
|
||||
changing = ("text_hash", "text", "created_at", "embedding", "dim")
|
||||
if dialect in ("postgresql", "sqlite"):
|
||||
if dialect == "postgresql":
|
||||
from sqlalchemy.dialects.postgresql import insert
|
||||
else:
|
||||
from sqlalchemy.dialects.sqlite import insert
|
||||
stmt = insert(t)
|
||||
stmt = stmt.on_conflict_do_update(
|
||||
index_elements=[t.c.source_type, t.c.source_id, t.c.model],
|
||||
set_={c: stmt.excluded[c] for c in changing},
|
||||
)
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(stmt, values)
|
||||
return
|
||||
async with engine.begin() as conn:
|
||||
for row in values:
|
||||
key = (row["source_type"], row["source_id"])
|
||||
if key in replace:
|
||||
await conn.execute(
|
||||
t.update()
|
||||
.where(
|
||||
t.c.source_type == row["source_type"],
|
||||
t.c.source_id == row["source_id"],
|
||||
t.c.model == row["model"],
|
||||
)
|
||||
.values({c: row[c] for c in changing})
|
||||
)
|
||||
else:
|
||||
await conn.execute(t.insert(), [row])
|
||||
|
||||
async def flush(self) -> None:
|
||||
"""Wait until every queued text has been embedded (or failed)."""
|
||||
if self._queue is not None and self._loop is asyncio.get_running_loop():
|
||||
await self._queue.join()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
"""Drain the queue for up to ``timeout`` seconds, then stop.
|
||||
|
||||
Texts still queued after that are dropped and counted. Disposes the
|
||||
DatabaseSink's engine again if it owns it, since the dispatcher may
|
||||
have closed that sink first.
|
||||
"""
|
||||
if self._queue is not None and self._loop is asyncio.get_running_loop():
|
||||
try:
|
||||
await asyncio.wait_for(self._queue.join(), self.timeout)
|
||||
except TimeoutError:
|
||||
left = self._queue.qsize()
|
||||
self.dropped += left
|
||||
logger.warning(
|
||||
"Embedding queue did not drain in %ss; %d texts dropped",
|
||||
self.timeout,
|
||||
left,
|
||||
)
|
||||
if self._task is not None:
|
||||
self._task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError, RuntimeError):
|
||||
await self._task
|
||||
self._task = None
|
||||
self._queue = None
|
||||
if getattr(self.database_sink, "_owns_engine", False):
|
||||
await self.database_sink.engine.dispose()
|
||||
|
||||
# Searching
|
||||
|
||||
async def similar(
|
||||
self,
|
||||
text: str,
|
||||
*,
|
||||
source_types: Iterable[str] | None = None,
|
||||
k: int = 10,
|
||||
min_score: float | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Stored texts closest to ``text``. Raises on failure; see the middleware."""
|
||||
text = (self.redactor or self._default_redactor).redact_text(text)[
|
||||
: self.max_chars
|
||||
]
|
||||
[vector] = await asyncio.wait_for(self.embedder.embed([text]), self.timeout)
|
||||
return await self.nearest(
|
||||
vector, source_types=source_types, k=k, min_score=min_score
|
||||
)
|
||||
|
||||
async def vector_for(self, source_type: str, source_id: str) -> list[float] | None:
|
||||
"""The stored vector for a source under this embedder's model, if any."""
|
||||
t = self.table
|
||||
async with self.database_sink.engine.connect() as conn:
|
||||
row = (
|
||||
await conn.execute(
|
||||
select(t.c.embedding).where(
|
||||
t.c.source_type == source_type,
|
||||
t.c.source_id == str(source_id),
|
||||
t.c.model == self.embedder.model,
|
||||
)
|
||||
)
|
||||
).first()
|
||||
return None if row is None else list(row.embedding)
|
||||
|
||||
async def nearest(
|
||||
self,
|
||||
vector: Sequence[float],
|
||||
*,
|
||||
source_types: Iterable[str] | None = None,
|
||||
k: int = 10,
|
||||
min_score: float | None = None,
|
||||
exclude: tuple[str, str] | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Rows closest to ``vector`` by cosine, best first, as dicts with
|
||||
``source_type``, ``source_id``, ``text``, ``score`` and ``created_at``.
|
||||
|
||||
``score`` is 1 - cosine distance, clamped to [0, 1]. Only rows made by
|
||||
this embedder's model are compared. ``exclude`` is a
|
||||
(source_type, source_id) left out of the results.
|
||||
"""
|
||||
if k < 1:
|
||||
return []
|
||||
t = self.table
|
||||
conditions = [t.c.model == self.embedder.model]
|
||||
if source_types is not None:
|
||||
conditions.append(t.c.source_type.in_(list(source_types)))
|
||||
if exclude is not None:
|
||||
conditions.append(
|
||||
not_(and_(t.c.source_type == exclude[0], t.c.source_id == exclude[1]))
|
||||
)
|
||||
engine = self.database_sink.engine
|
||||
cols = (t.c.source_type, t.c.source_id, t.c.text, t.c.created_at)
|
||||
results: list[dict[str, Any]] = []
|
||||
if engine.dialect.name == "postgresql":
|
||||
distance = t.c.embedding.op("<=>", return_type=Float)(
|
||||
literal(list(vector), type_=EmbeddingVector(len(vector)))
|
||||
)
|
||||
query = (
|
||||
select(*cols, distance.label("distance"))
|
||||
.where(*conditions)
|
||||
.order_by(distance)
|
||||
.limit(k)
|
||||
)
|
||||
async with engine.connect() as conn:
|
||||
for row in (await conn.execute(query)).mappings():
|
||||
row = dict(row)
|
||||
distance_value = row.pop("distance")
|
||||
if distance_value is None or math.isnan(distance_value):
|
||||
continue
|
||||
row["score"] = _clamp(1.0 - distance_value)
|
||||
results.append(row)
|
||||
else:
|
||||
query = (
|
||||
select(*cols, t.c.embedding)
|
||||
.where(*conditions)
|
||||
.order_by(t.c.created_at.desc())
|
||||
.limit(self.scan_limit + 1)
|
||||
)
|
||||
async with engine.connect() as conn:
|
||||
rows = [dict(r) for r in (await conn.execute(query)).mappings()]
|
||||
if len(rows) > self.scan_limit:
|
||||
rows = rows[: self.scan_limit]
|
||||
if not self._warned_scan:
|
||||
self._warned_scan = True
|
||||
logger.info(
|
||||
"Similarity search without pgvector compares only the "
|
||||
"newest %d embeddings (scan_limit)",
|
||||
self.scan_limit,
|
||||
)
|
||||
for row in rows:
|
||||
score = cosine_score(vector, row.pop("embedding"))
|
||||
if score is None:
|
||||
continue
|
||||
row["score"] = score
|
||||
results.append(row)
|
||||
results.sort(key=lambda r: r["score"], reverse=True)
|
||||
results = results[:k]
|
||||
if min_score is not None:
|
||||
results = [r for r in results if r["score"] >= min_score]
|
||||
for row in results:
|
||||
created = row["created_at"]
|
||||
if created is not None and created.tzinfo is None:
|
||||
row["created_at"] = created.replace(tzinfo=UTC)
|
||||
return results
|
||||
@@ -27,6 +27,7 @@ from .correlation import (
|
||||
select_calls,
|
||||
)
|
||||
from .dispatch import RecordDispatcher
|
||||
from .embeddings import EmbeddingSink
|
||||
from .events import (
|
||||
MAX_KEY_CHARS,
|
||||
MAX_KIND_CHARS,
|
||||
@@ -161,6 +162,12 @@ class InstrumentationMiddleware(Middleware):
|
||||
when it returns less than the tool's rate. For tests.
|
||||
recent_events: how many recent events to keep in memory for
|
||||
``events_for``, including ones still waiting to reach the sinks.
|
||||
capture_llm_text: store ``prompt`` and ``completion`` given to
|
||||
``record_llm_call``. Off by default, since chat text can hold
|
||||
anything a user typed; when off they are discarded, including
|
||||
``prompt``/``completion`` keys passed in its ``attrs``.
|
||||
max_llm_text_chars: longest prompt or completion kept. Defaults to
|
||||
``8 * max_error_chars`` (4000).
|
||||
|
||||
Exclusion and sampling apply to tool calls only; events are always
|
||||
recorded. A sampled tool's call is kept whenever an event was recorded
|
||||
@@ -189,6 +196,8 @@ class InstrumentationMiddleware(Middleware):
|
||||
sample_rates: Mapping[str, float] | None = None,
|
||||
random: Callable[[], float] = _random.random,
|
||||
recent_events: int = 2000,
|
||||
capture_llm_text: bool = False,
|
||||
max_llm_text_chars: int | None = None,
|
||||
):
|
||||
mode = (mode or os.getenv(MODE_ENV_VAR) or "meta").lower()
|
||||
if mode not in MODES:
|
||||
@@ -219,8 +228,19 @@ class InstrumentationMiddleware(Middleware):
|
||||
self._call_scope: contextvars.ContextVar[CallScope | None] = (
|
||||
contextvars.ContextVar(f"ffb_call_scope_{id(self):x}", default=None)
|
||||
)
|
||||
self.capture_llm_text = capture_llm_text
|
||||
self.max_llm_text_chars = (
|
||||
max_llm_text_chars
|
||||
if max_llm_text_chars is not None
|
||||
else 8 * max_error_chars
|
||||
)
|
||||
sinks = list(sinks) if sinks is not None else [JsonLinesSink()]
|
||||
for sink in sinks:
|
||||
# Embedded text gets this middleware's redaction rules too.
|
||||
if isinstance(sink, EmbeddingSink) and sink.redactor is None:
|
||||
sink.redactor = self.redactor
|
||||
self.dispatcher = RecordDispatcher(
|
||||
sinks if sinks is not None else [JsonLinesSink()],
|
||||
sinks,
|
||||
max_queue=max_queue,
|
||||
batch_size=batch_size,
|
||||
sink_timeout=sink_timeout,
|
||||
@@ -410,6 +430,9 @@ class InstrumentationMiddleware(Middleware):
|
||||
and session; ``call_id`` given here wins. ``attrs`` is redacted like
|
||||
call arguments and made JSON-safe. ``kind`` must be a non-empty string
|
||||
of at most 128 characters and ``key`` at most 255 once made a string.
|
||||
A ``key`` used for joining must stay unique for the retention period
|
||||
and across restarts, so not a per-boot counter: use a UUID or a
|
||||
startup token plus a counter.
|
||||
|
||||
Returns False when the event was not recorded: mode ``off``, an invalid
|
||||
``kind``, ``key`` or ``attrs``, a full queue, or no event loop to hand
|
||||
@@ -469,6 +492,8 @@ class InstrumentationMiddleware(Middleware):
|
||||
error: str | BaseException | None = None,
|
||||
key: Any = None,
|
||||
attrs: Mapping[str, Any] | None = None,
|
||||
prompt: str | None = None,
|
||||
completion: str | None = None,
|
||||
) -> bool:
|
||||
"""Record one LLM request as an ``llm.call`` event. Never raises.
|
||||
|
||||
@@ -477,9 +502,22 @@ class InstrumentationMiddleware(Middleware):
|
||||
``error`` (redacted, truncated to ``max_error_chars``) plus
|
||||
``error_type`` when ``error`` is an exception. ``ok`` defaults to
|
||||
``error is None``. Extra ``attrs`` are kept, but these names win.
|
||||
|
||||
``prompt`` and ``completion`` are stored, redacted and truncated to
|
||||
``max_llm_text_chars``, only with ``capture_llm_text=True``;
|
||||
otherwise they are discarded.
|
||||
"""
|
||||
try:
|
||||
fields: dict[str, Any] = dict(attrs or {})
|
||||
fields.pop("prompt", None)
|
||||
fields.pop("completion", None)
|
||||
if self.capture_llm_text:
|
||||
for name, value in (("prompt", prompt), ("completion", completion)):
|
||||
if value is None and attrs is not None:
|
||||
value = attrs.get(name)
|
||||
if value is not None:
|
||||
text = self.redactor.redact_text(str(value))
|
||||
fields[name] = text[: self.max_llm_text_chars]
|
||||
fields["model"] = None if model is None else str(model)
|
||||
for name, value in (
|
||||
("provider", provider),
|
||||
@@ -513,6 +551,7 @@ class InstrumentationMiddleware(Middleware):
|
||||
memory, so an event is visible as soon as ``record_event`` returns (on
|
||||
the loop thread). Without a DatabaseSink, memory only. ``occurred_at``
|
||||
is a UTC datetime. A database failure is logged and memory is used.
|
||||
Events are ordered by ``occurred_at``; their ids are random UUIDs.
|
||||
"""
|
||||
key = str(key)
|
||||
kinds = None if kinds is None else frozenset(kinds)
|
||||
@@ -533,6 +572,77 @@ class InstrumentationMiddleware(Middleware):
|
||||
rows.sort(key=lambda r: r["occurred_at"])
|
||||
return rows[-limit:] if limit > 0 else []
|
||||
|
||||
@property
|
||||
def embedding_sink(self) -> EmbeddingSink | None:
|
||||
"""The first EmbeddingSink among the sinks, if any."""
|
||||
return next(
|
||||
(s for s in self.dispatcher.sinks if isinstance(s, EmbeddingSink)), None
|
||||
)
|
||||
|
||||
async def similar(
|
||||
self,
|
||||
text: str,
|
||||
*,
|
||||
source_types: Iterable[str] | None = None,
|
||||
k: int = 10,
|
||||
min_score: float | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Embedded texts most like ``text``, best first. Never raises.
|
||||
|
||||
Needs an EmbeddingSink. Each result has ``source_type``,
|
||||
``source_id``, ``text``, ``score`` (1 - cosine distance, in [0, 1])
|
||||
and ``created_at``. ``source_types`` limits the search to some of
|
||||
``feedback``, ``call_error``, ``event`` and ``llm``; ``min_score``
|
||||
drops weaker matches. Returns [] on any failure, logged.
|
||||
"""
|
||||
sink = self.embedding_sink
|
||||
if sink is None:
|
||||
logger.warning("similar() needs an EmbeddingSink")
|
||||
return []
|
||||
try:
|
||||
return await sink.similar(
|
||||
str(text),
|
||||
source_types=None if source_types is None else list(source_types),
|
||||
k=k,
|
||||
min_score=min_score,
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("Similarity search failed", exc_info=True)
|
||||
return []
|
||||
|
||||
async def similar_feedback(
|
||||
self,
|
||||
feedback_ref: Any,
|
||||
*,
|
||||
k: int = 10,
|
||||
source_types: Iterable[str] | None = None,
|
||||
min_score: float | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Texts most like an embedded feedback item, excluding the item itself.
|
||||
|
||||
Uses the feedback's stored vector, so nothing is embedded. Returns []
|
||||
when the feedback is not embedded yet, and on any failure. Never raises.
|
||||
"""
|
||||
sink = self.embedding_sink
|
||||
if sink is None:
|
||||
logger.warning("similar_feedback() needs an EmbeddingSink")
|
||||
return []
|
||||
ref = str(feedback_ref)
|
||||
try:
|
||||
vector = await sink.vector_for("feedback", ref)
|
||||
if vector is None:
|
||||
return []
|
||||
return await sink.nearest(
|
||||
vector,
|
||||
source_types=None if source_types is None else list(source_types),
|
||||
k=k,
|
||||
min_score=min_score,
|
||||
exclude=("feedback", ref),
|
||||
)
|
||||
except Exception:
|
||||
logger.warning("similar_feedback(%r) failed", ref, exc_info=True)
|
||||
return []
|
||||
|
||||
@property
|
||||
def database_sink(self) -> DatabaseSink | None:
|
||||
"""The first DatabaseSink among the sinks, if any."""
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
"""Async SQLAlchemy storage for tool-call records, events and feedback links.
|
||||
"""Async SQLAlchemy storage for tool-call records, events, feedback links and
|
||||
embeddings.
|
||||
|
||||
Hosts that manage their own schema include ``build_metadata(prefix)`` in their
|
||||
Alembic ``env.py`` (``target_metadata``) and pass their existing AsyncEngine to
|
||||
@@ -7,6 +8,7 @@ Alembic ``env.py`` (``target_metadata``) and pass their existing AsyncEngine to
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import random as _random
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from datetime import UTC, datetime, timedelta
|
||||
@@ -18,20 +20,25 @@ from sqlalchemy import (
|
||||
Column,
|
||||
DateTime,
|
||||
Float,
|
||||
Index,
|
||||
Integer,
|
||||
MetaData,
|
||||
String,
|
||||
Table,
|
||||
Text,
|
||||
UniqueConstraint,
|
||||
and_,
|
||||
delete,
|
||||
exists,
|
||||
inspect,
|
||||
or_,
|
||||
select,
|
||||
text,
|
||||
)
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
from sqlalchemy.exc import DBAPIError
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
||||
from sqlalchemy.types import TypeDecorator
|
||||
|
||||
from .events import MAX_KEY_CHARS, MAX_KIND_CHARS, EventRecord
|
||||
from .record import ToolCallRecord
|
||||
@@ -44,19 +51,62 @@ DEFAULT_TABLE_PREFIX = "ffb_"
|
||||
# none_as_null: a missing payload is SQL NULL, so "WHERE args IS NULL" works.
|
||||
_JSON = JSON(none_as_null=True).with_variant(JSONB(none_as_null=True), "postgresql")
|
||||
|
||||
DEFAULT_EMBEDDING_DIM = 1024
|
||||
# source_type values in {prefix}embeddings. Feedback rows survive pruning.
|
||||
EMBEDDING_SOURCES = ("feedback", "call_error", "event", "llm")
|
||||
FEEDBACK_SOURCE = "feedback"
|
||||
|
||||
|
||||
class EmbeddingVector(TypeDecorator):
|
||||
"""pgvector ``VECTOR(dim)`` on PostgreSQL, a JSON list of floats elsewhere.
|
||||
|
||||
pgvector is imported only when a PostgreSQL dialect needs the type, so
|
||||
SQLite hosts never need it installed.
|
||||
"""
|
||||
|
||||
impl = JSON
|
||||
cache_ok = True
|
||||
|
||||
def __init__(self, dim: int = DEFAULT_EMBEDDING_DIM):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
|
||||
def load_dialect_impl(self, dialect: Any) -> Any:
|
||||
if dialect.name == "postgresql":
|
||||
try:
|
||||
from pgvector.sqlalchemy import VECTOR
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"Embeddings on PostgreSQL need pgvector: install "
|
||||
'"fastmcp-feedback[embeddings]"'
|
||||
) from exc
|
||||
return dialect.type_descriptor(VECTOR(self.dim))
|
||||
return dialect.type_descriptor(JSON())
|
||||
|
||||
def process_bind_param(self, value: Any, dialect: Any) -> Any:
|
||||
return None if value is None else [float(v) for v in value]
|
||||
|
||||
def process_result_value(self, value: Any, dialect: Any) -> Any:
|
||||
return None if value is None else [float(v) for v in value]
|
||||
|
||||
|
||||
def build_metadata(
|
||||
prefix: str = DEFAULT_TABLE_PREFIX, metadata: MetaData | None = None
|
||||
prefix: str = DEFAULT_TABLE_PREFIX,
|
||||
metadata: MetaData | None = None,
|
||||
embedding_dim: int | None = DEFAULT_EMBEDDING_DIM,
|
||||
) -> MetaData:
|
||||
"""Return MetaData holding ``{prefix}tool_calls``, ``{prefix}events`` and
|
||||
``{prefix}feedback_call_links``.
|
||||
"""Return MetaData holding ``{prefix}tool_calls``, ``{prefix}events``,
|
||||
``{prefix}feedback_call_links`` and ``{prefix}embeddings``.
|
||||
|
||||
Pass an existing MetaData to add the tables to it instead.
|
||||
``embedding_dim`` sizes the vector column; None leaves the embeddings
|
||||
table out, for hosts that do not embed and have no pgvector.
|
||||
"""
|
||||
metadata = metadata if metadata is not None else MetaData()
|
||||
calls = f"{prefix}tool_calls"
|
||||
events = f"{prefix}events"
|
||||
links = f"{prefix}feedback_call_links"
|
||||
embeddings = f"{prefix}embeddings"
|
||||
if calls not in metadata.tables:
|
||||
Table(
|
||||
calls,
|
||||
@@ -117,6 +167,34 @@ def build_metadata(
|
||||
Column("rule", String(32), nullable=False),
|
||||
Column("linked_at", DateTime(timezone=True), nullable=False),
|
||||
)
|
||||
if embedding_dim is not None and embeddings not in metadata.tables:
|
||||
table = Table(
|
||||
embeddings,
|
||||
metadata,
|
||||
Column("id", String(36), primary_key=True),
|
||||
# feedback, call_error, event or llm
|
||||
Column("source_type", String(16), nullable=False, index=True),
|
||||
# The feedback ref, call id or event id the text came from.
|
||||
Column("source_id", String(255), nullable=False, index=True),
|
||||
Column("model", String(128), nullable=False),
|
||||
Column("dim", Integer, nullable=False),
|
||||
# sha256 hex of ``text``; an unchanged text is not embedded again.
|
||||
Column("text_hash", String(64), nullable=False),
|
||||
# What was embedded: redacted and truncated.
|
||||
Column("text", Text, nullable=False),
|
||||
Column("created_at", DateTime(timezone=True), nullable=False, index=True),
|
||||
Column("embedding", EmbeddingVector(embedding_dim), nullable=False),
|
||||
UniqueConstraint(
|
||||
"source_type", "source_id", "model", name=f"{embeddings}_source_uq"
|
||||
),
|
||||
)
|
||||
# Cosine HNSW index; other databases search by brute force.
|
||||
Index(
|
||||
f"{embeddings}_hnsw",
|
||||
table.c.embedding,
|
||||
postgresql_using="hnsw",
|
||||
postgresql_ops={"embedding": "vector_cosine_ops"},
|
||||
).ddl_if(dialect="postgresql")
|
||||
return metadata
|
||||
|
||||
|
||||
@@ -138,6 +216,8 @@ class DatabaseSink:
|
||||
prune_interval: minimum time between automatic prunes.
|
||||
prune_batch: rows deleted per statement, so no single delete holds
|
||||
locks for long.
|
||||
embedding_dim: vector size of ``{prefix}embeddings``; must match the
|
||||
embedder of any EmbeddingSink writing through this sink.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -149,6 +229,7 @@ class DatabaseSink:
|
||||
retention: timedelta | None = None,
|
||||
prune_interval: timedelta = timedelta(hours=1),
|
||||
prune_batch: int = 5000,
|
||||
embedding_dim: int = DEFAULT_EMBEDDING_DIM,
|
||||
):
|
||||
if retention is not None and retention <= timedelta(0):
|
||||
raise ValueError(f"retention must be positive, got {retention!r}")
|
||||
@@ -156,12 +237,19 @@ class DatabaseSink:
|
||||
raise ValueError(f"prune_batch must be at least 1, got {prune_batch!r}")
|
||||
self._owns_engine = isinstance(engine, str)
|
||||
self.engine = create_async_engine(engine) if isinstance(engine, str) else engine
|
||||
self.metadata = build_metadata(prefix)
|
||||
self.metadata = build_metadata(prefix, embedding_dim=embedding_dim)
|
||||
self.table = self.metadata.tables[f"{prefix}tool_calls"]
|
||||
self.events = self.metadata.tables[f"{prefix}events"]
|
||||
self.links = self.metadata.tables[f"{prefix}feedback_call_links"]
|
||||
self.embeddings = self.metadata.tables[f"{prefix}embeddings"]
|
||||
self.embedding_dim = embedding_dim
|
||||
self.create_tables = create_tables
|
||||
self._tables_ready = False
|
||||
# The embeddings table (and on Postgres the vector extension) is only
|
||||
# created once an EmbeddingSink uses this sink; see ensure_embeddings.
|
||||
self._embeddings_ready = False
|
||||
self._embeddings_failed_at: float | None = None
|
||||
self._embeddings_lock = asyncio.Lock()
|
||||
# The background writer and link_feedback can both reach _ready() at
|
||||
# once; without the lock both run create_all and one fails.
|
||||
self._tables_lock = asyncio.Lock()
|
||||
@@ -172,11 +260,96 @@ class DatabaseSink:
|
||||
self._last_prune: float | None = None
|
||||
|
||||
async def ensure_tables(self) -> None:
|
||||
"""Create the tables if missing. Only called when create_tables=True."""
|
||||
async with self.engine.begin() as conn:
|
||||
await conn.run_sync(self.metadata.create_all)
|
||||
"""Create the tables if missing. Only called when create_tables=True.
|
||||
|
||||
Another process starting against the same empty database can create a
|
||||
table between our existence check and our CREATE, which fails ours
|
||||
(SQLite: "already exists"; Postgres: DuplicateTable or a unique
|
||||
violation in pg_type). So on a database error, wait a moment and run
|
||||
create_all once more; it skips whatever the other process made. A
|
||||
second failure is raised.
|
||||
"""
|
||||
try:
|
||||
await self._create_all()
|
||||
except DBAPIError as first:
|
||||
await asyncio.sleep(_random.uniform(0.05, 0.25))
|
||||
await self._create_all()
|
||||
logger.info(
|
||||
"table creation failed once (%s), retry succeeded; another "
|
||||
"process was probably creating the same tables",
|
||||
type(first).__name__,
|
||||
)
|
||||
self._tables_ready = True
|
||||
|
||||
async def _create_all(self) -> None:
|
||||
# The embeddings table is left to ensure_embeddings, so hosts that do
|
||||
# not embed never need pgvector or the vector extension.
|
||||
tables = [t for t in self.metadata.sorted_tables if t is not self.embeddings]
|
||||
async with self.engine.begin() as conn:
|
||||
await conn.run_sync(self.metadata.create_all, tables=tables)
|
||||
|
||||
async def ensure_embeddings(self) -> bool:
|
||||
"""Create ``{prefix}embeddings`` if missing. Only when create_tables=True.
|
||||
|
||||
On PostgreSQL this first runs ``CREATE EXTENSION IF NOT EXISTS vector``.
|
||||
Never raises: a missing privilege, a missing pgvector package or any
|
||||
other failure is logged, and False is returned. A failed attempt is
|
||||
retried at most once a minute. Called by EmbeddingSink before it writes.
|
||||
"""
|
||||
if not self.create_tables or self._embeddings_ready:
|
||||
return True
|
||||
async with self._embeddings_lock:
|
||||
if self._embeddings_ready:
|
||||
return True
|
||||
now = time.monotonic()
|
||||
if (
|
||||
self._embeddings_failed_at is not None
|
||||
and now - self._embeddings_failed_at < 60
|
||||
):
|
||||
return False
|
||||
try:
|
||||
await self._ready()
|
||||
if self.engine.dialect.name == "postgresql":
|
||||
await self._create_vector_extension()
|
||||
try:
|
||||
await self._create_embeddings_table()
|
||||
except DBAPIError:
|
||||
# Another process creating it at the same moment.
|
||||
await asyncio.sleep(_random.uniform(0.05, 0.25))
|
||||
await self._create_embeddings_table()
|
||||
except Exception as exc:
|
||||
self._embeddings_failed_at = now
|
||||
logger.warning(
|
||||
"Could not create %s (%s: %s); embeddings are not stored until "
|
||||
"it exists. Create it with a migration (see the README) or "
|
||||
"fix the cause; retrying in a minute.",
|
||||
self.embeddings.name,
|
||||
type(exc).__name__,
|
||||
str(exc).splitlines()[0] if str(exc) else "",
|
||||
)
|
||||
return False
|
||||
self._embeddings_ready = True
|
||||
return True
|
||||
|
||||
async def _create_vector_extension(self) -> None:
|
||||
try:
|
||||
async with self.engine.begin() as conn:
|
||||
await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
|
||||
except DBAPIError as exc:
|
||||
# Usually a missing privilege. The table step then says whether
|
||||
# the extension exists anyway (a DBA may have created it).
|
||||
logger.warning(
|
||||
"CREATE EXTENSION vector failed (%s); a superuser or database "
|
||||
"owner must run it once in this database",
|
||||
str(exc.orig).splitlines()[0] if exc.orig else type(exc).__name__,
|
||||
)
|
||||
|
||||
async def _create_embeddings_table(self) -> None:
|
||||
async with self.engine.begin() as conn:
|
||||
await conn.run_sync(
|
||||
lambda sync: self.embeddings.create(sync, checkfirst=True)
|
||||
)
|
||||
|
||||
async def _ready(self) -> None:
|
||||
if self.create_tables and not self._tables_ready:
|
||||
async with self._tables_lock:
|
||||
@@ -216,15 +389,16 @@ class DatabaseSink:
|
||||
logger.warning("Pruning %s failed", self.table.name, exc_info=True)
|
||||
|
||||
async def prune(self, older_than: datetime | None = None) -> int:
|
||||
"""Delete tool-call and event rows from before ``older_than``.
|
||||
"""Delete tool-call, event and embedding rows from before ``older_than``.
|
||||
|
||||
Defaults to now minus ``retention``; with neither, deletes nothing.
|
||||
Calls linked to feedback in ``{prefix}feedback_call_links`` are kept,
|
||||
and links are never deleted. Events go by ``occurred_at``, whatever
|
||||
call they belong to. Deletes ``prune_batch`` rows per transaction until
|
||||
none are left. Returns the total rows deleted from both tables and
|
||||
logs each table's count. A database without the events table (not yet
|
||||
migrated) has its calls pruned and the events step skipped.
|
||||
call they belong to. Embeddings go by ``created_at``, except those of
|
||||
feedback (``source_type="feedback"``), which are kept. Deletes
|
||||
``prune_batch`` rows per transaction until none are left. Returns the
|
||||
total rows deleted from all three tables and logs each table's count.
|
||||
A table that does not exist yet (not migrated) is skipped.
|
||||
"""
|
||||
if older_than is None:
|
||||
if self.retention is None:
|
||||
@@ -245,15 +419,26 @@ class DatabaseSink:
|
||||
events = await self._delete_batches(
|
||||
e, select(e.c.id).where(e.c.occurred_at < older_than)
|
||||
)
|
||||
if calls or events:
|
||||
vectors = 0
|
||||
if await self._has_table(self.embeddings.name):
|
||||
v = self.embeddings
|
||||
vectors = await self._delete_batches(
|
||||
v,
|
||||
select(v.c.id).where(
|
||||
v.c.created_at < older_than, v.c.source_type != FEEDBACK_SOURCE
|
||||
),
|
||||
)
|
||||
if calls or events or vectors:
|
||||
logger.info(
|
||||
"Pruned %d rows from %s and %d from %s",
|
||||
"Pruned %d rows from %s, %d from %s and %d from %s",
|
||||
calls,
|
||||
t.name,
|
||||
events,
|
||||
self.events.name,
|
||||
vectors,
|
||||
self.embeddings.name,
|
||||
)
|
||||
return calls + events
|
||||
return calls + events + vectors
|
||||
|
||||
async def _delete_batches(self, table: Any, ids_query: Any) -> int:
|
||||
batch = ids_query.limit(self.prune_batch)
|
||||
|
||||
@@ -21,7 +21,7 @@ class SubmissionMixin(BaseFeedbackMixin):
|
||||
"""Mixin for feedback submission functionality.
|
||||
|
||||
Pass an InstrumentationMiddleware to link each submission to the tool
|
||||
calls that led up to it.
|
||||
calls that led up to it and record a ``feedback.submitted`` event for it.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -86,6 +86,17 @@ class SubmissionMixin(BaseFeedbackMixin):
|
||||
|
||||
linked_calls = 0
|
||||
if self.instrumentation is not None:
|
||||
# Feeds an EmbeddingSink, so similar reports can be found.
|
||||
self.instrumentation.record_event(
|
||||
"feedback.submitted",
|
||||
key=feedback_id,
|
||||
attrs={
|
||||
"type": request.type,
|
||||
"title": request.title,
|
||||
"description": request.description,
|
||||
"submitter": request.submitter,
|
||||
},
|
||||
)
|
||||
links = await self.instrumentation.capture_recent_calls()
|
||||
if await self.instrumentation.link_feedback(feedback_id, links):
|
||||
linked_calls = len(links)
|
||||
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
"""One writer process for the cross-process create_tables race test.
|
||||
|
||||
Usage: python create_tables_race_worker.py <db url> <trigger file> <events>
|
||||
|
||||
Builds a DatabaseSink with create_tables=True, waits until the trigger file
|
||||
exists so every worker hits the empty database at the same moment, records
|
||||
events, flushes, and prints {"sink_errors": n} as JSON. Prints "ready" once
|
||||
imports and setup are done, so the parent can release all workers together.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
from fastmcp_feedback.instrumentation import DatabaseSink, InstrumentationMiddleware
|
||||
|
||||
|
||||
async def main(url: str, trigger: str, count: int) -> dict:
|
||||
sink = DatabaseSink(url, create_tables=True)
|
||||
mw = InstrumentationMiddleware([sink])
|
||||
print("ready", flush=True)
|
||||
deadline = time.monotonic() + 30
|
||||
while not os.path.exists(trigger):
|
||||
if time.monotonic() > deadline:
|
||||
raise SystemExit("trigger file never appeared")
|
||||
await asyncio.sleep(0.001)
|
||||
for i in range(count):
|
||||
mw.record_event("race.write", key=f"{os.getpid()}-{i}")
|
||||
await mw.flush()
|
||||
errors = mw.dispatcher.sink_errors
|
||||
await mw.aclose() # also disposes the sink's engine
|
||||
return {"sink_errors": errors}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
url, trigger, count = sys.argv[1], sys.argv[2], int(sys.argv[3])
|
||||
print(json.dumps(asyncio.run(main(url, trigger, count))), flush=True)
|
||||
@@ -0,0 +1,565 @@
|
||||
"""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)
|
||||
@@ -0,0 +1,162 @@
|
||||
"""create_tables=True when several processes start against the same empty database.
|
||||
|
||||
Both processes see the tables missing and both CREATE them; one fails. The
|
||||
sink retries create_all once, which skips what the other process made.
|
||||
"""
|
||||
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import create_engine, func, select
|
||||
from sqlalchemy.exc import OperationalError
|
||||
|
||||
from fastmcp_feedback.instrumentation import (
|
||||
DatabaseSink,
|
||||
InstrumentationMiddleware,
|
||||
build_metadata,
|
||||
)
|
||||
|
||||
EVENTS = build_metadata().tables["ffb_events"]
|
||||
WORKER = Path(__file__).parent.parent / "fixtures" / "create_tables_race_worker.py"
|
||||
|
||||
|
||||
def _already_exists():
|
||||
return OperationalError(
|
||||
"CREATE TABLE ffb_tool_calls",
|
||||
{},
|
||||
Exception("table ffb_tool_calls already exists"),
|
||||
)
|
||||
|
||||
|
||||
def race(url: str, trigger: Path, workers: int = 2, events: int = 20) -> list[dict]:
|
||||
"""Start the workers, release them together, return their JSON results."""
|
||||
procs = [
|
||||
subprocess.Popen(
|
||||
[sys.executable, str(WORKER), url, str(trigger), str(events)],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
for _ in range(workers)
|
||||
]
|
||||
try:
|
||||
for p in procs:
|
||||
assert p.stdout.readline().strip() == "ready"
|
||||
trigger.touch()
|
||||
results = []
|
||||
for p in procs:
|
||||
out, err = p.communicate(timeout=30)
|
||||
assert p.returncode == 0, err
|
||||
results.append(json.loads(out.strip().splitlines()[-1]))
|
||||
return results
|
||||
finally:
|
||||
for p in procs:
|
||||
if p.poll() is None:
|
||||
p.kill()
|
||||
p.wait()
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
def test_two_processes_cold_start(tmp_path):
|
||||
workers, events = 2, 20
|
||||
for run in range(5):
|
||||
db = tmp_path / f"shared{run}.db"
|
||||
results = race(
|
||||
f"sqlite+aiosqlite:///{db}", tmp_path / f"go{run}", workers, events
|
||||
)
|
||||
engine = create_engine(f"sqlite:///{db}")
|
||||
try:
|
||||
with engine.connect() as conn:
|
||||
stored = conn.execute(
|
||||
select(func.count()).select_from(EVENTS)
|
||||
).scalar_one()
|
||||
finally:
|
||||
engine.dispose()
|
||||
assert [r["sink_errors"] for r in results] == [0] * workers, f"run {run}"
|
||||
assert stored == workers * events, f"run {run}"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def sink(tmp_path):
|
||||
s = DatabaseSink(f"sqlite+aiosqlite:///{tmp_path}/calls.db", create_tables=True)
|
||||
yield s
|
||||
await s.engine.dispose()
|
||||
|
||||
|
||||
class TestEnsureTablesRetry:
|
||||
async def test_first_create_fails_retry_succeeds(self, sink, monkeypatch, caplog):
|
||||
real = sink.metadata.create_all
|
||||
calls = []
|
||||
|
||||
def create_all(bind, **kwargs):
|
||||
calls.append(1)
|
||||
if len(calls) == 1:
|
||||
raise _already_exists()
|
||||
return real(bind, **kwargs)
|
||||
|
||||
monkeypatch.setattr(sink.metadata, "create_all", create_all)
|
||||
mw = InstrumentationMiddleware([sink])
|
||||
caplog.set_level("DEBUG", logger="fastmcp_feedback.instrumentation.storage")
|
||||
assert mw.record_event("a.b", key="k") is True
|
||||
await mw.flush()
|
||||
assert mw.dispatcher.sink_errors == 0
|
||||
assert len(calls) == 2
|
||||
assert sink._tables_ready
|
||||
async with sink.engine.connect() as conn:
|
||||
count = (
|
||||
await conn.execute(select(func.count()).select_from(sink.events))
|
||||
).scalar_one()
|
||||
assert count == 1
|
||||
retry_logs = [r for r in caplog.records if "retry succeeded" in r.getMessage()]
|
||||
assert retry_logs and all(r.levelname == "INFO" for r in retry_logs)
|
||||
assert not [
|
||||
r
|
||||
for r in caplog.records
|
||||
if r.levelname == "WARNING" and r.name.endswith("storage")
|
||||
]
|
||||
|
||||
async def test_persistent_failure_surfaces(self, sink, monkeypatch):
|
||||
calls = []
|
||||
|
||||
def create_all(bind, **kwargs):
|
||||
calls.append(1)
|
||||
raise _already_exists()
|
||||
|
||||
monkeypatch.setattr(sink.metadata, "create_all", create_all)
|
||||
mw = InstrumentationMiddleware([sink])
|
||||
assert mw.record_event("a.b") is True
|
||||
await mw.flush()
|
||||
assert mw.dispatcher.sink_errors == 1
|
||||
assert len(calls) == 2
|
||||
assert not sink._tables_ready
|
||||
# The next batch tries again rather than assuming the tables exist.
|
||||
assert mw.record_event("a.c") is True
|
||||
await mw.flush()
|
||||
assert mw.dispatcher.sink_errors == 2
|
||||
assert len(calls) == 4
|
||||
await mw.aclose()
|
||||
|
||||
async def test_non_database_error_not_retried(self, sink, monkeypatch):
|
||||
calls = []
|
||||
|
||||
def create_all(bind, **kwargs):
|
||||
calls.append(1)
|
||||
raise RuntimeError("bug")
|
||||
|
||||
monkeypatch.setattr(sink.metadata, "create_all", create_all)
|
||||
with pytest.raises(RuntimeError):
|
||||
await sink.ensure_tables()
|
||||
assert len(calls) == 1
|
||||
|
||||
async def test_create_tables_false_never_creates(self, tmp_path, monkeypatch):
|
||||
s = DatabaseSink(f"sqlite+aiosqlite:///{tmp_path}/none.db")
|
||||
calls = []
|
||||
monkeypatch.setattr(s.metadata, "create_all", lambda *a, **k: calls.append(1))
|
||||
try:
|
||||
await s._ready()
|
||||
finally:
|
||||
await s.engine.dispose()
|
||||
assert calls == []
|
||||
@@ -0,0 +1,848 @@
|
||||
"""Tests for EmbeddingSink, OpenAIEmbedder and similarity search.
|
||||
|
||||
DB tests take the ``db_url`` fixture, so they run on SQLite and, with
|
||||
FFB_TEST_DATABASE_URL set, on Postgres with pgvector.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from datetime import UTC, datetime, timedelta
|
||||
|
||||
import pytest
|
||||
from fastmcp import Client, FastMCP
|
||||
from fastmcp.exceptions import ToolError
|
||||
from sqlalchemy import func, select, text, update
|
||||
|
||||
from fastmcp_feedback import add_feedback_tools
|
||||
from fastmcp_feedback.instrumentation import (
|
||||
REDACTED,
|
||||
DatabaseSink,
|
||||
EmbeddingSink,
|
||||
EventRecord,
|
||||
OpenAIEmbedder,
|
||||
Redactor,
|
||||
ToolCallRecord,
|
||||
build_metadata,
|
||||
extract_texts,
|
||||
instrument,
|
||||
)
|
||||
from fastmcp_feedback.instrumentation.embeddings import cosine_score
|
||||
|
||||
DIM = 256
|
||||
SECRET = "sk-abcdefghijklmnopqrstuvwx"
|
||||
|
||||
|
||||
def token_vector(text_: str, dim: int = DIM) -> list[float]:
|
||||
"""Bag of hashed words, normalized: texts sharing words are close."""
|
||||
v = [0.0] * dim
|
||||
for tok in re.findall(r"[a-z0-9]+", text_.lower()):
|
||||
h = int.from_bytes(hashlib.sha256(tok.encode()).digest()[:4], "big")
|
||||
v[h % dim] += 1.0
|
||||
n = math.sqrt(sum(x * x for x in v)) or 1.0
|
||||
return [x / n for x in v]
|
||||
|
||||
|
||||
class FakeEmbedder:
|
||||
def __init__(
|
||||
self,
|
||||
dim=DIM,
|
||||
model="fake-embed",
|
||||
*,
|
||||
fail=False,
|
||||
delay=0.0,
|
||||
gate=None,
|
||||
out_dim=None,
|
||||
):
|
||||
self.dim = dim
|
||||
self.model = model
|
||||
self.fail = fail
|
||||
self.delay = delay
|
||||
self.gate = gate
|
||||
self.out_dim = out_dim or dim
|
||||
self.calls: list[list[str]] = []
|
||||
|
||||
@property
|
||||
def seen(self) -> list[str]:
|
||||
return [t for batch in self.calls for t in batch]
|
||||
|
||||
async def embed(self, texts):
|
||||
self.calls.append(list(texts))
|
||||
if self.gate is not None:
|
||||
await self.gate.wait()
|
||||
if self.delay:
|
||||
await asyncio.sleep(self.delay)
|
||||
if self.fail:
|
||||
raise RuntimeError("embedding backend down")
|
||||
return [token_vector(t, self.out_dim) for t in texts]
|
||||
|
||||
|
||||
@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 call_record(tool="render", outcome="error", message="render timed out", **kw):
|
||||
return ToolCallRecord(
|
||||
tool=tool,
|
||||
started_at=datetime.now(UTC),
|
||||
duration_ms=1.0,
|
||||
ok=outcome == "ok",
|
||||
outcome=outcome,
|
||||
mode="meta",
|
||||
error_type=None if outcome == "ok" else kw.pop("error_type", "TimeoutError"),
|
||||
error_message=None if outcome == "ok" else message,
|
||||
**kw,
|
||||
)
|
||||
|
||||
|
||||
def event(kind, key=None, **attrs):
|
||||
return EventRecord(kind=kind, occurred_at=datetime.now(UTC), key=key, attrs=attrs)
|
||||
|
||||
|
||||
async def rows(sink, *where):
|
||||
t = sink.embeddings
|
||||
if not await sink._has_table(t.name): # created on the first embedding
|
||||
return []
|
||||
async with sink.engine.connect() as conn:
|
||||
result = await conn.execute(select(t).where(*where).order_by(t.c.source_id))
|
||||
return [dict(r) for r in result.mappings()]
|
||||
|
||||
|
||||
async def settle(mw, emb):
|
||||
await mw.flush()
|
||||
await emb.flush()
|
||||
|
||||
|
||||
class TestExtractTexts:
|
||||
def test_call_errors(self):
|
||||
ok = call_record(outcome="ok")
|
||||
err = call_record(message="render timed out", error_type="TimeoutError")
|
||||
soft = call_record(
|
||||
tool="find",
|
||||
outcome="soft_error",
|
||||
message="status=unknown_target",
|
||||
error_type="SoftError",
|
||||
)
|
||||
items = extract_texts([ok, err, soft])
|
||||
assert [(i.source_type, i.source_id, i.text) for i in items] == [
|
||||
("call_error", err.id, "render: TimeoutError: render timed out"),
|
||||
("call_error", soft.id, "find: SoftError: status=unknown_target"),
|
||||
]
|
||||
|
||||
def test_feedback(self):
|
||||
ev = event(
|
||||
"feedback.submitted",
|
||||
key="bug-7Q2X",
|
||||
title="Render hangs",
|
||||
description="the render tool hung on scene a",
|
||||
type="bug",
|
||||
)
|
||||
[item] = extract_texts([ev])
|
||||
assert (item.source_type, item.source_id, item.text) == (
|
||||
"feedback",
|
||||
"bug-7Q2X",
|
||||
"Render hangs\nthe render tool hung on scene a",
|
||||
)
|
||||
|
||||
def test_event_text_keys(self):
|
||||
ev = event(
|
||||
"job.failed",
|
||||
key="job-1",
|
||||
error_excerpt="solver diverged",
|
||||
reason="residual NaN",
|
||||
nproc=8,
|
||||
)
|
||||
quiet = event("job.finished", key="job-2", duration_s=3)
|
||||
[item] = extract_texts([ev, quiet])
|
||||
assert (item.source_type, item.source_id, item.text) == (
|
||||
"event",
|
||||
ev.id,
|
||||
"solver diverged\nresidual NaN",
|
||||
)
|
||||
[custom] = extract_texts([quiet], event_text_keys=("duration_s",))
|
||||
assert custom.text == "3"
|
||||
|
||||
def test_llm(self):
|
||||
ev = event("llm.call", model="m", prompt="why red?", completion="the light")
|
||||
[item] = extract_texts([ev])
|
||||
assert (item.source_type, item.source_id, item.text) == (
|
||||
"llm",
|
||||
ev.id,
|
||||
"prompt: why red?\ncompletion: the light",
|
||||
)
|
||||
|
||||
def test_llm_without_text_is_a_plain_event(self):
|
||||
ev = event("llm.call", model="m", ok=False, error="upstream 529 overloaded")
|
||||
[item] = extract_texts([ev])
|
||||
assert (item.source_type, item.text) == ("event", "upstream 529 overloaded")
|
||||
|
||||
def test_sources_filter(self):
|
||||
records = [
|
||||
call_record(),
|
||||
event("feedback.submitted", key="1", title="t", description="d"),
|
||||
event("job.failed", error="boom"),
|
||||
event("llm.call", prompt="p"),
|
||||
]
|
||||
assert {i.source_type for i in extract_texts(records)} == {
|
||||
"call_error",
|
||||
"feedback",
|
||||
"event",
|
||||
"llm",
|
||||
}
|
||||
assert [
|
||||
i.source_type for i in extract_texts(records, sources=["feedback"])
|
||||
] == ["feedback"]
|
||||
|
||||
|
||||
class TestEmbeddingSink:
|
||||
async def test_every_source_lands(self, db):
|
||||
sink = db()
|
||||
fake = FakeEmbedder()
|
||||
emb = EmbeddingSink(fake, sink)
|
||||
err = call_record()
|
||||
records = [
|
||||
err,
|
||||
event(
|
||||
"feedback.submitted",
|
||||
key="42",
|
||||
title="Render hangs",
|
||||
description="scene a",
|
||||
),
|
||||
event("job.failed", key="job-1", error_excerpt="solver diverged"),
|
||||
event("llm.call", prompt="hello", completion="hi"),
|
||||
]
|
||||
await emb.write(records)
|
||||
await emb.flush()
|
||||
stored = await rows(sink)
|
||||
assert {r["source_type"] for r in stored} == {
|
||||
"call_error",
|
||||
"feedback",
|
||||
"event",
|
||||
"llm",
|
||||
}
|
||||
for r in stored:
|
||||
assert r["dim"] == DIM and r["model"] == "fake-embed"
|
||||
assert len(r["embedding"]) == DIM
|
||||
assert r["text_hash"] == hashlib.sha256(r["text"].encode()).hexdigest()
|
||||
assert r["created_at"] is not None
|
||||
assert emb.embedded == 4 and emb.errors == 0
|
||||
await emb.aclose()
|
||||
|
||||
async def test_submit_feedback_is_embedded(self, db, tmp_path):
|
||||
app = FastMCP("app")
|
||||
sink = db("fb")
|
||||
fake = FakeEmbedder()
|
||||
emb = EmbeddingSink(fake, sink)
|
||||
mw = instrument(app, [sink, emb])
|
||||
add_feedback_tools(
|
||||
app, database_url=f"sqlite:///{tmp_path}/fb.db", instrumentation=mw
|
||||
)
|
||||
async with Client(app) as client:
|
||||
result = await client.call_tool(
|
||||
"submit_feedback",
|
||||
{
|
||||
"request": {
|
||||
"type": "bug",
|
||||
"title": "Render hangs",
|
||||
"description": "the render tool hung on scene a",
|
||||
"submitter": "model",
|
||||
}
|
||||
},
|
||||
)
|
||||
ref = result.data["feedback_id"]
|
||||
await settle(mw, emb)
|
||||
[row] = await rows(sink, sink.embeddings.c.source_type == "feedback")
|
||||
assert row["source_id"] == ref
|
||||
assert row["text"] == "Render hangs\nthe render tool hung on scene a"
|
||||
[ev] = await mw.events_for(ref, kinds=["feedback.submitted"])
|
||||
assert ev["attrs"]["type"] == "bug" and ev["attrs"]["submitter"] == "model"
|
||||
await mw.aclose()
|
||||
|
||||
async def test_redaction(self, db):
|
||||
sink = db()
|
||||
fake = FakeEmbedder()
|
||||
emb = EmbeddingSink(
|
||||
fake, sink, redactor=Redactor(extra_patterns=[r"acme_\w{8}"])
|
||||
)
|
||||
# Straight to the sink, so upstream redaction cannot hide a gap here.
|
||||
await emb.write(
|
||||
[
|
||||
call_record(message=f"401 with api_key={SECRET} and acme_12345678"),
|
||||
event("job.failed", error=f"Bearer {SECRET}"),
|
||||
]
|
||||
)
|
||||
await emb.flush()
|
||||
assert fake.seen and all(
|
||||
SECRET not in t and "acme_12345678" not in t for t in fake.seen
|
||||
)
|
||||
stored = await rows(sink)
|
||||
assert len(stored) == 2
|
||||
for r in stored:
|
||||
assert SECRET not in r["text"] and REDACTED in r["text"]
|
||||
await emb.aclose()
|
||||
|
||||
async def test_instrument_hands_over_its_redactor(self, db):
|
||||
app = FastMCP("app")
|
||||
sink = db()
|
||||
emb = EmbeddingSink(FakeEmbedder(), sink)
|
||||
redactor = Redactor(extra_patterns=[r"acme_\w{8}"])
|
||||
mw = instrument(app, [sink, emb], redactor=redactor)
|
||||
assert emb.redactor is redactor
|
||||
await mw.aclose()
|
||||
|
||||
async def test_max_chars(self, db):
|
||||
sink = db()
|
||||
fake = FakeEmbedder()
|
||||
emb = EmbeddingSink(fake, sink, max_chars=50)
|
||||
await emb.write([event("job.failed", error="word " * 100)])
|
||||
await emb.flush()
|
||||
[row] = await rows(sink)
|
||||
assert len(row["text"]) == 50 and fake.seen == [row["text"]]
|
||||
await emb.aclose()
|
||||
|
||||
async def test_dedupe_and_replace(self, db):
|
||||
sink = db()
|
||||
fake = FakeEmbedder()
|
||||
emb = EmbeddingSink(fake, sink)
|
||||
ev = event("feedback.submitted", key="7", title="Render hangs", description="a")
|
||||
await emb.write([ev])
|
||||
await emb.flush()
|
||||
await emb.write([ev, ev])
|
||||
await emb.flush()
|
||||
assert fake.seen == ["Render hangs\na"]
|
||||
assert emb.skipped == 2
|
||||
changed = event(
|
||||
"feedback.submitted", key="7", title="Render hangs", description="b"
|
||||
)
|
||||
await emb.write([changed])
|
||||
await emb.flush()
|
||||
[row] = await rows(sink)
|
||||
assert row["text"] == "Render hangs\nb"
|
||||
assert fake.seen == ["Render hangs\na", "Render hangs\nb"]
|
||||
# Another model's vectors are kept apart.
|
||||
other = EmbeddingSink(FakeEmbedder(model="other"), sink)
|
||||
await other.write([changed])
|
||||
await other.flush()
|
||||
assert len(await rows(sink)) == 2
|
||||
await emb.aclose()
|
||||
await other.aclose()
|
||||
|
||||
async def test_queue_full_drops(self, db):
|
||||
sink = db()
|
||||
gate = asyncio.Event()
|
||||
fake = FakeEmbedder(gate=gate)
|
||||
emb = EmbeddingSink(fake, sink, max_queue=2, batch_size=1)
|
||||
start = time.perf_counter()
|
||||
await emb.write([event("job.failed", error=f"boom {i}") for i in range(6)])
|
||||
assert time.perf_counter() - start < 0.5
|
||||
await asyncio.sleep(0.05) # the worker takes one and blocks on the gate
|
||||
await emb.write([event("job.failed", error=f"late {i}") for i in range(3)])
|
||||
assert emb.dropped >= 4
|
||||
gate.set()
|
||||
await emb.flush()
|
||||
assert emb.errors == 0
|
||||
assert len(await rows(sink)) == 9 - emb.dropped
|
||||
await emb.aclose()
|
||||
|
||||
async def test_embedder_errors_never_reach_tool_calls(self, db):
|
||||
app = FastMCP("app")
|
||||
|
||||
@app.tool
|
||||
def render(scene: str) -> str:
|
||||
raise RuntimeError(f"render of {scene} timed out")
|
||||
|
||||
@app.tool
|
||||
def ping() -> str:
|
||||
return "pong"
|
||||
|
||||
sink = db()
|
||||
emb = EmbeddingSink(FakeEmbedder(fail=True), sink)
|
||||
mw = instrument(app, [sink, emb])
|
||||
async with Client(app) as client:
|
||||
with contextlib.suppress(ToolError):
|
||||
await client.call_tool("render", {"scene": "a"})
|
||||
assert (await client.call_tool("ping", {})).data == "pong"
|
||||
await settle(mw, emb)
|
||||
assert emb.errors == 1
|
||||
assert mw.dispatcher.sink_errors == 0
|
||||
assert await rows(sink) == []
|
||||
async with sink.engine.connect() as conn:
|
||||
calls = (
|
||||
await conn.execute(select(func.count()).select_from(sink.table))
|
||||
).scalar_one()
|
||||
assert calls == 2
|
||||
assert await mw.similar("render timed out") == []
|
||||
await mw.aclose()
|
||||
|
||||
async def test_slow_embedder_stays_off_the_call_path(self, db):
|
||||
app = FastMCP("app")
|
||||
|
||||
@app.tool
|
||||
def render() -> str:
|
||||
raise RuntimeError("render timed out")
|
||||
|
||||
sink = db()
|
||||
emb = EmbeddingSink(FakeEmbedder(delay=2.0), sink, timeout=5.0)
|
||||
mw = instrument(app, [sink, emb], sink_timeout=0.5)
|
||||
async with Client(app) as client:
|
||||
start = time.perf_counter()
|
||||
for _ in range(3):
|
||||
with contextlib.suppress(ToolError):
|
||||
await client.call_tool("render", {})
|
||||
elapsed = time.perf_counter() - start
|
||||
await mw.flush()
|
||||
assert elapsed < 1.5
|
||||
assert mw.dispatcher.sink_errors == 0
|
||||
await emb.flush()
|
||||
assert emb.embedded == 3
|
||||
await mw.aclose()
|
||||
|
||||
async def test_wrong_vector_length_is_an_error(self, db):
|
||||
sink = db()
|
||||
emb = EmbeddingSink(FakeEmbedder(out_dim=DIM + 1), sink)
|
||||
await emb.write([event("job.failed", error="boom")])
|
||||
await emb.flush()
|
||||
assert emb.errors == 1 and await rows(sink) == []
|
||||
await emb.aclose()
|
||||
|
||||
async def test_dim_mismatch_refused_at_construction(self, db):
|
||||
with pytest.raises(ValueError, match="embedding_dim=512"):
|
||||
EmbeddingSink(FakeEmbedder(dim=512), db())
|
||||
|
||||
async def test_unknown_source_refused(self, db):
|
||||
with pytest.raises(ValueError, match="unknown sources"):
|
||||
EmbeddingSink(FakeEmbedder(), db(), sources=["feedback", "chat"])
|
||||
|
||||
async def test_aclose_drains(self, db):
|
||||
sink = db()
|
||||
emb = EmbeddingSink(FakeEmbedder(delay=0.1), sink, batch_size=1)
|
||||
await emb.write([event("job.failed", error=f"e{i}") for i in range(3)])
|
||||
await emb.aclose()
|
||||
assert emb.embedded == 3 and emb.dropped == 0
|
||||
|
||||
async def test_aclose_timeout_drops(self, db):
|
||||
sink = db()
|
||||
emb = EmbeddingSink(
|
||||
FakeEmbedder(gate=asyncio.Event()), sink, batch_size=1, timeout=0.2
|
||||
)
|
||||
await emb.write([event("job.failed", error=f"e{i}") for i in range(3)])
|
||||
await emb.aclose()
|
||||
assert emb.dropped == 2 and emb.embedded == 0
|
||||
|
||||
|
||||
class TestLlmText:
|
||||
def app(self, sinks, **kwargs):
|
||||
app = FastMCP("llm")
|
||||
holder = {}
|
||||
|
||||
@app.tool
|
||||
async def ask(q: str) -> str:
|
||||
holder["mw"].record_llm_call(
|
||||
model="m",
|
||||
prompt=f"{q} token={SECRET}",
|
||||
completion="because",
|
||||
attrs={"prompt": "smuggled", "temperature": 0.2},
|
||||
)
|
||||
return "because"
|
||||
|
||||
holder["mw"] = mw = instrument(app, sinks, **kwargs)
|
||||
return app, mw
|
||||
|
||||
async def test_dropped_by_default(self, db):
|
||||
sink = db()
|
||||
fake = FakeEmbedder()
|
||||
emb = EmbeddingSink(fake, sink)
|
||||
app, mw = self.app([sink, emb])
|
||||
async with Client(app) as client:
|
||||
await client.call_tool("ask", {"q": "why is the sky blue"})
|
||||
await settle(mw, emb)
|
||||
[ev] = [e for e in mw.recent_events if e.kind == "llm.call"]
|
||||
assert "prompt" not in ev.attrs and "completion" not in ev.attrs
|
||||
assert ev.attrs["temperature"] == 0.2
|
||||
assert fake.seen == []
|
||||
assert await rows(sink) == []
|
||||
await mw.aclose()
|
||||
|
||||
async def test_embedded_when_captured(self, db):
|
||||
sink = db()
|
||||
fake = FakeEmbedder()
|
||||
emb = EmbeddingSink(fake, sink)
|
||||
app, mw = self.app([sink, emb], capture_llm_text=True, max_error_chars=10)
|
||||
async with Client(app) as client:
|
||||
await client.call_tool("ask", {"q": "why is the sky blue " * 10})
|
||||
await settle(mw, emb)
|
||||
[ev] = [e for e in mw.recent_events if e.kind == "llm.call"]
|
||||
assert len(ev.attrs["prompt"]) == 80 # 8 * max_error_chars
|
||||
assert ev.attrs["completion"] == "because"
|
||||
[row] = await rows(sink)
|
||||
assert row["source_type"] == "llm" and row["source_id"] == ev.id
|
||||
assert row["text"].startswith("prompt: why is the sky blue")
|
||||
assert row["text"].endswith("completion: because")
|
||||
await mw.aclose()
|
||||
|
||||
async def test_captured_text_is_redacted(self, db):
|
||||
sink = db()
|
||||
emb = EmbeddingSink(FakeEmbedder(), sink)
|
||||
app, mw = self.app([sink, emb], capture_llm_text=True)
|
||||
async with Client(app) as client:
|
||||
await client.call_tool("ask", {"q": "q"})
|
||||
await settle(mw, emb)
|
||||
[ev] = [e for e in mw.recent_events if e.kind == "llm.call"]
|
||||
assert SECRET not in ev.attrs["prompt"] and REDACTED in ev.attrs["prompt"]
|
||||
[row] = await rows(sink)
|
||||
assert SECRET not in row["text"]
|
||||
await mw.aclose()
|
||||
|
||||
|
||||
FIXTURE = [
|
||||
("call_error", "c1", "render: TimeoutError: render timed out after 30 s"),
|
||||
("call_error", "c2", "upload: OSError: disk full on volume data"),
|
||||
("event", "e1", "render job timed out waiting for the gpu"),
|
||||
("feedback", "bug-1", "Render hangs\nthe render tool timed out on scene a"),
|
||||
("feedback", "bug-2", "Dark mode please\nthe ui is too bright at night"),
|
||||
("llm", "l1", "prompt: summarise the log\ncompletion: nothing notable"),
|
||||
]
|
||||
|
||||
|
||||
async def seed(sink, emb):
|
||||
from fastmcp_feedback.instrumentation import EmbeddingText
|
||||
|
||||
for source_type, source_id, text_ in FIXTURE:
|
||||
emb._enqueue(EmbeddingText(source_type, source_id, text_))
|
||||
await emb.flush()
|
||||
assert len(await rows(sink)) == len(FIXTURE)
|
||||
|
||||
|
||||
class TestSimilar:
|
||||
async def test_ranking_matches_brute_force(self, db, db_backend):
|
||||
sink = db()
|
||||
app = FastMCP("s")
|
||||
emb = EmbeddingSink(FakeEmbedder(), sink)
|
||||
mw = instrument(app, [sink, emb])
|
||||
await seed(sink, emb)
|
||||
|
||||
query = "render timed out"
|
||||
q = token_vector(query)
|
||||
expected = sorted(
|
||||
((cosine_score(q, token_vector(t)), sid) for _, sid, t in FIXTURE),
|
||||
reverse=True,
|
||||
)
|
||||
found = await mw.similar(query, k=len(FIXTURE))
|
||||
# Rows sharing no word with the query all score 0; their order is free.
|
||||
positive = [sid for score, sid in expected if score > 1e-6]
|
||||
assert len(positive) >= 3
|
||||
assert [r["source_id"] for r in found][: len(positive)] == positive
|
||||
assert {r["source_id"] for r in found} == {sid for _, sid in expected}
|
||||
want = {sid: score for score, sid in expected}
|
||||
for r in found:
|
||||
assert r["score"] == pytest.approx(want[r["source_id"]], abs=1e-4)
|
||||
assert 0.0 <= r["score"] <= 1.0
|
||||
assert r["created_at"].tzinfo is not None
|
||||
assert set(r) == {"source_type", "source_id", "text", "score", "created_at"}
|
||||
assert found[0]["source_id"] == "c1"
|
||||
|
||||
top = await mw.similar(query, k=2)
|
||||
assert [r["source_id"] for r in top] == positive[:2]
|
||||
only = await mw.similar(query, source_types=["feedback"])
|
||||
assert [r["source_id"] for r in only] == ["bug-1", "bug-2"]
|
||||
strong = await mw.similar(query, min_score=0.4)
|
||||
assert strong and all(r["score"] >= 0.4 for r in strong)
|
||||
assert len(strong) < len(FIXTURE)
|
||||
await mw.aclose()
|
||||
|
||||
async def test_similar_feedback_excludes_itself(self, db):
|
||||
sink = db()
|
||||
app = FastMCP("s")
|
||||
emb = EmbeddingSink(FakeEmbedder(), sink)
|
||||
mw = instrument(app, [sink, emb])
|
||||
await seed(sink, emb)
|
||||
found = await mw.similar_feedback("bug-1", k=3)
|
||||
assert "bug-1" not in [r["source_id"] for r in found]
|
||||
assert found[0]["source_id"] in ("c1", "e1")
|
||||
feedback_only = await mw.similar_feedback("bug-1", source_types=["feedback"])
|
||||
assert [r["source_id"] for r in feedback_only] == ["bug-2"]
|
||||
assert await mw.similar_feedback("bug-404") == []
|
||||
await mw.aclose()
|
||||
|
||||
async def test_scan_limit_without_pgvector(self, db, db_backend, caplog):
|
||||
if db_backend == "postgres":
|
||||
pytest.skip("brute force only runs without pgvector")
|
||||
sink = db()
|
||||
app = FastMCP("s")
|
||||
emb = EmbeddingSink(FakeEmbedder(), sink, scan_limit=2)
|
||||
mw = instrument(app, [sink, emb])
|
||||
await seed(sink, emb)
|
||||
caplog.set_level(logging.INFO, "fastmcp_feedback.instrumentation.embeddings")
|
||||
assert len(await mw.similar("render", k=10)) <= 2
|
||||
await mw.similar("render", k=10)
|
||||
assert len([r for r in caplog.records if "scan_limit" in r.getMessage()]) == 1
|
||||
await mw.aclose()
|
||||
|
||||
async def test_never_raises(self, db):
|
||||
app = FastMCP("s")
|
||||
mw = instrument(app, [])
|
||||
assert await mw.similar("x") == []
|
||||
assert await mw.similar_feedback("1") == []
|
||||
await mw.aclose()
|
||||
sink = db()
|
||||
emb = EmbeddingSink(FakeEmbedder(fail=True), sink)
|
||||
mw = instrument(FastMCP("t"), [sink, emb])
|
||||
assert await mw.similar("x") == []
|
||||
# No table yet: the lookup itself fails.
|
||||
assert await mw.similar_feedback("1") == []
|
||||
await mw.aclose()
|
||||
|
||||
|
||||
class TestPrune:
|
||||
async def test_keeps_feedback_embeddings(self, db):
|
||||
sink = db(prune_batch=2)
|
||||
emb = EmbeddingSink(FakeEmbedder(), sink)
|
||||
await seed(sink, emb)
|
||||
t = sink.embeddings
|
||||
old = datetime.now(UTC) - timedelta(days=40)
|
||||
async with sink.engine.begin() as conn:
|
||||
await conn.execute(
|
||||
update(t).where(t.c.source_id != "c2").values(created_at=old)
|
||||
)
|
||||
deleted = await sink.prune(older_than=datetime.now(UTC) - timedelta(days=30))
|
||||
assert deleted == 3
|
||||
left = sorted(r["source_id"] for r in await rows(sink))
|
||||
assert left == ["bug-1", "bug-2", "c2"]
|
||||
await emb.aclose()
|
||||
|
||||
async def test_prune_without_embeddings_table(self, db):
|
||||
sink = db()
|
||||
await sink.write([call_record()])
|
||||
assert not await sink._has_table(sink.embeddings.name)
|
||||
assert await sink.prune(older_than=datetime.now(UTC) + timedelta(days=1)) == 1
|
||||
|
||||
|
||||
class TestCreateTables:
|
||||
async def test_embeddings_table_only_with_an_embedding_sink(self, db):
|
||||
sink = db()
|
||||
await sink.write([call_record()])
|
||||
assert not await sink._has_table(sink.embeddings.name)
|
||||
emb = EmbeddingSink(FakeEmbedder(), sink)
|
||||
await emb.write([call_record()])
|
||||
await emb.flush()
|
||||
assert await sink._has_table(sink.embeddings.name)
|
||||
await emb.aclose()
|
||||
|
||||
async def test_create_tables_false_creates_nothing(self, db):
|
||||
sink = db(create_tables=False)
|
||||
assert await sink.ensure_embeddings() is True
|
||||
assert not await sink._has_table(sink.embeddings.name)
|
||||
|
||||
async def test_failure_is_logged_not_raised(self, db, monkeypatch, caplog):
|
||||
sink = db()
|
||||
|
||||
async def boom():
|
||||
raise RuntimeError("no room")
|
||||
|
||||
monkeypatch.setattr(sink, "_create_embeddings_table", boom)
|
||||
caplog.set_level(logging.WARNING)
|
||||
emb = EmbeddingSink(FakeEmbedder(), sink)
|
||||
assert await sink.ensure_embeddings() is False
|
||||
assert any("Could not create" in r.getMessage() for r in caplog.records)
|
||||
await emb.write([call_record()])
|
||||
await emb.flush()
|
||||
assert emb.errors == 1 # the insert fails; counted, not raised
|
||||
await emb.aclose()
|
||||
|
||||
@pytest.mark.postgres
|
||||
async def test_postgres_extension_table_and_index(self, db, db_backend):
|
||||
if db_backend != "postgres":
|
||||
pytest.skip("Postgres only")
|
||||
sink = db()
|
||||
await sink._ready()
|
||||
async with sink.engine.begin() as conn:
|
||||
await conn.execute(text("DROP EXTENSION IF EXISTS vector"))
|
||||
assert await sink.ensure_embeddings() is True
|
||||
async with sink.engine.connect() as conn:
|
||||
ext = (
|
||||
await conn.execute(
|
||||
text("SELECT count(*) FROM pg_extension WHERE extname = 'vector'")
|
||||
)
|
||||
).scalar_one()
|
||||
index = (
|
||||
await conn.execute(
|
||||
text(
|
||||
"SELECT indexdef FROM pg_indexes "
|
||||
"WHERE indexname = 'ffb_embeddings_hnsw'"
|
||||
)
|
||||
)
|
||||
).scalar_one()
|
||||
coltype = (
|
||||
await conn.execute(
|
||||
text(
|
||||
"SELECT format_type(atttypid, atttypmod) FROM pg_attribute "
|
||||
"WHERE attrelid = 'ffb_embeddings'::regclass "
|
||||
"AND attname = 'embedding'"
|
||||
)
|
||||
)
|
||||
).scalar_one()
|
||||
assert ext == 1
|
||||
assert "hnsw" in index and "vector_cosine_ops" in index
|
||||
assert coltype == f"vector({DIM})"
|
||||
|
||||
@pytest.mark.postgres
|
||||
async def test_postgres_without_privilege(self, db, db_backend, db_url, caplog):
|
||||
if db_backend != "postgres":
|
||||
pytest.skip("Postgres only")
|
||||
from sqlalchemy.engine import make_url
|
||||
|
||||
admin = db("priv", embedding_dim=DIM)
|
||||
role = f"ffb_limited_{uuid.uuid4().hex[:8]}"
|
||||
password = uuid.uuid4().hex
|
||||
url = make_url(db_url("priv"))
|
||||
async with admin.engine.begin() as conn:
|
||||
await conn.execute(text("DROP EXTENSION IF EXISTS vector"))
|
||||
await conn.execute(text(f"CREATE ROLE {role} LOGIN PASSWORD '{password}'"))
|
||||
await conn.execute(text(f"GRANT ALL ON SCHEMA public TO {role}"))
|
||||
limited = DatabaseSink(
|
||||
url.set(username=role, password=password).render_as_string(
|
||||
hide_password=False
|
||||
),
|
||||
create_tables=True,
|
||||
embedding_dim=DIM,
|
||||
)
|
||||
try:
|
||||
caplog.set_level(logging.WARNING)
|
||||
emb = EmbeddingSink(FakeEmbedder(), limited)
|
||||
assert await limited.ensure_embeddings() is False
|
||||
messages = [r.getMessage() for r in caplog.records]
|
||||
assert any("CREATE EXTENSION vector failed" in m for m in messages)
|
||||
assert any("Could not create" in m for m in messages)
|
||||
await limited.write([call_record()]) # the rest still works
|
||||
await emb.write([call_record()])
|
||||
await emb.flush()
|
||||
assert emb.errors == 1
|
||||
await emb.aclose()
|
||||
finally:
|
||||
await limited.engine.dispose()
|
||||
async with admin.engine.begin() as conn:
|
||||
await conn.execute(text(f"DROP OWNED BY {role}"))
|
||||
await conn.execute(text(f"DROP ROLE {role}"))
|
||||
|
||||
|
||||
class TestMetadata:
|
||||
def test_build_metadata_tables(self):
|
||||
md = build_metadata("x_", embedding_dim=768)
|
||||
t = md.tables["x_embeddings"]
|
||||
assert t.c.embedding.type.dim == 768
|
||||
assert {c.name for c in t.c} == {
|
||||
"id",
|
||||
"source_type",
|
||||
"source_id",
|
||||
"model",
|
||||
"dim",
|
||||
"text_hash",
|
||||
"text",
|
||||
"created_at",
|
||||
"embedding",
|
||||
}
|
||||
without = build_metadata("y_", embedding_dim=None)
|
||||
assert "y_embeddings" not in without.tables
|
||||
assert "y_tool_calls" in without.tables
|
||||
|
||||
|
||||
class TestOpenAIEmbedder:
|
||||
def transport(self, handler):
|
||||
from fastmcp_feedback.instrumentation.embeddings import _http_module
|
||||
|
||||
return _http_module().MockTransport(handler)
|
||||
|
||||
async def test_request_batches_and_order(self):
|
||||
seen = []
|
||||
|
||||
def handler(request):
|
||||
body = json.loads(request.content)
|
||||
seen.append((request.url.path, request.headers.get("authorization"), body))
|
||||
data = [
|
||||
{"index": i, "embedding": [float(len(t)), 0.0, 1.0]}
|
||||
for i, t in enumerate(body["input"])
|
||||
]
|
||||
from fastmcp_feedback.instrumentation.embeddings import _http_module
|
||||
|
||||
return _http_module().Response(200, json={"data": list(reversed(data))})
|
||||
|
||||
e = OpenAIEmbedder(
|
||||
"https://gw.example/v1/",
|
||||
api_key="k",
|
||||
model="m",
|
||||
dim=3,
|
||||
batch_size=2,
|
||||
transport=self.transport(handler),
|
||||
)
|
||||
out = await e.embed(["a", "bb", "ccc"])
|
||||
assert out == [[1.0, 0.0, 1.0], [2.0, 0.0, 1.0], [3.0, 0.0, 1.0]]
|
||||
assert [s[0] for s in seen] == ["/v1/embeddings", "/v1/embeddings"]
|
||||
assert seen[0][1] == "Bearer k"
|
||||
assert [s[2]["input"] for s in seen] == [["a", "bb"], ["ccc"]]
|
||||
assert all(s[2]["model"] == "m" for s in seen)
|
||||
assert await e.embed([]) == []
|
||||
|
||||
async def test_dim_mismatch_raises(self):
|
||||
from fastmcp_feedback.instrumentation.embeddings import _http_module
|
||||
|
||||
def handler(request):
|
||||
return _http_module().Response(
|
||||
200, json={"data": [{"index": 0, "embedding": [0.1, 0.2]}]}
|
||||
)
|
||||
|
||||
e = OpenAIEmbedder(
|
||||
"https://gw.example/v1", dim=3, transport=self.transport(handler)
|
||||
)
|
||||
with pytest.raises(ValueError, match="2-dim vector, expected 3"):
|
||||
await e.embed(["a"])
|
||||
|
||||
async def test_http_error_raises(self):
|
||||
from fastmcp_feedback.instrumentation.embeddings import _http_module
|
||||
|
||||
e = OpenAIEmbedder(
|
||||
"https://gw.example/v1",
|
||||
dim=3,
|
||||
transport=self.transport(
|
||||
lambda r: _http_module().Response(403, json={"error": "nope"})
|
||||
),
|
||||
)
|
||||
with pytest.raises(Exception, match="403"):
|
||||
await e.embed(["a"])
|
||||
|
||||
|
||||
LIVE_URL = os.environ.get("FFB_TEST_EMBED_URL")
|
||||
|
||||
|
||||
@pytest.mark.skipif(not LIVE_URL, reason="FFB_TEST_EMBED_URL is not set")
|
||||
async def test_live_gateway():
|
||||
"""FFB_TEST_EMBED_URL=https://.../v1 FFB_TEST_EMBED_KEY=... [FFB_TEST_EMBED_MODEL]."""
|
||||
e = OpenAIEmbedder(
|
||||
LIVE_URL,
|
||||
api_key=os.environ.get("FFB_TEST_EMBED_KEY"),
|
||||
model=os.environ.get("FFB_TEST_EMBED_MODEL", "mxbai-embed-large"),
|
||||
dim=int(os.environ.get("FFB_TEST_EMBED_DIM", "1024")),
|
||||
)
|
||||
a, b, c = await e.embed(
|
||||
["render timed out", "the render took too long and timed out", "dark mode"]
|
||||
)
|
||||
assert len(a) == e.dim
|
||||
assert cosine_score(a, b) > cosine_score(a, c)
|
||||
@@ -448,6 +448,7 @@ class TestDatabaseSink:
|
||||
assert build_metadata("ffb_", host) is host
|
||||
build_metadata("ffb_", host)
|
||||
assert sorted(host.tables) == [
|
||||
"ffb_embeddings",
|
||||
"ffb_events",
|
||||
"ffb_feedback_call_links",
|
||||
"ffb_tool_calls",
|
||||
|
||||
@@ -57,7 +57,7 @@ def make_app(sinks=None, **kwargs):
|
||||
|
||||
@app.tool
|
||||
def render(status: str) -> dict:
|
||||
return {"status": status, "error": f"render said {status}"}
|
||||
return {"status": status, "detail": f"render said {status}"}
|
||||
|
||||
@app.tool
|
||||
def crash() -> dict:
|
||||
|
||||
@@ -59,7 +59,7 @@ def make_app(sinks=None, **kwargs):
|
||||
|
||||
@app.tool
|
||||
def render(status: str) -> dict:
|
||||
return {"status": status, "error": f"render said {status}"}
|
||||
return {"status": status, "detail": f"render said {status}"}
|
||||
|
||||
@app.tool
|
||||
def respond(payload: dict) -> dict:
|
||||
@@ -119,6 +119,19 @@ class TestClassifier:
|
||||
{"is_error": True, "content": [{"type": "text", "text": "a"}]},
|
||||
"isError: a",
|
||||
),
|
||||
# mcspeak's failures: an error string and no status or flag key.
|
||||
(
|
||||
{"played": False, "error": "audio output stalled after 30s"},
|
||||
"error: audio output stalled after 30s",
|
||||
),
|
||||
(
|
||||
{"error": "Not listening: microphone stopped"},
|
||||
"error: Not listening: microphone stopped",
|
||||
),
|
||||
({"error": " padded "}, "error: padded"),
|
||||
({"status": "unhealthy", "error": "db gone"}, "error: db gone"),
|
||||
# Earlier checks still win, with their own message.
|
||||
({"success": False, "error": "x"}, "success=False: x"),
|
||||
],
|
||||
)
|
||||
def test_flagged(self, result, expected):
|
||||
@@ -135,6 +148,16 @@ class TestClassifier:
|
||||
{"result": False},
|
||||
{"state": 3},
|
||||
{"isError": False},
|
||||
{"played": True, "error": None},
|
||||
{"error": ""},
|
||||
{"error": " "},
|
||||
{"error": 0.02},
|
||||
{"error": 0},
|
||||
{"error": False},
|
||||
# A structured error is left to custom classifiers.
|
||||
{"error": {"code": 1}},
|
||||
{"error": ["a"]},
|
||||
{"failure_reason": "x"},
|
||||
[{"status": "failed"}],
|
||||
"failed",
|
||||
None,
|
||||
@@ -159,6 +182,88 @@ class TestClassifier:
|
||||
assert classify("t", {"passed": False}) == "passed=False"
|
||||
assert classify("t", {"status": "failed", "ok": False}) is None
|
||||
|
||||
def test_unknown_status_with_error_string(self):
|
||||
"""blender-mcp's shape: flagged by default, labelled by status when known."""
|
||||
result = {"status": "unknown_target", "error": "no object named Cube"}
|
||||
assert DEFAULT_RESULT_CLASSIFIER("t", result) == "error: no object named Cube"
|
||||
classify = soft_error_classifier(extra_statuses=("unknown_target",))
|
||||
assert classify("t", result) == "status=unknown_target: no object named Cube"
|
||||
|
||||
def test_custom_error_keys(self):
|
||||
classify = soft_error_classifier(error_keys=("failure_reason",))
|
||||
assert classify("t", {"failure_reason": "no gpu"}) == "failure_reason: no gpu"
|
||||
assert classify("t", {"error": "ignored"}) is None
|
||||
assert classify("t", {"failure_reason": ""}) is None
|
||||
|
||||
def test_empty_error_keys_disable(self):
|
||||
classify = soft_error_classifier(error_keys=())
|
||||
assert classify("t", {"played": False, "error": "stalled"}) is None
|
||||
assert classify("t", {"ok": False, "error": "x"}) == "ok=False: x"
|
||||
|
||||
|
||||
class TestExplicitSuccess:
|
||||
"""A flag key that is exactly True turns an error string into a warning."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"result",
|
||||
[
|
||||
{"ok": True, "error": "deprecated arg"},
|
||||
{"success": True, "error": "x"},
|
||||
{"ok": True, "status": "unknown_target", "error": "fell back"},
|
||||
],
|
||||
)
|
||||
def test_true_flag_with_error_string_is_ok(self, result):
|
||||
assert DEFAULT_RESULT_CLASSIFIER("t", result) is None
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"result, expected",
|
||||
[
|
||||
({"ok": False, "error": "x"}, "ok=False: x"),
|
||||
# Not exactly True, so not an explicit success.
|
||||
({"ok": 1, "error": "x"}, "error: x"),
|
||||
({"ok": "yes", "error": "x"}, "error: x"),
|
||||
({"success": 1.0, "error": "x"}, "error: x"),
|
||||
# No flag key: mcspeak's shape, as in 2026.10.01.
|
||||
(
|
||||
{"played": False, "error": "audio output stalled"},
|
||||
"error: audio output stalled",
|
||||
),
|
||||
# A False flag beats a True one.
|
||||
({"ok": True, "success": False, "error": "x"}, "success=False: x"),
|
||||
# A failing status is more specific than the flag, so it still flags.
|
||||
({"success": True, "status": "failed"}, "status=failed"),
|
||||
({"ok": True, "state": "error", "error": "x"}, "state=error: x"),
|
||||
({"ok": True, "isError": True, "content": ["boom"]}, "isError: boom"),
|
||||
],
|
||||
)
|
||||
def test_still_flagged(self, result, expected):
|
||||
assert DEFAULT_RESULT_CLASSIFIER("t", result) == expected
|
||||
|
||||
def test_custom_flag_keys(self):
|
||||
classify = soft_error_classifier(flag_keys=("passed",))
|
||||
assert classify("t", {"passed": True, "error": "slow"}) is None
|
||||
assert classify("t", {"passed": False, "error": "x"}) == "passed=False: x"
|
||||
# ok is not a flag key here, so it does not suppress the error check.
|
||||
assert classify("t", {"ok": True, "error": "x"}) == "error: x"
|
||||
|
||||
def test_no_flag_keys(self):
|
||||
classify = soft_error_classifier(flag_keys=())
|
||||
assert classify("t", {"ok": True, "error": "x"}) == "error: x"
|
||||
|
||||
async def test_recorded_as_ok(self):
|
||||
record = await record_for(
|
||||
"respond", {"payload": {"ok": True, "error": "deprecated arg"}}
|
||||
)
|
||||
assert (record.outcome, record.ok, record.error_message) == ("ok", True, None)
|
||||
|
||||
async def test_truthy_flag_recorded_as_soft_error(self):
|
||||
record = await record_for("respond", {"payload": {"ok": 1, "error": "x"}})
|
||||
assert (record.outcome, record.ok, record.error_message) == (
|
||||
"soft_error",
|
||||
False,
|
||||
"error: x",
|
||||
)
|
||||
|
||||
|
||||
class TestMiddleware:
|
||||
async def test_status_failed_is_soft_error(self):
|
||||
@@ -186,6 +291,33 @@ class TestMiddleware:
|
||||
assert REDACTED in record.error_message
|
||||
assert len(record.error_message) == 60
|
||||
|
||||
async def test_error_string_is_soft_error(self):
|
||||
record = await record_for(
|
||||
"respond",
|
||||
{"payload": {"played": False, "error": f"stalled, key {BMCP_TOKEN}"}},
|
||||
)
|
||||
assert (record.outcome, record.ok, record.error_type) == (
|
||||
"soft_error",
|
||||
False,
|
||||
"SoftError",
|
||||
)
|
||||
assert record.error_message == f"error: stalled, key {REDACTED}"
|
||||
|
||||
@pytest.mark.parametrize("value", [None, "", 0.02])
|
||||
async def test_empty_or_numeric_error_stays_ok(self, value):
|
||||
record = await record_for(
|
||||
"respond", {"payload": {"played": True, "error": value}}
|
||||
)
|
||||
assert (record.outcome, record.ok, record.error_message) == ("ok", True, None)
|
||||
|
||||
async def test_error_keys_disabled(self):
|
||||
record = await record_for(
|
||||
"respond",
|
||||
{"payload": {"error": "stalled"}},
|
||||
result_classifier=soft_error_classifier(error_keys=()),
|
||||
)
|
||||
assert (record.outcome, record.ok) == ("ok", True)
|
||||
|
||||
async def test_extra_statuses(self):
|
||||
record = await record_for(
|
||||
"render",
|
||||
@@ -393,3 +525,27 @@ class TestStorageAndCorrelation:
|
||||
assert record.error_message == (
|
||||
f"success=False: database down, password={REDACTED}"
|
||||
)
|
||||
|
||||
async def test_list_feedback_failure_is_soft_error(self, monkeypatch):
|
||||
"""list_feedback reports failure as {"feedback": [], "error": ...}."""
|
||||
app = FastMCP("fb")
|
||||
sink = MemorySink()
|
||||
mw = instrument(app, [sink])
|
||||
add_feedback_tools(app, database_url="sqlite:///:memory:", instrumentation=mw)
|
||||
|
||||
def broken(*args, **kwargs):
|
||||
raise RuntimeError("database down")
|
||||
|
||||
from fastmcp_feedback.database import FeedbackDatabase
|
||||
|
||||
monkeypatch.setattr(FeedbackDatabase, "session", broken)
|
||||
async with Client(app) as client:
|
||||
result = await client.call_tool("list_feedback", {})
|
||||
await mw.flush()
|
||||
assert result.data["error"] == "database down"
|
||||
[record] = sink.records
|
||||
assert (record.tool, record.outcome, record.error_message) == (
|
||||
"list_feedback",
|
||||
"soft_error",
|
||||
"error: database down",
|
||||
)
|
||||
|
||||
@@ -601,7 +601,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "fastmcp-feedback"
|
||||
version = "2026.9.29"
|
||||
version = "2026.10.1.1"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "fastmcp" },
|
||||
@@ -638,6 +638,9 @@ dev = [
|
||||
{ name = "ruff" },
|
||||
{ name = "sqlalchemy", extra = ["asyncio"] },
|
||||
]
|
||||
embeddings = [
|
||||
{ name = "pgvector" },
|
||||
]
|
||||
instrumentation = [
|
||||
{ name = "aiosqlite" },
|
||||
{ name = "sqlalchemy", extra = ["asyncio"] },
|
||||
@@ -671,6 +674,7 @@ requires-dist = [
|
||||
{ name = "fastmcp-feedback", extras = ["postgresql", "mysql", "dev"], marker = "extra == 'all'" },
|
||||
{ name = "mypy", marker = "extra == 'dev'", specifier = ">=1.7.0" },
|
||||
{ name = "pgvector", marker = "extra == 'dev'", specifier = ">=0.5.0" },
|
||||
{ name = "pgvector", marker = "extra == 'embeddings'", specifier = ">=0.5.0" },
|
||||
{ name = "pgvector", marker = "extra == 'test'", specifier = ">=0.5.0" },
|
||||
{ name = "psycopg2-binary", marker = "extra == 'postgresql'", specifier = ">=2.9.0" },
|
||||
{ name = "pydantic", specifier = ">=2.11.7" },
|
||||
@@ -691,7 +695,7 @@ requires-dist = [
|
||||
{ name = "sqlalchemy", extras = ["asyncio"], marker = "extra == 'instrumentation'", specifier = ">=2.0.43" },
|
||||
{ name = "sqlalchemy", extras = ["asyncio"], marker = "extra == 'test'", specifier = ">=2.0.43" },
|
||||
]
|
||||
provides-extras = ["postgresql", "instrumentation", "mysql", "dev", "test", "all"]
|
||||
provides-extras = ["postgresql", "instrumentation", "embeddings", "mysql", "dev", "test", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "fastmcp-slim"
|
||||
|
||||
Reference in New Issue
Block a user