Compare commits
453 commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d2af1a2e62 | ||
|
|
5cecd7b50e | ||
|
|
4b1ec096c6 | ||
|
|
95fdb46c3a | ||
|
|
bae359e3bf | ||
|
|
d7c131bc28 | ||
|
|
110adf23ee | ||
|
|
b374d5eb83 | ||
|
|
0026142e7d | ||
|
|
2582f2c05a | ||
|
|
494294046f | ||
|
|
eb9995e9b7 | ||
|
|
f45ed90b33 | ||
|
|
5653e876a8 | ||
|
|
15afb97a6e | ||
|
|
2a6d72171d | ||
|
|
88cd2b1d6c | ||
|
|
40d40bcbf2 | ||
|
|
ce69a8c989 | ||
|
|
7eeddd1a7c | ||
|
|
cf674b93ce | ||
|
|
4290111333 | ||
|
|
a18bfd9781 | ||
|
|
7f7223e0ac | ||
|
|
660c86f991 | ||
|
|
376dcb2e6f | ||
|
|
4bf4699ee0 | ||
|
|
880551deba | ||
|
|
9c56f56660 | ||
|
|
f4c657e094 | ||
|
|
5bf4cb50aa | ||
|
|
75547f1cbd | ||
|
|
187946a7ae | ||
|
|
3814b2bc78 | ||
|
|
a04a16c717 | ||
|
|
34180a0fd1 | ||
|
|
9baa213b34 | ||
|
|
4f7f69aaf0 | ||
|
|
31fc8aba29 | ||
|
|
d9f489dcc8 | ||
|
|
ddac328d05 | ||
|
|
f63dec5f75 | ||
|
|
2bb169867c | ||
|
|
2d8b2cb9ed | ||
|
|
f7e9535d7b | ||
|
|
7be71ebbac | ||
|
|
82fe91bb76 | ||
|
|
d06277408b | ||
|
|
f927998643 | ||
|
|
c0cf612d68 | ||
|
|
a7efd76016 | ||
|
|
b0818f4d3d | ||
|
|
e49f4264d3 | ||
|
|
02bccb1b90 | ||
|
|
35f8d721b1 | ||
|
|
61df6b65a6 | ||
|
|
2eaf2c6aea | ||
|
|
d2b2acef2c | ||
|
|
07421a7dd6 | ||
|
|
b93270adce | ||
|
|
22a5b7d17c | ||
|
|
b87ae16910 | ||
|
|
e0bde9d59a | ||
|
|
6133e5ff4a | ||
|
|
a1738c9d3e | ||
|
|
eb50c9c587 | ||
|
|
0c67db4459 | ||
|
|
ad100ecd4d | ||
|
|
d0c66eef8f | ||
|
|
afdef92b5b | ||
|
|
09a7076b7e | ||
|
|
d0df382ce8 | ||
|
|
ee835b77c2 | ||
|
|
b601489896 | ||
|
|
2bfb661c10 | ||
|
|
688d7e708e | ||
|
|
81e05cba7a | ||
|
|
1de1b662ae | ||
|
|
7dc853606f | ||
|
|
590719fca7 | ||
|
|
2325f187b5 | ||
|
|
fa319596cc | ||
|
|
bfd09f6880 | ||
|
|
225a37ca73 | ||
|
|
9dc79e7fda | ||
|
|
dafc15978e | ||
|
|
5e81054a1e | ||
|
|
bee67e736a | ||
|
|
249c51c65a | ||
|
|
b5aa0e7122 | ||
|
|
bd8c1a6d15 | ||
|
|
202cbd2d3f | ||
|
|
cad472e999 | ||
|
|
4d7fdd2d26 | ||
|
|
0dec3d4e63 | ||
|
|
95c02addd7 | ||
|
|
5ecffdedf2 | ||
|
|
18c22f1ddb | ||
|
|
b6bc54c69e | ||
|
|
cb9f945c79 | ||
|
|
105b2628de | ||
|
|
2fef67d9fd | ||
|
|
8791f12c37 | ||
|
|
64b1096865 | ||
|
|
de6da8b375 | ||
|
|
5f795cb0cb | ||
|
|
acea8cbaac | ||
|
|
c778fbf527 | ||
|
|
db7f0e0af6 | ||
|
|
d93cb5c891 | ||
|
|
45a5e68da5 | ||
|
|
474245ab59 | ||
|
|
9ed66e3a06 | ||
|
|
87367759f9 | ||
|
|
206d29b74a | ||
|
|
b8bf846bb7 | ||
|
|
8a7cfb949f | ||
|
|
0d7810c78a | ||
|
|
746b663a2f | ||
|
|
c8a9b5df50 | ||
|
|
9210642f65 | ||
|
|
cfb880e8a0 | ||
|
|
b2c617f9f2 | ||
|
|
dfa5f8027e | ||
|
|
0e0ce8d3f0 | ||
|
|
ea6b864f6e | ||
|
|
8ed24d0e24 | ||
|
|
1995a7da36 | ||
|
|
b76c5c7db0 | ||
|
|
ab0bdff14c | ||
|
|
db5f61d740 | ||
|
|
8f41d42ab2 | ||
|
|
ca2e28559e | ||
|
|
4c2bfc4fc1 | ||
|
|
1b3334b5af | ||
|
|
503db3271c | ||
|
|
2000098e16 | ||
|
|
2b2de5a61f | ||
|
|
de8075fcd1 | ||
|
|
59e05d2c22 | ||
|
|
71e4e4a40a | ||
|
|
5b420a9b23 | ||
|
|
49580228c2 | ||
|
|
9a3e2b86ca | ||
|
|
d9ac221ca0 | ||
|
|
05c204071d | ||
|
|
d367b1eb5a | ||
|
|
8db447e095 | ||
|
|
16b7319c48 | ||
|
|
028217b7c0 | ||
|
|
e8390ca747 | ||
|
|
79a63a46d4 | ||
|
|
fcfa4aefd8 | ||
|
|
73d04ddba5 | ||
|
|
9fefcdb629 | ||
|
|
bddb32b469 | ||
|
|
a9e66b001b | ||
|
|
c284d86885 | ||
|
|
f35be76604 | ||
|
|
307e250f29 | ||
|
|
16e0add64b | ||
|
|
c3182d0fb6 | ||
|
|
bc04dfdb7a | ||
|
|
e94623ec37 | ||
|
|
1d09b4e31b | ||
|
|
e1fd68e8cb | ||
|
|
2bbd949a29 | ||
|
|
eed820df1b | ||
|
|
fa242c1c94 | ||
|
|
9a17ff7457 | ||
|
|
bb92ccf67e | ||
|
|
5c9639bbbb | ||
|
|
95e38e7137 | ||
|
|
e15f73a387 | ||
|
|
fdb5710491 | ||
|
|
f33b789a31 | ||
|
|
397b553528 | ||
|
|
569947b28d | ||
|
|
735489d723 | ||
|
|
ab910eb776 | ||
|
|
4df31ff16e | ||
|
|
bc6c23cbf3 | ||
|
|
5e10846292 | ||
|
|
a57a73b6d0 | ||
|
|
aa49ec6cb3 | ||
|
|
c00ccb7322 | ||
|
|
152d57d6e2 | ||
|
|
86f254d5f2 | ||
|
|
959cf700ae | ||
|
|
2310b7a8b3 | ||
|
|
e9f6fea598 | ||
|
|
d7d27cc1a3 | ||
|
|
61a756da7f | ||
|
|
971ec0a5b0 | ||
|
|
def24a2434 | ||
|
|
da6cdfbc51 | ||
|
|
88925e86d9 | ||
|
|
97d9c6ed49 | ||
|
|
2a7c6510d0 | ||
|
|
e687c73906 | ||
|
|
4967765878 | ||
|
|
476f8d07a0 | ||
|
|
02b6104b4b | ||
|
|
483c0ec354 | ||
|
|
d6cfda22b5 | ||
|
|
0592206f16 | ||
|
|
19d5b2e7f6 | ||
|
|
d429ac4996 | ||
|
|
7f54eb4dbc | ||
|
|
3dcd463792 | ||
|
|
629e1ba4ea | ||
|
|
7f3590e881 | ||
|
|
d51356f476 | ||
|
|
ab19f78507 | ||
|
|
a896ac9eec | ||
|
|
d1691d3942 | ||
|
|
9d64a0d9f0 | ||
|
|
cab510d0d5 | ||
|
|
78b3231dcd | ||
|
|
48a94f7793 | ||
|
|
ed5519b38d | ||
|
|
058a820e14 | ||
|
|
72ef18e39d | ||
|
|
6533c28377 | ||
|
|
e8f00fcff4 | ||
|
|
7465026b6a | ||
|
|
2da294b850 | ||
|
|
489b8a65f0 | ||
|
|
15868762b0 | ||
|
|
73944f91ea | ||
|
|
895b5f0532 | ||
|
|
2402f324ba | ||
|
|
13bd69b908 | ||
|
|
62da6086b8 | ||
|
|
28217fcf82 | ||
|
|
65f6d72cd5 | ||
|
|
460215158d | ||
|
|
5b0444043a | ||
|
|
af6a6b0bbe | ||
|
|
d177b5883b | ||
|
|
cdd1b99e6b | ||
|
|
4eb72d6a49 | ||
|
|
adcf27f650 | ||
|
|
0882dc9fed | ||
|
|
8a4d488a72 | ||
|
|
da207da106 | ||
|
|
6f976ef2a9 | ||
|
|
c5a8e5571d | ||
|
|
198d9d7b73 | ||
|
|
4bcfb9cc3b | ||
|
|
d28e2662e5 | ||
|
|
bd7946178d | ||
|
|
2839a4434a | ||
|
|
b221b4ac03 | ||
|
|
47feeb32ee | ||
|
|
a5c4647005 | ||
|
|
11644c7f43 | ||
|
|
3a9ded2218 | ||
|
|
09ff101349 | ||
|
|
8e93b639bc | ||
|
|
c184a25d68 | ||
|
|
980b1985ab | ||
|
|
73dfc62f33 | ||
|
|
587ba75a62 | ||
|
|
3b9d6ae2c6 | ||
|
|
0634964e64 | ||
|
|
bf94efc655 | ||
|
|
f1d5918d43 | ||
|
|
db2b8fb883 | ||
|
|
73d9d93db9 | ||
|
|
cc04f92f28 | ||
|
|
b07426a883 | ||
|
|
a75d89122e | ||
|
|
9b2ae347d2 | ||
|
|
90e53b7348 | ||
|
|
92cab7c846 | ||
|
|
db6996b629 | ||
|
|
a68cb23b6e | ||
|
|
93c21272d1 | ||
|
|
981353eb98 | ||
|
|
e522d8cfa8 | ||
|
|
d864cf8008 | ||
|
|
c5ecf718d9 | ||
|
|
d1e3b828f1 | ||
|
|
e1dd8517f9 | ||
|
|
53bbf52697 | ||
|
|
771ac9c96c | ||
|
|
21e261f608 | ||
|
|
3ec72363bc | ||
|
|
7bdd11db39 | ||
|
|
0a0ddbd3b6 | ||
|
|
cee0824744 | ||
|
|
eb9934a6e5 | ||
|
|
a79f348ac7 | ||
|
|
ae2755e88f | ||
|
|
0f38417c60 | ||
|
|
e1e7936d15 | ||
|
|
89b3a9186c | ||
|
|
f4d8e744a2 | ||
|
|
97a29d168f | ||
|
|
11b1bfbc94 | ||
|
|
a69f3a8a98 | ||
|
|
2cd568847e | ||
|
|
0aa6d79f88 | ||
|
|
42d923fe4b | ||
|
|
85594a1fd1 | ||
|
|
fa0c4a4f50 | ||
|
|
8eef0b3734 | ||
|
|
46f7ab8d97 | ||
|
|
d9bd3a701f | ||
|
|
5d76461b9f | ||
|
|
8762df5e75 | ||
|
|
bf744fc72a | ||
|
|
10fd89f547 | ||
|
|
af77f836d5 | ||
|
|
6186dae83f | ||
|
|
614be8b2a7 | ||
|
|
d4d6414bc4 | ||
|
|
5e6284fd02 | ||
|
|
432c81149c | ||
|
|
d63199d96d | ||
|
|
96475bba8c | ||
|
|
57531f1653 | ||
|
|
ba963864f3 | ||
|
|
605367ee5a | ||
|
|
eb11a165b8 | ||
|
|
ac9b2cbf81 | ||
|
|
788a2fe731 | ||
|
|
a4839c0ad2 | ||
|
|
a2afdddc63 | ||
|
|
292835190a | ||
|
|
5a023a24ad | ||
|
|
c1cb9fe6c6 | ||
|
|
8b45f81464 | ||
|
|
adae9eff2e | ||
|
|
a9afb176d6 | ||
|
|
c137305468 | ||
|
|
27ed0b3bb3 | ||
|
|
b62d6e920b | ||
|
|
721acbcf38 | ||
|
|
68238f0b6a | ||
|
|
485a8f901e | ||
|
|
a528ab912f | ||
|
|
5752f61f2c | ||
|
|
685a7c393d | ||
|
|
749752820c | ||
|
|
2ac2594d0d | ||
|
|
e2b930d469 | ||
|
|
8e7e494fd1 | ||
|
|
ebd1fc7f3e | ||
|
|
de102b7e34 | ||
|
|
d1ad14f44b | ||
|
|
221c90af72 | ||
|
|
d9e4d1e57c | ||
|
|
94befc0dfd | ||
|
|
522959d9b4 | ||
|
|
feaca386d3 | ||
|
|
5c8df37af1 | ||
|
|
53084d6fdd | ||
|
|
88651a42d2 | ||
|
|
47015758a0 | ||
|
|
044da7ae99 | ||
|
|
f13a3fb677 | ||
|
|
ae345cc39f | ||
|
|
9168bc15ef | ||
|
|
eb09dc6d2c | ||
|
|
04ad3823b9 | ||
|
|
e186720a5e | ||
|
|
ab88188529 | ||
|
|
5200b71e2b | ||
|
|
a602e0fbf5 | ||
|
|
f96a428ef1 | ||
|
|
7ba78fdde3 | ||
|
|
f6acb65e95 | ||
|
|
7c120587f0 | ||
|
|
2afe1bd28c | ||
|
|
dc5ad8f699 | ||
|
|
7c76c3399e | ||
|
|
1e8e5e9f6f | ||
|
|
dddaf0f84c | ||
|
|
e2b273ee2a | ||
|
|
94e03d777a | ||
|
|
1e76e66b28 | ||
|
|
c716be9640 | ||
|
|
081856c6b1 | ||
|
|
cdaeaa94e6 | ||
|
|
93824cdee5 | ||
|
|
5f4c73f89f | ||
|
|
c62fd78c7a | ||
|
|
4c5050d161 | ||
|
|
acbd66afbd | ||
|
|
98a9c763ba | ||
|
|
55351a8829 | ||
|
|
629cf665a6 | ||
|
|
c2c38dd08b | ||
|
|
f87fabe556 | ||
|
|
6c5bc0aae1 | ||
|
|
597808c56e | ||
|
|
b8dcb066dc | ||
|
|
175929f23a | ||
|
|
77585ef26a | ||
|
|
43c17a6777 | ||
|
|
22ffd5e92e | ||
|
|
9e7db72997 | ||
|
|
9deb1f2bd4 | ||
|
|
ec78641cc4 | ||
|
|
a72c94aa77 | ||
|
|
71737095ea | ||
|
|
35b8b046ce | ||
|
|
f31060e741 | ||
|
|
543aba7547 | ||
|
|
29ccb9a035 | ||
|
|
0308220769 | ||
|
|
9f49bdde6e | ||
|
|
1afd615a60 | ||
|
|
2c5849b5d2 | ||
|
|
d9cc9a4eea | ||
|
|
c9466af986 | ||
|
|
0de434ab24 | ||
|
|
294ced57cd | ||
|
|
3cb591f60b | ||
|
|
75a21c82c2 | ||
|
|
4f0fb7322e | ||
|
|
82693f7974 | ||
|
|
8bf569ef89 | ||
|
|
7e4fecc13e | ||
|
|
3034280c78 | ||
|
|
cb3e30b66a | ||
|
|
3dd234ecfc | ||
|
|
11303aa714 | ||
|
|
0e8fa551f3 | ||
|
|
b5c0ea630c | ||
|
|
aa7a4ff91a | ||
|
|
1ed424b50d | ||
|
|
f02cde5ddf | ||
|
|
dd2817ff6d | ||
|
|
1f88944ada | ||
|
|
ce1b9ac88e | ||
|
|
5453c00c95 | ||
|
|
4522fbdf1b | ||
|
|
f41ec379df | ||
|
|
a1f3435df3 | ||
|
|
f997d7adc9 | ||
|
|
0e271eaf4b | ||
|
|
0813a1c980 | ||
|
|
0cbde6b7a2 | ||
|
|
69319c9390 | ||
|
|
ae603d9b4d | ||
|
|
aa620aeb60 | ||
|
|
472eb089fe | ||
|
|
ee38b3db42 | ||
|
|
5d5eee54bf | ||
|
|
d2f5ecfc58 |
404 changed files with 64699 additions and 95495 deletions
20
.agents/plugins/marketplace.json
Normal file
20
.agents/plugins/marketplace.json
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
{
|
||||
"name": "haiku-rag",
|
||||
"interface": {
|
||||
"displayName": "haiku.rag"
|
||||
},
|
||||
"plugins": [
|
||||
{
|
||||
"name": "haiku-rag",
|
||||
"source": {
|
||||
"source": "local",
|
||||
"path": "./plugins/haiku-rag"
|
||||
},
|
||||
"policy": {
|
||||
"installation": "AVAILABLE",
|
||||
"authentication": "ON_INSTALL"
|
||||
},
|
||||
"category": "Productivity"
|
||||
}
|
||||
]
|
||||
}
|
||||
14
.claude-plugin/marketplace.json
Normal file
14
.claude-plugin/marketplace.json
Normal file
|
|
@ -0,0 +1,14 @@
|
|||
{
|
||||
"name": "haiku-rag",
|
||||
"description": "The haiku.rag knowledge base as Claude Code tools and a skill.",
|
||||
"owner": {
|
||||
"name": "Yiorgis Gozadinos"
|
||||
},
|
||||
"plugins": [
|
||||
{
|
||||
"name": "haiku-rag",
|
||||
"source": "./plugins/haiku-rag",
|
||||
"description": "Search, read and analyze your haiku.rag knowledge base from Claude Code."
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
---
|
||||
name: debug-evals
|
||||
description: Debug haiku.rag evaluation runs in Logfire. Use when asked to look at Logfire for an eval run, find failing or low-scoring eval cases, compare runs, check citation quality (cited_map) or judge pass rate (answer_equivalent), or explain why an eval case failed. Drives the Logfire MCP against the `evals` service.
|
||||
description: Debug haiku.rag evaluation runs in Logfire. Use when asked to look at Logfire for an eval run, find failing or low-scoring eval cases, compare runs, check citation quality (cited_map) or judge pass rate (answer_equivalent), or explain why an eval case failed. Drives the Logfire MCP against the `evals` service, or the Logfire HTTP query API when the MCP is not loaded. Also covers monitoring a run that is still in flight.
|
||||
---
|
||||
|
||||
# Debug eval runs in Logfire
|
||||
|
|
@ -11,13 +11,82 @@ single case. Read-only.
|
|||
|
||||
## How to query
|
||||
|
||||
1. Confirm the current schema with `mcp__logfire__schema_reference` (spans and
|
||||
logs share the `records` table).
|
||||
2. Run SQL with `mcp__logfire__arbitrary_query` (`query` + `age` in minutes, max
|
||||
30 days). The same SQL works pasted into Logfire's Explore UI.
|
||||
1. Confirm the current schema with `mcp__logfire__query_schema_reference` (spans
|
||||
and logs share the `records` table).
|
||||
2. Run SQL with `mcp__logfire__query_run` (`query` + `project: "haiku"` +
|
||||
`start_timestamp`/`end_timestamp`, max 14 days). The remote MCP is
|
||||
org-scoped, so `project` is required; eval runs land in project `haiku`.
|
||||
The same SQL works pasted into Logfire's Explore UI.
|
||||
3. Read span attributes as JSON: `attributes->>'key'`, nested as
|
||||
`attributes->'a'->'b'->>'c'`. Cast when needed: `(...)::float`, `(...)::int`.
|
||||
4. Hand back a clickable trace with `mcp__logfire__logfire_link(trace_id)`.
|
||||
4. Hand back a clickable trace with
|
||||
`mcp__logfire__project_logfire_link(trace_id, project="haiku")`.
|
||||
|
||||
## When the Logfire MCP is not available
|
||||
|
||||
The `mcp__logfire__*` tools are not loaded in every session. The HTTP query API is
|
||||
the fallback and needs no MCP:
|
||||
|
||||
```
|
||||
POST https://logfire-eu.pydantic.dev/v2/query # EU projects
|
||||
POST https://logfire-us.pydantic.dev/v2/query # US projects
|
||||
Authorization: Bearer <api-key>
|
||||
Content-Type: application/json
|
||||
{"sql": "...", "min_timestamp": "2026-08-23T00:00:00Z"}
|
||||
```
|
||||
|
||||
- **`min_timestamp` is mandatory** and silently bounds every result. Too recent a
|
||||
value is indistinguishable from "no data".
|
||||
- Read tokens are being replaced by **API keys** (`pylf_v2_<region>_...`). Same
|
||||
`Authorization: Bearer` header; the region is in the prefix.
|
||||
- Keep the key in `~/.logfire-read-key` (mode 600) and read it from there so it never
|
||||
lands in a transcript. Helper next to this skill: `.claude/skills/debug-evals/lf-query.sh "<SQL>" [min_ts]`.
|
||||
- The API is **project-scoped**. A key for the wrong project authenticates fine and
|
||||
returns zero rows — it does not error. Diagnose in this order:
|
||||
1. wrong region → `HTTP 401 Invalid read token` on the other host;
|
||||
2. wrong project → auth succeeds, `count(*)` over months is 0;
|
||||
3. right project → `SELECT service_name, count(*) ... GROUP BY service_name` shows
|
||||
`evals`, `haiku-rag`, `haiku-ingester`.
|
||||
Eval runs live in project **`haiku`**. There is an empty project named `evals`,
|
||||
which is the natural wrong guess.
|
||||
- **The API caps returned rows and does not say so.** Aggregate server-side
|
||||
(`count(*)`, `avg(...)`, `sum(CASE WHEN ...)`) rather than pulling rows and counting
|
||||
them locally. A pass rate computed from a clipped page is wrong and looks fine.
|
||||
|
||||
### JSON access via the HTTP API
|
||||
|
||||
`assertions`, `scores`, `metrics` and `case_name` are **keys inside the `attributes`
|
||||
column, not columns** — `SELECT assertions` fails with `column not found`. Both of
|
||||
these work:
|
||||
|
||||
```sql
|
||||
attributes->'assertions'->'answer_equivalent' -- returns JSON
|
||||
json_get_bool(attributes,'assertions','answer_equivalent','value')
|
||||
json_get_float(attributes,'scores','cited_map','value')
|
||||
json_get_int(attributes,'attributes','n_searches')
|
||||
json_get_str(attributes,'attributes','citation_status')
|
||||
```
|
||||
|
||||
Prefer the `json_get_*` form inside aggregates — it yields a typed value, so no cast
|
||||
is needed and `sum(CASE WHEN ...)` behaves.
|
||||
|
||||
## Counting cases, not spans
|
||||
|
||||
**One exception appears once per span level.** A single failing case emits the same
|
||||
`exception_type` on `case: {case_name}`, `execute {task}` and `invoke_agent agent`
|
||||
(and often `chat {model}`), so a raw count over-reports by 3-4x. Always add
|
||||
`AND span_name='case: {case_name}'` when counting failures. Cross-check that the
|
||||
number equals the count of unjudged cases.
|
||||
|
||||
## Always report floor as well as judged
|
||||
|
||||
`assertion_pass_rate` **excludes unjudged cases from the denominator**, so cases that
|
||||
died produce no verdict and silently inflate the headline. Report both:
|
||||
|
||||
- judged rate = passed / (cases - unjudged)
|
||||
- floor = passed / cases
|
||||
|
||||
A run with 4.5% deaths reads ~3pp better than it is. Quote them together, always.
|
||||
|
||||
## Vocabulary
|
||||
|
||||
|
|
@ -28,7 +97,7 @@ A run is one experiment span; its cases are direct children sharing its
|
|||
- `attributes->>'name'` — run label (the `--name` arg, or `{dataset}_qa_evaluation` / `{dataset}_retrieval_evaluation`).
|
||||
- `attributes->>'dataset_name'` — dataset.
|
||||
- `(attributes->>'assertion_pass_rate')::float` — overall judge pass rate (QA runs).
|
||||
- `attributes->'logfire.experiment.metadata'->'metadata'` — run config: `target` (`rag-skill`|`analysis-skill`), `qa_model`, `embedder_model`, `chunk_size`, `search_limit`, `rerank_model`, `judge_model`, `skill_model`, etc.
|
||||
- `attributes->'logfire.experiment.metadata'->'metadata'` — run config: `target` (`rag-capability`|`analysis-capability`), `qa_model`, `embedder_model`, `chunk_size`, `search_limit`, `rerank_model`, `judge_model`, `qa_max_searches`, etc.
|
||||
- `trace_id` — scopes the whole run.
|
||||
- Case span: `span_name = 'case: {case_name}'` (scope `pydantic-evals`).
|
||||
- `message` — `case: <id>`.
|
||||
|
|
@ -36,8 +105,8 @@ A run is one experiment span; its cases are direct children sharing its
|
|||
- `attributes->'scores'->'cited_map'->>'value'` — citation average precision (0..1).
|
||||
- `attributes->'scores'->'number_match'->>'value'` — numeric-answer match (datasets that use it).
|
||||
- `duration` — task time in seconds.
|
||||
- Inside each case the skill under test emits agent spans (scope `pydantic-ai`):
|
||||
`execute {task}`, `agent run`, `running tool`, `chat {model}`.
|
||||
- Inside each case the capability under test emits agent spans (scope `pydantic-ai`):
|
||||
`execute {task}`, `invoke_agent agent`, `execute_tool {tool_name}`, `chat {model}`.
|
||||
|
||||
The service is `evals` regardless of model, so filter on `service_name = 'evals'`
|
||||
first. `otel_scope_name` separates the layers (`pydantic-evals` for run/case,
|
||||
|
|
@ -137,8 +206,8 @@ ORDER BY start_timestamp;
|
|||
mean_task_seconds — always report task time).
|
||||
3. Pull failing and low-citation cases, read the judge `reason`.
|
||||
4. To understand one case, take its start/end from the case query and run the
|
||||
agent-activity query, then `logfire_link(trace_id)` so the user can expand
|
||||
that case in the UI.
|
||||
agent-activity query, then `project_logfire_link(trace_id)` so the user can
|
||||
expand that case in the UI.
|
||||
|
||||
## When a query returns nothing
|
||||
|
||||
|
|
@ -149,3 +218,78 @@ SELECT DISTINCT otel_scope_name, span_name
|
|||
FROM records WHERE service_name='evals'
|
||||
ORDER BY 1,2;
|
||||
```
|
||||
|
||||
## Monitoring a run that is still in flight
|
||||
|
||||
An eval prints nothing until it finishes, so a live run's only progress signal is its
|
||||
case spans. Everything below works mid-run.
|
||||
|
||||
Progress and ETA:
|
||||
|
||||
```sql
|
||||
SELECT count(*) AS cases_done,
|
||||
min(start_timestamp) AS first_case,
|
||||
max(start_timestamp) AS latest_case,
|
||||
avg(duration) AS avg_case_s
|
||||
FROM records
|
||||
WHERE service_name='evals' AND span_name='case: {case_name}'
|
||||
AND start_timestamp > '<RUN_LAUNCH_TS>';
|
||||
```
|
||||
|
||||
**Cases run serially**, so `ETA_total = total_cases * avg_case_s`. Verify rather than
|
||||
assume: wall-clock per case (`latest_case - first_case` over `cases_done`) should equal
|
||||
`avg_case_s`. If it does, concurrency is 1 and the multiplication is valid. Concurrent
|
||||
requests seen on the model endpoint (`vllm:num_requests_running` > 1) are parallel
|
||||
searches *within* one case, not parallel cases.
|
||||
|
||||
**Never estimate a run's length from a `--limit N` smoke.** Its "avg task time per
|
||||
case" is a per-case duration on the easiest N cases of a deterministic prefix; the full
|
||||
set ran 32% slower (72.7s vs 55.2s) on FRAMES. Smokes validate wiring, not wall-clock.
|
||||
|
||||
Live headline, behaviour and failure composition in one pass:
|
||||
|
||||
```sql
|
||||
SELECT count(*) AS cases,
|
||||
sum(CASE WHEN json_get_bool(attributes,'assertions','answer_equivalent','value')
|
||||
THEN 1 ELSE 0 END) AS passed,
|
||||
sum(CASE WHEN json_get(attributes,'assertions','answer_equivalent') IS NULL
|
||||
THEN 1 ELSE 0 END) AS unjudged,
|
||||
avg(json_get_float(attributes,'scores','cited_map','value')) AS cited_map,
|
||||
avg(json_get_int(attributes,'attributes','n_requests')) AS req_per_case,
|
||||
avg(json_get_int(attributes,'attributes','n_searches')) AS searches,
|
||||
avg(json_get_int(attributes,'attributes','n_executions')) AS execs,
|
||||
sum(json_get_int(attributes,'attributes','n_rejected_searches')) AS rejected,
|
||||
sum(CASE WHEN json_get_str(attributes,'attributes','citation_status')='grounded'
|
||||
THEN 1 ELSE 0 END) AS grounded
|
||||
FROM records
|
||||
WHERE service_name='evals' AND span_name='case: {case_name}'
|
||||
AND start_timestamp > '<RUN_LAUNCH_TS>';
|
||||
```
|
||||
|
||||
Why a case died, counted correctly:
|
||||
|
||||
```sql
|
||||
SELECT sum(CASE WHEN exception_message LIKE '%token limit%' THEN 1 ELSE 0 END) AS token_limit,
|
||||
sum(CASE WHEN exception_message NOT LIKE '%token limit%' THEN 1 ELSE 0 END) AS other,
|
||||
count(*) AS dead_cases
|
||||
FROM records
|
||||
WHERE service_name='evals' AND is_exception
|
||||
AND span_name='case: {case_name}'
|
||||
AND start_timestamp > '<RUN_LAUNCH_TS>';
|
||||
```
|
||||
|
||||
`ToolFailedError` is mostly **not** a defect — an exhausted search or code budget
|
||||
reports failure to the model on purpose. `UnexpectedModelBehavior` is the one that
|
||||
kills a case.
|
||||
|
||||
### The per-case diagnostic attributes
|
||||
|
||||
`attributes->'attributes'` on a case span carries what the capability actually did:
|
||||
`n_requests`, `n_searches`, `n_search_calls`, `n_rejected_searches`, `n_failed_tools`,
|
||||
`n_executions`, `cited_uris`, `cited_chunk_ids`, `searched_uris`, `citation_status`
|
||||
(`grounded` | `missing` | `ungrounded`). `attributes->'metrics'` carries `requests`,
|
||||
`input_tokens`, `output_tokens` for the case.
|
||||
|
||||
Cite rate comes from `citation_status`, not from `cited_map` — they answer different
|
||||
questions, and conflating them has produced wrong claims before. And never steer on
|
||||
raw cite rate: it is confounded by task success, so measure it among *correct* answers.
|
||||
|
|
|
|||
25
.claude/skills/debug-evals/lf-query.sh
Executable file
25
.claude/skills/debug-evals/lf-query.sh
Executable file
|
|
@ -0,0 +1,25 @@
|
|||
#!/usr/bin/env bash
|
||||
# Query the Logfire API. Key is read from ~/.logfire-read-key and never echoed.
|
||||
# usage: lf-query.sh "<SQL>" [min_timestamp]
|
||||
set -uo pipefail
|
||||
KEY_FILE="$HOME/.logfire-read-key"
|
||||
[ -r "$KEY_FILE" ] || { echo "missing $KEY_FILE"; exit 2; }
|
||||
SQL="${1:?need SQL}"
|
||||
MIN="${2:-$(date -u -v-2d +%Y-%m-%dT%H:%M:%SZ 2>/dev/null || date -u -d '2 days ago' +%Y-%m-%dT%H:%M:%SZ)}"
|
||||
python3 - "$SQL" "$MIN" <<'PY'
|
||||
import json, os, sys, urllib.request, urllib.error
|
||||
sql, min_ts = sys.argv[1], sys.argv[2]
|
||||
key = open(os.path.expanduser("~/.logfire-read-key")).read().strip()
|
||||
# region comes from the key prefix: pylf_v2_<region>_...
|
||||
parts = key.split("_")
|
||||
region = parts[2] if len(parts) > 3 and parts[0] == "pylf" else "eu"
|
||||
req = urllib.request.Request(
|
||||
f"https://logfire-{region}.pydantic.dev/v2/query",
|
||||
data=json.dumps({"sql": sql, "min_timestamp": min_ts}).encode(),
|
||||
headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
|
||||
)
|
||||
try:
|
||||
print(json.dumps(json.load(urllib.request.urlopen(req, timeout=120)), indent=2)[:6000])
|
||||
except urllib.error.HTTPError as e:
|
||||
print(f"HTTP {e.code}: {e.read().decode()[:600]}")
|
||||
PY
|
||||
|
|
@ -12,15 +12,19 @@ that served a request. Read-only.
|
|||
|
||||
## How to query
|
||||
|
||||
1. Confirm the current schema with `mcp__logfire__schema_reference` (spans and
|
||||
logs share the `records` table).
|
||||
2. Run SQL with `mcp__logfire__arbitrary_query` (`query` + `age` in minutes, max
|
||||
30 days). The same SQL works pasted into Logfire's Explore UI.
|
||||
1. Confirm the current schema with `mcp__logfire__query_schema_reference` (spans
|
||||
and logs share the `records` table).
|
||||
2. Run SQL with `mcp__logfire__query_run` (`query` + `project: "haiku"` +
|
||||
`start_timestamp`/`end_timestamp`, max 14 days). The remote MCP is
|
||||
org-scoped, so `project` is required; the ingester ships to project `haiku`.
|
||||
The same SQL works pasted into Logfire's Explore UI.
|
||||
3. Read span attributes as JSON: `attributes->>'key'`, cast when needed
|
||||
(`(attributes->>'attempt')::int`).
|
||||
4. Hand back a clickable trace with `mcp__logfire__logfire_link(trace_id)`.
|
||||
5. For recent exceptions tied to a file, `mcp__logfire__find_exceptions_in_file`
|
||||
accepts `client/documents.py`, `ingester/workers/pool.py`, or
|
||||
4. Hand back a clickable trace with
|
||||
`mcp__logfire__project_logfire_link(trace_id, project="haiku")`.
|
||||
5. For recent exceptions tied to a file,
|
||||
`mcp__logfire__query_find_exceptions_in_file` accepts
|
||||
`client/documents.py`, `ingester/workers/pool.py`, or
|
||||
`ingester/pollers/base.py`.
|
||||
|
||||
Adjust `service_name` if the operator set `OTEL_SERVICE_NAME` (e.g. per tenant).
|
||||
|
|
@ -185,7 +189,8 @@ ORDER BY n DESC;
|
|||
`ingester.worker breaker opened` event flags a source the pool paused. The
|
||||
per-job dead/reschedule narration stays in the ingester console and the queue
|
||||
API, not Logfire.
|
||||
5. `logfire_link(trace_id)` for a document lets the user expand the full tree.
|
||||
5. `project_logfire_link(trace_id)` for a document lets the user expand the
|
||||
full tree.
|
||||
|
||||
## When a query returns nothing
|
||||
|
||||
|
|
|
|||
10
.github/workflows/build-docs.yml
vendored
10
.github/workflows/build-docs.yml
vendored
|
|
@ -3,27 +3,33 @@ on:
|
|||
push:
|
||||
branches:
|
||||
- main
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
permissions:
|
||||
contents: read
|
||||
pages: write
|
||||
id-token: write
|
||||
concurrency:
|
||||
group: pages
|
||||
group: pages-${{ github.ref }}
|
||||
cancel-in-progress: false
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: astral-sh/setup-uv@v5
|
||||
- uses: astral-sh/setup-uv@v9.0.0
|
||||
- run: uv sync --group dev
|
||||
- run: uv run zensical build
|
||||
- uses: actions/configure-pages@v5
|
||||
if: github.event_name == 'push'
|
||||
- uses: actions/upload-pages-artifact@v3
|
||||
if: github.event_name == 'push'
|
||||
with:
|
||||
path: ./site
|
||||
deploy:
|
||||
needs: build
|
||||
if: github.event_name == 'push'
|
||||
runs-on: ubuntu-latest
|
||||
environment:
|
||||
name: github-pages
|
||||
|
|
|
|||
2
.github/workflows/build-publish-slim.yml
vendored
2
.github/workflows/build-publish-slim.yml
vendored
|
|
@ -9,7 +9,7 @@ jobs:
|
|||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: astral-sh/setup-uv@v4
|
||||
- uses: astral-sh/setup-uv@v9.0.0
|
||||
with:
|
||||
enable-cache: true
|
||||
- name: Set up Python
|
||||
|
|
|
|||
2
.github/workflows/build-publish.yml
vendored
2
.github/workflows/build-publish.yml
vendored
|
|
@ -12,7 +12,7 @@ jobs:
|
|||
if: ${{ github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success' }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: astral-sh/setup-uv@v4
|
||||
- uses: astral-sh/setup-uv@v9.0.0
|
||||
with:
|
||||
enable-cache: true
|
||||
- name: Set up Python
|
||||
|
|
|
|||
31
.github/workflows/test.yml
vendored
31
.github/workflows/test.yml
vendored
|
|
@ -11,13 +11,13 @@ jobs:
|
|||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: astral-sh/setup-uv@v4
|
||||
- uses: astral-sh/setup-uv@v9.0.0
|
||||
with:
|
||||
enable-cache: true
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version-file: "pyproject.toml"
|
||||
python-version-file: ".python-version"
|
||||
- name: Install dependencies
|
||||
run: uv sync --all-extras
|
||||
- name: Lint
|
||||
|
|
@ -45,6 +45,25 @@ jobs:
|
|||
working-directory: app/frontend
|
||||
run: pnpm run check
|
||||
|
||||
test-uri-platforms:
|
||||
name: URI paths (${{ matrix.os }}, py${{ matrix.python }})
|
||||
runs-on: ${{ matrix.os }}
|
||||
strategy:
|
||||
matrix:
|
||||
os: [ubuntu-latest, macos-latest, windows-latest]
|
||||
python: ["3.13", "3.14"]
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: astral-sh/setup-uv@v9.0.0
|
||||
# --noconftest: tests/conftest.py imports the project's dependencies,
|
||||
# which this job deliberately does not install. test_uri.py is stdlib-only.
|
||||
- name: Test platform URI paths
|
||||
env:
|
||||
PYTHONPATH: haiku_rag_slim
|
||||
run: >
|
||||
uv run --no-project --python ${{ matrix.python }} --with pytest
|
||||
pytest tests/test_uri.py -q --noconftest -o addopts=
|
||||
|
||||
test:
|
||||
needs: [lint, lint-frontend]
|
||||
runs-on: ubuntu-latest
|
||||
|
|
@ -52,13 +71,13 @@ jobs:
|
|||
HF_TOKEN: ${{ secrets.HF_TOKEN }}
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: astral-sh/setup-uv@v4
|
||||
- uses: astral-sh/setup-uv@v9.0.0
|
||||
with:
|
||||
enable-cache: true
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version-file: "pyproject.toml"
|
||||
python-version-file: ".python-version"
|
||||
- name: Install dependencies
|
||||
run: uv sync --all-extras
|
||||
- name: Cache HuggingFace models
|
||||
|
|
@ -66,7 +85,7 @@ jobs:
|
|||
uses: actions/cache@v4
|
||||
with:
|
||||
path: ~/.cache/huggingface
|
||||
key: huggingface-${{ runner.os }}-test-models-v1
|
||||
key: huggingface-${{ runner.os }}-test-models-v2
|
||||
- name: Pre-download tokenizer
|
||||
if: steps.hf-cache.outputs.cache-hit != 'true'
|
||||
run: uv run python -c "from transformers import AutoTokenizer; AutoTokenizer.from_pretrained('Qwen/Qwen3-Embedding-0.6B')"
|
||||
|
|
@ -77,7 +96,7 @@ jobs:
|
|||
env:
|
||||
HF_HUB_OFFLINE: ${{ steps.hf-cache.outputs.cache-hit == 'true' && '1' || '0' }}
|
||||
TRANSFORMERS_OFFLINE: ${{ steps.hf-cache.outputs.cache-hit == 'true' && '1' || '0' }}
|
||||
run: uv run pytest -m "not integration" --cov=haiku --cov-report=xml
|
||||
run: uv run pytest -m "not integration" --cov --cov-report=xml --cov-report=term-missing:skip-covered
|
||||
- name: Upload coverage to Codecov
|
||||
uses: codecov/codecov-action@v5
|
||||
with:
|
||||
|
|
|
|||
522
CHANGELOG.md
522
CHANGELOG.md
|
|
@ -1,6 +1,502 @@
|
|||
# Changelog
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- Claude Code and Codex plugin under `plugins/haiku-rag/`: two client manifests
|
||||
sharing the server configuration and the `haiku-rag` Agent Skill.
|
||||
- MCP tool `execute_code(code, filter, sources)`: runs a program in the
|
||||
analysis sandbox over the selected documents and returns what it printed;
|
||||
one sandbox per call.
|
||||
- In the analysis sandbox, `search()` results carry `chunk_meta`,
|
||||
`list_documents()` rows and `metadata.json` carry the document `metadata`,
|
||||
and `/documents/{id}/chunks.jsonl` lists chunk ids with their metadata.
|
||||
`recovery_hint` in `haiku.rag.sandbox`.
|
||||
- MCP tools `get_document_outline` (heading tree with page numbers) and
|
||||
`get_document_section` (one section's text, subsections included), built
|
||||
on `document_items`. `build_toc` in `haiku.rag.context`.
|
||||
- MCP server `instructions`, `version`, and read-only `ToolAnnotations` on
|
||||
every tool; every parameter carries a description. `filter` on
|
||||
`search_documents` and `search_documents_by_image`. `DocumentInfo.metadata`.
|
||||
|
||||
### Changed
|
||||
|
||||
- `pydantic-monty>=0.0.23`. The analysis sandbox gains `collections`,
|
||||
`itertools`, `functools`, `dataclasses`, function decorators and
|
||||
`str.format`.
|
||||
- `fastmcp>=4.0.2,<5.0.0`, on MCP Python SDK 2. The MCP server answers both the
|
||||
session-based and the sessionless (2026-07-28) protocol.
|
||||
- Default models are `ollama:qwen3.8`: `ModelConfig`, `qa.model`,
|
||||
`processing.title_model` (was `ollama:gpt-oss`) and
|
||||
`processing.conversion_options.picture_description.model` (was
|
||||
`ollama:ministral-3`). Run `ollama pull qwen3.8`.
|
||||
- `qa.model.vision` defaults to `true`, matching `qwen3.8`. Set it `false` when
|
||||
pointing `qa.model` at a text-only model.
|
||||
- `enable_thinking` on `provider: ollama` maps to `reasoning_effort` for every
|
||||
model, not only `gpt-oss`: `false` sends `none`, `true` sends `high`.
|
||||
`gpt-oss` keeps `low` for `false`.
|
||||
- `processing.conversion_options.picture_description.model` defaults to
|
||||
`enable_thinking: false`, and the field now reaches the VLM: docling's
|
||||
picture-description request carries `reasoning_effort` in `params`.
|
||||
- MCP `search_documents` and `search_documents_by_image` expand results to
|
||||
their section (`HaikuRAG.expand_context`) and return the agent rendering
|
||||
as text (rank, `Document ID`, `Collection` over several databases, title,
|
||||
headings, the matched chunk's metadata, passage) and pictures as
|
||||
`ImageContent` blocks, with no structured content.
|
||||
`SearchResult.format_for_agent(include_document_id=, include_chunk_meta=)`;
|
||||
`collect_pictures` in `haiku.rag.tools.search`.
|
||||
- MCP tools raise on failure, with the error's message; an empty result no
|
||||
longer doubles as an error.
|
||||
- `haiku-rag mcp` covers the configured `lancedb.databases` set. `sources` on
|
||||
`search_documents`, `search_documents_by_image` and `execute_code`; `source`
|
||||
on `get_document`; an unknown name is a tool error. `DocumentInfo.source`.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `toc.json` `item_range` in the analysis sandbox is a line slice into
|
||||
`items.jsonl`, as documented; it held item positions.
|
||||
- Past `analysis.code_timeout` a sandbox program starts no further host call.
|
||||
Files served from memory and in-code `search()` / `list_documents()` were
|
||||
not checked against the deadline.
|
||||
|
||||
### Removed
|
||||
|
||||
- MCP tools `ask_question` and `analyze`.
|
||||
- `format_citations` in `haiku.rag.utils`; `format_citations_rich` stays.
|
||||
- MCP write tools `add_document_from_file`, `add_document_from_url`,
|
||||
`add_document_from_text` and `delete_document`. The server opens the
|
||||
database read-only; ingest with `haiku-rag add`, `add-src`, `delete` or
|
||||
`haiku-ingester`. `create_mcp_server` loses `read_only`.
|
||||
|
||||
## [0.82.1] - 2026-09-03
|
||||
|
||||
### Fixed
|
||||
|
||||
- Document timestamps include timezone information and satisfy the MCP
|
||||
date-time output schema. Stored timestamps without an offset are read as
|
||||
host-local time; a database written under a different host timezone shifts
|
||||
by that offset.
|
||||
|
||||
## [0.82.0] - 2026-09-03
|
||||
|
||||
### Removed
|
||||
|
||||
- `lancedb.uri`. Write `lancedb.databases: {NAME: <location>}`; a config carrying
|
||||
`uri` fails to load with that message. Configurations generated by
|
||||
`init-config` through 0.81 carry `uri: ""` and must drop the key.
|
||||
- `HAIKU_RAG_DB`. Capabilities cover the databases the configuration places, or
|
||||
the `db_path` argument.
|
||||
- `DB_PATH` in the `app/` backend and `examples/custom_agent_agui.py`. Both load
|
||||
the configuration as the CLI does (`HAIKU_RAG_CONFIG_PATH`, `./haiku.rag.yaml`,
|
||||
the platform directory); the compose files set
|
||||
`HAIKU_RAG_CONFIG_PATH=/app/haiku.rag.yaml` and mount `DB_VOLUME` at `/data`.
|
||||
|
||||
### Changed
|
||||
|
||||
- `qa.max_searches` counts search units: searches a model emits in one
|
||||
response share a unit, up to 3 per unit; sequential searches pay one unit
|
||||
each.
|
||||
- Searches in one model response deduplicate their results: evidence a sibling
|
||||
search already showed collapses to a reference line, and a picture attaches
|
||||
once per response.
|
||||
- Every database has a name: the key in `lancedb.databases`, or the path's stem
|
||||
for `--db PATH`, `db_path=` and the default database, which is the entry
|
||||
`haiku.rag` under `storage.data_dir` and selectable by that name.
|
||||
`SearchResult.source`, `Document.source` and `Citation.source` carry it on
|
||||
every value a database produces.
|
||||
- `db_path=` beside a configured `lancedb.databases` raises
|
||||
`AmbiguousDatabaseError` (`HaikuRAG`, `create_capability`, `create_mcp_server`,
|
||||
`Sandbox`). `haiku-rag --db PATH` and `haiku-ingester --db PATH` open that path
|
||||
whatever is configured.
|
||||
- `DatabaseRef(name, location, given)` replaces `DatabaseRef(name, uri, db_path)`;
|
||||
`DatabaseScope.at(path)` added; `locate_database` returns `Path | str`.
|
||||
`IngesterApp(config, scope)` takes a resolved scope in place of `db_path`.
|
||||
`SingleDatabaseSession(ref, config)` replaces `SingleDatabaseSession(db_path,
|
||||
config, source=)`. `Store.db_path` is `None` for a database behind a URI.
|
||||
- Opening a configured or default database that does not exist raises
|
||||
`SourceUnavailableError` naming the database and the remedy (`haiku-rag init`
|
||||
or `create=True`), where the default database raised `FileNotFoundError` with
|
||||
its path. A database given as a path still raises `FileNotFoundError`.
|
||||
- `Store(location, config)`, `connect_lancedb(location, config)`,
|
||||
`gather_database_info(location, config)` and `run_doctor(config, location, ...)`
|
||||
take the database location, a path or a URI. `ConnectionMode.of(location)`
|
||||
replaces `ConnectionMode.from_config`. `DatabaseRef.connection()` and
|
||||
`default_db_path` removed.
|
||||
|
||||
## [0.81.0] - 2026-09-01
|
||||
|
||||
### Added
|
||||
|
||||
- `search.vector_nprobes` (default 20) sets the IVF partitions each vector query searches.
|
||||
|
||||
### Documentation
|
||||
|
||||
- `docs/configuration/storage.md` "Vector Indexing" carries the measured with/without IVF_PQ retrieval comparison; `docs/benchmarks.md` states the published numbers are measured without a vector index.
|
||||
- Re-indexing note corrected: `optimize()` adds new chunks to an existing vector index, a rebuild retrains centroids.
|
||||
|
||||
### Fixed
|
||||
|
||||
- Cross-database fusion without a reranker orders the union by cosine
|
||||
similarity to the query, with within-database rank breaking ties, instead
|
||||
of round-robin by database declaration order. Full-text-only searches order
|
||||
by retrieval score. Fused results carry the ordering score.
|
||||
|
||||
## [0.80.0] - 2026-08-31
|
||||
|
||||
### Changed
|
||||
|
||||
- lancedb 0.37.1.
|
||||
- An invalid search `filter` raises `ValueError`.
|
||||
|
||||
### Added
|
||||
|
||||
- `--full-citations` on `haiku-rag ask` and `haiku-rag analyze` renders citation
|
||||
text untruncated. `format_citations_rich` takes a `full` argument.
|
||||
- `doctor` fails when the chunks FTS index covers no rows.
|
||||
- FTS and hybrid searches log a warning when the FTS index covers no rows.
|
||||
|
||||
### Removed
|
||||
|
||||
- `dot` removed from `search.vector_index_metric`; switch to `cosine` or `l2`
|
||||
and rerun `create-index`.
|
||||
|
||||
### Fixed
|
||||
|
||||
- FTS and hybrid search on a database whose FTS index covers no rows. Chunk
|
||||
writes now build the index and rebuild it if it covers none; `haiku-rag
|
||||
vacuum` also repairs it.
|
||||
- Migration to 0.38.0 no longer fails with `UnicodeDecodeError` on a
|
||||
`docling_document` blob written as zstd.
|
||||
- Vector search sets `search.vector_index_metric` on every query.
|
||||
|
||||
## [0.79.0] - 2026-08-28
|
||||
|
||||
### Changed
|
||||
|
||||
- Chat model provider `gemini` renamed to `google`, matching pydantic-ai. Update `provider: gemini` to `provider: google`.
|
||||
|
||||
### Added
|
||||
|
||||
- `lancedb.databases` configures a named set of local or remote databases.
|
||||
`search`, `ask` and `analyze` accept a `sources` subset; results, documents and
|
||||
citations carry the originating database in `source`, and model context names
|
||||
it as `Collection:` when a search spans more than one. Candidates are combined
|
||||
with the configured reranker, or reciprocal rank fusion when reranking is
|
||||
disabled. Operations that need one database raise `AmbiguousDatabaseError`, an
|
||||
unknown name raises `UnknownDatabaseError`, and a cited chunk ID retrieved from
|
||||
or previously cited from more than one selected database raises
|
||||
`AmbiguousCitationError`. An unavailable configured database raises
|
||||
`SourceUnavailableError`, which names the database and not its location.
|
||||
- `haiku-rag search`, `ask`, `analyze` and `chat` cover a configured set.
|
||||
Commands that access one database select it with `--db-name NAME` or
|
||||
`--db PATH`.
|
||||
- `HaikuRAG.aclose()` releases a client whatever it covers. `close()` remains
|
||||
limited to clients covering one database.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `haiku-rag settings` prints YAML.
|
||||
- The chat document filter pages results and lists the selected separately.
|
||||
Selection is by document ID and database, and a typed search applies on enter.
|
||||
- `haiku-rag list` prints only the fields a document has.
|
||||
- `haiku-rag` and `haiku-ingester` exit with a message on an embedder mismatch.
|
||||
- Capabilities created without a client honor `lancedb.uri`.
|
||||
- A `lancedb.uri` without a scheme is treated as a local path. `--db PATH`
|
||||
overrides it.
|
||||
- Inspector search results mark truncated previews with an ellipsis.
|
||||
- Document titles, URIs, headings and database names render as text, not Rich
|
||||
markup, in `search` output, chat citations and the chat document filter.
|
||||
- `reranking.model.base_url` accepts the endpoint with or without the `/v1` path, matching the `vllm` embedder. Writing `/v1` produced a request to `/v1/v1/rerank`.
|
||||
- An unrecognized chat model provider raises `Unknown model provider '<name>'` instead of reaching pydantic-ai as a `provider:name` string, and outranks the `api_key` check, so an unusable provider is no longer reported as a missing vendor environment variable.
|
||||
|
||||
## [0.78.0] - 2026-08-24
|
||||
|
||||
### Added
|
||||
|
||||
- `api_key` on model and embedding-model config, overriding the provider's environment variable. Honored on the `openai` and `ollama` providers, `vllm` embedders and rerankers, the picture-description VLM endpoint, and `doctor`'s endpoint probes; other providers raise.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `doctor`'s docling-serve probe sends `X-Api-Key`, so an instance requiring a key is reported reachable rather than unreachable.
|
||||
- The picture-description request to the public OpenAI endpoint sends `OPENAI_API_KEY`; it carried no authorization header.
|
||||
|
||||
## [0.77.0] - 2026-08-21
|
||||
|
||||
### Added
|
||||
|
||||
- The four capabilities support Pydantic AI agent specs via `from_spec`, registered with
|
||||
`Agent.from_spec(..., custom_capability_types=[...])`. `RAGCapability` and `AnalysisCapability` take
|
||||
`db_path`, `config`, `defer_loading`, `request_limit` and `vision`.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `Store`, `HaikuRAG` and `create_capability` coerce a string `db_path` to `Path`, as the documented
|
||||
`HaikuRAG("knowledge.lancedb")` and `rag(db_path="my.lancedb")` forms require.
|
||||
- A capability search that matches nothing returns `No results found.` instead of an empty string.
|
||||
- Repeating a search query within a question accumulates the results of both calls instead of replacing the earlier ones.
|
||||
- `file://` URIs resolve to a Windows path through `url2pathname`: `file:///C:/docs/a.pdf` was read as `\C:\docs\a.pdf`, so ingestion reported `File does not exist` for every discovered file. A URI authority is kept as a UNC server/share (`file://server/share/a.pdf`) except `localhost`, which is dropped.
|
||||
- A Windows drive letter is no longer read as a URI scheme, so `C:\docs\a.pdf` is accepted by `create_document_from_source`, `convert` and `resolve_adhoc_fetcher`.
|
||||
- `convert` and `check_source_accessible` decode percent-escapes in `file://` URIs, so a path containing `[`, `]` or a space resolves.
|
||||
|
||||
## [0.76.0] - 2026-08-20
|
||||
|
||||
### Added
|
||||
|
||||
- `haiku.rag.capabilities.EvidenceState`: the state base `RAGState` and `AnalysisState` derive from, with `begin_invocation()` for the per-question reset. `RAGCapabilityBase.evidence_record()` and `citation_index()` expose what a capability recorded, so a host reads it without reaching into `capability.state`.
|
||||
- The `haiku.rag` package declares the `jina` extra, so `provider: jina-local` is supported by declaration rather than through `cross-encoder`'s transitive `transformers` and `torch`. Raises the full package's torch floor to 2.0.
|
||||
- `providers.docling_serve.timeout` (default 300 seconds), forwarded to the docling-serve client's per-request timeout.
|
||||
|
||||
### Changed
|
||||
|
||||
- `haiku.rag.store.engine` split: table records, Arrow schemas, `index_specs`, `ensure_indexes`, `REQUIRED_TABLES` and `query_to_pydantic` are in `haiku.rag.store.schema`; `gather_database_info`, `get_database_stats`, `DatabaseInfo` and its result models are in `haiku.rag.store.info`. `Store` keeps lifecycle, locks, migration coordination, vacuuming and tags. Update imports.
|
||||
- Source adapters moved from `haiku.rag.ingester.sources` to `haiku.rag.sources`: `FetchResult`, `SourceEvent`, `SourceEventKind`, `RevisionSnapshot`, `Source` and the `FSSource`/`HTTPSource`/`S3Source`/`WebDAVSource` adapters. One-shot client ingestion uses them too, so they were never ingester-only. Update imports; the `haiku.rag.sources` plugin entry-point group is unchanged.
|
||||
- `HaikuRAG.convert(url)` fetches through `HTTPSource`, the same adapter the ingester uses, instead of its own httpx client. `_write_fetch_body` moved from `client.documents` to `client.processing`.
|
||||
- Chunk embedding is owned by the persistence funnels: `create_document`, `update_document` and source ingestion no longer embed eagerly before handing chunks to a check that would embed them anyway. The `document.embed` span moved onto `ensure_chunks_embedded`, so every path is instrumented rather than only ingest.
|
||||
- One-shot directory ingestion and `FSSource.discover` share `walk_files`, so the symlink-escape guard lives in one place.
|
||||
- Configuration sections reject unknown keys. A typo or a setting that has been renamed or removed now fails validation with its path (`providers.docling_serve.bogus: Extra inputs are not permitted`) instead of being silently ignored.
|
||||
- `processing.converter`, `processing.chunker` and `processing.chunker_type` are constrained to their supported values, so an unsupported one fails at load rather than at first use.
|
||||
- Numeric settings carry bounds: sizes, limits, dimensions, token budgets, attempt counts, breaker thresholds and `min_chunks` must be positive; retention, delays, intervals and cooldowns non-negative; `doctor.duplicates.similarity_threshold` within 0-1; `ingester.api.port` within 0-65535 (0 keeps its OS-assigned meaning); `ingester.workers.worker_count` allows 0 for an API-and-reaper-only process.
|
||||
- A configured reranker whose optional dependency is missing raises instead of silently disabling reranking, and names the extra to install (`uv pip install 'haiku.rag-slim[cohere]'`). A failure raised from inside an installed dependency propagates untouched rather than being reported as a missing package.
|
||||
- Multi-table writes (document create, update, batch import, cascade delete) go through `Store.write_transaction()`. Rollback restores in `RESTORE_TABLE_ORDER` and is shielded from cancellation, so a cancelled write rolls back instead of committing part of itself. `Store.restore_table_versions()` is removed.
|
||||
|
||||
### Removed
|
||||
|
||||
- `haiku.rag.config.Config`, the eagerly loaded configuration instance. Use `get_config()` for the current global config, or pass an `AppConfig`. Every internal default (`get_embedder`, `get_converter`, `get_chunker`, `get_reranker`, `embed_chunks`, `HaikuRAG`, `Store`, `HaikuRAGApp`, `create_mcp_server`) now takes `config: AppConfig | None = None` and resolves it per call, so `set_config()` reaches them. `RerankerBase._model` no longer defaults to the configured reranker name; `CohereReranker` takes its model name as an argument.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `haiku-rag settings` masked only top-level secret-named fields, printing nested ones in full (`lancedb.api_key`, `providers.docling_serve.api_key`, WebDAV source passwords). It redacts the whole dump.
|
||||
- `haiku-rag chat` reports a missing `tui` extra instead of raising `ImportError`, matching `haiku-rag inspect`.
|
||||
- `storage.data_dir: ""` resolves to the platform data directory, as documented; it coerced to `Path("")`, so the database was created in the process's working directory. Set `data_dir: .` to keep the old placement.
|
||||
- Documented `search.limit` default is 5, was stated as 10.
|
||||
- The documented way to disable reranking is omitting `reranking.model` or setting it to `null`; the previous `provider: ""` example raised `Unknown reranking provider`.
|
||||
- `prompts.picture_description: null` is not valid; the documented example sets a string or omits the key.
|
||||
- `create_document_from_source` closes the source adapter it builds for the call; adapters passed in through `sources` are left to their owner.
|
||||
- Directory ingestion skips symlinked files resolving outside the given directory, matching `FSSource.discover`.
|
||||
- A FULL rebuild no longer deletes a source-backed document before re-ingesting it: it refreshes the document in place, so the document id is preserved and a failed fetch or conversion falls back to rebuilding from stored content instead of losing the document.
|
||||
|
||||
## [0.75.0] - 2026-08-19
|
||||
|
||||
### Added
|
||||
|
||||
- `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata.
|
||||
- `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` / `mtrag_clapnq_live_uncompacted` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, live arms with and without `EvidenceCompactionCapability`, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, per-turn tool-traffic attributes, and `citation_status` / `turn_citation_status` eval attributes.
|
||||
- Raw chunk metadata is now exposed to search and citation results, through `SearchResult.chunk_meta` and `Citation.chunk_meta`. For context-expanded results, the metadata is that of the anchor chunk.
|
||||
- BTree indexes on `chunks.id`, `chunks.document_id` and `documents.id`, and a Bitmap index on `document_items.label`. Existing databases need `haiku-rag migrate`.
|
||||
- `lancedb.read_consistency_interval_seconds` (default 30), `lancedb.index_cache_size_bytes` and `lancedb.metadata_cache_size_bytes`. The LanceDB session is shared across connections in a process, so its index and metadata caches survive a connection being closed.
|
||||
|
||||
### Changed
|
||||
|
||||
- Search enrichment, the multimodal reranker's picture fetch, and context expansion each issue a fixed number of `document_items` queries regardless of how many documents a result set spans, instead of one set per document. `expand_with_items` takes the items to expand from rather than fetching them, and `DocumentItemRepository` gains `resolve_refs_grouped`, `get_items_in_ranges`, `get_pictures_grouped` and `get_caption_picture_refs_grouped`. The superseded methods are removed: `resolve_refs`, `get_items_in_range`, `get_caption_picture_refs`, `get_text_for_refs` and `get_all_items_grouped`. `get_pictures_grouped` returns each picture's text alongside its bytes under `with_text`, off by default so the reranker's blob fetch does not read a column it discards.
|
||||
- Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged.
|
||||
- `create_capability(rag=...)` lends a capability an open client rather than having it open its own; `client.ask`/`client.analyze` now pass theirs.
|
||||
- The MCP server opens one database client for its lifetime instead of one per tool call, and fails at startup if the database cannot be opened. `delete_document` no longer opens its own connection with `skip_validation=True`, so the server no longer opts out of embedding-config validation: drift that validation rejects (any `vector_dim` mismatch, or identity drift on a writable server) now fails MCP startup instead of serving a delete-only server. Use the CLI to delete under drift.
|
||||
- `import_documents` embeds chunks across the whole batch in one pass instead of per document.
|
||||
- `haiku.rag` and `haiku.rag-slim` summaries and keywords; both packages now publish `[project.urls]`.
|
||||
- `server.json` declares `title` and `websiteUrl`, and drops the `keywords` and `license` keys, which are not in the server schema.
|
||||
|
||||
### Removed
|
||||
|
||||
- `wix` evaluation dataset and its reference config `evaluations/configs/wix.yaml`.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `DocumentRepository.delete_all` recreated `document_items` from `DocumentItemRecord` instead of `get_document_items_arrow_schema()`, returning `picture_data` as `binary` rather than `large_binary`.
|
||||
- `server.json` runtime arguments are `mcp --stdio`, was `serve --mcp`.
|
||||
|
||||
## [0.74.0] - 2026-08-13
|
||||
|
||||
### Added
|
||||
|
||||
- `CitationPolicyCapability` (`haiku.rag.capabilities.policy.create_capability`): registering it requires every answer to declare its grounding, in any conversation that has something to declare — this question retrieved evidence, or something was cited earlier. A question that ends undeclared is sent back once to record what grounded the answer already given, and is recorded in `CitationPolicyState.violations` if it finishes undeclared regardless. A conversation with neither a current-question evidence outcome nor any earlier citation is not enforced.
|
||||
- `haiku.rag.capabilities.evidence.discover_evidence()` and `DiscoveredEvidence`, moved out of `compaction` so both optional capabilities share them. `RAGCapabilityBase.cite_available`.
|
||||
- `EvidenceCompactionCapability` (`haiku.rag.capabilities.compaction.create_capability`): registering it replaces earlier questions' evidence on the model request with the evidence that was cited, grouped by the question that cited it, cited page images re-attached, other earlier evidence returns reduced to a receipt. Requests only; `all_messages()` is untouched. No configuration.
|
||||
- `RAGState.evidence` / `AnalysisState.evidence` (`CapabilityEvidenceRecord`): which evidence a capability retrieved and cited, per question, keyed by message-count question identities and epochs, and whether that question is still being answered. `haiku.rag.capabilities.ledger.citation_status(records, question=...)` derives `missing` / `grounded` / `ungrounded` across capabilities.
|
||||
- `RAGCapabilityBase.evidence_tool_names()` and `get_picture_bytes()`.
|
||||
- `haiku.rag.tools.search.decode_picture()`.
|
||||
|
||||
### Changed
|
||||
|
||||
- `rag_cite` / `analysis_cite` accept an empty `chunk_ids`, recording the answer as ungrounded rather than failing the call, and the instructions no longer exempt a refusal or a corpus-level computation from citing.
|
||||
- `RAGCapability` and `AnalysisCapability` no longer rewrite the model request. Register `create_capability()` from `haiku.rag.capabilities.compaction` alongside them to keep earlier questions compacted.
|
||||
- Resuming a run (deferred tool results, an unfinished history tail) raises `RuntimeError` unless the host carries the capability state from the run being resumed.
|
||||
|
||||
### Removed
|
||||
|
||||
- `PRIOR_TURN_NOTICE` and `_compact_old_tool_returns` from `haiku.rag.capabilities._base`, and `RAGCapabilityBase.turn_start`.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `rag_cite` and `analysis_cite` no longer ask for another call when a call resolved some ids and not others.
|
||||
- `cross-encoder` reranking no longer ties the scores of strongly-relevant candidates, which left their order to the sort. Scores remain 0-1.
|
||||
- `haiku-rag` and `haiku-ingester` CLI startup no longer imports `lancedb`, `pyarrow` and `pydantic_ai`.
|
||||
- `haiku.rag.store` no longer re-exports `Store`; import it from `haiku.rag.store.engine`.
|
||||
- `docling-local` reuses one docling `DocumentConverter` per set of conversion options instead of building one per document, so local layout, table and OCR models are no longer loaded per document. Conversions through a shared converter are serialized.
|
||||
- A resumed run keeps the searches, citations and executions of the question in progress instead of clearing them.
|
||||
- Each page image attached to a search result is preceded by a line giving its position and the chunk id it came from. `build_binary_parts_from_results` is now `build_image_content_from_results` and returns those labels interleaved with the pictures.
|
||||
|
||||
## [0.73.0] - 2026-08-06
|
||||
|
||||
### Added
|
||||
|
||||
- `evaluations run` records `judge_extra_body`, `qa_extra_body` and `capability_extra_body` in experiment metadata.
|
||||
|
||||
### Changed
|
||||
|
||||
- `get_document_by_id` / `get_document_by_uri` no longer load the docling structure and page-image blobs. Load them with `DocumentRepository.get_docling_data` / `get_pages_data`, or `get_by_id(..., include_blobs=True)`.
|
||||
- Pinned judge sampling in every reference config under `evaluations/configs/` whose dataset is judged: `temperature` 0.6, `max_tokens` 16384, `extra_body` `top_p` 0.95 / `top_k` 20 / `min_p` 0 / `chat_template_kwargs.enable_thinking` true. `DEFAULT_JUDGE_MODEL` takes `temperature` 0.6, `max_tokens` 16384 and `top_p` 0.95, the subset ollama honours.
|
||||
- Reranker in the `orb_text`, `wix` and `t2_finqa` reference configs and in the reranking provider docs: `mixedbread-ai/mxbai-rerank-base-v2` → `Qwen/Qwen3-Reranker-4B` on the `vllm` provider.
|
||||
|
||||
### Fixed
|
||||
|
||||
- Ingester jobs failing with an `obstore` `PermissionDeniedError`, `UnauthenticatedError`, `UnknownConfigurationKeyError` or `InvalidPathError` are dead-lettered instead of retried to `max_attempts`.
|
||||
|
||||
## [0.72.1] - 2026-07-31
|
||||
|
||||
### Changed
|
||||
|
||||
- `list_documents(include_content=True)` returns `content` only; the docling structure and page-image blobs are no longer loaded by a listing.
|
||||
|
||||
## [0.72.0] - 2026-07-30
|
||||
|
||||
### Added
|
||||
|
||||
- `evaluations run` records `cited_chunk_ids`, `searched_uris`, `n_searches`, `n_search_calls`, `n_rejected_searches`, `n_failed_tools`, `n_executions` and `n_requests` as eval attributes alongside `cited_uris`.
|
||||
|
||||
### Changed
|
||||
|
||||
- A capability whose search or code-execution budget is spent says so in its instructions on every following request, naming the exhausted tools.
|
||||
|
||||
### Fixed
|
||||
|
||||
- A failed `analysis_execute_code` call that iterated a file object reports the `.readlines()` workaround alongside the `TypeError`.
|
||||
- A capability that reaches its request limit keeps its cite tool for two further requests that call one of its tools, while its other tools are removed.
|
||||
- A cited chunk id that does not match a retrieved id exactly resolves to the closest one above a 0.75 similarity cutoff.
|
||||
- Dotfiles are parsed as their actual format instead of a single unstructured text block. Docling ignores the extension of a name starting with a dot, so converters strip leading dots from the name they hand it.
|
||||
|
||||
### Documentation
|
||||
|
||||
- `qa.max_searches` is documented as defaulting to 5 in `docs/configuration/qa.md` and `docs/configuration/index.md`, was 3.
|
||||
- New "Vacuum Memory Requirements" subsection in `docs/configuration/storage.md` documenting the ~5x peak memory of vacuum compaction relative to the `documents` table, the mitigations, and upstream issue lancedb/lancedb#2325.
|
||||
|
||||
## [0.71.0] - 2026-07-29
|
||||
|
||||
### Added
|
||||
|
||||
- Analysis sandbox supports `open()` and `with` blocks for reading document files, including `.read()`, `.readline()`, and `.readlines()`.
|
||||
|
||||
### Changed
|
||||
|
||||
- Require `pydantic-ai-slim>=2.18,<3`.
|
||||
- `pydantic-monty` bumped to `>=0.0.19`; sandbox code now runs in a subprocess worker pool. A crashed or timed-out worker fails one call and the next call gets a replacement session.
|
||||
- The analysis sandbox gives Monty a duration budget of `analysis.code_timeout * analysis.max_executions` for the session, was `analysis.code_timeout`.
|
||||
- `enable_thinking` maps onto Pydantic AI's unified `thinking` setting for the `anthropic`, `gemini`, `groq` and `bedrock` providers. The Anthropic thinking budget is now Pydantic AI's default of 10000 tokens, was 4096, and `max_tokens` must exceed it on budget-based Claude models; `enable_thinking: false` disables Gemini thinking rather than only hiding thoughts; Groq maps reasoning effort rather than `groq_reasoning_format`; Bedrock Qwen with `enable_thinking: false` no longer sends `reasoning_config`, while Bedrock-served Claude keeps an explicit `thinking: disabled`. The `openai` and `ollama` providers still map to `openai_reasoning_effort`.
|
||||
- `provider: bedrock` with a proprietary OpenAI model such as `openai.o3-mini-v1:0` raises `UserError`: Bedrock Converse serves only the `gpt-oss` family. Use `provider: bedrock-mantle`.
|
||||
- Tool failures raise `pydantic_ai.ToolFailed` instead of returning failure text: search and code-execution limits, sandbox execution errors, and `get_document`/`summarize_document` misses.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `Store.set_haiku_version` stamps the store's own config into a recreated settings row instead of the process-global `Config`.
|
||||
- `check_source_accessible` returns `False` for a URI it cannot resolve (unparseable host, unreadable path) instead of raising and aborting a full rebuild.
|
||||
- `evaluations run` opens the database read-only outside the population phase, so an embedder identity differing from the stored one warns instead of aborting the run.
|
||||
- `docs/installation.md` documents the `jina`, `s3` and `ingester` extras, names the extras the full package actually pulls, drops the removed MixedBread AI reranker, and no longer lists Anthropic as a built-in provider.
|
||||
- `docs/tuning.md` no longer points at the removed `claim_timeout_s` setting.
|
||||
- `analysis.code_timeout` is enforced before each document read, bounding a call that reads in a loop.
|
||||
- `metadata.json` in the document VFS rejects writes, matching `content.txt`, `items.jsonl` and `toc.json`.
|
||||
|
||||
### Removed
|
||||
|
||||
- `SettingsRepository.create`, `get_by_id`, `update`, `delete`, `list_all` and `ChunkRepository.update`, `delete`, `get_chunks_in_range`.
|
||||
|
||||
## [0.70.0] - 2026-07-25
|
||||
|
||||
### Added
|
||||
|
||||
- `HaikuRAG.ask` and `HaikuRAG.analyze` accept `images: Sequence[bytes]`, attached to the question as model input; requires `vision: true` on the driving model.
|
||||
- `haiku-rag ask` and `haiku-rag analyze` accept `--image PATH` (repeatable).
|
||||
- MCP `ask_question` and `analyze` tools accept `images_base64`.
|
||||
- Chat TUI: `Ctrl+I` opens an image picker; attached images insert `[Image #N]` tokens in a multi-line prompt and are sent to the model with the message.
|
||||
|
||||
### Fixed
|
||||
|
||||
- Chat and inspector TUIs report the actual database-open error instead of an `AttributeError` from teardown.
|
||||
|
||||
## [0.69.0] - 2026-07-24
|
||||
|
||||
### Added
|
||||
|
||||
- Native deferred Pydantic AI `RAGCapability` and `AnalysisCapability` implementations under `haiku.rag.capabilities`, with namespaced host state and lazy per-run database and sandbox resources.
|
||||
- Prior-turn RAG and analysis tool results are compacted before model requests while current-turn evidence remains intact.
|
||||
- Per-question capability request limits force a final answer from gathered evidence by removing only the exhausted capability's tools; unrelated agent and capability tools remain available.
|
||||
|
||||
### Changed
|
||||
|
||||
- Require `pydantic-ai-slim>=2.11,<3`; the `vertexai` optional extra now installs Pydantic AI's `google` extra.
|
||||
- The chat TUI consumes native Pydantic AI stream events. The web example uses the standard `AGUIAdapter` and emits one final state snapshot instead of forwarding sub-agent activity and per-tool state events.
|
||||
- Chat capability selection is now `haiku-rag chat --capability/-c {rag,analysis}`. Migrate from `--skill/-s`.
|
||||
- Evaluation targets are now `rag-capability` and `analysis-capability`, and the model override is `--capability-model`. Migrate from `rag-skill`, `analysis-skill`, and `--skill-model`.
|
||||
|
||||
### Removed
|
||||
|
||||
- The `haiku.skills` dependency, `haiku.rag.skills` modules, Python entry-point discovery, and sub-agent execution layer. Migrate `create_skill(...)` plus `SkillToolset` usage to `haiku.rag.capabilities.*.create_capability(...)` passed through `Agent(capabilities=[...])`.
|
||||
- The `haiku-rag create-skill` package generator. Compose native capabilities directly and package application-specific instructions and data in the consuming project.
|
||||
- Legacy sub-agent `ActivitySnapshotEvent` plumbing and per-tool `StateDeltaEvent` generation.
|
||||
|
||||
## [0.68.0] - 2026-07-24
|
||||
|
||||
### Added
|
||||
|
||||
- `reranking.multimodal` config flag: picture chunks are sent to a vllm reranker as image documents.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `list`, `get`, and `visualize` no longer require an embeddings config matching the database.
|
||||
|
||||
## [0.67.3] - 2026-07-23
|
||||
|
||||
### Changed
|
||||
|
||||
- Security dependency updates: pillow 12.3.0, mcp 1.28.1, pydantic-ai 1.102.0, torch 2.13.0, transformers 5.5.0, soupsieve 2.9.1, joserfc 1.7.4, pyasn1 0.6.4, setuptools 83.0.0; frontend next 16.2.11, sharp 0.35.3, fast-uri 3.1.4, dompurify 3.4.12, body-parser 1.20.6.
|
||||
|
||||
## [0.67.2] - 2026-07-23
|
||||
|
||||
### Added
|
||||
|
||||
- `SearchResult.document_meta` and `Citation.document_meta` carry the parent document's metadata.
|
||||
|
||||
## [0.67.1] - 2026-07-22
|
||||
|
||||
### Added
|
||||
|
||||
- `hotpotqa` evaluation dataset.
|
||||
|
||||
### Changed
|
||||
|
||||
- `cite` skill tool names unresolvable chunk ids in its response when at least one id resolves.
|
||||
|
||||
### Fixed
|
||||
|
||||
- Document deletion no longer raises `ConfigMismatchError` on embedding config drift.
|
||||
|
||||
## [0.67.0] - 2026-07-16
|
||||
|
||||
### Added
|
||||
|
||||
- Database tags: `haiku-rag tag create/list/delete/restore`, tags shown in `history`. `tag restore` creates a `before-restore-*` safety tag before changing live state. Vacuum retains versions back to the oldest tag.
|
||||
|
||||
### Changed
|
||||
|
||||
- `lancedb` bumped to 0.34.0.
|
||||
|
||||
### Removed
|
||||
|
||||
- `--before` global flag and the `before` constructor arguments on `HaikuRAG`, `Store`, `HaikuRAGApp`, `ChatApp`/`run_chat`, and `InspectorApp`/`run_inspector`. There is no read-only replacement; create tags prospectively before important changes and use `tag restore` during a maintenance window.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `docling-local` text conversion no longer misroutes markdown/HTML content whose first bytes collide with a binary magic signature (e.g. `BM`, `ID3`) to an image or audio backend.
|
||||
|
||||
## [0.66.0] - 2026-07-14
|
||||
|
||||
### Changed
|
||||
|
|
@ -1937,7 +2433,31 @@ Existing documents without DoclingDocument data will work but won't have provena
|
|||
|
||||
- Initial version tracking
|
||||
|
||||
[Unreleased]: https://github.com/ggozad/haiku.rag/compare/0.66.0...HEAD
|
||||
[Unreleased]: https://github.com/ggozad/haiku.rag/compare/0.82.1...HEAD
|
||||
[0.82.1]: https://github.com/ggozad/haiku.rag/compare/0.82.0...0.82.1
|
||||
[0.82.0]: https://github.com/ggozad/haiku.rag/compare/0.81.0...0.82.0
|
||||
[0.81.0]: https://github.com/ggozad/haiku.rag/compare/0.80.0...0.81.0
|
||||
[0.80.0]: https://github.com/ggozad/haiku.rag/compare/0.79.0...0.80.0
|
||||
[0.79.0]: https://github.com/ggozad/haiku.rag/compare/0.78.0...0.79.0
|
||||
[0.78.0]: https://github.com/ggozad/haiku.rag/compare/0.77.0...0.78.0
|
||||
[0.77.0]: https://github.com/ggozad/haiku.rag/compare/0.77.0...0.77.0
|
||||
[0.77.0]: https://github.com/ggozad/haiku.rag/compare/0.76.0...0.77.0
|
||||
[0.76.0]: https://github.com/ggozad/haiku.rag/compare/0.75.0...0.76.0
|
||||
[0.75.0]: https://github.com/ggozad/haiku.rag/compare/0.74.0...0.75.0
|
||||
[0.74.0]: https://github.com/ggozad/haiku.rag/compare/0.73.0...0.74.0
|
||||
[0.73.0]: https://github.com/ggozad/haiku.rag/compare/0.72.1...0.73.0
|
||||
[0.72.1]: https://github.com/ggozad/haiku.rag/compare/0.72.0...0.72.1
|
||||
[0.72.0]: https://github.com/ggozad/haiku.rag/compare/0.71.0...0.72.0
|
||||
[0.71.0]: https://github.com/ggozad/haiku.rag/compare/0.70.0...0.71.0
|
||||
[0.70.0]: https://github.com/ggozad/haiku.rag/compare/0.69.0...0.70.0
|
||||
[0.69.0]: https://github.com/ggozad/haiku.rag/compare/0.68.0...0.69.0
|
||||
[0.68.0]: https://github.com/ggozad/haiku.rag/compare/0.67.3...0.68.0
|
||||
[0.67.3]: https://github.com/ggozad/haiku.rag/compare/0.67.2...0.67.3
|
||||
[0.67.2]: https://github.com/ggozad/haiku.rag/compare/0.67.1...0.67.2
|
||||
[0.67.1]: https://github.com/ggozad/haiku.rag/compare/0.67.0...0.67.1
|
||||
[0.67.0]: https://github.com/ggozad/haiku.rag/compare/0.67.0...0.67.0
|
||||
[0.67.0]: https://github.com/ggozad/haiku.rag/compare/0.67.0...0.67.0
|
||||
[0.67.0]: https://github.com/ggozad/haiku.rag/compare/0.66.0...0.67.0
|
||||
[0.66.0]: https://github.com/ggozad/haiku.rag/compare/0.65.1...0.66.0
|
||||
[0.65.1]: https://github.com/ggozad/haiku.rag/compare/0.65.0...0.65.1
|
||||
[0.65.0]: https://github.com/ggozad/haiku.rag/compare/0.64.0...0.65.0
|
||||
|
|
|
|||
46
README.md
46
README.md
|
|
@ -1,29 +1,36 @@
|
|||
# Haiku RAG
|
||||
# haiku.rag
|
||||
|
||||
[](https://pypi.org/project/haiku.rag/)
|
||||
[](https://pypi.org/project/haiku.rag/)
|
||||
[](https://pepy.tech/projects/haiku-rag-slim)
|
||||
[](https://ggozad.github.io/haiku.rag/)
|
||||
[](https://github.com/ggozad/haiku.rag/actions/workflows/test.yml)
|
||||
[](https://codecov.io/gh/ggozad/haiku.rag)
|
||||
|
||||
Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.pydantic.dev/), and [Docling](https://docling-project.github.io/docling/).
|
||||
Agentic RAG that answers questions about your own documents with citations to page numbers and section headings. Runs locally on an embedded database, no server required.
|
||||
|
||||
> **New: vision and multimodal search.** Picture-aware ingestion captures embedded figure bytes; vision-capable QA models receive them alongside text. Multimodal embedders put picture vectors in the same space as text, enabling text-as-query → figure hits and image-as-query retrieval.
|
||||
Built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.pydantic.dev/), and [Docling](https://docling-project.github.io/docling/). Full documentation at [ggozad.github.io/haiku.rag](https://ggozad.github.io/haiku.rag/).
|
||||
|
||||
## Features
|
||||
|
||||
- **Hybrid search** — Vector + full-text with Reciprocal Rank Fusion
|
||||
- **Multimodal & cross-modal search** — Multimodal embedders (vLLM, VoyageAI, Cohere) put picture vectors in the same space as text; supports text-as-query → figure hits and image-as-query
|
||||
- **Question answering** — RAG skill with citations (page numbers, section headings)
|
||||
- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text
|
||||
- **Question answering** — RAG capability with citations (page numbers, section headings)
|
||||
- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text; attach your own images to questions in `ask`, `analyze` and the chat TUI
|
||||
- **Reranking** — local cross-encoders, Cohere, Zero Entropy, or vLLM
|
||||
- **Analysis skill** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
|
||||
- **Analysis capability** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
|
||||
- **Evidence compaction** — Optional capability that replaces earlier questions' search results on the request with the evidence they cited, so long conversations stop resending everything they retrieved
|
||||
- **Citation policy** — Optional capability that requires every answer to declare what grounds it, including declaring that nothing does
|
||||
- **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory
|
||||
- **Document structure** — Stores full [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/), enabling structure-aware context expansion
|
||||
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, Cohere, LM Studio, vLLM (multimodal via `multimodal: true` on vLLM/VoyageAI/Cohere). QA: any model supported by Pydantic AI
|
||||
- **Multi-database search** — Search, ask, analyze, or chat across named databases with source attribution on results and citations
|
||||
- **Local-first** — Embedded LanceDB, no servers required. Also supports S3, GCS, Azure, and LanceDB Cloud
|
||||
- **CLI & Python API** — Full functionality from command line or code
|
||||
- **MCP server** — Expose as tools for AI assistants (Claude Desktop, etc.)
|
||||
- **Visual grounding** — View chunks highlighted on original page images
|
||||
- **Production ingester** — Long-lived `haiku-ingester` service with persistent SQLite queue, async worker pool with retries and a dead-letter queue, FS / HTTP / S3 / WebDAV source adapters, FastAPI control plane, and a browser dashboard for operators. See [docs/ingester.md](docs/ingester.md).
|
||||
- **Time travel** — Query the database at any historical point with `--before`
|
||||
- **Tags** — Name database states with `haiku-rag tag` and roll back to them
|
||||
- **Inspector** — TUI for browsing documents, chunks, and search results
|
||||
|
||||
## Installation
|
||||
|
|
@ -62,6 +69,9 @@ haiku-rag search "attention mechanism"
|
|||
# Ask questions with citations
|
||||
haiku-rag ask "What datasets were used for evaluation?"
|
||||
|
||||
# Ask about an image (vision-capable model)
|
||||
haiku-rag ask "Does this figure match the spec in the design doc?" --image figure.png
|
||||
|
||||
# Analyze — complex analytical tasks via code execution
|
||||
haiku-rag analyze "How many documents mention transformers?"
|
||||
|
||||
|
|
@ -96,16 +106,30 @@ async with HaikuRAG("knowledge.lancedb", create=True) as rag:
|
|||
print(f" [{cite.chunk_id}] p.{cite.page_numbers}: {cite.content[:80]}")
|
||||
```
|
||||
|
||||
For details on the skills the client wraps, see the [Skills docs](https://ggozad.github.io/haiku.rag/skills/).
|
||||
For direct agent composition, see the [capabilities documentation](https://ggozad.github.io/haiku.rag/capabilities/).
|
||||
|
||||
## MCP Server
|
||||
|
||||
Use with AI assistants like Claude Desktop:
|
||||
Use with AI assistants like Claude Code, Codex, and Claude Desktop:
|
||||
|
||||
```bash
|
||||
haiku-rag mcp --stdio
|
||||
```
|
||||
|
||||
In Claude Code, install the plugin, which registers the server and a skill:
|
||||
|
||||
```bash
|
||||
claude plugin marketplace add ggozad/haiku.rag
|
||||
claude plugin install haiku-rag
|
||||
```
|
||||
|
||||
In Codex, install the same plugin from its marketplace:
|
||||
|
||||
```bash
|
||||
codex plugin marketplace add ggozad/haiku.rag
|
||||
codex plugin add haiku-rag@haiku-rag
|
||||
```
|
||||
|
||||
Add to your Claude Desktop configuration:
|
||||
|
||||
```json
|
||||
|
|
@ -119,7 +143,7 @@ Add to your Claude Desktop configuration:
|
|||
}
|
||||
```
|
||||
|
||||
Provides tools for document management, search, QA, and analysis directly in your AI assistant.
|
||||
Provides search, document reading, and analysis tools directly in your AI assistant.
|
||||
|
||||
## Examples
|
||||
|
||||
|
|
@ -137,7 +161,7 @@ Full documentation at: https://ggozad.github.io/haiku.rag/
|
|||
- [Configuration](https://ggozad.github.io/haiku.rag/configuration/) - YAML reference
|
||||
- [CLI](https://ggozad.github.io/haiku.rag/cli/) - Command reference
|
||||
- [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs
|
||||
- [Skills](https://ggozad.github.io/haiku.rag/skills/) - The RAG and analysis skills the client wraps
|
||||
- [Capabilities](https://ggozad.github.io/haiku.rag/capabilities/) - Native Pydantic AI RAG and analysis capabilities
|
||||
- [Tuning](https://ggozad.github.io/haiku.rag/tuning/) - Retrieval and answer-quality tuning
|
||||
- [Ingester](https://ggozad.github.io/haiku.rag/ingester/) - Production ingester for continuous indexing from FS, HTTP, S3, and WebDAV
|
||||
- [MCP](https://ggozad.github.io/haiku.rag/mcp/) - Model Context Protocol integration
|
||||
|
|
|
|||
|
|
@ -2,8 +2,9 @@
|
|||
ANTHROPIC_API_KEY=your-anthropic-key
|
||||
OPENAI_API_KEY=your-openai-key
|
||||
|
||||
# Database path
|
||||
DB_PATH=/path/to/your/haiku.rag.lancedb
|
||||
# Host path of the LanceDB database, mounted at /data where haiku.rag.yaml
|
||||
# places it
|
||||
DB_VOLUME=./data/haiku.rag.lancedb
|
||||
|
||||
# Optional: Ollama base URL (if using local models)
|
||||
# Use host.docker.internal to reach Ollama running on the host machine
|
||||
|
|
|
|||
|
|
@ -40,7 +40,8 @@ A conversational RAG interface built with [CopilotKit](https://copilotkit.ai/) a
|
|||
|
||||
| Variable | Description | Required |
|
||||
|----------|-------------|----------|
|
||||
| `DB_PATH` | Path to your haiku.rag LanceDB database | Yes |
|
||||
| `DB_VOLUME` | Host path of the LanceDB database the compose files mount at `/data`, where `haiku.rag.yaml` places it (default `./data/haiku.rag.lancedb`) | No |
|
||||
| `HAIKU_RAG_CONFIG_PATH` | The configuration file; the compose files set it to the mounted `/app/haiku.rag.yaml` | No |
|
||||
| `ANTHROPIC_API_KEY` | Anthropic API key | One LLM key required |
|
||||
| `OPENAI_API_KEY` | OpenAI API key | One LLM key required |
|
||||
| `OLLAMA_BASE_URL` | Ollama server URL (default: `http://host.docker.internal:11434`) | For local models |
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from ag_ui.core import EventType, StateSnapshotEvent
|
||||
from dotenv import find_dotenv, load_dotenv
|
||||
|
|
@ -16,18 +16,18 @@ from starlette.requests import Request
|
|||
from starlette.responses import JSONResponse, Response, StreamingResponse
|
||||
from starlette.routing import Route
|
||||
|
||||
from haiku.rag.capabilities.compaction import (
|
||||
create_capability as create_compaction,
|
||||
)
|
||||
from haiku.rag.capabilities.policy import (
|
||||
create_capability as create_citation_policy,
|
||||
)
|
||||
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState, create_capability
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import load_yaml_config
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.skills.rag import create_skill, get_agent_preamble
|
||||
from haiku.rag.client.scope import DatabaseScope
|
||||
from haiku.rag.config import get_config
|
||||
from haiku.rag.telemetry import configure as configure_telemetry
|
||||
from haiku.rag.utils import get_model
|
||||
from haiku.skills import (
|
||||
SkillDeps,
|
||||
SkillToolset,
|
||||
run_agui_stream,
|
||||
)
|
||||
from haiku.skills.prompts import build_system_prompt
|
||||
|
||||
load_dotenv(find_dotenv(usecwd=True))
|
||||
|
||||
|
|
@ -38,20 +38,24 @@ logging.basicConfig(
|
|||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Load config
|
||||
config_path = Path("/app/haiku.rag.yaml")
|
||||
if config_path.exists():
|
||||
yaml_data = load_yaml_config(config_path)
|
||||
Config = AppConfig.model_validate(yaml_data)
|
||||
else:
|
||||
Config = AppConfig()
|
||||
# The configuration places the database. This app serves one.
|
||||
config = get_config()
|
||||
scope = DatabaseScope.resolve(config)
|
||||
if scope.covers_multiple:
|
||||
raise SystemExit(
|
||||
f"lancedb.databases names {', '.join(scope.names)}; this app serves one "
|
||||
"database: configure exactly one entry"
|
||||
)
|
||||
[database] = scope.databases
|
||||
|
||||
# Get DB path from environment
|
||||
db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
|
||||
db_path = Path(db_path_str)
|
||||
|
||||
logger.info(f"Database path: {db_path}")
|
||||
logger.info(f"QA Provider: {Config.qa.model.provider}, Model: {Config.qa.model.name}")
|
||||
def _database_exists() -> bool:
|
||||
"""A database behind a URI has no path to check."""
|
||||
return database.db_path is None or database.db_path.exists()
|
||||
|
||||
|
||||
logger.info(f"Database: {database.name} at {database.location}")
|
||||
logger.info(f"QA Provider: {config.qa.model.provider}, Model: {config.qa.model.name}")
|
||||
|
||||
# Only HaikuRAG client is a singleton (expensive to create)
|
||||
_client: HaikuRAG | None = None
|
||||
|
|
@ -69,23 +73,27 @@ async def get_client() -> HaikuRAG:
|
|||
if _client is None:
|
||||
async with _client_lock:
|
||||
if _client is None:
|
||||
client = HaikuRAG(db_path=db_path, config=Config, create=True)
|
||||
client = HaikuRAG(config=config, create=True)
|
||||
await client.__aenter__()
|
||||
_client = client
|
||||
return _client
|
||||
|
||||
|
||||
# Create skill, toolset, and agent
|
||||
skill = create_skill(db_path=db_path, config=Config)
|
||||
toolset = SkillToolset(skills=[skill])
|
||||
@dataclass
|
||||
class AppDeps:
|
||||
state: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
capability = create_capability(config=config, defer_loading=False)
|
||||
|
||||
agent = Agent(
|
||||
get_model(Config.qa.model, Config),
|
||||
instructions=build_system_prompt(
|
||||
toolset.skill_catalog, preamble=get_agent_preamble(Config)
|
||||
),
|
||||
toolsets=[toolset],
|
||||
deps_type=SkillDeps,
|
||||
get_model(config.qa.model, config),
|
||||
instructions=AGENT_PREAMBLE,
|
||||
# Conversations here are multi-turn, so earlier questions are reduced to the
|
||||
# evidence they cited rather than carried whole, and every answer declares
|
||||
# what grounds it so the UI can show citations for all of them.
|
||||
capabilities=[capability, create_compaction(), create_citation_policy()],
|
||||
deps_type=AppDeps,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -98,33 +106,21 @@ async def stream_chat(request: Request) -> Response:
|
|||
adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
|
||||
|
||||
incoming_state = run_input.state if isinstance(run_input.state, dict) else {}
|
||||
incoming_state.setdefault("rag", RAGState().model_dump(mode="json"))
|
||||
deps = AppDeps(state=incoming_state)
|
||||
|
||||
async def event_stream():
|
||||
async with run_agui_stream(
|
||||
adapter, toolset=toolset, deps=SkillDeps(state=incoming_state)
|
||||
) as stream:
|
||||
# Emit a STATE_SNAPSHOT after RUN_STARTED so the client holds every
|
||||
# namespace object before any STATE_DELTA patches into it. Without it,
|
||||
# the first `add /rag/<field>/...` fails against a missing parent.
|
||||
if incoming_state:
|
||||
toolset.restore_state_snapshot(incoming_state)
|
||||
snapshot = StateSnapshotEvent(
|
||||
type=EventType.STATE_SNAPSHOT,
|
||||
snapshot=toolset.build_state_snapshot(),
|
||||
)
|
||||
async def with_final_state():
|
||||
async for event in adapter.run_stream(deps=deps):
|
||||
if getattr(event, "type", None) == EventType.RUN_FINISHED:
|
||||
yield StateSnapshotEvent(
|
||||
type=EventType.STATE_SNAPSHOT,
|
||||
snapshot=deps.state,
|
||||
)
|
||||
yield event
|
||||
|
||||
async def with_state_snapshot():
|
||||
emitted = False
|
||||
async for event in stream:
|
||||
yield event
|
||||
if not emitted and getattr(event, "type", None) == (
|
||||
EventType.RUN_STARTED
|
||||
):
|
||||
yield snapshot
|
||||
emitted = True
|
||||
|
||||
async for chunk in adapter.encode_stream(with_state_snapshot()):
|
||||
yield chunk
|
||||
async for chunk in adapter.encode_stream(with_final_state()):
|
||||
yield chunk
|
||||
|
||||
return StreamingResponse(
|
||||
event_stream(),
|
||||
|
|
@ -142,17 +138,17 @@ async def health_check(_: Request) -> JSONResponse:
|
|||
return JSONResponse(
|
||||
{
|
||||
"status": "healthy",
|
||||
"qa_provider": Config.qa.model.provider,
|
||||
"qa_model": Config.qa.model.name,
|
||||
"db_path": str(db_path),
|
||||
"db_exists": db_path.exists(),
|
||||
"qa_provider": config.qa.model.provider,
|
||||
"qa_model": config.qa.model.name,
|
||||
"db_path": str(database.location),
|
||||
"db_exists": _database_exists(),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def list_documents(_: Request) -> JSONResponse:
|
||||
"""List all documents in the database."""
|
||||
if not db_path.exists():
|
||||
if not _database_exists():
|
||||
return JSONResponse({"documents": [], "error": "Database not found"})
|
||||
|
||||
client = await get_client()
|
||||
|
|
@ -168,17 +164,17 @@ async def list_documents(_: Request) -> JSONResponse:
|
|||
|
||||
async def db_info(_: Request) -> JSONResponse:
|
||||
"""Get database info and statistics."""
|
||||
if not db_path.exists():
|
||||
if not _database_exists():
|
||||
return JSONResponse(
|
||||
{
|
||||
"exists": False,
|
||||
"path": str(db_path),
|
||||
"path": str(database.location),
|
||||
"documents": 0,
|
||||
"chunks": 0,
|
||||
}
|
||||
)
|
||||
|
||||
from haiku.rag.store.engine import get_database_stats
|
||||
from haiku.rag.store.info import get_database_stats
|
||||
|
||||
client = await get_client()
|
||||
stats = await get_database_stats(client.store.db)
|
||||
|
|
@ -186,7 +182,7 @@ async def db_info(_: Request) -> JSONResponse:
|
|||
return JSONResponse(
|
||||
{
|
||||
"exists": True,
|
||||
"path": str(db_path),
|
||||
"path": str(database.location),
|
||||
"documents": stats["documents"].get("num_rows", 0),
|
||||
"chunks": stats["chunks"].get("num_rows", 0),
|
||||
"documents_bytes": stats["documents"].get("total_bytes", 0),
|
||||
|
|
@ -220,7 +216,7 @@ async def visualize_chunk(request: Request) -> JSONResponse:
|
|||
if isinstance(parsed, list):
|
||||
refs = [str(x) for x in parsed]
|
||||
|
||||
if not db_path.exists():
|
||||
if not _database_exists():
|
||||
return JSONResponse({"error": "Database not found"}, status_code=404)
|
||||
|
||||
client = await get_client()
|
||||
|
|
|
|||
|
|
@ -6,14 +6,14 @@ requires-python = ">=3.12"
|
|||
dependencies = [
|
||||
"starlette>=0.50.0",
|
||||
"uvicorn[standard]>=0.40.0",
|
||||
"pydantic-ai-slim[ag-ui,anthropic,openai]>=1.81.0",
|
||||
"pydantic-ai-slim[ag-ui,anthropic,openai]>=2.18.0,<3.0.0",
|
||||
"python-dotenv>=1.2.1",
|
||||
"haiku.rag-slim>=0.66.0",
|
||||
"haiku.rag-slim>=0.82.1",
|
||||
"logfire[pydantic-ai]>=3.17.0",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
dev = ["pyright>=1.1.407", "ruff>=0.14.10"]
|
||||
dev = ["ty>=0.0.28", "ruff>=0.14.10"]
|
||||
|
||||
[tool.hatch.metadata]
|
||||
allow-direct-references = true
|
||||
|
|
|
|||
|
|
@ -11,13 +11,14 @@ services:
|
|||
ports:
|
||||
- "127.0.0.1:8001:8000"
|
||||
environment:
|
||||
- DB_PATH=/data
|
||||
- HAIKU_RAG_CONFIG_PATH=/app/haiku.rag.yaml
|
||||
- ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY:-}
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY:-}
|
||||
- OLLAMA_BASE_URL=${OLLAMA_BASE_URL:-http://host.docker.internal:11434}
|
||||
- LOGFIRE_TOKEN=${LOGFIRE_TOKEN:-}
|
||||
volumes:
|
||||
- ${DB_PATH:-./data/haiku.rag.lancedb}:/data
|
||||
# haiku.rag.yaml places the database at /data.
|
||||
- ${DB_VOLUME:-./data/haiku.rag.lancedb}:/data
|
||||
- ./backend:/app/src:ro
|
||||
- ./haiku.rag.yaml:/app/haiku.rag.yaml:ro
|
||||
extra_hosts:
|
||||
|
|
|
|||
|
|
@ -7,13 +7,14 @@ services:
|
|||
ports:
|
||||
- "127.0.0.1:8001:8000"
|
||||
environment:
|
||||
- DB_PATH=/data
|
||||
- HAIKU_RAG_CONFIG_PATH=/app/haiku.rag.yaml
|
||||
- ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY:-}
|
||||
- OPENAI_API_KEY=${OPENAI_API_KEY:-}
|
||||
- OLLAMA_BASE_URL=${OLLAMA_BASE_URL:-http://host.docker.internal:11434}
|
||||
- LOGFIRE_TOKEN=${LOGFIRE_TOKEN:-}
|
||||
volumes:
|
||||
- ${DB_PATH:-./data/haiku.rag.lancedb}:/data
|
||||
# haiku.rag.yaml places the database at /data.
|
||||
- ${DB_VOLUME:-./data/haiku.rag.lancedb}:/data
|
||||
- ./haiku.rag.yaml:/app/haiku.rag.yaml:ro
|
||||
extra_hosts:
|
||||
- "host.docker.internal:host-gateway"
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ RUN corepack enable && corepack prepare pnpm@10.15.1 --activate
|
|||
|
||||
FROM base AS deps
|
||||
WORKDIR /app
|
||||
COPY package.json pnpm-lock.yaml ./
|
||||
COPY package.json pnpm-lock.yaml pnpm-workspace.yaml ./
|
||||
RUN pnpm install --frozen-lockfile
|
||||
|
||||
FROM base AS builder
|
||||
|
|
|
|||
|
|
@ -20,6 +20,8 @@ import {
|
|||
import { FilterIcon } from "../lib/icons";
|
||||
import type { RAGState } from "../lib/sessionStorage";
|
||||
import {
|
||||
AGUI_STATE_KEY,
|
||||
agentStateOf,
|
||||
createSession,
|
||||
getActiveSessionId,
|
||||
getLatestCitations,
|
||||
|
|
@ -32,10 +34,7 @@ import DbInfo from "./DbInfo";
|
|||
import DocumentFilter from "./DocumentFilter";
|
||||
import SessionManager from "./SessionManager";
|
||||
|
||||
// Must match state_namespace from haiku.rag.skills.rag
|
||||
const AGUI_STATE_KEY = "rag";
|
||||
|
||||
// AG-UI state is namespaced under AGUI_STATE_KEY
|
||||
// AG-UI state is namespaced under AGUI_STATE_KEY (see sessionStorage).
|
||||
interface AgentState {
|
||||
[AGUI_STATE_KEY]?: RAGState;
|
||||
}
|
||||
|
|
@ -115,24 +114,6 @@ function MessageIcon() {
|
|||
);
|
||||
}
|
||||
|
||||
function FileIcon() {
|
||||
return (
|
||||
<svg
|
||||
width="14"
|
||||
height="14"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
strokeWidth="2"
|
||||
strokeLinecap="round"
|
||||
strokeLinejoin="round"
|
||||
>
|
||||
<path d="M15 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V7Z" />
|
||||
<path d="M14 2v4a2 2 0 0 0 2 2h4" />
|
||||
</svg>
|
||||
);
|
||||
}
|
||||
|
||||
function ToolCallIndicator({
|
||||
toolName,
|
||||
status,
|
||||
|
|
@ -146,13 +127,9 @@ function ToolCallIndicator({
|
|||
|
||||
const getToolIcon = () => {
|
||||
switch (toolName) {
|
||||
case "search":
|
||||
case "rag_search":
|
||||
return <SearchIcon />;
|
||||
case "get_document":
|
||||
return <FileIcon />;
|
||||
case "execute_skill":
|
||||
case "execute_code":
|
||||
case "cite":
|
||||
case "rag_cite":
|
||||
return <MessageIcon />;
|
||||
default:
|
||||
return <SearchIcon />;
|
||||
|
|
@ -161,18 +138,10 @@ function ToolCallIndicator({
|
|||
|
||||
const getToolLabel = () => {
|
||||
switch (toolName) {
|
||||
case "search":
|
||||
case "rag_search":
|
||||
return "Search";
|
||||
case "get_document":
|
||||
return "Document";
|
||||
case "execute_skill":
|
||||
return "Skill";
|
||||
case "execute_code":
|
||||
return "Code";
|
||||
case "cite":
|
||||
case "rag_cite":
|
||||
return "Cite";
|
||||
case "list_documents":
|
||||
return "Documents";
|
||||
default:
|
||||
return toolName;
|
||||
}
|
||||
|
|
@ -180,34 +149,12 @@ function ToolCallIndicator({
|
|||
|
||||
const getDescription = () => {
|
||||
switch (toolName) {
|
||||
case "execute_skill": {
|
||||
const skill = args.skill_name as string | undefined;
|
||||
const request = args.request as string | undefined;
|
||||
return (
|
||||
<span className="tool-query">
|
||||
{skill ? `${skill}: ` : ""}
|
||||
{request ?? "Processing..."}
|
||||
</span>
|
||||
);
|
||||
}
|
||||
case "search": {
|
||||
case "rag_search": {
|
||||
const query = args.query as string;
|
||||
return <span className="tool-query">{query}</span>;
|
||||
}
|
||||
case "get_document":
|
||||
return <span className="tool-query">{args.query as string}</span>;
|
||||
case "execute_code": {
|
||||
const code = args.code as string | undefined;
|
||||
return (
|
||||
<span className="tool-query">
|
||||
{code ? code.slice(0, 80) : "Running code..."}
|
||||
</span>
|
||||
);
|
||||
}
|
||||
case "cite":
|
||||
case "rag_cite":
|
||||
return <span className="tool-query">Registering citations</span>;
|
||||
case "list_documents":
|
||||
return <span className="tool-query">Listing documents</span>;
|
||||
default:
|
||||
return <span>Processing...</span>;
|
||||
}
|
||||
|
|
@ -234,39 +181,6 @@ function ToolCallIndicator({
|
|||
);
|
||||
}
|
||||
|
||||
// Render an activity message from a skill sub-agent tool call/result
|
||||
function ActivityIndicator({
|
||||
message,
|
||||
isComplete,
|
||||
}: {
|
||||
// biome-ignore lint/suspicious/noExplicitAny: AG-UI activity message shape
|
||||
message: any;
|
||||
isComplete: boolean;
|
||||
}) {
|
||||
const content = message.content ?? {};
|
||||
const toolName = content.tool_name ?? "tool";
|
||||
|
||||
let args: Record<string, unknown> = {};
|
||||
if (content.args) {
|
||||
try {
|
||||
args =
|
||||
typeof content.args === "string"
|
||||
? JSON.parse(content.args)
|
||||
: content.args;
|
||||
} catch {
|
||||
// ignore parse errors
|
||||
}
|
||||
}
|
||||
|
||||
return (
|
||||
<ToolCallIndicator
|
||||
toolName={toolName}
|
||||
status={isComplete ? "complete" : "loading"}
|
||||
args={args}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
// Context for sharing chat state with the message view
|
||||
const ChatStateContext = createContext<RAGState | null>(null);
|
||||
|
||||
|
|
@ -298,21 +212,6 @@ function MessageViewWithCitations({
|
|||
const ragState = useContext(ChatStateContext);
|
||||
const latestCitations = ragState ? getLatestCitations(ragState) : [];
|
||||
|
||||
// Collect completed tool_call_ids from skill_tool_result activity messages
|
||||
const completedToolCallIds = useMemo(() => {
|
||||
const ids = new Set<string>();
|
||||
for (const msg of messages) {
|
||||
if (
|
||||
msg.role === "activity" &&
|
||||
msg.activityType === "skill_tool_result" &&
|
||||
msg.content?.tool_call_id
|
||||
) {
|
||||
ids.add(msg.content.tool_call_id);
|
||||
}
|
||||
}
|
||||
return ids;
|
||||
}, [messages]);
|
||||
|
||||
const cursor = isRunning ? (
|
||||
<div key="cursor" className="streaming-cursor">
|
||||
<span className="dot" />
|
||||
|
|
@ -322,8 +221,7 @@ function MessageViewWithCitations({
|
|||
) : null;
|
||||
|
||||
// CopilotChatMessageView renders one element per user/assistant message.
|
||||
// We interleave activity indicators (skill sub-agent tool calls) and
|
||||
// optionally inject CitationBlocks after assistant responses that
|
||||
// Inject CitationBlocks after assistant responses that
|
||||
// followed tool calls.
|
||||
return (
|
||||
<CopilotChatMessageView messages={messages} isRunning={isRunning}>
|
||||
|
|
@ -345,24 +243,6 @@ function MessageViewWithCitations({
|
|||
seenToolCalls = true;
|
||||
}
|
||||
|
||||
// Activity messages are not rendered by CopilotKit —
|
||||
// render them ourselves without consuming messageElements
|
||||
if (msg.role === "activity") {
|
||||
if (msg.activityType === "skill_tool_call") {
|
||||
const toolCallId = msg.content?.tool_call_id;
|
||||
result.push(
|
||||
<ActivityIndicator
|
||||
key={`activity-${msg.id}`}
|
||||
message={msg}
|
||||
isComplete={
|
||||
toolCallId ? completedToolCallIds.has(toolCallId) : false
|
||||
}
|
||||
/>,
|
||||
);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
if (msg.role !== "user" && msg.role !== "assistant") continue;
|
||||
|
||||
if (elemIdx < messageElements.length) {
|
||||
|
|
@ -438,11 +318,10 @@ function ChatContentInner({
|
|||
useEffect(() => {
|
||||
if (agent.messages.length > 0) return;
|
||||
const session = getSession(sessionId);
|
||||
// Seed the namespaced AG-UI state so the backend's first STATE_DELTA
|
||||
// (e.g. add /rag/searches/...) has a namespace object to patch into.
|
||||
agent.setState({
|
||||
[AGUI_STATE_KEY]: normalizeRAGState(session?.ragState),
|
||||
});
|
||||
// Seed state for the capabilities; the backend replaces it after each run.
|
||||
// The whole namespace map goes back, not just the fields this UI reads:
|
||||
// compaction and the citation policy read what earlier questions recorded.
|
||||
agent.setState({ ...agent.state, ...agentStateOf(session ?? undefined) });
|
||||
if (session && session.messages.length > 0) {
|
||||
// biome-ignore lint/suspicious/noExplicitAny: AG-UI Message type is a broad union
|
||||
agent.setMessages(session.messages as any[]);
|
||||
|
|
@ -455,21 +334,15 @@ function ChatContentInner({
|
|||
// biome-ignore lint/correctness/useExhaustiveDependencies: JSON.stringify tracks content changes
|
||||
useEffect(() => {
|
||||
if (sessionId && agent.messages.length > 0) {
|
||||
const currentRagState = normalizeRAGState(
|
||||
(agent.state as AgentState)?.[AGUI_STATE_KEY],
|
||||
);
|
||||
updateSessionMessages(
|
||||
sessionId,
|
||||
serializeMessages(agent.messages),
|
||||
currentRagState,
|
||||
(agent.state ?? {}) as Record<string, unknown>,
|
||||
);
|
||||
}
|
||||
}, [JSON.stringify(agent.messages), ragState, sessionId]);
|
||||
|
||||
// Deduplicate messages by id, keeping the last occurrence.
|
||||
// haiku.skills 0.10.0+ sends activity snapshots with the same id
|
||||
// and replace=true — CopilotKit doesn't deduplicate these, so we
|
||||
// must do it to avoid React duplicate key warnings.
|
||||
// Deduplicate messages by id to avoid React duplicate key warnings.
|
||||
// biome-ignore lint/correctness/useExhaustiveDependencies: stable identity via agent ref
|
||||
const messages = useMemo(() => {
|
||||
const seen = new Map<string, number>();
|
||||
|
|
|
|||
|
|
@ -11,12 +11,16 @@ export interface Citation {
|
|||
doc_item_refs?: string[];
|
||||
}
|
||||
|
||||
// Matches RAGState from the backend skill
|
||||
// Matches RAGState from the backend capability. The fields named here are the
|
||||
// ones this UI reads; the capability owns the rest of its namespace, including
|
||||
// the evidence record that compaction builds its capsule from, so the state has
|
||||
// to round-trip whole rather than be rebuilt from known keys.
|
||||
export interface RAGState {
|
||||
citation_index: Record<string, Citation>;
|
||||
citations: string[];
|
||||
document_filter: string | null;
|
||||
searches: Record<string, unknown[]>;
|
||||
[key: string]: unknown;
|
||||
}
|
||||
|
||||
export interface StoredMessage {
|
||||
|
|
@ -30,16 +34,38 @@ export interface StoredSession {
|
|||
id: string;
|
||||
title: string;
|
||||
messages: StoredMessage[];
|
||||
ragState: RAGState;
|
||||
// The whole AG-UI state. The rag namespace is not the only one a capability
|
||||
// writes: the citation policy records violations beside it.
|
||||
agentState: AgentState;
|
||||
// Sessions stored before agentState existed.
|
||||
ragState?: RAGState;
|
||||
createdAt: string;
|
||||
updatedAt: string;
|
||||
}
|
||||
|
||||
export const AGUI_STATE_KEY = "rag";
|
||||
|
||||
export type AgentState = Record<string, unknown>;
|
||||
|
||||
// Reads the rag namespace out of a stored session, whichever way it was stored.
|
||||
export function ragStateOf(session?: StoredSession): RAGState {
|
||||
const namespaced = session?.agentState?.[AGUI_STATE_KEY] as
|
||||
| Partial<RAGState>
|
||||
| undefined;
|
||||
return normalizeRAGState(namespaced ?? session?.ragState);
|
||||
}
|
||||
|
||||
// The state to seed an agent with when a session is resumed.
|
||||
export function agentStateOf(session?: StoredSession): AgentState {
|
||||
return session?.agentState ?? { [AGUI_STATE_KEY]: ragStateOf(session) };
|
||||
}
|
||||
|
||||
const SESSIONS_KEY = "haiku.rag.sessions";
|
||||
const ACTIVE_SESSION_KEY = "haiku.rag.activeSession";
|
||||
|
||||
export function normalizeRAGState(state?: Partial<RAGState>): RAGState {
|
||||
return {
|
||||
...state,
|
||||
citation_index: state?.citation_index ?? {},
|
||||
citations: state?.citations ?? [],
|
||||
document_filter: state?.document_filter ?? null,
|
||||
|
|
@ -81,7 +107,7 @@ export function createSession(): StoredSession {
|
|||
id: crypto.randomUUID(),
|
||||
title: "New Session",
|
||||
messages: [],
|
||||
ragState: normalizeRAGState(),
|
||||
agentState: { [AGUI_STATE_KEY]: normalizeRAGState() },
|
||||
createdAt: now,
|
||||
updatedAt: now,
|
||||
};
|
||||
|
|
@ -106,7 +132,7 @@ export function saveSession(session: StoredSession): void {
|
|||
export function updateSessionMessages(
|
||||
id: string,
|
||||
messages: StoredMessage[],
|
||||
ragState: RAGState,
|
||||
agentState: AgentState,
|
||||
): void {
|
||||
const sessions = getAllSessions();
|
||||
const idx = sessions.findIndex((s) => s.id === id);
|
||||
|
|
@ -114,7 +140,7 @@ export function updateSessionMessages(
|
|||
|
||||
const session = sessions[idx];
|
||||
session.messages = messages;
|
||||
session.ragState = ragState;
|
||||
session.agentState = agentState;
|
||||
session.updatedAt = new Date().toISOString();
|
||||
|
||||
// Derive title from first user message
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@
|
|||
"@ag-ui/client": "^0.0.57",
|
||||
"@copilotkit/react-core": "^1.61.1",
|
||||
"@copilotkit/runtime": "^1.61.1",
|
||||
"next": "^16.2.9",
|
||||
"next": "^16.2.11",
|
||||
"openai": "^6",
|
||||
"react": "^19.0.0",
|
||||
"react-dom": "^19.0.0",
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
5
app/frontend/pnpm-workspace.yaml
Normal file
5
app/frontend/pnpm-workspace.yaml
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
allowBuilds:
|
||||
'@scarf/scarf': false
|
||||
sharp: true
|
||||
overrides:
|
||||
sharp: '>=0.35.0'
|
||||
|
|
@ -1,6 +1,12 @@
|
|||
# haiku.rag configuration for the chat app
|
||||
# Copy to haiku.rag.yaml and customize as needed
|
||||
|
||||
# The database. The compose files mount DB_VOLUME (default
|
||||
# ./data/haiku.rag.lancedb) at /data.
|
||||
lancedb:
|
||||
databases:
|
||||
haiku.rag: /data
|
||||
|
||||
# QA model configuration
|
||||
qa:
|
||||
model:
|
||||
|
|
|
|||
|
|
@ -40,4 +40,4 @@ EXPOSE 8001 8765
|
|||
|
||||
# Default command: read-only MCP server. The companion ingester service is
|
||||
# launched via docker-compose against the same image.
|
||||
CMD ["haiku-rag", "--config", "/app/haiku.rag.yaml", "--read-only", "mcp", "--port", "8001"]
|
||||
CMD ["haiku-rag", "--config", "/app/haiku.rag.yaml", "mcp", "--port", "8001"]
|
||||
|
|
|
|||
|
|
@ -39,4 +39,4 @@ EXPOSE 8001 8765
|
|||
|
||||
# Default command: read-only MCP server. The companion ingester service is
|
||||
# launched via docker-compose against the same image.
|
||||
CMD ["haiku-rag", "--config", "/app/haiku.rag.yaml", "--read-only", "mcp", "--port", "8001"]
|
||||
CMD ["haiku-rag", "--config", "/app/haiku.rag.yaml", "mcp", "--port", "8001"]
|
||||
|
|
|
|||
12
docs/apps.md
12
docs/apps.md
|
|
@ -40,8 +40,8 @@ Create a `.env` file in the `app/` directory:
|
|||
ANTHROPIC_API_KEY=your-anthropic-key
|
||||
OPENAI_API_KEY=your-openai-key
|
||||
|
||||
# Database path
|
||||
DB_PATH=/path/to/your/haiku.rag.lancedb
|
||||
# Host path of the LanceDB database, mounted at /data
|
||||
DB_VOLUME=/path/to/your/haiku.rag.lancedb
|
||||
|
||||
# Optional: Ollama base URL (if using local models)
|
||||
OLLAMA_BASE_URL=http://localhost:11434
|
||||
|
|
@ -50,16 +50,22 @@ OLLAMA_BASE_URL=http://localhost:11434
|
|||
LOGFIRE_TOKEN=your-logfire-token
|
||||
```
|
||||
|
||||
For full configuration, mount a `haiku.rag.yaml` file:
|
||||
The mounted `haiku.rag.yaml` places the database at `/data` and configures the models; the compose files point `HAIKU_RAG_CONFIG_PATH` at it:
|
||||
|
||||
```yaml
|
||||
# app/haiku.rag.yaml
|
||||
lancedb:
|
||||
databases:
|
||||
haiku.rag: /data
|
||||
|
||||
qa:
|
||||
model:
|
||||
provider: anthropic
|
||||
name: claude-sonnet-4-20250514
|
||||
```
|
||||
|
||||
Outside compose, the backend loads its configuration like the CLI: `HAIKU_RAG_CONFIG_PATH`, then `./haiku.rag.yaml`, then the platform directory.
|
||||
|
||||
## API endpoints
|
||||
|
||||
| Endpoint | Method | Description |
|
||||
|
|
|
|||
|
|
@ -1,82 +1,161 @@
|
|||
# Benchmarks
|
||||
|
||||
We evaluate `haiku.rag` on a small set of datasets that exercise different parts of the pipeline. OpenRAG Bench (ORB), T²-RAGBench, and Wix are the datasets we currently track. Retrieval, QA accuracy, and citation retrieval are scored end-to-end through the rag and rag-analysis skills.
|
||||
We evaluate `haiku.rag` on a small set of datasets that exercise different parts of the pipeline. OpenRAG Bench (ORB), T²-RAGBench, HotpotQA, FRAMES, and MTRAG are the datasets we currently track. Retrieval, QA accuracy, and citation retrieval are scored end-to-end through the RAG and analysis capabilities.
|
||||
|
||||
## Running Evaluations
|
||||
## Current results
|
||||
|
||||
You can run evaluations with the `evaluations` CLI:
|
||||
Numbers below were measured on a recent `haiku.rag` version. Most rows were judged by `Qwen3.6-35B-A3B-NVFP4`; the `Qwen3.8-27B` rows were judged by the currently pinned `qwen3.8`, as their footnote states. Rows are not re-judged when the pinned judge changes, so compare rows judged by the same judge and treat cross-judge differences as unmeasured.
|
||||
|
||||
```bash
|
||||
evaluations run wix
|
||||
evaluations run orb_text
|
||||
```
|
||||
No benchmark database carries a vector index, so every number below reflects exact brute-force kNN rather than approximate search. A vector index is never built automatically. `haiku-rag create-index` builds one, and `haiku-rag doctor` reports whether a database has it. For the measured effect of indexing on retrieval, see [Vector Indexing](configuration/storage.md#vector-indexing).
|
||||
|
||||
The evaluation flow is orchestrated with [`pydantic-evals`](https://github.com/pydantic/pydantic-ai/tree/main/libs/pydantic-evals), which we leverage for dataset management, scoring, and report generation.
|
||||
### OpenRAG Bench (ORB)
|
||||
|
||||
### Pre-built Databases
|
||||
[OpenRAG Bench](https://huggingface.co/datasets/vectara/open_ragbench) contains ArXiv research papers with multimodal question-answering pairs. Queries include both text-based and image-based questions, testing retrieval and reasoning over visual content like figures, charts, and diagrams. Each query maps to one relevant document.
|
||||
|
||||
Building evaluation databases from scratch can take a long time, especially for large datasets like OpenRAG Bench. Pre-built databases are available on HuggingFace:
|
||||
Two approaches are benchmarked separately:
|
||||
|
||||
```bash
|
||||
# Download a specific dataset
|
||||
evaluations download wix
|
||||
- **Multimodal embedder** (`Qwen/Qwen3-VL-Embedding-8B`, served via vLLM): picture bytes and text live in a shared vector space, no VLM is run at ingest.
|
||||
- **Text embedder + VLM picture descriptions** (`qwen3-embedding:4b` + `ollama/ministral-3`): pictures are described at ingest and the descriptions are woven into chunk text. Retrieval runs over text only. See [Picture handling configuration](configuration/processing.md#picture-handling).
|
||||
|
||||
# Download all datasets
|
||||
evaluations download all
|
||||
#### Multimodal embedder
|
||||
|
||||
# Force re-download (overwrite existing)
|
||||
evaluations download wix --force
|
||||
```
|
||||
##### Retrieval (MAP)
|
||||
|
||||
Active datasets:
|
||||
| Embedding Model | Reranker | Cases | MAP |
|
||||
|------------------------------------------|------------------------------------------------------|------:|-------:|
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | none | 3045 | 0.9774 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | none | 3045 | 0.9798 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `nvidia/llama-nemotron-rerank-vl-1b-v2` (multimodal) | 3045 | 0.9913 |
|
||||
|
||||
| Dataset | Size |
|
||||
|---------|------|
|
||||
| `wix` | ~511MB |
|
||||
| `orb_text` — OpenRAG Bench, text embedder (`qwen3-embedding:4b`) with VLM picture descriptions baked into chunk content | ~18 GB |
|
||||
| `orb_multimodal` — OpenRAG Bench, multimodal embedder (`qwen3-vl-embedding-8b`); picture vectors live in the same space as text for cross-modal retrieval | ~16 GB |
|
||||
| `orb_multimodal_nemotron` — OpenRAG Bench, multimodal embedder (`nvidia/llama-nemotron-embed-vl-1b-v2`), the embedder behind the published headline results | ~16 GB |
|
||||
| `t2_finqa` — T²-RAGBench (FinQA) financial QA, text embedder (`qwen3-embedding:4b`); scored by exact numeric match, run with `--target analysis-skill` | ~2 GB |
|
||||
*The nemotron row without a reranker is measured on this release. The reranked row uses `reranking.multimodal: true`: picture chunks reach the vision reranker as images alongside their description text, measured on haiku.rag main post-v0.67.3.*
|
||||
|
||||
After downloading, run benchmarks with `--skip-db`. Each database is built with a specific embedder, so pass its reference config from `evaluations/configs/` (a database only opens against a config whose embedder matches):
|
||||
##### QA accuracy + citation retrieval
|
||||
|
||||
```bash
|
||||
evaluations run orb_multimodal_nemotron --skip-db --config configs/orb_multimodal_nemotron.yaml
|
||||
```
|
||||
| Embedding Model | Target | Capability model | Cases | QA accuracy | Mean `cited_map` |
|
||||
|------------------------------------------|-----------------|-----------------------------------|------:|-------------|------------------|
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 1409 | 0.89 | — |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 3039 | 0.9263 | 0.9761 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Gemma-4-26B-A4B-NVFP4` | 3040 | 0.9362 | 0.9343 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Qwen3.6-35B-A3B-NVFP4` | 3045 | 0.95 | 0.93 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-capability` | `vllm:Muse-Glimmer-30B-NVFP4` | 3045 | 0.9494 | 0.9771 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Muse-Glimmer-30B-NVFP4` | 3017 | 0.9718 | 0.9837 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-capability` | `vllm:Qwen3.8-27B-NVFP4` | 3045 | 0.9514 | 0.9817 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Qwen3.8-27B-NVFP4` | 3042 | 0.9629 | 0.9835 |
|
||||
|
||||
The configs use `vllm` as the model host. Point `base_url` at your own OpenAI-compatible endpoints to reproduce the numbers.
|
||||
*The `Muse-Glimmer-30B` rows run at `chat_template_kwargs.reasoning_strength: high`, no reranker, same judge.*
|
||||
|
||||
### Configuration
|
||||
*The `Qwen3.8-27B` rows run at `chat_template_kwargs.reasoning_effort: low`, no reranker, and are **judged by `Qwen3.8-27B` itself** — the pinned judge, and the same model that produced the answers. A 120-case cross-check by an independent judge agreed on 95% and was never stricter, but the incumbent `Qwen3.6` judge is no longer hosted, so the older rows cannot be re-judged for a like-for-like comparison. `rag-capability` cites on 99.34% of cases with 1.09 citations each; `analysis-capability` on 99.70% with 1.12, at 0.13 code executions per case. Case counts exclude 0 and 3 provider errors respectively.*
|
||||
|
||||
The benchmark script accepts several options:
|
||||
*Both nemotron `Gemma-4` rows are measured on this release, no reranker, judged by `vllm:Qwen3.6-35B-A3B-NVFP4` with thinking on, and exclude the cases that errored (6 of 3045 for `rag-capability`, 5 for `analysis-capability`). The `rag-capability` row cites at 99.64% with a mean of 1.08 citations per case, at a median 4.7s per case against 5.0s for `analysis-capability`. Citation coverage is what moved on this release: 4.9% of analysis cases register no citation, against 26.3% before, at unchanged searches and code executions per case. The remaining rows are from haiku.rag v0.52.0, where Qwen3-VL covered 1409 / 3045 cases.*
|
||||
|
||||
```bash
|
||||
evaluations run wix --config /path/to/haiku.rag.yaml --db /path/to/custom.lancedb
|
||||
```
|
||||
#### Text embedder + VLM picture descriptions
|
||||
|
||||
**Options:**
|
||||
##### Retrieval (MAP)
|
||||
|
||||
- `--config PATH` - Specify a custom `haiku.rag.yaml` configuration file
|
||||
- `--db PATH` - Override the database path (default: platform-specific user data directory)
|
||||
- `--skip-db` - Skip updating the evaluation database
|
||||
- `--skip-retrieval` - Skip retrieval benchmark
|
||||
- `--skip-qa` - Skip QA benchmark
|
||||
- `--limit N` - Limit number of test cases
|
||||
- `--name NAME` - Override the evaluation name
|
||||
- `--target {rag-skill,analysis-skill}` - Choose which [skill](skills/index.md) to benchmark end-to-end (default: `rag-skill`).
|
||||
- `--skill-model PROVIDER:NAME` - Override the skill model independently from the judge (default: `config.qa.model`, or `config.analysis.model` when set for `--target analysis-skill`).
|
||||
| Embedding Model | VLM | Reranker | Cases | MAP |
|
||||
|------------------------------------------|----------------------|------------------------|------:|-------:|
|
||||
| `qwen3-embedding:4b` | Ollama / ministral-3 | `mxbai-rerank-base-v2` | 3045 | 0.9834 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | Ollama / ministral-3 | `mxbai-rerank-base-v2` | 3045 | 0.9863 |
|
||||
|
||||
If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults.
|
||||
*Measured on haiku.rag v0.50.0.*
|
||||
|
||||
To pin the LLM judge in YAML (rather than the default `ollama:qwen3.6`):
|
||||
##### QA accuracy + citation retrieval
|
||||
|
||||
```yaml
|
||||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: gpt-4o-mini
|
||||
base_url: http://localhost:8000/v1 # optional, for OpenAI-compatible servers (vLLM, LM Studio, etc.)
|
||||
```
|
||||
| Embedding Model | VLM | Capability model | Cases | QA accuracy | Mean `cited_map` |
|
||||
|------------------------------------------|----------------------|------------------------------|------:|-------------|------------------|
|
||||
| `qwen3-embedding:4b` | Ollama / ministral-3 | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.80 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | Ollama / ministral-3 | `vllm:Gemma-4-26B-A4B-NVFP4` | 2836 | 0.96 | 0.81 |
|
||||
|
||||
*Measured on haiku.rag v0.50.0 with `mxbai-rerank-base-v2`, judged by `vllm:Qwen3.6-35B-A3B-NVFP4`. Nemotron covered 2836 / 3045 cases.*
|
||||
|
||||
### T²-RAGBench (FinQA)
|
||||
|
||||
[T²-RAGBench](https://huggingface.co/datasets/G4KMU/t2-ragbench) reformulates financial-report QA into context-independent questions with short numeric answers and a 1:1 gold document mapping. The FinQA subset is 2,789 single-page PDFs / 8,281 questions, ingested via docling. Unlike the other datasets, QA is scored deterministically with `NumberMatchEvaluator` (relative tolerance 0.01) instead of an LLM judge, so QA accuracy here is exact numeric match rather than a judged fraction.
|
||||
|
||||
##### QA accuracy + citation retrieval
|
||||
|
||||
| Embedding Model | Reranker | Target | Capability model | Cases | QA accuracy | Mean `cited_map` |
|
||||
|----------------------|------------------------|------------------|------------------------------|------:|-------------|------------------|
|
||||
| `qwen3-embedding:4b` | `mxbai-rerank-base-v2` | `analysis-capability` | `vllm:Qwen3.6-35B-A3B-NVFP4` | 7939 | 0.77 | 0.78 |
|
||||
|
||||
*Measured on haiku.rag v0.55.0, deterministic Number-Match scoring (ε=0.01), 2560-dim `qwen3-embedding:4b` (vLLM) with `mxbai-rerank-base-v2`. 341 / 8281 cases excluded as nulls (analysis spirals from the request limit and in-generation loops). Accuracy and `cited_map` are over the 7939 scored cases. Mean 16.0s/case.*
|
||||
|
||||
### HotpotQA
|
||||
|
||||
[HotpotQA](https://huggingface.co/datasets/hotpotqa/hotpot_qa) is multi-hop question answering over Wikipedia: each question requires combining facts from two supporting paragraphs, with distractor paragraphs in the corpus. We use the distractor validation split: 7,405 questions over ~66k unique paragraphs, each question mapping to two gold documents.
|
||||
|
||||
##### Retrieval (MAP)
|
||||
|
||||
| Embedding Model | Reranker | Cases | MAP |
|
||||
|----------------------|---------------------|------:|-------:|
|
||||
| `qwen3-embedding:4b` | `Qwen3-Reranker-4B` | 7405 | 0.8202 |
|
||||
| `qwen3-embedding:4b` | none | 7405 | 0.6995 |
|
||||
|
||||
The reranker's contribution is larger here than on the single-doc datasets: hybrid search usually surfaces the first-hop document at rank 1, while the second-hop document often needs the reranker to climb into the result window.
|
||||
|
||||
##### QA accuracy + citation retrieval
|
||||
|
||||
| Skill model | Reranker | QA accuracy | Mean `cited_map` |
|
||||
|------------------------------|---------------------|-------------|------------------|
|
||||
| `vllm:Gemma-4-26B-A4B-NVFP4` | `Qwen3-Reranker-4B` | 0.85 | 0.80 |
|
||||
| `vllm:Gemma-4-26B-A4B-NVFP4` | none | 0.83 | 0.75 |
|
||||
|
||||
*Measured on haiku.rag v0.66.0 with `qwen3-embedding:4b` (vLLM, dim 2560), judged by `vllm:Qwen3.6-35B-A3B-NVFP4`, 7,405 cases. The reranker lifts QA accuracy +2.7pts and `cited_map` +4.6pts. Without a reranker, `cited_map` (0.75) still exceeds the no-reranker retrieval MAP (0.70): the skill reformulates queries across search calls, partially recovering second-hop documents that a single query misses.*
|
||||
|
||||
### FRAMES
|
||||
|
||||
[FRAMES](https://huggingface.co/datasets/google/frames-benchmark) is Google's multi-hop QA benchmark: 824 questions, each grounded in 2–23 Wikipedia articles, exercising temporal, numerical, and tabular reasoning across documents. We evaluate 822 questions (2 excluded: a linked article was deleted from Wikipedia) over a fixed corpus of the 2,521 linked articles fetched at current revision. There is no official FRAMES evaluation setup; our protocol — fixed corpus, agentic retrieval, judged accuracy — corresponds to the paper's *multi-step retrieval* setting, where [the paper](https://arxiv.org/abs/2409.12941) reports 0.66 with Gemini-Pro-1.5 (0.729 in its oracle setting, with gold articles provided). Answers were authored against ~2024 revisions and may have drifted with article content.
|
||||
|
||||
##### Retrieval (MAP)
|
||||
|
||||
| Embedding Model | Reranker | Cases | MAP |
|
||||
|----------------------|---------------------|------:|-------:|
|
||||
| `qwen3-embedding:4b` | `Qwen3-Reranker-4B` | 822 | 0.5631 |
|
||||
|
||||
*Single-query retrieval is capped by FRAMES' indirection: in the zero-MAP queries the gold article's subject is never named in the question ("the year the Titanic sank" → `1912_Summer_Olympics`). The agentic targets recover these through iterative search, passing 55% of the very cases single-shot retrieval scores zero on.*
|
||||
|
||||
##### QA accuracy + citation retrieval
|
||||
|
||||
| Capability model | Target | QA accuracy | Mean `cited_map` |
|
||||
|------------------|--------|-------------|------------------|
|
||||
| `vllm:Muse-Glimmer-30B-NVFP4` | `analysis-capability` | 0.7506 | 0.5852 |
|
||||
| `vllm:Qwen3.8-27B-NVFP4` | `analysis-capability` | 0.8095 | 0.6847 |
|
||||
|
||||
*Both rows use `qwen3-embedding:4b` (vLLM, dim 2560), `Qwen3-Reranker-4B`, and are judged by `Qwen3.8-27B`. QA accuracy is over judged cases; counting unanswered cases as failures, the floors are 0.7397 (`Muse-Glimmer`, 1.5% lost to provider errors) and 0.7701 (`Qwen3.8`, 4.87% lost to answers truncated at `max_tokens: 16384`). The `Qwen3.8` row is self-judged — a 100-case paired cross-judge (Glimmer as judge, difference-in-differences) measured the self-preference at +1.0pp with 99% judge agreement. Cite rates are 90.1% (`Muse-Glimmer`) and 99.3% (`Qwen3.8`); `cited_map` is structurally capped below 1 on FRAMES because gold sets span 2–23 articles while answering typically uses a subset. `Qwen3.8` reaches its score on substantially less tool traffic than `Muse-Glimmer` (4.0 searches and 3.9 code executions per case vs 7.0 and 9.5, measured identically from tool spans).*
|
||||
|
||||
### MTRAG (ClapNQ)
|
||||
|
||||
[MTRAG](https://github.com/IBM/mt-rag-benchmark) is IBM's multi-turn RAG benchmark (TACL 2025, SemEval-2026 Task 8): human-authored conversations with per-turn answerability labels and binary relevance judgments. We evaluate the ClapNQ (Wikipedia) domain: 183,408 passages, 29 conversations, 224 turns, 208 retrieval queries.
|
||||
|
||||
Four dataset keys share one database. `mtrag_clapnq` retrieves with the raw last user turn and runs QA by replaying each task's reference conversation prefix as message history. `mtrag_clapnq_rewrite` retrieves with the human standalone rewrites. `mtrag_clapnq_live` replays whole conversations through a single capability session, carrying the model's own answers, tool history and capability state across turns, with `EvidenceCompactionCapability` registered. `mtrag_clapnq_live_uncompacted` is the same replay without compaction, isolating what compaction contributes. This is the only multi-turn evaluation, so it is the only one where compaction acts at all.
|
||||
|
||||
##### Retrieval (Recall@k / nDCG@k)
|
||||
|
||||
Directly comparable with [IBM's published results](https://github.com/IBM/mt-rag-benchmark/tree/main/mtrag-human/retrieval_tasks). Elser is IBM's strongest reported retriever.
|
||||
|
||||
| Retriever | Queries | R@5 | R@10 | nDCG@5 | nDCG@10 |
|
||||
|-----------|---------|----:|-----:|-------:|--------:|
|
||||
| Elser (IBM) | lastturn | 0.49 | 0.58 | 0.45 | 0.49 |
|
||||
| `haiku.rag` | lastturn | 0.501 | 0.600 | 0.455 | 0.497 |
|
||||
| Elser (IBM) | rewrite | 0.52 | 0.64 | 0.48 | 0.54 |
|
||||
| `haiku.rag` | rewrite | 0.548 | 0.668 | 0.503 | 0.556 |
|
||||
|
||||
##### QA accuracy + citation retrieval
|
||||
|
||||
| Mode | Capability model | Turns | QA accuracy | Mean `cited_map` |
|
||||
|------|------------------|------:|-------------|------------------|
|
||||
| Gold-prefix (`mtrag_clapnq`) | `vllm:Muse-Glimmer-30B-NVFP4` | 224 | 0.76 | 0.35 |
|
||||
| Live compacted (`mtrag_clapnq_live`) | `vllm:Muse-Glimmer-30B-NVFP4` | 224/224 scored | 0.83 micro / 0.84 macro | 0.42 |
|
||||
| Live uncompacted (`mtrag_clapnq_live_uncompacted`) | `vllm:Muse-Glimmer-30B-NVFP4` | 224/224 scored | 0.78 micro / 0.79 macro | 0.42 |
|
||||
|
||||
*Measured on haiku.rag v0.74.0 with `qwen3-embedding:4b` (vLLM, dim 2560), `Qwen3-Reranker-4B`, stock capability instructions with `reasoning_strength: high`, judged by the pinned `vllm:Qwen3.6-35B-A3B-NVFP4` (temperature 0.6, thinking). QA numbers are internal (our judge and rubric) and not comparable with IBM's published generation metrics. Gold-prefix and live rates answer different judge questions and are not comparable with each other. The dataset is text-only (ClapNQ passages), so it exercises no multimodal paths.*
|
||||
|
||||
The two live arms replay the same 29 conversations (224 turns) and differ only in registering `EvidenceCompactionCapability`, so they are compared as paired observations:
|
||||
|
||||
- Input tokens per model request, computed as total input tokens divided by model requests across the whole arm: 7,461 compacted (5,207,627 tokens over 698 requests) vs 13,539 uncompacted (9,707,530 over 717 requests). The uncompacted arm used 1.81x as many tokens per request, a 44.9% reduction under compaction.
|
||||
- Answer pass rate: 185/224 vs 175/224 turns. Of the 18 turns where the arms disagree, 14 pass only compacted and 4 only uncompacted. McNemar exact two-sided p = 0.031. The paired difference is +4.5pp with a Wald 95% CI of +0.8 to +8.1pp, so the honest claim is an improvement of roughly 1 to 8 points, not the point estimate.
|
||||
- Citation MAP, macro-averaged over conversations with 208 of 224 turns eligible (turns with gold passages) in each arm: 0.4174 compacted vs 0.4230 uncompacted. The gold-prefix 0.35 is over 208 of 224 eligible cases.
|
||||
- Refusal precision and recall against the answerability labels (16 UNANSWERABLE turns per arm): compacted 0.33 precision and 0.44 recall (21 refusals), uncompacted 0.23 and 0.31 (22 refusals). Gold-prefix: 0.24 and 0.44 (29 refusals).
|
||||
|
||||
## Methodology
|
||||
|
||||
|
|
@ -92,89 +171,121 @@ evaluations:
|
|||
|
||||
### QA Accuracy
|
||||
|
||||
`pydantic-evals` coordinates an LLM judge to determine whether the skill's answer is correct. The default judge is `ollama:qwen3.6`, pinned so changes to the skill model don't change the judge underneath. Set `evaluations.judge` in `haiku.rag.yaml` to override (including a custom `base_url` for any OpenAI-compatible endpoint). Accuracy is the fraction of correctly answered questions.
|
||||
`pydantic-evals` coordinates an LLM judge to determine whether the capability's answer is correct. The default judge is `ollama:qwen3.8`, pinned so changes to the capability model don't change the judge underneath. Set `evaluations.judge` in `haiku.rag.yaml` to override (including a custom `base_url` for any OpenAI-compatible endpoint). Accuracy is the fraction of correctly answered questions.
|
||||
|
||||
We picked `qwen3.6` over the previously-pinned `gpt-oss` after a 4-cell calibration (gpt-oss / qwen3.6 as both answerer and judge, with Claude Opus 4.7 as a reference). `qwen3.6` had κ ≥ 0.66 vs the reference on both same-family and cross-family answerers (vs ~0.39–0.55 for `gpt-oss`) and showed no measurable self-preference bias, while `gpt-oss` was ~10 pp more lenient on its own outputs.
|
||||
A dataset that brings its own deterministic evaluator is scored by that evaluator instead, and no judge runs. T²-RAGBench is the only such dataset today, scored by `NumberMatchEvaluator`.
|
||||
|
||||
`qwen3.8` replaced `qwen3.6` after a 120-case calibration on ORB, stratified 60 pass / 60 fail: agreement 0.950, Cohen's κ 0.900, and in all 6 disagreements it matched or beat `qwen3.6` (4 were `qwen3.6` failing answers that were equivalent in different notation). It emits no reasoning content, so it avoids the thinking spirals that made `qwen3.6` exceed its output budget and drop verdicts. `reasoning_effort` changes its verdicts in 1 case per 120, so the cheaper `low` is pinned.
|
||||
|
||||
Before that, we picked `qwen3.6` over the previously-pinned `gpt-oss` after a 4-cell calibration (gpt-oss / qwen3.6 as both answerer and judge, with Claude Opus 4.7 as a reference). `qwen3.6` had κ ≥ 0.66 vs the reference on both same-family and cross-family answerers (vs ~0.39–0.55 for `gpt-oss`) and showed no measurable self-preference bias, while `gpt-oss` was ~10 pp more lenient on its own outputs.
|
||||
|
||||
### Citation Retrieval
|
||||
|
||||
Alongside QA accuracy, a second metric scores the URIs the skill registered via the `cite` tool against each dataset's gold `expected_uris`, using the same MAP math as raw retrieval. The score key is `cited_map`. Console output also includes the cite rate (% of cases with at least one citation) and the mean number of citations per case.
|
||||
Alongside QA accuracy, a second metric scores the URIs the capability registered via the `cite` tool against each dataset's gold `expected_uris`, using the same MAP math as raw retrieval. The score key is `cited_map`. Console output also includes the cite rate (% of cases with at least one citation) and the mean number of citations per case.
|
||||
|
||||
This is computed alongside QA accuracy from the same skill run, no extra invocations. The signal complements raw retrieval: where raw retrieval measures whether the retriever surfaced the gold document at any rank, citation retrieval measures whether the skill grounded its answer on it.
|
||||
This is computed alongside QA accuracy from the same capability run, no extra invocations. The signal complements raw retrieval: where raw retrieval measures whether the retriever surfaced the gold document at any rank, citation retrieval measures whether the capability grounded its answer on it.
|
||||
|
||||
## Current results
|
||||
## Running Evaluations
|
||||
|
||||
Numbers measured under the current pinned judge (`ollama:qwen3.6`) on a recent `haiku.rag` version.
|
||||
You can run evaluations with the `evaluations` CLI:
|
||||
|
||||
### OpenRAG Bench (ORB)
|
||||
```bash
|
||||
evaluations run hotpotqa
|
||||
evaluations run orb_text
|
||||
```
|
||||
|
||||
[OpenRAG Bench](https://huggingface.co/datasets/vectara/open_ragbench) contains ArXiv research papers with multimodal question-answering pairs. Queries include both text-based and image-based questions, testing retrieval and reasoning over visual content like figures, charts, and diagrams. Each query maps to one relevant document.
|
||||
The evaluation flow is orchestrated with [`pydantic-evals`](https://github.com/pydantic/pydantic-ai/tree/main/libs/pydantic-evals), which we leverage for dataset management, scoring, and report generation.
|
||||
|
||||
Two approaches are benchmarked separately:
|
||||
### Pre-built Databases
|
||||
|
||||
- **Multimodal embedder** (`Qwen/Qwen3-VL-Embedding-8B`, served via vLLM): picture bytes and text live in a shared vector space, no VLM is run at ingest.
|
||||
- **Text embedder + VLM picture descriptions** (`qwen3-embedding:4b` + `ollama/ministral-3`): pictures are described at ingest and the descriptions are woven into chunk text. Retrieval runs over text only. See [Picture handling configuration](configuration/processing.md#picture-handling).
|
||||
Building evaluation databases from scratch can take a long time, especially for large datasets like OpenRAG Bench. Pre-built databases are available on HuggingFace:
|
||||
|
||||
#### Multimodal embedder
|
||||
```bash
|
||||
# Download a specific dataset
|
||||
evaluations download hotpotqa
|
||||
|
||||
##### Retrieval (MAP)
|
||||
# Download all datasets
|
||||
evaluations download all
|
||||
|
||||
| Embedding Model | Cases | MAP |
|
||||
|------------------------------------------|------:|-------:|
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | 3045 | 0.9774 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | 3045 | 0.9709 |
|
||||
# Force re-download (overwrite existing)
|
||||
evaluations download hotpotqa --force
|
||||
```
|
||||
|
||||
##### QA accuracy + citation retrieval
|
||||
Active datasets:
|
||||
|
||||
| Embedding Model | Target | Skill model | Cases | QA accuracy | Mean `cited_map` |
|
||||
|------------------------------------------|-----------------|-----------------------------------|------:|-------------|------------------|
|
||||
| `Qwen/Qwen3-VL-Embedding-8B` | `rag-skill` | `vllm:Gemma-4-26B-A4B-NVFP4` | 1409 | 0.89 | — |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-skill` | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.93 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-skill`| `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.94 | 0.78 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-skill`| `vllm:Qwen3.6-35B-A3B-NVFP4` | 3045 | 0.95 | 0.93 |
|
||||
| Dataset | Size |
|
||||
|---------|------|
|
||||
| `orb_text` — OpenRAG Bench, text embedder (`qwen3-embedding:4b`) with VLM picture descriptions baked into chunk content | ~18 GB |
|
||||
| `orb_multimodal` — OpenRAG Bench, multimodal embedder (`qwen3-vl-embedding-8b`); picture vectors live in the same space as text for cross-modal retrieval | ~16 GB |
|
||||
| `orb_multimodal_nemotron` — OpenRAG Bench, multimodal embedder (`nvidia/llama-nemotron-embed-vl-1b-v2`), the embedder behind the published headline results | ~16 GB |
|
||||
| `t2_finqa` — T²-RAGBench (FinQA) financial QA, text embedder (`qwen3-embedding:4b`); scored by exact numeric match, run with `--target analysis-capability` | ~2 GB |
|
||||
| `hotpotqa` — HotpotQA multi-hop QA over Wikipedia paragraphs, text embedder (`qwen3-embedding:4b`) | ~1.5 GB |
|
||||
| `mtrag_clapnq` — MTRAG multi-turn RAG, ClapNQ (Wikipedia) passages, text embedder (`qwen3-embedding:4b`); also serves the `mtrag_clapnq_rewrite`, `mtrag_clapnq_live` and `mtrag_clapnq_live_uncompacted` keys | ~2.8 GB |
|
||||
|
||||
*Measured on haiku.rag v0.52.0, no reranker, judged by `vllm:Qwen3.6-35B-A3B-NVFP4`. Qwen3-VL covered 1409 / 3045 cases.*
|
||||
After downloading, run benchmarks with `--skip-db`. Each database is built with a specific embedder, so pass its reference config from `evaluations/configs/` (a database only opens against a config whose embedder matches):
|
||||
|
||||
#### Text embedder + VLM picture descriptions
|
||||
```bash
|
||||
evaluations run orb_multimodal_nemotron --skip-db --config configs/orb_multimodal_nemotron.yaml
|
||||
```
|
||||
|
||||
##### Retrieval (MAP)
|
||||
The configs use `vllm` as the model host. Point `base_url` at your own OpenAI-compatible endpoints to reproduce the numbers.
|
||||
|
||||
| Embedding Model | VLM | Reranker | Cases | MAP |
|
||||
|------------------------------------------|----------------------|------------------------|------:|-------:|
|
||||
| `qwen3-embedding:4b` | Ollama / ministral-3 | `mxbai-rerank-base-v2` | 3045 | 0.9834 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | Ollama / ministral-3 | `mxbai-rerank-base-v2` | 3045 | 0.9863 |
|
||||
### Configuration
|
||||
|
||||
*Measured on haiku.rag v0.50.0.*
|
||||
The benchmark script accepts several options:
|
||||
|
||||
##### QA accuracy + citation retrieval
|
||||
```bash
|
||||
evaluations run hotpotqa --config /path/to/haiku.rag.yaml --db /path/to/custom.lancedb
|
||||
```
|
||||
|
||||
| Embedding Model | VLM | Skill model | Cases | QA accuracy | Mean `cited_map` |
|
||||
|------------------------------------------|----------------------|------------------------------|------:|-------------|------------------|
|
||||
| `qwen3-embedding:4b` | Ollama / ministral-3 | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.80 |
|
||||
| `nvidia/llama-nemotron-embed-vl-1b-v2` | Ollama / ministral-3 | `vllm:Gemma-4-26B-A4B-NVFP4` | 2836 | 0.96 | 0.81 |
|
||||
**Options:**
|
||||
|
||||
*Measured on haiku.rag v0.50.0 with `mxbai-rerank-base-v2`, judged by `vllm:Qwen3.6-35B-A3B-NVFP4`. Nemotron covered 2836 / 3045 cases.*
|
||||
- `--config PATH` - Specify a custom `haiku.rag.yaml` configuration file
|
||||
- `--db PATH` - Override the database path (default: platform-specific user data directory)
|
||||
- `--skip-db` - Skip updating the evaluation database
|
||||
- `--skip-retrieval` - Skip retrieval benchmark
|
||||
- `--skip-qa` - Skip QA benchmark
|
||||
- `--limit N` - Limit number of test cases
|
||||
- `--name NAME` - Override the evaluation name
|
||||
- `--target {rag-capability,analysis-capability}` - Choose which [capability](capabilities/index.md) to benchmark end-to-end (default: `rag-capability`). The target names remain stable dataset identifiers.
|
||||
- `--capability-model PROVIDER:NAME` - Override the capability model independently from the judge (default: `config.qa.model`, or `config.analysis.model` when set for `--target analysis-capability`).
|
||||
- `--filter CLAUSE` / `-f CLAUSE` - Restrict every benchmark search to a subset of the database (see [Restricting the corpus](#restricting-the-corpus)).
|
||||
|
||||
### T²-RAGBench (FinQA)
|
||||
If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults.
|
||||
|
||||
[T²-RAGBench](https://huggingface.co/datasets/G4KMU/t2-ragbench) reformulates financial-report QA into context-independent questions with short numeric answers and a 1:1 gold document mapping. The FinQA subset is 2,789 single-page PDFs / 8,281 questions, ingested via docling. Unlike the other datasets, QA is scored deterministically with `NumberMatchEvaluator` (relative tolerance 0.01) instead of an LLM judge, so QA accuracy here is exact numeric match rather than a judged fraction.
|
||||
To pin the LLM judge in YAML (rather than the default `ollama:qwen3.8`). These are the recommended settings:
|
||||
|
||||
##### QA accuracy + citation retrieval
|
||||
```yaml
|
||||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://localhost:8000/v1 # optional, for OpenAI-compatible servers (vLLM, LM Studio, etc.)
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
top_p: 0.95
|
||||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
reasoning_effort: low # qwen3.8: low | medium | xhigh (default)
|
||||
```
|
||||
|
||||
| Embedding Model | Reranker | Target | Skill model | Cases | QA accuracy | Mean `cited_map` |
|
||||
|----------------------|------------------------|------------------|------------------------------|------:|-------------|------------------|
|
||||
| `qwen3-embedding:4b` | `mxbai-rerank-base-v2` | `analysis-skill` | `vllm:Qwen3.6-35B-A3B-NVFP4` | 7939 | 0.77 | 0.78 |
|
||||
### Restricting the corpus
|
||||
|
||||
*Measured on haiku.rag v0.55.0, deterministic Number-Match scoring (ε=0.01), 2560-dim `qwen3-embedding:4b` (vLLM) with `mxbai-rerank-base-v2`. 341 / 8281 cases excluded as nulls (analysis spirals from the request limit and in-generation loops). Accuracy and `cited_map` are over the 7939 scored cases. Mean 16.0s/case.*
|
||||
When a database holds documents from several corpora — only some of which a dataset's questions are drawn from — `--filter` restricts every benchmark search to a subset. It takes the same SQL `WHERE` clause as `haiku-rag search --filter`, over document columns (`id`, `uri`, `title`, `created_at`, `updated_at`, `metadata`). Each dataset writes its own URIs: `orb_text` uses bare arXiv ids such as `2407.01528v3`, `hotpotqa` uses page titles.
|
||||
|
||||
### Wix
|
||||
```bash
|
||||
evaluations run orb_text --skip-db --config haiku.rag.s3.yaml \
|
||||
--filter "uri LIKE '2407%'"
|
||||
```
|
||||
|
||||
[WixQA](https://huggingface.co/datasets/Wix/WixQA) is real customer support questions paired with curated answers. 200 cases.
|
||||
If the corpora are distinguished by a tag rather than by URI, attach it at ingest time as document metadata and match it with `LIKE`. `metadata` is stored as a `json.dumps` string, so there is no JSON subfield access — match the serialized key/value, including the space after the colon:
|
||||
|
||||
`evaluations run wix --target rag-skill` runs the RAG skill end-to-end and produces both QA accuracy and a citation retrieval metric (`cited_map`) computed from the URIs the skill registered via the `cite` tool against the gold `expected_uris`.
|
||||
```bash
|
||||
evaluations run orb_text --skip-db --filter "metadata LIKE '%\"corpus\": \"orb_text\"%'"
|
||||
```
|
||||
|
||||
| Skill model | Reranker | QA accuracy | Mean `cited_map` |
|
||||
|------------------------------|------------------------|-------------|------------------|
|
||||
| `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | 0.87 | 0.38 |
|
||||
The clause applies to both benchmark phases — the retrieval benchmark's searches and every search the capability runs during QA — so the two score the same subset. It is recorded as `document_filter` in the run's experiment metadata, so a filtered run is never mistaken for an unfiltered one when comparing results.
|
||||
|
||||
*Measured on haiku.rag v0.48.0 with `qwen3-embedding:4b` (vLLM, dim 2560), `chunk_size=256`, `search.limit=5`. Judged by `vllm:Qwen3.6-35B-A3B-NVFP4` (qwen3.6 family, NVFP4 quant served via vLLM rather than the default Ollama). 172 / 198 completed cases (2 errored).*
|
||||
Filtering affects searches only — a run without `--skip-db` still populates the database with the dataset's full corpus.
|
||||
|
|
|
|||
59
docs/capabilities/analysis.md
Normal file
59
docs/capabilities/analysis.md
Normal file
|
|
@ -0,0 +1,59 @@
|
|||
# Analysis Capability
|
||||
|
||||
`AnalysisCapability` adds search, citations, and sandboxed Python computation over the document corpus. Use it for counts, aggregation, comparison, structural traversal, and section-scoped reading.
|
||||
|
||||
It is deferred by default, keeping its substantial instructions and tool schemas out of context until the model chooses to load it.
|
||||
|
||||
The default request limit is 30 model requests per question. Override it with `create_capability(request_limit=...)`, or set `request_limit=None` to disable it. As with the RAG capability, `create_capability(vision=...)` overrides the image-attachment gate, defaulting to the configured analysis model's `vision` flag. At the limit, `analysis_search` and `analysis_execute_code` are removed while `analysis_cite` remains for two further requests that call an analysis tool, so the model can register citations before answering from gathered evidence. Requests spent on other capabilities do not count against that window. Other agent and capability tools remain available, and the budget resets for every agent run.
|
||||
|
||||
When `qa.max_searches` or `analysis.max_executions` runs out, the exhausted tool keeps failing rather than disappearing, and the instructions name it on every following request. Searching from inside `analysis_execute_code` does not count against `qa.max_searches`.
|
||||
|
||||
## Tools
|
||||
|
||||
| Tool | Purpose |
|
||||
|---|---|
|
||||
| `analysis_search(query, limit?)` | Search the corpus for evidence. |
|
||||
| `analysis_execute_code(code)` | Run Python against the virtual document filesystem. |
|
||||
| `analysis_cite(chunk_ids)` | Register retrieved or filesystem-derived chunk IDs. |
|
||||
|
||||
The sandbox exposes documents under `/documents/{document_id}/` with `metadata.json`, `content.txt`, `items.jsonl`, `chunks.jsonl` (chunk ids with their metadata) and `toc.json`. In code, `await search()` results carry `chunk_meta` and `await list_documents()` rows carry `metadata`. The interpreter's limits and the per-call budgets are listed under [MCP, Code](../mcp.md#code).
|
||||
|
||||
## Compose an agent
|
||||
|
||||
Register it on its own, not alongside `RAGCapability`: it already searches and cites,
|
||||
and the two together give the model duplicate tools and separate budgets. See
|
||||
[Capabilities](index.md#compose-an-agent).
|
||||
|
||||
```python
|
||||
from pydantic_ai import Agent
|
||||
from haiku.rag.capabilities.analysis import create_capability as analysis
|
||||
from haiku.rag.capabilities.compaction import create_capability as compaction
|
||||
from haiku.rag.capabilities.policy import create_capability as citation_policy
|
||||
|
||||
agent = Agent(
|
||||
"openai:gpt-5",
|
||||
capabilities=[
|
||||
analysis(db_path="my.lancedb"),
|
||||
compaction(),
|
||||
citation_policy(),
|
||||
],
|
||||
)
|
||||
```
|
||||
|
||||
For the high-level convenience API:
|
||||
|
||||
```python
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
async with HaikuRAG("my.lancedb") as client:
|
||||
result = await client.analyze("Which quarter had the highest revenue?")
|
||||
print(result.answer)
|
||||
```
|
||||
|
||||
## State
|
||||
|
||||
When dependencies expose a state dictionary, `AnalysisState` is stored under `"analysis"`. It contains the document filter, code execution log, searches, citations, and the `evidence` record of what was retrieved and cited per question. Searches and executions are cleared when a new question starts, and a resumed question keeps them; the filter, citation index and evidence record persist.
|
||||
|
||||
This capability does not alter the message history either. Register the [compaction capability](compaction.md) to compact earlier questions.
|
||||
|
||||
The capability lazily opens both LanceDB and the sandbox only after it is loaded and a tool requires them. Resources close at the end of the agent run.
|
||||
50
docs/capabilities/compaction.md
Normal file
50
docs/capabilities/compaction.md
Normal file
|
|
@ -0,0 +1,50 @@
|
|||
# Evidence compaction capability
|
||||
|
||||
`EvidenceCompactionCapability` keeps a multi-turn conversation from carrying every
|
||||
search result it ever produced. Every question adds its evidence to the history, so
|
||||
requests grow turn after turn, which degrades answers and can exceed a provider's
|
||||
limits.
|
||||
|
||||
Register it alongside an evidence capability:
|
||||
|
||||
```python
|
||||
from pydantic_ai import Agent
|
||||
from haiku.rag.capabilities.compaction import create_capability as compaction
|
||||
from haiku.rag.capabilities.rag import create_capability as rag
|
||||
|
||||
agent = Agent(
|
||||
"openai:gpt-5",
|
||||
capabilities=[rag(db_path="my.lancedb"), compaction()],
|
||||
)
|
||||
```
|
||||
|
||||
It exposes no tools and takes no configuration. Registering it is the only switch:
|
||||
leave it out and the transcript reaches the model untouched.
|
||||
|
||||
The host must carry the capability state between runs, alongside the message
|
||||
history: the capsule is built from what earlier questions recorded there. Given
|
||||
only a message history, every run starts from an empty record, and compaction
|
||||
refuses rather than replace evidence it cannot retain. See
|
||||
[Compose an agent](index.md#compose-an-agent) for the shape.
|
||||
|
||||
## What it does
|
||||
|
||||
On each request, evidence from earlier questions is replaced by the evidence those
|
||||
questions actually cited. Cited text and cited page images are kept in full, grouped by
|
||||
the question that cited them, and stay citable by the same chunk ids. Evidence spanning
|
||||
more than one collection carries a `Collection:` line naming the one it came from. Every
|
||||
other earlier evidence return becomes a short receipt. The current question is untouched.
|
||||
|
||||
Compaction rewrites the request, never the stored history, so `all_messages()` still
|
||||
holds everything the run gathered.
|
||||
|
||||
This reduces what a request carries. It does not bound it: retained evidence still
|
||||
grows with the conversation. A host that needs more aggressive pruning can compact its
|
||||
own requests further, on the wire only.
|
||||
|
||||
## Resuming a question
|
||||
|
||||
Resuming a question (deferred tool results, an interruption, a suspension) requires the
|
||||
host to carry the capability state from the run being resumed, alongside the message
|
||||
history. Without it the identity of the question in progress is unknowable, and the run
|
||||
fails rather than silently treating it as a new question.
|
||||
150
docs/capabilities/index.md
Normal file
150
docs/capabilities/index.md
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
# Capabilities
|
||||
|
||||
haiku.rag provides native [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/):
|
||||
|
||||
| Capability | Use it for |
|
||||
|---|---|
|
||||
| [`RAGCapability`](rag.md) | Grounded document search and citations. |
|
||||
| [`AnalysisCapability`](analysis.md) | Corpus computation and structural analysis with sandboxed Python. |
|
||||
| [`EvidenceCompactionCapability`](compaction.md) | Optional. Shrinking a conversation's history to the evidence that was cited. |
|
||||
| [`CitationPolicyCapability`](policy.md) | Optional. Requiring every answer to declare what grounds it. |
|
||||
|
||||
The two evidence capabilities are deferred by default. An agent initially sees only their descriptions and the standard `load_capability` tool. Instructions and tools enter the model context only when the model loads a capability.
|
||||
|
||||
## Compose an agent
|
||||
|
||||
Pick one evidence capability, and add both optional capabilities to it:
|
||||
|
||||
```python
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.messages import ModelMessage
|
||||
|
||||
from haiku.rag.capabilities.compaction import create_capability as compaction
|
||||
from haiku.rag.capabilities.policy import create_capability as citation_policy
|
||||
from haiku.rag.capabilities.rag import create_capability as rag
|
||||
|
||||
|
||||
@dataclass
|
||||
class Deps:
|
||||
state: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
agent = Agent(
|
||||
"openai:gpt-5",
|
||||
capabilities=[
|
||||
rag(db_path="my.lancedb"),
|
||||
compaction(),
|
||||
citation_policy(),
|
||||
],
|
||||
deps_type=Deps,
|
||||
)
|
||||
|
||||
# One Deps and one history for the conversation: the capabilities read both.
|
||||
deps = Deps()
|
||||
history: list[ModelMessage] = []
|
||||
|
||||
result = await agent.run("What does the knowledge base say about X?", deps=deps, message_history=history)
|
||||
history = list(result.all_messages())
|
||||
print(result.output)
|
||||
```
|
||||
|
||||
!!! warning "Both optional capabilities need the host to carry state"
|
||||
|
||||
They read what earlier questions retrieved and cited from the capability's
|
||||
state, so the host must expose a `state` dict on its agent dependencies and
|
||||
hand the same dict back on every run of a conversation, alongside the message
|
||||
history. With only the message history, every run starts from an empty record:
|
||||
compaction refuses rather than replace evidence it cannot retain, and the
|
||||
citation policy cannot enforce a follow-up about evidence cited earlier.
|
||||
|
||||
Swap `rag` for `analysis` for an analysis agent. Both optional capabilities work the
|
||||
same way with either one, and neither exposes tools or takes configuration.
|
||||
|
||||
!!! note "Register one evidence capability, not both"
|
||||
|
||||
`RAGCapability` and `AnalysisCapability` overlap. Both search the same corpus and
|
||||
both register citations, so an agent holding both must choose between two
|
||||
near-identical search tools, and its citations land in whichever capability it
|
||||
happened to call. Each also carries its own request limit and its own search
|
||||
budget, so registering both doubles what a question may spend.
|
||||
|
||||
Choose by what the questions need. `RAGCapability` answers questions from retrieved
|
||||
passages. `AnalysisCapability` adds a Python sandbox and a document filesystem, for
|
||||
questions that compute over many documents or read their structure, and it can
|
||||
search too. If you need computation, register the analysis capability alone rather
|
||||
than adding it to the RAG one.
|
||||
|
||||
## Agent specs
|
||||
|
||||
The capabilities can be declared in a Pydantic AI [agent spec](https://ai.pydantic.dev/agent-spec/):
|
||||
|
||||
```yaml title="agent.yaml"
|
||||
model: openai:gpt-5
|
||||
instructions: You are a research assistant with access to a document knowledge base.
|
||||
capabilities:
|
||||
- RAGCapability:
|
||||
db_path: /data/kb.lancedb
|
||||
defer_loading: false
|
||||
- EvidenceCompactionCapability
|
||||
- CitationPolicyCapability
|
||||
```
|
||||
|
||||
Pydantic AI does not discover third-party capabilities, so the caller names the classes:
|
||||
|
||||
```python
|
||||
from pydantic_ai import Agent
|
||||
|
||||
from haiku.rag.capabilities.compaction import EvidenceCompactionCapability
|
||||
from haiku.rag.capabilities.policy import CitationPolicyCapability
|
||||
from haiku.rag.capabilities.rag import RAGCapability
|
||||
|
||||
agent = Agent.from_file(
|
||||
"agent.yaml",
|
||||
deps_type=Deps,
|
||||
custom_capability_types=[
|
||||
RAGCapability,
|
||||
EvidenceCompactionCapability,
|
||||
CitationPolicyCapability,
|
||||
],
|
||||
)
|
||||
```
|
||||
|
||||
`deps_type` stays a Python argument, since the capabilities read and write their state
|
||||
through `deps.state` (see [State](#state)). `Agent.from_file` reads YAML, which needs
|
||||
`pydantic-ai-slim[spec]`; `Agent.from_spec` takes a dict and needs no YAML parser.
|
||||
|
||||
Set `defer_loading: false` when the agent registers a single evidence capability, so its
|
||||
tools are visible immediately. Leave it at the default when the model should route among
|
||||
multiple capabilities.
|
||||
|
||||
A `config:` block accepts a whole `AppConfig`, for agents in one process that need
|
||||
different databases or embedding models:
|
||||
|
||||
```yaml
|
||||
capabilities:
|
||||
- RAGCapability:
|
||||
db_path: /data/kb.lancedb
|
||||
config:
|
||||
embeddings:
|
||||
model: {provider: ollama, name: embeddinggemma, vector_dim: 2048}
|
||||
```
|
||||
|
||||
The block is read like a `haiku.rag.yaml` file: keys it omits take `AppConfig` defaults
|
||||
rather than values from the configuration file on disk. The embedding model must match the
|
||||
database; a mismatch may prevent opening it or produce invalid retrieval. Write the block in
|
||||
full or omit it and let the [configuration file](../configuration/index.md) apply.
|
||||
|
||||
## State
|
||||
|
||||
Capabilities use a plain `state: dict[str, Any]` attribute on agent dependencies when one is available. RAG state lives under `"rag"`; analysis state lives under `"analysis"`. This keeps state independent of any transport or UI protocol.
|
||||
|
||||
Applications serving AG-UI should adapt the agent with Pydantic AI's `AGUIAdapter`. Native model and tool events require no haiku.rag-specific bridge.
|
||||
|
||||
## Database Selection
|
||||
|
||||
RAG and analysis capabilities cover the databases the configuration places: [`lancedb.databases`](../configuration/storage.md#multiple-databases), or with nothing configured the default database `haiku.rag` under `storage.data_dir`. The `db_path` argument places one database where the configuration places none; beside `lancedb.databases` it raises `AmbiguousDatabaseError`.
|
||||
|
||||
Passing a client through `rag=` bypasses this selection. The capability uses the databases covered by that client and does not close it.
|
||||
54
docs/capabilities/policy.md
Normal file
54
docs/capabilities/policy.md
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
# Citation policy capability
|
||||
|
||||
`CitationPolicyCapability` requires every answer to declare what grounds it. Citing is
|
||||
always available and always recorded without it, but nothing makes the model do it.
|
||||
|
||||
Register it alongside an evidence capability:
|
||||
|
||||
```python
|
||||
from pydantic_ai import Agent
|
||||
from haiku.rag.capabilities.policy import create_capability as citation_policy
|
||||
from haiku.rag.capabilities.rag import create_capability as rag
|
||||
|
||||
agent = Agent(
|
||||
"openai:gpt-5",
|
||||
capabilities=[rag(db_path="my.lancedb"), citation_policy()],
|
||||
)
|
||||
```
|
||||
|
||||
It exposes no tools and takes no configuration. Exactly one policy capability makes
|
||||
the decision, however many evidence capabilities are registered, so two of them
|
||||
cannot each demand a citation for one answer.
|
||||
|
||||
The host must carry the capability state between runs, alongside the message
|
||||
history. Enforcement reads what the conversation has already cited, so without it
|
||||
a follow-up about evidence cited earlier goes unenforced. See
|
||||
[Compose an agent](index.md#compose-an-agent) for the shape.
|
||||
|
||||
## Declaring nothing is a valid answer
|
||||
|
||||
A model that finds nothing relevant calls the cite tool with an empty list. That records
|
||||
the answer as *ungrounded*, which is distinct from an answer that declared nothing at
|
||||
all (*missing*). The distinction is what makes a declaration requirable without forcing
|
||||
the model to invent grounding.
|
||||
|
||||
## What happens when a question ends undeclared
|
||||
|
||||
The model is asked once to record what grounded the answer it already gave. It is not
|
||||
asked to change the answer. If the cite tool is no longer available by then, or the
|
||||
question finishes undeclared anyway, the question is recorded in
|
||||
`CitationPolicyState.violations` under the `"citation_policy"` state key. Pointing a
|
||||
model at a tool that is gone costs it retries, so the capability records the failure
|
||||
instead.
|
||||
|
||||
## Which answers are enforced
|
||||
|
||||
Every answer in a conversation that has something to declare: either this question
|
||||
retrieved evidence, or the conversation has already cited something, which stays
|
||||
available to later answers. A follow-up about evidence cited earlier is enforced even
|
||||
though it searched nothing, which is the case the capability exists for.
|
||||
|
||||
Once anything has been cited, later turns are enforced too, a greeting included. The
|
||||
model satisfies the policy by citing an empty list, at the cost of one extra request. A
|
||||
conversation with neither a current-question evidence outcome nor any earlier citation
|
||||
is not enforced.
|
||||
69
docs/capabilities/rag.md
Normal file
69
docs/capabilities/rag.md
Normal file
|
|
@ -0,0 +1,69 @@
|
|||
# RAG Capability
|
||||
|
||||
`RAGCapability` adds grounded document search and citations to a Pydantic AI agent. It is deferred by default, so its instructions and tools do not consume model context until loaded.
|
||||
|
||||
## Tools
|
||||
|
||||
| Tool | Purpose |
|
||||
|---|---|
|
||||
| `rag_search(query, limit?)` | Hybrid vector and full-text search with context expansion. |
|
||||
| `rag_cite(chunk_ids)` | Register exact result chunk IDs as answer citations. |
|
||||
|
||||
The distinct `rag_` prefix lets this capability coexist with analysis and other search providers.
|
||||
|
||||
## Create and compose
|
||||
|
||||
Register it on its own rather than alongside `AnalysisCapability`, which searches and
|
||||
cites as well. See [Capabilities](index.md#compose-an-agent).
|
||||
|
||||
```python
|
||||
from pydantic_ai import Agent
|
||||
from haiku.rag.capabilities.compaction import create_capability as compaction
|
||||
from haiku.rag.capabilities.policy import create_capability as citation_policy
|
||||
from haiku.rag.capabilities.rag import create_capability as rag
|
||||
|
||||
agent = Agent(
|
||||
"openai:gpt-5",
|
||||
capabilities=[
|
||||
rag(db_path="my.lancedb"),
|
||||
compaction(),
|
||||
citation_policy(),
|
||||
],
|
||||
)
|
||||
|
||||
result = await agent.run("What safety equipment does the manual require?")
|
||||
print(result.output)
|
||||
```
|
||||
|
||||
`create_capability` accepts `db_path`, `config`, `defer_loading`, `request_limit`, and `vision`. Set `defer_loading=False` for a dedicated RAG agent where routing is unnecessary. The default request limit is 20 model requests per question; set `request_limit=None` to disable it. `vision` controls whether picture results are attached to search returns as images and should reflect the model the hosting agent runs; it defaults to the configured QA model's `vision` flag.
|
||||
|
||||
When the limit is reached, `rag_search` is removed while `rag_cite` remains for two further requests that call a RAG tool, so the model can register citations before answering from evidence already gathered. Requests spent on other capabilities do not count against that window. Unrelated agent and capability tools remain available. A new agent run starts a fresh limit, so multi-turn chat does not consume one shared budget.
|
||||
|
||||
## State
|
||||
|
||||
When agent dependencies expose a `state` dictionary, the capability maintains a `RAGState` under `"rag"`:
|
||||
|
||||
```python
|
||||
class RAGState(BaseModel):
|
||||
citation_index: dict[str, Citation]
|
||||
citations: list[str]
|
||||
document_filter: str | None
|
||||
evidence: CapabilityEvidenceRecord
|
||||
searches: dict[str, list[SearchResult]]
|
||||
```
|
||||
|
||||
`document_filter`, `citation_index` and `evidence` persist across runs. Citations and searches are cleared when a new question starts; a run that resumes a question keeps the evidence it is still answering from.
|
||||
|
||||
`evidence` records which chunks this capability retrieved and cited, and in which question. `haiku.rag.capabilities.ledger.citation_status(records, question=...)` derives `missing`, `grounded` or `ungrounded` from it, across capabilities.
|
||||
|
||||
State is ordinary application state; the capability does not depend on AG-UI. An AG-UI application can expose it using Pydantic AI's standard adapter.
|
||||
|
||||
## Context management
|
||||
|
||||
This capability does not alter the message history. To stop long conversations resending old retrieved content, register the [compaction capability](compaction.md) alongside it.
|
||||
|
||||
## Domain context and vision
|
||||
|
||||
`prompts.domain_preamble` is prepended to the packaged capability instructions. When the capability's `vision` gate is on (by default, when the configured QA model has `vision: true`), picture results are attached to search returns as `BinaryContent`.
|
||||
|
||||
See [Search and question answering](../configuration/qa.md) and [picture processing](../configuration/processing.md#picture-handling).
|
||||
32
docs/chat.md
32
docs/chat.md
|
|
@ -13,11 +13,11 @@ haiku-rag chat --db /path/to/database.lancedb
|
|||
haiku-rag chat --model openai:gpt-4o
|
||||
```
|
||||
|
||||

|
||||

|
||||
|
||||
## How it works
|
||||
|
||||
The chat is a Pydantic AI agent with the `rag` [skill](skills/rag.md) attached by default. Each turn the agent decides which tool to call next, runs hybrid search against your documents, expands context around the hits, may issue further searches, and answers with citations. You see streaming text and a live indicator of which tool is running.
|
||||
The chat is a Pydantic AI agent with the [RAG capability](capabilities/rag.md) attached by default. A single capability loads eagerly; when both RAG and analysis are enabled, they remain deferred until the model chooses which one to load. Each turn the agent decides which tool to call next, runs hybrid search against your documents, expands context around the hits, may issue further searches, and answers with citations. You see streaming text and native tool events directly from Pydantic AI.
|
||||
|
||||
The session is in-memory for the lifetime of the TUI. Conversation history is kept across turns so follow-up questions reuse prior context. Citations are tracked per turn and inspectable via the command palette. Clearing the chat resets the session and the agent's memory.
|
||||
|
||||
|
|
@ -39,6 +39,12 @@ You can also render visual grounding from the CLI without launching the TUI:
|
|||
haiku-rag visualize <chunk_id>
|
||||
```
|
||||
|
||||
## Attaching images
|
||||
|
||||
Press `Ctrl+I` to open the image picker: a directory tree filtered to image files with a live preview. Selecting an image inserts an `[Image #N]` token at the cursor and attaches the image to your next message. Tokens delete as a unit with backspace or delete, and you can place them anywhere in the text to control where each image appears relative to your words.
|
||||
|
||||
Retrieval stays text-based; the images are sent to the model alongside your message, so the driving model needs `vision: true` in its configuration.
|
||||
|
||||
## Command palette
|
||||
|
||||
`Ctrl+P` opens the palette.
|
||||
|
|
@ -49,31 +55,33 @@ haiku-rag visualize <chunk_id>
|
|||
| Filter documents | Restrict searches to selected documents |
|
||||
| Show visual grounding | Visual grounding for a citation |
|
||||
| Database info | Document and chunk counts, storage stats |
|
||||
| View state | Current session state, citations, and intermediate tool results |
|
||||
|
||||
## Skills
|
||||
## Capabilities
|
||||
|
||||
The default skill is `rag`. Enable `analysis` when the question needs computation, aggregation, comparison across documents, or section-scoped reading that a single search can't deliver:
|
||||
The default capability is `rag`. Enable `analysis` when the question needs computation, aggregation, comparison across documents, or section-scoped reading that a single search can't deliver:
|
||||
|
||||
```bash
|
||||
# both skills (the agent routes between them)
|
||||
haiku-rag chat -s rag -s analysis
|
||||
# analysis instead of rag
|
||||
haiku-rag chat -c analysis
|
||||
|
||||
# analysis only
|
||||
haiku-rag chat -s analysis
|
||||
# both, which gives the model duplicate search and cite tools
|
||||
haiku-rag chat -c rag -c analysis
|
||||
```
|
||||
|
||||
The `analysis` skill mounts every document as a virtual filesystem at `/documents/{id}/` (with `metadata.json`, `content.txt`, `items.jsonl`, and `toc.json`) and runs Python in a sandboxed interpreter with `search` and `list_documents` as awaitable functions. It's the right choice for questions like:
|
||||
Prefer one. `analysis` searches and cites as well as computing, so pairing it with `rag`
|
||||
duplicates tools and budgets. See [Capabilities](capabilities/index.md#compose-an-agent).
|
||||
|
||||
The `analysis` capability mounts every document as a virtual filesystem at `/documents/{id}/` (with `metadata.json`, `content.txt`, `items.jsonl`, and `toc.json`) and runs Python in a sandboxed interpreter with `search` and `list_documents` as awaitable functions. It's the right choice for questions like:
|
||||
|
||||
- "How many of these documents mention X?"
|
||||
- "Summarize Section 5 of paper Y."
|
||||
- "Compare the experimental sections across these three reports."
|
||||
- "Which section discusses the proof of Theorem 4.10?"
|
||||
|
||||
For everyday Q&A, the rag skill alone is faster and cheaper. Attaching both lets the agent pick. See [Analysis skill](skills/analysis.md) for the full sandbox capabilities and worked code patterns.
|
||||
For everyday Q&A, RAG alone is faster and cheaper. Attaching both lets the agent pick. See [Analysis capability](capabilities/analysis.md).
|
||||
|
||||
## Document filter
|
||||
|
||||
Run "Filter documents" from the command palette to restrict searches to a subset. The filter applies to every search the agent runs for the rest of the session.
|
||||
|
||||
Chat also honors the global `--read-only` and `--before` flags. See the [CLI reference](cli.md) for details.
|
||||
Chat also honors the global `--read-only` flag. See the [CLI reference](cli.md) for details.
|
||||
|
|
|
|||
157
docs/cli.md
157
docs/cli.md
|
|
@ -7,12 +7,12 @@ The `haiku-rag` CLI provides complete document management functionality.
|
|||
|
||||
- `--config` - Specify custom configuration file
|
||||
- `--read-only` - Open database in read-only mode (blocks writes, skips upgrades)
|
||||
- `--before` - Query database as it existed before a datetime (implies `--read-only`)
|
||||
- `--db-name` - Name of a database from `lancedb.databases` to work on
|
||||
- `--version` / `-v` - Show version and exit
|
||||
|
||||
Per-command options:
|
||||
|
||||
- `--db` - Specify custom database path
|
||||
- `--db` - Open the database at this path, named by its stem, whatever the configuration places
|
||||
- `-h` - Show help for specific command
|
||||
|
||||
Example:
|
||||
|
|
@ -20,10 +20,12 @@ The `haiku-rag` CLI provides complete document management functionality.
|
|||
haiku-rag --config /path/to/config.yaml list
|
||||
haiku-rag --config /path/to/config.yaml list --db /path/to/custom.db
|
||||
haiku-rag --read-only search "query"
|
||||
haiku-rag --before "2025-01-15" search "query"
|
||||
haiku-rag --db-name papers list
|
||||
haiku-rag add -h
|
||||
```
|
||||
|
||||
With `lancedb.databases` configured, `search`, `ask`, `analyze`, `chat`, and `mcp` use the full set by default. Select one database for other commands with `--db-name` or `--db`. `settings`, `init-config`, and `download-models` do not open a database. See [Multiple Databases](configuration/storage.md#multiple-databases).
|
||||
|
||||
## Document Management
|
||||
|
||||
### Add Documents
|
||||
|
|
@ -163,11 +165,23 @@ Filter to specific documents:
|
|||
haiku-rag ask "What are the main findings?" --filter "uri LIKE '%paper%'"
|
||||
```
|
||||
|
||||
`ask` runs the [rag skill](skills/index.md) and always renders citations under the answer. When available, citations use the document title, otherwise they fall back to the URI.
|
||||
Attach images to the question, for example to check an image against indexed documents:
|
||||
```bash
|
||||
haiku-rag ask "Does this photo satisfy the spec in the design document?" --image photo.jpg
|
||||
```
|
||||
|
||||
`ask` runs the [RAG capability](capabilities/rag.md) and always renders citations under the answer. When available, citations use the document title, otherwise they fall back to the URI.
|
||||
|
||||
Citation text is truncated to a 300-character preview. To read the whole passage the model saw:
|
||||
```bash
|
||||
haiku-rag ask "What are the main findings?" --full-citations
|
||||
```
|
||||
|
||||
Flags:
|
||||
|
||||
- `--filter` / `-f`: Restrict searches to documents matching the filter (see [Filtering Search Results](python.md#filtering-search-results))
|
||||
- `--image`: Path to an image attached to the question (repeatable). Retrieval stays text-based; the model must have `vision: true` configured.
|
||||
- `--full-citations`: Show the full text of each citation instead of a truncated preview
|
||||
|
||||
## Analyze
|
||||
|
||||
|
|
@ -186,8 +200,10 @@ haiku-rag analyze "What is the total revenue?" --filter "title LIKE '%Financial%
|
|||
Flags:
|
||||
|
||||
- `--filter` / `-f`: SQL WHERE clause to restrict document access
|
||||
- `--image`: Path to an image attached to the question (repeatable). Requires `vision: true` on the analysis model.
|
||||
- `--full-citations`: Show the full text of each citation instead of a truncated preview
|
||||
|
||||
See [Analysis skill](skills/analysis.md) for details on capabilities and configuration.
|
||||
See [Analysis capability](capabilities/analysis.md) for details and configuration.
|
||||
|
||||
## Chat
|
||||
|
||||
|
|
@ -197,8 +213,8 @@ Launch an interactive chat session for multi-turn conversations:
|
|||
haiku-rag chat
|
||||
haiku-rag chat --db /path/to/database.lancedb
|
||||
|
||||
# Enable analysis skill (code execution)
|
||||
haiku-rag chat -s rag -s analysis
|
||||
# Enable the analysis capability (code execution)
|
||||
haiku-rag chat -c rag -c analysis
|
||||
```
|
||||
|
||||
!!! note
|
||||
|
|
@ -206,7 +222,7 @@ haiku-rag chat -s rag -s analysis
|
|||
|
||||
Flags:
|
||||
|
||||
- `--skill` / `-s`: Skills to enable. `rag` (default), `analysis`. Can be repeated for multiple skills.
|
||||
- `--capability` / `-c`: Capabilities to enable. `rag` (default), `analysis`. Can be repeated.
|
||||
|
||||
The chat interface provides:
|
||||
|
||||
|
|
@ -313,6 +329,7 @@ Checks include:
|
|||
- the configured embedding identity matches the stored settings
|
||||
- no database migrations are pending
|
||||
- the vector index covers all chunks
|
||||
- the full-text index covers the chunks it searches
|
||||
- near-identical documents (by embedding-centroid similarity) are grouped and reported, with the largest member flagged as the likely one to keep (advisory only, never deleted, tuned via `doctor.duplicates` in config)
|
||||
- API keys are set for configured providers
|
||||
|
||||
|
|
@ -324,7 +341,7 @@ It also probes the external endpoints the config uses and reports them under a P
|
|||
|
||||
SaaS providers (OpenAI, Anthropic, Cohere, Jina, ZeroEntropy, Voyage) are covered by the API-key check rather than a network probe. In-process local models (sentence-transformers, cross-encoder, jina-local) have no endpoint and are reported as such.
|
||||
|
||||
Each failure prints the command that fixes it (`rebuild`, `create-index`, `migrate`, `rebuild --set-embedder`). `doctor` makes no changes. It exits with status 1 when any check fails, so it can gate CI or monitoring.
|
||||
Each failure prints the command that fixes it (`rebuild`, `create-index`, `vacuum`, `migrate`, `rebuild --set-embedder`). `doctor` makes no changes. It exits with status 1 when any check fails, so it can gate CI or monitoring.
|
||||
|
||||
### Migrate Database
|
||||
|
||||
|
|
@ -382,18 +399,19 @@ haiku-rag create-index [--db /path/to/your.lancedb]
|
|||
|
||||
**Requirements:**
|
||||
- Minimum 256 chunks required for index creation (LanceDB training data requirement)
|
||||
- Creates an IVF_PQ index using the configured `search.vector_index_metric` (cosine/l2/dot)
|
||||
- Creates an IVF_PQ index using the configured `search.vector_index_metric` (cosine/l2)
|
||||
|
||||
**When to use:**
|
||||
- After ingesting documents (indexes are not created automatically)
|
||||
- After adding significant new data to rebuild the index
|
||||
- Use `haiku-rag info` to check index status and see how many chunks are indexed/unindexed
|
||||
- On a collection over 100,000 chunks (below that, brute-force kNN is exact and fast enough)
|
||||
- After substantial corpus growth, to retrain the centroids
|
||||
- Use `haiku-rag info` to check index status and see how many chunks are indexed/unindexed, or `haiku-rag doctor` for the same as a health check
|
||||
|
||||
See [Vector Indexing](configuration/storage.md#vector-indexing) for the measured accuracy and build cost.
|
||||
|
||||
**Search behavior:**
|
||||
- Without index: Brute-force kNN search (exact nearest neighbors, slower for large datasets)
|
||||
- With index: Fast ANN (approximate nearest neighbors) using IVF_PQ
|
||||
- With stale index: LanceDB combines indexed results (fast ANN) + brute-force kNN on unindexed rows
|
||||
- Performance degrades as more unindexed data accumulates
|
||||
- With index: ANN (approximate nearest neighbors) using IVF_PQ, tuned by `search.vector_nprobes`
|
||||
- Between a write and the next `optimize()`: LanceDB combines ANN over indexed rows with brute-force kNN over the remainder
|
||||
|
||||
### Rebuild Database
|
||||
|
||||
|
|
@ -459,9 +477,6 @@ haiku-rag mcp --port 9000
|
|||
|
||||
# Bind to all interfaces (containers, trusted LAN)
|
||||
haiku-rag mcp --host 0.0.0.0
|
||||
|
||||
# Read-only mode (no write tools)
|
||||
haiku-rag --read-only mcp
|
||||
```
|
||||
|
||||
See [MCP](mcp.md) for details. For continuous document ingestion
|
||||
|
|
@ -484,93 +499,53 @@ haiku-rag init-config [output_path]
|
|||
|
||||
If no path is specified, creates `haiku.rag.yaml` in the current directory.
|
||||
|
||||
## Create Skill
|
||||
## Tags
|
||||
|
||||
Generate a standalone skill package with an embedded database:
|
||||
A tag names the current database state. It is a logical snapshot composed of one LanceDB tag on each of the five tables, created from a single version snapshot.
|
||||
|
||||
```bash
|
||||
haiku-rag create-skill --name myskill --db /path/to/database.lancedb
|
||||
# Tag the current state, e.g. at deploy time or after an ingestion run
|
||||
haiku-rag tag create release-1
|
||||
|
||||
# List tags with the versions they point to
|
||||
haiku-rag tag list
|
||||
|
||||
# Delete a tag, releasing its versions for cleanup
|
||||
haiku-rag tag delete release-1
|
||||
```
|
||||
|
||||
The generated package is a pip-installable Python package that registers as a `haiku.skills` entry point.
|
||||
A tag present on every table is complete. A tag missing from some tables (created outside haiku.rag, or left behind by a failure) is partial. `tag list` marks partial tags. Partial tags can be listed and deleted but never restored.
|
||||
|
||||
### Options
|
||||
Create tags with other writers stopped. Tag creation coordinates writers within one process only; a writer in another process can commit between the per-table snapshot reads, and the tag then captures a mixed state.
|
||||
|
||||
| Flag | Description | Default |
|
||||
|------|-------------|---------|
|
||||
| `--name` | Skill name (lowercase alphanumeric and hyphens, required) | — |
|
||||
| `--db` | Path to LanceDB database to embed (required) | — |
|
||||
| `--description` | Skill description | Standard RAG description |
|
||||
| `--tools` | Comma-separated tool names, or `all` | `all` |
|
||||
| `--preamble` | Custom preamble for skill instructions | Standard RAG preamble |
|
||||
| `--config-file` | Path to `haiku.rag.yaml` to embed | None |
|
||||
| `--output` / `-o` | Output directory | Current directory |
|
||||
Tagged versions survive `vacuum`. Vacuum retains the oldest tagged version and every newer version; versions older than the oldest tag remain eligible for cleanup. Delete tags you no longer need so cleanup can advance.
|
||||
|
||||
### Available Tools
|
||||
### Restore
|
||||
|
||||
`cite`, `execute_code`, `get_document`, `list_documents`, `search`
|
||||
|
||||
### Example
|
||||
`tag restore` brings the database back to a tagged state:
|
||||
|
||||
```bash
|
||||
# Generate a skill with specific tools and custom preamble
|
||||
haiku-rag create-skill \
|
||||
--name medic \
|
||||
--db /path/to/medic.lancedb \
|
||||
--tools search,cite \
|
||||
--config-file /path/to/haiku.rag.yaml \
|
||||
--description "Military medic knowledge base" \
|
||||
--preamble "You are a military medic expert."
|
||||
|
||||
# Install the generated package
|
||||
uv pip install -e ./medic-skill
|
||||
|
||||
# Use with haiku-skills
|
||||
haiku-skills chat --use-entrypoints --skill medic
|
||||
haiku-rag tag restore release-1
|
||||
```
|
||||
|
||||
### Generated Package Structure
|
||||
Restore changes the live state. It is not a read-only view: each table gets a new latest version equal to the tagged one, and reads and writes continue from there. Versions written after the tag remain in history until vacuum removes them.
|
||||
|
||||
```
|
||||
{name}-skill/
|
||||
├── pyproject.toml
|
||||
└── {name}_skill/
|
||||
├── __init__.py # create_skill() entry point
|
||||
├── SKILL.md # Skill metadata and instructions
|
||||
└── assets/
|
||||
├── {name}.lancedb/ # Embedded database
|
||||
└── haiku.rag.yaml # Optional config
|
||||
```
|
||||
|
||||
## Time Travel
|
||||
|
||||
LanceDB maintains version history for tables, enabling you to query the database as it existed at a previous point in time. This is useful for:
|
||||
|
||||
- **Debugging**: Investigate data before a problematic change
|
||||
- **Auditing**: Verify what knowledge was available when a support ticket was filed
|
||||
|
||||
### Query Historical State
|
||||
|
||||
Use `--before` to query the database as it existed before a specific datetime:
|
||||
Before changing anything, restore creates a complete safety tag (`before-restore-<timestamp>`) for the current state and reports it, so you always have a named path back:
|
||||
|
||||
```bash
|
||||
# Query documents as of January 15, 2025
|
||||
haiku-rag --before "2025-01-15" list
|
||||
|
||||
# Search historical state
|
||||
haiku-rag --before "2025-01-15T14:30:00" search "machine learning"
|
||||
|
||||
# Ask questions against historical data
|
||||
haiku-rag --before "2025-01-15" ask "What documents existed?"
|
||||
haiku-rag tag create release-1 --db /path/to/db.lancedb
|
||||
# Stop all writers before either restore.
|
||||
haiku-rag tag restore release-1 --db /path/to/db.lancedb --yes
|
||||
haiku-rag tag list --db /path/to/db.lancedb
|
||||
haiku-rag tag restore before-restore-YYYYMMDDTHHMMSSZ --db /path/to/db.lancedb --yes
|
||||
```
|
||||
|
||||
Supported datetime formats:
|
||||
Restore is a maintenance operation:
|
||||
|
||||
- ISO 8601: `2025-01-15T14:30:00`, `2025-01-15T14:30:00Z`, `2025-01-15T14:30:00+00:00`
|
||||
- Date only: `2025-01-15` (interpreted as start of day)
|
||||
|
||||
!!! note
|
||||
Time travel mode automatically enables read-only mode. You cannot modify the database while viewing historical state.
|
||||
- Stop all ingestion and other writers before restoring and keep them stopped until it finishes.
|
||||
- The operation is coordinated but not transactionally atomic across tables. On failure it attempts to roll back to the pre-restore state and reports whether the rollback succeeded.
|
||||
- `--yes` only skips the confirmation prompt. It provides no locking and no concurrent-writer protection.
|
||||
- Restore never migrates. Restoring a tag from an older haiku.rag version completes normally, and the next open reports the required migration. Run `haiku-rag migrate` explicitly.
|
||||
|
||||
### Version History
|
||||
|
||||
|
|
@ -587,20 +562,18 @@ haiku-rag history --table documents
|
|||
haiku-rag history --limit 10
|
||||
```
|
||||
|
||||
Output shows version numbers and timestamps, sorted newest first:
|
||||
Output shows version numbers and timestamps, sorted newest first, with tags marked:
|
||||
|
||||
```
|
||||
Version History
|
||||
|
||||
documents
|
||||
v5: 2025-01-15 14:30:00
|
||||
v5: 2025-01-15 14:30:00 <- release-1
|
||||
v4: 2025-01-14 10:00:00
|
||||
v3: 2025-01-13 09:15:00
|
||||
|
||||
chunks
|
||||
v8: 2025-01-15 14:30:00
|
||||
v8: 2025-01-15 14:30:00 <- release-1
|
||||
v7: 2025-01-14 10:00:00
|
||||
...
|
||||
```
|
||||
|
||||
Use the timestamps from `history` to construct `--before` queries.
|
||||
|
|
|
|||
|
|
@ -61,7 +61,7 @@ embeddings:
|
|||
qa:
|
||||
model:
|
||||
provider: ollama
|
||||
name: gpt-oss
|
||||
name: qwen3.8
|
||||
enable_thinking: true
|
||||
```
|
||||
|
||||
|
|
@ -85,8 +85,8 @@ ingester:
|
|||
delete_orphans: true
|
||||
|
||||
lancedb:
|
||||
uri: "" # Empty for local, or db://, s3://, az://, gs://
|
||||
api_key: ""
|
||||
databases: {} # Name-to-location map; empty places haiku.rag under data_dir
|
||||
api_key: "" # LanceDB Cloud (db://) credentials
|
||||
region: ""
|
||||
|
||||
embeddings:
|
||||
|
|
@ -96,22 +96,24 @@ embeddings:
|
|||
vector_dim: 2560
|
||||
|
||||
reranking:
|
||||
# Omit this section, or set `model: null`, to disable reranking.
|
||||
model:
|
||||
provider: "" # Empty to disable, or cross-encoder, cohere, zeroentropy, vllm
|
||||
name: ""
|
||||
provider: cross-encoder # cross-encoder, cohere, zeroentropy, vllm, jina, jina-local
|
||||
name: cross-encoder/ms-marco-MiniLM-L-6-v2
|
||||
multimodal: false # vllm only: send picture chunks to the reranker as images
|
||||
|
||||
qa:
|
||||
model:
|
||||
provider: ollama
|
||||
name: gpt-oss
|
||||
name: qwen3.8
|
||||
enable_thinking: true
|
||||
temperature: 0.3
|
||||
max_searches: 3
|
||||
max_searches: 5
|
||||
|
||||
search:
|
||||
limit: 10 # Default number of results to return
|
||||
limit: 5 # Default number of results to return
|
||||
max_context_chars: 5000 # Maximum characters in expanded context
|
||||
vector_index_metric: cosine # cosine, l2, or dot
|
||||
vector_index_metric: cosine # cosine or l2
|
||||
vector_refine_factor: 30
|
||||
|
||||
doctor:
|
||||
|
|
@ -120,7 +122,7 @@ doctor:
|
|||
min_chunks: 3 # documents with fewer chunks are excluded
|
||||
|
||||
prompts:
|
||||
domain_preamble: "" # Prepended to skill instructions
|
||||
domain_preamble: "" # Prepended to capability instructions
|
||||
|
||||
processing:
|
||||
converter: docling-local # docling-local or docling-serve
|
||||
|
|
@ -133,7 +135,7 @@ processing:
|
|||
auto_title: false # Auto-generate titles on ingestion
|
||||
title_model:
|
||||
provider: ollama
|
||||
name: gpt-oss
|
||||
name: qwen3.8
|
||||
enable_thinking: false
|
||||
temperature: 0.3
|
||||
max_tokens: 100
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ processing:
|
|||
auto_title: false # Auto-generate titles on ingestion
|
||||
title_model: # LLM for title generation (fallback)
|
||||
provider: ollama
|
||||
name: gpt-oss
|
||||
name: qwen3.8
|
||||
enable_thinking: false
|
||||
|
||||
# Conversion options (works with both local and remote converters)
|
||||
|
|
@ -54,7 +54,7 @@ processing:
|
|||
picture_description:
|
||||
model:
|
||||
provider: ollama
|
||||
name: ministral-3
|
||||
name: qwen3.8
|
||||
pictures: image # none | description | image
|
||||
```
|
||||
|
||||
|
|
@ -270,7 +270,7 @@ processing:
|
|||
picture_description: # only consulted when pictures == "description"
|
||||
model:
|
||||
provider: ollama # any OpenAI-compatible /v1/chat/completions provider
|
||||
name: ministral-3
|
||||
name: qwen3.8
|
||||
timeout: 90
|
||||
max_tokens: 200
|
||||
```
|
||||
|
|
@ -294,7 +294,7 @@ Three independent settings drive ingest, retrieval, and QA:
|
|||
|---|---|---|
|
||||
| `processing.pictures` | Generate and/or describe pictures at ingest? | `none` / `description` / `image` (default) |
|
||||
| `embeddings.model.multimodal` | Can the embedder index image content? | `false` (default, text-only) / `true` (supported on `vllm`, `voyageai`, `cohere`) |
|
||||
| `qa.model.vision` | Can the QA model interpret images? | `false` (default) / `true` |
|
||||
| `qa.model.vision` | Can the QA model interpret images? | `false` / `true` (default) |
|
||||
|
||||
The Embedder column below is driven by `embeddings.model.multimodal`, not the provider name — a vision-capable model under a text-only configuration still indexes no images, and an image-only document then produces zero chunks. See [Multimodal embedders](providers.md#multimodal-embedders).
|
||||
|
||||
|
|
@ -313,7 +313,7 @@ The Embedder column below is driven by `embeddings.model.multimodal`, not the pr
|
|||
- `qa.model.vision: false` — text chunks only (descriptions, when present, answer figure questions in prose).
|
||||
- `qa.model.vision: true` — text chunks + raw picture bytes via `BinaryContent`. The model reads figures directly. Requires `pictures != none` so the bytes exist.
|
||||
|
||||
`qa.model.vision` is independent of ingestion. Flipping it never requires reingesting. Setting `vision: true` against a text-only model causes silent acceptance and confabulation on Ollama and a 400 on OpenAI. Default `false` is the safe choice.
|
||||
`qa.model.vision` is independent of ingestion. Flipping it never requires reingesting. It declares what the model can read: the default `qwen3.8` is vision-capable, so the default is `true`. Set it `false` when pointing `qa.model` at a text-only model, where `true` causes silent acceptance and confabulation on Ollama and a 400 on OpenAI.
|
||||
|
||||
**Recommended combinations:**
|
||||
|
||||
|
|
@ -368,7 +368,7 @@ processing:
|
|||
auto_title: true
|
||||
title_model:
|
||||
provider: ollama
|
||||
name: gpt-oss
|
||||
name: qwen3.8
|
||||
enable_thinking: false
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -1,30 +1,32 @@
|
|||
# Prompt Customization
|
||||
|
||||
Customize the prompts used by haiku.rag's skills to better match your domain and use case.
|
||||
Customize the prompts used by haiku.rag's capabilities to match your domain.
|
||||
|
||||
## Configuration
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
# Domain context prepended to skill instructions
|
||||
# Domain context prepended to capability instructions
|
||||
domain_preamble: |
|
||||
This knowledge base contains technical documentation for the Helios solar panel
|
||||
system, including installation manuals, maintenance procedures, and safety guidelines.
|
||||
Questions about "the system" or unqualified specs refer to the Helios panel.
|
||||
|
||||
# VLM prompt for image description during conversion (optional)
|
||||
picture_description: null # Uses default prompt
|
||||
# VLM prompt for image description during conversion.
|
||||
# Omit the key to use the built-in prompt.
|
||||
picture_description: |
|
||||
Describe this figure in two sentences, naming any axis labels and units.
|
||||
```
|
||||
|
||||
## Domain Preamble
|
||||
|
||||
The `domain_preamble` field provides **domain context** prepended to the rag and rag-analysis skill instructions. Use this to:
|
||||
The `domain_preamble` field provides **domain context** prepended to the RAG and analysis capability instructions. Use this to:
|
||||
|
||||
- Describe what the knowledge base contains
|
||||
- Clarify domain-specific terminology
|
||||
- Provide context that helps the model interpret ambiguous queries
|
||||
|
||||
**Important:** `domain_preamble` is for domain context, not behavioral instructions. Descriptions of subject matter, terminology, and content scope belong here. Behavioral guidance (tone, response style, formatting rules) lives in the skill's SKILL.md. Fork the skill via `haiku-rag create-skill` to customize behavior.
|
||||
**Important:** `domain_preamble` is for domain context, not behavioral instructions. Descriptions of subject matter, terminology, and content scope belong here. Applications can add behavioral guidance through normal Pydantic AI agent instructions.
|
||||
|
||||
**Example:**
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ haiku.rag supports multiple AI providers for embeddings, question answering, and
|
|||
|
||||
## Model Settings
|
||||
|
||||
Configure model behavior for the `qa` and `analysis` skills. These settings apply to any provider that supports them.
|
||||
Configure model behavior for the `qa` and `analysis` capabilities. These settings apply to any provider that supports them.
|
||||
|
||||
### Basic Settings
|
||||
|
||||
|
|
@ -15,7 +15,7 @@ Configure model behavior for the `qa` and `analysis` skills. These settings appl
|
|||
qa:
|
||||
model:
|
||||
provider: ollama
|
||||
name: gpt-oss
|
||||
name: qwen3.8
|
||||
temperature: 0.3
|
||||
max_tokens: 500
|
||||
```
|
||||
|
|
@ -29,8 +29,32 @@ qa:
|
|||
- **max_tokens**: Maximum tokens in response. Default: unset (provider default), except title generation (100).
|
||||
- **enable_thinking**: Control reasoning behavior (see below)
|
||||
- **base_url**: Custom endpoint for OpenAI-compatible servers (vLLM, LM Studio, etc.)
|
||||
- **api_key**: Key for this endpoint, overriding the provider's environment variable (see [Per-endpoint API keys](#per-endpoint-api-keys))
|
||||
- **extra_body**: Raw dict forwarded to the model SDK (see [Raw Provider Pass-through](#raw-provider-pass-through))
|
||||
|
||||
### Per-endpoint API keys
|
||||
|
||||
The `openai` provider reads `OPENAI_API_KEY`, so several `openai`-compatible endpoints in one config would otherwise share a single key. Set `api_key` per model to give each its own, and keep the secret in the environment with [variable expansion](index.md#environment-variables):
|
||||
|
||||
```yaml
|
||||
qa:
|
||||
model:
|
||||
provider: openai
|
||||
name: some-model
|
||||
base_url: https://vendor-a.example/v1
|
||||
api_key: ${VENDOR_A_KEY}
|
||||
|
||||
embeddings:
|
||||
model:
|
||||
provider: openai
|
||||
name: some-embedding-model
|
||||
vector_dim: 1024
|
||||
base_url: https://vendor-b.example/v1
|
||||
api_key: ${VENDOR_B_KEY}
|
||||
```
|
||||
|
||||
`api_key` is honored on the `openai` and `ollama` providers, on `vllm` embedders and rerankers, and on the picture-description VLM endpoint (which otherwise falls back to `OPENAI_API_KEY` only for the public OpenAI endpoint, never for a custom `base_url`). Other providers (`anthropic`, `cohere`, `voyageai`, …) reach their vendor SDK by name and read their own environment variable; setting `api_key` there raises rather than being dropped silently.
|
||||
|
||||
### Thinking Control
|
||||
|
||||
The `enable_thinking` setting controls whether models use explicit reasoning steps before answering.
|
||||
|
|
@ -54,8 +78,8 @@ See the [Pydantic AI thinking documentation](https://ai.pydantic.dev/thinking/)
|
|||
- **Anthropic**: All Claude models
|
||||
- **Google**: Gemini models with thinking support
|
||||
- **Groq**: Models with reasoning capabilities
|
||||
- **Bedrock**: Claude, OpenAI, and Qwen models
|
||||
- **Ollama**: Models supporting reasoning (gpt-oss, etc.)
|
||||
- **Bedrock**: Claude, Qwen, and `gpt-oss` models. Bedrock Converse does not serve the proprietary OpenAI models, so configuring one raises an error. Reach those through `provider: bedrock-mantle`.
|
||||
- **Ollama**: Any model with a thinking capability. `enable_thinking` maps to `reasoning_effort`: `false` sends `none` (`low` for `gpt-oss`, whose template has no `none` level), `true` sends `high`.
|
||||
- **vLLM**: Models with a pydantic-ai reasoning profile (gpt-oss). Qwen3, Gemma, and similar templates ignore the OpenAI `reasoning_effort` that `enable_thinking` translates to — use [`extra_body`](#raw-provider-pass-through) to drive them.
|
||||
- **LM Studio**: Models supporting reasoning (gpt-oss, etc.)
|
||||
|
||||
|
|
@ -63,6 +87,9 @@ See the [Pydantic AI thinking documentation](https://ai.pydantic.dev/thinking/)
|
|||
- Enable for QA, complex reasoning, and mathematical problems
|
||||
- Disable for speed-critical applications, title generation, and simple tasks
|
||||
|
||||
!!! note "Anthropic thinking and max_tokens"
|
||||
Anthropic requires `max_tokens` to exceed the thinking budget, and `enable_thinking: true` requests Pydantic AI's default budget of 10000 tokens. Set `max_tokens` above 10000 on Claude models that use budget-based thinking, or leave it unset on Sonnet 4.6+ and Opus 4.6+, which use adaptive thinking instead of a budget.
|
||||
|
||||
!!! note "vLLM-served models without a reasoning profile"
|
||||
On `provider: openai` with a custom `base_url`, `enable_thinking` only takes effect for models whose pydantic-ai profile advertises reasoning support (o-series, gpt-5, gpt-oss). For other vLLM-served models (Qwen3, Gemma family, …) the field is a silent no-op. Reach the chat template's thinking switch directly via [`extra_body`](#raw-provider-pass-through).
|
||||
|
||||
|
|
@ -100,11 +127,11 @@ qa:
|
|||
|
||||
Same mechanism, opposite direction. Without `extra_body` the Gemma-4 chat template defaults to non-thinking and dumps a verbose answer straight into `content`. With it on, vLLM (started with `--reasoning-parser`) populates the parsed `reasoning` field and leaves `content` as the concise final answer.
|
||||
|
||||
**Provider support:** honored by openai, ollama, anthropic, and groq via pydantic-ai's `ModelSettings.extra_body`. Silently ignored by gemini and bedrock.
|
||||
**Provider support:** honored by openai, ollama, anthropic, and groq via pydantic-ai's `ModelSettings.extra_body`. Silently ignored by google and bedrock.
|
||||
|
||||
## Embedding Providers
|
||||
|
||||
Embedding models require three settings: `provider`, `name`, and `vector_dim`. Optionally, use `base_url` for OpenAI-compatible servers.
|
||||
Embedding models require three settings: `provider`, `name`, and `vector_dim`. Optionally, use `base_url` for OpenAI-compatible servers and [`api_key`](#per-endpoint-api-keys) for the key that endpoint expects.
|
||||
|
||||
### Batch Size
|
||||
|
||||
|
|
@ -284,7 +311,7 @@ Configure which LLM provider to use for question answering. Any provider and mod
|
|||
qa:
|
||||
model:
|
||||
provider: ollama
|
||||
name: gpt-oss
|
||||
name: qwen3.8
|
||||
```
|
||||
|
||||
The Ollama base URL can be configured via the `OLLAMA_BASE_URL` environment variable, config file, or defaults to `http://localhost:11434`:
|
||||
|
|
@ -366,7 +393,7 @@ Any provider supported by Pydantic AI can be used. Examples:
|
|||
# Google Gemini
|
||||
qa:
|
||||
model:
|
||||
provider: gemini
|
||||
provider: google
|
||||
name: gemini-1.5-flash
|
||||
|
||||
# Groq
|
||||
|
|
@ -388,7 +415,7 @@ See the [Pydantic AI documentation](https://ai.pydantic.dev/models/) for the com
|
|||
|
||||
Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (10x the requested limit) and then reranks them to return the most relevant results.
|
||||
|
||||
Reranking is **disabled by default** (`provider: ""`) for faster searches. You can enable it by configuring one of the providers below.
|
||||
Reranking is **disabled by default** for faster searches: there is no `reranking.model`. Enable it by configuring one of the providers below, and disable it again by removing the section or setting `model: null`.
|
||||
|
||||
### Cohere
|
||||
|
||||
|
|
@ -444,11 +471,26 @@ For high-performance local reranking using dedicated reranking models:
|
|||
reranking:
|
||||
model:
|
||||
provider: vllm
|
||||
name: mixedbread-ai/mxbai-rerank-base-v2
|
||||
base_url: http://localhost:8001
|
||||
name: Qwen/Qwen3-Reranker-4B
|
||||
base_url: http://localhost:8001/v1
|
||||
```
|
||||
|
||||
**Note:** vLLM reranking uses the `/v1/rerank` API endpoint. You need to run a vLLM server separately with a reranking model loaded.
|
||||
**Note:** vLLM reranking posts to the `/v1/rerank` endpoint. As with the embedder, `base_url` may be written with or without the `/v1` path. You need to run a vLLM server separately with a reranking model loaded.
|
||||
|
||||
#### Multimodal reranking
|
||||
|
||||
When serving a vision reranker (for example `nvidia/llama-nemotron-rerank-vl-1b-v2`), set `multimodal: true` to score picture chunks by their image bytes in addition to their description text:
|
||||
|
||||
```yaml
|
||||
reranking:
|
||||
multimodal: true
|
||||
model:
|
||||
provider: vllm
|
||||
name: nvidia/llama-nemotron-rerank-vl-1b-v2
|
||||
base_url: http://localhost:8001/v1
|
||||
```
|
||||
|
||||
Picture chunks are sent as image documents (base64 data URIs) alongside plain text documents in the same rerank request. The flag is supported on the vllm provider only, and the served model must accept multimodal inputs.
|
||||
|
||||
### Jina AI
|
||||
|
||||
|
|
@ -506,7 +548,7 @@ Then configure with any HuggingFace model id:
|
|||
reranking:
|
||||
model:
|
||||
provider: cross-encoder
|
||||
name: mixedbread-ai/mxbai-rerank-base-v2
|
||||
name: Qwen/Qwen3-Reranker-0.6B
|
||||
```
|
||||
|
||||
Other tested models: `BAAI/bge-reranker-v2-m3`, `Qwen/Qwen3-Reranker-0.6B`, `cross-encoder/ms-marco-MiniLM-L-6-v2`. Any model exposed as a `sentence_transformers.CrossEncoder` works.
|
||||
Other tested models: `BAAI/bge-reranker-v2-m3`, `cross-encoder/ms-marco-MiniLM-L-6-v2`. Any model exposed as a `sentence_transformers.CrossEncoder` works.
|
||||
|
|
|
|||
|
|
@ -6,11 +6,11 @@ Configure search behavior and context expansion:
|
|||
|
||||
```yaml
|
||||
search:
|
||||
limit: 10 # Default number of results to return
|
||||
limit: 5 # Default number of results to return
|
||||
max_context_chars: 5000 # Maximum characters in expanded context
|
||||
```
|
||||
|
||||
- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, and QA. Default: 10
|
||||
- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, and QA. Default: 5
|
||||
- **max_context_chars**: Hard limit on total characters in expanded content. Default: 5000.
|
||||
|
||||
Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. Picture and table matches are exempt: they return their enclosing section as-is and never cross section boundaries. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
|
||||
|
|
@ -20,29 +20,29 @@ Context expansion is automatic and section-aware. For structured documents (with
|
|||
|
||||
## Question Answering Configuration
|
||||
|
||||
Configure the rag skill (used by `client.ask`, `haiku-rag ask`, and the MCP `ask_question` tool):
|
||||
Configure the RAG capability (used by `client.ask` and `haiku-rag ask`):
|
||||
|
||||
```yaml
|
||||
qa:
|
||||
model:
|
||||
provider: ollama
|
||||
name: gpt-oss
|
||||
name: qwen3.8
|
||||
enable_thinking: true
|
||||
temperature: 0.3 # Default: 0.3
|
||||
vision: false # Set true for vision-capable models
|
||||
max_searches: 3 # Maximum search tool calls per question
|
||||
vision: true # Set false for text-only models
|
||||
max_searches: 5 # Maximum search units per question
|
||||
```
|
||||
|
||||
- **model**: LLM configuration (see [Providers](providers.md#model-settings))
|
||||
- **model.vision**: Set to `true` for vision-capable models (`qwen2.5vl`, `qwen3.6`, `gpt-4o`, `claude-sonnet`, …). The skill's `search` tool only attaches picture bytes (`BinaryContent`) to its `ToolReturn` when this is `true`, otherwise picture bytes are withheld. See [Pictures × embedder × QA model](processing.md#pictures-embedder-qa-model-how-the-pieces-compose) for the full matrix.
|
||||
- **max_searches**: Maximum number of search tool calls the rag skill can make per question (default: 3)
|
||||
- **model.vision**: Set to `true` for vision-capable models (`qwen2.5vl`, `qwen3.6`, `gpt-4o`, `claude-sonnet`, …). The capability's `search` tool only attaches picture bytes (`BinaryContent`) to its `ToolReturn` when this is `true`, otherwise picture bytes are withheld. See [Pictures × embedder × QA model](processing.md#pictures-embedder-qa-model-how-the-pieces-compose) for the full matrix.
|
||||
- **max_searches**: Maximum number of search units a capability can spend per question (default: 5). Up to three searches emitted in the same model response share one unit, so a model that rephrases its query in one response spends one unit. A search in a later response starts a new unit, as does each further group of three within one response. Shared by the RAG and analysis capabilities. Searches in one response also deduplicate their returns: evidence a sibling search already showed collapses to a reference line, and each picture attaches once per response.
|
||||
|
||||
!!! note "Thinking on vLLM"
|
||||
`enable_thinking` only applies to models with a pydantic-ai reasoning profile (o-series, gpt-5, gpt-oss). For other vLLM-served models such as Qwen3 or the Gemma family, the field is a silent no-op — set the chat template switch via [`extra_body`](providers.md#raw-provider-pass-through) instead.
|
||||
|
||||
## Analysis Configuration
|
||||
|
||||
Configure the analysis skill:
|
||||
Configure the analysis capability:
|
||||
|
||||
```yaml
|
||||
analysis:
|
||||
|
|
@ -50,14 +50,14 @@ analysis:
|
|||
provider: anthropic
|
||||
name: claude-sonnet-4-20250514
|
||||
temperature: 0.0 # Default: 0.0 (deterministic for code generation)
|
||||
code_timeout: 60.0 # Max seconds for code execution
|
||||
code_timeout: 60.0 # Per call: compute stops, no read or search starts past it
|
||||
max_output_chars: 50000 # Truncate output after this many chars
|
||||
max_executions: 15 # Max execute_code calls per question
|
||||
```
|
||||
|
||||
- **model**: LLM configuration (see [Providers](providers.md#model-settings)). When unset, falls back to `qa.model`.
|
||||
- **code_timeout**: Maximum seconds for each code execution (default: 60)
|
||||
- **code_timeout**: Seconds a single `execute_code` call has (default: 60). Past it the sandbox starts no further host call, a document read or an in-code `search()` / `list_documents()`; one already running finishes. Code that computes without host calls is killed by the worker watchdog at the same limit. `code_timeout * max_executions` is the cumulative ceiling across all calls in one question.
|
||||
- **max_output_chars**: Truncate code output after this many characters (default: 50000)
|
||||
- **max_executions**: Maximum `execute_code` calls per question before the skill is told to answer from what it has (default: 15)
|
||||
- **max_executions**: Maximum `execute_code` calls per question before the capability is told to answer from what it has (default: 15)
|
||||
|
||||
See [Analysis skill](../skills/analysis.md) for usage details.
|
||||
See [Analysis capability](../capabilities/analysis.md) for usage details.
|
||||
|
|
|
|||
|
|
@ -1,5 +1,33 @@
|
|||
# Database and Storage
|
||||
|
||||
## Operational constraints
|
||||
|
||||
Four things to know before deploying.
|
||||
|
||||
**Run one writer per database.** This is a haiku.rag constraint, not a LanceDB
|
||||
one. A write that spans several tables is serialized by an in-process lock and
|
||||
rolled back by restoring each table to the version it had when the write started.
|
||||
Both are process-local: a second writing process can commit between that snapshot
|
||||
and the mutation, and a rollback would then revert its work along with ours. Run
|
||||
a single writer, either the [`haiku-ingester`](../ingester.md) service or your own
|
||||
application. Read-only consumers are unrestricted.
|
||||
|
||||
**Readers lag by an interval.** A connection always sees its own writes. It sees
|
||||
another process's writes after `lancedb.read_consistency_interval_seconds`
|
||||
(default 30).
|
||||
|
||||
**Migrate after an upgrade that changes the schema.** `haiku-rag migrate` applies
|
||||
pending migrations in place, and `haiku-rag info` lists what is pending. A
|
||||
release that needs it says so in the [changelog](../changelog.md).
|
||||
|
||||
**The embedding dimension is fixed per database.** Every chunk vector has the
|
||||
dimension the database was created with. Changing `embeddings.model.vector_dim`
|
||||
raises `ConfigMismatchError` on open, because stored vectors cannot be compared
|
||||
against new ones. Changing the provider or model name while keeping the dimension
|
||||
warns on a read-only open and raises on a writable one. `haiku-rag rebuild
|
||||
--set-embedder` adopts the new identity without re-embedding, and `haiku-rag
|
||||
rebuild --embed-only` re-embeds against the new model.
|
||||
|
||||
## Local Storage
|
||||
|
||||
By default, `haiku.rag` uses a local LanceDB database:
|
||||
|
|
@ -18,6 +46,44 @@ storage:
|
|||
!!! warning "Vacuum Retention Threshold"
|
||||
The `vacuum_retention_seconds` value should be larger than the typical time it takes to process and write a document. If a concurrent operation is in progress while vacuum runs, setting this value too low can cause race conditions where vacuum removes table versions that an in-flight operation still needs. The default of 86400 seconds (1 day) is conservative and safe for most use cases.
|
||||
|
||||
### Vacuum Memory Requirements
|
||||
|
||||
Vacuum compacts small data files into larger ones. LanceDB targets roughly one million rows per fragment, which a `documents` table holding multi-megabyte docling blobs never reaches, so each vacuum that follows new documents re-merges the whole existing fragment rather than only the new ones. Peak memory therefore scales with the total size of the `documents` table, not with how much was added.
|
||||
|
||||
Measured peak resident memory is about 5x the size of the `documents` table's data files. An 8.8 GB table peaked at 48.7 GB. Plan for **6x the size of `documents/` on disk** as available RAM, or the vacuum will be killed by the OOM killer partway through.
|
||||
|
||||
Check the current size with:
|
||||
|
||||
```bash
|
||||
du -sh /path/to/database.lancedb/documents.lance
|
||||
```
|
||||
|
||||
If that number times six exceeds available RAM, use one of:
|
||||
|
||||
- Reduce `images_scale` (see [Image Settings](processing.md#image-settings)). Rendered page rasters dominate the size of `documents`, and their byte cost falls with the square of the scale factor.
|
||||
- Set `generate_page_images: false` if visual grounding through `visualize_chunk()` is not needed. This removes page rasters entirely.
|
||||
- Set `auto_vacuum: false` and run `haiku-rag vacuum` manually when the machine is otherwise idle, so the peak does not land alongside ingestion.
|
||||
|
||||
Vacuum also folds new rows into the full-text index. Search stays correct without it but scans the uncovered rows on every query. `haiku-rag doctor` reports the coverage.
|
||||
|
||||
This is an upstream limitation rather than a `haiku.rag` setting. Compaction bounds itself by row count instead of bytes, and LanceDB's async API exposes no batch size or fragment target to override it. Tracked at [lancedb/lancedb#2325](https://github.com/lancedb/lancedb/issues/2325). The requirement above will drop once compaction batches by bytes.
|
||||
|
||||
### Placing the Database
|
||||
|
||||
`lancedb.databases` maps a name to a location, a local path or a URI, and is the one way to place databases. With nothing configured, the database is the entry `haiku.rag` at `<storage.data_dir>/haiku.rag.lancedb`. To put one database somewhere else, name it:
|
||||
|
||||
```yaml
|
||||
lancedb:
|
||||
databases:
|
||||
notes: /data/notes.lancedb
|
||||
```
|
||||
|
||||
The name is what `source` carries in search results, citations and documents, and what `--db-name` and `sources` select. The default database answers to `haiku.rag`.
|
||||
|
||||
An explicit `--db PATH` on the command line opens that database instead, named by the path's stem, whatever is configured. From Python, `db_path` places the database only where the configuration places none: beside `lancedb.databases` it raises `AmbiguousDatabaseError`.
|
||||
|
||||
A value with no scheme is a local path wherever it is configured, so `haiku-rag init` creates it and every command that opens an existing database requires it to exist. A mistyped path fails rather than becoming a new empty database.
|
||||
|
||||
## Database Creation
|
||||
|
||||
Databases must be explicitly created before use:
|
||||
|
|
@ -44,28 +110,31 @@ async with HaikuRAG(create=True) as client:
|
|||
|
||||
The [default location](index.md#configuration-file-locations) is platform-specific (e.g., `~/Library/Application Support/haiku.rag/` on macOS).
|
||||
|
||||
Operations on non-existent databases raise `FileNotFoundError`. This prevents accidental database creation from typos or misconfigured paths.
|
||||
Opening a nonexistent local database given as a path raises `FileNotFoundError`, naming the path. This prevents accidental database creation from typos or misconfigured paths. A configured or default database raises `SourceUnavailableError` instead, naming the database and not its location.
|
||||
|
||||
## Remote Storage
|
||||
|
||||
For remote storage, use the `lancedb` settings with various backends:
|
||||
For remote storage, give the database a URI as its location. Credentials and storage options are connection settings, shared by every database in the configuration:
|
||||
|
||||
```yaml
|
||||
# LanceDB Cloud
|
||||
lancedb:
|
||||
uri: db://your-database-name
|
||||
databases:
|
||||
papers: db://your-database-name
|
||||
api_key: your-api-key
|
||||
region: us-west-2 # optional
|
||||
region: us-west-2
|
||||
|
||||
# Amazon S3
|
||||
lancedb:
|
||||
uri: s3://my-bucket/my-table
|
||||
databases:
|
||||
papers: s3://my-bucket/my-table
|
||||
storage_options:
|
||||
region: us-east-1
|
||||
|
||||
# Amazon S3 with explicit credentials
|
||||
lancedb:
|
||||
uri: s3://my-bucket/my-table
|
||||
databases:
|
||||
papers: s3://my-bucket/my-table
|
||||
storage_options:
|
||||
aws_access_key_id: YOUR_ACCESS_KEY
|
||||
aws_secret_access_key: YOUR_SECRET_KEY
|
||||
|
|
@ -73,7 +142,8 @@ lancedb:
|
|||
|
||||
# S3-compatible (SeaweedFS, Tigris, etc.)
|
||||
lancedb:
|
||||
uri: s3://my-bucket/my-table
|
||||
databases:
|
||||
papers: s3://my-bucket/my-table
|
||||
storage_options:
|
||||
endpoint: http://localhost:8333
|
||||
aws_access_key_id: YOUR_ACCESS_KEY
|
||||
|
|
@ -83,44 +153,162 @@ lancedb:
|
|||
|
||||
# Azure Blob Storage
|
||||
lancedb:
|
||||
uri: az://my-container/my-table
|
||||
databases:
|
||||
papers: az://my-container/my-table
|
||||
|
||||
# Google Cloud Storage
|
||||
lancedb:
|
||||
uri: gs://my-bucket/my-table
|
||||
databases:
|
||||
papers: gs://my-bucket/my-table
|
||||
|
||||
# HDFS
|
||||
lancedb:
|
||||
uri: hdfs://namenode:port/path/to/table
|
||||
databases:
|
||||
papers: hdfs://namenode:port/path/to/table
|
||||
```
|
||||
|
||||
- **LanceDB Cloud** (`db://`): Requires `api_key` and `region`. Table optimization and indexing are managed server-side.
|
||||
- **Object storage** (`s3://`, `gs://`, `az://`, `hdfs://`): Uses `storage_options` for credentials and endpoint configuration. Authentication can also be provided via environment variables (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, etc.) or cloud provider SDK defaults (AWS CLI, Azure CLI, gcloud).
|
||||
- **S3-compatible stores** (MinIO, Tigris, etc.): Set `endpoint` in `storage_options`. When using `http://` endpoints, also set `allow_http: "true"`.
|
||||
- **Local path** (no scheme): a location without a scheme is a local path. See [Placing the Database](#placing-the-database).
|
||||
|
||||
The `storage_options` keys are case-insensitive and passed directly to the underlying object store library. Available keys depend on the backend. See the [LanceDB storage docs](https://lancedb.com/docs/storage/) for details.
|
||||
|
||||
**Note:** Table optimization is automatically handled by LanceDB Cloud (`db://` URIs) and is disabled for better performance. For object storage backends (S3, Azure, GCS), optimization and vector indexing are still performed normally.
|
||||
|
||||
### Caching and Read Consistency
|
||||
|
||||
```yaml
|
||||
lancedb:
|
||||
read_consistency_interval_seconds: 30 # null to never re-check
|
||||
index_cache_size_bytes: 536870912 # null for the LanceDB default
|
||||
metadata_cache_size_bytes: 268435456
|
||||
```
|
||||
|
||||
- **read_consistency_interval_seconds**: how often a connection checks for writes from another process. `null` never checks, so a long-lived reader never sees the ingester's writes. `0` checks on every read.
|
||||
- **index_cache_size_bytes** / **metadata_cache_size_bytes**: sizes for the caches held by the LanceDB session, which is shared across every connection in the process. The first vector query loads the index into it, so on object storage the cache is what stops the next connection refetching it. Size it for the total set of indexes a process keeps warm, against the memory available to it.
|
||||
|
||||
### Deployment Pattern: One Writer, Many Readers
|
||||
|
||||
LanceDB on S3 supports **exactly one writer + N readers** per database URI. Multiple writers against the same URI can race on the manifest commit and corrupt state. This is a LanceDB property, not something `haiku.rag` enforces.
|
||||
The [one-writer constraint](#operational-constraints) shapes the deployment: one
|
||||
writing process per database URI, any number of read-only consumers.
|
||||
|
||||
The recommended layout for production is "different buckets, same account, separate IAM roles per process":
|
||||
|
||||
- **Ingestion process** — IAM role with `s3:Get/List` on the documents bucket and `s3:Get/Put/Delete` on the LanceDB bucket. Runs `haiku-ingester serve` (with `ingester.sources[type=s3]` pointing at the documents bucket). Exactly one such process per LanceDB URI.
|
||||
- **Consumer processes** (1..N) — IAM role with `s3:Get/List` on the LanceDB bucket only. Run `haiku-rag --read-only mcp`, the chat TUI, etc. They never see the documents bucket.
|
||||
- **Consumer processes** (1..N) — IAM role with `s3:Get/List` on the LanceDB bucket only. Run `haiku-rag mcp`, the chat TUI, etc. They never see the documents bucket.
|
||||
|
||||
Each process picks up its own credentials from the AWS default chain (env vars, IAM instance role, AWS profile), so no credentials are hard-coded in the configuration files.
|
||||
|
||||
`haiku-ingester` writes the database the configuration places, so a one-entry `lancedb.databases` needs no further option. `--db PATH` overrides it. When `lancedb.databases` contains more than one database the ingester has no way to name which it writes, and refuses to start with `AmbiguousDatabaseError`: give each database its own ingester process, each with a configuration naming a single database, or select one with `--db PATH`.
|
||||
|
||||
## Multiple Databases
|
||||
|
||||
Use `lancedb.databases` to name local or remote databases that should be searched together:
|
||||
|
||||
```yaml
|
||||
lancedb:
|
||||
databases:
|
||||
papers: s3://my-bucket/papers.lancedb
|
||||
wiki: s3://my-bucket/wiki.lancedb
|
||||
notes: /data/notes.lancedb
|
||||
```
|
||||
|
||||
A location can be a URI or local path.
|
||||
|
||||
Results, documents, and citations use the configured name as `source`. An unavailable configured database raises `SourceUnavailableError`, which names the database and not its location, so a location never travels in an error a consumer might render or log. A migration, configuration or read-only failure keeps its own type, with the database named in the message. Commands that report on a database, such as `info`, still show where it is.
|
||||
|
||||
Searches spanning multiple databases identify each result with a model-facing `Collection:` line. Searches over one database omit it. Structured `source` fields on results, documents, citations, and analysis dictionaries are unchanged.
|
||||
|
||||
Embedding compatibility is checked against two different things.
|
||||
|
||||
On open, each database is compared with the current configuration. A dimension mismatch raises `ConfigMismatchError`. A provider or model-name mismatch at the same dimension warns in read-only mode and raises in writable mode.
|
||||
|
||||
Across a selection, the databases are compared with each other. Vector and hybrid search embed the query once, so every database answering it must record the same provider, model, and dimension. A disagreement raises `ConfigMismatchError` in read-only mode as well. Only the databases searched together have to agree, and full-text search embeds nothing, so it is unaffected.
|
||||
|
||||
### Search and Provenance
|
||||
|
||||
`search`, `ask`, and `analyze` use the full set by default. Pass `sources` to select a subset:
|
||||
|
||||
```python
|
||||
results = await client.search("query") # every database
|
||||
results = await client.search("query", sources=["papers"]) # one of them
|
||||
```
|
||||
|
||||
Candidates are combined into one ranked list with the configured reranker, or by cosine similarity to the query when reranking is disabled, with within-database rank breaking ties (full-text-only searches order by retrieval score). `SearchResult.source`, `Citation.source`, and `Document.source` carry the database name, for a set and for one database alike.
|
||||
|
||||
The CLI labels results and citations only when the operation spans multiple databases. A command already narrowed with `--db-name` does not repeat the name on every result.
|
||||
|
||||
#### Duplicate IDs
|
||||
|
||||
IDs are unique within a database, not across databases. Copies of a database therefore retain the same IDs.
|
||||
|
||||
Citation ambiguity is evaluated against evidence available to the run. A cited chunk ID is rejected with `AmbiguousCitationError` if search returned it from multiple databases, or it was previously cited from another database. If only one retrieved result has the ID, that result is cited. For an ID absent from search results, the fallback checks every selected database and rejects multiple holders. A shared ID that nothing cites is ignored.
|
||||
|
||||
`get_document_by_id`, `get_chunk_by_id` and `get_picture_bytes` take an optional `source`, and ask that database alone. A name the client does not cover raises `UnknownDatabaseError`. Without one, the document and chunk lookups ask every covered database and answer from the first that holds the ID; `get_picture_bytes` requires one whenever the client covers a set.
|
||||
|
||||
The analysis sandbox rejects shared document IDs because its mount path is `/documents/{id}/`.
|
||||
|
||||
The chat document filter selects by document and database: the search is narrowed to the databases the selection names, and the ID filter applies within them. An ID that copies share still matches in every selected database that holds it.
|
||||
|
||||
#### Ranking
|
||||
|
||||
Without a reranker, the fused list is ordered by cosine similarity between the query vector and each candidate. The databases in a selection share an embedder, so similarity in that one space is comparable across databases, where retrieval scores are each database's own arithmetic. Ties resolve by the candidate's rank within its own database, and configured order decides only when both tie. Similarities rarely tie exactly, so declaration order decides almost nothing: on MTRAG retrieval benchmarks, reversing it left recall unchanged in every cell. Full-text-only searches have no query vector and order by retrieval score instead.
|
||||
|
||||
Results are not guaranteed to spread across databases: a database with nothing relevant to a query contributes nothing, and a strong database can fill every slot. On MTRAG retrieval benchmarks over two to eight collections, cosine fusion holds recall roughly flat as collections are added, where position-based fusion lost up to half its recall at eight.
|
||||
|
||||
A configured reranker scores the combined candidate set directly, ignoring which database each candidate came from, and remains the strongest option: roughly 6 to 8 recall points above cosine fusion on the same benchmarks. Its cost grows with the number of databases because each contributes candidates.
|
||||
|
||||
Image queries are vector-only and skip the reranker: the reranker interface takes a text query, and multimodal reranking applies to pictures on the candidate side, not to image queries. Their fused list is ordered by cosine similarity like any other vector search.
|
||||
|
||||
If a selected database is unavailable, the operation fails with `SourceUnavailableError`, which names that database.
|
||||
|
||||
### Python Operations
|
||||
|
||||
Creating, writing, rebuilding, and vacuuming require one database. Calling these operations on a client that covers multiple raises `AmbiguousDatabaseError`. Select one at creation time or obtain a single-database client:
|
||||
|
||||
```python
|
||||
async with HaikuRAG(config=config, create=True, sources=["papers"]) as papers:
|
||||
...
|
||||
|
||||
async with HaikuRAG(config=config) as client:
|
||||
papers = (await client.clients_for(["papers"]))[0]
|
||||
```
|
||||
|
||||
Conversion, chunking, and title generation do not access a database and remain available on a multi-database client.
|
||||
|
||||
### CLI Commands
|
||||
|
||||
Commands use database sets as follows:
|
||||
|
||||
- **Set-capable**: `search`, `ask`, `analyze`, `chat`, and `mcp` use the full configured set, or the single database selected by `--db-name`.
|
||||
- **Config-only**: `settings`, `init-config`, and `download-models` do not open a database.
|
||||
- **Single-database**: everything else — document writes, `rebuild`, `vacuum`, `migrate`, `init`, `info`, `history`, `tag`, `doctor`, `list`, `inspect`, and `visualize` — works on one database, selected with the global `--db-name` option.
|
||||
|
||||
```bash
|
||||
haiku-rag search "query" # every configured database
|
||||
haiku-rag --db-name papers list # one of them
|
||||
haiku-rag --db-name papers migrate
|
||||
```
|
||||
|
||||
`--db-name` selects an entry from `lancedb.databases`, including remote entries, and `haiku.rag` when nothing is configured. `--db` opens a local path, named by its stem, whatever is configured. A single-database command requires one of these options when multiple databases are configured. A configured set of one is selected automatically.
|
||||
|
||||
Each database is created, migrated and vacuumed on its own:
|
||||
|
||||
```bash
|
||||
haiku-rag --db-name papers init
|
||||
haiku-rag --db-name wiki init
|
||||
```
|
||||
|
||||
## Vector Indexing
|
||||
|
||||
Configure vector search settings:
|
||||
|
||||
```yaml
|
||||
search:
|
||||
vector_index_metric: cosine # cosine, l2, or dot
|
||||
vector_index_metric: cosine # cosine or l2
|
||||
vector_refine_factor: 30 # Re-ranking factor for accuracy
|
||||
vector_nprobes: 20 # IVF partitions searched per query
|
||||
```
|
||||
|
||||
For search behavior settings (`limit`, `max_context_chars`), see [Search and Question Answering](qa.md#search-settings).
|
||||
|
|
@ -128,13 +316,24 @@ For search behavior settings (`limit`, `max_context_chars`), see [Search and Que
|
|||
- **vector_index_metric**: Distance metric for vector similarity:
|
||||
- `cosine`: Cosine similarity (default, best for most embeddings)
|
||||
- `l2`: Euclidean distance
|
||||
- `dot`: Dot product similarity
|
||||
- **vector_refine_factor**: Improves accuracy when using a vector index by retrieving `refine_factor * limit` candidates (using approximate search) and re-ranking them with exact distances. Higher values increase accuracy but slow down queries. Default: 30
|
||||
- **Only applies with a vector index** - has no effect on brute-force search, which already returns exact results
|
||||
- **vector_nprobes**: How many IVF partitions each query searches. Higher values increase recall and latency. A larger corpus holds more partitions, so the same value covers a smaller fraction of it. Default: 20
|
||||
- **Only applies with a vector index** - ignored by brute-force search
|
||||
|
||||
!!! note
|
||||
Vector indexes are only necessary for large datasets with over 100,000 chunks. For smaller datasets, LanceDB's brute-force kNN search provides exact results with good performance. Only create an index if you notice search performance degradation on large datasets.
|
||||
|
||||
Retrieval MAP with and without an index, measured on copies of the benchmark databases with no reranker:
|
||||
|
||||
| Dataset | Chunks | Dim | Exact | Indexed | Delta | Build | Peak RSS |
|
||||
|---------|-------:|----:|------:|--------:|------:|------:|---------:|
|
||||
| `hotpotqa` | 70,527 | 2560 | 0.6978 | 0.6979 | +0.0001 | 29.3 s | 3.19 GB |
|
||||
| `orb_multimodal_nemotron` | 121,168 | 2048 | 0.9799 | 0.9800 | +0.0001 | 25.8 s | 3.38 GB |
|
||||
| `frames` | 425,940 | 2560 | 0.5431 | 0.5387 | -0.0044 | 34.1 s | 4.02 GB |
|
||||
|
||||
An index costs no accuracy at 70k and 121k chunks and 0.0044 MAP at 426k. A larger corpus holds more IVF partitions, so the default number of probes covers a smaller fraction of the space, and `vector_refine_factor` can only re-score what those probes returned. Raise `vector_nprobes` to trade latency for recall on a large corpus. Build cost is near-flat in row count because training samples the data rather than scanning it, and vector dimension drives it more than corpus size.
|
||||
|
||||
**Index creation:**
|
||||
|
||||
Vector indexes are **not created automatically** during document ingestion to avoid slowing down the process. After you've added documents (at least 256 chunks required), create the index manually:
|
||||
|
|
@ -150,12 +349,12 @@ This command:
|
|||
|
||||
**Re-indexing:**
|
||||
|
||||
Indexes are not automatically updated when you add new documents. After adding a significant amount of new data:
|
||||
New chunks reach the index without a rebuild. `optimize()`, which runs after writes while `auto_vacuum` is on, adds them as a delta part. Between a write and the next optimize, LanceDB serves ANN over the indexed rows and a brute-force scan over the remainder, then combines the results.
|
||||
|
||||
A rebuild retrains the centroids, which are fitted once at build time and never recomputed. As a corpus grows past the distribution it was trained on the partitioning fits it less well, and delta parts accumulate. Rebuild after substantial growth:
|
||||
|
||||
```bash
|
||||
haiku-rag create-index # Rebuilds the index with all data
|
||||
haiku-rag create-index
|
||||
```
|
||||
|
||||
Searches still work with stale indexes - LanceDB uses the index for old data (fast ANN) and brute-force kNN for new unindexed rows, then combines the results. However, performance degrades as more unindexed data accumulates.
|
||||
|
||||
For datasets with fewer than 256 chunks, searches use brute-force kNN scans (exact nearest neighbors, 100% recall) which work well for small datasets but don't scale beyond a few hundred thousand vectors.
|
||||
|
|
|
|||
|
|
@ -105,6 +105,8 @@ for chunk in chunks:
|
|||
print(f"Headings: {meta.headings}")
|
||||
print(f"Page numbers: {meta.page_numbers}")
|
||||
print(f"Labels: {meta.labels}")
|
||||
# Access raw metadata (including headings, page_numbers and labels)
|
||||
print(f"Raw metadata: {chunk.metadata}")
|
||||
```
|
||||
|
||||
Chunks are returned with:
|
||||
|
|
@ -115,6 +117,8 @@ Chunks are returned with:
|
|||
- `embedding` - `None` (not yet embedded)
|
||||
- `document_id` - `None` (not yet stored)
|
||||
|
||||
A custom `DocumentChunker` can provide other keys and values in `metadata`. They will be stored with the chunk and are accessible when it is returned in a search result or citation, within `SearchResult.chunk_meta` / `Citation.chunk_meta`.
|
||||
|
||||
## Embed
|
||||
|
||||
`embed_chunks()` generates embeddings for chunks using the client's embedder. It automatically contextualizes chunks (prepends section headings) before embedding for better semantic search, without modifying the stored content:
|
||||
|
|
|
|||
|
|
@ -109,8 +109,15 @@ uv run ty check
|
|||
|
||||
Tests automatically set mock API keys for providers that require them during client initialization. When running with VCR playback, these mock keys are sufficient since no real API calls are made.
|
||||
|
||||
When recording new cassettes, set real API keys via environment variables:
|
||||
Recording reaches the real service, so the recording command needs network
|
||||
access and the keys that service reads. Name the exact test and pass `-n0`:
|
||||
a module-wide `--record-mode=rewrite` re-records every cassette in it,
|
||||
including ones whose service you do not have running.
|
||||
|
||||
```bash
|
||||
ANTHROPIC_API_KEY=sk-ant-... uv run pytest tests/test_qa.py::test_qa_anthropic --record-mode=rewrite
|
||||
# Ollama-backed cassettes need no key, only a running Ollama
|
||||
uv run pytest tests/test_embedder.py::test_ollama_embedder -n0 --record-mode=rewrite
|
||||
|
||||
# A keyed provider reads its own variable. Cohere's SDK reads CO_API_KEY
|
||||
CO_API_KEY=... uv run pytest tests/test_reranker.py::test_cohere_reranker -n0 --record-mode=rewrite
|
||||
```
|
||||
|
|
|
|||
|
|
@ -1,3 +1,38 @@
|
|||
---
|
||||
title: haiku.rag
|
||||
description: Local-first agentic RAG. Index PDFs, web pages, and whole directories, then ask questions and get answers cited to page numbers and section headings. Hybrid search, reranking, and multimodal retrieval on embedded LanceDB.
|
||||
---
|
||||
|
||||
haiku.rag indexes PDFs, web pages, and whole directories, retrieves with hybrid search, and answers with citations down to the page number and section heading. It runs on an embedded database with open models, so your documents stay on your machine and there is no server to operate.
|
||||
|
||||
```bash
|
||||
uv pip install haiku.rag
|
||||
|
||||
haiku-rag init
|
||||
haiku-rag add-src ~/Documents/some-paper.pdf
|
||||
haiku-rag ask "what does it conclude?"
|
||||
```
|
||||
|
||||
[Quickstart](tutorial.md) covers provider setup and the first ingestion.
|
||||
|
||||
## Why haiku.rag
|
||||
|
||||
**Answers you can check.** Every answer carries citations with page numbers and section headings. Visual grounding shows the cited chunk highlighted on the original page image. Optional capabilities require an answer to declare what grounds it, including declaring that nothing does.
|
||||
|
||||
**Local-first, no server.** Embedded [LanceDB](https://lancedb.com/) and open models through [Ollama](https://ollama.com/) by default. No database to run and no API keys required. The same code runs against S3, GCS, Azure, LanceDB Cloud, or any provider Pydantic AI supports.
|
||||
|
||||
**Built for agents.** Native [Pydantic AI](https://ai.pydantic.dev/) capabilities compose into your own agents. An [MCP server](mcp.md) exposes the same database to Claude Desktop and other assistants. The analysis capability runs sandboxed Python across documents for questions that need computation rather than retrieval.
|
||||
|
||||
**Measured, not asserted.** Retrieval and answer quality are tracked against public benchmarks with runnable configs. See [Benchmarks](benchmarks.md).
|
||||
|
||||
## Start here
|
||||
|
||||
- [Quickstart](tutorial.md): install, index, chat.
|
||||
- [Installation](installation.md): packages and extras.
|
||||
- [Architecture](overview.md): how a document becomes a cited answer.
|
||||
- [Capabilities](capabilities/index.md): native RAG and analysis capabilities for Pydantic AI agents.
|
||||
- [Python API](python.md): use haiku.rag from code.
|
||||
- [MCP server](mcp.md): expose haiku.rag to Claude Desktop or other AI assistants.
|
||||
- [Configuration](configuration/index.md): every setting.
|
||||
|
||||
MIT licensed. Source on [GitHub](https://github.com/ggozad/haiku.rag).
|
||||
|
|
|
|||
|
|
@ -206,7 +206,7 @@ so a class is its own factory:
|
|||
# example_pkg/__init__.py
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from haiku.rag.ingester.sources import FetchResult
|
||||
from haiku.rag.sources import FetchResult
|
||||
|
||||
|
||||
class Provider:
|
||||
|
|
@ -306,7 +306,7 @@ class Source(Protocol):
|
|||
```
|
||||
|
||||
`FetchResult`, `SourceEvent`, `SourceEventKind`, and `RevisionSnapshot`
|
||||
live in `haiku.rag.ingester.sources`.
|
||||
live in `haiku.rag.sources`.
|
||||
|
||||
```toml
|
||||
# in the source package's pyproject.toml
|
||||
|
|
@ -351,7 +351,8 @@ job at a time. `worker_count` is therefore also the maximum number of
|
|||
concurrent in-flight jobs. Jobs that hit a `TransientError` are
|
||||
rescheduled with exponential backoff plus jitter, up to `max_attempts`,
|
||||
then land in the dead-letter queue. `PermanentError` (unsupported
|
||||
extension, 4xx HTTP except 408/429, etc.) skips retry entirely.
|
||||
extension, 4xx HTTP except 408/429, object-store credential and
|
||||
configuration errors, etc.) skips retry entirely.
|
||||
|
||||
While a worker processes a job it renews the job's lease every
|
||||
`heartbeat_interval_s`. A reaper task resets any claim whose lease has not
|
||||
|
|
@ -436,7 +437,9 @@ server, then pollers, then in-flight workers.
|
|||
|
||||
### Single-writer constraint
|
||||
|
||||
LanceDB supports exactly one writer + N readers per database URI. Run
|
||||
haiku.rag serializes multi-table writes with a process-local lock and rolls
|
||||
them back by restoring table versions, so a second writing process can
|
||||
commit inside another's transaction and be reverted by its rollback. Run
|
||||
exactly one `haiku-ingester serve` against a given LanceDB. Multiple
|
||||
MCP servers or read-only consumers against the same DB are fine. Sharing
|
||||
the Postgres queue across processes is safe (the claim/lease lifecycle is
|
||||
|
|
|
|||
|
|
@ -10,37 +10,54 @@
|
|||
uv pip install haiku.rag
|
||||
```
|
||||
|
||||
The full package includes **all features and extras**:
|
||||
- **Document processing** (Docling) - PDF, DOCX, PPTX, images, and 40+ file formats
|
||||
- **All embedding providers** - VoyageAI
|
||||
- **All rerankers** - MixedBread AI, Cohere, Zero Entropy
|
||||
The full package pulls the `docling`, `voyageai`, `cohere`, `zeroentropy`,
|
||||
`cross-encoder`, `jina` and `tui` extras. It does not include `s3` or `ingester`:
|
||||
|
||||
```bash
|
||||
uv pip install 'haiku.rag[ingester]' # the haiku-ingester service
|
||||
uv pip install 'haiku.rag[s3]' # S3 and object storage
|
||||
```
|
||||
|
||||
### Slim Package (Minimal Dependencies)
|
||||
|
||||
```bash
|
||||
# Minimal installation (no document processing)
|
||||
uv pip install haiku.rag-slim
|
||||
|
||||
# With document processing
|
||||
uv pip install haiku.rag-slim[docling]
|
||||
|
||||
# With specific providers
|
||||
uv pip install haiku.rag-slim[docling,voyageai,cross-encoder]
|
||||
uv pip install 'haiku.rag-slim[docling]'
|
||||
uv pip install 'haiku.rag-slim[docling,voyageai,cross-encoder]'
|
||||
```
|
||||
|
||||
The slim package has minimal dependencies and lets you install only what you need:
|
||||
### Extras
|
||||
|
||||
- `docling` - PDF, DOCX, PPTX, images, and other document formats
|
||||
- `voyageai` - VoyageAI embeddings
|
||||
- `cross-encoder` - Local reranking via sentence-transformers
|
||||
- `cohere` - Cohere reranking
|
||||
- `zeroentropy` - Zero Entropy reranking
|
||||
- `tui` - Terminal UI for `chat` and `inspect` commands
|
||||
Every extra `haiku.rag-slim` defines. The right-hand column marks the ones the
|
||||
full `haiku.rag` package already includes.
|
||||
|
||||
| Extra | Provides | In `haiku.rag` |
|
||||
|---|---|---|
|
||||
| `docling` | PDF, DOCX, PPTX, images and 40+ formats, converted locally | yes |
|
||||
| `tui` | Terminal UI for `chat` and `inspect` | yes |
|
||||
| `voyageai` | VoyageAI embeddings | yes |
|
||||
| `cohere` | Cohere embeddings and reranking | yes |
|
||||
| `zeroentropy` | Zero Entropy reranking | yes |
|
||||
| `cross-encoder` | Local reranking via sentence-transformers | yes |
|
||||
| `jina` | Local Jina reranking (`provider: jina-local`) | yes |
|
||||
| `s3` | S3 and object-storage access | no |
|
||||
| `ingester` | The `haiku-ingester` service (also pulls `s3`) | no |
|
||||
| `anthropic` | Anthropic Claude models | no |
|
||||
| `google` | Google Gemini models | no |
|
||||
| `groq` | Groq models | no |
|
||||
| `mistral` | Mistral models | no |
|
||||
| `bedrock` | AWS Bedrock models | no |
|
||||
| `vertexai` | Google Vertex AI models | no |
|
||||
|
||||
Ollama and any OpenAI-compatible endpoint work with no extra at all.
|
||||
|
||||
**Built-in providers** (no extras needed):
|
||||
- **Ollama** (default embedding provider)
|
||||
- **OpenAI** (GPT models for QA and embeddings)
|
||||
- **Anthropic** (Claude models for QA)
|
||||
- **vLLM** and other OpenAI-compatible endpoints (embeddings, QA, reranking)
|
||||
- **Jina** reranking via `provider: jina`, which calls the Jina HTTP API
|
||||
|
||||
Other providers come from the extras above, which pull the matching Pydantic AI extra. For Claude models, `uv pip install 'haiku.rag-slim[anthropic]'`.
|
||||
|
||||
See [Configuration](configuration/index.md) for configuring providers including advanced options like vLLM.
|
||||
|
||||
|
|
@ -74,7 +91,7 @@ See [Remote processing](remote-processing.md) for setup instructions and [Docume
|
|||
|
||||
## Docker
|
||||
|
||||
Two Docker images are available:
|
||||
Only the slim image is published. Build the full image yourself:
|
||||
|
||||
### Slim Image (Minimal)
|
||||
|
||||
|
|
|
|||
180
docs/mcp.md
180
docs/mcp.md
|
|
@ -19,15 +19,71 @@ haiku-rag mcp --host 0.0.0.0 --port 8001
|
|||
# stdio transport (for Claude Desktop)
|
||||
haiku-rag mcp --stdio
|
||||
|
||||
# Read-only mode (excludes write tools)
|
||||
haiku-rag --read-only mcp --stdio
|
||||
```
|
||||
|
||||
`--host` defaults to `127.0.0.1` (loopback only). Bind to `0.0.0.0` only
|
||||
when you want the MCP server reachable from outside the local machine —
|
||||
e.g. inside a Docker container with port mapping, or on a trusted LAN.
|
||||
|
||||
**Read-only mode:** When `--read-only` is specified, write tools (`add_document_from_file`, `add_document_from_url`, `add_document_from_text`, `delete_document`) are not registered. Only search and query tools remain available.
|
||||
The server opens the database read-only. Ingestion goes through the CLI
|
||||
(`haiku-rag add`, `add-src`, `delete`) or [`haiku-ingester`](ingester.md).
|
||||
|
||||
## Collections
|
||||
|
||||
With several databases in `lancedb.databases`, the server covers all of
|
||||
them, as `haiku-rag search` does. Results and documents name theirs in
|
||||
`source`. `sources` on `search_documents`, `search_documents_by_image`
|
||||
and `execute_code` restricts a call to a subset; `source` on `get_document` names the database holding the
|
||||
document. A name the server does not cover is an error.
|
||||
`haiku-rag --db-name NAME mcp` serves one. See
|
||||
[Multiple Databases](configuration/storage.md#multiple-databases).
|
||||
|
||||
## Claude Code
|
||||
|
||||
The repository ships a plugin that registers the server and a skill telling
|
||||
Claude when and how to use it:
|
||||
|
||||
```bash
|
||||
claude plugin marketplace add ggozad/haiku.rag
|
||||
claude plugin install haiku-rag
|
||||
```
|
||||
|
||||
The plugin runs `haiku-rag mcp --stdio`, so `haiku-rag` must be on the PATH
|
||||
and the configuration decides the database. The skill pre-approves every tool
|
||||
and is also invocable as `/haiku-rag`. To register the server without the
|
||||
plugin:
|
||||
|
||||
```bash
|
||||
claude mcp add haiku-rag -- haiku-rag mcp --stdio
|
||||
```
|
||||
|
||||
The skill works with that registration too: copy `plugins/haiku-rag/skills/haiku-rag`
|
||||
into `~/.claude/skills/` and change the tool prefix in its `allowed-tools` from
|
||||
`mcp__plugin_haiku-rag_haiku-rag__` to `mcp__haiku-rag__`.
|
||||
|
||||
## Codex
|
||||
|
||||
The repository's Codex plugin registers the server and installs the same Agent
|
||||
Skill:
|
||||
|
||||
```bash
|
||||
codex plugin marketplace add ggozad/haiku.rag
|
||||
codex plugin add haiku-rag@haiku-rag
|
||||
```
|
||||
|
||||
The plugin runs `haiku-rag mcp --stdio`, so `haiku-rag` must be on the PATH.
|
||||
Invoke the skill as `$haiku-rag`. Codex can also select it automatically from
|
||||
its description. To register the server without the plugin:
|
||||
|
||||
```bash
|
||||
codex mcp add haiku-rag -- haiku-rag mcp --stdio
|
||||
```
|
||||
|
||||
The skill works with that registration too: copy
|
||||
`plugins/haiku-rag/skills/haiku-rag` into `~/.agents/skills/`.
|
||||
The `allowed-tools` field supplies Claude Code's tool pre-approval and may be
|
||||
ignored by other Agent Skills clients. Codex configures MCP tool approvals
|
||||
separately in `config.toml`.
|
||||
|
||||
## Claude Desktop Integration
|
||||
|
||||
|
|
@ -57,63 +113,95 @@ With a custom database path:
|
|||
}
|
||||
```
|
||||
|
||||
After restarting Claude Desktop, you can ask Claude to search your documents, add new content, or answer questions using your knowledge base.
|
||||
After restarting Claude Desktop, you can ask Claude to search your documents or answer questions using your knowledge base.
|
||||
|
||||
## Available Tools
|
||||
## Tools
|
||||
|
||||
### Document Management
|
||||
Every tool is read-only and says so in its annotations. Each parameter carries
|
||||
a description in the tool schema, so the listing below names them without
|
||||
repeating it.
|
||||
|
||||
- **`add_document_from_file`** - Add documents from local file paths
|
||||
- `file_path` (required): Path to the file
|
||||
- `metadata` (optional): Key-value metadata
|
||||
- `title` (optional): Human-readable title
|
||||
| Tool | Registered | Parameters |
|
||||
|---|---|---|
|
||||
| `search_documents` | always | `query`, `limit`, `include_images`, `filter`, `sources` |
|
||||
| `search_documents_by_image` | multimodal embedder only | `image_base64`, `limit`, `include_images`, `filter`, `sources` |
|
||||
| `get_document` | always | `document_id`, `source` |
|
||||
| `get_document_outline` | always | `document_id`, `source` |
|
||||
| `get_document_section` | always | `document_id`, `section_id`, `source` |
|
||||
| `list_documents` | always | `limit`, `offset`, `filter` |
|
||||
| `execute_code` | always | `code`, `filter`, `sources` |
|
||||
|
||||
- **`add_document_from_url`** - Add documents from URLs
|
||||
- `url` (required): URL to fetch
|
||||
- `metadata` (optional): Key-value metadata
|
||||
- `title` (optional): Human-readable title
|
||||
`search_documents` runs hybrid search, vector and full-text. Its text content
|
||||
is the rendering the in-process agents read: results best first, each with its
|
||||
rank, `Document ID`, `Collection` when the server covers several, the document
|
||||
title, section headings, the matched chunk's metadata when it has any, and the
|
||||
passage expanded to its section the way the agents get it
|
||||
(`search.max_context_chars` caps it). Pictures in the results follow as
|
||||
image blocks, one per distinct picture, each preceded by a line naming its
|
||||
result; `include_images: false` leaves them out. Search results carry no
|
||||
structured content, so every client shows the model the same text and
|
||||
images. Scores are not comparable across
|
||||
queries or search types, so rank is the signal. `search_documents_by_image`
|
||||
embeds the query image and searches by vector similarity alone.
|
||||
|
||||
- **`add_document_from_text`** - Add documents from raw text content
|
||||
- `content` (required): Text content
|
||||
- `uri` (optional): URI identifier
|
||||
- `metadata` (optional): Key-value metadata
|
||||
- `title` (optional): Human-readable title
|
||||
`get_document` returns a document whole, in reading order. For a long one,
|
||||
`get_document_outline` returns the heading tree with page numbers and
|
||||
`get_document_section` the text of one section, subsections included; a
|
||||
node's `id` in the outline is the `section_id`. A document without headings
|
||||
has an empty outline. `list_documents` returns titles, URIs and metadata,
|
||||
which is how a client learns what a filter can match.
|
||||
|
||||
- **`get_document`** - Retrieve a document by ID
|
||||
- `document_id` (required): The document ID
|
||||
### Code
|
||||
|
||||
- **`list_documents`** - List documents with pagination and filtering
|
||||
- `limit` (optional): Maximum number to return
|
||||
- `offset` (optional): Number to skip
|
||||
- `filter` (optional): SQL WHERE clause for filtering
|
||||
`execute_code` runs a Python program in the sandbox of the
|
||||
[analysis capability](capabilities/analysis.md), over the documents `filter`
|
||||
and `sources` select, and returns what it printed. The program reads
|
||||
`/documents/{document_id}/` (`metadata.json`, `content.txt`, `items.jsonl`,
|
||||
`chunks.jsonl`, `toc.json`) and can `await search()` and
|
||||
`await list_documents()`; the tool description spells out the fields and the
|
||||
patterns that matter. Each call is one program: nothing carries over between
|
||||
calls, and the sandbox is created and closed per call. A failing program is a
|
||||
tool error carrying the interpreter's message and any output printed before
|
||||
it. No model runs on the server. Claude Code moves a call still running after
|
||||
about two minutes to a background task.
|
||||
|
||||
- **`delete_document`** - Delete a document by ID
|
||||
- `document_id` (required): The document ID
|
||||
The interpreter is [Monty](https://github.com/pydantic/monty), a Python subset.
|
||||
Useful modules include `json`, `re`, `math`, `pathlib`, `datetime`,
|
||||
`collections`, `itertools`, `functools` and `dataclasses`. Absent, and often
|
||||
reached for: `decimal` and `statistics`. No generator functions, class
|
||||
inheritance or `match` statements, and a file object cannot be iterated. Files are read-only, and
|
||||
there is no network and no filesystem
|
||||
beyond `/documents`. `analysis.code_timeout` is the call's budget: compute is
|
||||
stopped at it, and past it no further host call starts, a file read or an
|
||||
in-code search alike, though one already running finishes.
|
||||
`analysis.max_output_chars` bounds the output.
|
||||
|
||||
### Search
|
||||
### Filters
|
||||
|
||||
- **`search_documents`** - Search using hybrid search (vector + full-text)
|
||||
- `query` (required): Search query
|
||||
- `limit` (optional): Maximum results (uses config default if not specified)
|
||||
- `include_images` (optional, default `true`): Attach base64-encoded picture bytes to picture-labeled results
|
||||
`filter` is a SQL WHERE clause over the document columns `id`, `uri`, `title`,
|
||||
`metadata`, `created_at`, `updated_at`. `metadata` is a JSON string, so match
|
||||
its keys with LIKE:
|
||||
|
||||
- **`search_documents_by_image`** - Search using an image as the query (registered only when the configured embedder supports images)
|
||||
- `image_base64` (required): Base64-encoded image (PNG/JPEG bytes)
|
||||
- `limit` (optional): Maximum results
|
||||
- `include_images` (optional, default `true`)
|
||||
```sql
|
||||
metadata LIKE '%"author": "Smith"%'
|
||||
uri LIKE '%.pdf'
|
||||
title = 'Q3 report'
|
||||
```
|
||||
|
||||
### Question Answering
|
||||
### Errors
|
||||
|
||||
- **`ask_question`** - Ask questions about your documents
|
||||
- `question` (required): The question to ask
|
||||
- `cite` (optional): Include source citations (default: false)
|
||||
- `deep` (optional): Use multi-agent deep QA for complex questions (default: false)
|
||||
A failure is an MCP error carrying its message, never an empty result: a
|
||||
document or section id that matches nothing, a collection the server does not
|
||||
cover, a filter the query engine rejects, invalid base64, a program that fails
|
||||
in `execute_code` with the error it hit, and anything unexpected with its own
|
||||
message.
|
||||
|
||||
- **`analyze`** - Answer complex analytical questions via code execution
|
||||
- `question` (required): The question to answer
|
||||
- `filter` (optional): SQL WHERE clause to restrict document access
|
||||
- `document` (optional): Document title/ID to pre-load (can repeat)
|
||||
- Best for aggregation, computation, and multi-document analysis
|
||||
### Instructions
|
||||
|
||||
The server publishes `instructions` describing the knowledge base: what it
|
||||
holds, when to reach for it, the collection names when it covers several, and
|
||||
`prompts.domain_preamble` when set. Claude Code and Codex show them to the
|
||||
model. Claude Desktop does not, so every tool description stands on its own.
|
||||
|
||||
## Continuous ingestion
|
||||
|
||||
|
|
|
|||
107
docs/overview.md
107
docs/overview.md
|
|
@ -1,47 +1,90 @@
|
|||
# Overview
|
||||
# Architecture
|
||||
|
||||
haiku.rag is an agentic RAG that runs locally and scales to production. Index PDFs, web pages, or whole directories. Ask questions and get cited answers. Build agents, skills, and MCP integrations on top.
|
||||
haiku.rag ingests documents, retrieves from them with hybrid search, and answers
|
||||
with citations. This page follows the data through the system. For a working
|
||||
setup, start with the [Quickstart](tutorial.md).
|
||||
|
||||
haiku.rag is open-source first. The defaults run open models through [Ollama](https://ollama.com/) so the full pipeline works without external API keys. Any provider Pydantic AI supports works in its place.
|
||||
## Ingestion
|
||||
|
||||
Built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.pydantic.dev/), and [Docling](https://docling-project.github.io/docling/). Embedded database, no servers required.
|
||||
|
||||
## See it work
|
||||
|
||||
```bash
|
||||
uv pip install haiku.rag
|
||||
|
||||
ollama pull qwen3-embedding:4b
|
||||
ollama pull gpt-oss
|
||||
|
||||
haiku-rag init
|
||||
haiku-rag add-src ~/Documents/some-paper.pdf
|
||||
haiku-rag chat
|
||||
```text
|
||||
source adapter -> converter -> chunker -> embedder -> LanceDB
|
||||
```
|
||||
|
||||
The chat TUI is one way to interact with the database. `haiku-rag ask` and `haiku-rag search` cover one-shot CLI usage. Python integrations, skills, and the MCP server work against the same database.
|
||||
A **source adapter** owns the I/O and the identity of a document: it fetches
|
||||
bytes, reports the backend's revision (mtime for a file, ETag for S3 or HTTP),
|
||||
and computes the content hash. The same adapters serve one-shot ingestion
|
||||
(`haiku-rag add-src`, `HaikuRAG.create_document_from_source`) and the continuous
|
||||
[`haiku-ingester`](ingester.md) service, so both agree on what a document is and
|
||||
when it has changed.
|
||||
|
||||
## What it does
|
||||
The **converter** turns those bytes into a `DoclingDocument`, the structured form
|
||||
that carries headings, tables, pictures and page provenance. It runs in-process
|
||||
with the `docling` extra, or against a [docling-serve](remote-processing.md)
|
||||
fleet.
|
||||
|
||||
**Ingest.** PDFs, DOCX, HTML, images, and 40+ formats via Docling. Add files, URLs, or whole directories with `haiku-rag add-src`, or run the [`haiku-ingester`](ingester.md) service for continuous, queue-backed ingestion from filesystem, HTTP, S3, or WebDAV sources.
|
||||
The **chunker** splits that structure into chunks, each keeping the headings it
|
||||
sits under, the page numbers it came from, and references to the document items
|
||||
it covers. With a multimodal embedder, pictures become chunks of their own.
|
||||
|
||||
**Search.** Hybrid retrieval (vector + full-text with reciprocal rank fusion), optional cross-encoder reranking, structure-aware context expansion. Image-as-query and cross-modal retrieval when configured with a multimodal embedder.
|
||||
The **embedder** vectorizes them in batches. The document, its mutable metadata,
|
||||
its chunks and its structural items are written under one process-local
|
||||
transaction: it takes a version snapshot, and on failure restores each table to
|
||||
it. A rollback that cannot complete raises rather than reporting success, and the
|
||||
snapshot is only meaningful while this process is the only writer.
|
||||
|
||||
**Answer.** RAG skill with citations including page numbers, section headings, and visual grounding. Vision-capable models receive figure bytes alongside chunk text. Analysis skill with a sandboxed Python interpreter for aggregation and computation across documents.
|
||||
## Storage
|
||||
|
||||
**Integrate.** Use it from Python, the CLI, the [MCP server](mcp.md), or as composable [skills](skills/index.md) built on haiku.skills. Skills bundle tools, prompts, and state for use inside any Pydantic AI agent.
|
||||
LanceDB is embedded, so there is no server. The same code runs against a local
|
||||
directory, S3, GCS, Azure or LanceDB Cloud by changing a database's location in
|
||||
`lancedb.databases`.
|
||||
|
||||
**Operate.** Embedded LanceDB by default. Also runs on S3, GCS, Azure, or LanceDB Cloud. Time-travel queries via LanceDB versioning. The [`haiku-ingester`](ingester.md) service runs continuously for production deployments.
|
||||
Tables are versioned. Vacuum collapses old versions on a retention window, and
|
||||
[tags](cli.md) name a state across all tables so a database can be restored to
|
||||
it later.
|
||||
|
||||
## Where to go next
|
||||
One process writes at a time. Reads are unrestricted, and a reader sees another
|
||||
process's writes after `lancedb.read_consistency_interval_seconds`.
|
||||
|
||||
- [Quickstart](tutorial.md): install, index, chat.
|
||||
- [Skills](skills/index.md): the rag and rag-analysis skills you compose into Pydantic AI agents.
|
||||
- [Python API](python.md): use haiku.rag from code.
|
||||
- [MCP server](mcp.md): expose haiku.rag to Claude Desktop or other AI assistants.
|
||||
- [Tuning](tuning.md): improve retrieval quality.
|
||||
- [Configuration](configuration/index.md): every setting.
|
||||
## Retrieval
|
||||
|
||||
## License
|
||||
```text
|
||||
query -> vector + full-text search -> fusion -> rerank -> context expansion
|
||||
```
|
||||
|
||||
MIT. Source on [GitHub](https://github.com/ggozad/haiku.rag).
|
||||
Search runs a vector query and a full-text query and fuses the rankings. With a
|
||||
reranker configured, it retrieves ten times the requested limit and reranks down
|
||||
to it, so quality improves without changing the caller's limit.
|
||||
|
||||
Results then expand: a chunk is returned with the section it belongs to, bounded
|
||||
by `search.max_context_chars`. Sections that fit come back whole, larger ones
|
||||
grow outward from the match, and small ones grow across boundaries. Every result
|
||||
carries its page numbers and headings, which is what makes a citation checkable.
|
||||
|
||||
## Answering
|
||||
|
||||
Two [capabilities](capabilities/index.md) sit on top, both native Pydantic AI
|
||||
capabilities you can attach to your own agent:
|
||||
|
||||
- The **RAG capability** searches and cites. Its citations carry page numbers and
|
||||
headings, and `haiku-rag visualize` draws the cited chunk on the page image.
|
||||
- The **analysis capability** adds a sandboxed Python interpreter with the
|
||||
documents mounted as a filesystem, for questions that need computation across
|
||||
documents rather than retrieval.
|
||||
|
||||
Two optional capabilities compose with them: evidence compaction replaces older
|
||||
turns' evidence with what was actually cited, and citation policy requires every
|
||||
answer to declare what grounds it.
|
||||
|
||||
The same database is reachable from [Python](python.md), the [CLI](cli.md), and
|
||||
the [MCP server](mcp.md).
|
||||
|
||||
## Running it
|
||||
|
||||
A laptop needs nothing but the package and Ollama. Production adds the
|
||||
[`haiku-ingester`](ingester.md) service, which polls its sources, queues work in
|
||||
SQLite or Postgres, and retries with a circuit breaker per source.
|
||||
|
||||
Before deploying, read the operational constraints in
|
||||
[Storage](configuration/storage.md): one writer per database, `haiku-rag migrate`
|
||||
after an upgrade that changes the schema, and a fixed embedding dimension per
|
||||
database.
|
||||
|
|
|
|||
124
docs/python.md
124
docs/python.md
|
|
@ -24,8 +24,10 @@ async with HaikuRAG("path/to/database.lancedb", read_only=True) as client:
|
|||
# await client.create_document(...) # Would raise ReadOnlyError
|
||||
```
|
||||
|
||||
`async with` is the lifecycle. A caller that owns the client some other way releases it with `await client.aclose()`, which does the same work for every client shape. `client.close()` closes the connection to one database and nothing else, since draining the background vacuum and releasing the embedder and reranker are awaitable; it refuses a client covering several.
|
||||
|
||||
!!! note
|
||||
Databases must be explicitly created with `create=True` or via `haiku-rag init` before use. Operations on non-existent databases will raise `FileNotFoundError`.
|
||||
Databases must be explicitly created with `create=True` or via `haiku-rag init` before use. Opening a nonexistent local database given as a path raises `FileNotFoundError`, naming the path; a configured or default database raises `SourceUnavailableError`, which names the database rather than its location. A path beside a configured `lancedb.databases` raises `AmbiguousDatabaseError`.
|
||||
|
||||
!!! note
|
||||
Read-only mode is useful for safely accessing databases without risk of modification. It blocks all write operations and downgrades an embedding provider/name mismatch to a warning instead of raising `ConfigMismatchError`.
|
||||
|
|
@ -87,6 +89,8 @@ PDFs that carry attachments via the `/EmbeddedFiles` table are split into one Do
|
|||
By ID:
|
||||
```python
|
||||
doc = await client.get_document_by_id("document-id-string")
|
||||
doc = await client.get_document_by_id("document-id-string", "papers")
|
||||
chunk = await client.get_chunk_by_id("chunk-id-string", "papers")
|
||||
```
|
||||
|
||||
By URI:
|
||||
|
|
@ -94,11 +98,20 @@ By URI:
|
|||
doc = await client.get_document_by_uri("file:///path/to/document.pdf")
|
||||
```
|
||||
|
||||
Both return content, uri, title and metadata. The multi-MB docling blobs are
|
||||
loaded separately:
|
||||
|
||||
```python
|
||||
docling = await client.document_repository.get_docling_data(doc.id)
|
||||
pages = await client.document_repository.get_pages_data(doc.id)
|
||||
```
|
||||
|
||||
List all documents:
|
||||
```python
|
||||
docs = await client.list_documents(limit=10, offset=0)
|
||||
|
||||
# Include full content and docling document (not loaded by default)
|
||||
# Include the text content (not loaded by default). A listing never loads the
|
||||
# docling blobs.
|
||||
docs = await client.list_documents(include_content=True)
|
||||
```
|
||||
|
||||
|
|
@ -187,6 +200,8 @@ for result in results:
|
|||
print(f"Document ID: {result.document_id}")
|
||||
```
|
||||
|
||||
Each result carries the parent document's metadata in `result.document_meta` and the relevant chunk's verbatim metadata in `result.chunk_meta`. Neither is shown to the model during QA.
|
||||
|
||||
Search with different search types:
|
||||
```python
|
||||
# Vector search only
|
||||
|
|
@ -219,6 +234,60 @@ for result in results:
|
|||
print(f"Document Title: {result.document_title}") # when available
|
||||
```
|
||||
|
||||
### Searching Multiple Databases
|
||||
|
||||
With [`lancedb.databases`](configuration/storage.md#multiple-databases) configured, a client covers the full set. Use `sources` to select a subset. Each result includes its database name:
|
||||
|
||||
```python
|
||||
results = await client.search("machine learning") # all of them
|
||||
results = await client.search("machine learning", sources=["papers"]) # one of them
|
||||
|
||||
for result in results:
|
||||
print(f"{result.source}: {result.content}")
|
||||
```
|
||||
|
||||
`ask` and `analyze` also accept `sources`. Citations include the database name:
|
||||
|
||||
```python
|
||||
answer, citations = await client.ask("What changed?", sources=["papers", "wiki"])
|
||||
for cite in citations:
|
||||
print(f"[{cite.source}] {cite.document_title or cite.document_uri}")
|
||||
|
||||
result = await client.analyze("How many documents mention it?", sources=["papers"])
|
||||
```
|
||||
|
||||
A scoped question can cite only the selected databases. Analysis mounts only their documents.
|
||||
|
||||
`sources=None` covers every database the client covers. `sources=[]` covers none: `search` returns no results, and `ask` and `analyze` run with no evidence from any database.
|
||||
|
||||
A name no client covers raises `UnknownDatabaseError`, a `KeyError`, wherever it is given: at construction, per query, and when placing a citation.
|
||||
|
||||
On the constructor `sources=[]` means something else. Passing `sources` alongside a database path raises `AmbiguousDatabaseError` immediately, since both say which database to open. Passing `sources=[]` alone raises `ValueError` on entering the client: a selection of nothing to search is a legitimate question, a client over no database is not.
|
||||
|
||||
#### Inspecting the client scope
|
||||
|
||||
```python
|
||||
client.covers_multiple # whether the client covers more than one database
|
||||
client.source_names # database names, in order; known before the client opens
|
||||
client.source # the one database's name, or None for a set
|
||||
|
||||
owner = await client.reader_for("papers") # the client reading that database
|
||||
papers, wiki = await client.clients_for(["papers", "wiki"])
|
||||
```
|
||||
|
||||
`reader_for` and `clients_for` open databases lazily and return borrowed clients. They remain valid while the covering client is open and inherit its read-only mode. The covering client owns and closes their database sessions.
|
||||
|
||||
To learn what a configuration covers without opening anything, resolve it:
|
||||
|
||||
```python
|
||||
from haiku.rag.client import DatabaseScope
|
||||
|
||||
for ref in DatabaseScope.resolve(config).databases:
|
||||
print(ref.name, ref.location) # "haiku.rag", Path(".../haiku.rag.lancedb") when nothing is configured
|
||||
```
|
||||
|
||||
`DatabaseScope.resolve` is pure: it reads the configuration and classifies each location as a local path or a URI.
|
||||
|
||||
### Filtering Search Results
|
||||
|
||||
Filter search results to only include chunks from documents matching specific criteria:
|
||||
|
|
@ -326,11 +395,22 @@ answer, citations = await client.ask(
|
|||
)
|
||||
```
|
||||
|
||||
`client.ask` runs the [rag skill](skills/index.md) and returns `(answer_text, list[Citation])`. Citations include page numbers, section headings, and document references.
|
||||
Attach images to the question, for example to check an image against indexed documents:
|
||||
|
||||
```python
|
||||
answer, citations = await client.ask(
|
||||
"Does this image satisfy the requirements in the design spec?",
|
||||
images=[Path("photo.jpg").read_bytes()],
|
||||
)
|
||||
```
|
||||
|
||||
Images are passed to the model alongside the question. Retrieval stays text-based. The QA model must have `vision: true` in its configuration.
|
||||
|
||||
`client.ask` runs the [RAG capability](capabilities/rag.md) and returns `(answer_text, list[Citation])`. Citations include page numbers, section headings, document references, the document's metadata (`document_meta`), and the cited chunk's raw, unparsed metadata (`chunk_meta`), so UIs can render metadata keys such as a public source URL alongside the citation.
|
||||
|
||||
The QA provider and model are configured in `haiku.rag.yaml` or can be passed directly to the client (see [Configuration](configuration/index.md)).
|
||||
|
||||
See also: [Skills](skills/index.md) for details on the skills the client wraps.
|
||||
See also: [Capabilities](capabilities/index.md) for direct agent composition.
|
||||
|
||||
## Analysis
|
||||
|
||||
|
|
@ -350,15 +430,17 @@ result = await client.analyze(
|
|||
)
|
||||
```
|
||||
|
||||
`client.analyze` runs the [analysis skill](skills/index.md), which writes and executes Python code in a sandboxed environment to solve problems that traditional RAG struggles with: aggregation, computation, and multi-document analysis.
|
||||
`client.analyze` runs the [analysis capability](capabilities/analysis.md), which writes and executes Python code in a sandboxed environment to solve problems that traditional RAG struggles with: aggregation, computation, and multi-document analysis.
|
||||
|
||||
See [Analysis skill](skills/analysis.md) for details on capabilities and configuration.
|
||||
`client.analyze` also accepts `images=` like `client.ask`, requiring `vision: true` on the analysis model (or the QA model when no analysis model is configured).
|
||||
|
||||
See [Analysis capability](capabilities/analysis.md) for details and configuration.
|
||||
|
||||
## Building custom agents
|
||||
|
||||
`client.ask` and `client.analyze` are the convenience wrappers. To build your own Pydantic AI agent against the same database, attach the rag and rag-analysis skills directly with `SkillToolset`. See [Skills](skills/index.md) for the full story and worked examples.
|
||||
`client.ask` and `client.analyze` are convenience wrappers. To build your own Pydantic AI agent, attach the native RAG and analysis capabilities directly. See [Capabilities](capabilities/index.md).
|
||||
|
||||
For the low-level toolset factories under `haiku.rag.tools` (one rung below the skill abstraction), see [Toolsets](tools.md).
|
||||
For the low-level toolset factories under `haiku.rag.tools` (one rung below the capability abstraction), see [Toolsets](tools.md).
|
||||
|
||||
## Importing Pre-Processed Documents
|
||||
|
||||
|
|
@ -434,6 +516,32 @@ await client.vacuum()
|
|||
|
||||
This compacts tables and removes historical versions to keep disk usage in check. It’s safe to run anytime, for example after bulk imports or periodically in long‑running apps.
|
||||
|
||||
### Tags
|
||||
|
||||
Tag the current database state and restore it later, for example after an ingestion run. A tag covers all five tables and is created from a single version snapshot. Create tags with other writers stopped: the snapshot is coordinated within one process only, and a writer in another process can commit between the per-table reads.
|
||||
|
||||
```python
|
||||
await client.store.create_tag("release-1")
|
||||
|
||||
tags = await client.store.list_tags()
|
||||
for name, info in tags.items():
|
||||
print(name, info.tables, info.complete)
|
||||
```
|
||||
|
||||
`restore_tag` brings the live database back to a tagged state. It creates a complete safety tag for the current state before changing any table and returns its name:
|
||||
|
||||
```python
|
||||
safety_tag = await client.store.restore_tag("release-1")
|
||||
```
|
||||
|
||||
Restore is a maintenance operation: stop all other writers first. A tag present on only some tables is partial; `list_tags` reports it via `missing_tables`, and partial tags can be deleted but never restored.
|
||||
|
||||
Delete tags you no longer need. Vacuum retains the oldest tagged version and everything newer:
|
||||
|
||||
```python
|
||||
await client.store.delete_tag("release-1")
|
||||
```
|
||||
|
||||
### Rebuilding the Database
|
||||
|
||||
```python
|
||||
|
|
|
|||
|
|
@ -64,6 +64,7 @@ providers:
|
|||
docling_serve:
|
||||
base_url: http://localhost:5001
|
||||
api_key: "" # Optional API key for authentication
|
||||
timeout: 300 # Per-request timeout in seconds
|
||||
```
|
||||
|
||||
For converter / chunker config options (chunking strategy, tokenizer,
|
||||
|
|
|
|||
|
|
@ -1,199 +0,0 @@
|
|||
# Analysis Skill
|
||||
|
||||
Plain RAG (search → cite → answer) works for questions whose answer sits in a chunk or two: "Who wrote this?", "What does X say about Y?". It struggles when the answer requires touching the whole corpus, reading a specific section in full, or doing arithmetic on the data.
|
||||
|
||||
The analysis skill (`rag-analysis`) gives the agent a second tool (`execute_code`) that runs Python in a sandboxed interpreter against a structured view of your documents. The agent can search, read, count, slice, and compare without leaving the tool call. Citations work the same way as the rag skill.
|
||||
|
||||
`client.analyze`, `haiku-rag analyze`, the MCP `analyze` tool, and the chat TUI (when `-s analysis` is enabled) all run through this skill.
|
||||
|
||||
## When to use it
|
||||
|
||||
Reach for the analysis skill when the question needs more than a search:
|
||||
|
||||
- **Aggregation across the corpus.** "How many documents mention security vulnerabilities?"
|
||||
- **Section-scoped reading.** "Summarize Section 5 of paper Y."
|
||||
- **Structural comparison.** "Do both papers have an Experimental Results section?"
|
||||
- **Computation on retrieved data.** "What's the average revenue across these quarterly reports?"
|
||||
- **Multi-step chains.** Search, filter the results in Python, search again, aggregate, all in one tool call.
|
||||
|
||||
For everyday Q&A, the [RAG skill](rag.md) is faster and cheaper. Attach both and the agent routes.
|
||||
|
||||
## How it works
|
||||
|
||||
Two things make the agent's programs short and the resulting analyses tractable:
|
||||
|
||||
1. **Search and document listing are awaitable inside the code.** `await search(query)` returns the same hits the rag skill sees: chunk IDs, text, source metadata, picture refs. The agent can immediately filter, sort, count, or follow up with another search without exiting the tool call.
|
||||
|
||||
2. **Every document is mounted as a virtual filesystem at `/documents/{id}/`.** The agent reads four files per document: identifiers and metadata, full text, a list of structured items (paragraphs, tables, figures, headings), and a section tree built from the document's headings. The structure exposes what search alone hides. The agent can navigate from a search hit to the section it lives in, slice a single section instead of pulling the whole document, or scan a document's text directly when keyword precision matters.
|
||||
|
||||
A search hit is always a starting point. The agent reads structure around it, drills into the right section, and cites the chunks it actually used. Chunk IDs from search results and chunk IDs surfaced through the VFS are both accepted by `cite`.
|
||||
|
||||
### Sandbox guarantees
|
||||
|
||||
The interpreter is [pydantic-monty](https://github.com/pydantic/monty), isolated from the host:
|
||||
|
||||
- **Virtual filesystem only.** `/documents/` is the entire FS.
|
||||
- **No network.** HTTP, sockets, and the `requests` family are unavailable.
|
||||
- **Limited imports.** Only `json`, `re`, `math`, `pathlib`.
|
||||
- **Execution timeout** (default 60s, configurable via `analysis.code_timeout`).
|
||||
- **Output truncation** (default 50000 chars, configurable via `analysis.max_output_chars`).
|
||||
- **Execution budget** (default 15 calls, configurable via `analysis.max_executions`). Past the budget, `execute_code` returns a notice telling the skill to answer from what it has instead of running more code.
|
||||
|
||||
Variables persist between `execute_code` calls within one invocation, so the agent can build state step by step. A fresh sandbox is built per `client.analyze` call.
|
||||
|
||||
## Tools
|
||||
|
||||
| Tool | Purpose |
|
||||
|------|---------|
|
||||
| `search(query, limit?)` | Hybrid search with context expansion. Same as the RAG skill's `search`. |
|
||||
| `execute_code(code)` | Run Python in a sandboxed interpreter with VFS access. |
|
||||
| `cite(chunk_ids)` | Register chunk IDs as citations. Call before producing the final answer. |
|
||||
|
||||
`list_documents` isn't exposed as a top-level tool but is available inside `execute_code` as `await list_documents()`.
|
||||
|
||||
## State
|
||||
|
||||
`AnalysisState` lives under the `"analysis"` namespace:
|
||||
|
||||
```python
|
||||
class AnalysisState(BaseModel):
|
||||
document_filter: str | None = None
|
||||
executions: list[CodeExecutionEntry] = []
|
||||
citation_index: dict[str, Citation] = {}
|
||||
citations: list[str] = []
|
||||
searches: dict[str, list[SearchResult]] = {}
|
||||
```
|
||||
|
||||
- **document_filter** — SQL WHERE clause applied to `search` and the VFS. The LLM can't bypass it: both views are scoped.
|
||||
- **executions** — Each `execute_code` call appends an entry with code, stdout, stderr, success. Cleared at the start of each invocation.
|
||||
- **citation_index** — All citations indexed by chunk ID. Accumulates across invocations.
|
||||
- **citations** — Chunk IDs registered via `cite` during the current invocation. Deduplicated, cleared per-invocation.
|
||||
- **searches** — Search results from both the `search` tool and sandbox-internal searches. Cleared per-invocation.
|
||||
|
||||
## `create_skill(db_path?, config?)`
|
||||
|
||||
```python
|
||||
from haiku.rag.skills.analysis import create_skill
|
||||
|
||||
skill = create_skill(db_path="my.lancedb")
|
||||
```
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `db_path` | `None` | Path to LanceDB database. Falls back to `HAIKU_RAG_DB` env var, then config default. |
|
||||
| `config` | `None` | `AppConfig` instance. Falls back to `get_config()`. |
|
||||
|
||||
## Use it
|
||||
|
||||
### From `client.analyze`
|
||||
|
||||
```python
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
async with HaikuRAG("my.lancedb") as client:
|
||||
result = await client.analyze("How many documents mention 'security'?")
|
||||
print(result.answer)
|
||||
for citation in result.citations:
|
||||
print(citation.uri, citation.title)
|
||||
```
|
||||
|
||||
`client.analyze` runs the skill end-to-end and returns an `AnalysisResult` with `answer` and `citations`. The executed Python programs live on `AnalysisState.executions` during the run, not on the returned result.
|
||||
|
||||
### Combine with the RAG skill
|
||||
|
||||
```python
|
||||
from haiku.rag.skills.rag import create_skill as create_rag_skill
|
||||
from haiku.rag.skills.analysis import create_skill as create_analysis_skill
|
||||
from haiku.skills.agent import SkillToolset
|
||||
|
||||
rag = create_rag_skill(db_path="my.lancedb")
|
||||
analysis = create_analysis_skill(db_path="my.lancedb")
|
||||
toolset = SkillToolset(skills=[rag, analysis])
|
||||
```
|
||||
|
||||
The agent routes Q&A to the rag skill and computational questions to rag-analysis.
|
||||
|
||||
## What the agent actually writes
|
||||
|
||||
You don't write these programs yourself. The agent does, inside `execute_code`. Seeing the shape helps when you tune prompts, debug a run via `AnalysisState.executions`, or design a custom skill.
|
||||
|
||||
**Aggregate across the corpus.** *"How many documents mention security vulnerabilities?"*
|
||||
|
||||
```python
|
||||
hits = await search("security vulnerability", limit=50)
|
||||
|
||||
doc_ids = {h['document_id'] for h in hits}
|
||||
print(f"{len(doc_ids)} documents mention security vulnerabilities")
|
||||
|
||||
# Cite the top hit per document
|
||||
seen = set()
|
||||
for hit in hits:
|
||||
if hit['document_id'] not in seen:
|
||||
seen.add(hit['document_id'])
|
||||
await cite(hit['chunk_id'])
|
||||
```
|
||||
|
||||
**Read one section in depth.** *"Summarize Section 5."*
|
||||
|
||||
```python
|
||||
from pathlib import Path
|
||||
import json
|
||||
|
||||
doc_id = "..." # from a prior search or list_documents
|
||||
toc = json.loads(Path(f'/documents/{doc_id}/toc.json').read_text())
|
||||
|
||||
section = next(n for n in toc['tree'] if n['title'].startswith('5'))
|
||||
start, end = section['item_range']
|
||||
lines = Path(f'/documents/{doc_id}/items.jsonl').read_text().splitlines()[start:end]
|
||||
|
||||
for line in lines:
|
||||
print(json.loads(line)['text'])
|
||||
|
||||
await cite(section['chunk_ids'])
|
||||
```
|
||||
|
||||
The section node already aggregates the chunks underneath it, so the agent cites the whole section without a separate search.
|
||||
|
||||
**Compare structure across documents.** *"Do both papers have an Experimental Results section?"*
|
||||
|
||||
```python
|
||||
from pathlib import Path
|
||||
import json
|
||||
|
||||
for doc_id in ["doc-a-id", "doc-b-id"]:
|
||||
toc = json.loads(Path(f'/documents/{doc_id}/toc.json').read_text())
|
||||
print(f"\n=== {toc['title']} ===")
|
||||
for node in toc['tree']:
|
||||
if 'experiment' in node['title'].lower():
|
||||
print(f" {node['title']} (pages {node['page_numbers']})")
|
||||
await cite(node['chunk_ids'])
|
||||
```
|
||||
|
||||
## Context filter
|
||||
|
||||
The `filter` parameter is enforced at the deps layer. The LLM can't bypass it: both the VFS and search results are scoped to the filter.
|
||||
|
||||
```python
|
||||
result = await client.analyze(
|
||||
"Summarize all findings",
|
||||
filter="uri LIKE '%confidential%'"
|
||||
)
|
||||
```
|
||||
|
||||
Useful for scoping to a corpus subset, enforcing access control, or restricting context.
|
||||
|
||||
## Configuration
|
||||
|
||||
```yaml
|
||||
analysis:
|
||||
model:
|
||||
provider: anthropic
|
||||
name: claude-sonnet-4-20250514
|
||||
code_timeout: 60.0 # Max seconds per code execution
|
||||
max_output_chars: 50000 # Truncate output after this many chars
|
||||
max_executions: 15 # Max execute_code calls per question
|
||||
```
|
||||
|
||||
When `analysis.model` is unset, the skill falls back to `qa.model`.
|
||||
|
||||
See [Search and question answering](../configuration/qa.md#analysis-configuration) for the full set.
|
||||
|
|
@ -1,124 +0,0 @@
|
|||
# Custom Skills
|
||||
|
||||
The two skills haiku.rag ships work against any LanceDB database. When you want a *domain-specific* skill that bundles its own data, prompt, and tool surface (for example, a "recipes" skill that knows about cooking and ships with a recipes database), generate one with `haiku-rag create-skill`.
|
||||
|
||||
The generated package is a regular pip-installable Python package that registers as a `haiku.skills` entry point. Any haiku.skills-aware host (haiku.skills CLI, your own agent, the AG-UI adapter) discovers it automatically.
|
||||
|
||||
## When to use a custom skill
|
||||
|
||||
- The model should consult a specific knowledge base for a specific kind of question, alongside other skills.
|
||||
- You want a different instruction prompt than the generic `rag` skill (different tone, refusal style, domain rules).
|
||||
- You want to ship a knowledge base plus its prompt as one distributable unit.
|
||||
- You're running multiple skills against different databases in the same agent.
|
||||
|
||||
If you just want to point a haiku.rag database at your own model and prompt, configure `haiku.rag.yaml` and use the built-in `rag` skill. No custom package needed.
|
||||
|
||||
## Generate
|
||||
|
||||
```bash
|
||||
haiku-rag create-skill \
|
||||
--name recipes \
|
||||
--db /path/to/recipes.lancedb \
|
||||
--tools search,cite \
|
||||
--description "Recipe and cooking knowledge base" \
|
||||
--preamble "You are a culinary expert helping with recipes and cooking techniques."
|
||||
```
|
||||
|
||||
Then install and use:
|
||||
|
||||
```bash
|
||||
uv pip install -e ./recipes-skill
|
||||
|
||||
haiku-skills list --use-entrypoints
|
||||
# recipes — Recipe and cooking knowledge base
|
||||
|
||||
haiku-skills chat --use-entrypoints --skill recipes
|
||||
```
|
||||
|
||||
### Flags
|
||||
|
||||
| Flag | Description | Default |
|
||||
|------|-------------|---------|
|
||||
| `--name` | Skill name (lowercase alphanumeric and hyphens). Required. | — |
|
||||
| `--db` | Path to the LanceDB database to embed. Required. | — |
|
||||
| `--description` | One-line skill description. The agent reads this to decide when to invoke. | Standard RAG description |
|
||||
| `--tools` | Comma-separated tool subset, or `all`. | `all` |
|
||||
| `--preamble` | Custom preamble for the skill's instructions. | Standard RAG preamble |
|
||||
| `--config-file` | Path to a `haiku.rag.yaml` to embed alongside the database. | None |
|
||||
| `--output` / `-o` | Output directory. | Current directory |
|
||||
|
||||
### Available tools
|
||||
|
||||
`cite`, `execute_code`, `get_document`, `list_documents`, `search`.
|
||||
|
||||
Drop `execute_code` from `--tools` if the skill shouldn't run sandboxed Python. That gives you a search-and-cite-only skill with no analysis capabilities.
|
||||
|
||||
## Anatomy of a generated skill
|
||||
|
||||
```
|
||||
{name}-skill/
|
||||
├── pyproject.toml
|
||||
└── {name}_skill/
|
||||
├── __init__.py # create_skill() entry point
|
||||
├── SKILL.md # Skill metadata and instructions
|
||||
└── assets/
|
||||
├── {name}.lancedb/ # The embedded database
|
||||
└── haiku.rag.yaml # Optional config (only if --config-file passed)
|
||||
```
|
||||
|
||||
- **`SKILL.md`** carries the instruction prompt the agent will follow. The frontmatter includes the skill name and description. Everything below is the prompt body. Edit this to change behavior.
|
||||
- **`__init__.py`** exposes `create_skill()` (the entry point) and `visualize_chunk()` for rendering visual grounding.
|
||||
- **`assets/{name}.lancedb/`** is the database, shipped inside the package.
|
||||
- **`assets/haiku.rag.yaml`** (optional) pins provider settings the skill needs.
|
||||
|
||||
The package can be installed locally with `uv pip install -e .` or published to PyPI.
|
||||
|
||||
## Generating visual grounding from a custom skill
|
||||
|
||||
Each generated skill exposes a `visualize_chunk()` function that returns the chunk's bounding boxes rendered onto its source page:
|
||||
|
||||
```python
|
||||
from recipes_skill import visualize_chunk
|
||||
|
||||
images = await visualize_chunk(chunk_id)
|
||||
# images is a list of PIL.Image objects, one per page the chunk covers
|
||||
images[0].save("citation.png")
|
||||
```
|
||||
|
||||
Pass chunk IDs from skill citations or search results. Same prerequisites as elsewhere in haiku.rag: documents need stored page images, and the chunk must come from a PDF or other docling-converted source.
|
||||
|
||||
## Multi-skill agents
|
||||
|
||||
Each generated skill is self-contained with its own database and instructions. Compose multiple skills in one agent and the model routes between them via their descriptions:
|
||||
|
||||
```python
|
||||
from recipes_skill import create_skill as create_recipes_skill
|
||||
from medic_skill import create_skill as create_medic_skill
|
||||
from haiku.skills.agent import SkillToolset
|
||||
from haiku.skills.prompts import build_system_prompt
|
||||
from pydantic_ai import Agent
|
||||
|
||||
recipes = create_recipes_skill()
|
||||
medic = create_medic_skill()
|
||||
toolset = SkillToolset(skills=[recipes, medic])
|
||||
|
||||
agent = Agent(
|
||||
"openai-chat:gpt-4o",
|
||||
instructions=build_system_prompt(toolset.skill_catalog),
|
||||
toolsets=[toolset],
|
||||
)
|
||||
|
||||
await agent.run("What's the optimal temperature for braising short ribs?")
|
||||
# Routes to recipes
|
||||
|
||||
await agent.run("What's the field treatment for tension pneumothorax?")
|
||||
# Routes to medic
|
||||
```
|
||||
|
||||
Each skill maintains state under its own namespace (`recipes`, `medic`, …), so citations and searches don't collide.
|
||||
|
||||
## Writing a skill from scratch
|
||||
|
||||
`create-skill` is the convenience path. If you need full control over the tools, state model, or instruction loading, write the skill against [haiku.skills](https://github.com/ggozad/haiku.skills) directly. The generated package in `{name}_skill/__init__.py` is a good reference. It composes haiku.rag's `_tools` factory with a `haiku.skills.Skill` and registers under the `haiku.skills` entry point group in `pyproject.toml`.
|
||||
|
||||
See the haiku.skills repository for the full Skill contract.
|
||||
|
|
@ -1,105 +0,0 @@
|
|||
# Skills
|
||||
|
||||
Skills put haiku.rag in front of a model. A skill bundles tools, an instruction prompt, and managed state into a unit that drops into any Pydantic AI agent via `SkillToolset`. haiku.rag ships two skills and supports custom skills.
|
||||
|
||||
Built on [haiku.skills](https://github.com/ggozad/haiku.skills).
|
||||
|
||||
## Available skills
|
||||
|
||||
| Skill | What it does | Reach for it when |
|
||||
|-------|--------------|-------------------|
|
||||
| [`rag`](rag.md) | Search, retrieve, and cite content from a knowledge base. | The model needs to find and quote evidence from documents. |
|
||||
| [`rag-analysis`](analysis.md) | Same as `rag`, plus a sandboxed Python interpreter mounting every document as a virtual filesystem. | The question requires computation, aggregation, structural traversal, or section-scoped reading. |
|
||||
|
||||
To ship your own skill (bundled with its own database), see [Custom skills](custom.md).
|
||||
|
||||
## Your first agent
|
||||
|
||||
```python
|
||||
from haiku.rag.skills.rag import create_skill
|
||||
from haiku.skills.agent import SkillToolset
|
||||
from haiku.skills.prompts import build_system_prompt
|
||||
from pydantic_ai import Agent
|
||||
|
||||
rag = create_skill(db_path="my.lancedb")
|
||||
toolset = SkillToolset(skills=[rag])
|
||||
|
||||
agent = Agent(
|
||||
"openai-chat:gpt-4o",
|
||||
instructions=build_system_prompt(toolset.skill_catalog),
|
||||
toolsets=[toolset],
|
||||
)
|
||||
|
||||
result = await agent.run("What does the knowledge base say about X?")
|
||||
print(result.output)
|
||||
```
|
||||
|
||||
The skill searches, cites, and answers. You supply the model and the question.
|
||||
|
||||
To run analysis against the same database, swap in the `rag-analysis` skill or attach both:
|
||||
|
||||
```python
|
||||
from haiku.rag.skills.rag import create_skill as create_rag_skill
|
||||
from haiku.rag.skills.analysis import create_skill as create_analysis_skill
|
||||
|
||||
rag = create_rag_skill(db_path="my.lancedb")
|
||||
analysis = create_analysis_skill(db_path="my.lancedb")
|
||||
toolset = SkillToolset(skills=[rag, analysis])
|
||||
```
|
||||
|
||||
The agent reads each skill's description and routes questions itself. See the individual skill pages for the tool surface, state model, and worked examples.
|
||||
|
||||
## State
|
||||
|
||||
Each skill manages its own state under a dedicated namespace. State is synced via the AG-UI protocol when using `AGUIAdapter`.
|
||||
|
||||
```python
|
||||
rag_state = toolset.get_namespace("rag")
|
||||
analysis_state = toolset.get_namespace("analysis")
|
||||
```
|
||||
|
||||
Both state models track citations, the current document filter, and per-turn searches. Analysis state also carries the sandbox execution log. See [RAG skill: state](rag.md#state) and [Analysis skill: state](analysis.md#state).
|
||||
|
||||
## Database path resolution
|
||||
|
||||
Both skills resolve the database path in the same order:
|
||||
|
||||
1. `db_path` argument passed to `create_skill()`
|
||||
2. `HAIKU_RAG_DB` environment variable
|
||||
3. Config default (`config.storage.data_dir / "haiku.rag.lancedb"`)
|
||||
|
||||
## AG-UI streaming for web apps
|
||||
|
||||
For browser apps, use pydantic-ai's `AGUIAdapter` to stream tool calls, text, and state deltas:
|
||||
|
||||
```python
|
||||
from pydantic_ai.ui.ag_ui import AGUIAdapter
|
||||
|
||||
adapter = AGUIAdapter(agent=agent, run_input=run_input)
|
||||
event_stream = adapter.run_stream()
|
||||
sse_event_stream = adapter.encode_stream(event_stream)
|
||||
```
|
||||
|
||||
See the [Web application](../apps.md) reference implementation.
|
||||
|
||||
## Exposing via MCP
|
||||
|
||||
To use a skill from Claude Desktop or another MCP-aware client, run the MCP server:
|
||||
|
||||
```bash
|
||||
haiku-rag mcp --stdio
|
||||
```
|
||||
|
||||
The server exposes the skill tools (search, ask, analyze) over MCP. See [MCP](../mcp.md).
|
||||
|
||||
## Discovery
|
||||
|
||||
Skills are registered as Python entry points under `haiku.skills`. They are discovered automatically:
|
||||
|
||||
```bash
|
||||
haiku-skills list --use-entrypoints
|
||||
# rag — Search, retrieve and analyze documents using RAG.
|
||||
# rag-analysis — Analyze documents using code execution in a sandboxed interpreter.
|
||||
```
|
||||
|
||||
This is what makes custom skills installable as plain pip packages. See [Custom skills](custom.md).
|
||||
|
|
@ -1,179 +0,0 @@
|
|||
# RAG Skill
|
||||
|
||||
The `rag` skill answers questions over a knowledge base with hybrid search, structure-aware context expansion, and explicit citations. `client.ask`, `haiku-rag ask`, the MCP `ask_question` tool, and the chat TUI all run through this skill.
|
||||
|
||||
## When to use it
|
||||
|
||||
- The model needs to find and quote evidence from a document corpus.
|
||||
- You want citations under every answer.
|
||||
- You're building a Q&A agent, a documentation chatbot, or any RAG-style integration.
|
||||
|
||||
If the question requires *computation* over the corpus (counts, aggregates, comparisons, section-scoped reading), reach for the [Analysis skill](analysis.md) instead, or attach both.
|
||||
|
||||
## Tools
|
||||
|
||||
| Tool | Purpose |
|
||||
|------|---------|
|
||||
| `search(query, limit?)` | Hybrid search (vector + full-text) with section-aware context expansion. Returns `chunk_id`, content, `doc_item_refs`, `picture_refs`, `picture_captions`, source metadata. |
|
||||
| `cite(chunk_ids)` | Register chunk IDs as citations for the current answer. The agent calls this before writing the final response. |
|
||||
|
||||
For corpus enumeration or full-document reads, reach for the [Analysis skill](analysis.md), which exposes `await list_documents()` and a `/documents/{id}/content.txt` virtual filesystem inside `execute_code`. Both are also available as opt-in tools when building a [custom skill](custom.md).
|
||||
|
||||
## State
|
||||
|
||||
The skill manages a `RAGState` under the `"rag"` namespace:
|
||||
|
||||
```python
|
||||
class RAGState(BaseModel):
|
||||
citation_index: dict[str, Citation] = {}
|
||||
citations: list[str] = []
|
||||
document_filter: str | None = None
|
||||
searches: dict[str, list[SearchResult]] = {}
|
||||
```
|
||||
|
||||
- **citation_index** — All citations indexed by chunk ID. Accumulates across invocations so historical chunk IDs stay resolvable in UI scrollback.
|
||||
- **citations** — Chunk IDs registered via `cite` during the current invocation. Deduplicated, cleared at the start of each invocation.
|
||||
- **document_filter** — SQL WHERE clause applied to `search`. Persists across invocations.
|
||||
- **searches** — Search results keyed by query string. Cleared at the start of each invocation.
|
||||
|
||||
## `create_skill(db_path?, config?)`
|
||||
|
||||
```python
|
||||
from haiku.rag.skills.rag import create_skill
|
||||
|
||||
skill = create_skill(db_path="my.lancedb")
|
||||
```
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `db_path` | `None` | Path to LanceDB database. Falls back to `HAIKU_RAG_DB` env var, then config default. |
|
||||
| `config` | `None` | `AppConfig` instance. Falls back to `get_config()`. |
|
||||
|
||||
## Examples
|
||||
|
||||
### Minimal agent
|
||||
|
||||
```python
|
||||
from haiku.rag.skills.rag import create_skill
|
||||
from haiku.skills.agent import SkillToolset
|
||||
from haiku.skills.prompts import build_system_prompt
|
||||
from pydantic_ai import Agent
|
||||
|
||||
rag = create_skill(db_path="my.lancedb")
|
||||
toolset = SkillToolset(skills=[rag])
|
||||
|
||||
agent = Agent(
|
||||
"openai-chat:gpt-4o",
|
||||
instructions=build_system_prompt(toolset.skill_catalog),
|
||||
toolsets=[toolset],
|
||||
)
|
||||
|
||||
result = await agent.run("What does the manual say about safety procedures?")
|
||||
print(result.output)
|
||||
|
||||
# Inspect what the model cited
|
||||
state = toolset.get_namespace("rag")
|
||||
for chunk_id in state.citations:
|
||||
citation = state.citation_index[chunk_id]
|
||||
print(f"- {citation.document_title}: {citation.content[:100]}…")
|
||||
```
|
||||
|
||||
### Domain customization
|
||||
|
||||
Set a domain preamble in `haiku.rag.yaml` and the skill picks it up:
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
domain_preamble: |
|
||||
The knowledge base contains the operations manual for the Helios solar array.
|
||||
"The array" or unqualified specs refer to Helios. Terminology like "string"
|
||||
refers to a series-connected panel chain, not text.
|
||||
```
|
||||
|
||||
To scope a session to a subset of documents, set the filter on the namespace state:
|
||||
|
||||
```python
|
||||
state = toolset.get_namespace("rag")
|
||||
state.document_filter = "uri LIKE '%helios/v4/%'"
|
||||
|
||||
result = await agent.run("What's the maintenance interval for the inverters?")
|
||||
```
|
||||
|
||||
The filter applies to every `search` call for the rest of the session, and the model can't bypass it from inside.
|
||||
|
||||
### Combining with the analysis skill
|
||||
|
||||
Attach both skills and the agent routes between them:
|
||||
|
||||
```python
|
||||
from haiku.rag.skills.rag import create_skill as create_rag_skill
|
||||
from haiku.rag.skills.analysis import create_skill as create_analysis_skill
|
||||
|
||||
rag = create_rag_skill(db_path="my.lancedb")
|
||||
analysis = create_analysis_skill(db_path="my.lancedb")
|
||||
toolset = SkillToolset(skills=[rag, analysis])
|
||||
|
||||
agent = Agent(
|
||||
"openai-chat:gpt-4o",
|
||||
instructions=build_system_prompt(toolset.skill_catalog),
|
||||
toolsets=[toolset],
|
||||
)
|
||||
|
||||
# Q&A → uses rag
|
||||
await agent.run("What safety equipment is required on-site?")
|
||||
|
||||
# Computational question → uses rag-analysis
|
||||
await agent.run("How many checklists mention torque specifications?")
|
||||
```
|
||||
|
||||
### Streaming to a web frontend
|
||||
|
||||
Wrap the agent with `AGUIAdapter` to stream tool calls, text deltas, and state changes to a CopilotKit-style frontend:
|
||||
|
||||
```python
|
||||
from pydantic_ai.ui.ag_ui import AGUIAdapter
|
||||
|
||||
adapter = AGUIAdapter(agent=agent, run_input=run_input)
|
||||
sse_stream = adapter.encode_stream(adapter.run_stream())
|
||||
```
|
||||
|
||||
See the [Web application](../apps.md) reference implementation for the full Starlette + Next.js setup.
|
||||
|
||||
### Exposing via MCP
|
||||
|
||||
To call the skill from Claude Desktop (or any MCP client), run the MCP server:
|
||||
|
||||
```bash
|
||||
haiku-rag mcp --stdio
|
||||
```
|
||||
|
||||
The exposed `ask_question` tool runs this skill. See [MCP](../mcp.md) for the configuration block.
|
||||
|
||||
## Configuration
|
||||
|
||||
The skill picks up its model and search behavior from the standard config sections:
|
||||
|
||||
```yaml
|
||||
qa:
|
||||
model:
|
||||
provider: ollama
|
||||
name: gpt-oss
|
||||
enable_thinking: true
|
||||
temperature: 0.3
|
||||
vision: false # set true for vision-capable QA models
|
||||
max_searches: 3
|
||||
|
||||
search:
|
||||
limit: 5
|
||||
max_context_chars: 5000
|
||||
```
|
||||
|
||||
See [Search and question answering](../configuration/qa.md) for every knob.
|
||||
|
||||
## Vision support
|
||||
|
||||
When `qa.model.vision: true` is set, the skill's `search` tool attaches picture bytes to its tool returns as `BinaryContent`. The model can then read figures, diagrams, and screenshots directly alongside the surrounding text. Requires `processing.pictures != none` so the bytes exist on disk. See the [pictures × embedder × QA model matrix](../configuration/processing.md#picture-handling) for the combinations that make sense.
|
||||
|
||||
## Customizing the skill prompt
|
||||
|
||||
The skill's instruction prompt lives in `SKILL.md` inside the package. For behavior changes (different phrasing, refusal style, additional rules), the supported path is to fork the skill with `haiku-rag create-skill` and edit the generated `SKILL.md`. The `domain_preamble` field above is for *what the corpus is about*, not for *how the agent should behave*. See [Custom skills](custom.md).
|
||||
|
|
@ -1,12 +1,12 @@
|
|||
# Toolsets
|
||||
|
||||
haiku.rag exposes its RAG capabilities as [haiku.skills](https://github.com/ggozad/haiku.skills) skills. For most integrations, see [Skills](skills/index.md).
|
||||
For agent integrations, use the native Pydantic AI [capabilities](capabilities/index.md). This page documents the lower-level toolsets used by other haiku.rag surfaces.
|
||||
|
||||
For lower-level access, `haiku.rag.tools` provides individual `FunctionToolset` factories that the skills themselves compose.
|
||||
For lower-level access, `haiku.rag.tools` provides individual `FunctionToolset` factories used across haiku.rag.
|
||||
|
||||
## Low-Level Toolsets
|
||||
|
||||
For advanced use cases, individual toolset factories are available in `haiku.rag.tools`. These are the same primitives the rag and rag-analysis skills compose, and can be reused to build custom agents.
|
||||
For advanced use cases, individual toolset factories are available in `haiku.rag.tools` and can be reused to build custom agents.
|
||||
|
||||
### RAGDeps Protocol
|
||||
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
How to adjust haiku.rag's pipeline for better retrieval and answer quality. For individual setting definitions and defaults, see [Configuration](configuration/index.md).
|
||||
|
||||
For ingester-side tuning (worker count, claim timeout, retry policy, backpressure, circuit breakers), see [Ingester → Workers and retry](ingester.md#workers-and-retry).
|
||||
For ingester-side tuning (worker count, lease TTL and heartbeat, retry policy, backpressure, circuit breakers), see [Ingester → Workers and retry](ingester.md#workers-and-retry).
|
||||
|
||||
## Pipeline Overview
|
||||
|
||||
|
|
@ -12,7 +12,7 @@ Documents flow through: **chunking → embedding → hybrid search (vector + FTS
|
|||
|
||||
### Chunking
|
||||
|
||||
`chunk_size` controls the granularity of retrieval. Smaller chunks match queries more precisely but carry less context each. Larger chunks provide more surrounding information but dilute relevance signals. On the Wix benchmark, increasing from 256 to 512 tokens raised MAP from 0.43 to 0.45 on plain text, a modest gain that also increases token cost per result. See [Processing](configuration/processing.md#chunk-size) for configuration.
|
||||
`chunk_size` controls the granularity of retrieval. Smaller chunks match queries more precisely but carry less context each. Larger chunks provide more surrounding information but dilute relevance signals. See [Processing](configuration/processing.md#chunk-size) for configuration.
|
||||
|
||||
`chunker_type` selects between `hybrid` (default) and `hierarchical` chunking. Hierarchical chunking preserves the document's heading structure and works better for deeply nested or structured content. See [Chunking Strategies](configuration/processing.md#chunking-strategies).
|
||||
|
||||
|
|
@ -22,7 +22,7 @@ Larger embedding models produce better representations at the cost of slower ind
|
|||
|
||||
### Reranking
|
||||
|
||||
When configured, a cross-encoder reranker re-scores 10x the requested candidates and returns the top results. This adds latency but improves precision. On the Wix benchmark, adding `mxbai-rerank-base-v2` raised MAP from 0.34 to 0.39 on HTML content. See [Search Settings](configuration/qa.md#search-settings) for how reranking integrates with search.
|
||||
When configured, a cross-encoder reranker re-scores 10x the requested candidates and returns the top results. This adds latency but improves precision. See [Search Settings](configuration/qa.md#search-settings) for how reranking integrates with search.
|
||||
|
||||
### Search Settings
|
||||
|
||||
|
|
@ -34,7 +34,7 @@ Context expansion is automatic and section-aware. Search results are expanded to
|
|||
|
||||
Model and temperature selection affect answer quality directly. See [Providers](configuration/providers.md#model-settings) for options.
|
||||
|
||||
`domain_preamble` prepends domain context to the rag and rag-analysis skill instructions. Use it to describe what the knowledge base contains and clarify domain-specific terminology. See [Prompt Customization](configuration/prompts.md).
|
||||
`domain_preamble` prepends domain context to the RAG and analysis capability instructions. Use it to describe what the knowledge base contains and clarify domain-specific terminology. See [Prompt Customization](configuration/prompts.md).
|
||||
|
||||
## What Requires a Rebuild
|
||||
|
||||
|
|
@ -46,7 +46,7 @@ Model and temperature selection affect answer quality directly. See [Providers](
|
|||
|
||||
## Inspector
|
||||
|
||||
The inspector shows what your model would actually receive for a given query. Run it against your database and step through the same hybrid search, context expansion, and chunk previews the rag skill uses at runtime. Press `c` on a chunk and you see the exact context the LLM would get back from a search hit.
|
||||
The inspector shows what your model would actually receive for a given query. Run it against your database and step through the same hybrid search, context expansion, and chunk previews the RAG capability uses at runtime. Press `c` on a chunk and you see the exact context the LLM would get back from a search hit.
|
||||
|
||||
```bash
|
||||
haiku-rag inspect
|
||||
|
|
@ -81,11 +81,11 @@ Mouse: click to select, scroll to view content.
|
|||
|
||||
### Search
|
||||
|
||||
Press `/` to open the search modal. Type a query and press `Enter`. The left panel lists results with relevance scores like `[0.95] content preview`. The right panel shows the full chunk and its metadata. `↑` / `↓` navigates results, `Enter` jumps to the document and chunk, `Esc` closes the modal. Search uses the same hybrid (vector + full-text) retrieval the rag skill uses.
|
||||
Press `/` to open the search modal. Type a query and press `Enter`. The left panel lists results with relevance scores like `[0.95] content preview`. The right panel shows the full chunk and its metadata. `↑` / `↓` navigates results, `Enter` jumps to the document and chunk, `Esc` closes the modal. Search uses the same hybrid (vector + full-text) retrieval the RAG capability uses.
|
||||
|
||||
### Context expansion (`c`)
|
||||
|
||||
Press `c` on a chunk to see the expanded context that would be fed to the rag skill. This is where you find out whether your `chunk_size`, `chunker_type`, and `max_context_chars` settings actually deliver the surrounding content the model needs. The modal shows:
|
||||
Press `c` on a chunk to see the expanded context that would be fed to the RAG capability. This is where you find out whether your `chunk_size`, `chunker_type`, and `max_context_chars` settings actually deliver the surrounding content the model needs. The modal shows:
|
||||
|
||||
- The expanded text. Section-aware expansion stays within section boundaries on structured documents and fills `max_context_chars` outward on unstructured ones.
|
||||
- Source document, content type, and relevance score.
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ You also need [Ollama](https://ollama.com/) for the default embedding and answer
|
|||
|
||||
```bash
|
||||
ollama pull qwen3-embedding:4b
|
||||
ollama pull gpt-oss
|
||||
ollama pull qwen3.8
|
||||
```
|
||||
|
||||
!!! note "Prefer OpenAI?"
|
||||
|
|
@ -81,6 +81,6 @@ haiku-rag ask "Who wrote haiku.rag?"
|
|||
- [Chat](chat.md): sessions, citations, and the full TUI.
|
||||
- [CLI reference](cli.md): every command.
|
||||
- [Python API](python.md): use haiku.rag in your own code.
|
||||
- [Skills](skills/index.md): the rag and rag-analysis skills the client wraps.
|
||||
- [Capabilities](capabilities/index.md): native RAG and analysis components used by the client.
|
||||
- [Tuning](tuning.md): better retrieval.
|
||||
- [Configuration](configuration/index.md): every setting.
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
# Haiku RAG - Evaluations
|
||||
# haiku.rag - Evaluations
|
||||
|
||||
Internal benchmarking and evaluation scripts for haiku.rag.
|
||||
|
||||
|
|
@ -8,7 +8,9 @@ This package is not published to PyPI and is only used for development and testi
|
|||
|
||||
Contains evaluation scripts for benchmarking RAG retrieval and QA performance. Available datasets:
|
||||
|
||||
- WiX (`wix`)
|
||||
- HotpotQA (`hotpotqa`) — multi-hop QA over Wikipedia paragraphs (distractor validation split, 7,405 questions, two gold documents per question)
|
||||
- MTRAG ClapNQ (`mtrag_clapnq`, `mtrag_clapnq_rewrite`) — IBM's multi-turn RAG benchmark, ClapNQ (Wikipedia) domain: 183,408 passages, 208 retrieval queries with binary qrels, 224 generation tasks. The base key retrieves with the raw last user turn; the `_rewrite` variant uses the human standalone rewrites (both share one database). Retrieval reports Recall@5/@10, nDCG@5/@10, and MAP against IBM's published setup. QA replays each task's reference conversation prefix as message history and answers the final turn; the judge sees the conversation as a transcript, citation MAP is scored only on turns with gold passages, and refusal precision/recall is reported against the answerability labels. Generation scores are internal (our judge and rubric), not comparable with IBM's published generation numbers. The `mtrag_clapnq_live` key replays whole conversations (one case per conversation, `--limit` counts conversations) through a single capability session, carrying the model's own answers and tool history across turns; it reports the same outcomes per turn plus micro (per-turn) and macro (per-conversation) aggregates.
|
||||
- FRAMES (`frames`) — multi-hop QA (822 questions, 2-23 gold Wikipedia articles per question; 2 of the original 824 questions are excluded because a linked article has been deleted from Wikipedia). The corpus is the union of the 2,521 linked articles, fetched from the Wikipedia REST API at current revision (revision id and fetch date recorded in the article cache) with navigation chrome stripped. There is no official FRAMES evaluation setup; numbers here correspond to the paper's multi-step retrieval setting (fixed corpus, agentic retrieval, judged accuracy) and are not comparable to its closed-book, oracle-prompt, or web-search settings. Answers were authored against ~2024 revisions and may have drifted with article content.
|
||||
- OpenRAG Bench, two variants:
|
||||
- `orb_text` — text embedder (`qwen3-embedding:4b`, 2560-dim) with VLM picture descriptions baked into chunk content at ingest. Use for text-only retrieval/QA against figure-rich corpora.
|
||||
- `orb_multimodal` — multimodal embedder (`qwen3-vl-embedding-8b`, 4096-dim) with picture vectors in the same space as text. Use for cross-modal retrieval (text-as-query → figure hits, image-as-query) and vision QA where the figure itself is the answer.
|
||||
|
|
@ -19,41 +21,41 @@ After installing the package, you can run evaluations using the `evaluations` co
|
|||
|
||||
```bash
|
||||
# Run retrieval + QA benchmarks
|
||||
evaluations run wix
|
||||
evaluations run hotpotqa
|
||||
evaluations run orb_text
|
||||
|
||||
# Use a custom config file
|
||||
evaluations run wix --config /path/to/haiku.rag.yaml
|
||||
evaluations run hotpotqa --config /path/to/haiku.rag.yaml
|
||||
|
||||
# Override the database path
|
||||
evaluations run wix --db /path/to/custom.lancedb
|
||||
evaluations run hotpotqa --db /path/to/custom.lancedb
|
||||
|
||||
# Skip database population and run only benchmarks
|
||||
evaluations run wix --skip-db
|
||||
evaluations run hotpotqa --skip-db
|
||||
|
||||
# Skip specific benchmarks
|
||||
evaluations run wix --skip-retrieval
|
||||
evaluations run wix --skip-qa
|
||||
evaluations run hotpotqa --skip-retrieval
|
||||
evaluations run hotpotqa --skip-qa
|
||||
|
||||
# Limit the number of test cases
|
||||
evaluations run wix --limit 100
|
||||
evaluations run hotpotqa --limit 100
|
||||
```
|
||||
|
||||
### Choosing the target
|
||||
|
||||
`evaluations run` benchmarks `--target rag-skill` by default. Use
|
||||
`--target analysis-skill` to benchmark the analysis skill against the same
|
||||
`evaluations run` benchmarks `--target rag-capability` by default. Use
|
||||
`--target analysis-capability` to benchmark the analysis capability against the same
|
||||
datasets and judge:
|
||||
|
||||
```bash
|
||||
evaluations run wix --target rag-skill
|
||||
evaluations run wix --target analysis-skill --skill-model ollama:gpt-oss
|
||||
evaluations run hotpotqa --target rag-capability
|
||||
evaluations run hotpotqa --target analysis-capability --capability-model ollama:qwen3.8
|
||||
```
|
||||
|
||||
`--skill-model "provider:name"` overrides the skill model independently from
|
||||
`--capability-model "provider:name"` overrides the capability model independently from
|
||||
the judge (defaults to `qa.model`, or `analysis.model` when set for the
|
||||
analysis-skill target). A citation retrieval metric (`cited_map`) is computed
|
||||
alongside QA accuracy from the URIs the skill registered via the `cite` tool.
|
||||
analysis-capability target). A citation retrieval metric (`cited_map`) is computed
|
||||
alongside QA accuracy from the URIs the capability registered via the `cite` tool.
|
||||
|
||||
### Debugging runs in Logfire
|
||||
|
||||
|
|
@ -67,15 +69,15 @@ cases) for use from Claude Code.
|
|||
Download pre-built evaluation databases from HuggingFace:
|
||||
|
||||
```bash
|
||||
evaluations download wix
|
||||
evaluations download hotpotqa
|
||||
evaluations download all
|
||||
evaluations download wix --force
|
||||
evaluations download hotpotqa --force
|
||||
```
|
||||
|
||||
Upload databases (maintainer only):
|
||||
|
||||
```bash
|
||||
evaluations upload wix
|
||||
evaluations upload hotpotqa
|
||||
evaluations upload all
|
||||
```
|
||||
|
||||
|
|
@ -87,3 +89,16 @@ By default, evaluation databases are stored in the haiku.rag data directory:
|
|||
- **Windows**: `C:/Users/<USER>/AppData/Roaming/haiku.rag/evaluations/dbs/`
|
||||
|
||||
You can override this with the `--db` option.
|
||||
|
||||
### Evaluating over Multiple Databases
|
||||
|
||||
With [`lancedb.databases`](https://ggozad.github.io/haiku.rag/configuration/storage/#multiple-databases) configured, `evaluations run <dataset> --skip-db` benchmarks the full set. Retrieval, QA, and live conversations preserve the database name on results and citations. A configured set of one follows the same path and retains its name.
|
||||
|
||||
Population writes one database and therefore requires `--db`:
|
||||
|
||||
```bash
|
||||
evaluations run hotpotqa --db /path/to/one.lancedb # populate, then benchmark
|
||||
evaluations run hotpotqa --skip-db # benchmark the configured set
|
||||
```
|
||||
|
||||
`--db` overrides the configured set for both population and benchmarks.
|
||||
|
|
|
|||
52
evaluations/configs/frames.yaml
Normal file
52
evaluations/configs/frames.yaml
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
# Reference config for the `frames` evaluation database.
|
||||
# FRAMES (google/frames-benchmark): 824 multi-hop questions over a corpus of
|
||||
# the ~2.5k Wikipedia articles linked per question, fetched at current
|
||||
# revision (revid + fetch date recorded in the article cache).
|
||||
# Run: evaluations run frames --config configs/frames.yaml
|
||||
# base_url uses the `vllm` host serving each model over an OpenAI-compatible API.
|
||||
|
||||
environment: development
|
||||
|
||||
storage:
|
||||
auto_vacuum: false
|
||||
|
||||
embeddings:
|
||||
model:
|
||||
provider: openai
|
||||
name: qwen3-embedding-4b
|
||||
vector_dim: 2560
|
||||
base_url: http://vllm:11431/v1
|
||||
|
||||
reranking:
|
||||
model:
|
||||
provider: vllm
|
||||
name: Qwen/Qwen3-Reranker-4B
|
||||
base_url: http://vllm:11455
|
||||
|
||||
analysis:
|
||||
# Bounds per-execution sandbox output so accumulated code returns cannot
|
||||
# outgrow the model's input budget.
|
||||
max_output_chars: 20000
|
||||
|
||||
qa:
|
||||
model:
|
||||
provider: openai
|
||||
name: gemma4-26b
|
||||
base_url: http://vllm:11432/v1
|
||||
# vLLM reserves max_tokens out of max_model_len; a large value starves
|
||||
# the input budget and 400s long agentic contexts.
|
||||
max_tokens: 8192
|
||||
|
||||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://vllm:11439/v1
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
top_p: 0.95
|
||||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
reasoning_effort: low
|
||||
37
evaluations/configs/hotpotqa.yaml
Normal file
37
evaluations/configs/hotpotqa.yaml
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
# Reference config for the `hotpotqa` pre-built evaluation database.
|
||||
# HotpotQA (distractor validation split) multi-hop QA over wiki paragraphs.
|
||||
# Run: evaluations run hotpotqa --config configs/hotpotqa.yaml
|
||||
# base_url uses the `vllm` host serving each model over an OpenAI-compatible API.
|
||||
|
||||
environment: development
|
||||
|
||||
storage:
|
||||
auto_vacuum: false
|
||||
|
||||
embeddings:
|
||||
model:
|
||||
provider: openai
|
||||
name: qwen3-embedding-4b
|
||||
vector_dim: 2560
|
||||
base_url: http://vllm:11431/v1
|
||||
|
||||
qa:
|
||||
model:
|
||||
provider: openai
|
||||
name: gemma4-26b
|
||||
base_url: http://vllm:11432/v1
|
||||
max_tokens: 49152
|
||||
|
||||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://vllm:11439/v1
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
top_p: 0.95
|
||||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
reasoning_effort: low
|
||||
56
evaluations/configs/mtrag_clapnq.yaml
Normal file
56
evaluations/configs/mtrag_clapnq.yaml
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
# Reference config for the `mtrag_clapnq` pre-built evaluation database.
|
||||
# IBM MTRAG, ClapNQ (Wikipedia) domain: multi-turn retrieval and QA over
|
||||
# 183,408 passages. Also serves mtrag_clapnq_rewrite, mtrag_clapnq_live and
|
||||
# mtrag_clapnq_live_uncompacted.
|
||||
# Run: evaluations run mtrag_clapnq --config configs/mtrag_clapnq.yaml
|
||||
# base_url uses the `vllm` host serving each model over an OpenAI-compatible API.
|
||||
# The corpus is text-only: no multimodal embedder, no vision paths. This eval
|
||||
# cannot exercise image or vision turn-boundary behavior.
|
||||
|
||||
environment: development
|
||||
|
||||
storage:
|
||||
auto_vacuum: false
|
||||
|
||||
embeddings:
|
||||
model:
|
||||
provider: openai
|
||||
name: qwen3-embedding-4b
|
||||
vector_dim: 2560
|
||||
base_url: http://vllm:11431/v1
|
||||
|
||||
reranking:
|
||||
model:
|
||||
provider: vllm
|
||||
name: Qwen/Qwen3-Reranker-4B
|
||||
base_url: http://vllm:11455
|
||||
|
||||
qa:
|
||||
model:
|
||||
provider: openai
|
||||
name: RedHatAI/Muse-Glimmer-30B-NVFP4
|
||||
base_url: http://vllm:11450/v1
|
||||
# vLLM enforces input + max_tokens <= max_model_len, so a large output
|
||||
# budget silently shrinks the input budget. MTRAG answers are sentences.
|
||||
max_tokens: 8192
|
||||
extra_body:
|
||||
chat_template_kwargs:
|
||||
# Part of the measured baseline. vLLM's reasoning parser consumes
|
||||
# enable_thinking before the chat template sees it; reasoning_strength
|
||||
# is the knob Muse Glimmer templates honour, and a template that
|
||||
# defaults it to low silently changes search behavior.
|
||||
reasoning_strength: high
|
||||
|
||||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://vllm:11439/v1
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
top_p: 0.95
|
||||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
reasoning_effort: low
|
||||
|
|
@ -25,3 +25,17 @@ qa:
|
|||
name: gemma4-26b
|
||||
base_url: http://vllm:11432/v1
|
||||
vision: true
|
||||
|
||||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://vllm:11439/v1
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
top_p: 0.95
|
||||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
reasoning_effort: low
|
||||
|
|
|
|||
|
|
@ -26,3 +26,17 @@ qa:
|
|||
name: gemma4-26b
|
||||
base_url: http://vllm:11432/v1
|
||||
vision: true
|
||||
|
||||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://vllm:11439/v1
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
top_p: 0.95
|
||||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
reasoning_effort: low
|
||||
|
|
|
|||
|
|
@ -17,8 +17,9 @@ embeddings:
|
|||
|
||||
reranking:
|
||||
model:
|
||||
provider: cross-encoder
|
||||
name: mixedbread-ai/mxbai-rerank-base-v2
|
||||
provider: vllm
|
||||
name: Qwen/Qwen3-Reranker-4B
|
||||
base_url: http://vllm:11455
|
||||
|
||||
qa:
|
||||
model:
|
||||
|
|
@ -26,3 +27,17 @@ qa:
|
|||
name: gemma4-26b
|
||||
base_url: http://vllm:11432/v1
|
||||
vision: true
|
||||
|
||||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://vllm:11439/v1
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
top_p: 0.95
|
||||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
reasoning_effort: low
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
# Reference config for the `t2_finqa` pre-built evaluation database.
|
||||
# T²-RAGBench (FinQA) financial QA, scored by exact numeric match.
|
||||
# Run: evaluations run t2_finqa --skip-db --target analysis-skill --config configs/t2_finqa.yaml
|
||||
# No `evaluations.judge` block: the spec sets `NumberMatchEvaluator`, which
|
||||
# replaces the evaluator list, so no LLM judge is constructed for this dataset.
|
||||
# Run: evaluations run t2_finqa --skip-db --target analysis-capability --config configs/t2_finqa.yaml
|
||||
# base_url uses the `vllm` host serving each model over an OpenAI-compatible API.
|
||||
|
||||
environment: development
|
||||
|
|
@ -22,8 +24,9 @@ processing:
|
|||
|
||||
reranking:
|
||||
model:
|
||||
provider: cross-encoder
|
||||
name: mixedbread-ai/mxbai-rerank-base-v2
|
||||
provider: vllm
|
||||
name: Qwen/Qwen3-Reranker-4B
|
||||
base_url: http://vllm:11455
|
||||
|
||||
qa:
|
||||
model:
|
||||
|
|
|
|||
|
|
@ -1,27 +0,0 @@
|
|||
# Reference config for the `wix` pre-built evaluation database.
|
||||
# Run: evaluations run wix --skip-db --config configs/wix.yaml
|
||||
# base_url uses the `vllm` host serving each model over an OpenAI-compatible API.
|
||||
|
||||
environment: development
|
||||
|
||||
storage:
|
||||
auto_vacuum: false
|
||||
|
||||
embeddings:
|
||||
model:
|
||||
provider: openai
|
||||
name: qwen3-embedding-4b
|
||||
vector_dim: 2560
|
||||
base_url: http://vllm:11431/v1
|
||||
|
||||
reranking:
|
||||
model:
|
||||
provider: cross-encoder
|
||||
name: mixedbread-ai/mxbai-rerank-base-v2
|
||||
|
||||
qa:
|
||||
model:
|
||||
provider: openai
|
||||
name: gemma4-26b
|
||||
base_url: http://vllm:11432/v1
|
||||
vision: true
|
||||
106
evaluations/evaluations/artifacts.py
Normal file
106
evaluations/evaluations/artifacts.py
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
"""Pre-built evaluation databases on HuggingFace."""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from huggingface_hub import HfApi, snapshot_download
|
||||
from rich.console import Console
|
||||
|
||||
from evaluations.config import DatasetSpec
|
||||
|
||||
console = Console()
|
||||
|
||||
HF_REPO_ID = "ggozad/haiku-rag-eval-dbs"
|
||||
|
||||
|
||||
def download_dataset_db(spec: DatasetSpec, force: bool = False) -> None:
|
||||
"""Fetch one dataset's database from HuggingFace into its local path."""
|
||||
db = spec.db_path()
|
||||
if db.exists() and not force:
|
||||
console.print(
|
||||
f"[yellow]Skipping {spec.key}: database already exists at {db}[/yellow]"
|
||||
)
|
||||
console.print("Use --force to overwrite.")
|
||||
return
|
||||
|
||||
console.print(f"[blue]Downloading {spec.key}...[/blue]")
|
||||
|
||||
try:
|
||||
downloaded_path = snapshot_download(
|
||||
repo_id=HF_REPO_ID,
|
||||
repo_type="dataset",
|
||||
allow_patterns=f"{spec.db_filename}/*",
|
||||
)
|
||||
except Exception as e:
|
||||
console.print(f"[red]Failed to download {spec.key}: {e}[/red]")
|
||||
return
|
||||
|
||||
source_path = Path(downloaded_path) / spec.db_filename
|
||||
if not source_path.exists():
|
||||
console.print(f"[red]Database {spec.key} not found in HuggingFace repo.[/red]")
|
||||
console.print(
|
||||
f"[yellow]The database may not have been uploaded yet. "
|
||||
f"Try running 'evaluations build {spec.key}' to create it locally.[/yellow]"
|
||||
)
|
||||
return
|
||||
|
||||
if db.exists():
|
||||
shutil.rmtree(db)
|
||||
|
||||
db.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copytree(source_path, db)
|
||||
|
||||
console.print(f"[green]Downloaded {spec.key} to {db}[/green]")
|
||||
|
||||
|
||||
def upload_dataset_db(spec: DatasetSpec) -> None:
|
||||
"""Push one dataset's database to HuggingFace (maintainer only).
|
||||
|
||||
Uses ``upload_large_folder`` for resumable, parallel transfer — important
|
||||
for the multi-GB ORB databases which would otherwise abort on any transient
|
||||
network failure under plain ``upload_folder``.
|
||||
|
||||
``upload_large_folder`` has no ``path_in_repo`` — it ships the contents of
|
||||
``folder_path`` to the repo root. Stage the db under a temp parent with
|
||||
hardlinks so the basename becomes the remote path, leaving everything else
|
||||
at the root undisturbed.
|
||||
"""
|
||||
db = spec.db_path()
|
||||
if not db.exists():
|
||||
console.print(f"[red]Database not found at {db}[/red]")
|
||||
return
|
||||
|
||||
api = HfApi()
|
||||
|
||||
# Wipe the existing remote path so we don't accumulate orphaned files from
|
||||
# prior uploads. upload_large_folder doesn't accept delete_patterns, so we
|
||||
# do this as a separate commit. Safe to run if the path is missing.
|
||||
try:
|
||||
api.delete_folder(
|
||||
path_in_repo=spec.db_filename,
|
||||
repo_id=HF_REPO_ID,
|
||||
repo_type="dataset",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
with tempfile.TemporaryDirectory() as staging:
|
||||
target = Path(staging) / spec.db_filename
|
||||
target.mkdir()
|
||||
for src in db.rglob("*"):
|
||||
if not src.is_file():
|
||||
continue
|
||||
dest = target / src.relative_to(db)
|
||||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||||
os.link(src, dest)
|
||||
|
||||
console.print(f"[blue]Uploading {spec.key} ({db})...[/blue]")
|
||||
api.upload_large_folder(
|
||||
folder_path=staging,
|
||||
repo_id=HF_REPO_ID,
|
||||
repo_type="dataset",
|
||||
)
|
||||
|
||||
console.print(f"[green]Uploaded {spec.key} to {HF_REPO_ID}[/green]")
|
||||
|
|
@ -1,45 +1,27 @@
|
|||
import asyncio
|
||||
import shutil
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, cast
|
||||
from typing import cast
|
||||
|
||||
import typer
|
||||
from dotenv import find_dotenv, load_dotenv
|
||||
from huggingface_hub import HfApi, snapshot_download
|
||||
from pydantic_evals import Case, Dataset as EvalDataset, set_eval_attribute
|
||||
from pydantic_evals.evaluators import Evaluator, LLMJudge
|
||||
from pydantic_evals.reporting import ReportCaseFailure
|
||||
from rich.console import Console
|
||||
from rich.progress import Progress
|
||||
|
||||
from evaluations.artifacts import download_dataset_db, upload_dataset_db
|
||||
from evaluations.config import DatasetSpec
|
||||
from evaluations.population import populate_db
|
||||
from evaluations.qa import TARGETS, Target, run_live_qa_benchmark, run_qa_benchmark
|
||||
from evaluations.retrieval import run_retrieval_benchmark
|
||||
from evaluations.datasets import DATASETS
|
||||
from evaluations.evaluators import (
|
||||
ANSWER_EQUIVALENCE_RUBRIC,
|
||||
CitationMAPEvaluator,
|
||||
MAPEvaluator,
|
||||
)
|
||||
from evaluations.skill_runner import SkillFactory, run_skill_question
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import AppConfig, find_config_file, load_yaml_config
|
||||
from haiku.rag.config.models import ModelConfig
|
||||
from haiku.rag.logging import configure_cli_logging
|
||||
from haiku.rag.telemetry import configure as configure_telemetry
|
||||
from haiku.rag.utils import get_model, parse_model_option
|
||||
from haiku.rag.utils import parse_model_option
|
||||
|
||||
Target = Literal["rag-skill", "analysis-skill"]
|
||||
TARGETS: tuple[Target, ...] = ("rag-skill", "analysis-skill")
|
||||
|
||||
# Pinned judge model. Decoupled from `config.qa.model` so a user changing
|
||||
# their QA model does not inadvertently change the judge — keeps cross-run
|
||||
# comparisons stable. Override per-run with `--judge-model provider:name`.
|
||||
DEFAULT_JUDGE_MODEL = ModelConfig(provider="ollama", name="qwen3.6")
|
||||
|
||||
load_dotenv(find_dotenv(usecwd=True))
|
||||
|
||||
HF_REPO_ID = "ggozad/haiku-rag-eval-dbs"
|
||||
|
||||
# Scrubbing off: eval outputs are financial answers with words like "authorized"
|
||||
# that trip Logfire's secret scrubber and redact the model's answer text.
|
||||
configure_telemetry(service_name="evals", scrubbing=False)
|
||||
|
|
@ -47,457 +29,6 @@ configure_cli_logging()
|
|||
console = Console()
|
||||
|
||||
|
||||
def build_experiment_metadata(
|
||||
dataset_key: str,
|
||||
test_cases: int,
|
||||
config: AppConfig,
|
||||
judge_config: ModelConfig | None = None,
|
||||
target: Target = "rag-skill",
|
||||
skill_config: ModelConfig | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build experiment metadata for Logfire tracking."""
|
||||
metadata: dict[str, Any] = {
|
||||
"dataset": dataset_key,
|
||||
"test_cases": test_cases,
|
||||
"target": target,
|
||||
"embedder_provider": config.embeddings.model.provider,
|
||||
"embedder_model": config.embeddings.model.name,
|
||||
"embedder_dim": config.embeddings.model.vector_dim,
|
||||
"chunk_size": config.processing.chunk_size,
|
||||
"search_limit": config.search.limit,
|
||||
"max_context_chars": config.search.max_context_chars,
|
||||
"rerank_provider": config.reranking.model.provider
|
||||
if config.reranking.model
|
||||
else None,
|
||||
"rerank_model": config.reranking.model.name if config.reranking.model else None,
|
||||
"qa_provider": config.qa.model.provider,
|
||||
"qa_model": config.qa.model.name,
|
||||
"qa_temperature": config.qa.model.temperature,
|
||||
"qa_max_tokens": config.qa.model.max_tokens,
|
||||
"qa_enable_thinking": config.qa.model.enable_thinking,
|
||||
"qa_max_searches": config.qa.max_searches,
|
||||
}
|
||||
if judge_config is not None:
|
||||
metadata.update(
|
||||
{
|
||||
"judge_provider": judge_config.provider,
|
||||
"judge_model": judge_config.name,
|
||||
"judge_temperature": judge_config.temperature,
|
||||
"judge_max_tokens": judge_config.max_tokens,
|
||||
"judge_enable_thinking": judge_config.enable_thinking,
|
||||
}
|
||||
)
|
||||
if skill_config is not None:
|
||||
metadata.update(
|
||||
{
|
||||
"skill_provider": skill_config.provider,
|
||||
"skill_model": skill_config.name,
|
||||
"skill_temperature": skill_config.temperature,
|
||||
"skill_max_tokens": skill_config.max_tokens,
|
||||
"skill_enable_thinking": skill_config.enable_thinking,
|
||||
}
|
||||
)
|
||||
return metadata
|
||||
|
||||
|
||||
async def populate_db(
|
||||
spec: DatasetSpec,
|
||||
config: AppConfig,
|
||||
db_path: Path | None = None,
|
||||
vacuum_interval: int = 100,
|
||||
) -> None:
|
||||
db = spec.db_path(db_path)
|
||||
db.parent.mkdir(parents=True, exist_ok=True)
|
||||
corpus = spec.document_loader()
|
||||
if spec.document_limit is not None:
|
||||
corpus = corpus.select(range(min(spec.document_limit, len(corpus))))
|
||||
|
||||
# Disable auto_vacuum - we'll vacuum periodically instead to prevent disk exhaustion
|
||||
config.storage.auto_vacuum = False
|
||||
|
||||
with Progress() as progress:
|
||||
task = progress.add_task("[green]Populating database...", total=len(corpus))
|
||||
async with HaikuRAG(db, config=config, create=True) as rag:
|
||||
docs_since_vacuum = 0
|
||||
for doc in corpus:
|
||||
doc_mapping = cast(Mapping[str, Any], doc)
|
||||
payload = spec.document_mapper(doc_mapping)
|
||||
if payload is None:
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
# `payload.uri` is the canonical document identifier and is now
|
||||
# honored by both `create_document` and (via the `uri=` override)
|
||||
# `create_document_from_source`, so it's also the right key to
|
||||
# look up an existing document, regardless of whether the source
|
||||
# is a file path or inline content.
|
||||
existing = await rag.get_document_by_uri(payload.uri)
|
||||
if existing is not None:
|
||||
assert existing.id
|
||||
chunks = await rag.chunk_repository.get_by_document_id(existing.id)
|
||||
if chunks:
|
||||
progress.advance(task)
|
||||
continue
|
||||
await rag.document_repository.delete(existing.id)
|
||||
|
||||
if payload.source_path is not None:
|
||||
await rag.create_document_from_source(
|
||||
source=payload.source_path,
|
||||
title=payload.title,
|
||||
metadata=payload.metadata,
|
||||
uri=payload.uri,
|
||||
)
|
||||
else:
|
||||
assert payload.content is not None
|
||||
await rag.create_document(
|
||||
content=payload.content,
|
||||
uri=payload.uri,
|
||||
title=payload.title,
|
||||
metadata=payload.metadata,
|
||||
format=payload.format,
|
||||
)
|
||||
docs_since_vacuum += 1
|
||||
progress.advance(task)
|
||||
|
||||
# Periodic vacuum to prevent disk exhaustion
|
||||
if docs_since_vacuum >= vacuum_interval:
|
||||
await rag.store.vacuum(retention_seconds=0)
|
||||
docs_since_vacuum = 0
|
||||
|
||||
# Final vacuum
|
||||
await rag.store.vacuum(retention_seconds=0)
|
||||
|
||||
|
||||
async def run_retrieval_benchmark(
|
||||
spec: DatasetSpec,
|
||||
config: AppConfig,
|
||||
limit: int | None = None,
|
||||
name: str | None = None,
|
||||
db_path: Path | None = None,
|
||||
multimodal_only: bool = False,
|
||||
) -> dict[str, float] | None:
|
||||
if spec.retrieval_loader is None or spec.retrieval_mapper is None:
|
||||
console.print("Skipping retrieval benchmark; no retrieval config.")
|
||||
return None
|
||||
|
||||
corpus = spec.retrieval_loader()
|
||||
if limit is not None:
|
||||
corpus = corpus.select(range(min(limit, len(corpus))))
|
||||
|
||||
cases = []
|
||||
with Progress() as progress:
|
||||
task = progress.add_task("[blue]Building retrieval cases...", total=len(corpus))
|
||||
for doc in corpus:
|
||||
doc_mapping = cast(Mapping[str, Any], doc)
|
||||
sample = spec.retrieval_mapper(doc_mapping)
|
||||
if sample is None or sample.skip:
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
# Filter for multimodal queries if requested
|
||||
if multimodal_only:
|
||||
if sample.source_type is None or "image" not in sample.source_type:
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
case = Case(
|
||||
inputs=sample.question,
|
||||
metadata={
|
||||
"relevant_uris": sample.expected_uris,
|
||||
"source_type": sample.source_type,
|
||||
},
|
||||
)
|
||||
cases.append(case)
|
||||
progress.advance(task)
|
||||
|
||||
if not cases:
|
||||
console.print("No retrieval cases to evaluate.")
|
||||
return None
|
||||
|
||||
if spec.retrieval_evaluator is None:
|
||||
raise ValueError(f"No retrieval evaluator configured for dataset: {spec.key}")
|
||||
|
||||
evaluator = spec.retrieval_evaluator
|
||||
metric_name = evaluator.__class__.__name__.replace("Evaluator", "").upper()
|
||||
|
||||
dataset = EvalDataset(
|
||||
cases=cases,
|
||||
evaluators=[evaluator],
|
||||
)
|
||||
|
||||
db = spec.db_path(db_path)
|
||||
async with HaikuRAG(db, config=config) as rag:
|
||||
|
||||
async def retrieval_target(question: str) -> list[str]:
|
||||
chunks = await rag.search(query=question, limit=5)
|
||||
|
||||
seen = set()
|
||||
identifiers = []
|
||||
for result in chunks:
|
||||
if result.document_id is None:
|
||||
continue
|
||||
doc = await rag.get_document_by_id(result.document_id)
|
||||
if doc is None:
|
||||
continue
|
||||
# Use arxiv_id from metadata if present, otherwise use URI
|
||||
doc_id = doc.metadata.get("arxiv_id") if doc.metadata else None
|
||||
if doc_id is None:
|
||||
doc_id = doc.uri
|
||||
if doc_id and doc_id not in seen:
|
||||
identifiers.append(doc_id)
|
||||
seen.add(doc_id)
|
||||
|
||||
return identifiers
|
||||
|
||||
eval_name = name if name is not None else f"{spec.key}_retrieval_evaluation"
|
||||
|
||||
experiment_metadata = build_experiment_metadata(
|
||||
dataset_key=spec.key,
|
||||
test_cases=len(cases),
|
||||
config=config,
|
||||
)
|
||||
|
||||
report = await dataset.evaluate(
|
||||
retrieval_target,
|
||||
name=eval_name,
|
||||
max_concurrency=1,
|
||||
progress=True,
|
||||
metadata=experiment_metadata,
|
||||
)
|
||||
|
||||
total_score = 0.0
|
||||
total_cases = 0
|
||||
for case in report.cases:
|
||||
if case.scores:
|
||||
for score_result in case.scores.values():
|
||||
total_score += score_result.value
|
||||
total_cases += 1
|
||||
|
||||
mean_score = total_score / total_cases if total_cases > 0 else 0.0
|
||||
|
||||
console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan")
|
||||
console.print(f"Dataset: {spec.key}")
|
||||
console.print(f"Total queries: {len(cases)}")
|
||||
console.print(f"{metric_name}: {mean_score:.4f}")
|
||||
|
||||
return {
|
||||
metric_name.lower(): mean_score,
|
||||
"queries": len(cases),
|
||||
}
|
||||
|
||||
|
||||
def _skill_factory_for_target(target: Target) -> SkillFactory:
|
||||
if target == "rag-skill":
|
||||
from haiku.rag.skills.rag import create_skill
|
||||
|
||||
return create_skill
|
||||
if target == "analysis-skill":
|
||||
from haiku.rag.skills.analysis import create_skill
|
||||
|
||||
return create_skill
|
||||
raise ValueError(f"target {target!r} is not a skill target")
|
||||
|
||||
|
||||
def _citation_evaluator_for(retrieval_evaluator: Evaluator | None) -> Evaluator | None:
|
||||
"""Return the citation-scoring twin of the dataset's retrieval evaluator."""
|
||||
if isinstance(retrieval_evaluator, MAPEvaluator):
|
||||
return CitationMAPEvaluator()
|
||||
return None
|
||||
|
||||
|
||||
def _attach_relevant_uris(
|
||||
cases: list[Case[str, str, dict[str, Any]]],
|
||||
spec: DatasetSpec,
|
||||
limit: int | None,
|
||||
) -> None:
|
||||
"""Augment QA cases with `relevant_uris` joined from retrieval samples.
|
||||
|
||||
Mutates each case's metadata in place. Cases with no matching retrieval
|
||||
sample (by question) are left untouched.
|
||||
"""
|
||||
if spec.retrieval_loader is None or spec.retrieval_mapper is None:
|
||||
return
|
||||
corpus = spec.retrieval_loader()
|
||||
if limit is not None:
|
||||
corpus = corpus.select(range(min(limit, len(corpus))))
|
||||
expected_by_question: dict[str, tuple[str, ...]] = {}
|
||||
for raw in corpus:
|
||||
sample = spec.retrieval_mapper(cast(Mapping[str, Any], raw))
|
||||
if sample is None or sample.skip:
|
||||
continue
|
||||
expected_by_question[sample.question] = sample.expected_uris
|
||||
for case in cases:
|
||||
uris = expected_by_question.get(case.inputs)
|
||||
if uris is None:
|
||||
continue
|
||||
metadata = case.metadata if case.metadata is not None else {}
|
||||
metadata["relevant_uris"] = list(uris)
|
||||
case.metadata = metadata
|
||||
|
||||
|
||||
def _filter_qa_corpus(corpus, case_ids: set[str] | None):
|
||||
"""Keep only rows whose ``id`` is in ``case_ids`` (failure-subset reruns).
|
||||
|
||||
Returns the corpus unchanged when ``case_ids`` is None.
|
||||
"""
|
||||
if case_ids is None:
|
||||
return corpus
|
||||
return corpus.filter(lambda row: row.get("id") in case_ids)
|
||||
|
||||
|
||||
async def run_qa_benchmark(
|
||||
spec: DatasetSpec,
|
||||
config: AppConfig,
|
||||
limit: int | None = None,
|
||||
name: str | None = None,
|
||||
db_path: Path | None = None,
|
||||
judge_model: ModelConfig | None = None,
|
||||
target: Target = "rag-skill",
|
||||
skill_model: ModelConfig | None = None,
|
||||
case_ids: set[str] | None = None,
|
||||
) -> ReportCaseFailure[str, str, dict[str, str]] | None:
|
||||
corpus = spec.qa_loader()
|
||||
corpus = _filter_qa_corpus(corpus, case_ids)
|
||||
if limit is not None:
|
||||
corpus = corpus.select(range(min(limit, len(corpus))))
|
||||
|
||||
cases = [
|
||||
spec.qa_case_builder(index, cast(Mapping[str, Any], doc))
|
||||
for index, doc in enumerate(corpus, start=1)
|
||||
]
|
||||
|
||||
judge_config = judge_model or DEFAULT_JUDGE_MODEL
|
||||
if target == "analysis-skill":
|
||||
# Mirror the skill-code resolver: explicit analysis.model wins,
|
||||
# else fall back to qa.model.
|
||||
skill_config = skill_model or config.analysis.model or config.qa.model
|
||||
else:
|
||||
skill_config = skill_model or config.qa.model
|
||||
db = spec.db_path(db_path)
|
||||
|
||||
_attach_relevant_uris(cases, spec, limit)
|
||||
citation_evaluator = _citation_evaluator_for(spec.retrieval_evaluator)
|
||||
|
||||
qa_evaluator = spec.qa_evaluator
|
||||
evaluators: list[Evaluator]
|
||||
if qa_evaluator is not None:
|
||||
evaluators = [qa_evaluator]
|
||||
else:
|
||||
evaluators = [
|
||||
LLMJudge(
|
||||
rubric=ANSWER_EQUIVALENCE_RUBRIC,
|
||||
include_input=True,
|
||||
include_expected_output=True,
|
||||
model=get_model(judge_config, config),
|
||||
assertion={
|
||||
"evaluation_name": "answer_equivalent",
|
||||
"include_reason": True,
|
||||
},
|
||||
),
|
||||
]
|
||||
if citation_evaluator is not None:
|
||||
evaluators.append(citation_evaluator)
|
||||
|
||||
evaluation_dataset = EvalDataset[str, str, dict[str, str]](
|
||||
name=spec.key, cases=cases, evaluators=evaluators
|
||||
)
|
||||
|
||||
eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
|
||||
experiment_metadata = build_experiment_metadata(
|
||||
dataset_key=spec.key,
|
||||
test_cases=len(cases),
|
||||
config=config,
|
||||
judge_config=judge_config,
|
||||
target=target,
|
||||
skill_config=skill_config,
|
||||
)
|
||||
|
||||
async def _evaluate(answer_fn: Callable[[str], Awaitable[str]]):
|
||||
return await evaluation_dataset.evaluate(
|
||||
answer_fn,
|
||||
name=eval_name,
|
||||
max_concurrency=1,
|
||||
progress=True,
|
||||
metadata=experiment_metadata,
|
||||
)
|
||||
|
||||
skill_factory = _skill_factory_for_target(target)
|
||||
resolved_skill_model = get_model(skill_config, config)
|
||||
|
||||
async def answer_question(question: str) -> str:
|
||||
result = await run_skill_question(
|
||||
skill_factory=skill_factory,
|
||||
db_path=db,
|
||||
config=config,
|
||||
question=question,
|
||||
skill_model=resolved_skill_model,
|
||||
)
|
||||
set_eval_attribute("cited_uris", result.cited_uris)
|
||||
return result.answer
|
||||
|
||||
report = await _evaluate(answer_question)
|
||||
|
||||
total_processed = len(report.cases)
|
||||
failures = report.failures
|
||||
if qa_evaluator is not None:
|
||||
score_key = qa_evaluator.get_default_evaluation_name()
|
||||
passing_cases = sum(
|
||||
1
|
||||
for case in report.cases
|
||||
if score_key in case.scores and case.scores[score_key].value >= 1.0
|
||||
)
|
||||
scoring = score_key
|
||||
else:
|
||||
passing_cases = sum(
|
||||
1
|
||||
for case in report.cases
|
||||
if case.assertions.get("answer_equivalent")
|
||||
and case.assertions["answer_equivalent"].value
|
||||
)
|
||||
scoring = "answer_equivalent"
|
||||
accuracy = passing_cases / total_processed if total_processed > 0 else 0
|
||||
|
||||
console.print("\n=== QA Benchmark Results ===", style="bold cyan")
|
||||
console.print(f"Scoring: {scoring}")
|
||||
console.print(f"Total questions: {total_processed}")
|
||||
console.print(f"Correct answers: {passing_cases}")
|
||||
console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)")
|
||||
|
||||
if citation_evaluator is not None:
|
||||
score_key = citation_evaluator.get_default_evaluation_name()
|
||||
scores = [
|
||||
case.scores[score_key].value
|
||||
for case in report.cases
|
||||
if score_key in case.scores
|
||||
]
|
||||
if scores:
|
||||
cited_count = sum(
|
||||
1 for case in report.cases if case.attributes.get("cited_uris")
|
||||
)
|
||||
mean_citations = sum(
|
||||
len(case.attributes.get("cited_uris") or []) for case in report.cases
|
||||
) / len(report.cases)
|
||||
mean_score = sum(scores) / len(scores)
|
||||
console.print(
|
||||
f"\n=== Citation Retrieval ({score_key}) ===", style="bold cyan"
|
||||
)
|
||||
console.print(f"Mean {score_key}: {mean_score:.4f}")
|
||||
console.print(
|
||||
f"Cite rate (≥1 citation): {cited_count / len(report.cases):.2%}"
|
||||
)
|
||||
console.print(f"Mean citations per case: {mean_citations:.2f}")
|
||||
|
||||
if failures:
|
||||
console.print("[red]\nSummary of failures:[/red]")
|
||||
for failure in failures:
|
||||
console.print(f"Case: {failure.name}")
|
||||
console.print(f"Question: {failure.inputs}")
|
||||
console.print(f"Error: {failure.error_message}")
|
||||
console.print("")
|
||||
|
||||
return failures[0] if failures else None
|
||||
|
||||
|
||||
async def evaluate_dataset(
|
||||
spec: DatasetSpec,
|
||||
config: AppConfig,
|
||||
|
|
@ -510,11 +41,29 @@ async def evaluate_dataset(
|
|||
vacuum_interval: int = 100,
|
||||
multimodal_only: bool = False,
|
||||
judge_model: ModelConfig | None = None,
|
||||
target: Target = "rag-skill",
|
||||
skill_model: ModelConfig | None = None,
|
||||
target: Target = "rag-capability",
|
||||
capability_model: ModelConfig | None = None,
|
||||
case_ids: set[str] | None = None,
|
||||
document_filter: str | None = None,
|
||||
) -> None:
|
||||
if document_filter is not None:
|
||||
console.print(f"Document filter: {document_filter}", style="dim")
|
||||
|
||||
if db_path is not None and config.lancedb.databases:
|
||||
raise ValueError(
|
||||
"--db PATH places the database where the configuration places none, "
|
||||
f"and this configuration names {', '.join(config.lancedb.databases)} "
|
||||
"in lancedb.databases. Drop --db to evaluate the configured set."
|
||||
)
|
||||
|
||||
if not skip_db:
|
||||
if spec.uses_configured_databases(config, db_path):
|
||||
raise ValueError(
|
||||
"lancedb.databases places the databases this run reads, and "
|
||||
"population writes to one, so it would ingest into a database "
|
||||
"the run does not read. Pass --skip-db to evaluate the "
|
||||
"configured set, or --db PATH to populate and evaluate one."
|
||||
)
|
||||
console.print(f"Using dataset: {spec.key}", style="bold magenta")
|
||||
await populate_db(
|
||||
spec, config, db_path=db_path, vacuum_interval=vacuum_interval
|
||||
|
|
@ -529,13 +78,15 @@ async def evaluate_dataset(
|
|||
name=name,
|
||||
db_path=db_path,
|
||||
multimodal_only=multimodal_only,
|
||||
document_filter=document_filter,
|
||||
)
|
||||
|
||||
if not skip_qa:
|
||||
console.print(
|
||||
f"\nRunning QA benchmarks (target={target})...", style="bold yellow"
|
||||
)
|
||||
await run_qa_benchmark(
|
||||
qa_benchmark = run_live_qa_benchmark if spec.live else run_qa_benchmark
|
||||
await qa_benchmark(
|
||||
spec,
|
||||
config,
|
||||
limit=limit,
|
||||
|
|
@ -543,8 +94,9 @@ async def evaluate_dataset(
|
|||
db_path=db_path,
|
||||
judge_model=judge_model,
|
||||
target=target,
|
||||
skill_model=skill_model,
|
||||
capability_model=capability_model,
|
||||
case_ids=case_ids,
|
||||
document_filter=document_filter,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -589,9 +141,20 @@ def _resolve_dataset(dataset: str) -> DatasetSpec:
|
|||
|
||||
|
||||
def _resolve_datasets(dataset: str) -> list[DatasetSpec]:
|
||||
"""Resolve 'all' or a single dataset key to a list of DatasetSpecs."""
|
||||
"""Resolve 'all' or a single dataset key to a list of DatasetSpecs.
|
||||
|
||||
'all' yields one spec per database: query variants sharing a db_filename
|
||||
would otherwise be downloaded/uploaded twice.
|
||||
"""
|
||||
if dataset.lower() == "all":
|
||||
return list(DATASETS.values())
|
||||
seen: set[str] = set()
|
||||
specs: list[DatasetSpec] = []
|
||||
for spec in DATASETS.values():
|
||||
if spec.db_filename in seen:
|
||||
continue
|
||||
seen.add(spec.db_filename)
|
||||
specs.append(spec)
|
||||
return specs
|
||||
return [_resolve_dataset(dataset)]
|
||||
|
||||
|
||||
|
|
@ -601,7 +164,11 @@ def run(
|
|||
config: Path | None = typer.Option(
|
||||
None, "--config", help="Path to haiku.rag YAML config file."
|
||||
),
|
||||
db: Path | None = typer.Option(None, "--db", help="Override the database path."),
|
||||
db: Path | None = typer.Option(
|
||||
None,
|
||||
"--db",
|
||||
help="Database path, where the configuration places no database.",
|
||||
),
|
||||
skip_db: bool = typer.Option(
|
||||
False, "--skip-db", help="Skip updating the evaluation db."
|
||||
),
|
||||
|
|
@ -622,16 +189,27 @@ def run(
|
|||
help="Only evaluate queries requiring image understanding.",
|
||||
),
|
||||
target: str = typer.Option(
|
||||
"rag-skill",
|
||||
"rag-capability",
|
||||
"--target",
|
||||
help="What to benchmark: rag-skill | analysis-skill.",
|
||||
help="What to benchmark: rag-capability | analysis-capability.",
|
||||
),
|
||||
skill_model: str | None = typer.Option(
|
||||
capability_model: str | None = typer.Option(
|
||||
None,
|
||||
"--skill-model",
|
||||
"--capability-model",
|
||||
help=(
|
||||
"Skill model as 'provider:name'. Defaults to qa.model (or "
|
||||
"analysis.model when --target is analysis-skill) from the config."
|
||||
"Capability model as 'provider:name'. Defaults to qa.model (or "
|
||||
"analysis.model when --target is analysis-capability) from the config."
|
||||
),
|
||||
),
|
||||
document_filter: str | None = typer.Option(
|
||||
None,
|
||||
"--filter",
|
||||
"-f",
|
||||
help=(
|
||||
"SQL WHERE clause over document columns (id, uri, title, "
|
||||
"created_at, updated_at, metadata) restricting every benchmark "
|
||||
"search, e.g. \"uri LIKE '%arxiv%'\". metadata is stored as a "
|
||||
"string, so match it with LIKE."
|
||||
),
|
||||
),
|
||||
filter_ids: Path | None = typer.Option(
|
||||
|
|
@ -651,7 +229,9 @@ def run(
|
|||
)
|
||||
target_value = cast(Target, target)
|
||||
judge_model_config = app_config.evaluations.judge
|
||||
skill_model_config = parse_model_option(skill_model) if skill_model else None
|
||||
capability_model_config = (
|
||||
parse_model_option(capability_model) if capability_model else None
|
||||
)
|
||||
|
||||
asyncio.run(
|
||||
evaluate_dataset(
|
||||
|
|
@ -667,8 +247,9 @@ def run(
|
|||
multimodal_only=multimodal_only,
|
||||
judge_model=judge_model_config,
|
||||
target=target_value,
|
||||
skill_model=skill_model_config,
|
||||
capability_model=capability_model_config,
|
||||
case_ids=_load_case_ids(filter_ids),
|
||||
document_filter=document_filter,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -679,111 +260,17 @@ def download(
|
|||
force: bool = typer.Option(False, "--force", help="Overwrite existing database."),
|
||||
) -> None:
|
||||
"""Download pre-built evaluation database from HuggingFace."""
|
||||
specs = _resolve_datasets(dataset)
|
||||
|
||||
for spec in specs:
|
||||
db = spec.db_path()
|
||||
if db.exists() and not force:
|
||||
console.print(
|
||||
f"[yellow]Skipping {spec.key}: database already exists at {db}[/yellow]"
|
||||
)
|
||||
console.print("Use --force to overwrite.")
|
||||
continue
|
||||
|
||||
console.print(f"[blue]Downloading {spec.key}...[/blue]")
|
||||
|
||||
try:
|
||||
downloaded_path = snapshot_download(
|
||||
repo_id=HF_REPO_ID,
|
||||
repo_type="dataset",
|
||||
allow_patterns=f"{spec.db_filename}/*",
|
||||
)
|
||||
except Exception as e:
|
||||
console.print(f"[red]Failed to download {spec.key}: {e}[/red]")
|
||||
continue
|
||||
|
||||
# Check if the expected database exists in the downloaded snapshot
|
||||
source_path = Path(downloaded_path) / spec.db_filename
|
||||
if not source_path.exists():
|
||||
console.print(
|
||||
f"[red]Database {spec.key} not found in HuggingFace repo.[/red]"
|
||||
)
|
||||
console.print(
|
||||
f"[yellow]The database may not have been uploaded yet. "
|
||||
f"Try running 'evaluations build {spec.key}' to create it locally.[/yellow]"
|
||||
)
|
||||
continue
|
||||
|
||||
# Remove existing database if force is set
|
||||
if db.exists():
|
||||
shutil.rmtree(db)
|
||||
|
||||
# Copy from cache to target location
|
||||
db.parent.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copytree(source_path, db)
|
||||
|
||||
console.print(f"[green]Downloaded {spec.key} to {db}[/green]")
|
||||
for spec in _resolve_datasets(dataset):
|
||||
download_dataset_db(spec, force=force)
|
||||
|
||||
|
||||
@app.command()
|
||||
def upload(
|
||||
dataset: str = typer.Argument(..., help="Dataset key or 'all' to upload all."),
|
||||
) -> None:
|
||||
"""Upload evaluation database to HuggingFace (maintainer only).
|
||||
|
||||
Uses ``upload_large_folder`` for resumable, parallel transfer — important
|
||||
for the multi-GB ORB databases which would otherwise abort on any transient
|
||||
network failure under plain ``upload_folder``.
|
||||
|
||||
``upload_large_folder`` has no ``path_in_repo`` — it ships the contents of
|
||||
``folder_path`` to the repo root. Stage the db under a temp parent with
|
||||
hardlinks so the basename becomes the remote path, leaving everything
|
||||
else at the root undisturbed.
|
||||
"""
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
specs = _resolve_datasets(dataset)
|
||||
|
||||
api = HfApi()
|
||||
|
||||
for spec in specs:
|
||||
db = spec.db_path()
|
||||
if not db.exists():
|
||||
console.print(f"[red]Database not found at {db}[/red]")
|
||||
continue
|
||||
|
||||
# Wipe the existing remote path so we don't accumulate orphaned files
|
||||
# from prior uploads. upload_large_folder doesn't accept delete_patterns,
|
||||
# so we do this as a separate commit. Safe to run if the path is missing.
|
||||
try:
|
||||
api.delete_folder(
|
||||
path_in_repo=spec.db_filename,
|
||||
repo_id=HF_REPO_ID,
|
||||
repo_type="dataset",
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
with tempfile.TemporaryDirectory() as staging:
|
||||
target = Path(staging) / spec.db_filename
|
||||
target.mkdir()
|
||||
for src in db.rglob("*"):
|
||||
if not src.is_file():
|
||||
continue
|
||||
rel = src.relative_to(db)
|
||||
dest = target / rel
|
||||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||||
os.link(src, dest)
|
||||
|
||||
console.print(f"[blue]Uploading {spec.key} ({db})...[/blue]")
|
||||
api.upload_large_folder(
|
||||
folder_path=staging,
|
||||
repo_id=HF_REPO_ID,
|
||||
repo_type="dataset",
|
||||
)
|
||||
|
||||
console.print(f"[green]Uploaded {spec.key} to {HF_REPO_ID}[/green]")
|
||||
"""Upload evaluation database to HuggingFace (maintainer only)."""
|
||||
for spec in _resolve_datasets(dataset):
|
||||
upload_dataset_db(spec)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
296
evaluations/evaluations/capability_runner.py
Normal file
296
evaluations/evaluations/capability_runner.py
Normal file
|
|
@ -0,0 +1,296 @@
|
|||
from collections.abc import Callable, Iterable
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
from pydantic_ai import Agent
|
||||
from pydantic_ai.messages import (
|
||||
ModelMessage,
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
RetryPromptPart,
|
||||
TextPart,
|
||||
ToolCallPart,
|
||||
ToolReturnPart,
|
||||
UserPromptPart,
|
||||
)
|
||||
from pydantic_ai.models import Model
|
||||
|
||||
from pydantic_ai.capabilities import AbstractCapability
|
||||
|
||||
from evaluations.config import Turn
|
||||
from haiku.rag.capabilities import EvidenceState, RAGCapabilityBase
|
||||
from haiku.rag.capabilities.compaction import create_capability as create_compaction
|
||||
from haiku.rag.capabilities.ledger import citation_status
|
||||
from haiku.rag.config.models import AppConfig
|
||||
|
||||
CapabilityFactory = Callable[..., RAGCapabilityBase[Any]]
|
||||
|
||||
|
||||
def prefix_to_messages(turns: Iterable[Turn]) -> list[ModelMessage]:
|
||||
"""Render a conversation prefix as pydantic-ai message history."""
|
||||
messages: list[ModelMessage] = []
|
||||
for turn in turns:
|
||||
if turn.speaker == "user":
|
||||
messages.append(ModelRequest(parts=[UserPromptPart(content=turn.text)]))
|
||||
else:
|
||||
messages.append(ModelResponse(parts=[TextPart(content=turn.text)]))
|
||||
return messages
|
||||
|
||||
|
||||
@dataclass
|
||||
class CapabilityRunResult:
|
||||
answer: str
|
||||
cited_uris: list[str] = field(default_factory=list)
|
||||
cited_chunk_ids: list[str] = field(default_factory=list)
|
||||
# The database each cited chunk came from, in the order they were cited.
|
||||
# Empty string for a citation built without a source.
|
||||
cited_sources: list[str] = field(default_factory=list)
|
||||
searched_uris: list[str] = field(default_factory=list)
|
||||
n_searches: int = 0
|
||||
n_executions: int = 0
|
||||
n_search_calls: int = 0
|
||||
n_rejected_searches: int = 0
|
||||
n_failed_tools: int = 0
|
||||
n_requests: int = 0
|
||||
citation_status: str | None = None
|
||||
|
||||
|
||||
class ToolTraffic(NamedTuple):
|
||||
n_search_calls: int
|
||||
n_rejected_searches: int
|
||||
n_failed_tools: int
|
||||
n_requests: int
|
||||
|
||||
|
||||
def _count_tool_traffic(
|
||||
messages: list[ModelMessage], namespace: str, tool_names: frozenset[str]
|
||||
) -> ToolTraffic:
|
||||
"""Count search calls, failed calls and model requests in a run.
|
||||
|
||||
The history is the only source: ``state.searches`` is keyed by query so it
|
||||
hides repeats and refusals, and ``for_run`` hands the run a ``replace()``
|
||||
copy, leaving the outer capability's counters at zero.
|
||||
|
||||
Only search failures mean an exhausted budget. A failed code call may be
|
||||
either the execution budget or any error in model-written Python, so
|
||||
``n_failed_tools`` covers both without claiming to tell them apart. It counts
|
||||
``RetryPromptPart`` too, since ``_cite`` rejects with ``ModelRetry`` and only
|
||||
``ToolFailed`` sets ``outcome="failed"``. Both are restricted to
|
||||
``tool_names``, excluding host tools and output-validation retries.
|
||||
|
||||
``n_requests`` counts the run's requests, which matches the capability's own
|
||||
budget only while it stays loaded — a deferred capability skips hooks until
|
||||
it loads.
|
||||
"""
|
||||
search_tool = f"{namespace}_search"
|
||||
search_calls = 0
|
||||
rejected_searches = 0
|
||||
failed_tools = 0
|
||||
requests = 0
|
||||
for message in messages:
|
||||
if isinstance(message, ModelResponse):
|
||||
requests += 1
|
||||
search_calls += sum(
|
||||
1
|
||||
for part in message.parts
|
||||
if isinstance(part, ToolCallPart) and part.tool_name == search_tool
|
||||
)
|
||||
continue
|
||||
for part in message.parts:
|
||||
if not isinstance(part, RetryPromptPart | ToolReturnPart):
|
||||
continue
|
||||
if part.tool_name not in tool_names:
|
||||
continue
|
||||
if isinstance(part, RetryPromptPart):
|
||||
failed_tools += 1
|
||||
elif part.outcome == "failed":
|
||||
failed_tools += 1
|
||||
if part.tool_name == search_tool:
|
||||
rejected_searches += 1
|
||||
return ToolTraffic(
|
||||
n_search_calls=search_calls,
|
||||
n_rejected_searches=rejected_searches,
|
||||
n_failed_tools=failed_tools,
|
||||
n_requests=requests,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _EvalDeps:
|
||||
state: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
def _prepare_agent(
|
||||
capability_factory: CapabilityFactory,
|
||||
db_path: Path | None,
|
||||
config: AppConfig,
|
||||
capability_model: str | Model,
|
||||
document_filter: str | None,
|
||||
request_limit: int | None,
|
||||
compaction: bool = False,
|
||||
) -> tuple[RAGCapabilityBase[Any], _EvalDeps, Agent[_EvalDeps, str]]:
|
||||
capability = capability_factory(
|
||||
db_path=db_path,
|
||||
config=config,
|
||||
defer_loading=False,
|
||||
)
|
||||
if request_limit is not None:
|
||||
capability.request_limit = request_limit
|
||||
state = capability.state_type()
|
||||
if document_filter is not None:
|
||||
state.document_filter = document_filter
|
||||
|
||||
capabilities: list[AbstractCapability] = [capability]
|
||||
if compaction:
|
||||
capabilities.append(create_compaction())
|
||||
deps = _EvalDeps(state={capability.state_namespace: state.model_dump(mode="json")})
|
||||
agent = Agent(
|
||||
capability_model,
|
||||
deps_type=_EvalDeps,
|
||||
capabilities=capabilities,
|
||||
)
|
||||
return capability, deps, agent
|
||||
|
||||
|
||||
def _state_after_run(
|
||||
capability: RAGCapabilityBase[Any], deps: _EvalDeps
|
||||
) -> EvidenceState:
|
||||
return capability.state_type.model_validate(deps.state[capability.state_namespace])
|
||||
|
||||
|
||||
async def run_capability_question(
|
||||
capability_factory: CapabilityFactory,
|
||||
db_path: Path | None,
|
||||
config: AppConfig,
|
||||
question: str,
|
||||
capability_model: str | Model,
|
||||
document_filter: str | None = None,
|
||||
request_limit: int | None = None,
|
||||
message_history: list[ModelMessage] | None = None,
|
||||
) -> CapabilityRunResult:
|
||||
"""Run a single question through a capability and return answer + retrieval data.
|
||||
|
||||
Builds a native capability via ``capability_factory(db_path=..., config=...)``.
|
||||
After the run, citations and searched documents
|
||||
are extracted from the state for downstream eval scoring.
|
||||
|
||||
The capability must produce a state with RAG-capability-shaped fields (citation
|
||||
index, searches, optional document filter) — i.e. ``RAGState`` or
|
||||
``AnalysisState`` from ``haiku.rag.capabilities``.
|
||||
"""
|
||||
capability, deps, agent = _prepare_agent(
|
||||
capability_factory,
|
||||
db_path,
|
||||
config,
|
||||
capability_model,
|
||||
document_filter,
|
||||
request_limit,
|
||||
)
|
||||
agent_result = await agent.run(question, deps=deps, message_history=message_history)
|
||||
traffic = _count_tool_traffic(
|
||||
agent_result.new_messages(), capability.state_namespace, capability.tool_names
|
||||
)
|
||||
return _result_from_run(
|
||||
agent_result.output, _state_after_run(capability, deps), traffic
|
||||
)
|
||||
|
||||
|
||||
async def run_capability_conversation(
|
||||
capability_factory: CapabilityFactory,
|
||||
db_path: Path | None,
|
||||
config: AppConfig,
|
||||
questions: list[str],
|
||||
capability_model: str | Model,
|
||||
document_filter: str | None = None,
|
||||
compaction: bool = False,
|
||||
) -> list[CapabilityRunResult]:
|
||||
"""Run a conversation's user turns sequentially through one capability.
|
||||
|
||||
Each turn runs with the previous turn's full ``all_messages()`` as history
|
||||
(tool calls and returns included) and the same state dict, which is what
|
||||
lets ``EvidenceCompactionCapability`` (registered when ``compaction`` is
|
||||
True) replace earlier questions' evidence on the request. Per-invocation
|
||||
state (citations, searches) is cleared by the capability on every run, so
|
||||
each returned result reflects only its turn.
|
||||
"""
|
||||
capability, deps, agent = _prepare_agent(
|
||||
capability_factory,
|
||||
db_path,
|
||||
config,
|
||||
capability_model,
|
||||
document_filter=document_filter,
|
||||
request_limit=None,
|
||||
compaction=compaction,
|
||||
)
|
||||
history: list[ModelMessage] | None = None
|
||||
results: list[CapabilityRunResult] = []
|
||||
for question in questions:
|
||||
agent_result = await agent.run(question, deps=deps, message_history=history)
|
||||
history = agent_result.all_messages()
|
||||
traffic = _count_tool_traffic(
|
||||
agent_result.new_messages(),
|
||||
capability.state_namespace,
|
||||
capability.tool_names,
|
||||
)
|
||||
results.append(
|
||||
_result_from_run(
|
||||
agent_result.output, _state_after_run(capability, deps), traffic
|
||||
)
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
def _result_from_run(
|
||||
answer: str, typed: EvidenceState, traffic: ToolTraffic
|
||||
) -> CapabilityRunResult:
|
||||
cited_chunk_ids: list[str] = list(typed.citations)
|
||||
seen_cited: set[str] = set()
|
||||
cited_uris: list[str] = []
|
||||
cited_sources: list[str] = []
|
||||
for chunk_id in cited_chunk_ids:
|
||||
citation = typed.citation_index.get(chunk_id)
|
||||
if citation is None:
|
||||
continue
|
||||
cited_sources.append(citation.source or "")
|
||||
if citation.document_uri not in seen_cited:
|
||||
seen_cited.add(citation.document_uri)
|
||||
cited_uris.append(citation.document_uri)
|
||||
|
||||
seen_searched: set[str] = set()
|
||||
searched_uris: list[str] = []
|
||||
for results in typed.searches.values():
|
||||
for search_result in results:
|
||||
uri = search_result.document_uri
|
||||
if uri and uri not in seen_searched:
|
||||
seen_searched.add(uri)
|
||||
searched_uris.append(uri)
|
||||
|
||||
executions = getattr(typed, "executions", None)
|
||||
n_executions = len(executions) if executions is not None else 0
|
||||
|
||||
record = typed.evidence
|
||||
status = (
|
||||
citation_status([record], question=record.question)
|
||||
if record.question is not None
|
||||
else None
|
||||
)
|
||||
|
||||
return CapabilityRunResult(
|
||||
answer=answer,
|
||||
cited_uris=cited_uris,
|
||||
cited_chunk_ids=cited_chunk_ids,
|
||||
cited_sources=cited_sources,
|
||||
searched_uris=searched_uris,
|
||||
# Distinct search keys, not searches. Analysis files every in-code
|
||||
# `search()` under one "_sandbox" key, so twenty sandbox searches read
|
||||
# as one here; `n_search_calls` is the true count of search *tool*
|
||||
# calls, and in-code searches are not counted anywhere.
|
||||
n_searches=len(typed.searches),
|
||||
n_executions=n_executions,
|
||||
n_search_calls=traffic.n_search_calls,
|
||||
n_rejected_searches=traffic.n_rejected_searches,
|
||||
n_failed_tools=traffic.n_failed_tools,
|
||||
n_requests=traffic.n_requests,
|
||||
citation_status=status,
|
||||
)
|
||||
|
|
@ -1,11 +1,42 @@
|
|||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import Any, Literal
|
||||
|
||||
from datasets import Dataset
|
||||
from pydantic import BaseModel, model_validator
|
||||
from pydantic_evals import Case
|
||||
from pydantic_evals.evaluators import Evaluator
|
||||
from haiku.rag.config.models import AppConfig
|
||||
|
||||
|
||||
class Turn(BaseModel):
|
||||
speaker: Literal["user", "agent"]
|
||||
text: str
|
||||
|
||||
|
||||
class ConversationInput(BaseModel):
|
||||
"""A conversation prefix plus the final user question (the last turn)."""
|
||||
|
||||
turns: list[Turn]
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _ends_with_user_turn(self) -> "ConversationInput":
|
||||
if not self.turns or self.turns[-1].speaker != "user":
|
||||
raise ValueError("conversation must end with a user turn")
|
||||
return self
|
||||
|
||||
@property
|
||||
def question(self) -> str:
|
||||
return self.turns[-1].text
|
||||
|
||||
@property
|
||||
def prefix(self) -> list[Turn]:
|
||||
return self.turns[:-1]
|
||||
|
||||
@property
|
||||
def transcript(self) -> str:
|
||||
return "\n".join(f"{turn.speaker}: {turn.text}" for turn in self.turns)
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -30,7 +61,7 @@ DocumentLoader = Callable[[], Dataset]
|
|||
DocumentMapper = Callable[[Mapping[str, Any]], DocumentPayload | None]
|
||||
RetrievalLoader = Callable[[], Dataset]
|
||||
RetrievalMapper = Callable[[Mapping[str, Any]], RetrievalSample | None]
|
||||
CaseBuilder = Callable[[int, Mapping[str, Any]], Case[str, str, dict[str, str]]]
|
||||
CaseBuilder = Callable[[int, Mapping[str, Any]], Case[Any, Any, dict[str, Any]]]
|
||||
|
||||
|
||||
@dataclass
|
||||
|
|
@ -43,9 +74,26 @@ class DatasetSpec:
|
|||
qa_case_builder: CaseBuilder
|
||||
retrieval_loader: RetrievalLoader | None = None
|
||||
retrieval_mapper: RetrievalMapper | None = None
|
||||
retrieval_evaluator: Evaluator | None = None
|
||||
retrieval_evaluators: list[Evaluator] | None = None
|
||||
citation_evaluator: Evaluator | None = None
|
||||
qa_evaluator: Evaluator | None = None
|
||||
document_limit: int | None = None
|
||||
retrieval_limit: int = 5
|
||||
ingest_batch_size: int | None = None
|
||||
live: bool = False
|
||||
compaction: bool = False
|
||||
experiment_metadata: dict[str, Any] | None = None
|
||||
|
||||
def uses_configured_databases(
|
||||
self, config: AppConfig, override_path: Path | None = None
|
||||
) -> bool:
|
||||
"""Whether `lancedb.databases` places the databases to evaluate over.
|
||||
|
||||
`--db PATH` places the database where the configuration places none;
|
||||
`evaluate_dataset` refuses the two together. True for a mapping of one,
|
||||
which is a configured database like any other and keeps its name.
|
||||
"""
|
||||
return bool(config.lancedb.databases) and override_path is None
|
||||
|
||||
def db_path(self, override_path: Path | None = None) -> Path:
|
||||
"""Get the database path.
|
||||
|
|
|
|||
|
|
@ -1,17 +1,29 @@
|
|||
from evaluations.config import DatasetSpec
|
||||
|
||||
from .frames import FRAMES_SPEC
|
||||
from .hotpotqa import HOTPOTQA_SPEC
|
||||
from .mtrag import (
|
||||
MTRAG_CLAPNQ_LIVE_SPEC,
|
||||
MTRAG_CLAPNQ_LIVE_UNCOMPACTED_SPEC,
|
||||
MTRAG_CLAPNQ_REWRITE_SPEC,
|
||||
MTRAG_CLAPNQ_SPEC,
|
||||
)
|
||||
from .open_rag_bench import (
|
||||
ORB_MULTIMODAL_NEMOTRON_SPEC,
|
||||
ORB_MULTIMODAL_SPEC,
|
||||
ORB_TEXT_SPEC,
|
||||
)
|
||||
from .t2_ragbench import T2_FINQA_SPEC, T2_TATDQA_SPEC
|
||||
from .wix import WIX_SPEC
|
||||
|
||||
DATASETS: dict[str, DatasetSpec] = {
|
||||
spec.key: spec
|
||||
for spec in (
|
||||
WIX_SPEC,
|
||||
FRAMES_SPEC,
|
||||
HOTPOTQA_SPEC,
|
||||
MTRAG_CLAPNQ_SPEC,
|
||||
MTRAG_CLAPNQ_REWRITE_SPEC,
|
||||
MTRAG_CLAPNQ_LIVE_SPEC,
|
||||
MTRAG_CLAPNQ_LIVE_UNCOMPACTED_SPEC,
|
||||
ORB_TEXT_SPEC,
|
||||
ORB_MULTIMODAL_SPEC,
|
||||
ORB_MULTIMODAL_NEMOTRON_SPEC,
|
||||
|
|
|
|||
331
evaluations/evaluations/datasets/frames.py
Normal file
331
evaluations/evaluations/datasets/frames.py
Normal file
|
|
@ -0,0 +1,331 @@
|
|||
"""FRAMES benchmark (google/frames-benchmark).
|
||||
|
||||
824 multi-hop questions, each grounded in two or more Wikipedia articles. The
|
||||
corpus is the union of the articles linked per question, fetched from the
|
||||
Wikipedia REST API at current revision and cached locally with the revision id
|
||||
and fetch date.
|
||||
"""
|
||||
|
||||
import ast
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs, quote, unquote, urlsplit
|
||||
|
||||
import httpx
|
||||
from bs4 import BeautifulSoup
|
||||
from datasets import Dataset, load_dataset
|
||||
from pydantic_evals import Case
|
||||
|
||||
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
|
||||
from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
USER_AGENT = "haiku.rag-evaluations (https://github.com/ggozad/haiku.rag)"
|
||||
FETCH_ATTEMPTS = 3
|
||||
THROTTLE_SECONDS = 1.0
|
||||
RATE_LIMIT_BACKOFF_SECONDS = 60.0
|
||||
|
||||
|
||||
# Articles deleted from Wikipedia since FRAMES was authored; the questions
|
||||
# linking them have lost their evidence and are excluded from the benchmark.
|
||||
_DELETED_ARTICLES = frozenset(
|
||||
{
|
||||
"https://en.wikipedia.org/wiki/Nemanja_Marković",
|
||||
"https://en.wikipedia.org/wiki/Jack_Vance_(tennis)",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def load_frames_test() -> Dataset:
|
||||
return load_dataset("google/frames-benchmark")["test"]
|
||||
|
||||
|
||||
def question_is_answerable(doc: Mapping[str, Any]) -> bool:
|
||||
return not _DELETED_ARTICLES & set(question_expected_uris(doc))
|
||||
|
||||
|
||||
def load_frames_questions() -> Dataset:
|
||||
"""Answerable questions with a stable `id` (the dataset row number)."""
|
||||
dataset = load_frames_test().filter(question_is_answerable)
|
||||
return dataset.map(lambda row: {"id": str(row["Unnamed: 0"])})
|
||||
|
||||
|
||||
def parse_wiki_links(raw: str) -> list[str]:
|
||||
"""Extract URLs from a `wiki_links` value.
|
||||
|
||||
The value is a Python-list-repr string. A single list element may pack
|
||||
several comma-separated URLs, and may carry trailing prose annotations;
|
||||
titles themselves can contain commas, so elements are split only where a
|
||||
new URL starts.
|
||||
"""
|
||||
links: list[str] = []
|
||||
for element in ast.literal_eval(raw):
|
||||
for part in re.split(r",\s*(?=http)", element):
|
||||
tokens = part.split()
|
||||
if not tokens:
|
||||
continue
|
||||
url = tokens[0].strip(", ")
|
||||
if url:
|
||||
links.append(url)
|
||||
return links
|
||||
|
||||
|
||||
def normalize_wiki_url(url: str) -> str | None:
|
||||
"""Canonical article URL, used both as document uri and expected uri.
|
||||
|
||||
Strips fragments, decodes percent-escapes, folds mobile hosts, resolves
|
||||
`index.php?title=` and `Special:Search` forms, and applies MediaWiki title
|
||||
canonicalization (underscores, first letter uppercased). Returns None for
|
||||
strings that don't point to an article.
|
||||
"""
|
||||
url = url.strip()
|
||||
if not url:
|
||||
return None
|
||||
if "://" not in url:
|
||||
url = "https://" + url
|
||||
parts = urlsplit(url)
|
||||
host = parts.netloc.replace(".m.wikipedia.org", ".wikipedia.org")
|
||||
if host == "w.wiki":
|
||||
return url
|
||||
if parts.path.startswith("/wiki/"):
|
||||
title = parts.path[len("/wiki/") :]
|
||||
elif parts.path.startswith("/w/index.php"):
|
||||
query = parse_qs(parts.query)
|
||||
title = query.get("title", [""])[0]
|
||||
if not title or title.startswith("Special:"):
|
||||
title = query.get("search", [""])[0]
|
||||
else:
|
||||
return None
|
||||
title = unquote(title).replace(" ", "_").strip("_")
|
||||
if not title:
|
||||
return None
|
||||
return f"https://{host}/wiki/{title[0].upper() + title[1:]}"
|
||||
|
||||
|
||||
def parse_revid(etag: str | None) -> str | None:
|
||||
"""Revision id from a Wikipedia REST ETag header (`W/"<revid>/<uuid>"`)."""
|
||||
if not etag:
|
||||
return None
|
||||
match = re.search(r'"([^/"]+)/', etag)
|
||||
return match.group(1) if match else None
|
||||
|
||||
|
||||
def strip_navigation(html: str) -> str:
|
||||
"""Drop navigation chrome (navboxes, succession boxes) from parsoid HTML.
|
||||
|
||||
These render as link-spam tables naming hundreds of related articles,
|
||||
polluting retrieval. Infoboxes carry no navigation role and are kept.
|
||||
"""
|
||||
soup = BeautifulSoup(html, "html.parser")
|
||||
for element in soup.find_all(attrs={"role": "navigation"}):
|
||||
element.decompose()
|
||||
return str(soup)
|
||||
|
||||
|
||||
def get_cache_dir() -> Path:
|
||||
cache_dir = Path.home() / ".cache" / "haiku.rag" / "evaluations" / "frames_articles"
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
return cache_dir
|
||||
|
||||
|
||||
def _fetch_category_page(
|
||||
host: str, title: str, client: httpx.Client
|
||||
) -> tuple[str, str, str | None]:
|
||||
"""Category pages render empty via page/html; synthesize a members list."""
|
||||
response = client.get(
|
||||
f"https://{host}/w/api.php",
|
||||
params={
|
||||
"action": "query",
|
||||
"list": "categorymembers",
|
||||
"cmtitle": title,
|
||||
"cmlimit": "500",
|
||||
"format": "json",
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
members = [m["title"] for m in response.json()["query"]["categorymembers"]]
|
||||
display = title.replace("_", " ")
|
||||
content = f"# {display}\n\nPages in this category:\n"
|
||||
content += "\n".join(f"- {member}" for member in members) + "\n"
|
||||
return content, "md", None
|
||||
|
||||
|
||||
def _fetch_article_page(
|
||||
uri: str, client: httpx.Client
|
||||
) -> tuple[str, str, str | None, str]:
|
||||
"""Fetch parsoid HTML for an article; returns (content, format, revid, title)."""
|
||||
parts = urlsplit(uri)
|
||||
host = parts.netloc
|
||||
if host == "w.wiki":
|
||||
resolved = urlsplit(str(client.get(uri).url))
|
||||
host = resolved.netloc
|
||||
title = unquote(resolved.path[len("/wiki/") :])
|
||||
else:
|
||||
title = unquote(parts.path[len("/wiki/") :])
|
||||
response = client.get(
|
||||
f"https://{host}/api/rest_v1/page/html/{quote(title, safe='')}"
|
||||
)
|
||||
response.raise_for_status()
|
||||
revid = parse_revid(response.headers.get("etag"))
|
||||
return response.text, "html", revid, title
|
||||
|
||||
|
||||
def _backoff_seconds(error: Exception, attempt: int) -> float:
|
||||
if isinstance(error, httpx.HTTPStatusError) and error.response.status_code == 429:
|
||||
retry_after = error.response.headers.get("retry-after")
|
||||
return float(retry_after) if retry_after else RATE_LIMIT_BACKOFF_SECONDS
|
||||
return 5.0 * attempt
|
||||
|
||||
|
||||
def fetch_article(
|
||||
uri: str, cache_dir: Path, client: httpx.Client | None
|
||||
) -> dict[str, Any] | None:
|
||||
"""Return a corpus row for `uri`, fetching and caching it if needed.
|
||||
|
||||
The cache holds the raw page plus a JSON sidecar with title, format,
|
||||
revision id, and fetch date; a present sidecar marks a complete entry and
|
||||
is served without network access.
|
||||
"""
|
||||
base = quote(uri, safe="")
|
||||
meta_path = cache_dir / f"{base}.json"
|
||||
if meta_path.exists():
|
||||
row = json.loads(meta_path.read_text())
|
||||
row["path"] = str(cache_dir / f"{base}.{row['format']}")
|
||||
return row
|
||||
|
||||
assert client is not None
|
||||
title = unquote(urlsplit(uri).path[len("/wiki/") :])
|
||||
# Wikimedia throttles sustained bot traffic; pace uncached fetches.
|
||||
time.sleep(THROTTLE_SECONDS)
|
||||
for attempt in range(1, FETCH_ATTEMPTS + 1):
|
||||
try:
|
||||
if title.startswith("Category:"):
|
||||
content, format, revid = _fetch_category_page(
|
||||
urlsplit(uri).netloc, title, client
|
||||
)
|
||||
else:
|
||||
content, format, revid, title = _fetch_article_page(uri, client)
|
||||
break
|
||||
except Exception as e:
|
||||
if attempt == FETCH_ATTEMPTS:
|
||||
logger.warning(f"Failed to fetch {uri}: {e}")
|
||||
return None
|
||||
logger.info(f"Retrying {uri} after error: {e}")
|
||||
time.sleep(_backoff_seconds(e, attempt))
|
||||
|
||||
row: dict[str, Any] = {
|
||||
"uri": uri,
|
||||
"title": title.replace("_", " "),
|
||||
"format": format,
|
||||
"revid": revid,
|
||||
"fetched_at": datetime.now(UTC).date().isoformat(),
|
||||
}
|
||||
content_path = cache_dir / f"{base}.{format}"
|
||||
content_path.write_text(content)
|
||||
meta_path.write_text(json.dumps(row))
|
||||
row["path"] = str(content_path)
|
||||
return row
|
||||
|
||||
|
||||
def question_expected_uris(doc: Mapping[str, Any]) -> tuple[str, ...]:
|
||||
uris: list[str] = []
|
||||
for link in parse_wiki_links(doc["wiki_links"]):
|
||||
normalized = normalize_wiki_url(link)
|
||||
if normalized is not None and normalized not in uris:
|
||||
uris.append(normalized)
|
||||
return tuple(uris)
|
||||
|
||||
|
||||
_cached_corpus: list[dict[str, Any]] | None = None
|
||||
|
||||
|
||||
def load_frames_corpus() -> list[dict[str, Any]]:
|
||||
"""Fetch (or read from cache) every article linked by any question."""
|
||||
global _cached_corpus
|
||||
if _cached_corpus is None:
|
||||
uris: dict[str, None] = {}
|
||||
for doc in load_frames_questions():
|
||||
for uri in question_expected_uris(doc):
|
||||
uris.setdefault(uri)
|
||||
cache_dir = get_cache_dir()
|
||||
rows: list[dict[str, Any]] = []
|
||||
with httpx.Client(
|
||||
headers={"User-Agent": USER_AGENT}, follow_redirects=True, timeout=60.0
|
||||
) as client:
|
||||
for index, uri in enumerate(uris, start=1):
|
||||
row = fetch_article(uri, cache_dir, client)
|
||||
if row is not None:
|
||||
rows.append(row)
|
||||
if index % 100 == 0:
|
||||
logger.info(f"Fetched {index}/{len(uris)} articles")
|
||||
logger.info(f"Fetched {len(rows)}/{len(uris)} articles")
|
||||
if len(rows) < len(uris):
|
||||
raise RuntimeError(
|
||||
f"Fetched only {len(rows)}/{len(uris)} FRAMES articles; "
|
||||
"refusing to build a partial corpus. Re-run to resume from cache."
|
||||
)
|
||||
_cached_corpus = rows
|
||||
return _cached_corpus
|
||||
|
||||
|
||||
def document_loader() -> Dataset:
|
||||
return Dataset.from_list(load_frames_corpus())
|
||||
|
||||
|
||||
def map_frames_document(doc: Mapping[str, Any]) -> DocumentPayload:
|
||||
content = Path(doc["path"]).read_text()
|
||||
if doc["format"] == "html":
|
||||
content = strip_navigation(content)
|
||||
metadata: dict[str, str] = {"fetched_at": doc["fetched_at"]}
|
||||
if doc.get("revid"):
|
||||
metadata["revid"] = doc["revid"]
|
||||
return DocumentPayload(
|
||||
uri=doc["uri"],
|
||||
content=content,
|
||||
title=doc["title"],
|
||||
metadata=metadata,
|
||||
format=doc["format"],
|
||||
)
|
||||
|
||||
|
||||
def map_frames_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
|
||||
uris = question_expected_uris(doc)
|
||||
if not uris:
|
||||
return None
|
||||
return RetrievalSample(question=doc["Prompt"], expected_uris=uris)
|
||||
|
||||
|
||||
def build_frames_case(
|
||||
index: int, doc: Mapping[str, Any]
|
||||
) -> Case[str, str, dict[str, str]]:
|
||||
return Case(
|
||||
name=f"{index}_{doc['id']}",
|
||||
inputs=doc["Prompt"],
|
||||
expected_output=doc["Answer"],
|
||||
metadata={
|
||||
"question_id": str(doc["id"]),
|
||||
"reasoning_types": str(doc["reasoning_types"]),
|
||||
"case_index": str(index),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
FRAMES_SPEC = DatasetSpec(
|
||||
key="frames",
|
||||
db_filename="frames.lancedb",
|
||||
document_loader=document_loader,
|
||||
document_mapper=map_frames_document,
|
||||
qa_loader=load_frames_questions,
|
||||
qa_case_builder=build_frames_case,
|
||||
retrieval_loader=load_frames_questions,
|
||||
retrieval_mapper=map_frames_retrieval,
|
||||
retrieval_evaluators=[MAPEvaluator()],
|
||||
citation_evaluator=CitationMAPEvaluator(),
|
||||
)
|
||||
109
evaluations/evaluations/datasets/hotpotqa.py
Normal file
109
evaluations/evaluations/datasets/hotpotqa.py
Normal file
|
|
@ -0,0 +1,109 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Any, cast
|
||||
|
||||
from datasets import Dataset, load_dataset
|
||||
from pydantic_evals import Case
|
||||
|
||||
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
|
||||
from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator
|
||||
|
||||
|
||||
def load_hotpotqa_validation() -> Dataset:
|
||||
dataset_dict = load_dataset("hotpotqa/hotpot_qa", "distractor")
|
||||
return dataset_dict["validation"]
|
||||
|
||||
|
||||
def extract_unique_documents(dataset: Dataset) -> list[dict[str, Any]]:
|
||||
"""Extract unique documents from all context paragraphs, deduplicated by title."""
|
||||
seen_titles: set[str] = set()
|
||||
documents: list[dict[str, Any]] = []
|
||||
|
||||
for sample in dataset:
|
||||
sample = cast(Mapping[str, Any], sample)
|
||||
context = sample["context"]
|
||||
titles = context["title"]
|
||||
sentences_list = context["sentences"]
|
||||
|
||||
for title, sentences in zip(titles, sentences_list):
|
||||
if title in seen_titles:
|
||||
continue
|
||||
seen_titles.add(title)
|
||||
content = " ".join(sentences)
|
||||
documents.append({"title": title, "content": content})
|
||||
|
||||
return documents
|
||||
|
||||
|
||||
_cached_documents: list[dict[str, Any]] | None = None
|
||||
|
||||
|
||||
def load_hotpotqa_documents() -> list[dict[str, Any]]:
|
||||
"""Load and cache unique documents from HotpotQA."""
|
||||
global _cached_documents
|
||||
if _cached_documents is None:
|
||||
dataset = load_hotpotqa_validation()
|
||||
_cached_documents = extract_unique_documents(dataset)
|
||||
return _cached_documents
|
||||
|
||||
|
||||
def document_loader() -> Dataset:
|
||||
"""Return documents as a Dataset-like iterable."""
|
||||
docs = load_hotpotqa_documents()
|
||||
return Dataset.from_list(docs)
|
||||
|
||||
|
||||
def map_hotpotqa_document(doc: Mapping[str, Any]) -> DocumentPayload:
|
||||
return DocumentPayload(
|
||||
uri=doc["title"],
|
||||
content=doc["content"],
|
||||
title=doc["title"],
|
||||
)
|
||||
|
||||
|
||||
def map_hotpotqa_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
|
||||
supporting_facts = doc["supporting_facts"]
|
||||
titles = supporting_facts["title"]
|
||||
if not titles:
|
||||
return None
|
||||
|
||||
unique_titles = tuple(dict.fromkeys(titles))
|
||||
return RetrievalSample(
|
||||
question=doc["question"],
|
||||
expected_uris=unique_titles,
|
||||
)
|
||||
|
||||
|
||||
def build_hotpotqa_case(
|
||||
index: int, doc: Mapping[str, Any]
|
||||
) -> Case[str, str, dict[str, str]]:
|
||||
question_id = doc["id"]
|
||||
question_type = doc["type"]
|
||||
level = doc["level"]
|
||||
|
||||
case_name = f"{index}_{question_id}"
|
||||
|
||||
return Case(
|
||||
name=case_name,
|
||||
inputs=doc["question"],
|
||||
expected_output=doc["answer"],
|
||||
metadata={
|
||||
"question_id": str(question_id),
|
||||
"type": str(question_type),
|
||||
"level": str(level),
|
||||
"case_index": str(index),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
HOTPOTQA_SPEC = DatasetSpec(
|
||||
key="hotpotqa",
|
||||
db_filename="hotpotqa.lancedb",
|
||||
document_loader=document_loader,
|
||||
document_mapper=map_hotpotqa_document,
|
||||
qa_loader=load_hotpotqa_validation,
|
||||
qa_case_builder=build_hotpotqa_case,
|
||||
retrieval_loader=load_hotpotqa_validation,
|
||||
retrieval_mapper=map_hotpotqa_retrieval,
|
||||
retrieval_evaluators=[MAPEvaluator()],
|
||||
citation_evaluator=CitationMAPEvaluator(),
|
||||
)
|
||||
304
evaluations/evaluations/datasets/mtrag.py
Normal file
304
evaluations/evaluations/datasets/mtrag.py
Normal file
|
|
@ -0,0 +1,304 @@
|
|||
import json
|
||||
import zipfile
|
||||
from collections.abc import Iterable, Mapping
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from datasets import Dataset
|
||||
from pydantic_evals import Case
|
||||
|
||||
from evaluations.config import (
|
||||
ConversationInput,
|
||||
DatasetSpec,
|
||||
DocumentPayload,
|
||||
RetrievalSample,
|
||||
Turn,
|
||||
)
|
||||
from evaluations.evaluators import (
|
||||
CitationMAPEvaluator,
|
||||
MAPEvaluator,
|
||||
NDCGEvaluator,
|
||||
RecallEvaluator,
|
||||
)
|
||||
|
||||
REPO_SHA = "cc5b1d481b391181b89f7ced860308482e785463"
|
||||
_BASE_URL = f"https://raw.githubusercontent.com/IBM/mt-rag-benchmark/{REPO_SHA}"
|
||||
|
||||
_CORPUS_FILE = "corpora/passage_level/clapnq.jsonl.zip"
|
||||
_QRELS_FILE = "mtrag-human/retrieval_tasks/clapnq/qrels/dev.tsv"
|
||||
_QUERY_FILES = {
|
||||
"lastturn": "mtrag-human/retrieval_tasks/clapnq/clapnq_lastturn.jsonl",
|
||||
"rewrite": "mtrag-human/retrieval_tasks/clapnq/clapnq_rewrite.jsonl",
|
||||
}
|
||||
_GEN_TASKS_FILE = "mtrag-human/generation_tasks/reference.jsonl"
|
||||
_CLAPNQ_COLLECTION = "mt-rag-clapnq-elser-512-100-20240503"
|
||||
|
||||
|
||||
def get_cache_dir() -> Path:
|
||||
cache_dir = Path.home() / ".cache" / "haiku.rag" / "evaluations" / "mtrag"
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
return cache_dir
|
||||
|
||||
|
||||
def _download(rel_path: str) -> Path:
|
||||
dest = get_cache_dir() / rel_path.replace("/", "_")
|
||||
if dest.exists():
|
||||
return dest
|
||||
|
||||
with httpx.stream(
|
||||
"GET", f"{_BASE_URL}/{rel_path}", timeout=120.0, follow_redirects=True
|
||||
) as response:
|
||||
response.raise_for_status()
|
||||
tmp = dest.with_suffix(dest.suffix + ".part")
|
||||
with tmp.open("wb") as fh:
|
||||
for data in response.iter_bytes():
|
||||
fh.write(data)
|
||||
tmp.rename(dest)
|
||||
return dest
|
||||
|
||||
|
||||
def _parse_qrels(lines: Iterable[str]) -> dict[str, list[str]]:
|
||||
"""Group qrel corpus-ids by query-id, preserving file order."""
|
||||
qrels: dict[str, list[str]] = {}
|
||||
rows = iter(lines)
|
||||
next(rows) # header: query-id / corpus-id / score
|
||||
for line in rows:
|
||||
if not line.strip():
|
||||
continue
|
||||
query_id, corpus_id, _score = line.rstrip("\n").split("\t")
|
||||
qrels.setdefault(query_id, []).append(corpus_id)
|
||||
return qrels
|
||||
|
||||
|
||||
def _validate_qrels_resolve(
|
||||
corpus_ids: set[str], qrels: Mapping[str, list[str]]
|
||||
) -> None:
|
||||
unresolved = sorted(
|
||||
{cid for ids in qrels.values() for cid in ids if cid not in corpus_ids}
|
||||
)
|
||||
if unresolved:
|
||||
raise ValueError(
|
||||
f"{len(unresolved)} qrel corpus-ids do not resolve to corpus "
|
||||
f"passages, e.g. {unresolved[:3]}"
|
||||
)
|
||||
|
||||
|
||||
def _join_queries_qrels(
|
||||
queries: Iterable[Mapping[str, Any]], qrels: Mapping[str, list[str]]
|
||||
) -> list[dict[str, Any]]:
|
||||
records = []
|
||||
for query in queries:
|
||||
query_id = query["_id"]
|
||||
expected = qrels.get(query_id)
|
||||
if expected is None:
|
||||
raise ValueError(f"query {query_id} has no qrels")
|
||||
records.append(
|
||||
{
|
||||
"query_id": query_id,
|
||||
"question": query["text"],
|
||||
"expected_uris": expected,
|
||||
}
|
||||
)
|
||||
return records
|
||||
|
||||
|
||||
def _load_qrels() -> dict[str, list[str]]:
|
||||
path = _download(_QRELS_FILE)
|
||||
return _parse_qrels(path.read_text().splitlines())
|
||||
|
||||
|
||||
def load_clapnq_corpus() -> Dataset:
|
||||
path = _download(_CORPUS_FILE)
|
||||
records: list[dict[str, str]] = []
|
||||
with zipfile.ZipFile(path) as zf:
|
||||
with zf.open(zf.namelist()[0]) as fh:
|
||||
for line in fh:
|
||||
rec = json.loads(line)
|
||||
records.append(
|
||||
{"_id": rec["_id"], "title": rec["title"], "text": rec["text"]}
|
||||
)
|
||||
_validate_qrels_resolve({rec["_id"] for rec in records}, _load_qrels())
|
||||
return Dataset.from_list(records)
|
||||
|
||||
|
||||
def map_mtrag_document(doc: Mapping[str, Any]) -> DocumentPayload:
|
||||
return DocumentPayload(uri=doc["_id"], content=doc["text"], title=doc["title"])
|
||||
|
||||
|
||||
def load_clapnq_retrieval(variant: str) -> Dataset:
|
||||
path = _download(_QUERY_FILES[variant])
|
||||
queries = [json.loads(line) for line in path.read_text().splitlines() if line]
|
||||
return Dataset.from_list(_join_queries_qrels(queries, _load_qrels()))
|
||||
|
||||
|
||||
def map_mtrag_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
|
||||
return RetrievalSample(
|
||||
question=doc["question"],
|
||||
expected_uris=tuple(doc["expected_uris"]),
|
||||
)
|
||||
|
||||
|
||||
def _task_to_record(
|
||||
task: Mapping[str, Any], qrels: Mapping[str, list[str]]
|
||||
) -> dict[str, Any] | None:
|
||||
"""Reduce a reference.jsonl generation task to the fields QA cases need.
|
||||
|
||||
Task `contexts` are the original system's retrievals, never gold relevance;
|
||||
gold passages come from the qrels keyed by task_id.
|
||||
"""
|
||||
if task["Collection"] != _CLAPNQ_COLLECTION:
|
||||
return None
|
||||
return {
|
||||
"id": task["task_id"],
|
||||
"turn": task["turn"],
|
||||
"turns": [
|
||||
{"speaker": message["speaker"], "text": message["text"]}
|
||||
for message in task["input"]
|
||||
],
|
||||
"answer": task["targets"][0]["text"],
|
||||
"answerability": task["Answerability"][0],
|
||||
"multi_turn_type": task["Multi-Turn"][0],
|
||||
"question_type": list(task["Question Type"]),
|
||||
"relevant_uris": qrels.get(task["task_id"]),
|
||||
}
|
||||
|
||||
|
||||
def _qa_records() -> list[dict[str, Any]]:
|
||||
path = _download(_GEN_TASKS_FILE)
|
||||
qrels = _load_qrels()
|
||||
records = []
|
||||
for line in path.read_text().splitlines():
|
||||
if not line.strip():
|
||||
continue
|
||||
record = _task_to_record(json.loads(line), qrels)
|
||||
if record is not None:
|
||||
records.append(record)
|
||||
return records
|
||||
|
||||
|
||||
def load_clapnq_qa() -> Dataset:
|
||||
return Dataset.from_list(_qa_records())
|
||||
|
||||
|
||||
def build_mtrag_case(
|
||||
index: int, doc: Mapping[str, Any]
|
||||
) -> Case[ConversationInput, str, dict[str, Any]]:
|
||||
metadata: dict[str, Any] = {
|
||||
"task_id": doc["id"],
|
||||
"turn": doc["turn"],
|
||||
"answerability": doc["answerability"],
|
||||
"multi_turn_type": doc["multi_turn_type"],
|
||||
"question_type": list(doc["question_type"]),
|
||||
}
|
||||
if doc["relevant_uris"]:
|
||||
metadata["relevant_uris"] = list(doc["relevant_uris"])
|
||||
return Case(
|
||||
name=f"{index}_{doc['id']}",
|
||||
inputs=ConversationInput(
|
||||
turns=[Turn(**turn) for turn in doc["turns"]],
|
||||
),
|
||||
expected_output=doc["answer"],
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
|
||||
def _group_conversations(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Group per-turn generation records into full conversations.
|
||||
|
||||
Turns are ordered numerically within each conversation; each turn carries
|
||||
its user question, reference answer, answerability label, and gold
|
||||
passages when the turn has qrels.
|
||||
"""
|
||||
grouped: dict[str, list[dict[str, Any]]] = {}
|
||||
for record in records:
|
||||
conversation_id = record["id"].split("<::>")[0]
|
||||
grouped.setdefault(conversation_id, []).append(record)
|
||||
|
||||
conversations = []
|
||||
for conversation_id, tasks in grouped.items():
|
||||
tasks.sort(key=lambda record: int(record["turn"]))
|
||||
turns = []
|
||||
for task in tasks:
|
||||
turn: dict[str, Any] = {
|
||||
"task_id": task["id"],
|
||||
"turn": task["turn"],
|
||||
"question": task["turns"][-1]["text"],
|
||||
"reference": task["answer"],
|
||||
"answerability": task["answerability"],
|
||||
"multi_turn_type": task["multi_turn_type"],
|
||||
"question_type": list(task["question_type"]),
|
||||
"relevant_uris": list(task["relevant_uris"] or []),
|
||||
}
|
||||
turns.append(turn)
|
||||
conversations.append({"id": conversation_id, "turns": turns})
|
||||
return conversations
|
||||
|
||||
|
||||
def load_clapnq_conversations() -> Dataset:
|
||||
return Dataset.from_list(_group_conversations(_qa_records()))
|
||||
|
||||
|
||||
def build_mtrag_live_case(
|
||||
index: int, doc: Mapping[str, Any]
|
||||
) -> Case[list[str], list[str], dict[str, Any]]:
|
||||
questions = [turn["question"] for turn in doc["turns"]]
|
||||
metadata_turns = [
|
||||
{key: value for key, value in turn.items() if key != "question"}
|
||||
for turn in doc["turns"]
|
||||
]
|
||||
return Case(
|
||||
name=f"{index}_{doc['id']}",
|
||||
inputs=questions,
|
||||
metadata={"conversation_id": doc["id"], "turns": metadata_turns},
|
||||
)
|
||||
|
||||
|
||||
def _mtrag_spec(key: str, variant: str) -> DatasetSpec:
|
||||
return DatasetSpec(
|
||||
key=key,
|
||||
db_filename="mtrag_clapnq.lancedb",
|
||||
document_loader=load_clapnq_corpus,
|
||||
document_mapper=map_mtrag_document,
|
||||
qa_loader=load_clapnq_qa,
|
||||
qa_case_builder=build_mtrag_case,
|
||||
retrieval_loader=partial(load_clapnq_retrieval, variant),
|
||||
retrieval_mapper=map_mtrag_retrieval,
|
||||
retrieval_evaluators=[
|
||||
RecallEvaluator(k=5),
|
||||
RecallEvaluator(k=10),
|
||||
NDCGEvaluator(k=5),
|
||||
NDCGEvaluator(k=10),
|
||||
MAPEvaluator(),
|
||||
],
|
||||
citation_evaluator=CitationMAPEvaluator(),
|
||||
retrieval_limit=10,
|
||||
ingest_batch_size=512,
|
||||
experiment_metadata={"mtrag_mode": "gold_prefix"},
|
||||
)
|
||||
|
||||
|
||||
MTRAG_CLAPNQ_SPEC = _mtrag_spec("mtrag_clapnq", "lastturn")
|
||||
MTRAG_CLAPNQ_REWRITE_SPEC = _mtrag_spec("mtrag_clapnq_rewrite", "rewrite")
|
||||
|
||||
|
||||
def _mtrag_live_spec(key: str, compaction: bool) -> DatasetSpec:
|
||||
return DatasetSpec(
|
||||
key=key,
|
||||
db_filename="mtrag_clapnq.lancedb",
|
||||
document_loader=load_clapnq_corpus,
|
||||
document_mapper=map_mtrag_document,
|
||||
qa_loader=load_clapnq_conversations,
|
||||
qa_case_builder=build_mtrag_live_case,
|
||||
ingest_batch_size=512,
|
||||
live=True,
|
||||
compaction=compaction,
|
||||
experiment_metadata={"mtrag_mode": "live_session", "compaction": compaction},
|
||||
)
|
||||
|
||||
|
||||
MTRAG_CLAPNQ_LIVE_SPEC = _mtrag_live_spec("mtrag_clapnq_live", compaction=True)
|
||||
MTRAG_CLAPNQ_LIVE_UNCOMPACTED_SPEC = _mtrag_live_spec(
|
||||
"mtrag_clapnq_live_uncompacted", compaction=False
|
||||
)
|
||||
|
|
@ -10,7 +10,7 @@ from huggingface_hub import hf_hub_download
|
|||
from pydantic_evals import Case
|
||||
|
||||
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
|
||||
from evaluations.evaluators import MAPEvaluator
|
||||
from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -226,7 +226,8 @@ def _orb_spec(key: str, db_filename: str) -> DatasetSpec:
|
|||
qa_case_builder=build_orb_case,
|
||||
retrieval_loader=load_orb_retrieval,
|
||||
retrieval_mapper=map_orb_retrieval,
|
||||
retrieval_evaluator=MAPEvaluator(),
|
||||
retrieval_evaluators=[MAPEvaluator()],
|
||||
citation_evaluator=CitationMAPEvaluator(),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,11 @@ from huggingface_hub import hf_hub_download
|
|||
from pydantic_evals import Case
|
||||
|
||||
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
|
||||
from evaluations.evaluators import MAPEvaluator, NumberMatchEvaluator
|
||||
from evaluations.evaluators import (
|
||||
CitationMAPEvaluator,
|
||||
MAPEvaluator,
|
||||
NumberMatchEvaluator,
|
||||
)
|
||||
|
||||
REPO_ID = "G4KMU/t2-ragbench"
|
||||
|
||||
|
|
@ -142,7 +146,8 @@ def _t2_spec(subset: str, key: str, db_filename: str) -> DatasetSpec:
|
|||
qa_case_builder=build_t2_case,
|
||||
retrieval_loader=partial(load_t2_qa, subset),
|
||||
retrieval_mapper=map_t2_retrieval,
|
||||
retrieval_evaluator=MAPEvaluator(),
|
||||
retrieval_evaluators=[MAPEvaluator()],
|
||||
citation_evaluator=CitationMAPEvaluator(),
|
||||
qa_evaluator=NumberMatchEvaluator(),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,84 +0,0 @@
|
|||
import json
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import Any
|
||||
|
||||
from datasets import Dataset, load_dataset
|
||||
from pydantic_evals import Case
|
||||
|
||||
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
|
||||
from evaluations.evaluators import MAPEvaluator
|
||||
|
||||
|
||||
def load_wix_corpus() -> Dataset:
|
||||
dataset_dict = load_dataset("Wix/WixQA", "wix_kb_corpus")
|
||||
return dataset_dict["train"]
|
||||
|
||||
|
||||
def map_wix_document(doc: Mapping[str, Any]) -> DocumentPayload:
|
||||
article_id = doc.get("id")
|
||||
url = doc.get("url")
|
||||
uri = str(article_id) if article_id is not None else str(url)
|
||||
|
||||
metadata: dict[str, str] = {}
|
||||
if article_id is not None:
|
||||
metadata["article_id"] = str(article_id)
|
||||
if url:
|
||||
metadata["url"] = str(url)
|
||||
|
||||
return DocumentPayload(
|
||||
uri=uri,
|
||||
content=doc["html_content"],
|
||||
title=doc.get("title"),
|
||||
metadata=metadata or None,
|
||||
format="html",
|
||||
)
|
||||
|
||||
|
||||
def load_wix_qa() -> Dataset:
|
||||
dataset_dict = load_dataset("Wix/WixQA", "wixqa_expertwritten")
|
||||
return dataset_dict["train"]
|
||||
|
||||
|
||||
def map_wix_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
|
||||
article_ids: Iterable[int | str] | None = doc.get("article_ids")
|
||||
if not article_ids:
|
||||
return None
|
||||
|
||||
expected_uris = tuple(str(article_id) for article_id in article_ids)
|
||||
return RetrievalSample(
|
||||
question=doc["question"],
|
||||
expected_uris=expected_uris,
|
||||
)
|
||||
|
||||
|
||||
def build_wix_case(
|
||||
index: int, doc: Mapping[str, Any]
|
||||
) -> Case[str, str, dict[str, str]]:
|
||||
article_ids = tuple(str(article_id) for article_id in doc.get("article_ids") or [])
|
||||
joined_ids = "-".join(article_ids)
|
||||
case_name = f"{index}_{joined_ids}" if joined_ids else f"case_{index}"
|
||||
|
||||
metadata = {
|
||||
"case_index": str(index),
|
||||
"document_ids": json.dumps(article_ids),
|
||||
}
|
||||
|
||||
return Case(
|
||||
name=case_name,
|
||||
inputs=doc["question"],
|
||||
expected_output=doc["answer"],
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
|
||||
WIX_SPEC = DatasetSpec(
|
||||
key="wix",
|
||||
db_filename="wix.lancedb",
|
||||
document_loader=load_wix_corpus,
|
||||
document_mapper=map_wix_document,
|
||||
qa_loader=load_wix_qa,
|
||||
qa_case_builder=build_wix_case,
|
||||
retrieval_loader=load_wix_qa,
|
||||
retrieval_mapper=map_wix_retrieval,
|
||||
retrieval_evaluator=MAPEvaluator(),
|
||||
)
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
from evaluations.evaluators.citation import CitationMAPEvaluator
|
||||
from evaluations.evaluators.conversation import ConversationEvaluator
|
||||
from evaluations.evaluators.judge import (
|
||||
ANSWER_EQUIVALENCE_RUBRIC,
|
||||
LLMJudge,
|
||||
|
|
@ -6,12 +7,26 @@ from evaluations.evaluators.judge import (
|
|||
)
|
||||
from evaluations.evaluators.map import MAPEvaluator
|
||||
from evaluations.evaluators.number_match import NumberMatchEvaluator
|
||||
from evaluations.evaluators.refusal import (
|
||||
REFUSAL_ELIGIBLE_LABELS,
|
||||
REFUSAL_RUBRIC,
|
||||
RefusalJudge,
|
||||
)
|
||||
from evaluations.evaluators.retrieval import NDCGEvaluator, RecallEvaluator
|
||||
from evaluations.evaluators.transcript import TranscriptLLMJudge
|
||||
|
||||
__all__ = [
|
||||
"ANSWER_EQUIVALENCE_RUBRIC",
|
||||
"REFUSAL_ELIGIBLE_LABELS",
|
||||
"REFUSAL_RUBRIC",
|
||||
"CitationMAPEvaluator",
|
||||
"ConversationEvaluator",
|
||||
"LLMJudge",
|
||||
"LLMJudgeResponseSchema",
|
||||
"MAPEvaluator",
|
||||
"NDCGEvaluator",
|
||||
"NumberMatchEvaluator",
|
||||
"RecallEvaluator",
|
||||
"RefusalJudge",
|
||||
"TranscriptLLMJudge",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_evals.evaluators import Evaluator, EvaluatorContext
|
||||
from pydantic_evals.evaluators.evaluator import EvaluatorOutput
|
||||
|
||||
|
||||
def _cited_uris(ctx: EvaluatorContext) -> list[str]:
|
||||
|
|
@ -13,28 +14,34 @@ def _relevant_uris(ctx: EvaluatorContext) -> set[str]:
|
|||
return set(ctx.metadata.get("relevant_uris", []))
|
||||
|
||||
|
||||
def average_precision(cited: list[str], relevant: set[str]) -> float:
|
||||
"""AP of the cited URIs against the relevant set (0.0 when nothing hits)."""
|
||||
precisions: list[float] = []
|
||||
found = 0
|
||||
for rank, uri in enumerate(cited, start=1):
|
||||
if uri in relevant:
|
||||
found += 1
|
||||
precisions.append(found / rank)
|
||||
if not precisions:
|
||||
return 0.0
|
||||
return sum(precisions) / len(relevant)
|
||||
|
||||
|
||||
@dataclass
|
||||
class CitationMAPEvaluator(Evaluator):
|
||||
"""Average precision over the URIs the skill cited via the `cite` tool.
|
||||
"""Average precision over the URIs the capability cited via the `cite` tool.
|
||||
|
||||
Reads ``cited_uris`` from ``ctx.attributes`` (recorded during the task run
|
||||
via :func:`pydantic_evals.set_eval_attribute`) and ``relevant_uris`` from
|
||||
``ctx.metadata``.
|
||||
``ctx.metadata``. Cases without relevant URIs (e.g. unanswerable turns)
|
||||
are ineligible and produce no score.
|
||||
"""
|
||||
|
||||
def get_default_evaluation_name(self) -> str:
|
||||
return "cited_map"
|
||||
|
||||
def evaluate(self, ctx: EvaluatorContext) -> float:
|
||||
def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput:
|
||||
relevant = _relevant_uris(ctx)
|
||||
if not relevant:
|
||||
return 0.0
|
||||
precisions: list[float] = []
|
||||
found = 0
|
||||
for rank, uri in enumerate(_cited_uris(ctx), start=1):
|
||||
if uri in relevant:
|
||||
found += 1
|
||||
precisions.append(found / rank)
|
||||
if not precisions:
|
||||
return 0.0
|
||||
return sum(precisions) / len(relevant)
|
||||
return {}
|
||||
return average_precision(_cited_uris(ctx), relevant)
|
||||
|
|
|
|||
118
evaluations/evaluations/evaluators/conversation.py
Normal file
118
evaluations/evaluations/evaluators/conversation.py
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_ai import models
|
||||
from pydantic_evals.evaluators import Evaluator, EvaluatorContext
|
||||
from pydantic_evals.evaluators.evaluator import EvaluationReason, EvaluatorOutput
|
||||
from pydantic_evals.evaluators.llm_as_a_judge import (
|
||||
judge_input_output_expected,
|
||||
judge_output,
|
||||
)
|
||||
|
||||
from evaluations.evaluators.citation import average_precision
|
||||
from evaluations.evaluators.refusal import REFUSAL_ELIGIBLE_LABELS, REFUSAL_RUBRIC
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConversationEvaluator(Evaluator):
|
||||
"""Score a live-session conversation turn by turn.
|
||||
|
||||
Expects the case output to be the list of per-turn answers, case inputs
|
||||
the list of user questions, ``metadata["turns"]`` the per-turn reference,
|
||||
answerability label, and optional gold ``relevant_uris``, and the
|
||||
``turn_cited_uris`` attribute the per-turn cited URIs.
|
||||
|
||||
Each turn's answer is judged against the reference with the conversation
|
||||
so far — including the model's own earlier answers — as context. Citation
|
||||
AP is computed on turns with gold passages; refusal on ANSWERABLE and
|
||||
UNANSWERABLE turns. Returned counts allow micro aggregation across
|
||||
conversations; ``turn_pass_rate`` is the per-conversation (macro) rate.
|
||||
Per-turn verdicts are returned as ``turn_{n}_pass`` (with the judge's
|
||||
reason), ``turn_{n}_refused``, and ``turn_{n}_cited_ap`` for diagnosis.
|
||||
"""
|
||||
|
||||
rubric: str
|
||||
model: models.Model | models.KnownModelName | str | None = None
|
||||
|
||||
async def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput:
|
||||
questions: list[str] = list(ctx.inputs)
|
||||
answers: list[str] = list(ctx.output)
|
||||
turns: list[dict] = (ctx.metadata or {}).get("turns", [])
|
||||
turn_cited: list[list[str]] = list(
|
||||
ctx.attributes.get("turn_cited_uris") or [[] for _ in answers]
|
||||
)
|
||||
if not (len(questions) == len(answers) == len(turns) == len(turn_cited)):
|
||||
raise ValueError(
|
||||
f"conversation arrays disagree: {len(questions)} questions, "
|
||||
f"{len(answers)} answers, {len(turns)} turn annotations, "
|
||||
f"{len(turn_cited)} citation lists"
|
||||
)
|
||||
|
||||
passed = 0
|
||||
judged = 0
|
||||
citation_scores: list[float] = []
|
||||
true_refusals = 0
|
||||
false_refusals = 0
|
||||
unanswerable = 0
|
||||
per_turn: dict[str, EvaluationReason | bool | float | str] = {}
|
||||
|
||||
transcript_lines: list[str] = []
|
||||
for index, (question, answer, turn) in enumerate(
|
||||
zip(questions, answers, turns)
|
||||
):
|
||||
number = index + 1
|
||||
transcript_lines.append(f"user: {question}")
|
||||
transcript = "\n".join(transcript_lines)
|
||||
transcript_lines.append(f"agent: {answer}")
|
||||
|
||||
try:
|
||||
grading = await judge_input_output_expected(
|
||||
transcript, answer, turn["reference"], self.rubric, self.model
|
||||
)
|
||||
except Exception as error:
|
||||
per_turn[f"turn_{number}_judge_error"] = str(error)[:200]
|
||||
else:
|
||||
judged += 1
|
||||
if grading.pass_:
|
||||
passed += 1
|
||||
per_turn[f"turn_{number}_pass"] = EvaluationReason(
|
||||
value=grading.pass_, reason=grading.reason
|
||||
)
|
||||
|
||||
label = turn.get("answerability")
|
||||
if label in REFUSAL_ELIGIBLE_LABELS:
|
||||
try:
|
||||
refused = (
|
||||
await judge_output(answer, REFUSAL_RUBRIC, self.model)
|
||||
).pass_
|
||||
except Exception as error:
|
||||
per_turn[f"turn_{number}_judge_error"] = str(error)[:200]
|
||||
else:
|
||||
per_turn[f"turn_{number}_refused"] = refused
|
||||
if label == "UNANSWERABLE":
|
||||
unanswerable += 1
|
||||
if refused:
|
||||
true_refusals += 1
|
||||
elif refused:
|
||||
false_refusals += 1
|
||||
|
||||
relevant = set(turn.get("relevant_uris") or [])
|
||||
if relevant:
|
||||
turn_ap = average_precision(turn_cited[index], relevant)
|
||||
citation_scores.append(turn_ap)
|
||||
per_turn[f"turn_{number}_cited_ap"] = turn_ap
|
||||
|
||||
total = len(answers)
|
||||
result: dict[str, EvaluationReason | bool | float | int | str] = {
|
||||
"turn_pass_rate": passed / judged if judged else 0.0,
|
||||
"turns_passed": passed,
|
||||
"turns_judged": judged,
|
||||
"turns_total": total,
|
||||
"cited_eligible": len(citation_scores),
|
||||
"true_refusals": true_refusals,
|
||||
"false_refusals": false_refusals,
|
||||
"unanswerable_turns": unanswerable,
|
||||
}
|
||||
if citation_scores:
|
||||
result["cited_map"] = sum(citation_scores) / len(citation_scores)
|
||||
result.update(per_turn)
|
||||
return result
|
||||
|
|
@ -2,6 +2,8 @@ from dataclasses import dataclass
|
|||
|
||||
from pydantic_evals.evaluators import Evaluator, EvaluatorContext
|
||||
|
||||
from evaluations.evaluators.citation import average_precision
|
||||
|
||||
|
||||
@dataclass
|
||||
class MAPEvaluator(Evaluator):
|
||||
|
|
@ -28,22 +30,6 @@ class MAPEvaluator(Evaluator):
|
|||
if ctx.metadata is None:
|
||||
return 0.0
|
||||
relevant_uris = set(ctx.metadata.get("relevant_uris", []))
|
||||
retrieved_uris = ctx.output
|
||||
|
||||
if not relevant_uris:
|
||||
return 0.0
|
||||
|
||||
num_relevant = len(relevant_uris)
|
||||
precisions = []
|
||||
num_relevant_found = 0
|
||||
|
||||
for rank, uri in enumerate(retrieved_uris, start=1):
|
||||
if uri in relevant_uris:
|
||||
num_relevant_found += 1
|
||||
precision_at_k = num_relevant_found / rank
|
||||
precisions.append(precision_at_k)
|
||||
|
||||
if not precisions:
|
||||
return 0.0
|
||||
|
||||
return sum(precisions) / num_relevant
|
||||
return average_precision(list(ctx.output), relevant_uris)
|
||||
|
|
|
|||
29
evaluations/evaluations/evaluators/refusal.py
Normal file
29
evaluations/evaluations/evaluators/refusal.py
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
from pydantic_evals.evaluators import EvaluatorContext, LLMJudge
|
||||
from pydantic_evals.evaluators.evaluator import EvaluatorOutput
|
||||
|
||||
REFUSAL_RUBRIC = (
|
||||
"The output declines to answer: it states that the information is not "
|
||||
"available, cannot be found, or is not known, rather than providing a "
|
||||
"substantive answer."
|
||||
)
|
||||
|
||||
REFUSAL_ELIGIBLE_LABELS: Final = ("ANSWERABLE", "UNANSWERABLE")
|
||||
|
||||
|
||||
@dataclass
|
||||
class RefusalJudge(LLMJudge):
|
||||
"""Judges whether the output declines to answer.
|
||||
|
||||
Runs only on turns whose answerability label makes refusal meaningful
|
||||
(ANSWERABLE/UNANSWERABLE); other turns produce no evaluation and cost no
|
||||
judge call.
|
||||
"""
|
||||
|
||||
async def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput:
|
||||
label = (ctx.metadata or {}).get("answerability")
|
||||
if label not in REFUSAL_ELIGIBLE_LABELS:
|
||||
return {}
|
||||
return await super().evaluate(ctx)
|
||||
54
evaluations/evaluations/evaluators/retrieval.py
Normal file
54
evaluations/evaluations/evaluators/retrieval.py
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
import math
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic_evals.evaluators import Evaluator, EvaluatorContext
|
||||
|
||||
|
||||
def _relevant_and_retrieved(ctx: EvaluatorContext) -> tuple[set[str], list[str]]:
|
||||
if ctx.metadata is None:
|
||||
return set(), []
|
||||
return set(ctx.metadata.get("relevant_uris", [])), list(ctx.output)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RecallEvaluator(Evaluator):
|
||||
"""Recall@k: fraction of relevant documents retrieved in the top k."""
|
||||
|
||||
k: int
|
||||
|
||||
def get_default_evaluation_name(self) -> str:
|
||||
return f"recall_{self.k}"
|
||||
|
||||
def evaluate(self, ctx: EvaluatorContext) -> float:
|
||||
relevant, retrieved = _relevant_and_retrieved(ctx)
|
||||
if not relevant:
|
||||
return 0.0
|
||||
found = sum(1 for uri in retrieved[: self.k] if uri in relevant)
|
||||
return found / len(relevant)
|
||||
|
||||
|
||||
@dataclass
|
||||
class NDCGEvaluator(Evaluator):
|
||||
"""Binary nDCG@k: DCG of relevant documents in the top k over the ideal DCG.
|
||||
|
||||
Gains are binary (relevant or not), matching qrels without graded scores.
|
||||
"""
|
||||
|
||||
k: int
|
||||
|
||||
def get_default_evaluation_name(self) -> str:
|
||||
return f"ndcg_{self.k}"
|
||||
|
||||
def evaluate(self, ctx: EvaluatorContext) -> float:
|
||||
relevant, retrieved = _relevant_and_retrieved(ctx)
|
||||
if not relevant:
|
||||
return 0.0
|
||||
dcg = sum(
|
||||
1 / math.log2(rank + 1)
|
||||
for rank, uri in enumerate(retrieved[: self.k], start=1)
|
||||
if uri in relevant
|
||||
)
|
||||
ideal = sum(
|
||||
1 / math.log2(rank + 1) for rank in range(1, min(len(relevant), self.k) + 1)
|
||||
)
|
||||
return dcg / ideal
|
||||
21
evaluations/evaluations/evaluators/transcript.py
Normal file
21
evaluations/evaluations/evaluators/transcript.py
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
from dataclasses import dataclass, replace
|
||||
|
||||
from pydantic_evals.evaluators import EvaluatorContext, LLMJudge
|
||||
from pydantic_evals.evaluators.evaluator import EvaluatorOutput
|
||||
|
||||
from evaluations.config import ConversationInput
|
||||
|
||||
|
||||
@dataclass
|
||||
class TranscriptLLMJudge(LLMJudge):
|
||||
"""LLMJudge that shows conversation inputs as a readable transcript.
|
||||
|
||||
pydantic-evals serializes custom input models as JSON in the judge prompt;
|
||||
a ConversationInput is rendered as `speaker: text` lines instead. Plain
|
||||
string inputs pass through unchanged.
|
||||
"""
|
||||
|
||||
async def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput:
|
||||
if isinstance(ctx.inputs, ConversationInput):
|
||||
ctx = replace(ctx, inputs=ctx.inputs.transcript)
|
||||
return await super().evaluate(ctx)
|
||||
86
evaluations/evaluations/experiment.py
Normal file
86
evaluations/evaluations/experiment.py
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
"""Experiment metadata recorded with every eval run."""
|
||||
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from haiku.rag.config import AppConfig
|
||||
from haiku.rag.config.models import ModelConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from evaluations.qa import Target
|
||||
|
||||
# Pinned judge model. Decoupled from `config.qa.model` so a user changing
|
||||
# their QA model does not inadvertently change the judge — keeps cross-run
|
||||
# comparisons stable. Override per-run with `--judge-model provider:name`.
|
||||
#
|
||||
# Sampling follows Qwen's recommendation for thinking mode; its model cards
|
||||
# forbid greedy decoding. Only the keys ollama honours are set: it silently
|
||||
# ignores `top_k`, `min_p` and `chat_template_kwargs`. The vLLM reference
|
||||
# configs under `evaluations/configs/` carry those too, plus
|
||||
# `reasoning_effort`, which qwen3.8 reads from `chat_template_kwargs`.
|
||||
DEFAULT_JUDGE_MODEL = ModelConfig(
|
||||
provider="ollama",
|
||||
name="qwen3.8",
|
||||
temperature=0.6,
|
||||
max_tokens=16384,
|
||||
extra_body={"top_p": 0.95},
|
||||
)
|
||||
|
||||
|
||||
def build_experiment_metadata(
|
||||
dataset_key: str,
|
||||
test_cases: int,
|
||||
config: AppConfig,
|
||||
judge_config: ModelConfig | None = None,
|
||||
target: "Target" = "rag-capability",
|
||||
capability_config: ModelConfig | None = None,
|
||||
document_filter: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""Build experiment metadata for Logfire tracking."""
|
||||
metadata: dict[str, Any] = {
|
||||
"dataset": dataset_key,
|
||||
"test_cases": test_cases,
|
||||
"target": target,
|
||||
"embedder_provider": config.embeddings.model.provider,
|
||||
"embedder_model": config.embeddings.model.name,
|
||||
"embedder_dim": config.embeddings.model.vector_dim,
|
||||
"chunk_size": config.processing.chunk_size,
|
||||
"search_limit": config.search.limit,
|
||||
"max_context_chars": config.search.max_context_chars,
|
||||
"rerank_provider": config.reranking.model.provider
|
||||
if config.reranking.model
|
||||
else None,
|
||||
"rerank_model": config.reranking.model.name if config.reranking.model else None,
|
||||
"qa_provider": config.qa.model.provider,
|
||||
"qa_model": config.qa.model.name,
|
||||
"qa_temperature": config.qa.model.temperature,
|
||||
"qa_max_tokens": config.qa.model.max_tokens,
|
||||
"qa_enable_thinking": config.qa.model.enable_thinking,
|
||||
"qa_extra_body": config.qa.model.extra_body,
|
||||
"qa_max_searches": config.qa.max_searches,
|
||||
"document_filter": document_filter,
|
||||
}
|
||||
if judge_config is not None:
|
||||
metadata.update(
|
||||
{
|
||||
"judge_provider": judge_config.provider,
|
||||
"judge_model": judge_config.name,
|
||||
"judge_temperature": judge_config.temperature,
|
||||
"judge_max_tokens": judge_config.max_tokens,
|
||||
"judge_enable_thinking": judge_config.enable_thinking,
|
||||
# Sampling and thinking reach vLLM through extra_body, so
|
||||
# without it a trace cannot tell which judge settings ran.
|
||||
"judge_extra_body": judge_config.extra_body,
|
||||
}
|
||||
)
|
||||
if capability_config is not None:
|
||||
metadata.update(
|
||||
{
|
||||
"capability_provider": capability_config.provider,
|
||||
"capability_model": capability_config.name,
|
||||
"capability_temperature": capability_config.temperature,
|
||||
"capability_max_tokens": capability_config.max_tokens,
|
||||
"capability_enable_thinking": capability_config.enable_thinking,
|
||||
"capability_extra_body": capability_config.extra_body,
|
||||
}
|
||||
)
|
||||
return metadata
|
||||
147
evaluations/evaluations/population.py
Normal file
147
evaluations/evaluations/population.py
Normal file
|
|
@ -0,0 +1,147 @@
|
|||
"""Populating an evaluation database from a dataset spec."""
|
||||
|
||||
from collections.abc import Callable, Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
|
||||
from rich.console import Console
|
||||
from rich.progress import Progress
|
||||
|
||||
from evaluations.config import DatasetSpec
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.client.documents import DocumentImport
|
||||
from haiku.rag.config import AppConfig
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
async def _ingest_batched(
|
||||
rag: HaikuRAG,
|
||||
spec: DatasetSpec,
|
||||
corpus,
|
||||
batch_size: int,
|
||||
on_document: Callable[[], None] = lambda: None,
|
||||
) -> None:
|
||||
"""Ingest inline-content documents via `import_documents` batches.
|
||||
|
||||
Each batch writes the documents/chunks/document_items tables once and
|
||||
embeds every chunk in one batched pass. A URI is skipped on resume only
|
||||
when its document has chunks; a chunkless document (crash between the
|
||||
document and chunk writes) is deleted and re-imported.
|
||||
"""
|
||||
uri_rows = await (
|
||||
rag.store.document_meta_table.query().select(["id", "uri"]).to_list()
|
||||
)
|
||||
chunk_rows = await rag.store.chunks_table.query().select(["document_id"]).to_list()
|
||||
chunked_ids = {row["document_id"] for row in chunk_rows}
|
||||
complete = {row["uri"] for row in uri_rows if row["id"] in chunked_ids}
|
||||
chunkless = {
|
||||
row["uri"]: row["id"] for row in uri_rows if row["id"] not in chunked_ids
|
||||
}
|
||||
|
||||
batch: list[DocumentImport] = []
|
||||
for doc in corpus:
|
||||
payload = spec.document_mapper(cast(Mapping[str, Any], doc))
|
||||
if payload is None or payload.uri in complete:
|
||||
on_document()
|
||||
continue
|
||||
if payload.uri in chunkless:
|
||||
await rag.delete_document(chunkless[payload.uri])
|
||||
assert payload.content is not None, "batched ingest requires inline content"
|
||||
docling_document = await rag.convert(payload.content, format=payload.format)
|
||||
chunks = await rag.chunk(docling_document)
|
||||
batch.append(
|
||||
DocumentImport(
|
||||
docling_document=docling_document,
|
||||
chunks=chunks,
|
||||
uri=payload.uri,
|
||||
title=payload.title,
|
||||
metadata=payload.metadata or {},
|
||||
)
|
||||
)
|
||||
if len(batch) >= batch_size:
|
||||
await rag.import_documents(batch)
|
||||
batch = []
|
||||
on_document()
|
||||
|
||||
if batch:
|
||||
await rag.import_documents(batch)
|
||||
|
||||
|
||||
async def populate_db(
|
||||
spec: DatasetSpec,
|
||||
config: AppConfig,
|
||||
db_path: Path | None = None,
|
||||
vacuum_interval: int = 100,
|
||||
) -> None:
|
||||
db = spec.db_path(db_path)
|
||||
db.parent.mkdir(parents=True, exist_ok=True)
|
||||
corpus = spec.document_loader()
|
||||
if spec.document_limit is not None:
|
||||
corpus = corpus.select(range(min(spec.document_limit, len(corpus))))
|
||||
|
||||
# Disable auto_vacuum - we'll vacuum periodically instead to prevent disk exhaustion
|
||||
config.storage.auto_vacuum = False
|
||||
|
||||
with Progress() as progress:
|
||||
task = progress.add_task("[green]Populating database...", total=len(corpus))
|
||||
async with HaikuRAG(db, config=config, create=True) as rag:
|
||||
if spec.ingest_batch_size is not None:
|
||||
await _ingest_batched(
|
||||
rag,
|
||||
spec,
|
||||
corpus,
|
||||
batch_size=spec.ingest_batch_size,
|
||||
on_document=lambda: progress.advance(task),
|
||||
)
|
||||
await rag.store.vacuum(retention_seconds=0)
|
||||
return
|
||||
|
||||
docs_since_vacuum = 0
|
||||
for doc in corpus:
|
||||
doc_mapping = cast(Mapping[str, Any], doc)
|
||||
payload = spec.document_mapper(doc_mapping)
|
||||
if payload is None:
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
# `payload.uri` is the canonical document identifier and is now
|
||||
# honored by both `create_document` and (via the `uri=` override)
|
||||
# `create_document_from_source`, so it's also the right key to
|
||||
# look up an existing document, regardless of whether the source
|
||||
# is a file path or inline content.
|
||||
existing = await rag.get_document_by_uri(payload.uri)
|
||||
if existing is not None:
|
||||
assert existing.id
|
||||
chunks = await rag.chunk_repository.get_by_document_id(existing.id)
|
||||
if chunks:
|
||||
progress.advance(task)
|
||||
continue
|
||||
await rag.document_repository.delete(existing.id)
|
||||
|
||||
if payload.source_path is not None:
|
||||
await rag.create_document_from_source(
|
||||
source=payload.source_path,
|
||||
title=payload.title,
|
||||
metadata=payload.metadata,
|
||||
uri=payload.uri,
|
||||
)
|
||||
else:
|
||||
assert payload.content is not None
|
||||
await rag.create_document(
|
||||
content=payload.content,
|
||||
uri=payload.uri,
|
||||
title=payload.title,
|
||||
metadata=payload.metadata,
|
||||
format=payload.format,
|
||||
)
|
||||
docs_since_vacuum += 1
|
||||
progress.advance(task)
|
||||
|
||||
# Periodic vacuum to prevent disk exhaustion
|
||||
if docs_since_vacuum >= vacuum_interval:
|
||||
await rag.store.vacuum(retention_seconds=0)
|
||||
docs_since_vacuum = 0
|
||||
|
||||
# Final vacuum
|
||||
await rag.store.vacuum(retention_seconds=0)
|
||||
570
evaluations/evaluations/qa.py
Normal file
570
evaluations/evaluations/qa.py
Normal file
|
|
@ -0,0 +1,570 @@
|
|||
"""QA benchmarks: single-question runs and live multi-turn conversations."""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, NamedTuple, cast
|
||||
|
||||
from pydantic_evals import Case, Dataset as EvalDataset, set_eval_attribute
|
||||
from pydantic_evals.evaluators import Evaluator
|
||||
from pydantic_evals.reporting import ReportCaseFailure
|
||||
from rich.console import Console
|
||||
|
||||
from evaluations.capability_runner import (
|
||||
CapabilityFactory,
|
||||
prefix_to_messages,
|
||||
run_capability_conversation,
|
||||
run_capability_question,
|
||||
)
|
||||
from evaluations.config import ConversationInput, DatasetSpec
|
||||
from evaluations.evaluators import (
|
||||
ANSWER_EQUIVALENCE_RUBRIC,
|
||||
REFUSAL_ELIGIBLE_LABELS,
|
||||
REFUSAL_RUBRIC,
|
||||
ConversationEvaluator,
|
||||
RefusalJudge,
|
||||
TranscriptLLMJudge,
|
||||
)
|
||||
from evaluations.experiment import DEFAULT_JUDGE_MODEL, build_experiment_metadata
|
||||
from haiku.rag.config import AppConfig
|
||||
from haiku.rag.config.models import ModelConfig
|
||||
from haiku.rag.utils import get_model
|
||||
|
||||
console = Console()
|
||||
|
||||
Target = Literal["rag-capability", "analysis-capability"]
|
||||
TARGETS: tuple[Target, ...] = ("rag-capability", "analysis-capability")
|
||||
|
||||
|
||||
def _capability_factory_for_target(target: Target) -> CapabilityFactory:
|
||||
if target == "rag-capability":
|
||||
from haiku.rag.capabilities.rag import create_capability
|
||||
|
||||
return create_capability
|
||||
if target == "analysis-capability":
|
||||
from haiku.rag.capabilities.analysis import create_capability
|
||||
|
||||
return create_capability
|
||||
raise ValueError(f"target {target!r} is not a capability target")
|
||||
|
||||
|
||||
def _attach_relevant_uris(
|
||||
cases: list[Case[str, str, dict[str, Any]]],
|
||||
spec: DatasetSpec,
|
||||
limit: int | None,
|
||||
) -> None:
|
||||
"""Augment QA cases with `relevant_uris` joined from retrieval samples.
|
||||
|
||||
Mutates each case's metadata in place. Cases with no matching retrieval
|
||||
sample (by question) are left untouched.
|
||||
"""
|
||||
if spec.retrieval_loader is None or spec.retrieval_mapper is None:
|
||||
return
|
||||
if not any(isinstance(case.inputs, str) for case in cases):
|
||||
return
|
||||
corpus = spec.retrieval_loader()
|
||||
if limit is not None:
|
||||
corpus = corpus.select(range(min(limit, len(corpus))))
|
||||
expected_by_question: dict[str, tuple[str, ...]] = {}
|
||||
for raw in corpus:
|
||||
sample = spec.retrieval_mapper(cast(Mapping[str, Any], raw))
|
||||
if sample is None or sample.skip:
|
||||
continue
|
||||
expected_by_question[sample.question] = sample.expected_uris
|
||||
for case in cases:
|
||||
if not isinstance(case.inputs, str):
|
||||
continue
|
||||
uris = expected_by_question.get(case.inputs)
|
||||
if uris is None:
|
||||
continue
|
||||
metadata = case.metadata if case.metadata is not None else {}
|
||||
metadata["relevant_uris"] = list(uris)
|
||||
case.metadata = metadata
|
||||
|
||||
|
||||
def _resolve_capability_config(
|
||||
target: Target, config: AppConfig, capability_model: ModelConfig | None
|
||||
) -> ModelConfig:
|
||||
if target == "analysis-capability":
|
||||
# Mirror the capability-code resolver: explicit analysis.model wins,
|
||||
# else fall back to qa.model.
|
||||
return capability_model or config.analysis.model or config.qa.model
|
||||
return capability_model or config.qa.model
|
||||
|
||||
|
||||
def _live_summary(report_cases, report_failures) -> dict[str, float | int] | None:
|
||||
"""Aggregate ConversationEvaluator scores across conversations.
|
||||
|
||||
Micro rates weight every turn equally (sums across conversations); macro
|
||||
rates average per-conversation means, so short conversations don't get
|
||||
overweighted by micro nor long ones by macro. Failed conversations are
|
||||
operational exclusions: they count toward the attempted coverage figures
|
||||
but never toward the rates.
|
||||
"""
|
||||
|
||||
def _score(case, key: str):
|
||||
result = case.scores.get(key)
|
||||
return result.value if result is not None else None
|
||||
|
||||
scored = [case for case in report_cases if _score(case, "turns_total") is not None]
|
||||
if not scored:
|
||||
return None
|
||||
|
||||
failed_turns = sum(
|
||||
len(failure.inputs) if isinstance(failure.inputs, list) else 0
|
||||
for failure in report_failures
|
||||
)
|
||||
turns_total = sum(_score(case, "turns_total") for case in scored)
|
||||
turns_judged = sum(_score(case, "turns_judged") for case in scored)
|
||||
turns_passed = sum(_score(case, "turns_passed") for case in scored)
|
||||
# A conversation with zero judged turns (its judge calls all failed)
|
||||
# reports turn_pass_rate 0.0; averaging that in would count a judge
|
||||
# outage as a failed conversation, against the exclusion policy.
|
||||
judged = [case for case in scored if _score(case, "turns_judged")]
|
||||
summary: dict[str, float | int] = {
|
||||
"conversations": len(scored),
|
||||
"conversations_attempted": len(report_cases) + len(report_failures),
|
||||
"turns_total": turns_total,
|
||||
"turns_judged": turns_judged,
|
||||
"turns_attempted": turns_total + failed_turns,
|
||||
"micro_pass_rate": turns_passed / turns_judged if turns_judged else 0.0,
|
||||
"macro_pass_rate": sum(_score(case, "turn_pass_rate") for case in judged)
|
||||
/ len(judged)
|
||||
if judged
|
||||
else 0.0,
|
||||
}
|
||||
|
||||
cited = [case for case in scored if _score(case, "cited_map") is not None]
|
||||
eligible = sum(_score(case, "cited_eligible") for case in scored)
|
||||
if cited and eligible:
|
||||
summary["cited_eligible"] = eligible
|
||||
summary["cited_map_micro"] = (
|
||||
sum(
|
||||
_score(case, "cited_map") * _score(case, "cited_eligible")
|
||||
for case in cited
|
||||
)
|
||||
/ eligible
|
||||
)
|
||||
summary["cited_map_macro"] = sum(
|
||||
_score(case, "cited_map") for case in cited
|
||||
) / len(cited)
|
||||
|
||||
true_refusals = sum(_score(case, "true_refusals") for case in scored)
|
||||
false_refusals = sum(_score(case, "false_refusals") for case in scored)
|
||||
unanswerable = sum(_score(case, "unanswerable_turns") for case in scored)
|
||||
refusals = true_refusals + false_refusals
|
||||
summary["unanswerable_turns"] = unanswerable
|
||||
summary["refusals"] = refusals
|
||||
summary["refusal_precision"] = true_refusals / refusals if refusals else 0.0
|
||||
summary["refusal_recall"] = true_refusals / unanswerable if unanswerable else 0.0
|
||||
return summary
|
||||
|
||||
|
||||
def _refusal_metrics(report_cases) -> tuple[float, float, int, int] | None:
|
||||
"""Refusal precision/recall against answerability labels.
|
||||
|
||||
Uses cases the refusal judge scored (ANSWERABLE/UNANSWERABLE turns).
|
||||
Returns (precision, recall, unanswerable_count, refusal_count), or None
|
||||
when no case was judged.
|
||||
"""
|
||||
outcomes: list[tuple[str, bool]] = []
|
||||
for case in report_cases:
|
||||
refused = case.assertions.get("refused")
|
||||
label = (case.metadata or {}).get("answerability")
|
||||
if refused is None or label not in REFUSAL_ELIGIBLE_LABELS:
|
||||
continue
|
||||
outcomes.append((label, bool(refused.value)))
|
||||
if not outcomes:
|
||||
return None
|
||||
refusals = [(label, r) for label, r in outcomes if r]
|
||||
true_refusals = sum(1 for label, _ in refusals if label == "UNANSWERABLE")
|
||||
unanswerable = sum(1 for label, _ in outcomes if label == "UNANSWERABLE")
|
||||
precision = true_refusals / len(refusals) if refusals else 0.0
|
||||
recall = true_refusals / unanswerable if unanswerable else 0.0
|
||||
return precision, recall, unanswerable, len(refusals)
|
||||
|
||||
|
||||
def _filter_qa_corpus(corpus, case_ids: set[str] | None):
|
||||
"""Keep only rows whose ``id`` is in ``case_ids`` (failure-subset reruns).
|
||||
|
||||
Returns the corpus unchanged when ``case_ids`` is None; matching nothing
|
||||
raises.
|
||||
"""
|
||||
if case_ids is None:
|
||||
return corpus
|
||||
filtered = corpus.filter(lambda row: row.get("id") in case_ids)
|
||||
if len(filtered) == 0:
|
||||
raise ValueError(
|
||||
f"--filter-ids matched none of the {len(corpus)} cases. "
|
||||
"Check that the ids belong to this dataset and that its rows are "
|
||||
"keyed by `id`."
|
||||
)
|
||||
return filtered
|
||||
|
||||
|
||||
class _QARun(NamedTuple):
|
||||
cases: list[Case[Any, Any, dict[str, Any]]]
|
||||
db: Path | None
|
||||
judge_config: ModelConfig
|
||||
eval_name: str
|
||||
experiment_metadata: dict[str, Any]
|
||||
capability_factory: CapabilityFactory
|
||||
capability_model: Any
|
||||
|
||||
|
||||
def _prepare_qa_run(
|
||||
spec: DatasetSpec,
|
||||
config: AppConfig,
|
||||
limit: int | None,
|
||||
name: str | None,
|
||||
db_path: Path | None,
|
||||
judge_model: ModelConfig | None,
|
||||
target: Target,
|
||||
capability_model: ModelConfig | None,
|
||||
case_ids: set[str] | None,
|
||||
document_filter: str | None,
|
||||
) -> _QARun:
|
||||
"""Shared setup for the QA runners: cases, models, name and metadata."""
|
||||
corpus = spec.qa_loader()
|
||||
corpus = _filter_qa_corpus(corpus, case_ids)
|
||||
if limit is not None:
|
||||
corpus = corpus.select(range(min(limit, len(corpus))))
|
||||
|
||||
cases = [
|
||||
spec.qa_case_builder(index, cast(Mapping[str, Any], doc))
|
||||
for index, doc in enumerate(corpus, start=1)
|
||||
]
|
||||
|
||||
judge_config = judge_model or DEFAULT_JUDGE_MODEL
|
||||
capability_config = _resolve_capability_config(target, config, capability_model)
|
||||
|
||||
eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
|
||||
experiment_metadata = build_experiment_metadata(
|
||||
dataset_key=spec.key,
|
||||
test_cases=len(cases),
|
||||
config=config,
|
||||
judge_config=judge_config,
|
||||
target=target,
|
||||
capability_config=capability_config,
|
||||
document_filter=document_filter,
|
||||
)
|
||||
experiment_metadata.update(spec.experiment_metadata or {})
|
||||
|
||||
return _QARun(
|
||||
cases=cases,
|
||||
db=None
|
||||
if spec.uses_configured_databases(config, db_path)
|
||||
else spec.db_path(db_path),
|
||||
judge_config=judge_config,
|
||||
eval_name=eval_name,
|
||||
experiment_metadata=experiment_metadata,
|
||||
capability_factory=_capability_factory_for_target(target),
|
||||
capability_model=get_model(capability_config, config),
|
||||
)
|
||||
|
||||
|
||||
def _print_mean_task_time(report_cases, unit: str = "case") -> None:
|
||||
if not report_cases:
|
||||
return
|
||||
mean = sum(case.task_duration for case in report_cases) / len(report_cases)
|
||||
console.print(f"Avg task time per {unit}: {mean:.2f}s")
|
||||
|
||||
|
||||
def _print_failures(failures, show_question: bool = False) -> None:
|
||||
if not failures:
|
||||
return
|
||||
console.print("[red]\nSummary of failures:[/red]")
|
||||
for failure in failures:
|
||||
console.print(f"Case: {failure.name}")
|
||||
if show_question:
|
||||
console.print(f"Question: {failure.inputs}")
|
||||
console.print(f"Error: {failure.error_message}")
|
||||
console.print("")
|
||||
|
||||
|
||||
async def run_qa_benchmark(
|
||||
spec: DatasetSpec,
|
||||
config: AppConfig,
|
||||
limit: int | None = None,
|
||||
name: str | None = None,
|
||||
db_path: Path | None = None,
|
||||
judge_model: ModelConfig | None = None,
|
||||
target: Target = "rag-capability",
|
||||
capability_model: ModelConfig | None = None,
|
||||
case_ids: set[str] | None = None,
|
||||
document_filter: str | None = None,
|
||||
) -> ReportCaseFailure[str, str, dict[str, str]] | None:
|
||||
run = _prepare_qa_run(
|
||||
spec,
|
||||
config,
|
||||
limit,
|
||||
name,
|
||||
db_path,
|
||||
judge_model,
|
||||
target,
|
||||
capability_model,
|
||||
case_ids,
|
||||
document_filter,
|
||||
)
|
||||
cases, judge_config = run.cases, run.judge_config
|
||||
|
||||
_attach_relevant_uris(cases, spec, limit)
|
||||
citation_evaluator = spec.citation_evaluator
|
||||
|
||||
qa_evaluator = spec.qa_evaluator
|
||||
evaluators: list[Evaluator]
|
||||
if qa_evaluator is not None:
|
||||
evaluators = [qa_evaluator]
|
||||
else:
|
||||
evaluators = [
|
||||
TranscriptLLMJudge(
|
||||
rubric=ANSWER_EQUIVALENCE_RUBRIC,
|
||||
include_input=True,
|
||||
include_expected_output=True,
|
||||
model=get_model(judge_config, config),
|
||||
assertion={
|
||||
"evaluation_name": "answer_equivalent",
|
||||
"include_reason": True,
|
||||
},
|
||||
),
|
||||
]
|
||||
if citation_evaluator is not None:
|
||||
evaluators.append(citation_evaluator)
|
||||
# RefusalJudge scores only cases whose metadata carries an answerability
|
||||
# label; on unlabeled datasets it returns no score without a judge call.
|
||||
evaluators.append(
|
||||
RefusalJudge(
|
||||
rubric=REFUSAL_RUBRIC,
|
||||
model=get_model(judge_config, config),
|
||||
assertion={"evaluation_name": "refused", "include_reason": False},
|
||||
)
|
||||
)
|
||||
|
||||
evaluation_dataset = EvalDataset[Any, str, dict[str, Any]](
|
||||
name=spec.key, cases=cases, evaluators=evaluators
|
||||
)
|
||||
|
||||
async def answer_question(inputs: str | ConversationInput) -> str:
|
||||
if isinstance(inputs, ConversationInput):
|
||||
question = inputs.question
|
||||
message_history = prefix_to_messages(inputs.prefix)
|
||||
else:
|
||||
question = inputs
|
||||
message_history = None
|
||||
result = await run_capability_question(
|
||||
capability_factory=run.capability_factory,
|
||||
db_path=run.db,
|
||||
config=config,
|
||||
question=question,
|
||||
capability_model=run.capability_model,
|
||||
document_filter=document_filter,
|
||||
message_history=message_history,
|
||||
)
|
||||
set_eval_attribute("cited_uris", result.cited_uris)
|
||||
set_eval_attribute("cited_chunk_ids", result.cited_chunk_ids)
|
||||
set_eval_attribute("cited_sources", result.cited_sources)
|
||||
set_eval_attribute("searched_uris", result.searched_uris)
|
||||
set_eval_attribute("n_searches", result.n_searches)
|
||||
set_eval_attribute("n_search_calls", result.n_search_calls)
|
||||
set_eval_attribute("n_rejected_searches", result.n_rejected_searches)
|
||||
set_eval_attribute("n_failed_tools", result.n_failed_tools)
|
||||
set_eval_attribute("n_executions", result.n_executions)
|
||||
set_eval_attribute("n_requests", result.n_requests)
|
||||
set_eval_attribute("citation_status", result.citation_status)
|
||||
return result.answer
|
||||
|
||||
report = await evaluation_dataset.evaluate(
|
||||
answer_question,
|
||||
name=run.eval_name,
|
||||
max_concurrency=1,
|
||||
progress=True,
|
||||
metadata=run.experiment_metadata,
|
||||
)
|
||||
|
||||
total_processed = len(report.cases)
|
||||
failures = report.failures
|
||||
if qa_evaluator is not None:
|
||||
score_key = qa_evaluator.get_default_evaluation_name()
|
||||
passing_cases = sum(
|
||||
1
|
||||
for case in report.cases
|
||||
if score_key in case.scores and case.scores[score_key].value >= 1.0
|
||||
)
|
||||
scoring = score_key
|
||||
else:
|
||||
passing_cases = sum(
|
||||
1
|
||||
for case in report.cases
|
||||
if case.assertions.get("answer_equivalent")
|
||||
and case.assertions["answer_equivalent"].value
|
||||
)
|
||||
scoring = "answer_equivalent"
|
||||
accuracy = passing_cases / total_processed if total_processed > 0 else 0
|
||||
|
||||
console.print("\n=== QA Benchmark Results ===", style="bold cyan")
|
||||
console.print(f"Scoring: {scoring}")
|
||||
console.print(f"Total questions: {total_processed}")
|
||||
console.print(f"Correct answers: {passing_cases}")
|
||||
console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)")
|
||||
_print_mean_task_time(report.cases)
|
||||
|
||||
if citation_evaluator is not None:
|
||||
score_key = citation_evaluator.get_default_evaluation_name()
|
||||
scores = [
|
||||
case.scores[score_key].value
|
||||
for case in report.cases
|
||||
if score_key in case.scores
|
||||
]
|
||||
if scores:
|
||||
cited_count = sum(
|
||||
1 for case in report.cases if case.attributes.get("cited_uris")
|
||||
)
|
||||
mean_citations = sum(
|
||||
len(case.attributes.get("cited_uris") or []) for case in report.cases
|
||||
) / len(report.cases)
|
||||
mean_score = sum(scores) / len(scores)
|
||||
console.print(
|
||||
f"\n=== Citation Retrieval ({score_key}) ===", style="bold cyan"
|
||||
)
|
||||
console.print(f"Mean {score_key}: {mean_score:.4f}")
|
||||
console.print(
|
||||
f"Eligible cases (gold passages known): {len(scores)}/{len(report.cases)}"
|
||||
)
|
||||
console.print(
|
||||
f"Cite rate (≥1 citation): {cited_count / len(report.cases):.2%}"
|
||||
)
|
||||
console.print(f"Mean citations per case: {mean_citations:.2f}")
|
||||
|
||||
if (metrics := _refusal_metrics(report.cases)) is not None:
|
||||
precision, recall, unanswerable, refusals = metrics
|
||||
console.print("\n=== Refusal vs answerability labels ===", style="bold cyan")
|
||||
console.print(f"Refusal precision: {precision:.2%} | recall: {recall:.2%}")
|
||||
console.print(
|
||||
f"UNANSWERABLE turns: {unanswerable} | refusals: {refusals} "
|
||||
"(PARTIAL excluded)"
|
||||
)
|
||||
|
||||
_print_failures(failures, show_question=True)
|
||||
|
||||
return failures[0] if failures else None
|
||||
|
||||
|
||||
async def run_live_qa_benchmark(
|
||||
spec: DatasetSpec,
|
||||
config: AppConfig,
|
||||
limit: int | None = None,
|
||||
name: str | None = None,
|
||||
db_path: Path | None = None,
|
||||
judge_model: ModelConfig | None = None,
|
||||
target: Target = "rag-capability",
|
||||
capability_model: ModelConfig | None = None,
|
||||
case_ids: set[str] | None = None,
|
||||
document_filter: str | None = None,
|
||||
) -> None:
|
||||
"""Replay conversations turn by turn through one capability session.
|
||||
|
||||
One case per conversation; ``limit`` counts conversations. Answers carry
|
||||
forward as real message history, so prior-turn compaction is exercised.
|
||||
"""
|
||||
run = _prepare_qa_run(
|
||||
spec,
|
||||
config,
|
||||
limit,
|
||||
name,
|
||||
db_path,
|
||||
judge_model,
|
||||
target,
|
||||
capability_model,
|
||||
case_ids,
|
||||
document_filter,
|
||||
)
|
||||
|
||||
evaluation_dataset = EvalDataset[Any, Any, dict[str, Any]](
|
||||
name=spec.key,
|
||||
cases=run.cases,
|
||||
evaluators=[
|
||||
ConversationEvaluator(
|
||||
rubric=ANSWER_EQUIVALENCE_RUBRIC,
|
||||
model=get_model(run.judge_config, config),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
async def answer_conversation(questions: list[str]) -> list[str]:
|
||||
results = await run_capability_conversation(
|
||||
capability_factory=run.capability_factory,
|
||||
db_path=run.db,
|
||||
config=config,
|
||||
questions=list(questions),
|
||||
capability_model=run.capability_model,
|
||||
document_filter=document_filter,
|
||||
compaction=spec.compaction,
|
||||
)
|
||||
set_eval_attribute("turn_cited_uris", [r.cited_uris for r in results])
|
||||
set_eval_attribute("turn_n_search_calls", [r.n_search_calls for r in results])
|
||||
set_eval_attribute(
|
||||
"turn_n_rejected_searches", [r.n_rejected_searches for r in results]
|
||||
)
|
||||
set_eval_attribute("turn_n_failed_tools", [r.n_failed_tools for r in results])
|
||||
set_eval_attribute("turn_n_requests", [r.n_requests for r in results])
|
||||
set_eval_attribute("turn_citation_status", [r.citation_status for r in results])
|
||||
return [r.answer for r in results]
|
||||
|
||||
report = await evaluation_dataset.evaluate(
|
||||
answer_conversation,
|
||||
name=run.eval_name,
|
||||
max_concurrency=1,
|
||||
progress=True,
|
||||
metadata=run.experiment_metadata,
|
||||
)
|
||||
|
||||
summary = _live_summary(report.cases, report.failures)
|
||||
console.print("\n=== Live Conversation Results ===", style="bold cyan")
|
||||
if summary is None:
|
||||
attempted = len(report.cases) + len(report.failures)
|
||||
console.print(f"No conversations were scored ({attempted} attempted).")
|
||||
else:
|
||||
console.print(
|
||||
f"Conversations scored: {summary['conversations']}"
|
||||
f"/{summary['conversations_attempted']} | turns scored: "
|
||||
f"{summary['turns_total']}/{summary['turns_attempted']}"
|
||||
)
|
||||
if summary["turns_judged"] < summary["turns_total"]:
|
||||
console.print(
|
||||
f"Turns judged: {summary['turns_judged']}/{summary['turns_total']} "
|
||||
"(per-turn judge errors excluded from rates)"
|
||||
)
|
||||
if report.failures:
|
||||
console.print(
|
||||
"Failed conversations are operational exclusions — "
|
||||
"not counted as wrong answers."
|
||||
)
|
||||
console.print(
|
||||
f"Answer pass rate — micro (per turn): {summary['micro_pass_rate']:.4f} | "
|
||||
f"macro (per conversation): {summary['macro_pass_rate']:.4f}"
|
||||
)
|
||||
if "cited_map_micro" in summary:
|
||||
console.print(
|
||||
f"cited_map — micro: {summary['cited_map_micro']:.4f} | "
|
||||
f"macro: {summary['cited_map_macro']:.4f} "
|
||||
f"(eligible turns: {summary['cited_eligible']})"
|
||||
)
|
||||
console.print(
|
||||
f"Refusal precision: {summary['refusal_precision']:.2%} | "
|
||||
f"recall: {summary['refusal_recall']:.2%} "
|
||||
f"(UNANSWERABLE turns: {summary['unanswerable_turns']}, "
|
||||
f"refusals: {summary['refusals']})"
|
||||
)
|
||||
if report.cases:
|
||||
mean_task_time = sum(case.task_duration for case in report.cases) / len(
|
||||
report.cases
|
||||
)
|
||||
turns = sum(len(case.output or []) for case in report.cases)
|
||||
per_turn = (
|
||||
sum(case.task_duration for case in report.cases) / turns if turns else 0.0
|
||||
)
|
||||
console.print(
|
||||
f"Avg task time: {mean_task_time:.2f}s per conversation | "
|
||||
f"{per_turn:.2f}s per turn"
|
||||
)
|
||||
|
||||
_print_failures(report.failures)
|
||||
132
evaluations/evaluations/retrieval.py
Normal file
132
evaluations/evaluations/retrieval.py
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
"""Retrieval benchmark: search the corpus and score the ranking."""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Any, cast
|
||||
|
||||
from pydantic_evals import Case, Dataset as EvalDataset
|
||||
from rich.console import Console
|
||||
from rich.progress import Progress
|
||||
|
||||
from evaluations.config import DatasetSpec
|
||||
from evaluations.experiment import build_experiment_metadata
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import AppConfig
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
async def run_retrieval_benchmark(
|
||||
spec: DatasetSpec,
|
||||
config: AppConfig,
|
||||
limit: int | None = None,
|
||||
name: str | None = None,
|
||||
db_path: Path | None = None,
|
||||
multimodal_only: bool = False,
|
||||
document_filter: str | None = None,
|
||||
) -> dict[str, float] | None:
|
||||
if spec.retrieval_loader is None or spec.retrieval_mapper is None:
|
||||
console.print("Skipping retrieval benchmark; no retrieval config.")
|
||||
return None
|
||||
|
||||
corpus = spec.retrieval_loader()
|
||||
if limit is not None:
|
||||
corpus = corpus.select(range(min(limit, len(corpus))))
|
||||
|
||||
cases = []
|
||||
with Progress() as progress:
|
||||
task = progress.add_task("[blue]Building retrieval cases...", total=len(corpus))
|
||||
for doc in corpus:
|
||||
doc_mapping = cast(Mapping[str, Any], doc)
|
||||
sample = spec.retrieval_mapper(doc_mapping)
|
||||
if sample is None or sample.skip:
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
# Filter for multimodal queries if requested
|
||||
if multimodal_only:
|
||||
if sample.source_type is None or "image" not in sample.source_type:
|
||||
progress.advance(task)
|
||||
continue
|
||||
|
||||
case = Case(
|
||||
inputs=sample.question,
|
||||
metadata={
|
||||
"relevant_uris": sample.expected_uris,
|
||||
"source_type": sample.source_type,
|
||||
},
|
||||
)
|
||||
cases.append(case)
|
||||
progress.advance(task)
|
||||
|
||||
if not cases:
|
||||
console.print("No retrieval cases to evaluate.")
|
||||
return None
|
||||
|
||||
if not spec.retrieval_evaluators:
|
||||
raise ValueError(f"No retrieval evaluators configured for dataset: {spec.key}")
|
||||
|
||||
dataset = EvalDataset(
|
||||
name=f"{spec.key}-retrieval",
|
||||
cases=cases,
|
||||
evaluators=list(spec.retrieval_evaluators),
|
||||
)
|
||||
|
||||
db = (
|
||||
None
|
||||
if spec.uses_configured_databases(config, db_path)
|
||||
else spec.db_path(db_path)
|
||||
)
|
||||
async with HaikuRAG(db, config=config, read_only=True) as rag:
|
||||
|
||||
async def retrieval_target(question: str) -> list[str]:
|
||||
chunks = await rag.search(
|
||||
query=question,
|
||||
limit=spec.retrieval_limit,
|
||||
include_images=False,
|
||||
filter=document_filter,
|
||||
)
|
||||
|
||||
seen = set()
|
||||
identifiers = []
|
||||
for result in chunks:
|
||||
uri = result.document_uri
|
||||
if uri and uri not in seen:
|
||||
identifiers.append(uri)
|
||||
seen.add(uri)
|
||||
|
||||
return identifiers
|
||||
|
||||
eval_name = name if name is not None else f"{spec.key}_retrieval_evaluation"
|
||||
|
||||
experiment_metadata = build_experiment_metadata(
|
||||
dataset_key=spec.key,
|
||||
test_cases=len(cases),
|
||||
config=config,
|
||||
document_filter=document_filter,
|
||||
)
|
||||
|
||||
report = await dataset.evaluate(
|
||||
retrieval_target,
|
||||
name=eval_name,
|
||||
max_concurrency=1,
|
||||
progress=True,
|
||||
metadata=experiment_metadata,
|
||||
)
|
||||
|
||||
per_metric: dict[str, list[float]] = {}
|
||||
for case in report.cases:
|
||||
for key, score_result in case.scores.items():
|
||||
per_metric.setdefault(key, []).append(score_result.value)
|
||||
|
||||
console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan")
|
||||
console.print(f"Dataset: {spec.key}")
|
||||
console.print(f"Total queries: {len(cases)}")
|
||||
results: dict[str, float] = {"queries": len(cases)}
|
||||
for key, values in per_metric.items():
|
||||
mean_score = sum(values) / len(values)
|
||||
metric_name = key.replace("Evaluator", "").upper()
|
||||
console.print(f"{metric_name}: {mean_score:.4f}")
|
||||
results[metric_name.lower()] = mean_score
|
||||
|
||||
return results
|
||||
|
|
@ -1,97 +0,0 @@
|
|||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Protocol, cast
|
||||
|
||||
from pydantic_ai.models import Model
|
||||
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
from haiku.skills import run_skill
|
||||
from haiku.skills.models import Skill
|
||||
|
||||
SkillFactory = Callable[..., Skill]
|
||||
|
||||
|
||||
class _RagLikeState(Protocol):
|
||||
document_filter: str | None
|
||||
citation_index: dict[str, Citation]
|
||||
citations: list[str]
|
||||
searches: dict[str, list[SearchResult]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class SkillRunResult:
|
||||
answer: str
|
||||
cited_uris: list[str] = field(default_factory=list)
|
||||
cited_chunk_ids: list[str] = field(default_factory=list)
|
||||
searched_uris: list[str] = field(default_factory=list)
|
||||
n_searches: int = 0
|
||||
n_executions: int = 0
|
||||
|
||||
|
||||
async def run_skill_question(
|
||||
skill_factory: SkillFactory,
|
||||
db_path: Path,
|
||||
config: AppConfig,
|
||||
question: str,
|
||||
skill_model: str | Model,
|
||||
document_filter: str | None = None,
|
||||
request_limit: int | None = None,
|
||||
) -> SkillRunResult:
|
||||
"""Run a single question through a skill and return answer + retrieval data.
|
||||
|
||||
Builds the skill via ``skill_factory(db_path=..., config=...)`` and
|
||||
invokes it with a fresh state instance derived from
|
||||
``skill.state_type``. After the run, citations and searched documents
|
||||
are extracted from the state for downstream eval scoring.
|
||||
|
||||
The skill must produce a state with RAG-skill-shaped fields (citation
|
||||
index, searches, optional document filter) — i.e. ``RAGState`` or
|
||||
``AnalysisState`` from ``haiku.rag.skills``.
|
||||
"""
|
||||
skill = skill_factory(db_path=db_path, config=config)
|
||||
if request_limit is not None:
|
||||
skill.request_limit = request_limit
|
||||
|
||||
if skill.state_type is None:
|
||||
raise ValueError(f"Skill {skill.metadata.name!r} has no state_type")
|
||||
state = skill.state_type()
|
||||
typed = cast(_RagLikeState, state)
|
||||
if document_filter is not None:
|
||||
typed.document_filter = document_filter
|
||||
|
||||
answer, _, _ = await run_skill(skill_model, skill, question, state=state)
|
||||
|
||||
cited_chunk_ids: list[str] = list(typed.citations)
|
||||
seen_cited: set[str] = set()
|
||||
cited_uris: list[str] = []
|
||||
for chunk_id in cited_chunk_ids:
|
||||
citation = typed.citation_index.get(chunk_id)
|
||||
if citation is None:
|
||||
continue
|
||||
if citation.document_uri not in seen_cited:
|
||||
seen_cited.add(citation.document_uri)
|
||||
cited_uris.append(citation.document_uri)
|
||||
|
||||
seen_searched: set[str] = set()
|
||||
searched_uris: list[str] = []
|
||||
for results in typed.searches.values():
|
||||
for result in results:
|
||||
uri = result.document_uri
|
||||
if uri and uri not in seen_searched:
|
||||
seen_searched.add(uri)
|
||||
searched_uris.append(uri)
|
||||
|
||||
executions = getattr(state, "executions", None)
|
||||
n_executions = len(executions) if executions is not None else 0
|
||||
|
||||
return SkillRunResult(
|
||||
answer=answer,
|
||||
cited_uris=cited_uris,
|
||||
cited_chunk_ids=cited_chunk_ids,
|
||||
searched_uris=searched_uris,
|
||||
n_searches=len(typed.searches),
|
||||
n_executions=n_executions,
|
||||
)
|
||||
|
|
@ -19,7 +19,7 @@ def _format_number(value: float) -> str:
|
|||
|
||||
|
||||
def extract_prediction(output: str | None) -> str:
|
||||
"""Pull the primary numeric answer from a skill output, for submission.
|
||||
"""Pull the primary numeric answer from a capability output, for submission.
|
||||
|
||||
Restricts to a declared ``ANSWER:`` line when present (via ``_answer_segment``)
|
||||
so reasoning numbers don't leak. Strips ``$`` and thousands separators,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
name = "haiku.rag-evals"
|
||||
description = "Benchmarking and evaluation scripts for haiku.rag"
|
||||
version = "0.66.0"
|
||||
version = "0.82.1"
|
||||
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
|
||||
license = { text = "MIT" }
|
||||
requires-python = ">=3.12"
|
||||
|
|
@ -10,6 +10,7 @@ requires-python = ">=3.12"
|
|||
dependencies = [
|
||||
"haiku.rag-slim",
|
||||
"pydantic-ai-slim[evals,logfire]>=1.81.0",
|
||||
"beautifulsoup4>=4.12.0",
|
||||
"datasets>=4.6.1",
|
||||
"huggingface_hub>=0.20.0",
|
||||
"typer>=0.21.0,<0.22.0",
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
511
evaluations/tests/test_capability_runner.py
Normal file
511
evaluations/tests/test_capability_runner.py
Normal file
|
|
@ -0,0 +1,511 @@
|
|||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic_ai.messages import (
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
RetryPromptPart,
|
||||
TextPart,
|
||||
ToolCallPart,
|
||||
ToolReturnPart,
|
||||
UserPromptPart,
|
||||
)
|
||||
from pydantic_ai.models.test import TestModel
|
||||
|
||||
from evaluations.capability_runner import (
|
||||
CapabilityRunResult,
|
||||
_count_tool_traffic,
|
||||
run_capability_question,
|
||||
)
|
||||
from haiku.rag.capabilities.analysis import create_capability as create_analysis
|
||||
from haiku.rag.capabilities.rag import create_capability as create_rag
|
||||
from haiku.rag.config.models import AppConfig
|
||||
|
||||
ANALYSIS_TOOLS = frozenset(
|
||||
{"analysis_search", "analysis_execute_code", "analysis_cite"}
|
||||
)
|
||||
|
||||
|
||||
def test_count_tool_traffic_sees_a_rejected_cite_call():
|
||||
"""`_cite` rejects with ModelRetry, which is not a failed ToolReturnPart."""
|
||||
messages = [
|
||||
ModelRequest(parts=[UserPromptPart(content="q")]),
|
||||
ModelResponse(parts=[ToolCallPart("analysis_cite", {"chunk_ids": []})]),
|
||||
ModelRequest(
|
||||
parts=[
|
||||
RetryPromptPart(
|
||||
tool_name="analysis_cite",
|
||||
content="No citations registered: chunk_ids was empty.",
|
||||
tool_call_id="1",
|
||||
)
|
||||
]
|
||||
),
|
||||
ModelResponse(parts=[TextPart("done")]),
|
||||
]
|
||||
|
||||
traffic = _count_tool_traffic(messages, "analysis", ANALYSIS_TOOLS)
|
||||
|
||||
assert traffic.n_failed_tools == 1
|
||||
assert traffic.n_rejected_searches == 0
|
||||
|
||||
|
||||
def test_count_tool_traffic_separates_search_rejections_from_code_errors():
|
||||
"""A crash in model-written Python must not read as budget exhaustion."""
|
||||
messages = [
|
||||
ModelRequest(parts=[UserPromptPart(content="q")]),
|
||||
ModelResponse(parts=[ToolCallPart("analysis_execute_code", {"code": "1/0"})]),
|
||||
ModelRequest(
|
||||
parts=[
|
||||
ToolReturnPart(
|
||||
tool_name="analysis_execute_code",
|
||||
content="ZeroDivisionError",
|
||||
tool_call_id="1",
|
||||
outcome="failed",
|
||||
)
|
||||
]
|
||||
),
|
||||
ModelResponse(parts=[TextPart("done")]),
|
||||
]
|
||||
|
||||
traffic = _count_tool_traffic(messages, "analysis", ANALYSIS_TOOLS)
|
||||
|
||||
assert traffic.n_search_calls == 0
|
||||
assert traffic.n_rejected_searches == 0
|
||||
assert traffic.n_failed_tools == 1
|
||||
assert traffic.n_requests == 2
|
||||
|
||||
|
||||
def test_count_tool_traffic_counts_attempts_not_distinct_queries():
|
||||
"""Rejected and repeated calls both count; `state.searches` hides them."""
|
||||
messages = [
|
||||
ModelRequest(parts=[UserPromptPart(content="q")]),
|
||||
ModelResponse(
|
||||
parts=[
|
||||
ToolCallPart("analysis_search", {"query": "same"}),
|
||||
ToolCallPart("analysis_search", {"query": "same"}),
|
||||
]
|
||||
),
|
||||
ModelRequest(
|
||||
parts=[
|
||||
ToolReturnPart(
|
||||
tool_name="analysis_search", content="results", tool_call_id="1"
|
||||
),
|
||||
ToolReturnPart(
|
||||
tool_name="analysis_search",
|
||||
content="Search limit reached.",
|
||||
tool_call_id="2",
|
||||
outcome="failed",
|
||||
),
|
||||
]
|
||||
),
|
||||
ModelResponse(parts=[TextPart("done")]),
|
||||
]
|
||||
|
||||
traffic = _count_tool_traffic(messages, "analysis", ANALYSIS_TOOLS)
|
||||
|
||||
assert traffic.n_search_calls == 2
|
||||
assert traffic.n_rejected_searches == 1
|
||||
assert traffic.n_requests == 2
|
||||
|
||||
|
||||
async def test_runs_rag_capability_without_legacy_capability_layer(tmp_path):
|
||||
result = await run_capability_question(
|
||||
create_rag,
|
||||
tmp_path / "rag.lancedb",
|
||||
AppConfig(),
|
||||
"hello",
|
||||
TestModel(call_tools=[]),
|
||||
document_filter="uri = 'manual.pdf'",
|
||||
)
|
||||
|
||||
assert result.answer == "success (no tool calls)"
|
||||
assert result.cited_uris == []
|
||||
assert result.n_searches == 0
|
||||
assert result.citation_status == "missing"
|
||||
|
||||
|
||||
async def test_runs_analysis_capability_without_legacy_capability_layer(tmp_path):
|
||||
result = await run_capability_question(
|
||||
create_analysis,
|
||||
tmp_path / "rag.lancedb",
|
||||
AppConfig(),
|
||||
"hello",
|
||||
TestModel(call_tools=[]),
|
||||
request_limit=5,
|
||||
)
|
||||
|
||||
assert result.answer == "success (no tool calls)"
|
||||
assert result.n_executions == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("override", "expected"), [(None, 30), (5, 5)])
|
||||
async def test_analysis_capability_applies_request_limit(tmp_path, override, expected):
|
||||
capability = create_analysis(
|
||||
db_path=tmp_path / "rag.lancedb",
|
||||
config=AppConfig(),
|
||||
defer_loading=False,
|
||||
)
|
||||
with patch(
|
||||
"evaluations.capability_runner.Agent.run", new_callable=AsyncMock
|
||||
) as run:
|
||||
run.return_value = SimpleNamespace(
|
||||
output="done", all_messages=lambda: [], new_messages=lambda: []
|
||||
)
|
||||
|
||||
await run_capability_question(
|
||||
lambda **_kwargs: capability,
|
||||
tmp_path / "rag.lancedb",
|
||||
AppConfig(),
|
||||
"hello",
|
||||
TestModel(call_tools=[]),
|
||||
request_limit=override,
|
||||
)
|
||||
|
||||
assert capability.request_limit == expected
|
||||
assert "usage_limits" not in run.call_args.kwargs
|
||||
|
||||
|
||||
class TestCitationStatusDerivation:
|
||||
"""`citation_status` distinguishes an answer that declared nothing
|
||||
(`missing`) from one that declared ungrounded (`ungrounded`) — refusals
|
||||
now cite an empty list."""
|
||||
|
||||
def _result(self, record) -> CapabilityRunResult:
|
||||
from evaluations.capability_runner import ToolTraffic, _result_from_run
|
||||
from haiku.rag.capabilities.rag import RAGState
|
||||
|
||||
state = RAGState(evidence=record)
|
||||
return _result_from_run("answer", state, ToolTraffic(0, 0, 0, 1))
|
||||
|
||||
def test_grounded(self) -> None:
|
||||
from haiku.rag.capabilities.ledger import (
|
||||
CapabilityEvidenceRecord,
|
||||
CitationDeclaration,
|
||||
EvidenceRef,
|
||||
)
|
||||
|
||||
record = CapabilityEvidenceRecord(
|
||||
question=2,
|
||||
latest_evidence_epoch=3,
|
||||
declaration=CitationDeclaration(
|
||||
question=2,
|
||||
epoch=5,
|
||||
refs=[EvidenceRef(capability="rag", chunk_id="c1")],
|
||||
),
|
||||
)
|
||||
assert self._result(record).citation_status == "grounded"
|
||||
|
||||
def test_ungrounded(self) -> None:
|
||||
from haiku.rag.capabilities.ledger import (
|
||||
CapabilityEvidenceRecord,
|
||||
CitationDeclaration,
|
||||
)
|
||||
|
||||
record = CapabilityEvidenceRecord(
|
||||
question=2,
|
||||
latest_evidence_epoch=3,
|
||||
declaration=CitationDeclaration(question=2, epoch=5, refs=[]),
|
||||
)
|
||||
assert self._result(record).citation_status == "ungrounded"
|
||||
|
||||
def test_missing(self) -> None:
|
||||
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
|
||||
|
||||
record = CapabilityEvidenceRecord(question=2, latest_evidence_epoch=3)
|
||||
assert self._result(record).citation_status == "missing"
|
||||
|
||||
def test_none_without_a_question(self) -> None:
|
||||
"""A record no run ever stamped (mocked runs) derives no status."""
|
||||
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
|
||||
|
||||
assert self._result(CapabilityEvidenceRecord()).citation_status is None
|
||||
|
||||
|
||||
class TestPrefixToMessages:
|
||||
def test_maps_turns_to_model_messages(self) -> None:
|
||||
from pydantic_ai.messages import (
|
||||
ModelRequest,
|
||||
ModelResponse,
|
||||
TextPart,
|
||||
UserPromptPart,
|
||||
)
|
||||
|
||||
from evaluations.capability_runner import prefix_to_messages
|
||||
from evaluations.config import Turn
|
||||
|
||||
messages = prefix_to_messages(
|
||||
[
|
||||
Turn(speaker="user", text="who takes photos of planes?"),
|
||||
Turn(speaker="agent", text="Ground-to-air photographers."),
|
||||
]
|
||||
)
|
||||
|
||||
assert len(messages) == 2
|
||||
assert isinstance(messages[0], ModelRequest)
|
||||
assert isinstance(messages[0].parts[0], UserPromptPart)
|
||||
assert messages[0].parts[0].content == "who takes photos of planes?"
|
||||
assert isinstance(messages[1], ModelResponse)
|
||||
assert isinstance(messages[1].parts[0], TextPart)
|
||||
assert messages[1].parts[0].content == "Ground-to-air photographers."
|
||||
|
||||
def test_empty_prefix(self) -> None:
|
||||
from evaluations.capability_runner import prefix_to_messages
|
||||
|
||||
assert prefix_to_messages([]) == []
|
||||
|
||||
|
||||
async def test_message_history_passed_to_agent_run(tmp_path):
|
||||
from evaluations.capability_runner import prefix_to_messages
|
||||
from evaluations.config import Turn
|
||||
|
||||
history = prefix_to_messages([Turn(speaker="user", text="earlier question")])
|
||||
capability = create_rag(
|
||||
db_path=tmp_path / "rag.lancedb",
|
||||
config=AppConfig(),
|
||||
defer_loading=False,
|
||||
)
|
||||
with patch(
|
||||
"evaluations.capability_runner.Agent.run", new_callable=AsyncMock
|
||||
) as run:
|
||||
run.return_value = SimpleNamespace(output="done", new_messages=lambda: [])
|
||||
|
||||
await run_capability_question(
|
||||
lambda **_kwargs: capability,
|
||||
tmp_path / "rag.lancedb",
|
||||
AppConfig(),
|
||||
"follow-up question",
|
||||
TestModel(call_tools=[]),
|
||||
message_history=history,
|
||||
)
|
||||
|
||||
assert run.call_args.kwargs["message_history"] is history
|
||||
|
||||
|
||||
async def test_conversation_threads_own_messages_across_turns(tmp_path):
|
||||
"""Each turn runs with the previous turn's full message history (including
|
||||
tool traffic), so prior-turn compaction operates on real history."""
|
||||
from evaluations.capability_runner import run_capability_conversation
|
||||
|
||||
capability = create_rag(
|
||||
db_path=tmp_path / "rag.lancedb",
|
||||
config=AppConfig(),
|
||||
defer_loading=False,
|
||||
)
|
||||
histories: list[object] = []
|
||||
|
||||
async def _run(question, deps=None, message_history=None):
|
||||
histories.append(message_history)
|
||||
return SimpleNamespace(
|
||||
output=f"answer to {question}",
|
||||
all_messages=lambda: [f"history after {question}"],
|
||||
new_messages=lambda: [],
|
||||
)
|
||||
|
||||
with patch("evaluations.capability_runner.Agent.run", side_effect=_run):
|
||||
result = await run_capability_conversation(
|
||||
lambda **_kwargs: capability,
|
||||
tmp_path / "rag.lancedb",
|
||||
AppConfig(),
|
||||
["q1", "q2", "q3"],
|
||||
TestModel(call_tools=[]),
|
||||
)
|
||||
|
||||
assert [t.answer for t in result] == [
|
||||
"answer to q1",
|
||||
"answer to q2",
|
||||
"answer to q3",
|
||||
]
|
||||
assert histories == [None, ["history after q1"], ["history after q2"]]
|
||||
|
||||
|
||||
async def test_conversation_applies_document_filter(tmp_path):
|
||||
"""The filter must reach the capability state so every search in the
|
||||
conversation is restricted, same as the single-question runner."""
|
||||
from evaluations.capability_runner import run_capability_conversation
|
||||
|
||||
deps_seen = []
|
||||
|
||||
async def _run(question, deps=None, message_history=None):
|
||||
deps_seen.append(deps)
|
||||
return SimpleNamespace(
|
||||
output="a", all_messages=lambda: [], new_messages=lambda: []
|
||||
)
|
||||
|
||||
with patch("evaluations.capability_runner.Agent.run", side_effect=_run):
|
||||
await run_capability_conversation(
|
||||
create_rag,
|
||||
tmp_path / "rag.lancedb",
|
||||
AppConfig(),
|
||||
["q1"],
|
||||
TestModel(call_tools=[]),
|
||||
document_filter="uri = 'manual.pdf'",
|
||||
)
|
||||
|
||||
assert deps_seen[0].state["rag"]["document_filter"] == "uri = 'manual.pdf'"
|
||||
|
||||
|
||||
async def test_conversation_carries_one_state_dict_across_turns(tmp_path):
|
||||
"""Capabilities read and write state through the deps dict, and the same
|
||||
dict is carried across every turn of a conversation."""
|
||||
from evaluations.capability_runner import run_capability_conversation
|
||||
|
||||
deps_seen: list[object] = []
|
||||
|
||||
async def _run(question, deps=None, message_history=None):
|
||||
deps_seen.append(deps)
|
||||
return SimpleNamespace(
|
||||
output="a", all_messages=lambda: [], new_messages=lambda: []
|
||||
)
|
||||
|
||||
with patch("evaluations.capability_runner.Agent.run", side_effect=_run):
|
||||
await run_capability_conversation(
|
||||
create_rag,
|
||||
tmp_path / "rag.lancedb",
|
||||
AppConfig(),
|
||||
["q1", "q2", "q3"],
|
||||
TestModel(call_tools=[]),
|
||||
)
|
||||
|
||||
assert deps_seen[0] is deps_seen[1] is deps_seen[2]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("compaction", "expected"), [(False, 0), (True, 1)])
|
||||
async def test_conversation_compaction_registration(tmp_path, compaction, expected):
|
||||
from haiku.rag.capabilities.compaction import EvidenceCompactionCapability
|
||||
|
||||
from evaluations.capability_runner import run_capability_conversation
|
||||
|
||||
with patch("evaluations.capability_runner.Agent") as agent_cls:
|
||||
agent_cls.return_value.run = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
output="a", all_messages=lambda: [], new_messages=lambda: []
|
||||
)
|
||||
)
|
||||
await run_capability_conversation(
|
||||
create_rag,
|
||||
tmp_path / "rag.lancedb",
|
||||
AppConfig(),
|
||||
["q1"],
|
||||
TestModel(call_tools=[]),
|
||||
compaction=compaction,
|
||||
)
|
||||
capabilities = agent_cls.call_args.kwargs["capabilities"]
|
||||
|
||||
compactors = [
|
||||
c for c in capabilities if isinstance(c, EvidenceCompactionCapability)
|
||||
]
|
||||
assert len(compactors) == expected
|
||||
assert len(capabilities) == 1 + expected
|
||||
|
||||
|
||||
async def test_conversation_end_to_end_with_compaction(tmp_path):
|
||||
from evaluations.capability_runner import run_capability_conversation
|
||||
|
||||
result = await run_capability_conversation(
|
||||
create_rag,
|
||||
tmp_path / "rag.lancedb",
|
||||
AppConfig(),
|
||||
["first question", "follow-up"],
|
||||
TestModel(call_tools=[]),
|
||||
compaction=True,
|
||||
)
|
||||
|
||||
assert [turn.answer for turn in result] == ["success (no tool calls)"] * 2
|
||||
|
||||
|
||||
async def test_conversation_end_to_end_with_test_model(tmp_path):
|
||||
from evaluations.capability_runner import run_capability_conversation
|
||||
|
||||
result = await run_capability_conversation(
|
||||
create_rag,
|
||||
tmp_path / "rag.lancedb",
|
||||
AppConfig(),
|
||||
["first question", "follow-up"],
|
||||
TestModel(call_tools=[]),
|
||||
)
|
||||
|
||||
assert len(result) == 2
|
||||
assert all(turn.answer == "success (no tool calls)" for turn in result)
|
||||
assert all(turn.cited_uris == [] for turn in result)
|
||||
|
||||
|
||||
async def test_gold_prefix_run_answers_with_history(tmp_path):
|
||||
"""End-to-end through a real Agent: the prefix rides along as history."""
|
||||
from evaluations.capability_runner import prefix_to_messages
|
||||
from evaluations.config import Turn
|
||||
|
||||
history = prefix_to_messages(
|
||||
[
|
||||
Turn(speaker="user", text="who takes photos of planes?"),
|
||||
Turn(speaker="agent", text="Ground-to-air photographers."),
|
||||
]
|
||||
)
|
||||
result = await run_capability_question(
|
||||
create_rag,
|
||||
tmp_path / "rag.lancedb",
|
||||
AppConfig(),
|
||||
"No, I meant photos in the air.",
|
||||
TestModel(call_tools=[]),
|
||||
message_history=history,
|
||||
)
|
||||
|
||||
assert result.answer == "success (no tool calls)"
|
||||
|
||||
|
||||
def test_records_the_database_each_citation_came_from():
|
||||
"""A run over several databases records which one grounded the answer."""
|
||||
from haiku.rag.capabilities._base import EvidenceState
|
||||
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
from evaluations.capability_runner import ToolTraffic, _result_from_run
|
||||
|
||||
def cited(chunk_id: str, source: str | None) -> Citation:
|
||||
return Citation(
|
||||
chunk_id=chunk_id,
|
||||
document_id=f"doc-{chunk_id}",
|
||||
document_uri=f"test://{chunk_id}",
|
||||
content="body",
|
||||
source=source,
|
||||
)
|
||||
|
||||
state = EvidenceState(
|
||||
citations=["a1", "b1", "a2"],
|
||||
citation_index={
|
||||
"a1": cited("a1", "alpha"),
|
||||
"b1": cited("b1", "beta"),
|
||||
"a2": cited("a2", "alpha"),
|
||||
},
|
||||
evidence=CapabilityEvidenceRecord(question=1),
|
||||
)
|
||||
|
||||
result = _result_from_run("answer", state, ToolTraffic(0, 0, 0, 0))
|
||||
|
||||
assert result.cited_sources == ["alpha", "beta", "alpha"]
|
||||
|
||||
|
||||
def test_a_hand_built_citation_without_a_source_records_an_empty_string():
|
||||
"""A citation built without a source is recorded as an empty string."""
|
||||
from haiku.rag.capabilities._base import EvidenceState
|
||||
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
|
||||
from evaluations.capability_runner import ToolTraffic, _result_from_run
|
||||
|
||||
state = EvidenceState(
|
||||
citations=["c1"],
|
||||
citation_index={
|
||||
"c1": Citation(
|
||||
chunk_id="c1",
|
||||
document_id="d1",
|
||||
document_uri="test://one",
|
||||
content="body",
|
||||
)
|
||||
},
|
||||
evidence=CapabilityEvidenceRecord(question=1),
|
||||
)
|
||||
|
||||
result = _result_from_run("answer", state, ToolTraffic(0, 0, 0, 0))
|
||||
|
||||
assert result.cited_sources == [""]
|
||||
|
|
@ -28,17 +28,25 @@ class TestCitationMAPEvaluator:
|
|||
def test_no_matches(self) -> None:
|
||||
assert self.evaluator.evaluate(_ctx(["x", "y"], ["a", "b"])) == 0.0
|
||||
|
||||
def test_no_relevant(self) -> None:
|
||||
assert self.evaluator.evaluate(_ctx(["a"], [])) == 0.0
|
||||
|
||||
def test_no_citations(self) -> None:
|
||||
assert self.evaluator.evaluate(_ctx([], ["a"])) == 0.0
|
||||
|
||||
def test_metadata_none(self) -> None:
|
||||
def test_ineligible_when_no_relevant_uris(self) -> None:
|
||||
"""Turns without gold passages (unanswerable) produce no score at all,
|
||||
not a penalizing zero."""
|
||||
assert self.evaluator.evaluate(_ctx(["a"], [])) == {}
|
||||
|
||||
def test_ineligible_when_relevant_uris_missing(self) -> None:
|
||||
ctx = MagicMock()
|
||||
ctx.metadata = {"answerability": "UNANSWERABLE"}
|
||||
ctx.attributes = {"cited_uris": ["a"]}
|
||||
assert self.evaluator.evaluate(ctx) == {}
|
||||
|
||||
def test_ineligible_when_metadata_none(self) -> None:
|
||||
ctx = MagicMock()
|
||||
ctx.metadata = None
|
||||
ctx.attributes = {"cited_uris": ["a"]}
|
||||
assert self.evaluator.evaluate(ctx) == 0.0
|
||||
assert self.evaluator.evaluate(ctx) == {}
|
||||
|
||||
def test_evaluation_name(self) -> None:
|
||||
assert self.evaluator.get_default_evaluation_name() == "cited_map"
|
||||
|
|
|
|||
|
|
@ -1,7 +1,16 @@
|
|||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from evaluations.config import (
|
||||
ConversationInput,
|
||||
DatasetSpec,
|
||||
DocumentPayload,
|
||||
RetrievalSample,
|
||||
Turn,
|
||||
)
|
||||
|
||||
|
||||
def _make_spec(**kwargs: object) -> DatasetSpec:
|
||||
|
|
@ -49,8 +58,53 @@ class TestDatasetSpecDefaults:
|
|||
spec = _make_spec()
|
||||
assert spec.retrieval_loader is None
|
||||
assert spec.retrieval_mapper is None
|
||||
assert spec.retrieval_evaluator is None
|
||||
assert spec.retrieval_evaluators is None
|
||||
assert spec.citation_evaluator is None
|
||||
assert spec.document_limit is None
|
||||
assert spec.retrieval_limit == 5
|
||||
|
||||
|
||||
class TestConversationInput:
|
||||
def _conversation(self) -> ConversationInput:
|
||||
return ConversationInput(
|
||||
turns=[
|
||||
Turn(speaker="user", text="who takes photos of planes?"),
|
||||
Turn(speaker="agent", text="Ground-to-air photographers."),
|
||||
Turn(speaker="user", text="No, I meant photos in the air."),
|
||||
]
|
||||
)
|
||||
|
||||
def test_question_is_last_turn(self) -> None:
|
||||
assert self._conversation().question == "No, I meant photos in the air."
|
||||
|
||||
def test_prefix_excludes_last_turn(self) -> None:
|
||||
prefix = self._conversation().prefix
|
||||
assert [t.speaker for t in prefix] == ["user", "agent"]
|
||||
|
||||
def test_transcript_renders_speaker_lines(self) -> None:
|
||||
assert self._conversation().transcript == (
|
||||
"user: who takes photos of planes?\n"
|
||||
"agent: Ground-to-air photographers.\n"
|
||||
"user: No, I meant photos in the air."
|
||||
)
|
||||
|
||||
def test_single_turn_has_empty_prefix(self) -> None:
|
||||
conversation = ConversationInput(turns=[Turn(speaker="user", text="hi")])
|
||||
assert conversation.prefix == []
|
||||
assert conversation.question == "hi"
|
||||
|
||||
def test_must_end_with_user_turn(self) -> None:
|
||||
with pytest.raises(ValidationError, match="user turn"):
|
||||
ConversationInput(
|
||||
turns=[
|
||||
Turn(speaker="user", text="q"),
|
||||
Turn(speaker="agent", text="a"),
|
||||
]
|
||||
)
|
||||
|
||||
def test_must_have_turns(self) -> None:
|
||||
with pytest.raises(ValidationError, match="user turn"):
|
||||
ConversationInput(turns=[])
|
||||
|
||||
|
||||
class TestDocumentPayload:
|
||||
|
|
@ -91,3 +145,56 @@ class TestRetrievalSample:
|
|||
)
|
||||
assert sample.skip is True
|
||||
assert sample.source_type == "image"
|
||||
|
||||
|
||||
class TestCoversASet:
|
||||
"""A run over `lancedb.databases` passes no path: a path names one
|
||||
database."""
|
||||
|
||||
def test_a_configured_set_is_covered(self):
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||
|
||||
from evaluations.datasets import DATASETS
|
||||
|
||||
spec = next(iter(DATASETS.values()))
|
||||
config = AppConfig(
|
||||
lancedb=LanceDBConfig(databases={"a": "/a.lancedb", "b": "/b.lancedb"})
|
||||
)
|
||||
|
||||
assert spec.uses_configured_databases(config) is True
|
||||
|
||||
def test_a_named_path_overrides_the_set(self):
|
||||
"""`--db` names the one database to evaluate, whatever the
|
||||
configuration names."""
|
||||
from pathlib import Path as _Path
|
||||
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||
|
||||
from evaluations.datasets import DATASETS
|
||||
|
||||
spec = next(iter(DATASETS.values()))
|
||||
config = AppConfig(
|
||||
lancedb=LanceDBConfig(databases={"a": "/a.lancedb", "b": "/b.lancedb"})
|
||||
)
|
||||
|
||||
assert spec.uses_configured_databases(config, _Path("/chosen.lancedb")) is False
|
||||
|
||||
def test_a_configured_set_of_one_is_still_configured(self):
|
||||
"""A mapping of one is a named database like any other."""
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||
|
||||
from evaluations.datasets import DATASETS
|
||||
|
||||
spec = next(iter(DATASETS.values()))
|
||||
config = AppConfig(lancedb=LanceDBConfig(databases={"a": "/a.lancedb"}))
|
||||
|
||||
assert spec.uses_configured_databases(config) is True
|
||||
|
||||
def test_naming_no_database_is_not_a_set(self):
|
||||
from haiku.rag.config.models import AppConfig
|
||||
|
||||
from evaluations.datasets import DATASETS
|
||||
|
||||
spec = next(iter(DATASETS.values()))
|
||||
|
||||
assert spec.uses_configured_databases(AppConfig()) is False
|
||||
|
|
|
|||
257
evaluations/tests/test_conversation_evaluator.py
Normal file
257
evaluations/tests/test_conversation_evaluator.py
Normal file
|
|
@ -0,0 +1,257 @@
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic_evals.evaluators import EvaluatorContext
|
||||
from pydantic_evals.evaluators.evaluator import EvaluationReason
|
||||
|
||||
from evaluations.evaluators.conversation import ConversationEvaluator
|
||||
|
||||
|
||||
def _ctx(
|
||||
questions: list[str],
|
||||
answers: list[str],
|
||||
turns: list[dict],
|
||||
turn_cited_uris: list[list[str]] | None = None,
|
||||
) -> EvaluatorContext:
|
||||
return EvaluatorContext(
|
||||
name="conv",
|
||||
inputs=questions,
|
||||
metadata={"conversation_id": "conv1", "turns": turns},
|
||||
expected_output=None,
|
||||
output=answers,
|
||||
duration=0.0,
|
||||
_span_tree=MagicMock(),
|
||||
attributes={"turn_cited_uris": turn_cited_uris or [[] for _ in answers]},
|
||||
metrics={},
|
||||
)
|
||||
|
||||
|
||||
def _grading(pass_: bool) -> MagicMock:
|
||||
return MagicMock(score=None, pass_=pass_, reason=None)
|
||||
|
||||
|
||||
class TestConversationEvaluator:
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_turn_scores_and_aggregates(self) -> None:
|
||||
evaluator = ConversationEvaluator(rubric="equivalence rubric", model="test")
|
||||
ctx = _ctx(
|
||||
questions=["q1", "q2", "q3"],
|
||||
answers=["a1", "a2", "a3"],
|
||||
turns=[
|
||||
{
|
||||
"reference": "r1",
|
||||
"answerability": "ANSWERABLE",
|
||||
"relevant_uris": ["p1", "p2"],
|
||||
},
|
||||
{"reference": "r2", "answerability": "UNANSWERABLE"},
|
||||
{
|
||||
"reference": "r3",
|
||||
"answerability": "PARTIAL",
|
||||
"relevant_uris": ["p3"],
|
||||
},
|
||||
],
|
||||
turn_cited_uris=[["p1"], [], ["p3"]],
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"evaluations.evaluators.conversation.judge_input_output_expected",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=[_grading(True), _grading(False), _grading(True)],
|
||||
) as judge_answer,
|
||||
patch(
|
||||
"evaluations.evaluators.conversation.judge_output",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=[_grading(False), _grading(True)],
|
||||
) as judge_refusal,
|
||||
):
|
||||
result = await evaluator.evaluate(ctx)
|
||||
|
||||
assert isinstance(result, dict)
|
||||
assert result == {
|
||||
"turn_pass_rate": pytest.approx(2 / 3),
|
||||
"turns_passed": 2,
|
||||
"turns_judged": 3,
|
||||
"turns_total": 3,
|
||||
"cited_map": pytest.approx((0.5 + 1.0) / 2),
|
||||
"cited_eligible": 2,
|
||||
"true_refusals": 1,
|
||||
"false_refusals": 0,
|
||||
"unanswerable_turns": 1,
|
||||
"turn_1_pass": EvaluationReason(value=True, reason=None),
|
||||
"turn_2_pass": EvaluationReason(value=False, reason=None),
|
||||
"turn_3_pass": EvaluationReason(value=True, reason=None),
|
||||
"turn_1_refused": False,
|
||||
"turn_2_refused": True,
|
||||
"turn_1_cited_ap": 0.5,
|
||||
"turn_3_cited_ap": 1.0,
|
||||
}
|
||||
# Refusal judged only on ANSWERABLE/UNANSWERABLE turns.
|
||||
assert judge_refusal.await_count == 2
|
||||
assert judge_answer.await_count == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_judge_sees_live_transcript(self) -> None:
|
||||
"""Turn 2 is judged against the conversation so far with OUR answer to
|
||||
turn 1, not the reference."""
|
||||
evaluator = ConversationEvaluator(rubric="rubric", model="test")
|
||||
ctx = _ctx(
|
||||
questions=["q1", "q2"],
|
||||
answers=["my a1", "my a2"],
|
||||
turns=[
|
||||
{"reference": "r1", "answerability": "ANSWERABLE"},
|
||||
{"reference": "r2", "answerability": "ANSWERABLE"},
|
||||
],
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"evaluations.evaluators.conversation.judge_input_output_expected",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_grading(True),
|
||||
) as judge_answer,
|
||||
patch(
|
||||
"evaluations.evaluators.conversation.judge_output",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_grading(False),
|
||||
),
|
||||
):
|
||||
await evaluator.evaluate(ctx)
|
||||
|
||||
second_call = judge_answer.await_args_list[1]
|
||||
transcript, answer, reference = second_call.args[:3]
|
||||
assert transcript == "user: q1\nagent: my a1\nuser: q2"
|
||||
assert answer == "my a2"
|
||||
assert reference == "r2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_citation_scores_without_eligible_turns(self) -> None:
|
||||
evaluator = ConversationEvaluator(rubric="rubric", model="test")
|
||||
ctx = _ctx(
|
||||
questions=["q1"],
|
||||
answers=["a1"],
|
||||
turns=[{"reference": "r1", "answerability": "UNANSWERABLE"}],
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"evaluations.evaluators.conversation.judge_input_output_expected",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_grading(False),
|
||||
),
|
||||
patch(
|
||||
"evaluations.evaluators.conversation.judge_output",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_grading(True),
|
||||
),
|
||||
):
|
||||
result = await evaluator.evaluate(ctx)
|
||||
|
||||
assert result == {
|
||||
"turn_pass_rate": 0.0,
|
||||
"turns_passed": 0,
|
||||
"turns_judged": 1,
|
||||
"turns_total": 1,
|
||||
"cited_eligible": 0,
|
||||
"true_refusals": 1,
|
||||
"false_refusals": 0,
|
||||
"unanswerable_turns": 1,
|
||||
"turn_1_pass": EvaluationReason(value=False, reason=None),
|
||||
"turn_1_refused": True,
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mismatched_arrays_raise(self) -> None:
|
||||
evaluator = ConversationEvaluator(rubric="rubric", model="test")
|
||||
ctx = _ctx(
|
||||
questions=["q1", "q2"],
|
||||
answers=["a1"],
|
||||
turns=[{"reference": "r1", "answerability": "ANSWERABLE"}],
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="conversation arrays disagree"):
|
||||
await evaluator.evaluate(ctx)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_judge_error_voids_one_turn_not_the_conversation(self) -> None:
|
||||
evaluator = ConversationEvaluator(rubric="rubric", model="test")
|
||||
ctx = _ctx(
|
||||
questions=["q1", "q2", "q3"],
|
||||
answers=["a1", "a2", "a3"],
|
||||
turns=[
|
||||
{"reference": "r1", "answerability": "ANSWERABLE"},
|
||||
{"reference": "r2", "answerability": "ANSWERABLE"},
|
||||
{"reference": "r3", "answerability": "ANSWERABLE"},
|
||||
],
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"evaluations.evaluators.conversation.judge_input_output_expected",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=[
|
||||
_grading(True),
|
||||
RuntimeError("token limit exceeded"),
|
||||
_grading(True),
|
||||
],
|
||||
),
|
||||
patch(
|
||||
"evaluations.evaluators.conversation.judge_output",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_grading(False),
|
||||
),
|
||||
):
|
||||
result = await evaluator.evaluate(ctx)
|
||||
|
||||
assert result == {
|
||||
"turn_pass_rate": 1.0,
|
||||
"turns_passed": 2,
|
||||
"turns_judged": 2,
|
||||
"turns_total": 3,
|
||||
"cited_eligible": 0,
|
||||
"true_refusals": 0,
|
||||
"false_refusals": 0,
|
||||
"unanswerable_turns": 0,
|
||||
"turn_1_pass": EvaluationReason(value=True, reason=None),
|
||||
"turn_3_pass": EvaluationReason(value=True, reason=None),
|
||||
"turn_2_judge_error": "token limit exceeded",
|
||||
"turn_1_refused": False,
|
||||
"turn_2_refused": False,
|
||||
"turn_3_refused": False,
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refusal_judge_error_skips_refusal_verdict_only(self) -> None:
|
||||
evaluator = ConversationEvaluator(rubric="rubric", model="test")
|
||||
ctx = _ctx(
|
||||
questions=["q1"],
|
||||
answers=["a1"],
|
||||
turns=[{"reference": "r1", "answerability": "UNANSWERABLE"}],
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"evaluations.evaluators.conversation.judge_input_output_expected",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_grading(True),
|
||||
),
|
||||
patch(
|
||||
"evaluations.evaluators.conversation.judge_output",
|
||||
new_callable=AsyncMock,
|
||||
side_effect=RuntimeError("boom"),
|
||||
),
|
||||
):
|
||||
result = await evaluator.evaluate(ctx)
|
||||
|
||||
assert result == {
|
||||
"turn_pass_rate": 1.0,
|
||||
"turns_passed": 1,
|
||||
"turns_judged": 1,
|
||||
"turns_total": 1,
|
||||
"cited_eligible": 0,
|
||||
"true_refusals": 0,
|
||||
"false_refusals": 0,
|
||||
"unanswerable_turns": 0,
|
||||
"turn_1_pass": EvaluationReason(value=True, reason=None),
|
||||
"turn_1_judge_error": "boom",
|
||||
}
|
||||
|
|
@ -1,5 +1,25 @@
|
|||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from evaluations.datasets.frames import (
|
||||
FETCH_ATTEMPTS,
|
||||
build_frames_case,
|
||||
fetch_article,
|
||||
map_frames_document,
|
||||
map_frames_retrieval,
|
||||
normalize_wiki_url,
|
||||
parse_revid,
|
||||
parse_wiki_links,
|
||||
question_is_answerable,
|
||||
strip_navigation,
|
||||
)
|
||||
from evaluations.datasets.hotpotqa import (
|
||||
build_hotpotqa_case,
|
||||
extract_unique_documents,
|
||||
map_hotpotqa_document,
|
||||
map_hotpotqa_retrieval,
|
||||
)
|
||||
from evaluations.datasets.open_rag_bench import (
|
||||
build_orb_case,
|
||||
download_pdf,
|
||||
|
|
@ -14,78 +34,78 @@ from evaluations.datasets.t2_ragbench import (
|
|||
map_t2_document,
|
||||
map_t2_retrieval,
|
||||
)
|
||||
from evaluations.datasets.wix import (
|
||||
build_wix_case,
|
||||
map_wix_document,
|
||||
map_wix_retrieval,
|
||||
)
|
||||
|
||||
|
||||
class TestWix:
|
||||
def test_map_document_with_all_fields(self) -> None:
|
||||
doc = {
|
||||
"id": 123,
|
||||
"url": "https://wix.com/article",
|
||||
"html_content": "<p>Content</p>",
|
||||
"title": "My Article",
|
||||
}
|
||||
payload = map_wix_document(doc)
|
||||
assert payload.uri == "123"
|
||||
assert payload.content == "<p>Content</p>"
|
||||
assert payload.title == "My Article"
|
||||
assert payload.format == "html"
|
||||
assert payload.metadata == {
|
||||
"article_id": "123",
|
||||
"url": "https://wix.com/article",
|
||||
}
|
||||
|
||||
def test_map_document_no_id(self) -> None:
|
||||
doc = {
|
||||
"id": None,
|
||||
"url": "https://wix.com/page",
|
||||
"html_content": "<p>Text</p>",
|
||||
"title": None,
|
||||
}
|
||||
payload = map_wix_document(doc)
|
||||
assert payload.uri == "https://wix.com/page"
|
||||
|
||||
def test_map_document_no_metadata(self) -> None:
|
||||
doc = {"id": None, "url": None, "html_content": "<p>X</p>", "title": None}
|
||||
payload = map_wix_document(doc)
|
||||
assert payload.metadata is None
|
||||
class TestHotpotQA:
|
||||
def test_map_document(self) -> None:
|
||||
doc = {"title": "Albert Einstein", "content": "Was a physicist."}
|
||||
payload = map_hotpotqa_document(doc)
|
||||
assert payload.uri == "Albert Einstein"
|
||||
assert payload.content == "Was a physicist."
|
||||
assert payload.title == "Albert Einstein"
|
||||
|
||||
def test_map_retrieval(self) -> None:
|
||||
doc = {"question": "How to add a page?", "article_ids": [10, 20]}
|
||||
sample = map_wix_retrieval(doc)
|
||||
doc = {
|
||||
"question": "Who was Einstein?",
|
||||
"supporting_facts": {"title": ["Albert Einstein", "Physics"]},
|
||||
}
|
||||
sample = map_hotpotqa_retrieval(doc)
|
||||
assert sample is not None
|
||||
assert sample.question == "How to add a page?"
|
||||
assert sample.expected_uris == ("10", "20")
|
||||
assert sample.expected_uris == ("Albert Einstein", "Physics")
|
||||
|
||||
def test_map_retrieval_no_article_ids(self) -> None:
|
||||
doc = {"question": "Q?", "article_ids": None}
|
||||
assert map_wix_retrieval(doc) is None
|
||||
def test_map_retrieval_deduplicates_titles(self) -> None:
|
||||
doc = {
|
||||
"question": "Q?",
|
||||
"supporting_facts": {"title": ["A", "B", "A"]},
|
||||
}
|
||||
sample = map_hotpotqa_retrieval(doc)
|
||||
assert sample is not None
|
||||
assert sample.expected_uris == ("A", "B")
|
||||
|
||||
def test_map_retrieval_empty_article_ids(self) -> None:
|
||||
doc = {"question": "Q?", "article_ids": []}
|
||||
assert map_wix_retrieval(doc) is None
|
||||
def test_map_retrieval_no_titles(self) -> None:
|
||||
doc = {"question": "Q?", "supporting_facts": {"title": []}}
|
||||
assert map_hotpotqa_retrieval(doc) is None
|
||||
|
||||
def test_build_case(self) -> None:
|
||||
doc = {
|
||||
"question": "How?",
|
||||
"answer": "Like this.",
|
||||
"article_ids": [5, 10],
|
||||
"id": "abc123",
|
||||
"question": "What is X?",
|
||||
"answer": "X is Y.",
|
||||
"type": "comparison",
|
||||
"level": "hard",
|
||||
}
|
||||
case = build_hotpotqa_case(5, doc)
|
||||
assert case.name == "5_abc123"
|
||||
assert case.inputs == "What is X?"
|
||||
assert case.expected_output == "X is Y."
|
||||
assert case.metadata == {
|
||||
"question_id": "abc123",
|
||||
"type": "comparison",
|
||||
"level": "hard",
|
||||
"case_index": "5",
|
||||
}
|
||||
case = build_wix_case(2, doc)
|
||||
assert case.name == "2_5-10"
|
||||
assert case.inputs == "How?"
|
||||
assert case.expected_output == "Like this."
|
||||
assert case.metadata is not None
|
||||
assert case.metadata["case_index"] == "2"
|
||||
|
||||
def test_build_case_no_article_ids(self) -> None:
|
||||
doc = {"question": "Q?", "answer": "A.", "article_ids": None}
|
||||
case = build_wix_case(1, doc)
|
||||
assert case.name == "case_1"
|
||||
def test_extract_unique_documents(self) -> None:
|
||||
# Simulate a minimal dataset with context
|
||||
dataset = [
|
||||
{
|
||||
"context": {
|
||||
"title": ["Doc A", "Doc B"],
|
||||
"sentences": [["Sentence 1."], ["Sentence 2.", " More."]],
|
||||
}
|
||||
},
|
||||
{
|
||||
"context": {
|
||||
"title": ["Doc A", "Doc C"],
|
||||
"sentences": [["Dupe."], ["Sentence 3."]],
|
||||
}
|
||||
},
|
||||
]
|
||||
docs = extract_unique_documents(dataset) # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
|
||||
assert len(docs) == 3
|
||||
titles = [d["title"] for d in docs]
|
||||
assert titles == ["Doc A", "Doc B", "Doc C"]
|
||||
assert docs[1]["content"] == "Sentence 2. More."
|
||||
|
||||
|
||||
class TestOpenRAGBench:
|
||||
|
|
@ -314,3 +334,389 @@ class TestT2RAGBench:
|
|||
|
||||
assert len(corpus) == 2
|
||||
assert {r["context_id"] for r in corpus} == {"ctx_a", "ctx_b"}
|
||||
|
||||
|
||||
class TestFrames:
|
||||
def test_parse_wiki_links_plain(self) -> None:
|
||||
raw = (
|
||||
"['https://en.wikipedia.org/wiki/James_Buchanan', "
|
||||
"'https://en.wikipedia.org/wiki/Harriet_Lane']"
|
||||
)
|
||||
assert parse_wiki_links(raw) == [
|
||||
"https://en.wikipedia.org/wiki/James_Buchanan",
|
||||
"https://en.wikipedia.org/wiki/Harriet_Lane",
|
||||
]
|
||||
|
||||
def test_parse_wiki_links_splits_comma_joined_urls(self) -> None:
|
||||
raw = (
|
||||
"['https://en.wikipedia.org/wiki/Tim_Salmon, "
|
||||
"https://en.wikipedia.org/wiki/Troy_Glaus, ']"
|
||||
)
|
||||
assert parse_wiki_links(raw) == [
|
||||
"https://en.wikipedia.org/wiki/Tim_Salmon",
|
||||
"https://en.wikipedia.org/wiki/Troy_Glaus",
|
||||
]
|
||||
|
||||
def test_parse_wiki_links_keeps_commas_inside_titles(self) -> None:
|
||||
raw = (
|
||||
"['https://en.wikipedia.org/wiki/Lincoln,_Nebraska', "
|
||||
"'https://en.wikipedia.org/wiki/Key_West#:~:text=The%20southernmost,"
|
||||
"apart%20at%20their%20closest%20points.']"
|
||||
)
|
||||
assert parse_wiki_links(raw) == [
|
||||
"https://en.wikipedia.org/wiki/Lincoln,_Nebraska",
|
||||
"https://en.wikipedia.org/wiki/Key_West#:~:text=The%20southernmost,"
|
||||
"apart%20at%20their%20closest%20points.",
|
||||
]
|
||||
|
||||
def test_parse_wiki_links_strips_trailing_annotation(self) -> None:
|
||||
raw = "['https://en.wikipedia.org/wiki/Pok%C3%A9mon (NOT REQUIRED, BUT HELPFUL) ']"
|
||||
assert parse_wiki_links(raw) == ["https://en.wikipedia.org/wiki/Pok%C3%A9mon"]
|
||||
|
||||
def test_normalize_strips_fragment_and_mobile_host(self) -> None:
|
||||
assert (
|
||||
normalize_wiki_url("https://en.m.wikipedia.org/wiki/World_War_I#Aftermath")
|
||||
== "https://en.wikipedia.org/wiki/World_War_I"
|
||||
)
|
||||
|
||||
def test_normalize_decodes_and_canonicalizes_title(self) -> None:
|
||||
assert (
|
||||
normalize_wiki_url("https://en.wikipedia.org/wiki/pain %26 Gain")
|
||||
== "https://en.wikipedia.org/wiki/Pain_&_Gain"
|
||||
)
|
||||
|
||||
def test_normalize_schemeless(self) -> None:
|
||||
assert (
|
||||
normalize_wiki_url("en.wikipedia.org/wiki/Grazia_Deledda")
|
||||
== "https://en.wikipedia.org/wiki/Grazia_Deledda"
|
||||
)
|
||||
|
||||
def test_normalize_index_php_title(self) -> None:
|
||||
assert (
|
||||
normalize_wiki_url(
|
||||
"https://en.wikipedia.org/w/index.php?title=Bronco&redirect=no"
|
||||
)
|
||||
== "https://en.wikipedia.org/wiki/Bronco"
|
||||
)
|
||||
|
||||
def test_normalize_search_url(self) -> None:
|
||||
url = (
|
||||
"https://en.wikipedia.org/w/index.php?search=Polytrichum+piliferum"
|
||||
"&title=Special:Search&profile=advanced&fulltext=1&ns0=1"
|
||||
)
|
||||
assert (
|
||||
normalize_wiki_url(url)
|
||||
== "https://en.wikipedia.org/wiki/Polytrichum_piliferum"
|
||||
)
|
||||
|
||||
def test_normalize_shortlink_passthrough(self) -> None:
|
||||
assert normalize_wiki_url("https://w.wiki/ASFv") == "https://w.wiki/ASFv"
|
||||
|
||||
def test_normalize_rejects_non_article(self) -> None:
|
||||
assert normalize_wiki_url("") is None
|
||||
assert normalize_wiki_url("https://en.wikipedia.org/foo") is None
|
||||
|
||||
def test_parse_revid(self) -> None:
|
||||
assert parse_revid('W/"1364811104/52cd04f4-864c-11f1"') == "1364811104"
|
||||
assert parse_revid('"1234/abc"') == "1234"
|
||||
assert parse_revid(None) is None
|
||||
assert parse_revid("") is None
|
||||
|
||||
def test_strip_navigation_removes_navboxes_keeps_infobox(self) -> None:
|
||||
html = (
|
||||
"<html><body>"
|
||||
'<table class="infobox"><tbody><tr><td>Born April 23, 1791</td></tr></tbody></table>'
|
||||
"<p>Some prose.</p>"
|
||||
'<div role="navigation"><table><tbody><tr><td>v t e Presidents</td></tr></tbody></table></div>'
|
||||
"</body></html>"
|
||||
)
|
||||
stripped = strip_navigation(html)
|
||||
assert "Born April 23, 1791" in stripped
|
||||
assert "Some prose." in stripped
|
||||
assert "v t e Presidents" not in stripped
|
||||
|
||||
def test_map_retrieval_normalizes_and_dedupes(self) -> None:
|
||||
row = {
|
||||
"Prompt": "Who was the 15th president?",
|
||||
"wiki_links": (
|
||||
"['https://en.wikipedia.org/wiki/James_Buchanan#Presidency', "
|
||||
"'https://en.m.wikipedia.org/wiki/James_Buchanan', "
|
||||
"'https://en.wikipedia.org/wiki/Harriet_Lane']"
|
||||
),
|
||||
}
|
||||
sample = map_frames_retrieval(row)
|
||||
assert sample is not None
|
||||
assert sample.question == "Who was the 15th president?"
|
||||
assert sample.expected_uris == (
|
||||
"https://en.wikipedia.org/wiki/James_Buchanan",
|
||||
"https://en.wikipedia.org/wiki/Harriet_Lane",
|
||||
)
|
||||
|
||||
def test_map_retrieval_empty_links(self) -> None:
|
||||
assert map_frames_retrieval({"Prompt": "Q", "wiki_links": "[]"}) is None
|
||||
|
||||
def test_map_document_html_strips_navigation(self, tmp_path: Path) -> None:
|
||||
page = tmp_path / "article.html"
|
||||
page.write_text(
|
||||
"<html><body><p>Buchanan was a president.</p>"
|
||||
'<div role="navigation">v t e spam</div></body></html>'
|
||||
)
|
||||
row = {
|
||||
"uri": "https://en.wikipedia.org/wiki/James_Buchanan",
|
||||
"title": "James Buchanan",
|
||||
"path": str(page),
|
||||
"format": "html",
|
||||
"revid": "1364811104",
|
||||
"fetched_at": "2026-07-23",
|
||||
}
|
||||
payload = map_frames_document(row)
|
||||
assert payload.uri == "https://en.wikipedia.org/wiki/James_Buchanan"
|
||||
assert payload.title == "James Buchanan"
|
||||
assert payload.format == "html"
|
||||
assert "Buchanan was a president." in (payload.content or "")
|
||||
assert "v t e spam" not in (payload.content or "")
|
||||
assert payload.metadata == {
|
||||
"revid": "1364811104",
|
||||
"fetched_at": "2026-07-23",
|
||||
}
|
||||
|
||||
def test_map_document_markdown_passthrough(self, tmp_path: Path) -> None:
|
||||
page = tmp_path / "category.md"
|
||||
page.write_text(
|
||||
"Pages in Category:Summer Olympics in London:\n- 1908 Summer Olympics\n"
|
||||
)
|
||||
row = {
|
||||
"uri": "https://en.wikipedia.org/wiki/Category:Summer_Olympics_in_London",
|
||||
"title": "Category:Summer Olympics in London",
|
||||
"path": str(page),
|
||||
"format": "md",
|
||||
"revid": None,
|
||||
"fetched_at": "2026-07-23",
|
||||
}
|
||||
payload = map_frames_document(row)
|
||||
assert payload.format == "md"
|
||||
assert "1908 Summer Olympics" in (payload.content or "")
|
||||
assert payload.metadata == {"fetched_at": "2026-07-23"}
|
||||
|
||||
def test_build_case(self) -> None:
|
||||
row = {
|
||||
"id": "7",
|
||||
"Prompt": "Who was the 15th president?",
|
||||
"Answer": "James Buchanan",
|
||||
"reasoning_types": "Multiple constraints | Temporal reasoning",
|
||||
}
|
||||
case = build_frames_case(3, row)
|
||||
assert case.name == "3_7"
|
||||
assert case.inputs == "Who was the 15th president?"
|
||||
assert case.expected_output == "James Buchanan"
|
||||
assert case.metadata == {
|
||||
"question_id": "7",
|
||||
"reasoning_types": "Multiple constraints | Temporal reasoning",
|
||||
"case_index": "3",
|
||||
}
|
||||
|
||||
def test_fetch_article_cache_hit_needs_no_network(self, tmp_path: Path) -> None:
|
||||
uri = "https://en.wikipedia.org/wiki/James_Buchanan"
|
||||
from urllib.parse import quote
|
||||
|
||||
base = quote(uri, safe="")
|
||||
(tmp_path / f"{base}.html").write_text("<html><body>cached</body></html>")
|
||||
(tmp_path / f"{base}.json").write_text(
|
||||
'{"uri": "https://en.wikipedia.org/wiki/James_Buchanan",'
|
||||
' "title": "James Buchanan", "format": "html",'
|
||||
' "revid": "123", "fetched_at": "2026-07-23"}'
|
||||
)
|
||||
row = fetch_article(uri, tmp_path, client=None)
|
||||
assert row is not None
|
||||
assert row["uri"] == uri
|
||||
assert row["revid"] == "123"
|
||||
assert row["format"] == "html"
|
||||
assert Path(row["path"]).read_text().startswith("<html>")
|
||||
|
||||
def test_fetch_article_category_synthesizes_members(self, tmp_path: Path) -> None:
|
||||
class StubResponse:
|
||||
def __init__(self, payload: dict) -> None:
|
||||
self._payload = payload
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
pass
|
||||
|
||||
def json(self) -> dict:
|
||||
return self._payload
|
||||
|
||||
class StubClient:
|
||||
def get(self, url: str, params: dict | None = None) -> StubResponse:
|
||||
assert params is not None
|
||||
assert params["list"] == "categorymembers"
|
||||
return StubResponse(
|
||||
{
|
||||
"query": {
|
||||
"categorymembers": [
|
||||
{"title": "1908 Summer Olympics"},
|
||||
{"title": "2012 Summer Olympics"},
|
||||
]
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
uri = "https://en.wikipedia.org/wiki/Category:Summer_Olympics_in_London"
|
||||
row = fetch_article(
|
||||
uri,
|
||||
tmp_path,
|
||||
client=StubClient(), # ty: ignore[invalid-argument-type]
|
||||
)
|
||||
assert row is not None
|
||||
assert row["format"] == "md"
|
||||
content = Path(row["path"]).read_text()
|
||||
assert "1908 Summer Olympics" in content
|
||||
assert "2012 Summer Olympics" in content
|
||||
|
||||
def test_fetch_article_retries_transient_failures(
|
||||
self, tmp_path: Path, monkeypatch
|
||||
) -> None:
|
||||
sleeps: list[float] = []
|
||||
monkeypatch.setattr(
|
||||
"evaluations.datasets.frames.time.sleep", lambda s: sleeps.append(s)
|
||||
)
|
||||
|
||||
class FlakyResponse:
|
||||
text = "<html><body>ok</body></html>"
|
||||
headers = {"etag": 'W/"42/uuid"'}
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
pass
|
||||
|
||||
class FlakyClient:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
def get(self, url: str, params: dict | None = None) -> FlakyResponse:
|
||||
self.calls += 1
|
||||
if self.calls < 3:
|
||||
raise OSError("connection reset")
|
||||
return FlakyResponse()
|
||||
|
||||
client = FlakyClient()
|
||||
row = fetch_article(
|
||||
"https://en.wikipedia.org/wiki/Capybara",
|
||||
tmp_path,
|
||||
client=client, # ty: ignore[invalid-argument-type]
|
||||
)
|
||||
assert row is not None
|
||||
assert row["revid"] == "42"
|
||||
assert client.calls == 3
|
||||
# One throttle sleep before fetching plus one backoff per failure.
|
||||
assert len(sleeps) == 3
|
||||
|
||||
def test_fetch_article_honors_retry_after_on_rate_limit(
|
||||
self, tmp_path: Path, monkeypatch
|
||||
) -> None:
|
||||
import httpx
|
||||
|
||||
sleeps: list[float] = []
|
||||
monkeypatch.setattr(
|
||||
"evaluations.datasets.frames.time.sleep", lambda s: sleeps.append(s)
|
||||
)
|
||||
request = httpx.Request("GET", "https://en.wikipedia.org/x")
|
||||
|
||||
class OkResponse:
|
||||
text = "<html><body>ok</body></html>"
|
||||
headers = {"etag": 'W/"42/uuid"'}
|
||||
|
||||
def raise_for_status(self) -> None:
|
||||
pass
|
||||
|
||||
class RateLimitedClient:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
def get(self, url: str, params: dict | None = None) -> OkResponse:
|
||||
self.calls += 1
|
||||
if self.calls == 1:
|
||||
raise httpx.HTTPStatusError(
|
||||
"429 too many requests",
|
||||
request=request,
|
||||
response=httpx.Response(
|
||||
429, headers={"retry-after": "13"}, request=request
|
||||
),
|
||||
)
|
||||
return OkResponse()
|
||||
|
||||
row = fetch_article(
|
||||
"https://en.wikipedia.org/wiki/Capybara",
|
||||
tmp_path,
|
||||
client=RateLimitedClient(), # ty: ignore[invalid-argument-type]
|
||||
)
|
||||
assert row is not None
|
||||
assert 13.0 in sleeps
|
||||
|
||||
def test_fetch_article_gives_up_after_max_attempts(
|
||||
self, tmp_path: Path, monkeypatch
|
||||
) -> None:
|
||||
monkeypatch.setattr("evaluations.datasets.frames.time.sleep", lambda s: None)
|
||||
|
||||
class DeadClient:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
def get(self, url: str, params: dict | None = None):
|
||||
self.calls += 1
|
||||
raise OSError("connection reset")
|
||||
|
||||
client = DeadClient()
|
||||
row = fetch_article(
|
||||
"https://en.wikipedia.org/wiki/Capybara",
|
||||
tmp_path,
|
||||
client=client, # ty: ignore[invalid-argument-type]
|
||||
)
|
||||
assert row is None
|
||||
assert client.calls == FETCH_ATTEMPTS
|
||||
|
||||
def test_load_corpus_raises_on_partial_fetch(self, monkeypatch) -> None:
|
||||
import evaluations.datasets.frames as frames
|
||||
|
||||
monkeypatch.setattr(frames, "_cached_corpus", None)
|
||||
monkeypatch.setattr(
|
||||
frames,
|
||||
"load_frames_questions",
|
||||
lambda: [
|
||||
{
|
||||
"wiki_links": "['https://en.wikipedia.org/wiki/A', "
|
||||
"'https://en.wikipedia.org/wiki/B']"
|
||||
}
|
||||
],
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
frames, "fetch_article", lambda uri, cache_dir, client: None
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="0/2"):
|
||||
frames.load_frames_corpus()
|
||||
|
||||
def test_question_with_deleted_article_is_excluded(self) -> None:
|
||||
gone = {
|
||||
"wiki_links": "['https://en.wikipedia.org/wiki/Jack_Vance_(tennis)', "
|
||||
"'https://en.wikipedia.org/wiki/Capybara']"
|
||||
}
|
||||
kept = {"wiki_links": "['https://en.wikipedia.org/wiki/Capybara']"}
|
||||
assert question_is_answerable(gone) is False
|
||||
assert question_is_answerable(kept) is True
|
||||
|
||||
def test_questions_carry_stable_ids(self, monkeypatch) -> None:
|
||||
import evaluations.datasets.frames as frames
|
||||
from datasets import Dataset
|
||||
|
||||
rows = Dataset.from_list(
|
||||
[
|
||||
{
|
||||
"Unnamed: 0": 7,
|
||||
"wiki_links": "['https://en.wikipedia.org/wiki/Capybara']",
|
||||
},
|
||||
{
|
||||
"Unnamed: 0": 8,
|
||||
"wiki_links": "['https://en.wikipedia.org/wiki/Jack_Vance_(tennis)']",
|
||||
},
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(frames, "load_frames_test", lambda: rows)
|
||||
questions = frames.load_frames_questions()
|
||||
assert [row["id"] for row in questions] == ["7"]
|
||||
|
|
|
|||
|
|
@ -1,9 +1,14 @@
|
|||
from unittest.mock import MagicMock
|
||||
import math
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic_evals.evaluators import EvaluatorContext
|
||||
|
||||
from evaluations.evaluators import REFUSAL_RUBRIC, RefusalJudge
|
||||
|
||||
from evaluations.evaluators.map import MAPEvaluator
|
||||
from evaluations.evaluators.number_match import NumberMatchEvaluator
|
||||
from evaluations.evaluators.retrieval import NDCGEvaluator, RecallEvaluator
|
||||
|
||||
|
||||
class TestMAPEvaluator:
|
||||
|
|
@ -57,6 +62,189 @@ class TestMAPEvaluator:
|
|||
assert self.evaluator.evaluate(ctx) == 0.0
|
||||
|
||||
|
||||
def _retrieval_ctx(relevant_uris: list[str], retrieved_uris: list[str]) -> MagicMock:
|
||||
ctx = MagicMock()
|
||||
ctx.metadata = {"relevant_uris": relevant_uris}
|
||||
ctx.output = retrieved_uris
|
||||
return ctx
|
||||
|
||||
|
||||
class TestRecallEvaluator:
|
||||
def test_evaluation_name_includes_k(self) -> None:
|
||||
assert RecallEvaluator(k=5).get_default_evaluation_name() == "recall_5"
|
||||
assert RecallEvaluator(k=10).get_default_evaluation_name() == "recall_10"
|
||||
|
||||
def test_all_relevant_within_k(self) -> None:
|
||||
ctx = _retrieval_ctx(["a", "b"], ["a", "b", "c"])
|
||||
assert RecallEvaluator(k=5).evaluate(ctx) == 1.0
|
||||
|
||||
def test_partial_recall(self) -> None:
|
||||
ctx = _retrieval_ctx(["a", "b"], ["a", "c", "d"])
|
||||
assert RecallEvaluator(k=3).evaluate(ctx) == 0.5
|
||||
|
||||
def test_relevant_beyond_k_not_counted(self) -> None:
|
||||
ctx = _retrieval_ctx(["a"], ["b", "c", "d", "e", "f", "a"])
|
||||
assert RecallEvaluator(k=5).evaluate(ctx) == 0.0
|
||||
assert RecallEvaluator(k=10).evaluate(ctx) == 1.0
|
||||
|
||||
def test_empty_relevant(self) -> None:
|
||||
ctx = _retrieval_ctx([], ["a"])
|
||||
assert RecallEvaluator(k=5).evaluate(ctx) == 0.0
|
||||
|
||||
def test_none_metadata(self) -> None:
|
||||
ctx = MagicMock()
|
||||
ctx.metadata = None
|
||||
ctx.output = ["a"]
|
||||
assert RecallEvaluator(k=5).evaluate(ctx) == 0.0
|
||||
|
||||
|
||||
class TestNDCGEvaluator:
|
||||
def test_evaluation_name_includes_k(self) -> None:
|
||||
assert NDCGEvaluator(k=5).get_default_evaluation_name() == "ndcg_5"
|
||||
|
||||
def test_perfect_ranking(self) -> None:
|
||||
ctx = _retrieval_ctx(["a", "b"], ["a", "b", "c"])
|
||||
assert NDCGEvaluator(k=5).evaluate(ctx) == pytest.approx(1.0)
|
||||
|
||||
def test_single_relevant_at_rank_two(self) -> None:
|
||||
# DCG = 1/log2(3); IDCG = 1/log2(2) = 1
|
||||
ctx = _retrieval_ctx(["a"], ["b", "a"])
|
||||
expected = 1 / math.log2(3)
|
||||
assert NDCGEvaluator(k=5).evaluate(ctx) == pytest.approx(expected)
|
||||
|
||||
def test_two_relevant_with_gap(self) -> None:
|
||||
# Relevant at ranks 1 and 3: DCG = 1 + 1/log2(4) = 1.5
|
||||
# IDCG = 1 + 1/log2(3)
|
||||
ctx = _retrieval_ctx(["a", "b"], ["a", "c", "b"])
|
||||
expected = 1.5 / (1 + 1 / math.log2(3))
|
||||
assert NDCGEvaluator(k=3).evaluate(ctx) == pytest.approx(expected)
|
||||
|
||||
def test_relevant_beyond_k_not_counted(self) -> None:
|
||||
ctx = _retrieval_ctx(["a"], ["b", "c", "d", "e", "f", "a"])
|
||||
assert NDCGEvaluator(k=5).evaluate(ctx) == 0.0
|
||||
|
||||
def test_ideal_dcg_capped_at_k(self) -> None:
|
||||
# 3 relevant but k=2: IDCG uses only the top-2 ideal ranks, so a
|
||||
# retrieval with both top-2 slots relevant scores 1.0.
|
||||
ctx = _retrieval_ctx(["a", "b", "c"], ["a", "b"])
|
||||
assert NDCGEvaluator(k=2).evaluate(ctx) == pytest.approx(1.0)
|
||||
|
||||
def test_empty_relevant(self) -> None:
|
||||
ctx = _retrieval_ctx([], ["a"])
|
||||
assert NDCGEvaluator(k=5).evaluate(ctx) == 0.0
|
||||
|
||||
def test_none_metadata(self) -> None:
|
||||
ctx = MagicMock()
|
||||
ctx.metadata = None
|
||||
ctx.output = ["a"]
|
||||
assert NDCGEvaluator(k=5).evaluate(ctx) == 0.0
|
||||
|
||||
|
||||
def _evaluator_ctx(inputs: object, metadata: dict | None = None) -> EvaluatorContext:
|
||||
return EvaluatorContext(
|
||||
name="case",
|
||||
inputs=inputs,
|
||||
metadata=metadata,
|
||||
expected_output="expected",
|
||||
output="answer",
|
||||
duration=0.0,
|
||||
_span_tree=MagicMock(),
|
||||
attributes={},
|
||||
metrics={},
|
||||
)
|
||||
|
||||
|
||||
class TestTranscriptLLMJudge:
|
||||
@pytest.mark.asyncio
|
||||
async def test_conversation_inputs_judged_as_transcript(self) -> None:
|
||||
from evaluations.config import ConversationInput, Turn
|
||||
from evaluations.evaluators import TranscriptLLMJudge
|
||||
|
||||
judge = TranscriptLLMJudge(
|
||||
rubric="rubric",
|
||||
include_input=True,
|
||||
include_expected_output=True,
|
||||
model="test",
|
||||
)
|
||||
conversation = ConversationInput(
|
||||
turns=[
|
||||
Turn(speaker="user", text="q1"),
|
||||
Turn(speaker="agent", text="a1"),
|
||||
Turn(speaker="user", text="q2"),
|
||||
]
|
||||
)
|
||||
grading = MagicMock(score=None, pass_=True, reason="ok")
|
||||
with patch(
|
||||
"pydantic_evals.evaluators.llm_as_a_judge.judge_input_output_expected",
|
||||
new_callable=AsyncMock,
|
||||
return_value=grading,
|
||||
) as judge_call:
|
||||
await judge.evaluate(_evaluator_ctx(conversation))
|
||||
|
||||
assert judge_call.await_args is not None
|
||||
assert judge_call.await_args.args[0] == "user: q1\nagent: a1\nuser: q2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_string_inputs_pass_through(self) -> None:
|
||||
from evaluations.evaluators import TranscriptLLMJudge
|
||||
|
||||
judge = TranscriptLLMJudge(
|
||||
rubric="rubric",
|
||||
include_input=True,
|
||||
include_expected_output=True,
|
||||
model="test",
|
||||
)
|
||||
grading = MagicMock(score=None, pass_=True, reason="ok")
|
||||
with patch(
|
||||
"pydantic_evals.evaluators.llm_as_a_judge.judge_input_output_expected",
|
||||
new_callable=AsyncMock,
|
||||
return_value=grading,
|
||||
) as judge_call:
|
||||
await judge.evaluate(_evaluator_ctx("plain question"))
|
||||
|
||||
assert judge_call.await_args is not None
|
||||
assert judge_call.await_args.args[0] == "plain question"
|
||||
|
||||
|
||||
class TestRefusalJudge:
|
||||
def _judge(self) -> RefusalJudge:
|
||||
return RefusalJudge(
|
||||
rubric=REFUSAL_RUBRIC,
|
||||
model="test",
|
||||
assertion={"evaluation_name": "refused", "include_reason": False},
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("label", ["ANSWERABLE", "UNANSWERABLE"])
|
||||
async def test_judges_eligible_labels(self, label: str) -> None:
|
||||
grading = MagicMock(score=None, pass_=True, reason=None)
|
||||
with patch(
|
||||
"pydantic_evals.evaluators.llm_as_a_judge.judge_output",
|
||||
new_callable=AsyncMock,
|
||||
return_value=grading,
|
||||
) as judge_call:
|
||||
result = await self._judge().evaluate(
|
||||
_evaluator_ctx("q", metadata={"answerability": label})
|
||||
)
|
||||
|
||||
judge_call.assert_awaited_once()
|
||||
assert result == {"refused": True}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"metadata", [{"answerability": "PARTIAL"}, {"answerability": None}, {}, None]
|
||||
)
|
||||
async def test_ineligible_turns_skip_the_judge(self, metadata) -> None:
|
||||
with patch(
|
||||
"pydantic_evals.evaluators.llm_as_a_judge.judge_output",
|
||||
new_callable=AsyncMock,
|
||||
) as judge_call:
|
||||
result = await self._judge().evaluate(_evaluator_ctx("q", metadata))
|
||||
|
||||
judge_call.assert_not_awaited()
|
||||
assert result == {}
|
||||
|
||||
|
||||
class TestNumberMatchEvaluator:
|
||||
def setup_method(self) -> None:
|
||||
self.evaluator = NumberMatchEvaluator()
|
||||
|
|
|
|||
278
evaluations/tests/test_mtrag.py
Normal file
278
evaluations/tests/test_mtrag.py
Normal file
|
|
@ -0,0 +1,278 @@
|
|||
import pytest
|
||||
|
||||
from evaluations.config import ConversationInput
|
||||
from evaluations.datasets import DATASETS
|
||||
from evaluations.datasets.mtrag import (
|
||||
MTRAG_CLAPNQ_LIVE_SPEC,
|
||||
MTRAG_CLAPNQ_REWRITE_SPEC,
|
||||
MTRAG_CLAPNQ_SPEC,
|
||||
_group_conversations,
|
||||
_join_queries_qrels,
|
||||
_parse_qrels,
|
||||
_task_to_record,
|
||||
_validate_qrels_resolve,
|
||||
build_mtrag_case,
|
||||
build_mtrag_live_case,
|
||||
map_mtrag_document,
|
||||
map_mtrag_retrieval,
|
||||
)
|
||||
from evaluations.evaluators import (
|
||||
CitationMAPEvaluator,
|
||||
MAPEvaluator,
|
||||
NDCGEvaluator,
|
||||
RecallEvaluator,
|
||||
)
|
||||
|
||||
GENERATION_TASK = {
|
||||
"task_id": "conv1<::>2",
|
||||
"conversation_id": "conv1",
|
||||
"turn": "2",
|
||||
"Collection": "mt-rag-clapnq-elser-512-100-20240503",
|
||||
"Answerability": ["ANSWERABLE"],
|
||||
"Multi-Turn": ["Follow-up"],
|
||||
"Question Type": ["Factoid"],
|
||||
"input": [
|
||||
{"speaker": "user", "text": "q1", "metadata": {}},
|
||||
{"speaker": "agent", "text": "a1", "metadata": {}},
|
||||
{"speaker": "user", "text": "q2", "metadata": {}},
|
||||
],
|
||||
"targets": [{"text": "reference answer"}],
|
||||
"contexts": [{"document_id": "retrieved-not-gold"}],
|
||||
}
|
||||
|
||||
|
||||
class TestDocumentMapper:
|
||||
def test_maps_passage_to_payload(self) -> None:
|
||||
payload = map_mtrag_document(
|
||||
{"_id": "837799097_6931-7548-0-617", "title": "T", "text": "body"}
|
||||
)
|
||||
assert payload.uri == "837799097_6931-7548-0-617"
|
||||
assert payload.title == "T"
|
||||
assert payload.content == "body"
|
||||
|
||||
|
||||
class TestQrels:
|
||||
QRELS_TSV = (
|
||||
"query-id\tcorpus-id\tscore\n"
|
||||
"conv1<::>2\tdoc1_0-10-0-10\t1\n"
|
||||
"conv1<::>2\tdoc2_5-20-0-15\t1\n"
|
||||
"conv2<::>1\tdoc3_0-9-0-9\t1\n"
|
||||
)
|
||||
|
||||
def test_parse_groups_by_query_preserving_order(self) -> None:
|
||||
qrels = _parse_qrels(self.QRELS_TSV.splitlines())
|
||||
assert qrels == {
|
||||
"conv1<::>2": ["doc1_0-10-0-10", "doc2_5-20-0-15"],
|
||||
"conv2<::>1": ["doc3_0-9-0-9"],
|
||||
}
|
||||
|
||||
def test_join_builds_records(self) -> None:
|
||||
qrels = _parse_qrels(self.QRELS_TSV.splitlines())
|
||||
queries = [
|
||||
{"_id": "conv1<::>2", "text": "q one"},
|
||||
{"_id": "conv2<::>1", "text": "q two"},
|
||||
]
|
||||
records = _join_queries_qrels(queries, qrels)
|
||||
assert records == [
|
||||
{
|
||||
"query_id": "conv1<::>2",
|
||||
"question": "q one",
|
||||
"expected_uris": ["doc1_0-10-0-10", "doc2_5-20-0-15"],
|
||||
},
|
||||
{
|
||||
"query_id": "conv2<::>1",
|
||||
"question": "q two",
|
||||
"expected_uris": ["doc3_0-9-0-9"],
|
||||
},
|
||||
]
|
||||
|
||||
def test_join_raises_on_query_without_qrels(self) -> None:
|
||||
with pytest.raises(ValueError, match="no qrels"):
|
||||
_join_queries_qrels([{"_id": "missing<::>1", "text": "q"}], {})
|
||||
|
||||
def test_validation_passes_when_all_resolve(self) -> None:
|
||||
qrels = {"q1": ["a", "b"]}
|
||||
_validate_qrels_resolve({"a", "b", "c"}, qrels)
|
||||
|
||||
def test_validation_raises_on_unresolved_id(self) -> None:
|
||||
qrels = {"q1": ["a", "ghost"]}
|
||||
with pytest.raises(ValueError, match="ghost"):
|
||||
_validate_qrels_resolve({"a"}, qrels)
|
||||
|
||||
|
||||
class TestRetrievalMapper:
|
||||
def test_maps_joined_record(self) -> None:
|
||||
sample = map_mtrag_retrieval(
|
||||
{
|
||||
"query_id": "conv1<::>2",
|
||||
"question": "who?",
|
||||
"expected_uris": ["u1", "u2"],
|
||||
}
|
||||
)
|
||||
assert sample is not None
|
||||
assert sample.question == "who?"
|
||||
assert sample.expected_uris == ("u1", "u2")
|
||||
|
||||
|
||||
class TestSpecs:
|
||||
def test_registered(self) -> None:
|
||||
assert DATASETS["mtrag_clapnq"] is MTRAG_CLAPNQ_SPEC
|
||||
assert DATASETS["mtrag_clapnq_rewrite"] is MTRAG_CLAPNQ_REWRITE_SPEC
|
||||
|
||||
def test_variants_share_db(self) -> None:
|
||||
assert MTRAG_CLAPNQ_SPEC.db_filename == MTRAG_CLAPNQ_REWRITE_SPEC.db_filename
|
||||
|
||||
def test_retrieval_configuration(self) -> None:
|
||||
for spec in (MTRAG_CLAPNQ_SPEC, MTRAG_CLAPNQ_REWRITE_SPEC):
|
||||
assert spec.retrieval_limit == 10
|
||||
assert spec.ingest_batch_size == 512
|
||||
assert spec.retrieval_evaluators is not None
|
||||
kinds = {
|
||||
(type(e), getattr(e, "k", None)) for e in spec.retrieval_evaluators
|
||||
}
|
||||
assert kinds == {
|
||||
(RecallEvaluator, 5),
|
||||
(RecallEvaluator, 10),
|
||||
(NDCGEvaluator, 5),
|
||||
(NDCGEvaluator, 10),
|
||||
(MAPEvaluator, None),
|
||||
}
|
||||
assert isinstance(spec.citation_evaluator, CitationMAPEvaluator)
|
||||
|
||||
|
||||
class TestGenerationTasks:
|
||||
def test_task_to_record(self) -> None:
|
||||
record = _task_to_record(GENERATION_TASK, {"conv1<::>2": ["p1", "p2"]})
|
||||
assert record == {
|
||||
"id": "conv1<::>2",
|
||||
"turn": "2",
|
||||
"turns": [
|
||||
{"speaker": "user", "text": "q1"},
|
||||
{"speaker": "agent", "text": "a1"},
|
||||
{"speaker": "user", "text": "q2"},
|
||||
],
|
||||
"answer": "reference answer",
|
||||
"answerability": "ANSWERABLE",
|
||||
"multi_turn_type": "Follow-up",
|
||||
"question_type": ["Factoid"],
|
||||
"relevant_uris": ["p1", "p2"],
|
||||
}
|
||||
|
||||
def test_task_without_qrels_has_no_relevant_uris(self) -> None:
|
||||
record = _task_to_record(GENERATION_TASK, {})
|
||||
assert record is not None
|
||||
assert record["relevant_uris"] is None
|
||||
|
||||
def test_other_collections_excluded(self) -> None:
|
||||
task = {**GENERATION_TASK, "Collection": "mt-rag-govt-elser-512-100-20240611"}
|
||||
assert _task_to_record(task, {}) is None
|
||||
|
||||
def test_build_case_conversation_and_metadata(self) -> None:
|
||||
record = _task_to_record(GENERATION_TASK, {"conv1<::>2": ["p1"]})
|
||||
assert record is not None
|
||||
case = build_mtrag_case(3, record)
|
||||
|
||||
assert isinstance(case.inputs, ConversationInput)
|
||||
assert case.inputs.question == "q2"
|
||||
assert [t.speaker for t in case.inputs.turns] == ["user", "agent", "user"]
|
||||
assert case.expected_output == "reference answer"
|
||||
assert case.metadata == {
|
||||
"task_id": "conv1<::>2",
|
||||
"turn": "2",
|
||||
"answerability": "ANSWERABLE",
|
||||
"multi_turn_type": "Follow-up",
|
||||
"question_type": ["Factoid"],
|
||||
"relevant_uris": ["p1"],
|
||||
}
|
||||
|
||||
def test_build_case_omits_relevant_uris_when_absent(self) -> None:
|
||||
record = _task_to_record(
|
||||
{**GENERATION_TASK, "Answerability": ["UNANSWERABLE"]}, {}
|
||||
)
|
||||
assert record is not None
|
||||
case = build_mtrag_case(1, record)
|
||||
|
||||
assert case.metadata is not None
|
||||
assert "relevant_uris" not in case.metadata
|
||||
assert case.metadata["answerability"] == "UNANSWERABLE"
|
||||
|
||||
|
||||
class TestLiveConversations:
|
||||
def _records(self) -> list[dict]:
|
||||
turn1 = _task_to_record(
|
||||
{
|
||||
**GENERATION_TASK,
|
||||
"task_id": "conv1<::>1",
|
||||
"turn": "1",
|
||||
"input": [{"speaker": "user", "text": "q1", "metadata": {}}],
|
||||
"targets": [{"text": "r1"}],
|
||||
},
|
||||
{"conv1<::>1": ["p1"]},
|
||||
)
|
||||
turn2 = _task_to_record(GENERATION_TASK, {"conv1<::>2": ["p2", "p3"]})
|
||||
other = _task_to_record(
|
||||
{
|
||||
**GENERATION_TASK,
|
||||
"task_id": "conv2<::>1",
|
||||
"turn": "1",
|
||||
"input": [{"speaker": "user", "text": "other q", "metadata": {}}],
|
||||
"targets": [{"text": "other r"}],
|
||||
"Answerability": ["UNANSWERABLE"],
|
||||
},
|
||||
{},
|
||||
)
|
||||
assert turn1 and turn2 and other
|
||||
# turn 2 first: grouping must sort turns numerically within a conversation
|
||||
return [turn2, turn1, other]
|
||||
|
||||
def test_grouping_sorts_turns_within_conversations(self) -> None:
|
||||
conversations = _group_conversations(self._records())
|
||||
assert [c["id"] for c in conversations] == ["conv1", "conv2"]
|
||||
conv1 = conversations[0]
|
||||
assert [t["question"] for t in conv1["turns"]] == ["q1", "q2"]
|
||||
assert [t["reference"] for t in conv1["turns"]] == ["r1", "reference answer"]
|
||||
assert conv1["turns"][1]["relevant_uris"] == ["p2", "p3"]
|
||||
|
||||
def test_build_live_case(self) -> None:
|
||||
conversations = _group_conversations(self._records())
|
||||
case = build_mtrag_live_case(1, conversations[0])
|
||||
|
||||
assert case.inputs == ["q1", "q2"]
|
||||
assert case.metadata is not None
|
||||
assert case.metadata["conversation_id"] == "conv1"
|
||||
turns = case.metadata["turns"]
|
||||
assert turns[0] == {
|
||||
"task_id": "conv1<::>1",
|
||||
"turn": "1",
|
||||
"reference": "r1",
|
||||
"answerability": "ANSWERABLE",
|
||||
"multi_turn_type": "Follow-up",
|
||||
"question_type": ["Factoid"],
|
||||
"relevant_uris": ["p1"],
|
||||
}
|
||||
other_case = build_mtrag_live_case(2, conversations[1])
|
||||
assert other_case.metadata is not None
|
||||
assert other_case.metadata["turns"][0]["relevant_uris"] == []
|
||||
|
||||
def test_live_spec(self) -> None:
|
||||
assert DATASETS["mtrag_clapnq_live"] is MTRAG_CLAPNQ_LIVE_SPEC
|
||||
assert MTRAG_CLAPNQ_LIVE_SPEC.db_filename == MTRAG_CLAPNQ_SPEC.db_filename
|
||||
assert MTRAG_CLAPNQ_LIVE_SPEC.live is True
|
||||
assert MTRAG_CLAPNQ_LIVE_SPEC.retrieval_loader is None
|
||||
assert MTRAG_CLAPNQ_LIVE_SPEC.experiment_metadata == {
|
||||
"mtrag_mode": "live_session",
|
||||
"compaction": True,
|
||||
}
|
||||
assert MTRAG_CLAPNQ_SPEC.experiment_metadata == {"mtrag_mode": "gold_prefix"}
|
||||
|
||||
def test_live_compaction_arms(self) -> None:
|
||||
assert MTRAG_CLAPNQ_LIVE_SPEC.compaction is True
|
||||
uncompacted = DATASETS["mtrag_clapnq_live_uncompacted"]
|
||||
assert uncompacted.compaction is False
|
||||
assert uncompacted.live is True
|
||||
assert uncompacted.db_filename == MTRAG_CLAPNQ_LIVE_SPEC.db_filename
|
||||
assert uncompacted.qa_case_builder is MTRAG_CLAPNQ_LIVE_SPEC.qa_case_builder
|
||||
assert uncompacted.experiment_metadata == {
|
||||
"mtrag_mode": "live_session",
|
||||
"compaction": False,
|
||||
}
|
||||
61
evaluations/tests/test_reference_configs.py
Normal file
61
evaluations/tests/test_reference_configs.py
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from evaluations.datasets import DATASETS
|
||||
from haiku.rag.config import load_yaml_config
|
||||
from haiku.rag.config.models import AppConfig
|
||||
|
||||
CONFIG_DIR = Path(__file__).parent.parent / "configs"
|
||||
|
||||
PINNED_JUDGE_SAMPLING = {
|
||||
"temperature": 0.6,
|
||||
"max_tokens": 16384,
|
||||
"extra_body": {
|
||||
"top_p": 0.95,
|
||||
"top_k": 20,
|
||||
"min_p": 0,
|
||||
"chat_template_kwargs": {"reasoning_effort": "low"},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _config_paths() -> list[Path]:
|
||||
return sorted(CONFIG_DIR.glob("*.yaml"))
|
||||
|
||||
|
||||
def _load(path: Path) -> AppConfig:
|
||||
return AppConfig.model_validate(load_yaml_config(path))
|
||||
|
||||
|
||||
def test_configs_present() -> None:
|
||||
assert _config_paths(), f"no reference configs found in {CONFIG_DIR}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", _config_paths(), ids=lambda p: p.stem)
|
||||
def test_config_validates(path: Path) -> None:
|
||||
_load(path)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", _config_paths(), ids=lambda p: p.stem)
|
||||
def test_filename_names_a_dataset(path: Path) -> None:
|
||||
assert path.stem in DATASETS
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", _config_paths(), ids=lambda p: p.stem)
|
||||
def test_judge_pinned_where_the_judge_runs(path: Path) -> None:
|
||||
"""Datasets without their own qa_evaluator are scored by the LLM judge.
|
||||
|
||||
Those configs must carry the frozen judge settings, so accuracy stays
|
||||
comparable across runs. Datasets that bring a deterministic evaluator
|
||||
never construct a judge, so a judge block there would be dead config.
|
||||
"""
|
||||
judge = _load(path).evaluations.judge
|
||||
if DATASETS[path.stem].qa_evaluator is not None:
|
||||
assert judge is None
|
||||
return
|
||||
|
||||
assert judge is not None
|
||||
assert judge.temperature == PINNED_JUDGE_SAMPLING["temperature"]
|
||||
assert judge.max_tokens == PINNED_JUDGE_SAMPLING["max_tokens"]
|
||||
assert judge.extra_body == PINNED_JUDGE_SAMPLING["extra_body"]
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Reference in a new issue