Compare commits

...

51 commits

Author SHA1 Message Date
Yiorgis Gozadinos
d2af1a2e62
Wait for drained jobs instead of sleeping in the breaker isolation test
Some checks failed
build-docs / build (push) Has been cancelled
Tests / lint (push) Has been cancelled
Tests / lint-frontend (push) Has been cancelled
Tests / URI paths (macos-latest, py3.13) (push) Has been cancelled
Tests / URI paths (macos-latest, py3.14) (push) Has been cancelled
Tests / URI paths (ubuntu-latest, py3.13) (push) Has been cancelled
Tests / URI paths (ubuntu-latest, py3.14) (push) Has been cancelled
Tests / URI paths (windows-latest, py3.13) (push) Has been cancelled
Tests / URI paths (windows-latest, py3.14) (push) Has been cancelled
build-docs / deploy (push) Has been cancelled
Tests / test (push) Has been cancelled
test_breaker_isolates_sources snapshotted job status after a fixed 0.2s
sleep and failed on a loaded runner while two good jobs were still in
flight. Poll for the drained jobs under a 5s deadline.

Assert the paused source's jobs are unattempted: the queued-uri assertion
alone passes with source isolation disabled, since a retried job returns
to QUEUED.
2026-09-07 14:42:03 +03:00
Yiorgis Gozadinos
5cecd7b50e
Merge pull request #600 from ggozad/feat/mcp-revisited
Make the MCP server read-only and multi-database; replace ask_question and analyze with execute_code; ship a Claude Code and Codex plugin with the haiku-rag skill
2026-09-07 05:50:09 -05:00
Yiorgis Gozadinos
4b1ec096c6
One error contract for the MCP server
Every failure reaches the client as an MCP error carrying its message;
mask_error_details is set off explicitly, since FastMCP also reads it from
the environment. The masking goes, and with it the filter pre-check that
ran a count before every filtered call and the UnknownDatabaseError
translations that existed only to survive it. Explicit domain errors stay.
2026-09-07 13:14:16 +03:00
Yiorgis Gozadinos
95fdb46c3a
Check the sandbox deadline on every host call, one error contract
analysis.code_timeout was enforced only in _run_on_loop, the bridge for
database-bound reads; metadata.json, the cached JSONL files and in-code
search() and list_documents() never looked at the clock, and Monty's
watchdog counts compute only. Every host call now checks the deadline
before it starts.

Host errors keep their message for every caller: the masking added for the
MCP server goes, and the sandbox is one path for the capability and the
server alike.
2026-09-07 12:47:28 +03:00
Yiorgis Gozadinos
bae359e3bf
Align the MCP guidance with the analysis instructions and move to Monty 0.0.23
The execute_code description gains the toc.json node shape, the patterns
the evals put into the analysis instructions (one list_documents call to
map a title to an id, files carry no source, toc before search for a known
document, doc_item_refs are items self_refs, chunk ids join files and are
not citations) and the sandbox's read-only, no-network, time-limit and
output facts, which the analysis instructions now state too. The skill
gains pictures as images, image search when offered, and citing chunk
metadata locators. docs/mcp.md lists the interpreter's limits under Code.

pydantic-monty>=0.0.23 brings collections, itertools, functools,
dataclasses, function decorators and str.format into the sandbox; every
layer names the same modules, and a test imports them. Monty now caps host
callbacks per checkout at 1000 by default; the sandbox raises it out of
reach, since the time budgets govern.
2026-09-07 12:02:33 +03:00
Yiorgis Gozadinos
d7c131bc28
Package the plugin for Codex from the same bundle
The Claude Code plugin moves to plugins/haiku-rag/ and gains a Codex
manifest at .codex-plugin/plugin.json; both share one .mcp.json and one
haiku-rag Agent Skill. .agents/plugins/marketplace.json serves Codex, the
Claude marketplace points at the new path, and scripts/bump_version.py
bumps both manifests. The skill declares its compatibility.
2026-09-07 11:21:25 +03:00
Yiorgis Gozadinos
110adf23ee
Make toc.json item_range the line slice it documents
build_toc stored item positions, while the sandbox instructions describe
item_range as a slice into items.jsonl; a gap in positions pulled the next
heading into a section. Ranges are now indices into the position-ordered
items, and get_document_section slices by index too. docs/mcp.md names the
tools that take sources.
2026-09-07 10:27:02 +03:00
Yiorgis Gozadinos
b374d5eb83
Replace ask_question and analyze with execute_code
In Claude Code the client is the model, so the server no longer runs one.
execute_code runs a Python program per call in the analysis sandbox over
the selected documents and returns what it printed; the sandbox is created
and closed per call so Monty's cumulative budget and a frozen mount never
outlive a program. --no-agents goes with the two tools, and format_citations
in haiku.rag.utils goes with its only caller.

The sandbox exposes chunk metadata to code: chunk_meta on search results,
metadata on list_documents rows and in metadata.json, and chunks.jsonl per
document. A host-side failure inside a program, a document read or an
in-code search raising, reaches the program by exception type only and is
logged with its traceback. recovery_hint moves to haiku.rag.sandbox.

Closes #604.
2026-09-07 10:02:42 +03:00
Yiorgis Gozadinos
0026142e7d
Give MCP search results one channel and tidy two messages
Search results carry text and image blocks only. Claude Code and the
Agent SDK do not forward text blocks when structuredContent is present
and Desktop forwards both, so sending both either hid the rendering or
doubled it. The invalid-filter error keeps the engine's diagnosis and
lists our columns instead of lance's internals. format_citations no
longer repeats the URI of an untitled document.

Refs #599
2026-09-04 15:50:24 +03:00
Yiorgis Gozadinos
2582f2c05a
Expand MCP search results and render the matched chunk's metadata
Both search tools pass their results through HaikuRAG.expand_context, as
every other consumer of search results already did, so a client reads
the hit in its section rather than the chunk that matched. The rendering
gains an opt-in include_chunk_meta that shows the metadata stored with
the matched chunk beyond haiku.rag's structural keys, labelled as the
matched chunk's because an expanded passage spans several chunks and
only the anchor's metadata survives expansion. The capabilities' rendering
is unchanged.

Refs #599
2026-09-04 14:12:12 +03:00
Yiorgis Gozadinos
494294046f
Move to fastmcp 4 and MCP SDK 2
fastmcp>=4.0.2,<5.0.0. Protocol types are snake_case (read_only_hint,
mime_type, input_schema, server_info); the camelCase names warn and go
in fastmcp 5, so the next major is an explicit upgrade. The client
defaults to the sessionless protocol, where initialize_result is None,
so the tests read client.instructions and client.server_info, which both
protocol modes populate. The server answers either mode.
2026-09-04 13:11:06 +03:00
Yiorgis Gozadinos
eb9995e9b7
Ship a Claude Code plugin with the haiku-rag skill
claude-plugin/ holds the plugin manifest, the server configuration
(haiku-rag mcp --stdio, the configuration decides the database) and a
skill that says when to reach for the knowledge base and how to move
from a search result to a document, a section, an answer or a
computation. A repo-root marketplace manifest makes
`claude plugin marketplace add ggozad/haiku.rag` work. The skill
pre-approves every tool the server registers, and a test keeps the two
in step.

The manifest carries the package version, which bump_version.py now
rewrites: a versioned plugin updates only on a bump, so the installed
skill stays in step with the haiku-rag release the user has.

Refs #599
2026-09-04 13:00:15 +03:00
Yiorgis Gozadinos
f45ed90b33
Add --no-agents to leave the model-backed MCP tools out
haiku-rag mcp --no-agents reaches create_mcp_server(agents=False):
ask_question and analyze are not registered and the instructions drop
the line describing them. qa.model always has a default, so a server
without a usable model cannot be detected from configuration; the flag
is how an operator says so.

Refs #599
2026-09-04 13:00:15 +03:00
Yiorgis Gozadinos
5653e876a8
Return search results as agent text, images and structured content
search_documents and search_documents_by_image return a ToolResult: the
format_for_agent rendering with rank, Document ID and Collection so the
text alone drives the document tools; one ImageContent per distinct
picture, labelled with its result; and the SearchResult list without
image_data as structured content. format_for_agent gains an opt-in
include_document_id, so the capabilities' rendering is unchanged.
collect_pictures is the one place pictures are deduplicated and
validated for both wire formats.

Refs #599
2026-09-04 13:00:15 +03:00
Yiorgis Gozadinos
15afb97a6e
Navigate documents by outline and section from the MCP server
build_toc moves from the sandbox into haiku.rag.context; the sandbox
keeps its toc.json unchanged. get_document_outline returns the heading
tree with page numbers and get_document_section one section's text,
subsections included, both resolved in the database holding the
document. Chunk ids never leave the server. ask_question drops `cite`
and always appends its citations.

Refs #599
2026-09-04 13:00:15 +03:00
Yiorgis Gozadinos
2a6d72171d
Make MCP failures errors on the wire
FastMCP masks unexpected exceptions and logs the traceback server-side,
so paths and provider URLs never cross the transport. Expected failures
raise ToolError with a message: unknown document, unknown collection,
invalid filter, invalid base64, and agent failures naming only the
exception type. No tool returns an empty value or an error string on
failure any more.

A filter is validated on its own before the read that would use it, with
a filtered count on one of the selected databases: that is the query
engine rejecting the filter and nothing else, so its message (columns and
the statement) can be forwarded, while a ValueError raised later in the
read stays masked and no database outside the selection is opened.

Refs #599
2026-09-04 13:00:15 +03:00
Yiorgis Gozadinos
88cd2b1d6c
Describe the MCP server and its tools to clients
FastMCP gets instructions, the haiku.rag-slim version and a lifespan.
Every tool carries read-only ToolAnnotations with a title, a description
that says when to use it, and a description on every parameter. The
filter description lists the document columns from DocumentMetaRecord
and how to match metadata; filter is accepted by both search tools.
One strict base64 decoder serves every image parameter.
DocumentInfo carries metadata.

Refs #599
2026-09-04 13:00:15 +03:00
Yiorgis Gozadinos
40d40bcbf2
Cover the configured database set from the MCP server
`haiku-rag mcp` passes covers_set=True and _covering no longer refuses a
scope over several databases. search_documents, search_documents_by_image,
ask_question and analyze take `sources`; get_document takes `source`;
DocumentInfo carries `source`. format_citations gains include_source, which
ask_question sets from covers_multiple so citations name their database
only when the server covers several.

Refs #599
2026-09-04 13:00:14 +03:00
Yiorgis Gozadinos
ce69a8c989
Remove the MCP write tools; the server opens read-only
add_document_from_file, add_document_from_url, add_document_from_text
and delete_document are gone. create_mcp_server and _covering lose
read_only; the client always opens with read_only=True, so the
--read-only flag before `mcp` is redundant and the docs, Dockerfiles and
compose example drop it.

Ingestion is `haiku-rag add`/`add-src`/`delete` and haiku-ingester. No
known consumer used the write tools; one stdio server per client window
made multi-writer the accidental default, and streamable HTTP has no auth.

The test_app stub drops a comment and __init__ that described run_mcp
constructing the client positionally; every path goes through _covering.

Refs #599
2026-09-04 12:59:48 +03:00
Yiorgis Gozadinos
7eeddd1a7c
Start MCP server revisited (#599) 2026-09-04 12:59:08 +03:00
Yiorgis Gozadinos
cf674b93ce
Merge pull request #601 from ggozad/chore/update-default-models
Update default models to qwen3.8, disable telemetry in tests.
2026-09-04 04:53:08 -05:00
Yiorgis Gozadinos
4290111333
Stop tests from sending telemetry to logfire
telemetry.configure() passes send_to_logfire="if-token-present", so on a
machine with a logfire token the CLI tests configured a live exporter:
spans and the periodic metrics export went to the developer's project, and
where the export fired while a cassette was open, VCR rejected it with
CannotOverwriteExistingCassetteException in record mode none.

A token resolves from LOGFIRE_TOKEN or from logfire_credentials.json under
LOGFIRE_CREDENTIALS_DIR (default .logfire, relative to the working
directory), so an authenticated developer uploads with the variable unset.
conftest drops the variable and points credentials discovery at an empty
temporary directory, alongside the existing HAIKU_RAG_CONFIG_PATH setup and
before any haiku.rag import. The explicit send_to_logfire argument beats
LOGFIRE_SEND_TO_LOGFIRE, so denying the token is what disables the exporter.
2026-09-04 12:36:36 +03:00
Yiorgis Gozadinos
a18bfd9781
Anchor the doclaynet fixture path on the test file
test_split_and_merge_matches_single_pass resolved tests/data/doclaynet.pdf
relative to the working directory, so it failed with FileNotFoundError for
any invocation outside the repository root. conftest.py resolves the same
file as Path(__file__).parent / "data"; this was the only site that did not.
2026-09-04 12:36:36 +03:00
Yiorgis Gozadinos
7f7223e0ac
Default to ollama:qwen3.8
Replaces gpt-oss on ModelConfig, qa.model and processing.title_model, and
ministral-3 on the picture-description model. qa.model.vision follows the
model and is now true.

enable_thinking was gated on the gpt-oss name, so it did nothing for
qwen3.8. With title_model's max_tokens of 100 the reasoning consumed the
whole budget and title generation returned an empty string. The mapping
now applies to any ollama model via reasoning_effort(): false sends
"none", true sends "high". Measured on qwen3.8:27b-mlx, "low" does not
disable thinking and "none" does; gpt-oss is the inverse, its template
has no "none" level, so it keeps "low".

Picture description bypasses get_model -- docling posts the request
itself from a params dict -- so the flag was inert on that path too.
vlm_api_params() carries reasoning_effort into both converters' request
bodies. At max_tokens 200 the description survived either way, but the
switch cut completion tokens from 141 to 45.

test_search_tool_skips_binary_content_when_qa_model_is_text_only asserted
the vision default rather than setting it; it now configures vision=False
itself.

docs/benchmarks.md keeps ministral-3: those are recorded measurements.
2026-09-04 12:36:36 +03:00
Yiorgis Gozadinos
660c86f991
vb 2026-09-03 18:00:43 +03:00
Yiorgis Gozadinos
376dcb2e6f
Merge pull request #596 from lawrenceakka/tz-fix
Fix for missing timezone info on ingested docs. mcp schema requires TZ
2026-09-03 09:59:14 -05:00
Yiorgis Gozadinos
4bf4699ee0
Note that offsetless stored timestamps are read as host-local time 2026-09-03 17:48:18 +03:00
Lawrence Akka
880551deba
Use astimezone, add timezone tests, linting 2026-09-03 17:48:09 +03:00
Lawrence Akka
9c56f56660
Fix missing timezone info on ingested docs. mcp schema requires TZ 2026-09-03 17:48:01 +03:00
Yiorgis Gozadinos
f4c657e094
Document LanceDB Cloud region as required 2026-09-03 17:44:32 +03:00
Yiorgis Gozadinos
5bf4cb50aa
vb 2026-09-03 15:43:56 +03:00
Yiorgis Gozadinos
75547f1cbd
Merge pull request #598 from ggozad/chore/db-entry-refactor
Replace lancedb.uri with lancedb.databases; name every database; remove HAIKU_RAG_DB and DB_PATH
2026-09-03 07:42:53 -05:00
Yiorgis Gozadinos
187946a7ae
Name the remedy for a missing configured database and list the databases an unknown name could have meant
A configured or default database that does not exist raises
SourceUnavailableError with the way to create it and without its location.
An unknown --db-name lists the databases there are, configured or default.
The CHANGELOG records Store.db_path, SingleDatabaseSession(ref, config),
the empty lancedb.uri migration and the error-type change.
2026-09-03 15:26:42 +03:00
Yiorgis Gozadinos
3814b2bc78
Report a stemless --db path as a usage error and separate the two lancedb.uri migrations
An empty lancedb.uri, which init-config used to emit, is told to remove the
key; only a non-empty value is told the databases spelling. haiku-rag --db
and haiku-ingester --db turn a path with no stem into typer.BadParameter.
2026-09-03 15:12:09 +03:00
Yiorgis Gozadinos
a04a16c717
Remove the environment overrides and the last unnamed-database wording
HAIKU_RAG_DB and DB_PATH are gone: a capability covers what the
configuration places or the db_path it is given, and the app backend and
the AG-UI example load their configuration as the CLI does. The compose
files point HAIKU_RAG_CONFIG_PATH at the mounted haiku.rag.yaml, which
places the database at /data where DB_VOLUME is mounted; the backend
refuses a configured set since it serves one database. The chat scopes a
selection by source only over a set and names databases on filter rows
only across several. Docstrings, docs and test fixtures stop describing an
unnamed database; every database a search, listing or citation reports
carries a name.
2026-09-03 15:12:09 +03:00
Yiorgis Gozadinos
34180a0fd1
One reference and one placement for a database
DatabaseRef is a name and a location. The configuration places databases
through lancedb.databases alone; with none configured the default is the
entry haiku.rag under storage.data_dir, selectable like any other.
lancedb.uri is removed, and a config carrying it fails to load with the
replacement spelled out. A path passed from Python is valid where the
configuration places nothing and raises AmbiguousDatabaseError beside
lancedb.databases; haiku-rag --db and haiku-ingester --db construct the
scope directly, so a human's override keeps working. Every database
answers to a name, and a database given as a path keeps its own errors.
2026-09-03 15:12:08 +03:00
Yiorgis Gozadinos
9baa213b34
Hand storage the database location, not the configuration that placed it
Store, connect_lancedb, gather_database_info and run_doctor take a
location, a path or a URI, and classify it with ConnectionMode.of.
SingleDatabaseSession owns the resolved DatabaseRef and passes its
location down. This removes DatabaseRef.connection(), default_db_path,
the placeholder path for URI-backed databases and the per-database
config copies, so the configuration a client holds is the one the caller
gave it. The chat hands its capabilities the scope it opened along with
the client it lends, and the v0.58.0 migration no longer checks local
free disk for a database behind a URI.
2026-09-03 15:11:37 +03:00
Yiorgis Gozadinos
4f7f69aaf0
Merge pull request #597 from ggozad/feat/search-fanout-units
Price search budget in units and deduplicate fan-out results
2026-09-03 07:09:13 -05:00
Yiorgis Gozadinos
31fc8aba29
Guard dedup against compaction and document search units
A picture deduplicated within a burst must still reach the model once in
its own question and, when cited, again via the capsule re-fetch after
compaction; uncited it is dropped. Two wire tests pin both paths. The
harness builds the rag capability with defer_loading=False: a deferred
evidence capability the model has not loaded presents an empty record,
so compaction computes boundary 0 and rewrites nothing.

qa.md describes max_searches in units.
2026-09-02 12:29:12 +03:00
Yiorgis Gozadinos
d9f489dcc8
Deduplicate search results within one model response
Sibling searches emitted in one response overlap heavily (40.6% of
returned chunk slots on Glimmer ORB fan-out cases). A result whose
rendered evidence a sibling already showed keeps its rank slot but
collapses to a reference line, and a picture attaches once per response
keyed on (source, document_id, self_ref). Equivalence is the
format_for_agent rendering at neutral rank/total plus picture keys,
bucketed under the qualified chunk id, so another database's copy or a
different expansion of the same anchor formats in full.

Search state now commits only after formatting and image construction
succeed: a raising image build no longer leaves results citable that
the model never saw, notes evidence for them, or suppresses a later
sibling.
2026-09-02 12:29:12 +03:00
Yiorgis Gozadinos
ddac328d05
Price qa.max_searches in search units
Searches a model emits in one response share a budget unit, up to
FREE_SIBLINGS_PER_ROUND (3) per unit; sequential searches pay one unit
each, as before. Grouping keys on RunContext.run_step, which pydantic-ai
increments once per model request. A budget-rejected round fails all its
remaining siblings, and tracking resets per run.

Glimmer opens most questions with a burst of ~3 rephrasings in a single
response (95.8% of its three-search ORB cases are one-response bursts),
spending 3 of 5 searches before reading anything. Pass rate at 3 calls
equals 1 call, so a burst is priced as one probe.
2026-09-02 12:29:12 +03:00
Yiorgis Gozadinos
f63dec5f75
vb 2026-09-01 16:07:03 +03:00
Yiorgis Gozadinos
2bb169867c
Merge pull request #595 from ggozad/feat/vector-nprobes
Expose search.vector_nprobes
2026-09-01 08:05:19 -05:00
Yiorgis Gozadinos
2d8b2cb9ed
Expose search.vector_nprobes
The IVF probe count was the one vector search parameter with no
setting, and it is what bounds recall on a large indexed corpus:
partitions grow with the corpus, so a fixed probe count covers less of
it, and vector_refine_factor can only re-score what the probes
returned. Defaults to 20, matching lance, so search is unchanged.

Also corrects the create-index docs, which predate the measurement:
optimize() covers new chunks, so a rebuild is about retraining
centroids rather than reaching unindexed rows.
2026-09-01 15:54:36 +03:00
Yiorgis Gozadinos
f7e9535d7b
Merge pull request #593 from ggozad/fix/cross-database-fusion
Order cross-database fusion by cosine similarity to the query
2026-09-01 07:36:12 -05:00
Yiorgis Gozadinos
7be71ebbac
Order cross-database fusion by cosine similarity to the query
Retrieval scores are each database's own rank arithmetic; the databases
in a selection share an embedder, so similarity in that one space is the
signal comparable across databases by construction. Measured product to
product against score ordering: +8.3 to +16.6pp recall@5 across five
cells on two corpora, flat in collection count and corpus shape where
score ordering dips with both, closing roughly 60% of the gap to a
reranker; order-sensitivity residual 0.00pp in every cell. Exact ties
collapse from 51-81% of candidates to under 1%. Full-text-only searches
keep retrieval-score order, having no query vector. The vector column
already travels with every search result, so the similarity costs no
additional transfer; per-chunk embeddings are materialized only for the
federated path that reads them.
2026-09-01 15:26:42 +03:00
Yiorgis Gozadinos
82fe91bb76
Order cross-database fusion by retrieval score
Rank interleaving guarantees every database slots regardless of content;
on domain-split collections it allocates no better than chance and costs
4.7pp recall@5 at four collections against score ordering (7.1pp at
eight). Hybrid scores are each database's own vector/FTS rank agreement,
which carries across databases; equal scores resolve by within-database
rank, and only a tie on both falls to configured order, leaving
permutation sensitivity at 0.02-0.26pp. Fused results carry the
candidate's own retrieval score, so the context-expansion re-sort
preserves fused order.
2026-09-01 15:26:41 +03:00
Yiorgis Gozadinos
d06277408b
Break cross-database RRF rank ties by retrieval score
Disjoint corpora give every database's rank-r candidate the same RRF
score, and the stable sort resolved those ties to lancedb.databases
declaration order, discarding the retrieval scores entirely. Ties now
break on the raw retrieval score, which is uncalibrated across indexes
but only ever orders candidates within one rank tier: the databases in
a fusion share an embedder and ran the same search type, and it can
never lift a candidate above another rank. Hybrid per-database scores
are themselves rank-derived, so exact agreement still ties and keeps
configured order, deterministically. The n > limit depth quota is
unchanged, pending the retrieval eval.
2026-09-01 15:26:41 +03:00
Yiorgis Gozadinos
f927998643
Merge pull request #594 from ggozad/docs/vector-index-state
Document vector index state and measured cost
2026-09-01 07:23:50 -05:00
Yiorgis Gozadinos
c0cf612d68
Document vector index state and measured cost
Benchmarks state that every published number is exact brute-force kNN;
no benchmark database carries a vector index. The Vector Indexing
section gains the measured with/without IVF_PQ comparison (hotpotqa,
orb_multimodal_nemotron, frames: free to ~121k chunks, 0.0044 MAP at
426k, ~30 s / 4 GB build) and corrects the re-indexing story: optimize()
folds new chunks into the index as delta parts via auto_vacuum, so a
rebuild is about retraining centroids, not covering new rows.
2026-09-01 15:19:13 +03:00
Yiorgis Gozadinos
a7efd76016
vb 2026-08-31 20:04:16 +03:00
125 changed files with 5024 additions and 2102 deletions

View 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"
}
]
}

View 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."
}
]
}

View file

@ -2,6 +2,148 @@
## [Unreleased] ## [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 ### Changed
- lancedb 0.37.1. - lancedb 0.37.1.
@ -2291,7 +2433,11 @@ Existing documents without DoclingDocument data will work but won't have provena
- Initial version tracking - Initial version tracking
[Unreleased]: https://github.com/ggozad/haiku.rag/compare/0.79.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.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.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.77.0...0.77.0

View file

@ -16,7 +16,7 @@ Built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.pydantic.dev/
- **Hybrid search** — Vector + full-text with Reciprocal Rank Fusion - **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 - **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 capability with citations (page numbers, section headings) - **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`, MCP, and the chat TUI - **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 - **Reranking** — local cross-encoders, Cohere, Zero Entropy, or vLLM
- **Analysis capability** — 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 - **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
@ -110,12 +110,26 @@ For direct agent composition, see the [capabilities documentation](https://ggoza
## MCP Server ## MCP Server
Use with AI assistants like Claude Desktop: Use with AI assistants like Claude Code, Codex, and Claude Desktop:
```bash ```bash
haiku-rag mcp --stdio 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: Add to your Claude Desktop configuration:
```json ```json
@ -129,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 ## Examples

View file

@ -2,8 +2,9 @@
ANTHROPIC_API_KEY=your-anthropic-key ANTHROPIC_API_KEY=your-anthropic-key
OPENAI_API_KEY=your-openai-key OPENAI_API_KEY=your-openai-key
# Database path # Host path of the LanceDB database, mounted at /data where haiku.rag.yaml
DB_PATH=/path/to/your/haiku.rag.lancedb # places it
DB_VOLUME=./data/haiku.rag.lancedb
# Optional: Ollama base URL (if using local models) # Optional: Ollama base URL (if using local models)
# Use host.docker.internal to reach Ollama running on the host machine # Use host.docker.internal to reach Ollama running on the host machine

View file

@ -40,7 +40,8 @@ A conversational RAG interface built with [CopilotKit](https://copilotkit.ai/) a
| Variable | Description | Required | | 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 | | `ANTHROPIC_API_KEY` | Anthropic API key | One LLM key required |
| `OPENAI_API_KEY` | OpenAI 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 | | `OLLAMA_BASE_URL` | Ollama server URL (default: `http://host.docker.internal:11434`) | For local models |

View file

@ -1,9 +1,7 @@
import asyncio import asyncio
import logging import logging
import os
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from dataclasses import dataclass, field from dataclasses import dataclass, field
from pathlib import Path
from typing import Any from typing import Any
from ag_ui.core import EventType, StateSnapshotEvent from ag_ui.core import EventType, StateSnapshotEvent
@ -26,8 +24,8 @@ from haiku.rag.capabilities.policy import (
) )
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState, create_capability from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState, create_capability
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config import load_yaml_config from haiku.rag.client.scope import DatabaseScope
from haiku.rag.config.models import AppConfig from haiku.rag.config import get_config
from haiku.rag.telemetry import configure as configure_telemetry from haiku.rag.telemetry import configure as configure_telemetry
from haiku.rag.utils import get_model from haiku.rag.utils import get_model
@ -40,19 +38,23 @@ logging.basicConfig(
) )
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# Load config # The configuration places the database. This app serves one.
config_path = Path("/app/haiku.rag.yaml") config = get_config()
if config_path.exists(): scope = DatabaseScope.resolve(config)
yaml_data = load_yaml_config(config_path) if scope.covers_multiple:
config = AppConfig.model_validate(yaml_data) raise SystemExit(
else: f"lancedb.databases names {', '.join(scope.names)}; this app serves one "
config = AppConfig() "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}") 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}") logger.info(f"QA Provider: {config.qa.model.provider}, Model: {config.qa.model.name}")
# Only HaikuRAG client is a singleton (expensive to create) # Only HaikuRAG client is a singleton (expensive to create)
@ -71,7 +73,7 @@ async def get_client() -> HaikuRAG:
if _client is None: if _client is None:
async with _client_lock: async with _client_lock:
if _client is None: if _client is None:
client = HaikuRAG(db_path=db_path, config=config, create=True) client = HaikuRAG(config=config, create=True)
await client.__aenter__() await client.__aenter__()
_client = client _client = client
return _client return _client
@ -82,7 +84,7 @@ class AppDeps:
state: dict[str, Any] = field(default_factory=dict) state: dict[str, Any] = field(default_factory=dict)
capability = create_capability(db_path=db_path, config=config, defer_loading=False) capability = create_capability(config=config, defer_loading=False)
agent = Agent( agent = Agent(
get_model(config.qa.model, config), get_model(config.qa.model, config),
@ -138,15 +140,15 @@ async def health_check(_: Request) -> JSONResponse:
"status": "healthy", "status": "healthy",
"qa_provider": config.qa.model.provider, "qa_provider": config.qa.model.provider,
"qa_model": config.qa.model.name, "qa_model": config.qa.model.name,
"db_path": str(db_path), "db_path": str(database.location),
"db_exists": db_path.exists(), "db_exists": _database_exists(),
} }
) )
async def list_documents(_: Request) -> JSONResponse: async def list_documents(_: Request) -> JSONResponse:
"""List all documents in the database.""" """List all documents in the database."""
if not db_path.exists(): if not _database_exists():
return JSONResponse({"documents": [], "error": "Database not found"}) return JSONResponse({"documents": [], "error": "Database not found"})
client = await get_client() client = await get_client()
@ -162,11 +164,11 @@ async def list_documents(_: Request) -> JSONResponse:
async def db_info(_: Request) -> JSONResponse: async def db_info(_: Request) -> JSONResponse:
"""Get database info and statistics.""" """Get database info and statistics."""
if not db_path.exists(): if not _database_exists():
return JSONResponse( return JSONResponse(
{ {
"exists": False, "exists": False,
"path": str(db_path), "path": str(database.location),
"documents": 0, "documents": 0,
"chunks": 0, "chunks": 0,
} }
@ -180,7 +182,7 @@ async def db_info(_: Request) -> JSONResponse:
return JSONResponse( return JSONResponse(
{ {
"exists": True, "exists": True,
"path": str(db_path), "path": str(database.location),
"documents": stats["documents"].get("num_rows", 0), "documents": stats["documents"].get("num_rows", 0),
"chunks": stats["chunks"].get("num_rows", 0), "chunks": stats["chunks"].get("num_rows", 0),
"documents_bytes": stats["documents"].get("total_bytes", 0), "documents_bytes": stats["documents"].get("total_bytes", 0),
@ -214,7 +216,7 @@ async def visualize_chunk(request: Request) -> JSONResponse:
if isinstance(parsed, list): if isinstance(parsed, list):
refs = [str(x) for x in parsed] 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) return JSONResponse({"error": "Database not found"}, status_code=404)
client = await get_client() client = await get_client()

View file

@ -8,7 +8,7 @@ dependencies = [
"uvicorn[standard]>=0.40.0", "uvicorn[standard]>=0.40.0",
"pydantic-ai-slim[ag-ui,anthropic,openai]>=2.18.0,<3.0.0", "pydantic-ai-slim[ag-ui,anthropic,openai]>=2.18.0,<3.0.0",
"python-dotenv>=1.2.1", "python-dotenv>=1.2.1",
"haiku.rag-slim>=0.79.0", "haiku.rag-slim>=0.82.1",
"logfire[pydantic-ai]>=3.17.0", "logfire[pydantic-ai]>=3.17.0",
] ]

View file

@ -11,13 +11,14 @@ services:
ports: ports:
- "127.0.0.1:8001:8000" - "127.0.0.1:8001:8000"
environment: environment:
- DB_PATH=/data - HAIKU_RAG_CONFIG_PATH=/app/haiku.rag.yaml
- ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY:-} - ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY:-}
- OPENAI_API_KEY=${OPENAI_API_KEY:-} - OPENAI_API_KEY=${OPENAI_API_KEY:-}
- OLLAMA_BASE_URL=${OLLAMA_BASE_URL:-http://host.docker.internal:11434} - OLLAMA_BASE_URL=${OLLAMA_BASE_URL:-http://host.docker.internal:11434}
- LOGFIRE_TOKEN=${LOGFIRE_TOKEN:-} - LOGFIRE_TOKEN=${LOGFIRE_TOKEN:-}
volumes: 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 - ./backend:/app/src:ro
- ./haiku.rag.yaml:/app/haiku.rag.yaml:ro - ./haiku.rag.yaml:/app/haiku.rag.yaml:ro
extra_hosts: extra_hosts:

View file

@ -7,13 +7,14 @@ services:
ports: ports:
- "127.0.0.1:8001:8000" - "127.0.0.1:8001:8000"
environment: environment:
- DB_PATH=/data - HAIKU_RAG_CONFIG_PATH=/app/haiku.rag.yaml
- ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY:-} - ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY:-}
- OPENAI_API_KEY=${OPENAI_API_KEY:-} - OPENAI_API_KEY=${OPENAI_API_KEY:-}
- OLLAMA_BASE_URL=${OLLAMA_BASE_URL:-http://host.docker.internal:11434} - OLLAMA_BASE_URL=${OLLAMA_BASE_URL:-http://host.docker.internal:11434}
- LOGFIRE_TOKEN=${LOGFIRE_TOKEN:-} - LOGFIRE_TOKEN=${LOGFIRE_TOKEN:-}
volumes: 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 - ./haiku.rag.yaml:/app/haiku.rag.yaml:ro
extra_hosts: extra_hosts:
- "host.docker.internal:host-gateway" - "host.docker.internal:host-gateway"

View file

@ -1,6 +1,12 @@
# haiku.rag configuration for the chat app # haiku.rag configuration for the chat app
# Copy to haiku.rag.yaml and customize as needed # 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 configuration
qa: qa:
model: model:

View file

@ -40,4 +40,4 @@ EXPOSE 8001 8765
# Default command: read-only MCP server. The companion ingester service is # Default command: read-only MCP server. The companion ingester service is
# launched via docker-compose against the same image. # 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"]

View file

@ -39,4 +39,4 @@ EXPOSE 8001 8765
# Default command: read-only MCP server. The companion ingester service is # Default command: read-only MCP server. The companion ingester service is
# launched via docker-compose against the same image. # 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"]

View file

@ -40,8 +40,8 @@ Create a `.env` file in the `app/` directory:
ANTHROPIC_API_KEY=your-anthropic-key ANTHROPIC_API_KEY=your-anthropic-key
OPENAI_API_KEY=your-openai-key OPENAI_API_KEY=your-openai-key
# Database path # Host path of the LanceDB database, mounted at /data
DB_PATH=/path/to/your/haiku.rag.lancedb DB_VOLUME=/path/to/your/haiku.rag.lancedb
# Optional: Ollama base URL (if using local models) # Optional: Ollama base URL (if using local models)
OLLAMA_BASE_URL=http://localhost:11434 OLLAMA_BASE_URL=http://localhost:11434
@ -50,16 +50,22 @@ OLLAMA_BASE_URL=http://localhost:11434
LOGFIRE_TOKEN=your-logfire-token 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 ```yaml
# app/haiku.rag.yaml # app/haiku.rag.yaml
lancedb:
databases:
haiku.rag: /data
qa: qa:
model: model:
provider: anthropic provider: anthropic
name: claude-sonnet-4-20250514 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 ## API endpoints
| Endpoint | Method | Description | | Endpoint | Method | Description |

View file

@ -6,6 +6,8 @@ We evaluate `haiku.rag` on a small set of datasets that exercise different parts
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. 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.
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).
### OpenRAG Bench (ORB) ### OpenRAG Bench (ORB)
[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. [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.

View file

@ -16,7 +16,7 @@ When `qa.max_searches` or `analysis.max_executions` runs out, the exhausted tool
| `analysis_execute_code(code)` | Run Python against the virtual document filesystem. | | `analysis_execute_code(code)` | Run Python against the virtual document filesystem. |
| `analysis_cite(chunk_ids)` | Register retrieved or filesystem-derived chunk IDs. | | `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`, and `toc.json`. 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 ## Compose an agent

View file

@ -145,12 +145,6 @@ Applications serving AG-UI should adapt the agent with Pydantic AI's `AGUIAdapte
## Database Selection ## Database Selection
RAG and analysis capabilities select databases in this order: 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`.
1. The `db_path` argument.
2. `HAIKU_RAG_DB`.
3. [`lancedb.databases`](../configuration/storage.md#multiple-databases), which selects the full configured set.
4. [`lancedb.uri`](../configuration/storage.md#changing-the-default-database-path), which selects one database.
5. `config.storage.data_dir / "haiku.rag.lancedb"`.
Passing a client through `rag=` bypasses this selection. The capability uses the databases covered by that client and does not close it. Passing a client through `rag=` bypasses this selection. The capability uses the databases covered by that client and does not close it.

View file

@ -12,7 +12,7 @@ The `haiku-rag` CLI provides complete document management functionality.
Per-command options: 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 - `-h` - Show help for specific command
Example: Example:
@ -24,7 +24,7 @@ The `haiku-rag` CLI provides complete document management functionality.
haiku-rag add -h haiku-rag add -h
``` ```
With `lancedb.databases` configured, `search`, `ask`, `analyze`, and `chat` 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). 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 ## Document Management
@ -402,15 +402,16 @@ haiku-rag create-index [--db /path/to/your.lancedb]
- Creates an IVF_PQ index using the configured `search.vector_index_metric` (cosine/l2) - Creates an IVF_PQ index using the configured `search.vector_index_metric` (cosine/l2)
**When to use:** **When to use:**
- After ingesting documents (indexes are not created automatically) - On a collection over 100,000 chunks (below that, brute-force kNN is exact and fast enough)
- After adding significant new data to rebuild the index - After substantial corpus growth, to retrain the centroids
- Use `haiku-rag info` to check index status and see how many chunks are indexed/unindexed - 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:** **Search behavior:**
- Without index: Brute-force kNN search (exact nearest neighbors, slower for large datasets) - Without index: Brute-force kNN search (exact nearest neighbors, slower for large datasets)
- With index: Fast ANN (approximate nearest neighbors) using IVF_PQ - With index: ANN (approximate nearest neighbors) using IVF_PQ, tuned by `search.vector_nprobes`
- With stale index: LanceDB combines indexed results (fast ANN) + brute-force kNN on unindexed rows - Between a write and the next `optimize()`: LanceDB combines ANN over indexed rows with brute-force kNN over the remainder
- Performance degrades as more unindexed data accumulates
### Rebuild Database ### Rebuild Database
@ -476,9 +477,6 @@ haiku-rag mcp --port 9000
# Bind to all interfaces (containers, trusted LAN) # Bind to all interfaces (containers, trusted LAN)
haiku-rag mcp --host 0.0.0.0 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 See [MCP](mcp.md) for details. For continuous document ingestion

View file

@ -61,7 +61,7 @@ embeddings:
qa: qa:
model: model:
provider: ollama provider: ollama
name: gpt-oss name: qwen3.8
enable_thinking: true enable_thinking: true
``` ```
@ -85,10 +85,9 @@ ingester:
delete_orphans: true delete_orphans: true
lancedb: lancedb:
uri: "" # Empty for local, or db://, s3://, az://, gs:// databases: {} # Name-to-location map; empty places haiku.rag under data_dir
api_key: "" api_key: "" # LanceDB Cloud (db://) credentials
region: "" region: ""
databases: {} # Name-to-location map to search multiple at once; excludes uri
embeddings: embeddings:
model: model:
@ -106,7 +105,7 @@ reranking:
qa: qa:
model: model:
provider: ollama provider: ollama
name: gpt-oss name: qwen3.8
enable_thinking: true enable_thinking: true
temperature: 0.3 temperature: 0.3
max_searches: 5 max_searches: 5
@ -136,7 +135,7 @@ processing:
auto_title: false # Auto-generate titles on ingestion auto_title: false # Auto-generate titles on ingestion
title_model: title_model:
provider: ollama provider: ollama
name: gpt-oss name: qwen3.8
enable_thinking: false enable_thinking: false
temperature: 0.3 temperature: 0.3
max_tokens: 100 max_tokens: 100

View file

@ -30,7 +30,7 @@ processing:
auto_title: false # Auto-generate titles on ingestion auto_title: false # Auto-generate titles on ingestion
title_model: # LLM for title generation (fallback) title_model: # LLM for title generation (fallback)
provider: ollama provider: ollama
name: gpt-oss name: qwen3.8
enable_thinking: false enable_thinking: false
# Conversion options (works with both local and remote converters) # Conversion options (works with both local and remote converters)
@ -54,7 +54,7 @@ processing:
picture_description: picture_description:
model: model:
provider: ollama provider: ollama
name: ministral-3 name: qwen3.8
pictures: image # none | description | image pictures: image # none | description | image
``` ```
@ -270,7 +270,7 @@ processing:
picture_description: # only consulted when pictures == "description" picture_description: # only consulted when pictures == "description"
model: model:
provider: ollama # any OpenAI-compatible /v1/chat/completions provider provider: ollama # any OpenAI-compatible /v1/chat/completions provider
name: ministral-3 name: qwen3.8
timeout: 90 timeout: 90
max_tokens: 200 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) | | `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`) | | `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). 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: 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: 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:** **Recommended combinations:**
@ -368,7 +368,7 @@ processing:
auto_title: true auto_title: true
title_model: title_model:
provider: ollama provider: ollama
name: gpt-oss name: qwen3.8
enable_thinking: false enable_thinking: false
``` ```

View file

@ -15,7 +15,7 @@ Configure model behavior for the `qa` and `analysis` capabilities. These setting
qa: qa:
model: model:
provider: ollama provider: ollama
name: gpt-oss name: qwen3.8
temperature: 0.3 temperature: 0.3
max_tokens: 500 max_tokens: 500
``` ```
@ -79,7 +79,7 @@ See the [Pydantic AI thinking documentation](https://ai.pydantic.dev/thinking/)
- **Google**: Gemini models with thinking support - **Google**: Gemini models with thinking support
- **Groq**: Models with reasoning capabilities - **Groq**: Models with reasoning capabilities
- **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`. - **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**: Models supporting reasoning (gpt-oss, etc.) - **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. - **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.) - **LM Studio**: Models supporting reasoning (gpt-oss, etc.)
@ -311,7 +311,7 @@ Configure which LLM provider to use for question answering. Any provider and mod
qa: qa:
model: model:
provider: ollama 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`: The Ollama base URL can be configured via the `OLLAMA_BASE_URL` environment variable, config file, or defaults to `http://localhost:11434`:

View file

@ -20,22 +20,22 @@ Context expansion is automatic and section-aware. For structured documents (with
## Question Answering Configuration ## Question Answering Configuration
Configure the RAG capability (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 ```yaml
qa: qa:
model: model:
provider: ollama provider: ollama
name: gpt-oss name: qwen3.8
enable_thinking: true enable_thinking: true
temperature: 0.3 # Default: 0.3 temperature: 0.3 # Default: 0.3
vision: false # Set true for vision-capable models vision: true # Set false for text-only models
max_searches: 5 # Maximum search tool calls per question max_searches: 5 # Maximum search units per question
``` ```
- **model**: LLM configuration (see [Providers](providers.md#model-settings)) - **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 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. - **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 tool calls a capability can make per question (default: 5). Shared by the RAG and analysis capabilities. - **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" !!! 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. `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.
@ -50,13 +50,13 @@ analysis:
provider: anthropic provider: anthropic
name: claude-sonnet-4-20250514 name: claude-sonnet-4-20250514
temperature: 0.0 # Default: 0.0 (deterministic for code generation) temperature: 0.0 # Default: 0.0 (deterministic for code generation)
code_timeout: 60.0 # Max seconds a call may spend reading documents 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_output_chars: 50000 # Truncate output after this many chars
max_executions: 15 # Max execute_code calls per question 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`. - **model**: LLM configuration (see [Providers](providers.md#model-settings)). When unset, falls back to `qa.model`.
- **code_timeout**: Seconds a single `execute_code` call may spend reading documents (default: 60). The sandbox refuses further reads past this point. Code that computes without reading is killed by the worker watchdog at the same limit. `code_timeout * max_executions` is the cumulative ceiling across all calls in one question. - **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_output_chars**: Truncate code output after this many characters (default: 50000)
- **max_executions**: Maximum `execute_code` calls per question before the capability 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)

View file

@ -68,18 +68,19 @@ Vacuum also folds new rows into the full-text index. Search stays correct withou
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. 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.
### Changing the Default Database Path ### Placing the Database
`storage.data_dir` holds the default database, always called `haiku.rag.lancedb`. To put the database somewhere else for every command, give `lancedb.uri` a local path: `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 ```yaml
lancedb: lancedb:
uri: /data/notes.lancedb databases:
notes: /data/notes.lancedb
``` ```
An explicit `--db PATH` overrides `lancedb.uri` for that invocation. 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`.
This places one database without naming it. Its `source` is `None` in search results, citations and documents, since only [`lancedb.databases`](#multiple-databases) assigns the names that carry provenance. A path here changes where the database lives, not what it is called. 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. 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.
@ -109,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). The [default location](index.md#configuration-file-locations) is platform-specific (e.g., `~/Library/Application Support/haiku.rag/` on macOS).
Opening a nonexistent unnamed local database raises `FileNotFoundError`, naming its path. This prevents accidental database creation from typos or misconfigured paths. A database named in `lancedb.databases` raises `SourceUnavailableError` instead, naming the database and not its location. 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 ## 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 ```yaml
# LanceDB Cloud # LanceDB Cloud
lancedb: lancedb:
uri: db://your-database-name databases:
papers: db://your-database-name
api_key: your-api-key api_key: your-api-key
region: us-west-2 # optional region: us-west-2
# Amazon S3 # Amazon S3
lancedb: lancedb:
uri: s3://my-bucket/my-table databases:
papers: s3://my-bucket/my-table
storage_options: storage_options:
region: us-east-1 region: us-east-1
# Amazon S3 with explicit credentials # Amazon S3 with explicit credentials
lancedb: lancedb:
uri: s3://my-bucket/my-table databases:
papers: s3://my-bucket/my-table
storage_options: storage_options:
aws_access_key_id: YOUR_ACCESS_KEY aws_access_key_id: YOUR_ACCESS_KEY
aws_secret_access_key: YOUR_SECRET_KEY aws_secret_access_key: YOUR_SECRET_KEY
@ -138,7 +142,8 @@ lancedb:
# S3-compatible (SeaweedFS, Tigris, etc.) # S3-compatible (SeaweedFS, Tigris, etc.)
lancedb: lancedb:
uri: s3://my-bucket/my-table databases:
papers: s3://my-bucket/my-table
storage_options: storage_options:
endpoint: http://localhost:8333 endpoint: http://localhost:8333
aws_access_key_id: YOUR_ACCESS_KEY aws_access_key_id: YOUR_ACCESS_KEY
@ -148,21 +153,24 @@ lancedb:
# Azure Blob Storage # Azure Blob Storage
lancedb: lancedb:
uri: az://my-container/my-table databases:
papers: az://my-container/my-table
# Google Cloud Storage # Google Cloud Storage
lancedb: lancedb:
uri: gs://my-bucket/my-table databases:
papers: gs://my-bucket/my-table
# HDFS # HDFS
lancedb: 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. - **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). - **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"`. - **S3-compatible stores** (MinIO, Tigris, etc.): Set `endpoint` in `storage_options`. When using `http://` endpoints, also set `allow_http: "true"`.
- **Local path** (no scheme): `uri` also takes a local path, which is how the default database is pointed elsewhere. See [Changing the Default Database Path](#changing-the-default-database-path). - **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. 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.
@ -188,11 +196,11 @@ 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": 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. - **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. 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 `lancedb.uri` 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`. A one-entry mapping is unambiguous and is accepted. `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 ## Multiple Databases
@ -206,7 +214,7 @@ lancedb:
notes: /data/notes.lancedb notes: /data/notes.lancedb
``` ```
A location can be a URI or local path. `databases` and `uri` are mutually exclusive. 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. 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.
@ -227,7 +235,7 @@ results = await client.search("query") # every database
results = await client.search("query", sources=["papers"]) # one of them results = await client.search("query", sources=["papers"]) # one of them
``` ```
Candidates are combined into one ranked list with the configured reranker, or with reciprocal rank fusion when reranking is disabled. `SearchResult.source`, `Citation.source`, and `Document.source` contain the configured database name. The name is retained when a client covers only one named database. Databases configured through `lancedb.uri` are unnamed, so their `source` is `None`. 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. 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.
@ -245,17 +253,13 @@ The chat document filter selects by document and database: the search is narrowe
#### Ranking #### Ranking
Reciprocal rank fusion compares positions rather than scores, so each database contributes top-ranked results even when another database has stronger matches. A reranker scores the combined candidate set directly, which has been measured to help aggregate retrieval and to hurt attribution between near-identical documents. 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.
Aggregate retrieval is stronger with a reranker. In a 3,045-query evaluation over a corpus split across three databases, reranking produced retrieval MAP 0.9914, compared with 0.9918 for the same corpus in one database. Without a reranker, MAP was 0.6044, compared with 0.9798 in one database. Reranking cost grows with the number of databases because each contributes candidates. 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 reranker scores the combined candidates with no notion of which database each came from, so on near-identical text it can pick the wrong database's chunk, where fusion keeps them apart because each database contributes its own top-ranked result. In two nine-case acceptance runs over a synthetic corpus holding one station in two databases under near-identical names, attribution was weaker with reranking: citing the right database succeeded 5 of 9 and 6 of 9 times with a reranker, against 8 of 9 and 9 of 9 without. 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.
Configure a reranker where retrieval breadth matters, and measure it where answers have to attribute between documents that read alike. 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.
Without a reranker, consider increasing `search.limit` with the number of databases. With three complete rankings and a limit of 5, a database may contribute only one or two results. A higher limit also sends more results to the caller and model.
Image queries are vector-only and skip the reranker: there is no query text to score a document against, so candidates keep their vector ranking and fusion ranks by position.
If a selected database is unavailable, the operation fails with `SourceUnavailableError`, which names that database. If a selected database is unavailable, the operation fails with `SourceUnavailableError`, which names that database.
@ -277,9 +281,9 @@ Conversion, chunking, and title generation do not access a database and remain a
Commands use database sets as follows: Commands use database sets as follows:
- **Set-capable**: `search`, `ask`, `analyze`, and `chat` use the full configured set, or the single database selected by `--db-name`. - **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. - **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`, `visualize`, and `mcp` — works on one database, selected with the global `--db-name` option. - **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 ```bash
haiku-rag search "query" # every configured database haiku-rag search "query" # every configured database
@ -287,7 +291,7 @@ haiku-rag --db-name papers list # one of them
haiku-rag --db-name papers migrate haiku-rag --db-name papers migrate
``` ```
`--db-name` selects an entry from `lancedb.databases`, including remote entries. `--db` selects a local path and overrides the configured location. A single-database command requires one of these options when multiple databases are configured. A configured set of one is selected automatically. `--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: Each database is created, migrated and vacuumed on its own:
@ -304,6 +308,7 @@ Configure vector search settings:
search: search:
vector_index_metric: cosine # cosine or l2 vector_index_metric: cosine # cosine or l2
vector_refine_factor: 30 # Re-ranking factor for accuracy 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). For search behavior settings (`limit`, `max_context_chars`), see [Search and Question Answering](qa.md#search-settings).
@ -313,10 +318,22 @@ For search behavior settings (`limit`, `max_context_chars`), see [Search and Que
- `l2`: Euclidean distance - `l2`: Euclidean distance
- **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 - **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 - **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 !!! 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. 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:** **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: 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:
@ -332,12 +349,12 @@ This command:
**Re-indexing:** **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 ```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. 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.

View file

@ -19,15 +19,71 @@ haiku-rag mcp --host 0.0.0.0 --port 8001
# stdio transport (for Claude Desktop) # stdio transport (for Claude Desktop)
haiku-rag mcp --stdio 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 `--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 — 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. 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 ## 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 | Tool | Registered | Parameters |
- `file_path` (required): Path to the file |---|---|---|
- `metadata` (optional): Key-value metadata | `search_documents` | always | `query`, `limit`, `include_images`, `filter`, `sources` |
- `title` (optional): Human-readable title | `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 `search_documents` runs hybrid search, vector and full-text. Its text content
- `url` (required): URL to fetch is the rendering the in-process agents read: results best first, each with its
- `metadata` (optional): Key-value metadata rank, `Document ID`, `Collection` when the server covers several, the document
- `title` (optional): Human-readable title 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 `get_document` returns a document whole, in reading order. For a long one,
- `content` (required): Text content `get_document_outline` returns the heading tree with page numbers and
- `uri` (optional): URI identifier `get_document_section` the text of one section, subsections included; a
- `metadata` (optional): Key-value metadata node's `id` in the outline is the `section_id`. A document without headings
- `title` (optional): Human-readable title 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 ### Code
- `document_id` (required): The document ID
- **`list_documents`** - List documents with pagination and filtering `execute_code` runs a Python program in the sandbox of the
- `limit` (optional): Maximum number to return [analysis capability](capabilities/analysis.md), over the documents `filter`
- `offset` (optional): Number to skip and `sources` select, and returns what it printed. The program reads
- `filter` (optional): SQL WHERE clause for filtering `/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 The interpreter is [Monty](https://github.com/pydantic/monty), a Python subset.
- `document_id` (required): The document ID 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) `filter` is a SQL WHERE clause over the document columns `id`, `uri`, `title`,
- `query` (required): Search query `metadata`, `created_at`, `updated_at`. `metadata` is a JSON string, so match
- `limit` (optional): Maximum results (uses config default if not specified) its keys with LIKE:
- `include_images` (optional, default `true`): Attach base64-encoded picture bytes to picture-labeled results
- **`search_documents_by_image`** - Search using an image as the query (registered only when the configured embedder supports images) ```sql
- `image_base64` (required): Base64-encoded image (PNG/JPEG bytes) metadata LIKE '%"author": "Smith"%'
- `limit` (optional): Maximum results uri LIKE '%.pdf'
- `include_images` (optional, default `true`) title = 'Q3 report'
```
### Question Answering ### Errors
- **`ask_question`** - Ask questions about your documents A failure is an MCP error carrying its message, never an empty result: a
- `question` (required): The question to ask document or section id that matches nothing, a collection the server does not
- `cite` (optional): Include source citations (default: false) cover, a filter the query engine rejects, invalid base64, a program that fails
- `images_base64` (optional): Base64-encoded images attached to the question (requires a vision-capable QA model) in `execute_code` with the error it hit, and anything unexpected with its own
message.
- **`analyze`** - Answer complex analytical questions via code execution ### Instructions
- `question` (required): The question to answer
- `filter` (optional): SQL WHERE clause to restrict document access The server publishes `instructions` describing the knowledge base: what it
- `images_base64` (optional): Base64-encoded images attached to the question (requires a vision-capable analysis model) holds, when to reach for it, the collection names when it covers several, and
- Best for aggregation, computation, and multi-document analysis `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 ## Continuous ingestion

View file

@ -35,7 +35,8 @@ snapshot is only meaningful while this process is the only writer.
## Storage ## Storage
LanceDB is embedded, so there is no server. The same code runs against a local LanceDB is embedded, so there is no server. The same code runs against a local
directory, S3, GCS, Azure or LanceDB Cloud by changing `lancedb.uri`. directory, S3, GCS, Azure or LanceDB Cloud by changing a database's location in
`lancedb.databases`.
Tables are versioned. Vacuum collapses old versions on a retention window, and 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 [tags](cli.md) name a state across all tables so a database can be restored to

View file

@ -27,7 +27,7 @@ async with HaikuRAG("path/to/database.lancedb", read_only=True) as client:
`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. `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 !!! note
Databases must be explicitly created with `create=True` or via `haiku-rag init` before use. Opening a nonexistent unnamed local database raises `FileNotFoundError`, naming its path; one named in `lancedb.databases` raises `SourceUnavailableError`, which names the database rather than its location. 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 !!! 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`. 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`.
@ -268,8 +268,8 @@ On the constructor `sources=[]` means something else. Passing `sources` alongsid
```python ```python
client.covers_multiple # whether the client covers more than one database client.covers_multiple # whether the client covers more than one database
client.source_names # configured names, in order client.source_names # database names, in order; known before the client opens
client.source # one configured name, or None for a set or unnamed database client.source # the one database's name, or None for a set
owner = await client.reader_for("papers") # the client reading that database owner = await client.reader_for("papers") # the client reading that database
papers, wiki = await client.clients_for(["papers", "wiki"]) papers, wiki = await client.clients_for(["papers", "wiki"])
@ -277,6 +277,17 @@ 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. `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 ### Filtering Search Results
Filter search results to only include chunks from documents matching specific criteria: Filter search results to only include chunks from documents matching specific criteria:

View file

@ -12,7 +12,7 @@ You also need [Ollama](https://ollama.com/) for the default embedding and answer
```bash ```bash
ollama pull qwen3-embedding:4b ollama pull qwen3-embedding:4b
ollama pull gpt-oss ollama pull qwen3.8
``` ```
!!! note "Prefer OpenAI?" !!! note "Prefer OpenAI?"

View file

@ -49,7 +49,7 @@ datasets and judge:
```bash ```bash
evaluations run hotpotqa --target rag-capability evaluations run hotpotqa --target rag-capability
evaluations run hotpotqa --target analysis-capability --capability-model ollama:gpt-oss evaluations run hotpotqa --target analysis-capability --capability-model ollama:qwen3.8
``` ```
`--capability-model "provider:name"` overrides the capability model independently from `--capability-model "provider:name"` overrides the capability model independently from

View file

@ -49,6 +49,13 @@ async def evaluate_dataset(
if document_filter is not None: if document_filter is not None:
console.print(f"Document filter: {document_filter}", style="dim") 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 not skip_db:
if spec.uses_configured_databases(config, db_path): if spec.uses_configured_databases(config, db_path):
raise ValueError( raise ValueError(
@ -157,7 +164,11 @@ def run(
config: Path | None = typer.Option( config: Path | None = typer.Option(
None, "--config", help="Path to haiku.rag YAML config file." 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( skip_db: bool = typer.Option(
False, "--skip-db", help="Skip updating the evaluation db." False, "--skip-db", help="Skip updating the evaluation db."
), ),

View file

@ -44,7 +44,7 @@ class CapabilityRunResult:
cited_uris: list[str] = field(default_factory=list) cited_uris: list[str] = field(default_factory=list)
cited_chunk_ids: 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. # The database each cited chunk came from, in the order they were cited.
# Empty string where the database is unnamed. # Empty string for a citation built without a source.
cited_sources: list[str] = field(default_factory=list) cited_sources: list[str] = field(default_factory=list)
searched_uris: list[str] = field(default_factory=list) searched_uris: list[str] = field(default_factory=list)
n_searches: int = 0 n_searches: int = 0

View file

@ -89,9 +89,9 @@ class DatasetSpec:
) -> bool: ) -> bool:
"""Whether `lancedb.databases` places the databases to evaluate over. """Whether `lancedb.databases` places the databases to evaluate over.
A path names one database and wins over the configuration, both when it `--db PATH` places the database where the configuration places none;
comes from `--db` and when the client resolves it. True for a mapping of `evaluate_dataset` refuses the two together. True for a mapping of one,
one, which is a configured database like any other and keeps its name. which is a configured database like any other and keeps its name.
""" """
return bool(config.lancedb.databases) and override_path is None return bool(config.lancedb.databases) and override_path is None

View file

@ -2,7 +2,7 @@
name = "haiku.rag-evals" name = "haiku.rag-evals"
description = "Benchmarking and evaluation scripts for haiku.rag" description = "Benchmarking and evaluation scripts for haiku.rag"
version = "0.79.0" version = "0.82.1"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" } license = { text = "MIT" }
requires-python = ">=3.12" requires-python = ">=3.12"

View file

@ -711,6 +711,37 @@ class TestEvaluateDatasetJudgeModel:
assert mock_qa.call_args[1]["judge_model"] is custom_judge assert mock_qa.call_args[1]["judge_model"] is custom_judge
@pytest.mark.asyncio
async def test_a_db_path_beside_configured_databases_is_refused(tmp_path) -> None:
"""`--db` places the database where the configuration places none; beside
`lancedb.databases` the run refuses before touching anything."""
from haiku.rag.config.models import LanceDBConfig
config = AppConfig(
lancedb=LanceDBConfig(databases={"alpha": str(tmp_path / "a.lancedb")})
)
spec = DatasetSpec(
key="test",
db_filename="test.lancedb",
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
document_mapper=lambda doc: None,
qa_loader=lambda: [], # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
qa_case_builder=lambda idx, doc: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
)
with pytest.raises(ValueError, match="alpha"):
await evaluate_dataset(
spec=spec,
config=config,
skip_db=True,
skip_retrieval=True,
skip_qa=True,
limit=None,
name=None,
db_path=tmp_path / "other.lancedb",
)
class TestExperimentMetadataTargets: class TestExperimentMetadataTargets:
def test_default_target_is_rag_capability(self) -> None: def test_default_target_is_rag_capability(self) -> None:
result = build_experiment_metadata( result = build_experiment_metadata(

View file

@ -485,8 +485,8 @@ def test_records_the_database_each_citation_came_from():
assert result.cited_sources == ["alpha", "beta", "alpha"] assert result.cited_sources == ["alpha", "beta", "alpha"]
def test_an_unnamed_database_records_no_source(): def test_a_hand_built_citation_without_a_source_records_an_empty_string():
"""One database names nothing: the field holds 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._base import EvidenceState
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
from haiku.rag.store.models.citation import Citation from haiku.rag.store.models.citation import Citation

View file

@ -26,8 +26,8 @@ uv run python examples/custom_agent.py /path/to/db.lancedb
**Script:** `custom_agent_agui.py` **Script:** `custom_agent_agui.py`
A Starlette app that adapts a native RAG-capable agent to AG-UI. A Starlette app that adapts a native RAG-capable agent to AG-UI. The configuration places the database (`HAIKU_RAG_CONFIG_PATH`, or `./haiku.rag.yaml`):
```bash ```bash
DB_PATH=/path/to/db.lancedb uv run uvicorn examples.custom_agent_agui:app --reload --port 8000 uv run uvicorn examples.custom_agent_agui:app --reload --port 8000
``` ```

View file

@ -8,13 +8,13 @@ Requirements:
- An Anthropic API key (for the QA model) or adjust the model below - An Anthropic API key (for the QA model) or adjust the model below
Usage: Usage:
DB_PATH=/path/to/db.lancedb uv run uvicorn examples.custom_agent_agui:app --reload --port 8000 uv run uvicorn examples.custom_agent_agui:app --reload --port 8000
The configuration places the database (HAIKU_RAG_CONFIG_PATH, or
./haiku.rag.yaml).
""" """
import os
import sys
from dataclasses import dataclass, field from dataclasses import dataclass, field
from pathlib import Path
from typing import Any from typing import Any
from ag_ui.core import EventType, StateSnapshotEvent from ag_ui.core import EventType, StateSnapshotEvent
@ -30,14 +30,7 @@ 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.policy import create_capability as citation_policy
from haiku.rag.capabilities.rag import RAGState, create_capability from haiku.rag.capabilities.rag import RAGState, create_capability
db_path = os.environ.get("DB_PATH") capability = create_capability(defer_loading=False)
if not db_path:
print(
"Set DB_PATH environment variable to your haiku.rag database", file=sys.stderr
)
sys.exit(1)
capability = create_capability(db_path=Path(db_path), defer_loading=False)
@dataclass @dataclass

View file

@ -103,7 +103,6 @@ services:
"haiku-rag", "haiku-rag",
"--config", "--config",
"/app/haiku.rag.yaml", "/app/haiku.rag.yaml",
"--read-only",
"mcp", "mcp",
"--host", "--host",
"0.0.0.0", "0.0.0.0",

View file

@ -1,5 +1,4 @@
import logging import logging
from functools import cached_property
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
@ -64,21 +63,10 @@ class HaikuRAGApp:
[ref] = self.scope.databases [ref] = self.scope.databases
return ref return ref
@cached_property
def _connection(self) -> "tuple[AppConfig, Path]":
"""How to open the one database this command works on, directly.
Derived per database from its configured location.
"""
from haiku.rag.client.session import default_db_path
config, db_path = self._one.connection(self.config)
return config, db_path or default_db_path(config)
@property @property
def _store_config(self) -> AppConfig: def _location(self) -> "Path | str":
"""The configuration for opening the one database directly.""" """Where the one database this command works on is."""
return self._connection[0] return self._one.location
@property @property
def _is_local(self) -> bool: def _is_local(self) -> bool:
@ -91,17 +79,14 @@ class HaikuRAGApp:
@property @property
def _path(self) -> Path: def _path(self) -> Path:
"""The path of the one database this command works on. """The path of the one local database this command works on."""
assert self._one.db_path is not None
A database behind a URI has none of its own, and the default stands in: return self._one.db_path
the URI in `_store_config` is what decides where it connects.
"""
return self._connection[1]
@property @property
def display_path(self) -> "Path | str": def display_path(self) -> "Path | str":
"""What a one-database command calls the database it opened.""" """What a one-database command calls the database it opened."""
return self._one.db_path or self._one.uri return self._one.location
@property @property
def database_missing(self) -> bool: def database_missing(self) -> bool:
@ -140,7 +125,7 @@ class HaikuRAGApp:
self.console.print("[red]Database path does not exist.[/red]") self.console.print("[red]Database path does not exist.[/red]")
return return
info = await gather_database_info(self._store_config, self._path) info = await gather_database_info(self._location, self.config)
if not info.exists: if not info.exists:
self.console.print( self.console.print(
@ -282,8 +267,8 @@ class HaikuRAGApp:
cm = status if status is not None else nullcontext() cm = status if status is not None else nullcontext()
with cm: with cm:
report = await run_doctor( report = await run_doctor(
self._store_config, self.config,
self._path, self._location,
dict(os.environ), dict(os.environ),
duplicates_out=duplicates_out, duplicates_out=duplicates_out,
on_progress=on_progress, on_progress=on_progress,
@ -340,8 +325,8 @@ class HaikuRAGApp:
return return
async with Store( async with Store(
self._path, self._location,
config=self._store_config, config=self.config,
skip_validation=True, skip_validation=True,
read_only=True, read_only=True,
skip_migration_check=True, skip_migration_check=True,
@ -415,15 +400,15 @@ class HaikuRAGApp:
""" """
from haiku.rag.store.engine import Store from haiku.rag.store.engine import Store
return Store(self._path, config=self._store_config, read_only=self.read_only) return Store(self._location, config=self.config, read_only=self.read_only)
def _tag_read_store(self) -> "Store": def _tag_read_store(self) -> "Store":
"""Read-only store for tag inspection; works on old or drifted DBs.""" """Read-only store for tag inspection; works on old or drifted DBs."""
from haiku.rag.store.engine import Store from haiku.rag.store.engine import Store
return Store( return Store(
self._path, self._location,
config=self._store_config, config=self.config,
skip_validation=True, skip_validation=True,
skip_migration_check=True, skip_migration_check=True,
read_only=True, read_only=True,
@ -760,8 +745,8 @@ class HaikuRAGApp:
from haiku.rag.store.engine import Store from haiku.rag.store.engine import Store
async with Store( async with Store(
self._path, self._location,
config=self._store_config, config=self.config,
skip_validation=True, skip_validation=True,
skip_migration_check=True, skip_migration_check=True,
read_only=self.read_only, read_only=self.read_only,
@ -949,7 +934,7 @@ class HaikuRAGApp:
# The resolved scope: a path overrides a configured URI, and a derived # The resolved scope: a path overrides a configured URI, and a derived
# single-database configuration drops the name results and citations # single-database configuration drops the name results and citations
# carry. # carry.
server = _mcp_server_covering(self.scope, self.config, self.read_only) server = _mcp_server_covering(self.scope, self.config)
try: try:
if transport == "stdio": if transport == "stdio":
await server.run_stdio_async() await server.run_stdio_async()

View file

@ -1,5 +1,4 @@
import asyncio import asyncio
import os
from dataclasses import dataclass, field, replace from dataclasses import dataclass, field, replace
from difflib import get_close_matches from difflib import get_close_matches
from pathlib import Path from pathlib import Path
@ -14,6 +13,7 @@ from pydantic_ai import (
) )
from pydantic_ai.capabilities import AbstractCapability from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import ( from pydantic_ai.messages import (
BinaryContent,
InstructionPart, InstructionPart,
ModelMessage, ModelMessage,
ModelRequest, ModelRequest,
@ -29,6 +29,7 @@ from pydantic_ai.toolsets import AgentToolset
from haiku.rag.capabilities._tools import ( from haiku.rag.capabilities._tools import (
CodeExecutionEntry, CodeExecutionEntry,
EvidenceKey,
merge_results, merge_results,
search_corpus, search_corpus,
) )
@ -43,7 +44,7 @@ from haiku.rag.store.models.citation import (
ambiguous_citation, ambiguous_citation,
resolve_citations, resolve_citations,
) )
from haiku.rag.tools.search import build_image_content_from_results from haiku.rag.tools.search import PictureKey, build_image_content_from_results
CITATION_GRACE_REQUESTS = 2 CITATION_GRACE_REQUESTS = 2
"""Requests calling this capability's tools that its cite tool outlives the rest by. """Requests calling this capability's tools that its cite tool outlives the rest by.
@ -61,6 +62,15 @@ Calibration knob. Two unrelated UUID4s reach about 0.5, while dropping or
duplicating a character or a whole group stays above 0.75, so the gap is wide. duplicating a character or a whole group stays above 0.75, so the gap is wide.
""" """
FREE_SIBLINGS_PER_ROUND = 3
"""Searches one budget unit covers when emitted in the same model response.
Calibration knob, sized to the measured modal burst. ``qa.max_searches``
counts units, so a model rephrasing its query a few times in one response
spends one unit, while every search of a sequential searcher is a unit of its
own.
"""
def _ambiguous_retry(error: AmbiguousCitationError) -> ModelRetry: def _ambiguous_retry(error: AmbiguousCitationError) -> ModelRetry:
"""The only way out of an id that names a chunk in two databases. """The only way out of an id that names a chunk in two databases.
@ -88,12 +98,7 @@ def _nearest_known_id(chunk_id: str, known_ids: list[str]) -> str:
def resolve_scope(db_path: Path | str | None, config: AppConfig) -> DatabaseScope: def resolve_scope(db_path: Path | str | None, config: AppConfig) -> DatabaseScope:
"""The databases a capability covers, resolved once at its entry point. """The databases a capability covers, resolved once at its entry point."""
``HAIKU_RAG_DB`` is read here and nowhere else.
"""
if db_path is None and (env_db := os.environ.get("HAIKU_RAG_DB")):
db_path = Path(env_db).expanduser()
return DatabaseScope.resolve(config, database_path=db_path) return DatabaseScope.resolve(config, database_path=db_path)
@ -177,6 +182,12 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False) resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
search_count: int = field(default=0, repr=False) search_count: int = field(default=0, repr=False)
search_step: int = field(default=0, repr=False)
"""The run_step whose searches are being priced and deduplicated."""
step_searches: int = field(default=0, repr=False)
step_rejected: bool = field(default=False, repr=False)
step_shown: set[EvidenceKey] = field(default_factory=set, repr=False)
step_pictures: set[PictureKey] = field(default_factory=set, repr=False)
request_count: int = field(default=0, repr=False) request_count: int = field(default=0, repr=False)
grace_requests_used: int = field(default=0, repr=False) grace_requests_used: int = field(default=0, repr=False)
epoch: int = field(default=0, repr=False) epoch: int = field(default=0, repr=False)
@ -226,6 +237,11 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
rag_lock=asyncio.Lock(), rag_lock=asyncio.Lock(),
resource_lock=asyncio.Lock(), resource_lock=asyncio.Lock(),
search_count=0, search_count=0,
search_step=0,
step_searches=0,
step_rejected=False,
step_shown=set(),
step_pictures=set(),
request_count=0, request_count=0,
grace_requests_used=0, grace_requests_used=0,
epoch=0, epoch=0,
@ -493,32 +509,54 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
retrieved_now=retrieved, retrieved_now=retrieved,
) )
async def _search(self, query: str, limit: int | None) -> str | ToolReturn: async def _search(
self, query: str, limit: int | None, run_step: int
) -> str | ToolReturn:
assert self.state is not None assert self.state is not None
self.search_count += 1 if run_step != self.search_step:
if self.search_count > self._max_searches: self.search_step = run_step
self.step_searches = 0
self.step_rejected = False
self.step_shown = set()
self.step_pictures = set()
self.step_searches += 1
if (self.step_searches - 1) % FREE_SIBLINGS_PER_ROUND == 0:
self.search_count += 1
if self.step_rejected or self.search_count > self._max_searches:
self.step_rejected = True
raise ToolFailed( raise ToolFailed(
"Search limit reached. Answer the question using " "Search limit reached. Answer the question using "
"the results you already have." "the results you already have."
) )
async with self.rag_lock: async with self.rag_lock:
formatted, results, include_collection = await search_corpus( formatted, results, rendered, include_collection = await search_corpus(
await self._ensure_rag(), await self._ensure_rag(),
query, query,
limit=limit, limit=limit,
document_filter=self.state.document_filter, document_filter=self.state.document_filter,
sources=self.state.sources, sources=self.state.sources,
shown=self.step_shown,
) )
parts: list[str | BinaryContent] = []
emitted: set[PictureKey] = set()
if self.vision:
parts, emitted = build_image_content_from_results(
results,
include_collection=include_collection,
exclude=self.step_pictures,
)
# Everything the search produced commits together, after formatting and
# image construction have both succeeded: a search that raises must not
# leave results citable, note evidence the model never received, or
# suppress a later sibling's results.
state = self.state state = self.state
# A model can search the same query twice with different limits, and the # A model can search the same query twice with different limits, and the
# narrower return must not drop what the wider one already showed it. # narrower return must not drop what the wider one already showed it.
merge_results(state.searches.setdefault(query, []), results) merge_results(state.searches.setdefault(query, []), results)
self._note_evidence() self._note_evidence()
if self.vision and ( self.step_shown |= rendered
parts := build_image_content_from_results( self.step_pictures |= emitted
results, include_collection=include_collection if parts:
)
):
return ToolReturn(return_value=formatted, content=parts) return ToolReturn(return_value=formatted, content=parts)
return formatted return formatted

View file

@ -1,9 +1,11 @@
from collections.abc import Iterable from collections.abc import Iterable
from collections.abc import Set as AbstractSet
from pydantic import BaseModel from pydantic import BaseModel
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.store.models.chunk import SearchResult, qualified_id from haiku.rag.store.models.chunk import SearchResult, qualified_id
from haiku.rag.tools.search import picture_keys
class CodeExecutionEntry(BaseModel): class CodeExecutionEntry(BaseModel):
@ -13,14 +15,48 @@ class CodeExecutionEntry(BaseModel):
success: bool = True success: bool = True
EvidenceKey = tuple[tuple[str | None, str | None], tuple[str, frozenset]]
"""What tells one rendered result from another: qualified id, then signature.
The qualified id comes first because the rendered string alone would conflate
identical renderings of the same chunk id held by two databases.
"""
def evidence_signature(result: SearchResult, include_collection: bool) -> tuple:
"""The rendered evidence a result shows the model, as an equivalence key.
Rank and total are held at neutral values: they vary with a result's
position, and position (like score) must not tell two renderings apart.
"""
return (
result.format_for_agent(rank=0, total=0, include_collection=include_collection),
picture_keys(result),
)
def evidence_key(result: SearchResult, include_collection: bool) -> EvidenceKey:
return (
qualified_id(result.source, result.chunk_id),
evidence_signature(result, include_collection),
)
async def search_corpus( async def search_corpus(
rag: HaikuRAG, rag: HaikuRAG,
query: str, query: str,
limit: int | None = None, limit: int | None = None,
document_filter: str | None = None, document_filter: str | None = None,
sources: list[str] | None = None, sources: list[str] | None = None,
) -> tuple[str, list[SearchResult], bool]: shown: AbstractSet[EvidenceKey] = frozenset(),
"""Search and context-expand results, and whether they name their collection.""" ) -> tuple[str, list[SearchResult], set[EvidenceKey], bool]:
"""Search and context-expand results, eliding evidence already shown.
Returns the formatted results, the full result list, the evidence keys the
formatting rendered in full, and whether results name their collection. A
result whose key is in ``shown`` keeps its slot but collapses to one line;
the result list is never filtered.
"""
results = await rag.search( results = await rag.search(
query, limit=limit, filter=document_filter, sources=sources query, limit=limit, filter=document_filter, sources=sources
) )
@ -29,13 +65,25 @@ async def search_corpus(
# two collections names them even when everything came back from one. # two collections names them even when everything came back from one.
selected = rag.source_names if sources is None else sources selected = rag.source_names if sources is None else sources
include_collection = len(set(selected)) > 1 include_collection = len(set(selected)) > 1
formatted = "\n\n---\n\n".join( rendered: set[EvidenceKey] = set()
result.format_for_agent( parts: list[str] = []
rank=index + 1, total=len(results), include_collection=include_collection total = len(results)
) for index, result in enumerate(results):
for index, result in enumerate(results) key = evidence_key(result, include_collection)
) if key in shown or key in rendered:
return formatted or "No results found.", list(results), include_collection parts.append(
f"Also matched, shown above: [{result.chunk_id}] "
f"[rank {index + 1} of {total}]"
)
else:
parts.append(
result.format_for_agent(
rank=index + 1, total=total, include_collection=include_collection
)
)
rendered.add(key)
formatted = "\n\n---\n\n".join(parts)
return formatted or "No results found.", list(results), rendered, include_collection
def merge_results( def merge_results(
@ -56,6 +104,9 @@ def merge_results(
__all__ = [ __all__ = [
"CodeExecutionEntry", "CodeExecutionEntry",
"EvidenceKey",
"evidence_key",
"evidence_signature",
"merge_results", "merge_results",
"search_corpus", "search_corpus",
] ]

View file

@ -19,7 +19,7 @@ from haiku.rag.capabilities._base import (
) )
from haiku.rag.capabilities._tools import merge_results from haiku.rag.capabilities._tools import merge_results
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.rag.sandbox import AnalysisContext, Sandbox from haiku.rag.sandbox import AnalysisContext, Sandbox, recovery_hint
STATE_NAMESPACE = "analysis" STATE_NAMESPACE = "analysis"
_CAPABILITY_ID = "haiku-rag-analysis" _CAPABILITY_ID = "haiku-rag-analysis"
@ -49,21 +49,6 @@ def multiple_collections_instructions() -> str:
return _multiple_collections_path.read_text().rstrip() return _multiple_collections_path.read_text().rstrip()
def _recovery_hint(stderr: str) -> str:
"""Name the workaround for sandbox limits models trip over repeatedly.
The instructions already say file objects are not iterable, and models write
``for line in open(...)`` regardless. Carrying the fix in the error gives
them something to act on for the retry.
"""
if "TextIOWrapper" in stderr and "not iterable" in stderr:
return (
"\n\nHint: file objects cannot be iterated here. Read lines with "
'.readlines() or .read().split("\\n").'
)
return ""
@dataclass @dataclass
class AnalysisCapability(RAGCapabilityBase[AnalysisState]): class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
"""Deferred capability for sandboxed computation over a RAG corpus.""" """Deferred capability for sandboxed computation over a RAG corpus."""
@ -139,7 +124,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
) )
if not result.success: if not result.success:
raise ToolFailed( raise ToolFailed(
f"{result.stderr}{_recovery_hint(result.stderr)}" f"{result.stderr}{recovery_hint(result.stderr)}"
f"\n\nOutput: {result.stdout}" f"\n\nOutput: {result.stdout}"
) )
return result.stdout or "No output." return result.stdout or "No output."
@ -172,7 +157,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
ctx: RunContext[Any], query: str, limit: int | None = None ctx: RunContext[Any], query: str, limit: int | None = None
) -> str | ToolReturn: ) -> str | ToolReturn:
"""Search the knowledge base for evidence to analyze.""" """Search the knowledge base for evidence to analyze."""
return await self._with_state(self._search(query, limit)) return await self._with_state(self._search(query, limit, ctx.run_step))
async def analysis_execute_code(ctx: RunContext[Any], code: str) -> Any: async def analysis_execute_code(ctx: RunContext[Any], code: str) -> Any:
"""Execute Python against the sandboxed document filesystem.""" """Execute Python against the sandboxed document filesystem."""

View file

@ -13,11 +13,11 @@ You can mix the two. The rule: always call `analysis_cite` before answering —
Execute Python code in a sandboxed interpreter. Variables persist between calls — you can build state incrementally. Use `print()` to output results. Execute Python code in a sandboxed interpreter. Variables persist between calls — you can build state incrementally. Use `print()` to output results.
Inside the code, these functions are available (use `await`): Inside the code, these functions are available (use `await`):
- `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels, picture_refs (subset of doc_item_refs labeled `picture`) - `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels, picture_refs (subset of doc_item_refs labeled `picture`), chunk_meta (the matched chunk's stored metadata, custom keys included)
- `await list_documents()` → list of dicts with keys: id, title, uri, created_at - `await list_documents()` → list of dicts with keys: id, title, uri, created_at, metadata
Available modules: `json`, `re`, `math`, `pathlib` Useful modules include `json`, `re`, `math`, `pathlib`, `datetime`, `collections`, `itertools`, `functools` and `dataclasses`. `decimal` and `statistics` do not exist.
Not supported: class inheritance and metaclasses, generators/yield, match statements, decorators, `collections`, iterating a file object (`for line in f`) Not supported: class inheritance and metaclasses, generators/yield, match statements, iterating a file object (`for line in f`)
### analysis_search ### analysis_search
Search the knowledge base directly (outside code execution). Each result has a `Type:` (paragraph, table, code, list_item, picture). When the Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text — use it directly to answer questions about figures, diagrams, charts, screenshots. Search the knowledge base directly (outside code execution). Each result has a `Type:` (paragraph, table, code, list_item, picture). When the Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text — use it directly to answer questions about figures, diagrams, charts, screenshots.
@ -39,16 +39,17 @@ All documents are mounted as a virtual filesystem at `/documents/`:
``` ```
/documents/{document_id}/ /documents/{document_id}/
metadata.json # {"id", "title", "uri", "created_at"} metadata.json # {"id", "title", "uri", "created_at", "metadata"}
content.txt # Full document text content.txt # Full document text
items.jsonl # Structured items (one JSON object per line) items.jsonl # Structured items (one JSON object per line)
chunks.jsonl # Chunks in order with their metadata (one JSON object per line)
toc.json # Section tree derived from heading_level toc.json # Section tree derived from heading_level
``` ```
`{document_id}` is an internal identifier, not the user-facing `uri` (filename, URL, etc.). When you only know a document by its URI or title, use `await list_documents()` to enumerate ids and match against `uri` / `title` — that's a single call to the host. Iterating `/documents/` and reading every `metadata.json` works too but is much slower on portal-scale corpora. `{document_id}` is an internal identifier, not the user-facing `uri` (filename, URL, etc.). When you only know a document by its URI or title, use `await list_documents()` to enumerate ids and match against `uri` / `title` — that's a single call to the host. Iterating `/documents/` and reading every `metadata.json` works too but is much slower on portal-scale corpora.
### Reading files ### Reading files
Read with `Path.read_text()` or `open()` (including `with` blocks); file objects support `.read()`, `.readline()`, and `.readlines()`. A file object cannot be iterated, so read line-wise with `.readlines()` or `.read().split("\n")` instead of `for line in f`. Files are read-only; writing raises `PermissionError`. Read with `Path.read_text()` or `open()` (including `with` blocks); file objects support `.read()`, `.readline()`, and `.readlines()`. A file object cannot be iterated, so read line-wise with `.readlines()` or `.read().split("\n")` instead of `for line in f`. Files are read-only; writing raises `PermissionError`. There is no network. A call has a time limit, named in the error when it is hit, and output past a size is cut with an `... (output truncated)` marker.
```python ```python
from pathlib import Path from pathlib import Path
@ -70,7 +71,7 @@ for line in Path(f'/documents/{doc_id}/items.jsonl').read_text().strip().split("
``` ```
### metadata.json ### metadata.json
Document metadata: `id`, `title`, `uri`, `created_at`. Document metadata: `id`, `title`, `uri`, `created_at`, and `metadata`, the keys stored with the document.
### content.txt ### content.txt
Full text content. Use for regex or keyword search across a whole document. Full text content. Use for regex or keyword search across a whole document.
@ -86,6 +87,9 @@ Each row carries:
- `chunk_ids`: chunks that contain this item — pass to `analysis_cite()` to ground an answer that read this item directly - `chunk_ids`: chunks that contain this item — pass to `analysis_cite()` to ground an answer that read this item directly
- `heading_level`: H-level for `section_header` rows; `0` on non-header rows - `heading_level`: H-level for `section_header` rows; `0` on non-header rows
### chunks.jsonl
The document's chunks in order, one JSON object per line: `chunk_id` and `metadata`, the chunk's stored metadata (`doc_item_refs`, `headings`, `labels`, `page_numbers`, and any custom keys such as paragraph or footnote numbers). To read by chunk metadata, keep the matching rows and take the `items.jsonl` rows whose `chunk_ids` name them.
### toc.json ### toc.json
Section tree derived from `heading_level`: `{"doc_id", "title", "tree": [...]}` where each node has `{self_ref, level, title, page_numbers, item_range: [start, end_exclusive], chunk_ids, children}`. `item_range` is a line slice into `items.jsonl``items[start:end]`. `chunk_ids` aggregates the citable chunks across all items in the section — pass directly to `analysis_cite()` to ground a section-scoped answer without a corpus-wide `search()` call. `tree: []` for docs with no headers. Section tree derived from `heading_level`: `{"doc_id", "title", "tree": [...]}` where each node has `{self_ref, level, title, page_numbers, item_range: [start, end_exclusive], chunk_ids, children}`. `item_range` is a line slice into `items.jsonl``items[start:end]`. `chunk_ids` aggregates the citable chunks across all items in the section — pass directly to `analysis_cite()` to ground a section-scoped answer without a corpus-wide `search()` call. `tree: []` for docs with no headers.
@ -112,6 +116,6 @@ You MUST call `analysis_cite` before producing your final answer, every time, wi
- Use `print()` to output results — the output is your only feedback - Use `print()` to output results — the output is your only feedback
- When you write code, execute it — don't describe what code would do. But not every question needs code; simple lookups are best answered by `analysis_search → analysis_cite`. - When you write code, execute it — don't describe what code would do. But not every question needs code; simple lookups are best answered by `analysis_search → analysis_cite`.
- Use `await` for all async functions inside `analysis_execute_code` (`search`, `list_documents`) - Use `await` for all async functions inside `analysis_execute_code` (`search`, `list_documents`)
- Read files with `Path.read_text()` or `open()`/`with`. For lines use `.readlines()` or `.read().split("\n")`, never `for line in f`. The `collections` module is unavailable. - Read files with `Path.read_text()` or `open()`/`with`. For lines use `.readlines()` or `.read().split("\n")`, never `for line in f`.
- Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `analysis_cite` tool separately to register citations. `cite{...}` markdown-style inline references do nothing; only an actual `analysis_cite` tool call registers a citation. - Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `analysis_cite` tool separately to register citations. `cite{...}` markdown-style inline references do nothing; only an actual `analysis_cite` tool call registers a citation.
- **Before you write your final answer, invoke the `analysis_cite` tool with the supporting chunk_ids, or with an empty list if there are none.** This is the last tool call before answering, every time. - **Before you write your final answer, invoke the `analysis_cite` tool with the supporting chunk_ids, or with an empty list if there are none.** This is the last tool call before answering, every time.

View file

@ -81,7 +81,7 @@ class RAGCapability(RAGCapabilityBase[RAGState]):
ctx: RunContext[Any], query: str, limit: int | None = None ctx: RunContext[Any], query: str, limit: int | None = None
) -> str | ToolReturn: ) -> str | ToolReturn:
"""Search the knowledge base using hybrid vector and full-text search.""" """Search the knowledge base using hybrid vector and full-text search."""
return await self._with_state(self._search(query, limit)) return await self._with_state(self._search(query, limit, ctx.run_step))
async def rag_cite(ctx: RunContext[Any], chunk_ids: list[str]) -> Any: async def rag_cite(ctx: RunContext[Any], chunk_ids: list[str]) -> Any:
"""Register exact search-result chunk IDs as citations for the answer.""" """Register exact search-result chunk IDs as citations for the answer."""

View file

@ -42,14 +42,8 @@ def run_chat(
config.qa.model = model_config config.qa.model = model_config
config.analysis.model = model_config config.analysis.model = model_config
# The capabilities read the databases the scope covers, not what the # The app opens the scope and lends that client to the capabilities, which
# configuration names: a `--db PATH` selection is outside the # read what `--db PATH` or `--db-name NAME` selected.
# configuration, and a `--db-name NAME` selection is narrower than it.
if scope.covers_multiple:
capability_config, capability_db_path = config, None
else:
capability_config, capability_db_path = scope.databases[0].connection(config)
enabled = capabilities or ["rag"] enabled = capabilities or ["rag"]
capability_list = [] capability_list = []
defer_loading = len(enabled) > 1 defer_loading = len(enabled) > 1
@ -68,8 +62,7 @@ def run_chat(
capability_list.append( capability_list.append(
create_capability( create_capability(
db_path=capability_db_path, config=config,
config=capability_config,
defer_loading=defer_loading, defer_loading=defer_loading,
vision=driving_model.vision, vision=driving_model.vision,
) )
@ -80,8 +73,7 @@ def run_chat(
capability_list.append( capability_list.append(
create_capability( create_capability(
db_path=capability_db_path, config=config,
config=capability_config,
defer_loading=defer_loading, defer_loading=defer_loading,
vision=driving_model.vision, vision=driving_model.vision,
) )

View file

@ -148,10 +148,12 @@ class ChatApp(App):
# a client whose __aenter__ failed. # a client whose __aenter__ failed.
await client.__aenter__() await client.__aenter__()
self.client = client self.client = client
# Lent to the capabilities: already the databases they were built for, # Lent to the capabilities, with the scope it covers: one connection
# and one connection per database however many capabilities read it. # per database however many capabilities read it, and the analysis
# sandbox is built over the same selection.
for capability in self._capabilities: for capability in self._capabilities:
capability.borrowed_rag = client capability.borrowed_rag = client
capability.scope = self.scope
self._agent = Agent( self._agent = Agent(
self._model, self._model,
@ -425,7 +427,8 @@ class ChatApp(App):
def on_document_filter_modal_filter_changed(self, event: Any) -> None: def on_document_filter_modal_filter_changed(self, event: Any) -> None:
"""Scope the conversation to the selection: the filter carries the ids, """Scope the conversation to the selection: the filter carries the ids,
and `sources` restricts the search to the databases the selection names. and over a set `sources` restricts the search to the databases the
selection names. One database needs no narrowing by source.
""" """
from haiku.rag.tools.filters import build_document_id_filter from haiku.rag.tools.filters import build_document_id_filter
@ -434,10 +437,9 @@ class ChatApp(App):
doc_filter = build_document_id_filter( doc_filter = build_document_id_filter(
sorted({doc_id for _, doc_id in event.selected}) sorted({doc_id for _, doc_id in event.selected})
) )
selected_sources = {source for source, _ in event.selected} selected_sources = sorted({source for source, _ in event.selected if source})
sources: list[str] | None = None covers_multiple = self.client is not None and self.client.covers_multiple
if selected_sources and None not in selected_sources: sources = selected_sources if covers_multiple and selected_sources else None
sources = sorted(s for s in selected_sources if s is not None)
for namespace, state_type in ( for namespace, state_type in (
(RAG_STATE_NAMESPACE, RAGState), (RAG_STATE_NAMESPACE, RAGState),
(ANALYSIS_STATE_NAMESPACE, AnalysisState), (ANALYSIS_STATE_NAMESPACE, AnalysisState),

View file

@ -24,15 +24,18 @@ class DocumentCheckbox(Checkbox):
self.doc_id = doc_id self.doc_id = doc_id
def _labelled(docs) -> list[tuple[str, str | None, str]]: def _labelled(
"""Each document's label, database and id, sorted by label. The database is docs, *, name_database: bool = False
named alongside the title, which a title alone does not say. Labels are ) -> list[tuple[str, str | None, str]]:
escaped: titles and database names are data, not Textual markup.""" """Each document's label, database and id, sorted by label. Across several
databases the database is named alongside the title, which a title alone
does not say. Labels are escaped: titles and database names are data, not
Textual markup."""
rows = [ rows = [
( (
escape( escape(
f"{doc.title or doc.uri or doc.id}" f"{doc.title or doc.uri or doc.id}"
+ (f" ({doc.source})" if doc.source else "") + (f" ({doc.source})" if name_database and doc.source else "")
), ),
doc.source, doc.source,
doc.id, doc.id,
@ -221,7 +224,9 @@ class DocumentFilterModal(ModalScreen):
DocumentCheckbox( DocumentCheckbox(
label, source, doc_id, value=(source, doc_id) in self._selected label, source, doc_id, value=(source, doc_id) in self._selected
) )
for label, source, doc_id in _labelled(docs) for label, source, doc_id in _labelled(
docs, name_database=self.client.covers_multiple
)
] ]
if boxes: if boxes:
await filter_list.mount_all(boxes) await filter_list.mount_all(boxes)

View file

@ -90,9 +90,10 @@ def resolve_scope(
"""The databases a command works on, resolved once. """The databases a command works on, resolved once.
The CLI decides only what it alone knows: that `--db` and `--db-name` are The CLI decides only what it alone knows: that `--db` and `--db-name` are
the same thing said twice, and whether this command can read more than one. the same thing said twice, that a human typing `--db PATH` means that
Everything else an unknown name, a `lancedb.uri`, the default location database whatever is configured, and whether this command can read more
is `DatabaseScope.resolve`'s to answer. than one. Everything else an unknown name, the default location is
`DatabaseScope.resolve`'s to answer.
""" """
from haiku.rag.client.scope import DatabaseScope from haiku.rag.client.scope import DatabaseScope
@ -100,9 +101,12 @@ def resolve_scope(
raise AmbiguousDatabaseError( raise AmbiguousDatabaseError(
"pass --db or --db-name, not both: they name the same thing" "pass --db or --db-name, not both: they name the same thing"
) )
scope = DatabaseScope.resolve( if db is not None:
get_config(), database_name=_db_name, database_path=db try:
) return DatabaseScope.at(db)
except ValueError as error:
raise typer.BadParameter(str(error), param_hint="--db") from error
scope = DatabaseScope.resolve(get_config(), database_name=_db_name)
if scope.covers_multiple and not covers_set: if scope.covers_multiple and not covers_set:
raise AmbiguousDatabaseError( raise AmbiguousDatabaseError(
f"lancedb.databases names {', '.join(sorted(scope.names))}; this " f"lancedb.databases names {', '.join(sorted(scope.names))}; this "
@ -882,7 +886,7 @@ def mcp(
), ),
) -> None: ) -> None:
"""Run the MCP server.""" """Run the MCP server."""
app = create_app(db) app = create_app(db, covers_set=True)
transport = "stdio" if stdio else None transport = "stdio" if stdio else None

View file

@ -21,7 +21,6 @@ from haiku.rag.client.session import (
FederatedSession, FederatedSession,
SingleDatabaseSession, SingleDatabaseSession,
aclose_quietly, aclose_quietly,
default_db_path,
) )
from haiku.rag.config import AppConfig, get_config from haiku.rag.config import AppConfig, get_config
from haiku.rag.converters import get_converter from haiku.rag.converters import get_converter
@ -129,25 +128,22 @@ class HaikuRAG:
"""Initialize the RAG client with a database path. """Initialize the RAG client with a database path.
Args: Args:
db_path: Path or string path to the database. When omitted, resolves db_path: Path or string path to the database, named by its stem.
``lancedb.databases``, then ``lancedb.uri``, then the default Valid where the configuration places no database; beside
path under ``storage.data_dir``. ``lancedb.databases`` it raises ``AmbiguousDatabaseError``.
When omitted, the configured databases are covered, or the
default database ``haiku.rag`` under ``storage.data_dir``.
config: Configuration to use. Defaults to the current global config. config: Configuration to use. Defaults to the current global config.
skip_validation: Whether to skip configuration validation on database load. skip_validation: Whether to skip configuration validation on database load.
create: Whether to create the database if it doesn't exist. create: Whether to create the database if it doesn't exist.
read_only: Whether to open the database in read-only mode. read_only: Whether to open the database in read-only mode.
sources: Names from ``config.lancedb.databases`` this client covers, sources: Names of the databases this client covers, None for all of
None for all of them. Only that setting names databases, so a them. Rejected alongside ``db_path``, which says the same thing
name raises when ``lancedb.uri`` placed the database, and is
rejected alongside ``db_path``, which says the same thing
another way. ``[]`` raises too: a client over no database can do another way. ``[]`` raises too: a client over no database can do
nothing, unlike ``sources=[]`` on a search, which is a selection nothing, unlike ``sources=[]`` on a search, which is a selection
of nothing to search. of nothing to search.
""" """
self._configured = config if config is not None else get_config() self._configured = config if config is not None else get_config()
# What the caller configured, kept intact: entering derives a
# single-database configuration from it, and every re-entry derives
# from the configured set.
self._config = self._configured self._config = self._configured
self._requested_db_path = Path(db_path) if db_path is not None else None self._requested_db_path = Path(db_path) if db_path is not None else None
if self._requested_db_path is not None and sources is not None: if self._requested_db_path is not None and sources is not None:
@ -169,24 +165,31 @@ class HaikuRAG:
@property @property
def covers_multiple(self) -> bool: def covers_multiple(self) -> bool:
"""Whether this client reads from more than one database.""" """Whether this client reads from more than one database.
return isinstance(self._session, FederatedSession)
Known before the client enters: coverage is a fact of the resolved
scope.
"""
if self._session is not None:
return isinstance(self._session, FederatedSession)
return self._resolve_scope().covers_multiple
@property @property
def source_names(self) -> tuple[str, ...]: def source_names(self) -> tuple[str, ...]:
"""The configured databases this client covers, in configured order. """The databases this client covers, by name, in configured order.
A single database contributes its own name, or nothing where the Known before the client enters: coverage is a fact of the resolved
configuration named none. scope.
""" """
if isinstance(self._session, FederatedSession): if isinstance(self._session, FederatedSession):
return self._session.names return self._session.names
return () if self.source is None else (self.source,) if isinstance(self._session, SingleDatabaseSession):
return (self._session.source,)
return self._resolve_scope().names
@property @property
def source(self) -> str | None: def source(self) -> str | None:
"""The configured database this client reads, or None while covering a """The database this client reads, or None while covering a set."""
set or reading a database the configuration did not name."""
if isinstance(self._session, SingleDatabaseSession): if isinstance(self._session, SingleDatabaseSession):
return self._session.source return self._session.source
return None return None
@ -339,15 +342,12 @@ class HaikuRAG:
return self return self
[ref] = scope.databases [ref] = scope.databases
self._config, db_path = ref.connection(self._configured)
self._session = await SingleDatabaseSession( self._session = await SingleDatabaseSession(
db_path if db_path is not None else default_db_path(self._config), ref,
self._config, self._config,
skip_validation=self._skip_validation, skip_validation=self._skip_validation,
create=self._create, create=self._create,
read_only=self._read_only, read_only=self._read_only,
source=ref.name,
).open() ).open()
return self return self
@ -410,9 +410,8 @@ class HaikuRAG:
`lender` is the client that opened it, whose reranker this one borrows. `lender` is the client that opened it, whose reranker this one borrows.
""" """
client = cls( client = cls(config=session.config, read_only=session.read_only)
session.db_path, config=session.config, read_only=session.read_only client._scope = DatabaseScope((session.ref,))
)
client._session = session client._session = session
client._owns_session = False client._owns_session = False
client._lender = lender client._lender = lender
@ -858,7 +857,7 @@ class HaikuRAG:
if unknown: if unknown:
raise UnknownDatabaseError( raise UnknownDatabaseError(
f"unknown database(s) {', '.join(sorted(set(unknown)))}; this " f"unknown database(s) {', '.join(sorted(set(unknown)))}; this "
f"client covers {', '.join(sorted(covered)) or 'a single unnamed database'}" f"client covers {', '.join(sorted(covered))}"
) )
async def clients_covering( async def clients_covering(
@ -882,8 +881,8 @@ class HaikuRAG:
return [] return []
if sources != [self.source]: if sources != [self.source]:
raise UnknownDatabaseError( raise UnknownDatabaseError(
f"unknown database(s) {', '.join(sources) or '(none)'}; this " f"unknown database(s) {', '.join(sources)}; this client covers "
f"client covers {self.source or 'a single unnamed database'}" f"{self.source}"
) )
return [self] return [self]
@ -905,9 +904,6 @@ class HaikuRAG:
if not await self.clients_covering(sources): if not await self.clients_covering(sources):
return [] return []
results = await search(self, query, limit, search_type, filter, include_images) results = await search(self, query, limit, search_type, filter, include_images)
# A database named in config keeps its name even when it is the only one
# this client covers. Only an unnamed `lancedb.uri` database leaves
# source unset.
for result in results: for result in results:
result.source = self.source result.source = self.source
return results return results

View file

@ -2,7 +2,7 @@ import asyncio
import json import json
import logging import logging
from collections.abc import AsyncGenerator from collections.abc import AsyncGenerator
from datetime import datetime from datetime import UTC, datetime
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from docling_core.types.doc.document import DescriptionMetaField, PictureMeta from docling_core.types.doc.document import DescriptionMetaField, PictureMeta
@ -531,7 +531,7 @@ async def _flush_rebuild_batch(
if not documents: if not documents:
return return
now = datetime.now().isoformat() now = datetime.now(UTC).isoformat()
# Batch update documents and document_meta using merge_insert (one LanceDB # Batch update documents and document_meta using merge_insert (one LanceDB
# version per table). Content+blobs go to documents; mutable attributes go # version per table). Content+blobs go to documents; mutable attributes go

View file

@ -1,3 +1,4 @@
from collections.abc import Mapping
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
@ -8,50 +9,60 @@ from haiku.rag.store.exceptions import (
) )
from haiku.rag.utils import locate_database from haiku.rag.utils import locate_database
DEFAULT_DATABASE_FILENAME = "haiku.rag.lancedb"
def database_name(path: Path) -> str:
"""The name a database at `path` answers to: the path's stem."""
if not path.stem:
raise ValueError(f"a database at {path} has no name: the path has no stem")
return path.stem
@dataclass(frozen=True) @dataclass(frozen=True)
class DatabaseRef: class DatabaseRef:
"""A resolved database location, and the configured name it answers to. """A resolved database: the name it answers to, and where it is.
Exactly one of ``uri`` and ``db_path`` is set. ``name`` is the key from ``name`` is the key from ``lancedb.databases``, or the stem of a path the
``lancedb.databases``, and the only identity that leaves the configuration: caller gave. It is the only identity that leaves the configuration: it
it travels in results, citations and errors, where a location must not. travels in results, citations and errors, where a location must not.
None where nothing names the database. ``location`` is a local path, or a URI. ``given`` marks a path the caller
gave, whose errors may name it: the caller already knows where it is.
""" """
name: str | None name: str
uri: str location: Path | str
db_path: Path | None given: bool = False
def __post_init__(self) -> None: def __post_init__(self) -> None:
if bool(self.uri) == (self.db_path is not None): if not self.name.strip():
raise ValueError(f"a database at {self.location} has no name")
if isinstance(self.location, str):
if not self.location.strip():
raise ValueError(f"database {self.name!r} has no location")
object.__setattr__(self, "location", locate_database(self.location))
if self.given and not isinstance(self.location, Path):
raise ValueError( raise ValueError(
"a database is either a URI or a local path: " f"database {self.name!r} is given as a path, and {self.location} "
f"got uri={self.uri!r} and db_path={self.db_path!r}" "is a URI"
) )
@classmethod @classmethod
def at(cls, path: Path | str, *, name: str | None = None) -> "DatabaseRef": def at(cls, path: Path | str) -> "DatabaseRef":
"""A database at a path the caller named, taken as given.""" """A database at a path the caller named, taken as given."""
return cls(name=name, uri="", db_path=Path(path)) path = Path(path)
return cls(name=database_name(path), location=path, given=True)
@classmethod @classmethod
def configured(cls, name: str | None, location: str) -> "DatabaseRef": def configured(cls, name: str, location: str | Path) -> "DatabaseRef":
"""A database the configuration placed, by ``lancedb.uri`` or an entry in """A database the configuration placed. A location carrying a scheme is
``lancedb.databases``. A location carrying a scheme is a URI, anything a URI, anything else a local path."""
else a local path.""" return cls(name=name, location=location)
uri, db_path = locate_database(location)
return cls(name=name, uri=uri, db_path=db_path)
def connection(self, config: AppConfig) -> tuple[AppConfig, Path | None]: @property
"""The configuration and path to open this one database with. def db_path(self) -> Path | None:
"""The local path, or None for a database behind a URI."""
A copy: the caller's configuration still names whatever set it named. return self.location if isinstance(self.location, Path) else None
"""
one = config.model_copy(deep=True)
one.lancedb.databases = {}
one.lancedb.uri = self.uri
return one, self.db_path
@dataclass(frozen=True) @dataclass(frozen=True)
@ -59,10 +70,7 @@ class DatabaseScope:
"""The databases an operation covers. """The databases an operation covers.
Resolved once, from configuration plus at most one selector, then passed Resolved once, from configuration plus at most one selector, then passed
down. Never empty. down. Never empty. Nothing here reads the environment.
Nothing here reads the environment: ``HAIKU_RAG_DB`` is the capability entry
point's to honour.
""" """
databases: tuple[DatabaseRef, ...] databases: tuple[DatabaseRef, ...]
@ -71,6 +79,14 @@ class DatabaseScope:
if not self.databases: if not self.databases:
raise ValueError("a scope covers at least one database") raise ValueError("a scope covers at least one database")
@classmethod
def at(cls, path: Path | str) -> "DatabaseScope":
"""One database at a path the caller named, whatever is configured.
The CLI's ``--db``: a human typing a path means that database.
"""
return cls((DatabaseRef.at(path),))
@classmethod @classmethod
def resolve( def resolve(
cls, cls,
@ -81,9 +97,10 @@ class DatabaseScope:
) -> "DatabaseScope": ) -> "DatabaseScope":
"""The databases named by `config` and at most one selector. """The databases named by `config` and at most one selector.
A path names one database that nothing calls anything; a name selects one The configuration places databases: ``lancedb.databases``, or where it
of the configured set and keeps its name. With no selector the configured names none, the default database under ``storage.data_dir`` as the entry
set is covered in configuration order, a set of one included. ``haiku.rag``. A name selects one of them. A path places a database
where the configuration places none, and is refused beside one it does.
""" """
if database_name is not None and database_path is not None: if database_name is not None and database_path is not None:
raise AmbiguousDatabaseError( raise AmbiguousDatabaseError(
@ -91,33 +108,37 @@ class DatabaseScope:
"pass one of them" "pass one of them"
) )
declared = config.lancedb.databases configured = config.lancedb.databases
if database_path is not None: if database_path is not None:
return cls((DatabaseRef.at(database_path),)) if configured:
raise AmbiguousDatabaseError(
"a database path and lancedb.databases both place the "
f"database: db_path={Path(database_path)} and databases "
f"name {', '.join(sorted(configured))}; pass one of them"
)
return cls.at(database_path)
declared: Mapping[str, str | Path] = configured or {
"haiku.rag": config.storage.data_dir / DEFAULT_DATABASE_FILENAME
}
if database_name is not None: if database_name is not None:
if database_name not in declared: if database_name not in declared:
raise UnknownDatabaseError( raise UnknownDatabaseError(
f"unknown database {database_name!r}; lancedb.databases names " f"unknown database {database_name!r}; the databases are "
f"{', '.join(sorted(declared)) or 'nothing'}" f"{', '.join(sorted(declared))}"
) )
return cls( return cls(
(DatabaseRef.configured(database_name, declared[database_name]),) (DatabaseRef.configured(database_name, declared[database_name]),)
) )
if declared: return cls(
return cls( tuple(
tuple( DatabaseRef.configured(name, location)
DatabaseRef.configured(name, location) for name, location in declared.items()
for name, location in declared.items()
)
) )
)
if config.lancedb.uri:
return cls((DatabaseRef.configured(None, config.lancedb.uri),))
return cls((DatabaseRef.at(config.storage.data_dir / "haiku.rag.lancedb"),))
def select(self, names: list[str]) -> "DatabaseScope": def select(self, names: list[str]) -> "DatabaseScope":
"""The databases in this scope named by `names`, in the order given. """The databases in this scope named by `names`, in the order given.
@ -128,7 +149,7 @@ class DatabaseScope:
raise ValueError( raise ValueError(
"sources=[] selects no database; pass None for all of them" "sources=[] selects no database; pass None for all of them"
) )
by_name = {ref.name: ref for ref in self.databases if ref.name is not None} by_name = {ref.name: ref for ref in self.databases}
missing = [name for name in names if name not in by_name] missing = [name for name in names if name not in by_name]
if missing: if missing:
raise UnknownDatabaseError( raise UnknownDatabaseError(
@ -144,5 +165,5 @@ class DatabaseScope:
@property @property
def names(self) -> tuple[str, ...]: def names(self) -> tuple[str, ...]:
"""The configured names covered, in order. Empty where none is named.""" """The names of the databases covered, in order."""
return tuple(ref.name for ref in self.databases if ref.name is not None) return tuple(ref.name for ref in self.databases)

View file

@ -105,6 +105,11 @@ async def search_sources(
fetch_limit = _fetch_limit(client, query, limit) fetch_limit = _fetch_limit(client, query, limit)
query_vector = await _embed_query(selected[0], query, resolved) query_vector = await _embed_query(selected[0], query, resolved)
text = query if isinstance(query, str) else "" text = query if isinstance(query, str) else ""
# Embeddings are read only by cosine fusion: a reranker scores the union
# itself, and its 10x over-fetch would materialize them for nothing.
uses_cosine = query_vector is not None and (
not isinstance(query, str) or client.reranker is None
)
per_source = await gather_all( per_source = await gather_all(
*( *(
c.chunk_repository.search( c.chunk_repository.search(
@ -113,12 +118,15 @@ async def search_sources(
search_type=resolved, search_type=resolved,
filter=filter, filter=filter,
query_vector=query_vector, query_vector=query_vector,
with_vectors=uses_cosine,
) )
for c in selected for c in selected
) )
) )
ranked = await _fuse(client, selected, query, per_source, limit) ranked = await _fuse(
client, selected, query, per_source, limit, query_vector=query_vector
)
results: list[SearchResult] = [] results: list[SearchResult] = []
for owner, chunk, score in ranked: for owner, chunk, score in ranked:
@ -149,13 +157,21 @@ async def _fuse(
query: "str | bytes | PILImage.Image", query: "str | bytes | PILImage.Image",
per_source: list[list[tuple[Chunk, float]]], per_source: list[list[tuple[Chunk, float]]],
limit: int, limit: int,
query_vector: list[float] | None = None,
) -> list[tuple["HaikuRAG", Chunk, float]]: ) -> list[tuple["HaikuRAG", Chunk, float]]:
"""One ranked list from several, keeping each candidate's owner. """One ranked list from several, keeping each candidate's owner.
A configured reranker scores the union directly, which is what makes ranking A configured reranker scores the union directly, which is what makes ranking
across databases tractable: it compares query against document and does not across databases tractable: it compares query against document and does not
care where a candidate came from. Without one, reciprocal rank fusion over the care where a candidate came from. Without one, the union is ordered by
per-database rankings, since scores from separate indexes are not comparable. cosine similarity to the query vector: the databases in a selection share an
embedder, so similarity in that one space is the signal that is comparable
across databases by construction, where retrieval scores are each database's
own rank arithmetic. A search with no query vector (full-text) orders by the
retrieval score instead. In both, ties resolve by within-database rank the
candidate nothing in its own database beat wins and only a tie on both
falls to configured order. The returned score is the one the union was
ordered by, so downstream re-sorts (context expansion) preserve this order.
""" """
owned = [ owned = [
(client, chunk, score) (client, chunk, score)
@ -165,8 +181,9 @@ async def _fuse(
if not owned: if not owned:
return [] return []
# An image query has no text for a reranker to score against, and the check # The reranker interface takes a text query, so an image query skips it,
# precedes `reranker`, which builds the reranker on first access. # and the check precedes `reranker`, which builds the reranker on first
# access.
if isinstance(query, str): if isinstance(query, str):
reranker = federator.reranker reranker = federator.reranker
if reranker is not None: if reranker is not None:
@ -191,12 +208,38 @@ async def _fuse(
) )
return [(owner_of[id(chunk)], chunk, score) for chunk, score in reranked] return [(owner_of[id(chunk)], chunk, score) for chunk, score in reranked]
scored: list[tuple[float, HaikuRAG, Chunk]] = [] scored: list[tuple[float, float, HaikuRAG, Chunk]] = []
for client, candidates in zip(clients, per_source, strict=True): for client, candidates in zip(clients, per_source, strict=True):
for rank, (chunk, _) in enumerate(candidates): for rank, (chunk, score) in enumerate(candidates):
scored.append((1.0 / (_RRF_K + rank + 1), client, chunk)) scored.append((1.0 / (_RRF_K + rank + 1), score, client, chunk))
scored.sort(key=lambda item: item[0], reverse=True)
return [(client, chunk, score) for score, client, chunk in scored[:limit]] embeddings = [chunk.embedding for _, _, _, chunk in scored]
if query_vector is not None and all(e is not None for e in embeddings):
similarities = _cosine_to(query_vector, embeddings) # ty: ignore[invalid-argument-type]
scored = [
(rank_score, similarity, client, chunk)
for (rank_score, _, client, chunk), similarity in zip(
scored, similarities, strict=True
)
]
scored.sort(key=lambda item: (item[1], item[0]), reverse=True)
return [(client, chunk, score) for _, score, client, chunk in scored[:limit]]
def _cosine_to(query_vector: list[float], embeddings: list[list[float]]) -> list[float]:
"""Cosine similarity of each embedding to the query vector.
A zero-norm vector has no direction, so its similarity is 0 rather than a
division error.
"""
import numpy as np
query = np.asarray(query_vector, dtype=np.float32)
matrix = np.asarray(embeddings, dtype=np.float32)
norms = np.linalg.norm(matrix, axis=1) * np.linalg.norm(query)
with np.errstate(divide="ignore", invalid="ignore"):
similarities = np.where(norms > 0, matrix @ query / norms, 0.0)
return [float(s) for s in similarities]
# Reciprocal rank fusion's smoothing constant, the value the literature uses. # Reciprocal rank fusion's smoothing constant, the value the literature uses.
@ -268,10 +311,9 @@ async def _rank(
) -> list[tuple[Chunk, float]]: ) -> list[tuple[Chunk, float]]:
"""Order candidates and cut them to `limit`. """Order candidates and cut them to `limit`.
An image query carries no text for a reranker to score against, so its The reranker interface takes a text query, so an image query keeps the
candidates keep the vector ranking. Its type is checked before vector ranking. Its type is checked before `client.reranker`, which builds
`client.reranker`, which builds the reranker on first access and loads model the reranker on first access and loads model weights for a local one.
weights for a local one.
""" """
if not isinstance(query, str): if not isinstance(query, str):
return candidates[:limit] return candidates[:limit]

View file

@ -43,36 +43,29 @@ async def aclose_quietly(closeable: Any, what: str) -> None:
logger.debug("Closing the %s failed on teardown", what, exc_info=True) logger.debug("Closing the %s failed on teardown", what, exc_info=True)
def default_db_path(config: AppConfig) -> Path:
"""Where a database lives when its location names no path."""
return config.storage.data_dir / "haiku.rag.lancedb"
class SingleDatabaseSession: class SingleDatabaseSession:
"""One database: its store, its repositories, and their lifecycle. """One database: its store, its repositories, and their lifecycle.
Everything that needs a store lives here, so nothing above has to ask whether Everything that needs a store lives here, so nothing above has to ask whether
it has one. ``source`` is the configured name this database answers to, or it has one. Built from the resolved reference: ``source`` is the name it
None where nothing names it. answers to, and the store receives its location.
``db_path``, ``config``, ``read_only`` and ``source`` are readable: a client ``ref``, ``config``, ``read_only`` and ``source`` are readable: a client
borrowing this session reports them as its own. borrowing this session reports them as its own.
""" """
def __init__( def __init__(
self, self,
db_path: Path | str, ref: DatabaseRef,
config: AppConfig, config: AppConfig,
*, *,
skip_validation: bool = False, skip_validation: bool = False,
create: bool = False, create: bool = False,
read_only: bool = False, read_only: bool = False,
source: str | None = None,
) -> None: ) -> None:
self.db_path = db_path self.ref = ref
self.config = config self.config = config
self.read_only = read_only self.read_only = read_only
self.source = source
self._skip_validation = skip_validation self._skip_validation = skip_validation
self._create = create self._create = create
self._vacuum_tasks: set[asyncio.Task] = set() self._vacuum_tasks: set[asyncio.Task] = set()
@ -80,19 +73,25 @@ class SingleDatabaseSession:
self._vacuum_dirty = False self._vacuum_dirty = False
@property @property
def location(self) -> Path | str: def source(self) -> str:
"""Configured URI or local path for this database. return self.ref.name
Not `db_path`, which is a placeholder where a URI holds the database. @property
""" def location(self) -> Path | str:
return self.config.lancedb.uri or self.db_path """Where this database is: its path, or its URI."""
return self.ref.location
@property
def db_path(self) -> Path | None:
"""The local path, or None for a database behind a URI."""
return self.ref.db_path
async def open(self) -> "SingleDatabaseSession": async def open(self) -> "SingleDatabaseSession":
"""Connect, validate, and build the repositories.""" """Connect, validate, and build the repositories."""
failure: str | None = None failure: str | None = None
try: try:
self.store = Store( self.store = Store(
self.db_path, self.location,
config=self.config, config=self.config,
skip_validation=self._skip_validation, skip_validation=self._skip_validation,
create=self._create, create=self._create,
@ -107,20 +106,22 @@ class SingleDatabaseSession:
raise raise
except _NAMEABLE_FAILURES as error: except _NAMEABLE_FAILURES as error:
# The message keeps its remedy and gains the database's name. # The message keeps its remedy and gains the database's name.
if self.source is None: if self.ref.given:
raise raise
raise type(error)(f"database {self.source!r}: {error}") from error raise type(error)(f"database {self.source!r}: {error}") from error
except Exception as error: except Exception as error:
# Without a name there is nothing to report in the location's place. # A path the caller gave may be named: the caller knows it already.
if self.source is None: if self.ref.given:
raise raise
failure = type(error).__name__ failure = (
"does not exist; create it with `haiku-rag init` or `create=True`"
if isinstance(error, FileNotFoundError)
else f"could not be opened: {type(error).__name__}"
)
if failure is not None: if failure is not None:
# Raised outside the handler: the exception carries neither a cause # Raised outside the handler: the exception carries neither a cause
# nor a location-bearing context. # nor a location-bearing context.
raise SourceUnavailableError( raise SourceUnavailableError(f"database {self.source!r} {failure}")
f"database {self.source!r} could not be opened: {failure}"
)
self.document_repository = DocumentRepository(self.store) self.document_repository = DocumentRepository(self.store)
self.chunk_repository = ChunkRepository(self.store) self.chunk_repository = ChunkRepository(self.store)
self.document_item_repository = DocumentItemRepository(self.store) self.document_item_repository = DocumentItemRepository(self.store)
@ -266,9 +267,7 @@ class FederatedSession:
skip_validation: bool = False, skip_validation: bool = False,
read_only: bool = False, read_only: bool = False,
) -> None: ) -> None:
self._refs: dict[str, DatabaseRef] = { self._refs: dict[str, DatabaseRef] = {ref.name: ref for ref in scope.databases}
ref.name: ref for ref in scope.databases if ref.name is not None
}
self._config = config self._config = config
self._skip_validation = skip_validation self._skip_validation = skip_validation
self._read_only = read_only self._read_only = read_only
@ -309,14 +308,11 @@ class FederatedSession:
Registered here because a cancelled `gather` discards its results. Registered here because a cancelled `gather` discards its results.
""" """
ref = self._refs[name]
one, db_path = ref.connection(self._config)
self._sessions[name] = await SingleDatabaseSession( self._sessions[name] = await SingleDatabaseSession(
db_path if db_path is not None else default_db_path(one), self._refs[name],
one, self._config,
skip_validation=self._skip_validation, skip_validation=self._skip_validation,
read_only=self._read_only, read_only=self._read_only,
source=ref.name,
).open() ).open()
async def aclose(self) -> None: async def aclose(self) -> None:

View file

@ -40,7 +40,7 @@ class ModelConfig(ConfigModel):
""" """
provider: str = "ollama" provider: str = "ollama"
name: str = "gpt-oss" name: str = "qwen3.8"
base_url: str | None = None base_url: str | None = None
api_key: str | None = None api_key: str | None = None
@ -102,12 +102,12 @@ class LanceDBConfig(ConfigModel):
The cache sizes are per process, since the session is shared across The cache sizes are per process, since the session is shared across
connections. connections.
`databases` maps a name to a location, for searching multiple at once. The `databases` maps a name to a location, a local path or a URI, and is the one
name is what results and citations carry, so a location never leaves the way to place databases. The name is what results and citations carry, so a
configuration. Mutually exclusive with `uri`. location never leaves the configuration. Empty means the default database,
`haiku.rag`, under `storage.data_dir`.
""" """
uri: str = ""
api_key: str = "" api_key: str = ""
region: str = "" region: str = ""
storage_options: dict[str, str] = Field(default_factory=dict) storage_options: dict[str, str] = Field(default_factory=dict)
@ -116,13 +116,23 @@ class LanceDBConfig(ConfigModel):
index_cache_size_bytes: int | None = Field(default=None, ge=0) index_cache_size_bytes: int | None = Field(default=None, ge=0)
metadata_cache_size_bytes: int | None = Field(default=None, ge=0) metadata_cache_size_bytes: int | None = Field(default=None, ge=0)
@model_validator(mode="after") @model_validator(mode="before")
def _one_way_of_naming_databases(self) -> "LanceDBConfig": @classmethod
if self.uri and self.databases: def _uri_names_its_replacement(cls, data: Any) -> Any:
if isinstance(data, dict) and "uri" in data:
if str(data["uri"]).strip():
raise ValueError(
"lancedb.uri was removed; write lancedb.databases: {NAME: "
f"{data['uri']!r}}} instead"
)
raise ValueError( raise ValueError(
"lancedb.uri and lancedb.databases are mutually exclusive: " "lancedb.uri was removed; remove the empty key. With no "
"use uri for one unnamed location, or databases for named ones" "lancedb.databases the database is haiku.rag under storage.data_dir"
) )
return data
@model_validator(mode="after")
def _every_database_is_named_and_placed(self) -> "LanceDBConfig":
for name, location in self.databases.items(): for name, location in self.databases.items():
# A blank name is falsy, so source routing reads it as absent; a # A blank name is falsy, so source routing reads it as absent; a
# blank location resolves to the working directory. # blank location resolves to the working directory.
@ -155,9 +165,10 @@ class QAConfig(ConfigModel):
model: ModelConfig = Field( model: ModelConfig = Field(
default_factory=lambda: ModelConfig( default_factory=lambda: ModelConfig(
provider="ollama", provider="ollama",
name="gpt-oss", name="qwen3.8",
enable_thinking=True, enable_thinking=True,
temperature=0.3, temperature=0.3,
vision=True,
) )
) )
max_searches: int = Field(default=5, ge=0) max_searches: int = Field(default=5, ge=0)
@ -206,7 +217,8 @@ class PictureDescriptionConfig(ConfigModel):
model: ModelConfig = Field( model: ModelConfig = Field(
default_factory=lambda: ModelConfig( default_factory=lambda: ModelConfig(
provider="ollama", provider="ollama",
name="ministral-3", name="qwen3.8",
enable_thinking=False,
temperature=0.0, temperature=0.0,
) )
) )
@ -293,7 +305,7 @@ class ProcessingConfig(ConfigModel):
title_model: ModelConfig = Field( title_model: ModelConfig = Field(
default_factory=lambda: ModelConfig( default_factory=lambda: ModelConfig(
provider="ollama", provider="ollama",
name="gpt-oss", name="qwen3.8",
enable_thinking=False, enable_thinking=False,
temperature=0.3, temperature=0.3,
max_tokens=100, max_tokens=100,
@ -306,6 +318,7 @@ class SearchConfig(ConfigModel):
max_context_chars: int = Field(default=5000, gt=0) max_context_chars: int = Field(default=5000, gt=0)
vector_index_metric: Literal["cosine", "l2"] = "cosine" vector_index_metric: Literal["cosine", "l2"] = "cosine"
vector_refine_factor: int = Field(default=30, gt=0) vector_refine_factor: int = Field(default=30, gt=0)
vector_nprobes: int = Field(default=20, gt=0)
class OllamaConfig(ConfigModel): class OllamaConfig(ConfigModel):

View file

@ -32,6 +32,8 @@ In both cases:
- Results without doc_item_refs pass through unexpanded - Results without doc_item_refs pass through unexpanded
""" """
from typing import Any
from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import DocumentItem from haiku.rag.store.models.document_item import DocumentItem
@ -488,3 +490,77 @@ def expand_with_items(
final_results.append(built) final_results.append(built)
return final_results + passthrough return final_results + passthrough
def build_toc(
items: list["DocumentItem"],
chunk_index: dict[str, list[str]],
) -> list[dict[str, Any]]:
"""Build a nested section tree from items in position order.
Each ``section_header`` with ``heading_level > 0`` becomes a node. Nesting
follows the explicit levels: a header pops the stack until the top is at
a strictly shallower level, then becomes a child of that top (or a root).
``item_range = [start, end_exclusive]`` indexes the position-ordered item
list, which is the line numbering of the sandbox's ``items.jsonl``: ``start``
is the header's index and ``end_exclusive`` the index of the next header
whose level is the same or shallower (the next sibling or ancestor that
ends this section), or the item count if no such header exists. Indices,
not positions: positions may have gaps.
``chunk_ids`` aggregates the chunks covered by all items in the section's
``item_range`` (deduped, order preserved). Pass directly to ``cite()`` to
ground a section-scoped answer without a corpus-wide ``search()`` call.
Items without a section_header label (or with ``heading_level == 0``) are
skipped. When all section_headers carry the same level the output is a
flat sibling list (see docling-project/docling#2121 for an upstream case
where every PDF section_header is emitted at level=1).
"""
# Defensive: every consumer is supposed to pass items in position order,
# but the end_exclusive lookahead below silently miscomputes section
# boundaries if it's not — better to sort once than trust the caller.
items = sorted(items, key=lambda i: i.position)
header_indices = [
idx
for idx, i in enumerate(items)
if i.label == "section_header" and i.heading_level > 0
]
if not header_indices:
return []
ends: list[int] = []
for n, idx in enumerate(header_indices):
end = len(items)
for later in header_indices[n + 1 :]:
if items[later].heading_level <= items[idx].heading_level:
end = later
break
ends.append(end)
roots: list[dict[str, Any]] = []
stack: list[tuple[int, dict[str, Any]]] = []
for idx, end in zip(header_indices, ends, strict=True):
h = items[idx]
seen: set[str] = set()
chunk_ids: list[str] = []
for item in items[idx:end]:
for cid in chunk_index.get(item.self_ref, []):
if cid not in seen:
seen.add(cid)
chunk_ids.append(cid)
node: dict[str, Any] = {
"self_ref": h.self_ref,
"level": h.heading_level,
"title": h.text,
"page_numbers": list(h.page_numbers),
"item_range": [idx, end],
"chunk_ids": chunk_ids,
"children": [],
}
while stack and stack[-1][0] >= h.heading_level:
stack.pop()
(stack[-1][1]["children"] if stack else roots).append(node)
stack.append((h.heading_level, node))
return roots

View file

@ -42,6 +42,20 @@ def vlm_api_headers(model: "ModelConfig") -> dict[str, str]:
return {} return {}
def vlm_api_params(model: "ModelConfig", max_tokens: int) -> dict[str, object]:
"""Request body fields docling posts alongside the picture."""
from haiku.rag.utils import reasoning_effort
params: dict[str, object] = {
"model": model.name,
"max_completion_tokens": max_tokens,
}
effort = reasoning_effort(model)
if effort is not None:
params["reasoning_effort"] = effort
return params
class DocumentConverter(ABC): class DocumentConverter(ABC):
"""Abstract base class for document converters. """Abstract base class for document converters.

View file

@ -12,6 +12,7 @@ from haiku.rag.config import AppConfig
from haiku.rag.converters.base import ( from haiku.rag.converters.base import (
DocumentConverter, DocumentConverter,
vlm_api_headers, vlm_api_headers,
vlm_api_params,
vlm_api_url, vlm_api_url,
) )
from haiku.rag.converters.text_utils import TextFileHandler, docling_safe_name from haiku.rag.converters.text_utils import TextFileHandler, docling_safe_name
@ -153,10 +154,7 @@ class DoclingLocalConverter(DocumentConverter):
pipeline_options.picture_description_options = PictureDescriptionApiOptions( pipeline_options.picture_description_options = PictureDescriptionApiOptions(
url=AnyUrl(vlm_api_url(self.config, pic_desc.model)), url=AnyUrl(vlm_api_url(self.config, pic_desc.model)),
headers=vlm_api_headers(pic_desc.model), headers=vlm_api_headers(pic_desc.model),
params=dict( params=vlm_api_params(pic_desc.model, pic_desc.max_tokens),
model=pic_desc.model.name,
max_completion_tokens=pic_desc.max_tokens,
),
prompt=self.config.prompts.picture_description, prompt=self.config.prompts.picture_description,
timeout=pic_desc.timeout, timeout=pic_desc.timeout,
) )

View file

@ -9,6 +9,7 @@ from haiku.rag.config import AppConfig
from haiku.rag.converters.base import ( from haiku.rag.converters.base import (
DocumentConverter, DocumentConverter,
vlm_api_headers, vlm_api_headers,
vlm_api_params,
vlm_api_url, vlm_api_url,
) )
from haiku.rag.converters.text_utils import TextFileHandler, docling_safe_name from haiku.rag.converters.text_utils import TextFileHandler, docling_safe_name
@ -110,10 +111,7 @@ class DoclingServeConverter(DocumentConverter):
picture_description_api = { picture_description_api = {
"url": vlm_api_url(self.config, pic_desc.model), "url": vlm_api_url(self.config, pic_desc.model),
"headers": vlm_api_headers(pic_desc.model), "headers": vlm_api_headers(pic_desc.model),
"params": { "params": vlm_api_params(pic_desc.model, pic_desc.max_tokens),
"model": pic_desc.model.name,
"max_completion_tokens": pic_desc.max_tokens,
},
"prompt": prompt, "prompt": prompt,
"timeout": pic_desc.timeout, "timeout": pic_desc.timeout,
} }

View file

@ -1080,7 +1080,7 @@ async def run_provider_checks(
async def run_doctor( async def run_doctor(
config: AppConfig, config: AppConfig,
db_path: Path, location: Path | str,
environ: dict[str, str], environ: dict[str, str],
duplicates_out: Path | None = None, duplicates_out: Path | None = None,
on_progress: Callable[[str], None] | None = None, on_progress: Callable[[str], None] | None = None,
@ -1092,7 +1092,7 @@ async def run_doctor(
""" """
notify = on_progress or (lambda _label: None) notify = on_progress or (lambda _label: None)
notify("Inspecting tables") notify("Inspecting tables")
db = await connect_lancedb(config, db_path) db = await connect_lancedb(location, config)
stats = await get_database_stats(db) stats = await get_database_stats(db)
results: list[CheckResult] = [] results: list[CheckResult] = []
@ -1110,7 +1110,7 @@ async def run_doctor(
missing = [name for name in REQUIRED_TABLES if not stats[name]["exists"]] missing = [name for name in REQUIRED_TABLES if not stats[name]["exists"]]
if not missing: if not missing:
async with Store( async with Store(
db_path, location,
config=config, config=config,
skip_validation=True, skip_validation=True,
read_only=True, read_only=True,

View file

@ -1,6 +1,5 @@
from fastapi import APIRouter, Depends, HTTPException, status from fastapi import APIRouter, Depends, HTTPException, status
from haiku.rag.client.session import default_db_path
from haiku.rag.ingester.api.server import APIState, get_state from haiku.rag.ingester.api.server import APIState, get_state
from haiku.rag.store.info import DatabaseInfo, gather_database_info from haiku.rag.store.info import DatabaseInfo, gather_database_info
@ -21,5 +20,4 @@ async def database(state: APIState = Depends(get_state)) -> DatabaseInfo:
detail="database not configured", detail="database not configured",
) )
[ref] = state.scope.databases [ref] = state.scope.databases
one, db_path = ref.connection(state.config) return await gather_database_info(ref.location, state.config)
return await gather_database_info(one, db_path or default_db_path(one))

View file

@ -4,7 +4,6 @@ import signal
from collections.abc import Callable from collections.abc import Callable
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from datetime import UTC, datetime from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from pydantic import BaseModel from pydantic import BaseModel
@ -22,6 +21,8 @@ from haiku.rag.ingester.workers.retry import RetryPolicy
if TYPE_CHECKING: if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncEngine from sqlalchemy.ext.asyncio import AsyncEngine
from haiku.rag.client.scope import DatabaseScope
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
_MANIFEST_EXTRA_KEY = "_manifest" _MANIFEST_EXTRA_KEY = "_manifest"
@ -72,14 +73,14 @@ class IngesterApp:
WorkerPool, and a HaikuRAG client for the worker pool to ingest through. WorkerPool, and a HaikuRAG client for the worker pool to ingest through.
""" """
def __init__(self, *, config: AppConfig, db_path: Path | None = None): def __init__(self, *, config: AppConfig, scope: "DatabaseScope | None" = None):
"""The ingester over the database `scope` covers, or the one the
configuration places when no scope is handed in."""
from haiku.rag.client.scope import DatabaseScope from haiku.rag.client.scope import DatabaseScope
from haiku.rag.store.exceptions import AmbiguousDatabaseError from haiku.rag.store.exceptions import AmbiguousDatabaseError
self._config = config self._config = config
# `--db` is an explicit override; None leaves placement to the self._scope = scope if scope is not None else DatabaseScope.resolve(config)
# configuration.
self._scope = DatabaseScope.resolve(config, database_path=db_path)
if self._scope.covers_multiple: if self._scope.covers_multiple:
raise AmbiguousDatabaseError( raise AmbiguousDatabaseError(
"haiku-ingester writes one database, and lancedb.databases " "haiku-ingester writes one database, and lancedb.databases "

View file

@ -42,6 +42,7 @@ from haiku.rag.store.exceptions import ( # noqa: E402
) )
if TYPE_CHECKING: if TYPE_CHECKING:
from haiku.rag.client.scope import DatabaseScope
from haiku.rag.ingester.app import BatchProgress, BatchProgressCallback from haiku.rag.ingester.app import BatchProgress, BatchProgressCallback
_cli = typer.Typer( _cli = typer.Typer(
@ -218,6 +219,19 @@ def _load_manifest(path: Path) -> BatchManifest:
return BatchManifest.model_validate(data) return BatchManifest.model_validate(data)
def _scope_for(db: Path | None) -> "DatabaseScope | None":
"""`--db PATH` is the operator's explicit override: that database, whatever
is configured. None leaves placement to the configuration."""
from haiku.rag.client.scope import DatabaseScope
if db is None:
return None
try:
return DatabaseScope.at(db)
except ValueError as error:
raise typer.BadParameter(str(error), param_hint="--db") from error
@_cli.command("serve") @_cli.command("serve")
def serve( def serve(
db: Path | None = typer.Option( db: Path | None = typer.Option(
@ -259,7 +273,7 @@ def serve(
app_config.ingester.api.port = port app_config.ingester.api.port = port
if root_path is not None: if root_path is not None:
app_config.ingester.api.root_path = root_path app_config.ingester.api.root_path = root_path
app = IngesterApp(config=app_config, db_path=db) app = IngesterApp(config=app_config, scope=_scope_for(db))
asyncio.run(app.serve(api=not no_api)) asyncio.run(app.serve(api=not no_api))
@ -318,7 +332,7 @@ async def _run_batch(
) -> None: ) -> None:
from haiku.rag.ingester.app import IngesterApp from haiku.rag.ingester.app import IngesterApp
app = IngesterApp(config=app_config, db_path=db_path) app = IngesterApp(config=app_config, scope=_scope_for(db_path))
if dry_run: if dry_run:
report = await app.run_batch_dry_run() report = await app.run_batch_dry_run()
if report.failed_sweeps: if report.failed_sweeps:

View file

@ -20,13 +20,12 @@ async def database_lines(client: "HaikuRAG") -> list[str]:
Reported through the connection the client already holds. A failure becomes Reported through the connection the client already holds. A failure becomes
a line of the report, and the other databases still report. a line of the report, and the other databases still report.
""" """
from haiku.rag.store.engine import ConnectionMode
from haiku.rag.store.info import get_database_stats from haiku.rag.store.info import get_database_stats
lines: list[str] = [] lines: list[str] = []
db_path = client.store.db_path db_path = client.store.db_path
if client.store._connection_mode == ConnectionMode.LOCAL and not db_path.exists(): if db_path is not None and not db_path.exists():
return ["[red]Database path does not exist.[/red]"] return ["[red]Database path does not exist.[/red]"]
try: try:

View file

@ -1,66 +1,168 @@
import asyncio import asyncio
import base64
from collections.abc import AsyncIterator from collections.abc import AsyncIterator
from contextlib import AsyncExitStack, asynccontextmanager from contextlib import AsyncExitStack, asynccontextmanager
from importlib import metadata
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Annotated
from fastmcp import FastMCP from fastmcp import FastMCP
from fastmcp.exceptions import ToolError
from fastmcp.tools import ToolResult
from mcp.types import ContentBlock, ImageContent, TextContent, ToolAnnotations
from pydantic import Field
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, get_config from haiku.rag.config import AppConfig, get_config
from haiku.rag.context import build_toc
from haiku.rag.sandbox import AnalysisContext, Sandbox, recovery_hint
from haiku.rag.store.models import Document, SearchResult from haiku.rag.store.models import Document, SearchResult
from haiku.rag.tools.document import DocumentInfo from haiku.rag.store.schema import DocumentMetaRecord
from haiku.rag.utils import format_citations from haiku.rag.tools.document import DocumentInfo, DocumentSection, OutlineNode
from haiku.rag.tools.search import collect_pictures
if TYPE_CHECKING: if TYPE_CHECKING:
from typing import Any
from haiku.rag.client.scope import DatabaseScope from haiku.rag.client.scope import DatabaseScope
from haiku.rag.store.models.document_item import DocumentItem
_FILTER_COLUMNS = ", ".join(DocumentMetaRecord.model_fields)
Filter = Annotated[
str | None,
Field(
description=(
f"SQL WHERE clause over the document columns {_FILTER_COLUMNS}, "
"restricting which documents are used. `metadata` is a JSON string, "
'so match its keys with LIKE: metadata LIKE \'%"author": "Smith"%\'. '
"Also uri LIKE '%.pdf', title = 'Q3 report'."
)
),
]
Sources = Annotated[
list[str] | None,
Field(description="Collections to use, by name. All of them by default."),
]
def _decode_images(images_base64: list[str] | None) -> list[bytes] | None: def _read_only(title: str) -> ToolAnnotations:
if not images_base64: return ToolAnnotations(title=title, read_only_hint=True, open_world_hint=False)
return None
import base64
return [base64.b64decode(b64, validate=True) for b64 in images_base64]
def _decode_image(image_base64: str) -> bytes:
try:
return base64.b64decode(image_base64, validate=True)
except ValueError as e:
# binascii.Error for characters outside the alphabet or bad padding,
# ValueError itself for non-ASCII input.
raise ToolError("Invalid base64 image") from e
def _instructions(scope: "DatabaseScope", config: AppConfig) -> str:
"""What the server is for, naming no tools: the client has every tool's
description from the listing."""
lines = [
"haiku-rag is the user's knowledge base: documents they ingested, "
"searchable by meaning and keyword, readable whole or section by section, "
"or computed across with code."
]
lines.append(
"Use it whenever a question could be answered from those documents, "
"before answering from memory, and say when it had nothing relevant."
)
if scope.covers_multiple:
lines.append(
f"It holds several collections: {', '.join(scope.names)}. Results "
"name theirs in `source`; pass `sources` to use a subset."
)
if config.prompts.domain_preamble:
lines.append(config.prompts.domain_preamble)
return "\n".join(lines)
def _search_result(results: list[SearchResult], covers_multiple: bool) -> ToolResult:
"""Results as the in-process agents read them, plus the matched chunk's
metadata, then each distinct picture as an image block labelled with its
result. No structured content: a client given both shows the model the
JSON and drops the text, or shows both."""
total = len(results)
text = "\n\n".join(
result.format_for_agent(
rank=rank,
total=total,
include_collection=covers_multiple,
include_document_id=True,
include_chunk_meta=True,
)
for rank, result in enumerate(results, 1)
)
content: list[ContentBlock] = [
TextContent(type="text", text=text or "No results found.")
]
pictures, _ = collect_pictures(results)
for source, chunk_id, self_ref, picture in pictures:
collection = f" in {source}" if covers_multiple and source else ""
content.append(
TextContent(
type="text",
text=f"Picture {self_ref} of search result [{chunk_id}]{collection}",
)
)
content.append(
ImageContent(
type="image",
data=base64.b64encode(picture.data).decode("ascii"),
mime_type="image/png",
)
)
return ToolResult(content=content)
def _node(toc: "dict[str, Any]") -> OutlineNode:
return OutlineNode(
id=toc["self_ref"],
title=toc["title"],
level=toc["level"],
page_numbers=toc["page_numbers"],
children=[_node(child) for child in toc["children"]],
)
def _find(toc: list["dict[str, Any]"], section_id: str) -> "dict[str, Any] | None":
for node in toc:
if node["self_ref"] == section_id:
return node
found = _find(node["children"], section_id)
if found is not None:
return found
return None
def create_mcp_server( def create_mcp_server(
db_path: Path | None = None, db_path: Path | None = None, config: AppConfig | None = None
config: AppConfig | None = None,
read_only: bool = False,
) -> FastMCP: ) -> FastMCP:
"""Create an MCP server over one database. """Create an MCP server over the databases the configuration places.
Args: Args:
db_path: Path to the database file, or None to let `config` place it. A db_path: Path to the database file, where `config` places none; or
path overrides a configured `lancedb.uri`: for a URI-backed None to serve the databases the configuration places. Beside
database, pass None. `lancedb.databases` a path raises `AmbiguousDatabaseError`.
config: Configuration to use. config: Configuration to use.
read_only: If True, write tools (add_document_*, delete_document) are not registered.
""" """
from haiku.rag.client.scope import DatabaseScope from haiku.rag.client.scope import DatabaseScope
config = config if config is not None else get_config() config = config if config is not None else get_config()
return _covering( return _covering(DatabaseScope.resolve(config, database_path=db_path), config)
DatabaseScope.resolve(config, database_path=db_path), config, read_only
)
def _covering(scope: "DatabaseScope", config: AppConfig, read_only: bool) -> FastMCP: def _covering(scope: "DatabaseScope", config: AppConfig) -> FastMCP:
"""An MCP server over databases someone already resolved. """An MCP server over databases someone already resolved.
Internal, as ``HaikuRAG._covering`` is: the public factory takes a path and Internal, as ``HaikuRAG._covering`` is: the public factory takes a path and
resolves it, which is its own job. A caller that resolved already passes the resolves it, which is its own job. A caller that resolved already passes the
scope, so the configured name survives, which results and citations carry as scope, so the configured name survives, which results carry as ``source``.
``source``.
""" """
from haiku.rag.store.exceptions import AmbiguousDatabaseError
if scope.covers_multiple:
raise AmbiguousDatabaseError(
"an MCP server serves one database, and this scope covers "
f"{', '.join(scope.names)}; name the one to serve"
)
client: HaikuRAG | None = None client: HaikuRAG | None = None
stack = AsyncExitStack() stack = AsyncExitStack()
client_lock = asyncio.Lock() client_lock = asyncio.Lock()
@ -76,7 +178,7 @@ def _covering(scope: "DatabaseScope", config: AppConfig, read_only: bool) -> Fas
async with client_lock: async with client_lock:
if client is None: if client is None:
client = await stack.enter_async_context( client = await stack.enter_async_context(
HaikuRAG._covering(scope, config, read_only=read_only) HaikuRAG._covering(scope, config, read_only=True)
) )
return client return client
@ -95,90 +197,53 @@ def _covering(scope: "DatabaseScope", config: AppConfig, read_only: bool) -> Fas
finally: finally:
client = None client = None
mcp = FastMCP("haiku-rag", lifespan=lifespan) # Explicit: the setting is also read from the environment, and the contract
# is that every failure reaches the client with its message.
mcp = FastMCP(
"haiku-rag",
instructions=_instructions(scope, config),
version=metadata.version("haiku.rag-slim"),
lifespan=lifespan,
mask_error_details=False,
)
# Write tools - only registered when not in read-only mode @mcp.tool(annotations=_read_only("Search documents"))
if not read_only:
@mcp.tool()
async def add_document_from_file(
file_path: str,
metadata: dict[str, Any] | None = None,
title: str | None = None,
) -> str | None:
"""Add a document to the RAG system from a file path."""
try:
rag = await _client()
result = await rag.create_document_from_source(
Path(file_path), title=title, metadata=metadata or {}
)
# Handle both single document and list of documents (directories)
if isinstance(result, list):
return result[0].id if result else None
return result.id
except Exception:
return None
@mcp.tool()
async def add_document_from_url(
url: str, metadata: dict[str, Any] | None = None, title: str | None = None
) -> str | None:
"""Add a document to the RAG system from a URL."""
try:
rag = await _client()
result = await rag.create_document_from_source(
url, title=title, metadata=metadata or {}
)
# Handle both single document and list of documents
if isinstance(result, list):
return result[0].id if result else None
return result.id
except Exception:
return None
@mcp.tool()
async def add_document_from_text(
content: str,
uri: str | None = None,
metadata: dict[str, Any] | None = None,
title: str | None = None,
) -> str | None:
"""Add a document to the RAG system from text content."""
try:
rag = await _client()
document = await rag.create_document(
content, uri, title=title, metadata=metadata or {}
)
return document.id
except Exception:
return None
@mcp.tool()
async def delete_document(document_id: str) -> bool:
"""Delete a document by its ID."""
try:
rag = await _client()
return await rag.delete_document(document_id)
except Exception:
return False
# Read tools - always registered
@mcp.tool()
async def search_documents( async def search_documents(
query: str, limit: int | None = None, include_images: bool = True query: str,
) -> list[SearchResult]: limit: int | None = None,
"""Search the RAG system for documents using hybrid search (vector similarity + full-text search). include_images: bool = True,
filter: Filter = None,
sources: Sources = None,
) -> ToolResult:
"""Search the knowledge base by meaning and keyword.
When include_images is True (default) and a picture-labeled chunk is Use this first for any question the documents might answer; it needs
in the result set, ``SearchResult.image_data`` carries base64-encoded no model and is the cheapest call. Results come best first, each with
PNG bytes keyed by self_ref. Set to False to omit the bytes from the its rank, `Document ID`, `Collection` when the server covers several,
response (smaller JSON payload for plain-text consumers). the document title, section headings, the matched chunk's metadata
when it has any, and the matching passage expanded to its section;
pass the id and collection to the document tools. Pictures in the
results follow as images, each labelled with its result. Ranks, not scores,
are the signal: scores are not comparable across queries. If nothing
relevant comes back, rephrase once or narrow with `filter` before
concluding the material is absent.
Args:
query: What to look for, in natural language or keywords.
limit: How many results to return; the server's configured default
when omitted.
include_images: Return the pictures in the results as images.
False for a smaller response.
""" """
try: rag = await _client()
rag = await _client() results = await rag.search(
return await rag.search(query, limit=limit, include_images=include_images) query,
except Exception: limit=limit,
return [] filter=filter,
include_images=include_images,
sources=sources,
)
return _search_result(await rag.expand_context(results), rag.covers_multiple)
# Image-as-query tool, only registered when the configured embedder # Image-as-query tool, only registered when the configured embedder
# supports image embeddings. Probed at server-build time when no Store is # supports image embeddings. Probed at server-build time when no Store is
@ -188,123 +253,208 @@ def _covering(scope: "DatabaseScope", config: AppConfig, read_only: bool) -> Fas
if get_embedder(config).supports_images: if get_embedder(config).supports_images:
@mcp.tool() @mcp.tool(annotations=_read_only("Search documents by image"))
async def search_documents_by_image( async def search_documents_by_image(
image_base64: str, image_base64: str,
limit: int | None = None, limit: int | None = None,
include_images: bool = True, include_images: bool = True,
) -> list[SearchResult]: filter: Filter = None,
"""Search the RAG system using an image as the query. sources: Sources = None,
) -> ToolResult:
"""Search the knowledge base with an image as the query.
``image_base64`` is a base64-encoded image (PNG/JPEG bytes). The Use this when the question is about a picture rather than words.
image is embedded via the configured multimodal embedder and the The image is embedded and matched against document text and
chunks table is searched vector-only. ``include_images`` controls figures by vector similarity alone. Results have the shape of
whether picture bytes are attached to picture-labeled results. `search_documents` results.
Args:
image_base64: The query image, PNG or JPEG bytes as base64.
limit: How many results to return; the server's configured
default when omitted.
include_images: Return the pictures in the results as images.
False for a smaller response.
""" """
import base64 raw = _decode_image(image_base64)
try:
raw = base64.b64decode(image_base64)
except Exception:
return []
try:
rag = await _client()
return await rag.search(raw, limit=limit, include_images=include_images)
except Exception:
return []
@mcp.tool()
async def get_document(document_id: str) -> Document | None:
"""Get a document by its ID."""
try:
rag = await _client() rag = await _client()
return await rag.get_document_by_id(document_id) results = await rag.search(
except Exception: raw,
return None limit=limit,
filter=filter,
include_images=include_images,
sources=sources,
)
return _search_result(
await rag.expand_context(results), rag.covers_multiple
)
@mcp.tool() @mcp.tool(annotations=_read_only("Get document"))
async def get_document(document_id: str, source: str | None = None) -> Document:
"""Read one document whole, in reading order.
Use this after a search when a passage is not enough. Returns the
document's content, title, uri and metadata. Ids come from search
results and `list_documents`.
Args:
document_id: The document's id.
source: The collection holding it. Without one every collection
is asked.
"""
rag = await _client()
document = await rag.get_document_by_id(document_id, source)
if document is None:
raise ToolError(f"No document with id {document_id!r}")
return document
async def _items_of(document_id: str, source: str | None) -> list["DocumentItem"]:
"""A document's items in reading order, from the database holding it."""
rag = await _client()
document = await rag.get_document_by_id(document_id, source)
if document is None:
raise ToolError(f"No document with id {document_id!r}")
owner = await rag.reader_for(source or document.source)
assert owner is not None, "a stored document names its database"
return await owner.document_item_repository.get_all_items(document_id)
@mcp.tool(annotations=_read_only("Document outline"))
async def get_document_outline(
document_id: str, source: str | None = None
) -> list[OutlineNode]:
"""The heading tree of a document, with page numbers.
Use this on a long document to see its structure before reading, then
pass a node's `id` to `get_document_section`. Returns the headings
nested by level; an empty list means the document has no headings,
so read it with `get_document`.
Args:
document_id: The document's id.
source: The collection holding it. Without one every collection
is asked.
"""
return [
_node(toc) for toc in build_toc(await _items_of(document_id, source), {})
]
@mcp.tool(annotations=_read_only("Document section"))
async def get_document_section(
document_id: str, section_id: str, source: str | None = None
) -> DocumentSection:
"""The text of one section of a document, subsections included.
Use this to read a part of a long document instead of the whole.
`section_id` is a node `id` from `get_document_outline`. Returns the
section's heading, page numbers and text in reading order, up to the
next heading of the same or a higher level.
Args:
document_id: The document's id.
section_id: The `id` of a node in the document's outline.
source: The collection holding it. Without one every collection
is asked.
"""
items = await _items_of(document_id, source)
node = _find(build_toc(items, {}), section_id)
if node is None:
raise ToolError(f"No section {section_id!r} in document {document_id!r}")
start, end = node["item_range"]
ordered = sorted(items, key=lambda item: item.position)
return DocumentSection(
id=node["self_ref"],
title=node["title"],
page_numbers=node["page_numbers"],
content="\n\n".join(item.text for item in ordered[start:end] if item.text),
)
@mcp.tool(annotations=_read_only("List documents"))
async def list_documents( async def list_documents(
limit: int | None = None, limit: int | None = None,
offset: int | None = None, offset: int | None = None,
filter: str | None = None, filter: Filter = None,
) -> list[DocumentInfo]: ) -> list[DocumentInfo]:
"""List all documents with optional pagination and filtering. """List what the knowledge base holds.
Use this to see which documents exist, their titles, URIs and
metadata, and so what a `filter` can match. Not a search: it returns
no passages.
Args: Args:
limit: Maximum number of documents to return. limit: How many documents to return.
offset: Number of documents to skip. offset: How many documents to skip, for paging.
filter: Optional SQL WHERE clause to filter documents.
""" """
try: rag = await _client()
rag = await _client() documents = await rag.list_documents(limit, offset, filter)
documents = await rag.list_documents(limit, offset, filter) return [
DocumentInfo(
id=doc.id,
title=doc.title or "Untitled",
uri=doc.uri or "",
created=doc.created_at.strftime("%Y-%m-%d"),
source=doc.source,
metadata=doc.metadata,
)
for doc in documents
]
return [ @mcp.tool(annotations=_read_only("Run code over the documents"))
DocumentInfo( async def execute_code(
id=doc.id, code: str, filter: Filter = None, sources: Sources = None
title=doc.title or "Untitled",
uri=doc.uri or "",
created=doc.created_at.strftime("%Y-%m-%d"),
)
for doc in documents
]
except Exception:
return []
@mcp.tool()
async def ask_question(
question: str,
cite: bool = False,
images_base64: list[str] | None = None,
) -> str: ) -> str:
"""Ask a question using the QA agent. """Run a Python program over the documents and return what it printed.
Use this when the answer is a count, an aggregate, a comparison across
many documents, a lookup by document or chunk metadata, or a pattern
over whole documents: whatever a search cannot rank. The program runs
in a sandboxed interpreter on the server. Each call is one program,
nothing carries over between calls, and `print` is the only output.
Inside the program, `/documents/{document_id}/` holds `metadata.json`
(id, title, uri, created_at, metadata), `content.txt` (the whole text),
`items.jsonl` (one item per line: self_ref, label, text, page_numbers,
heading_level, chunk_ids), `chunks.jsonl` (one chunk per line: chunk_id,
metadata) and `toc.json` (`doc_id`, `title`, `tree`; each node has
self_ref, level, title, page_numbers, item_range as a slice into
items.jsonl, chunk_ids and children; an empty tree means no headings).
Read files with `Path.read_text()` or `open()`; a file object cannot be
iterated, use `.readlines()`. `await search(query, limit=10)` returns
dicts with chunk_id, content, document_id, document_title, document_uri,
source, score, page_numbers, headings, doc_item_refs, labels,
picture_refs (the doc_item_refs that are pictures) and chunk_meta.
`await list_documents()` returns dicts with id, title, uri, created_at,
source and metadata. Both see the documents `filter` and `sources`
select. Useful modules include json, re, math, pathlib, datetime,
collections, itertools, functools and dataclasses; decimal and
statistics do not exist. No generator functions, match statements or
class inheritance.
Files are read-only, there is no network, a call has a time limit named
in the error when it is hit, and output past a size is truncated.
Map a title or URI to a document id with one `list_documents()` call
rather than reading every `metadata.json`. The files carry no `source`,
so over several collections group by the `source` of `list_documents()`
rows. For a known document's structure read its `toc.json` before
searching: `search()` ranks across every document. A hit's
`doc_item_refs` are `self_ref` values in `items.jsonl`, which places it
in its section. `chunk_ids` on items and `chunk_id` in `chunks.jsonl`
join the two files; they are not citations.
Args: Args:
question: The question to ask. code: The program. Use `await` on search and list_documents.
cite: Whether to include citations in the response.
images_base64: Base64-encoded images attached to the question
(requires a vision-capable QA model).
Returns:
The answer as a string.
""" """
rag = await _client()
sandbox = Sandbox._covering(
scope, config, AnalysisContext(filter=filter, sources=sources), rag=rag
)
try: try:
images = _decode_images(images_base64) result = await sandbox.execute(code)
rag = await _client() finally:
answer, citations = await rag.ask(question, images=images) await sandbox.close()
if cite and citations: if not result.success:
answer += "\n\n" + format_citations(citations) raise ToolError(
return answer f"{result.stderr}{recovery_hint(result.stderr)}"
except Exception as e: f"\n\nOutput: {result.stdout}"
return f"Error answering question: {e!s}" )
return result.stdout or "No output."
@mcp.tool()
async def analyze(
question: str,
filter: str | None = None,
images_base64: list[str] | None = None,
) -> str:
"""Answer complex questions using the analysis capability.
Use this for questions requiring computation, aggregation, or
structural traversal across documents. The capability can write and
execute Python code in a sandboxed interpreter.
Args:
question: The question to answer.
filter: Optional SQL WHERE clause to filter documents.
images_base64: Base64-encoded images attached to the question
(requires a vision-capable analysis model).
Returns:
The answer as a string.
"""
try:
images = _decode_images(images_base64)
rag = await _client()
result = await rag.analyze(question, filter=filter, images=images)
return result.answer
except Exception as e:
return f"Error running analysis capability: {e!s}"
return mcp return mcp

View file

@ -1,10 +1,11 @@
from haiku.rag.sandbox.dependencies import AnalysisContext from haiku.rag.sandbox.dependencies import AnalysisContext
from haiku.rag.sandbox.models import AnalysisResult from haiku.rag.sandbox.models import AnalysisResult
from haiku.rag.sandbox.sandbox import Sandbox, SandboxResult from haiku.rag.sandbox.sandbox import Sandbox, SandboxResult, recovery_hint
__all__ = [ __all__ = [
"AnalysisContext", "AnalysisContext",
"AnalysisResult", "AnalysisResult",
"Sandbox", "Sandbox",
"SandboxResult", "SandboxResult",
"recovery_hint",
] ]

View file

@ -17,8 +17,9 @@ from pydantic_monty import (
) )
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.rag.context import build_toc
from haiku.rag.sandbox.dependencies import AnalysisContext from haiku.rag.sandbox.dependencies import AnalysisContext
from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.chunk import Chunk, SearchResult
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentItem from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentItem
from haiku.rag.utils import gather_all from haiku.rag.utils import gather_all
@ -29,6 +30,9 @@ if TYPE_CHECKING:
from haiku.rag.client.scope import DatabaseScope from haiku.rag.client.scope import DatabaseScope
_MAX_HOST_CALLS = 10_000_000
@dataclass @dataclass
class SandboxResult: class SandboxResult:
"""Result of executing code in the sandbox.""" """Result of executing code in the sandbox."""
@ -38,79 +42,19 @@ class SandboxResult:
success: bool success: bool
def _build_toc( def recovery_hint(stderr: str) -> str:
items: list["DocumentItem"], """Name the workaround for sandbox limits models trip over repeatedly.
chunk_index: dict[str, list[str]],
) -> list[dict[str, Any]]:
"""Build a nested section tree from items in position order.
Each ``section_header`` with ``heading_level > 0`` becomes a node. Nesting The instructions already say file objects are not iterable, and models write
follows the explicit levels: a header pops the stack until the top is at ``for line in open(...)`` regardless. Carrying the fix in the error gives
a strictly shallower level, then becomes a child of that top (or a root). them something to act on for the retry.
``item_range = [position, end_exclusive]`` where ``end_exclusive`` is the
position of the next header whose level is the same or shallower (i.e.
the next sibling or ancestor that ends this section), or the total item
count if no such header exists.
``chunk_ids`` aggregates the chunks covered by all items in the section's
``item_range`` (deduped, order preserved). Pass directly to ``cite()`` to
ground a section-scoped answer without a corpus-wide ``search()`` call.
Items without a section_header label (or with ``heading_level == 0``) are
skipped. When all section_headers carry the same level the output is a
flat sibling list (see docling-project/docling#2121 for an upstream case
where every PDF section_header is emitted at level=1).
""" """
# Defensive: every consumer is supposed to pass items in position order, if "TextIOWrapper" in stderr and "not iterable" in stderr:
# but the end_exclusive lookahead below silently miscomputes section return (
# boundaries if it's not — better to sort once than trust the caller. "\n\nHint: file objects cannot be iterated here. Read lines with "
items = sorted(items, key=lambda i: i.position) '.readlines() or .read().split("\\n").'
headers: list[DocumentItem] = [ )
i for i in items if i.label == "section_header" and i.heading_level > 0 return ""
]
if not headers:
return []
total = max((i.position for i in items), default=-1) + 1
items_by_position: dict[int, DocumentItem] = {i.position: i for i in items}
ends: list[int] = []
for idx, h in enumerate(headers):
end = total
for j in range(idx + 1, len(headers)):
if headers[j].heading_level <= h.heading_level:
end = headers[j].position
break
ends.append(end)
roots: list[dict[str, Any]] = []
stack: list[tuple[int, dict[str, Any]]] = []
for h, end in zip(headers, ends, strict=True):
seen: set[str] = set()
chunk_ids: list[str] = []
for pos in range(h.position, end):
item = items_by_position.get(pos)
if item is None:
continue
for cid in chunk_index.get(item.self_ref, []):
if cid not in seen:
seen.add(cid)
chunk_ids.append(cid)
node: dict[str, Any] = {
"self_ref": h.self_ref,
"level": h.heading_level,
"title": h.text,
"page_numbers": list(h.page_numbers),
"item_range": [h.position, end],
"chunk_ids": chunk_ids,
"children": [],
}
while stack and stack[-1][0] >= h.heading_level:
stack.pop()
(stack[-1][1]["children"] if stack else roots).append(node)
stack.append((h.heading_level, node))
return roots
class Sandbox: class Sandbox:
@ -120,7 +64,8 @@ class Sandbox:
The interpreter runs in a subprocess worker checked out of an ``AsyncMonty`` The interpreter runs in a subprocess worker checked out of an ``AsyncMonty``
pool. External functions (search, list_documents) are called by Monty code pool. External functions (search, list_documents) are called by Monty code
using ``await`` and resolved asynchronously on the host. Documents are using ``await`` and resolved asynchronously on the host. Documents are
exposed via a virtual filesystem at ``/documents/{id}/``. exposed via a virtual filesystem at ``/documents/{id}/``: ``metadata.json``,
``content.txt``, ``items.jsonl``, ``chunks.jsonl`` and ``toc.json``.
The session persists across ``execute()`` calls within the same Sandbox The session persists across ``execute()`` calls within the same Sandbox
instance variables carry over. Call ``close()`` to return the worker to instance variables carry over. Call ``close()`` to return the worker to
@ -150,6 +95,7 @@ class Sandbox:
_doc_items: dict[str, list["DocumentItem"]] _doc_items: dict[str, list["DocumentItem"]]
_doc_chunk_index: dict[str, dict[str, list[str]]] _doc_chunk_index: dict[str, dict[str, list[str]]]
_items_jsonl_cache: dict[str, str] _items_jsonl_cache: dict[str, str]
_chunks_jsonl_cache: dict[str, str]
_toc_json_cache: dict[str, str] _toc_json_cache: dict[str, str]
_opened: "HaikuRAG | None" _opened: "HaikuRAG | None"
_pool: AsyncMonty | None _pool: AsyncMonty | None
@ -216,6 +162,7 @@ class Sandbox:
self._doc_items = {} self._doc_items = {}
self._doc_chunk_index = {} self._doc_chunk_index = {}
self._items_jsonl_cache = {} self._items_jsonl_cache = {}
self._chunks_jsonl_cache = {}
self._toc_json_cache = {} self._toc_json_cache = {}
self._pool = None self._pool = None
self._session = None self._session = None
@ -328,14 +275,43 @@ class Sandbox:
assert self._loop is not None, ( assert self._loop is not None, (
"VFS reads happen during execute(); the loop must be captured first." "VFS reads happen during execute(); the loop must be captured first."
) )
if self._deadline is not None and self._loop.time() > self._deadline: if self._past_deadline():
coro.close() coro.close()
raise TimeoutError( raise self._time_limit()
"time limit exceeded: no further document reads after "
f"{self._config.analysis.code_timeout}s"
)
return asyncio.run_coroutine_threadsafe(coro, self._loop).result() return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
def _past_deadline(self) -> bool:
return (
self._deadline is not None
and self._loop is not None
and self._loop.time() > self._deadline
)
def _time_limit(self) -> TimeoutError:
return TimeoutError(
"time limit exceeded: no further document reads or calls after "
f"{self._config.analysis.code_timeout}s"
)
def _check_deadline(self) -> None:
"""Refuse a host call once the call's time is up.
Monty's watchdog counts only time the worker spends computing, so every
host call, a file served from memory and an in-code search included,
checks the deadline before it runs.
"""
if self._past_deadline():
raise self._time_limit()
def _timed(
self, read: Callable[["PurePosixPath"], str]
) -> Callable[["PurePosixPath"], str]:
def call(path: "PurePosixPath") -> str:
self._check_deadline()
return read(path)
return call
async def _discard_session(self) -> None: async def _discard_session(self) -> None:
"""Drop a session whose worker is gone. """Drop a session whose worker is gone.
@ -371,6 +347,7 @@ class Sandbox:
context = self._context context = self._context
async def search(query: str, limit: int = 10) -> list[dict[str, Any]]: async def search(query: str, limit: int = 10) -> list[dict[str, Any]]:
self._check_deadline()
# Picture bytes are deliberately not attached to in-code search # Picture bytes are deliberately not attached to in-code search
# results: the Monty interpreter has no PIL/base64/hashlib, so the # results: the Monty interpreter has no PIL/base64/hashlib, so the
# agent's Python can't do anything with them. The driving model # agent's Python can't do anything with them. The driving model
@ -404,11 +381,13 @@ class Sandbox:
"doc_item_refs": r.doc_item_refs, "doc_item_refs": r.doc_item_refs,
"labels": r.labels, "labels": r.labels,
"picture_refs": picture_refs, "picture_refs": picture_refs,
"chunk_meta": r.chunk_meta,
} }
) )
return out return out
async def list_documents() -> list[dict[str, Any]]: async def list_documents() -> list[dict[str, Any]]:
self._check_deadline()
docs, _ = await self._documents() docs, _ = await self._documents()
return [ return [
{ {
@ -417,6 +396,7 @@ class Sandbox:
"uri": d.uri, "uri": d.uri,
"created_at": str(d.created_at), "created_at": str(d.created_at),
"source": d.source, "source": d.source,
"metadata": d.metadata,
} }
for d in docs for d in docs
] ]
@ -433,6 +413,7 @@ class Sandbox:
- metadata.json: CallbackFile (eager, small) - metadata.json: CallbackFile (eager, small)
- content.txt: CallbackFile (lazy, can be large) - content.txt: CallbackFile (lazy, can be large)
- items.jsonl: CallbackFile (lazy, bulk-cached) - items.jsonl: CallbackFile (lazy, bulk-cached)
- chunks.jsonl: CallbackFile (lazy, bulk-cached)
- toc.json: CallbackFile (lazy, bulk-cached) - toc.json: CallbackFile (lazy, bulk-cached)
""" """
files: list[CallbackFile] = [] files: list[CallbackFile] = []
@ -507,6 +488,31 @@ class Sandbox:
return read_items return read_items
def _make_chunks_reader(
did: str,
) -> Callable[["PurePosixPath"], str]:
def read_chunks(_path: "PurePosixPath") -> str:
cached = sandbox._chunks_jsonl_cache.get(did)
if cached is not None:
return cached
async def _fetch() -> list[Chunk]:
async with sandbox._connection(sandbox._owners.get(did)) as rag:
return await rag.chunk_repository.get_by_document_id(did)
chunks = sandbox._run_on_loop(_fetch())
jsonl = "\n".join(
json.dumps(
{"chunk_id": chunk.id, "metadata": chunk.metadata},
ensure_ascii=False,
)
for chunk in chunks
)
sandbox._chunks_jsonl_cache[did] = jsonl
return jsonl
return read_chunks
def _make_toc_reader( def _make_toc_reader(
did: str, did: str,
) -> Callable[["PurePosixPath"], str]: ) -> Callable[["PurePosixPath"], str]:
@ -520,7 +526,7 @@ class Sandbox:
{ {
"doc_id": did, "doc_id": did,
"title": doc_titles.get(did), "title": doc_titles.get(did),
"tree": _build_toc(items, chunk_index), "tree": build_toc(items, chunk_index),
}, },
ensure_ascii=False, ensure_ascii=False,
) )
@ -541,6 +547,7 @@ class Sandbox:
"title": doc.title, "title": doc.title,
"uri": doc.uri, "uri": doc.uri,
"created_at": str(doc.created_at), "created_at": str(doc.created_at),
"metadata": doc.metadata,
}, },
ensure_ascii=False, ensure_ascii=False,
) )
@ -550,7 +557,7 @@ class Sandbox:
files.append( files.append(
CallbackFile( CallbackFile(
f"{doc_dir}/metadata.json", f"{doc_dir}/metadata.json",
read=lambda _path, text=metadata: text, read=self._timed(lambda _path, text=metadata: text),
write=_deny_write, write=_deny_write,
) )
) )
@ -571,14 +578,21 @@ class Sandbox:
files.append( files.append(
CallbackFile( CallbackFile(
f"{doc_dir}/content.txt", f"{doc_dir}/content.txt",
read=_make_content_reader(doc_id), read=self._timed(_make_content_reader(doc_id)),
write=_deny_write, write=_deny_write,
) )
) )
files.append( files.append(
CallbackFile( CallbackFile(
f"{doc_dir}/items.jsonl", f"{doc_dir}/items.jsonl",
read=_make_items_reader(doc_id), read=self._timed(_make_items_reader(doc_id)),
write=_deny_write,
)
)
files.append(
CallbackFile(
f"{doc_dir}/chunks.jsonl",
read=self._timed(_make_chunks_reader(doc_id)),
write=_deny_write, write=_deny_write,
) )
) )
@ -589,7 +603,7 @@ class Sandbox:
files.append( files.append(
CallbackFile( CallbackFile(
f"{doc_dir}/toc.json", f"{doc_dir}/toc.json",
read=_make_toc_reader(doc_id), read=self._timed(_make_toc_reader(doc_id)),
write=_deny_write, write=_deny_write,
) )
) )
@ -601,12 +615,19 @@ class Sandbox:
Monty spends ``max_duration_secs`` across the session's whole life, and Monty spends ``max_duration_secs`` across the session's whole life, and
the session is reused so variables persist between calls: the budget the session is reused so variables persist between calls: the budget
covers the whole run. ``code_timeout`` is enforced per call elsewhere: the read covers the whole run. ``code_timeout`` is enforced per call elsewhere: past
deadline in ``_run_on_loop`` bounds a call that reads, and the pool's its deadline no further host call starts (``_check_deadline``), and the
``request_timeout`` bounds one that computes. pool's ``request_timeout`` bounds compute.
``max_suspensions`` counts host callbacks per session, document reads
included, defaults to 1000 and cannot be disabled. The time budgets are
the governors here, so it is set where no program reaches it.
""" """
analysis = self._config.analysis analysis = self._config.analysis
return {"max_duration_secs": analysis.code_timeout * analysis.max_executions} return {
"max_duration_secs": analysis.code_timeout * analysis.max_executions,
"max_suspensions": _MAX_HOST_CALLS,
}
async def _ensure_initialized(self) -> tuple[AsyncMontySession, OSAccess]: async def _ensure_initialized(self) -> tuple[AsyncMontySession, OSAccess]:
"""Check out a worker session and build the VFS on first use.""" """Check out a worker session and build the VFS on first use."""

View file

@ -38,11 +38,12 @@ class ConnectionMode(Enum):
OBJECT_STORAGE = "object_storage" OBJECT_STORAGE = "object_storage"
@staticmethod @staticmethod
def from_config(config: AppConfig) -> "ConnectionMode": def of(location: Path | str) -> "ConnectionMode":
uri = config.lancedb.uri """How a location is connected to: a path is local, `db://` is LanceDB
if not uri: Cloud, any other scheme is object storage."""
if isinstance(location, Path) or "://" not in location:
return ConnectionMode.LOCAL return ConnectionMode.LOCAL
if uri.startswith("db://"): if location.startswith("db://"):
return ConnectionMode.CLOUD return ConnectionMode.CLOUD
return ConnectionMode.OBJECT_STORAGE return ConnectionMode.OBJECT_STORAGE
@ -72,8 +73,10 @@ def _session(config: AppConfig) -> lancedb.Session:
async def connect_lancedb( async def connect_lancedb(
config: AppConfig, db_path: Path | None = None location: Path | str, config: AppConfig
) -> lancedb.AsyncConnection: ) -> lancedb.AsyncConnection:
"""Connect to the database at `location`, with the connection settings
(credentials, storage options, caches, consistency) from `config`."""
interval = config.lancedb.read_consistency_interval_seconds interval = config.lancedb.read_consistency_interval_seconds
kwargs: dict[str, Any] = { kwargs: dict[str, Any] = {
"session": _session(config), "session": _session(config),
@ -81,22 +84,19 @@ async def connect_lancedb(
timedelta(seconds=interval) if interval is not None else None timedelta(seconds=interval) if interval is not None else None
), ),
} }
mode = ConnectionMode.from_config(config) mode = ConnectionMode.of(location)
if mode == ConnectionMode.CLOUD: if mode == ConnectionMode.CLOUD:
return await lancedb.connect_async( return await lancedb.connect_async(
uri=config.lancedb.uri, uri=str(location),
api_key=config.lancedb.api_key, api_key=config.lancedb.api_key,
region=config.lancedb.region, region=config.lancedb.region,
**kwargs, **kwargs,
) )
elif mode == ConnectionMode.OBJECT_STORAGE: if mode == ConnectionMode.OBJECT_STORAGE:
if config.lancedb.storage_options: if config.lancedb.storage_options:
kwargs["storage_options"] = config.lancedb.storage_options kwargs["storage_options"] = config.lancedb.storage_options
return await lancedb.connect_async(uri=config.lancedb.uri, **kwargs) return await lancedb.connect_async(uri=str(location), **kwargs)
else: return await lancedb.connect_async(Path(location).absolute(), **kwargs)
if db_path is None:
raise ValueError("No lancedb.uri configured and no db_path provided")
return await lancedb.connect_async(db_path.absolute(), **kwargs)
def _stored_vector_dim(settings: dict) -> int | None: def _stored_vector_dim(settings: dict) -> int | None:
@ -180,14 +180,24 @@ class TagInfo:
class Store: class Store:
def __init__( def __init__(
self, self,
db_path: Path | str, location: Path | str,
config: AppConfig | None = None, config: AppConfig | None = None,
skip_validation: bool = False, skip_validation: bool = False,
create: bool = False, create: bool = False,
read_only: bool = False, read_only: bool = False,
skip_migration_check: bool = False, skip_migration_check: bool = False,
): ):
self.db_path: Path = Path(db_path) """A store over the database at `location`, a local path or a URI.
`config` supplies connection settings; where the database is comes
from `location` alone.
"""
self._location: Path | str = location
self.db_path: Path | None = (
Path(location)
if ConnectionMode.of(location) == ConnectionMode.LOCAL
else None
)
self._config = config if config is not None else get_config() self._config = config if config is not None else get_config()
self._read_only = read_only self._read_only = read_only
self._create = create self._create = create
@ -200,7 +210,7 @@ class Store:
self._rebuild_lock = asyncio.Lock() self._rebuild_lock = asyncio.Lock()
self._is_new_db = False self._is_new_db = False
if self._connection_mode == ConnectionMode.LOCAL: if self.db_path is not None:
if not self.db_path.exists(): if not self.db_path.exists():
if not create: if not create:
raise FileNotFoundError( raise FileNotFoundError(
@ -231,7 +241,7 @@ class Store:
async def _initialize(self): async def _initialize(self):
"""Perform async initialization: connect to LanceDB, init tables, validate.""" """Perform async initialization: connect to LanceDB, init tables, validate."""
self.db: lancedb.AsyncConnection = await connect_lancedb( self.db: lancedb.AsyncConnection = await connect_lancedb(
self._config, self.db_path self.location, self._config
) )
# Read once and thread onward: on object storage each of these is a # Read once and thread onward: on object storage each of these is a
@ -392,9 +402,14 @@ class Store:
needed = datetime.now() - oldest + TAG_RETENTION_MARGIN needed = datetime.now() - oldest + TAG_RETENTION_MARGIN
return max(retention, needed) return max(retention, needed)
@property
def location(self) -> Path | str:
"""Where this store connected: a local path, or a URI."""
return self._location
@property @property
def _connection_mode(self) -> ConnectionMode: def _connection_mode(self) -> ConnectionMode:
return ConnectionMode.from_config(self._config) return ConnectionMode.of(self._location)
async def _ensure_vector_index(self) -> None: async def _ensure_vector_index(self) -> None:
"""Create or rebuild vector index on chunks table. """Create or rebuild vector index on chunks table.

View file

@ -96,15 +96,15 @@ class DatabaseInfo(BaseModel):
packages: dict[str, str] = Field(default_factory=dict) packages: dict[str, str] = Field(default_factory=dict)
async def gather_database_info(config: AppConfig, db_path: Path) -> DatabaseInfo: async def gather_database_info(location: Path | str, config: AppConfig) -> DatabaseInfo:
"""Collect read-only database state without going through Store, so a """Collect read-only database state without going through Store, so a
database missing tables (e.g. pre-migration) still reports what it can.""" database missing tables (e.g. pre-migration) still reports what it can."""
from haiku.rag.store.upgrades import get_pending_upgrades from haiku.rag.store.upgrades import get_pending_upgrades
from haiku.rag.utils import get_package_versions from haiku.rag.utils import get_package_versions
display_path = config.lancedb.uri or str(db_path) display_path = str(location)
db = await connect_lancedb(config, db_path) db = await connect_lancedb(location, config)
stats = await get_database_stats(db) stats = await get_database_stats(db)
if not any(entry["exists"] for entry in stats.values()): if not any(entry["exists"] for entry in stats.values()):

View file

@ -1,3 +1,4 @@
import json
from typing import TYPE_CHECKING, Literal from typing import TYPE_CHECKING, Literal
from pydantic import BaseModel, PrivateAttr from pydantic import BaseModel, PrivateAttr
@ -143,13 +144,14 @@ class SearchResult(BaseModel):
consumers (UIs). Never part of ``format_for_agent`` output. consumers (UIs). Never part of ``format_for_agent`` output.
``chunk_meta`` is the anchor chunk's unparsed ``Chunk.metadata`` and does not ``chunk_meta`` is the anchor chunk's unparsed ``Chunk.metadata`` and does not
include the metadata of any other chunks merged with it. Never part of include the metadata of any other chunks merged with it. Left out of
``format_for_agent`` output. ``format_for_agent`` output unless ``include_chunk_meta`` asks for its custom
keys.
``source`` names the configured database a result came from: the name from ``source`` names the database a result came from: the name from
``lancedb.databases``, never a path or URI, so a location cannot travel in a ``lancedb.databases`` or a path's stem, never a path or URI, so a location
result, a citation or a log. It is None only where no database is named, as cannot travel in a result, a citation or a log. Every result a search
with the single ``lancedb.uri``. produces carries it; None only on a result built by hand.
""" """
content: str content: str
@ -202,6 +204,8 @@ class SearchResult(BaseModel):
total: int | None = None, total: int | None = None,
*, *,
include_collection: bool = False, include_collection: bool = False,
include_document_id: bool = False,
include_chunk_meta: bool = False,
) -> str: ) -> str:
"""Format this search result for inclusion in agent context. """Format this search result for inclusion in agent context.
@ -215,7 +219,11 @@ class SearchResult(BaseModel):
`include_collection` is the caller's decision, not this result's: a `include_collection` is the caller's decision, not this result's: a
search spanning one collection has nothing to distinguish, whether or search spanning one collection has nothing to distinguish, whether or
not that collection is named. not that collection is named. `include_document_id` is for a reader
that will fetch the document by id from the text alone.
`include_chunk_meta` renders the metadata stored with the matched
chunk beyond haiku.rag's own structural keys; on an expanded result it
locates the hit, not the whole passage.
""" """
if rank is not None and total is not None: if rank is not None and total is not None:
parts = [f"[{self.chunk_id}] [rank {rank} of {total}]"] parts = [f"[{self.chunk_id}] [rank {rank} of {total}]"]
@ -224,6 +232,9 @@ class SearchResult(BaseModel):
else: else:
parts = [f"[{self.chunk_id}] (score: {self.score:.2f})"] parts = [f"[{self.chunk_id}] (score: {self.score:.2f})"]
if include_document_id and self.document_id:
parts.append(f"Document ID: {self.document_id}")
if include_collection and self.source: if include_collection and self.source:
parts.append(f"Collection: {self.source}") parts.append(f"Collection: {self.source}")
@ -242,6 +253,16 @@ class SearchResult(BaseModel):
if primary_label: if primary_label:
parts.append(f"Type: {primary_label}") parts.append(f"Type: {primary_label}")
if include_chunk_meta:
custom = {
key: value
for key, value in self.chunk_meta.items()
if key not in ChunkMetadata.model_fields
}
if custom:
rendered = json.dumps(custom, ensure_ascii=False, sort_keys=True)
parts.append(f"Matched chunk metadata: {rendered}")
# Surface picture captions when present. Order matches the binary # Surface picture captions when present. Order matches the binary
# attachments emitted by build_image_content_from_results, so the model # attachments emitted by build_image_content_from_results, so the model
# can correlate caption ↔ attached image by position (BinaryContent # can correlate caption ↔ attached image by position (BinaryContent

View file

@ -24,9 +24,9 @@ class Citation(BaseModel):
``chunk_ids`` lists the ids of all chunks whose expansion ranges merged ``chunk_ids`` lists the ids of all chunks whose expansion ranges merged
into the cited result (always includes ``chunk_id``). into the cited result (always includes ``chunk_id``).
``source`` names the configured database the cited chunk came from: the name ``source`` names the database the cited chunk came from: the name from
from ``lancedb.databases``, never a path or URI. It is None only where no ``lancedb.databases`` or a path's stem, never a path or URI. None only on a
database is named, as with the single ``lancedb.uri``. citation resolved from a hand-built result.
``doc_item_refs`` are the ``self_ref`` values of every item in the cited ``doc_item_refs`` are the ``self_ref`` values of every item in the cited
content the exact items the model saw. Visual grounding resolves bounding content the exact items the model saw. Visual grounding resolves bounding

View file

@ -1,8 +1,8 @@
import json import json
from datetime import datetime from datetime import UTC, datetime
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
from pydantic import BaseModel, Field from pydantic import BaseModel, Field, field_validator
from haiku.rag.store.compression import compress_docling_split, decompress_json from haiku.rag.store.compression import compress_docling_split, decompress_json
@ -14,10 +14,10 @@ class Document(BaseModel):
""" """
Represents a document with an ID, content, and metadata. Represents a document with an ID, content, and metadata.
``source`` names the configured database a document came from: the name ``source`` names the database a document came from: the name from
from ``lancedb.databases``, never a path or URI. It is None where no ``lancedb.databases`` or a path's stem, never a path or URI. Every document
database is named, as with the single ``lancedb.uri``, and is never a database returns carries it; it is never persisted, and None only on a
persisted. document built by hand.
""" """
id: str | None = None id: str | None = None
@ -29,8 +29,13 @@ class Document(BaseModel):
docling_document: bytes | None = Field(default=None, exclude=True) docling_document: bytes | None = Field(default=None, exclude=True)
docling_pages: bytes | None = Field(default=None, exclude=True) docling_pages: bytes | None = Field(default=None, exclude=True)
docling_version: str | None = Field(default=None, exclude=True) docling_version: str | None = Field(default=None, exclude=True)
created_at: datetime = Field(default_factory=datetime.now) created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
updated_at: datetime = Field(default_factory=datetime.now) updated_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
@field_validator("created_at", "updated_at")
@classmethod
def _to_utc(cls, value: datetime) -> datetime:
return value if value.tzinfo else value.astimezone(UTC)
def set_docling(self, docling_doc: "DoclingDocument") -> None: def set_docling(self, docling_doc: "DoclingDocument") -> None:
"""Serialize and store a DoclingDocument, splitting structure and pages. """Serialize and store a DoclingDocument, splitting structure and pages.

View file

@ -240,6 +240,7 @@ class ChunkRepository:
search_type: SearchType = "hybrid", search_type: SearchType = "hybrid",
filter: str | None = None, filter: str | None = None,
query_vector: list[float] | None = None, query_vector: list[float] | None = None,
with_vectors: bool = False,
) -> list[tuple[Chunk, float]]: ) -> list[tuple[Chunk, float]]:
"""Search for relevant chunks using the specified search method. """Search for relevant chunks using the specified search method.
@ -294,6 +295,7 @@ class ChunkRepository:
.column("vector") .column("vector")
.distance_type(self.store._config.search.vector_index_metric) .distance_type(self.store._config.search.vector_index_metric)
.refine_factor(self.store._config.search.vector_refine_factor) .refine_factor(self.store._config.search.vector_refine_factor)
.nprobes(self.store._config.search.vector_nprobes)
) )
# An image query has no text to match, so it stays vector-only. # An image query has no text to match, so it stays vector-only.
if search_type != "vector" and query.strip(): if search_type != "vector" and query.strip():
@ -304,7 +306,7 @@ class ChunkRepository:
if chunk_filter is not None: if chunk_filter is not None:
results = results.where(chunk_filter) results = results.where(chunk_filter)
results = results.limit(limit) results = results.limit(limit)
return await self._process_search_results(results) return await self._process_search_results(results, with_vectors=with_vectors)
async def get_by_document_id( async def get_by_document_id(
self, self,
@ -405,7 +407,7 @@ class ChunkRepository:
return len(df) return len(df)
async def _process_search_results( async def _process_search_results(
self, query_result: "AsyncQueryBase" self, query_result: "AsyncQueryBase", with_vectors: bool = False
) -> list[tuple[Chunk, float]]: ) -> list[tuple[Chunk, float]]:
"""Process search results into chunks with document info and scores.""" """Process search results into chunks with document info and scores."""
import pandas as pd import pandas as pd
@ -456,6 +458,13 @@ class ChunkRepository:
) )
documents_map = {str(row["id"]): row for row in doc_rows} documents_map = {str(row["id"]): row for row in doc_rows}
# The query projects no columns, so the vectors are already in the
# frame; only the federated fusion path reads them, so materializing
# per-chunk lists is gated on the caller asking.
vectors = (
df["vector"].tolist() if with_vectors and "vector" in df.columns else None
)
chunks_with_scores = [] chunks_with_scores = []
for i, chunk_record in enumerate(pydantic_results): for i, chunk_record in enumerate(pydantic_results):
doc = documents_map.get(chunk_record.document_id) doc = documents_map.get(chunk_record.document_id)
@ -468,6 +477,7 @@ class ChunkRepository:
document_uri=doc["uri"] if doc else None, document_uri=doc["uri"] if doc else None,
document_title=doc["title"] if doc else None, document_title=doc["title"] if doc else None,
document_meta=json.loads(doc.get("metadata", "{}") if doc else "{}"), document_meta=json.loads(doc.get("metadata", "{}") if doc else "{}"),
embedding=list(vectors[i]) if vectors is not None else None,
) )
score = scores[i] if i < len(scores) else 1.0 score = scores[i] if i < len(scores) else 1.0
chunks_with_scores.append((chunk, score)) chunks_with_scores.append((chunk, score))

View file

@ -1,5 +1,5 @@
import json import json
from datetime import datetime from datetime import UTC, datetime
from typing import overload from typing import overload
from uuid import uuid4 from uuid import uuid4
@ -77,8 +77,12 @@ class DocumentRepository:
docling_document=doc.docling_document, docling_document=doc.docling_document,
docling_pages=doc.docling_pages, docling_pages=doc.docling_pages,
docling_version=doc.docling_version, docling_version=doc.docling_version,
created_at=datetime.fromisoformat(created) if created else datetime.now(), created_at=datetime.fromisoformat(created)
updated_at=datetime.fromisoformat(updated) if updated else datetime.now(), if created
else datetime.now(UTC),
updated_at=datetime.fromisoformat(updated)
if updated
else datetime.now(UTC),
) )
def _to_documents_record(self, entity: Document, doc_id: str) -> DocumentRecord: def _to_documents_record(self, entity: Document, doc_id: str) -> DocumentRecord:
@ -138,7 +142,7 @@ class DocumentRepository:
# document_meta) would surface. # document_meta) would surface.
if isinstance(entity, Document): if isinstance(entity, Document):
doc_id = str(uuid4()) doc_id = str(uuid4())
now = datetime.now().isoformat() now = datetime.now(UTC).isoformat()
await self.store.document_meta_table.add( await self.store.document_meta_table.add(
[self._to_meta_record(entity, doc_id, now, now)] [self._to_meta_record(entity, doc_id, now, now)]
) )
@ -159,7 +163,7 @@ class DocumentRepository:
if not documents: if not documents:
return [] return []
now = datetime.now().isoformat() now = datetime.now(UTC).isoformat()
created_at = datetime.fromisoformat(now) created_at = datetime.fromisoformat(now)
doc_records = [] doc_records = []
meta_records = [] meta_records = []
@ -272,7 +276,7 @@ class DocumentRepository:
self.store._assert_writable() self.store._assert_writable()
assert entity.id, "Document ID is required for update" assert entity.id, "Document ID is required for update"
now = datetime.now().isoformat() now = datetime.now(UTC).isoformat()
entity.updated_at = datetime.fromisoformat(now) entity.updated_at = datetime.fromisoformat(now)
created = entity.created_at.isoformat() if entity.created_at else now created = entity.created_at.isoformat() if entity.created_at else now
record = self._to_meta_record(entity, entity.id, created, now) record = self._to_meta_record(entity, entity.id, created, now)

View file

@ -77,8 +77,11 @@ async def _apply_split_document_meta(store: Store) -> None:
Exception Exception
): # pragma: no cover - defensive; stats() failure shouldn't block the split ): # pragma: no cover - defensive; stats() failure shouldn't block the split
live_bytes = 0 live_bytes = 0
free_bytes = shutil.disk_usage(store.db_path).free # A database behind a URI has no local disk to run out of.
if live_bytes and free_bytes < live_bytes: free_bytes = (
shutil.disk_usage(store.db_path).free if store.db_path is not None else None
)
if live_bytes and free_bytes is not None and free_bytes < live_bytes:
logger.warning( logger.warning(
"Skipping post-migration vacuum: need ~%.2f GB free to compact the " "Skipping post-migration vacuum: need ~%.2f GB free to compact the "
"documents table, have %.2f GB. Run `haiku-rag vacuum` once you have " "documents table, have %.2f GB. Run `haiku-rag vacuum` once you have "

View file

@ -27,6 +27,27 @@ class DocumentInfo(BaseModel):
title: str title: str
uri: str uri: str
created: str created: str
source: str | None = None
metadata: dict = {}
class OutlineNode(BaseModel):
"""A heading in a document's outline. `id` is the heading item's self_ref."""
id: str
title: str
level: int
page_numbers: list[int] = []
children: list["OutlineNode"] = []
class DocumentSection(BaseModel):
"""One section's text in reading order, subsections included."""
id: str
title: str
page_numbers: list[int] = []
content: str
class DocumentListResponse(BaseModel): class DocumentListResponse(BaseModel):

View file

@ -1,5 +1,6 @@
import base64 import base64
from collections.abc import Callable from collections.abc import Callable
from collections.abc import Set as AbstractSet
from io import BytesIO from io import BytesIO
from PIL import Image from PIL import Image
@ -20,6 +21,22 @@ their own picture had it removed, along with their text.
""" """
PictureKey = tuple[str | None, str | None, str]
"""Identity of one attached picture: (source, document_id, self_ref).
``self_ref`` alone collides across documents, and a copy of a document in
another collection carries its own pictures.
"""
def picture_keys(result: SearchResult) -> frozenset[PictureKey]:
"""The identity of every picture this result carries."""
return frozenset(
(result.source, result.document_id, self_ref)
for self_ref in (result.image_data or {})
)
def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None: def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None:
"""Wrap picture bytes for the wire, or return nothing if they will not decode. """Wrap picture bytes for the wire, or return nothing if they will not decode.
@ -35,31 +52,20 @@ def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None:
return BinaryContent(data=data, media_type="image/png", identifier=self_ref) return BinaryContent(data=data, media_type="image/png", identifier=self_ref)
def build_image_content_from_results( def collect_pictures(
results: list[SearchResult], results: list[SearchResult], exclude: AbstractSet[PictureKey] = frozenset()
include_collection: bool = False, ) -> tuple[list[tuple[str | None, str | None, str, BinaryContent]], set[PictureKey]]:
) -> list[str | BinaryContent]: """Every distinct, decodable picture attached to ``results``, in order.
"""Decode and validate picture bytes attached to search results, labelled.
Dedup keyed on ``(source, document_id, self_ref)`` so the same picture in Returns ``(source, chunk_id, self_ref, picture)`` per picture and the
different chunks is sent once, and a copy in another collection is its own. Pictures that fail ``PictureKey`` of each. Dedup keyed on ``PictureKey`` so the same picture in
``PIL.Image.verify()`` are skipped the model adapter renders one different chunks is emitted once, and a copy in another collection is its
vision placeholder per ``BinaryContent``, so emitting one for an own; ``exclude`` seeds that dedup with pictures already sent. Pictures that
image the server can't decode leaves the processor with an fail ``PIL.Image.verify()`` are skipped.
off-by-one count.
Every picture is preceded by a line naming the result it belongs to.
``ToolReturn.content`` reaches the model as a user-role message, so
retrieved pictures are otherwise indistinguishable from ones the user
attached, and models narrate them as part of the question: unlabelled,
gemma4-26b answered about a figure from an unrelated document, and with a
single note ahead of the batch it still called them "images in the prompt".
The label also names the chunk to cite for a figure, which
``BinaryContent.identifier`` cannot do it does not survive serialization
to the vision API.
""" """
collected: list[tuple[str | None, str | None, str, BinaryContent]] = [] collected: list[tuple[str | None, str | None, str, BinaryContent]] = []
seen: set[tuple[str | None, str | None, str]] = set() seen: set[PictureKey] = set(exclude)
emitted: set[PictureKey] = set()
for result in results: for result in results:
if not result.image_data: if not result.image_data:
continue continue
@ -72,7 +78,34 @@ def build_image_content_from_results(
continue continue
collected.append((result.source, result.chunk_id, self_ref, picture)) collected.append((result.source, result.chunk_id, self_ref, picture))
seen.add(key) seen.add(key)
emitted.add(key)
return collected, emitted
def build_image_content_from_results(
results: list[SearchResult],
include_collection: bool = False,
exclude: AbstractSet[PictureKey] = frozenset(),
) -> tuple[list[str | BinaryContent], set[PictureKey]]:
"""Decode and validate picture bytes attached to search results, labelled.
Returns the labelled content and the ``PictureKey`` of every picture it
emitted, as ``collect_pictures`` decides them. An undecodable picture is
skipped because the model adapter renders one vision placeholder per
``BinaryContent``, so emitting one for an image the server can't decode
leaves the processor with an off-by-one count.
Every picture is preceded by a line naming the result it belongs to.
``ToolReturn.content`` reaches the model as a user-role message, so
retrieved pictures are otherwise indistinguishable from ones the user
attached, and models narrate them as part of the question: unlabelled,
gemma4-26b answered about a figure from an unrelated document, and with a
single note ahead of the batch it still called them "images in the prompt".
The label also names the chunk to cite for a figure, which
``BinaryContent.identifier`` cannot do it does not survive serialization
to the vision API.
"""
collected, emitted = collect_pictures(results, exclude)
content: list[str | BinaryContent] = [] content: list[str | BinaryContent] = []
total = len(collected) total = len(collected)
for position, (source, chunk_id, self_ref, picture) in enumerate(collected, 1): for position, (source, chunk_id, self_ref, picture) in enumerate(collected, 1):
@ -83,7 +116,7 @@ def build_image_content_from_results(
f"Not provided by the user. {RETRIEVED_IMAGE_TAG}" f"Not provided by the user. {RETRIEVED_IMAGE_TAG}"
) )
content.append(picture) content.append(picture)
return content return content, emitted
def create_search_toolset( def create_search_toolset(
@ -174,7 +207,7 @@ def create_search_toolset(
if not config.qa.model.vision: if not config.qa.model.vision:
return text return text
image_content = build_image_content_from_results( image_content, _ = build_image_content_from_results(
results_list, include_collection=include_collection results_list, include_collection=include_collection
) )
if image_content: if image_content:

View file

@ -4,7 +4,7 @@ import sys
from collections.abc import Awaitable from collections.abc import Awaitable
from importlib import metadata from importlib import metadata
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any, NoReturn, cast from typing import TYPE_CHECKING, Any, Literal, NoReturn, cast
from packaging.version import Version, parse from packaging.version import Version, parse
@ -41,7 +41,7 @@ def parse_model_option(value: str) -> "ModelConfig":
parts = value.split(":", 1) parts = value.split(":", 1)
if len(parts) != 2 or not parts[0] or not parts[1]: if len(parts) != 2 or not parts[0] or not parts[1]:
raise ValueError( raise ValueError(
f"Invalid model format '{value}'. Expected 'provider:name' (e.g. 'ollama:gpt-oss')." f"Invalid model format '{value}'. Expected 'provider:name' (e.g. 'ollama:qwen3.8')."
) )
return ModelConfig(provider=parts[0], name=parts[1]) return ModelConfig(provider=parts[0], name=parts[1])
@ -182,6 +182,20 @@ _OPENAI_COMPAT_PROFILE: "OpenAIModelProfile" = {
} }
def reasoning_effort(
model_config: "ModelConfig",
) -> Literal["none", "low", "high"] | None:
"""OpenAI `reasoning_effort` for a model config, or None when unset.
"low" is gpt-oss's floor; its template rejects "none".
"""
if model_config.enable_thinking is None:
return None
if model_config.enable_thinking:
return "high"
return "low" if model_config.name == "gpt-oss" else "none"
def get_model( def get_model(
model_config: "ModelConfig", model_config: "ModelConfig",
app_config: "AppConfig | None" = None, app_config: "AppConfig | None" = None,
@ -213,12 +227,9 @@ def get_model(
if provider == "ollama": if provider == "ollama":
model_settings = None model_settings = None
# Apply thinking control for gpt-oss effort = reasoning_effort(model_config)
if model == "gpt-oss" and model_config.enable_thinking is not None: if effort is not None:
if model_config.enable_thinking is False: model_settings = OpenAIChatModelSettings(openai_reasoning_effort=effort)
model_settings = OpenAIChatModelSettings(openai_reasoning_effort="low")
else:
model_settings = OpenAIChatModelSettings(openai_reasoning_effort="high")
model_settings = apply_common_settings( model_settings = apply_common_settings(
model_settings, model_config, map_thinking=False model_settings, model_config, map_thinking=False
@ -382,43 +393,6 @@ def _citation_label(c: "Citation") -> str:
return c.document_title or c.document_uri return c.document_title or c.document_uri
def format_citations(citations: "list[Citation]") -> str:
"""Format citations as plain text with preserved formatting.
Used by things like the MCP server where Rich renderables are not available.
Pictures referenced by the chunk are surfaced as ``[Figure: <ref>]`` markers.
"""
if not citations:
return ""
lines = ["## Citations\n"]
for i, c in enumerate(citations):
idx = c.index if c.index is not None else (i + 1)
title = c.document_title or c.document_uri
header = f"[{idx}] {title}"
location_parts = []
pages = _citation_pages(c)
if pages:
location_parts.append(pages)
section = _citation_section(c)
if section:
location_parts.append(f"Section: {section}")
source = c.document_uri
if location_parts:
source += f" - {', '.join(location_parts)}"
lines.append(f"{header} {source}")
for ref in c.picture_refs:
lines.append(f"[Figure: {ref}]")
lines.append(c.content)
lines.append("")
return "\n".join(lines)
def truncated(text: str, limit: int) -> str: def truncated(text: str, limit: int) -> str:
"""The first `limit` characters of `text`, with `…` appended when anything """The first `limit` characters of `text`, with `…` appended when anything
was dropped. A cut result is `limit` characters plus the mark.""" was dropped. A cut result is `limit` characters plus the mark."""
@ -549,17 +523,15 @@ def raise_missing_extra(module: str, extra: str, exc: ModuleNotFoundError) -> No
) from exc ) from exc
def locate_database(location: str) -> tuple[str, Path | None]: def locate_database(location: str) -> Path | str:
"""Split a configured location into (uri, db_path). """A configured location as a URI, or as a local path.
A value with a scheme is a `lancedb.uri`; anything else is a local path. A value with a scheme is a URI, which `ConnectionMode` opens without the
`ConnectionMode` classifies a `uri` as object storage and opens it without existence check a local database gets; anything else is a local path.
the existence check a local database gets, so a local path never travels
as one.
""" """
if "://" in location: if "://" in location:
return location, None return location
return "", Path(location) return Path(location)
def get_default_data_dir() -> Path: def get_default_data_dir() -> Path:

View file

@ -2,7 +2,7 @@
name = "haiku.rag-slim" name = "haiku.rag-slim"
description = "Local-first agentic RAG with citations - hybrid search, reranking and multimodal retrieval over your own documents, no database server required - Minimal dependencies" description = "Local-first agentic RAG with citations - hybrid search, reranking and multimodal retrieval over your own documents, no database server required - Minimal dependencies"
version = "0.79.0" version = "0.82.1"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" } license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" } readme = { file = "README.md", content-type = "text/markdown" }
@ -40,12 +40,12 @@ dependencies = [
"docling-core>=2.82.0,<3.0.0", "docling-core>=2.82.0,<3.0.0",
"httpx>=0.28.1", "httpx>=0.28.1",
"jinja2>=3.1.0", "jinja2>=3.1.0",
"fastmcp>=3.3.0", "fastmcp>=4.0.2,<5.0.0",
"lancedb==0.37.1", "lancedb==0.37.1",
"pathspec>=1.0.4", "pathspec>=1.0.4",
"pydantic>=2.12.5", "pydantic>=2.12.5",
"pydantic-ai-slim[openai,logfire,ag-ui]>=2.18.0,<3.0.0", "pydantic-ai-slim[openai,logfire,ag-ui]>=2.18.0,<3.0.0",
"pydantic-monty>=0.0.19", "pydantic-monty>=0.0.23",
"pypdfium2>=5.0", "pypdfium2>=5.0",
"python-dotenv>=1.2.2", "python-dotenv>=1.2.2",
"pyyaml>=6.0.3", "pyyaml>=6.0.3",

View file

@ -0,0 +1,12 @@
{
"name": "haiku-rag",
"version": "0.82.1",
"description": "Search, read and analyze your haiku.rag knowledge base from Claude Code.",
"author": {
"name": "Yiorgis Gozadinos",
"email": "ggozadinos@gmail.com"
},
"homepage": "https://ggozad.github.io/haiku.rag/mcp/",
"repository": "https://github.com/ggozad/haiku.rag",
"license": "MIT"
}

View file

@ -0,0 +1,26 @@
{
"name": "haiku-rag",
"version": "0.82.1",
"description": "Search, read and analyze your haiku.rag knowledge base from Codex.",
"author": {
"name": "Yiorgis Gozadinos",
"email": "ggozadinos@gmail.com",
"url": "https://github.com/ggozad"
},
"homepage": "https://ggozad.github.io/haiku.rag/mcp/",
"repository": "https://github.com/ggozad/haiku.rag",
"license": "MIT",
"keywords": ["rag", "knowledge-base", "search", "documents", "mcp"],
"skills": "./skills/",
"mcpServers": "./.mcp.json",
"interface": {
"displayName": "haiku.rag",
"shortDescription": "Search and analyze your haiku.rag knowledge base",
"longDescription": "Search, read, and compute over documents in your local haiku.rag knowledge base through MCP tools.",
"developerName": "Yiorgis Gozadinos",
"category": "Productivity",
"capabilities": ["Interactive", "Read"],
"websiteURL": "https://ggozad.github.io/haiku.rag/",
"defaultPrompt": "Search my haiku.rag knowledge base and cite the relevant documents."
}
}

View file

@ -0,0 +1,8 @@
{
"mcpServers": {
"haiku-rag": {
"command": "haiku-rag",
"args": ["mcp", "--stdio"]
}
}
}

View file

@ -0,0 +1,77 @@
---
name: haiku-rag
description: Search, read and compute over the user's haiku.rag knowledge base
through the haiku-rag MCP tools. Use whenever a request could be answered
from the user's ingested documents, when asked to find, look up, check or
cite something in their documents or knowledge base, or when the question is
about the user's own material rather than general knowledge.
compatibility: Requires the haiku-rag MCP server to be registered in the client.
allowed-tools:
- mcp__plugin_haiku-rag_haiku-rag__search_documents
- mcp__plugin_haiku-rag_haiku-rag__search_documents_by_image
- mcp__plugin_haiku-rag_haiku-rag__get_document
- mcp__plugin_haiku-rag_haiku-rag__get_document_outline
- mcp__plugin_haiku-rag_haiku-rag__get_document_section
- mcp__plugin_haiku-rag_haiku-rag__list_documents
- mcp__plugin_haiku-rag_haiku-rag__execute_code
---
# Working with the knowledge base
Check the knowledge base before answering from memory whenever the question
could be about the user's documents. Say so when it has nothing relevant.
## Find
`search_documents` is the first call. Results come best first with the document
title, section headings, the matched chunk's metadata when it has any, and the
passage in its section. Pictures in the results arrive as images: answer
figure questions from them. `filter` restricts which documents are searched,
`limit` how many results come back. If it misses, rephrase once or narrow with
a filter before concluding the material is not there. When the question is
about an image rather than words and the server offers
`search_documents_by_image`, it takes the image as the query.
## Read
Every search result shows its `Document ID` (and `Collection` when there are
several); pass them to the read tools. `get_document` returns a document's
whole text in reading order. For a long one, `get_document_outline` gives the
heading tree with page numbers and `get_document_section` the text of one
section, subsections included.
## Compute
`execute_code` runs a Python program on the server over the same documents.
Under `/documents/{id}/` each has `metadata.json`, `content.txt`, `items.jsonl`,
`chunks.jsonl` and `toc.json`, and the program can `await search(query)` and
`await list_documents()`. Write code when the answer is a count, an aggregate, a
comparison across many documents, a lookup by document or chunk metadata, or a
pattern over whole documents: whatever search cannot rank. Each call is one
program and variables do not carry over, so gather, compute and `print` a
compact result in the same program. `filter` and `sources` select the documents
it sees. For a known document's structure read its `toc.json` first; `search()`
ranks across every document. Map a title or URI to an id with one
`list_documents()` call rather than reading every `metadata.json`; the files
carry no `source`, so over several collections group by its rows. Answer and
cite from what it printed.
## Explore
`list_documents` shows what is stored: titles, URIs and metadata. It is how you
learn what a filter can match.
## Filters
A SQL WHERE clause over the document columns `id`, `uri`, `title`,
`created_at`, `updated_at`, `metadata`. `metadata` is a JSON string, so match
it with LIKE: `metadata LIKE '%"author": "Smith"%'`. Also `uri LIKE '%.pdf'`,
`title = 'Q3 report'`.
## Results and citations
Rank is the signal; scores are not comparable across queries and are never
confidence. Cite the document title or URI, the section heading and page
numbers when present, and the matched chunk's metadata when it carries locators
such as paragraph or footnote numbers. When results carry `source`, the server
covers several collections: name it, and pass `sources` to search a subset.

View file

@ -2,7 +2,7 @@
name = "haiku.rag" name = "haiku.rag"
description = "Local-first agentic RAG with citations - hybrid search, reranking and multimodal retrieval over your own documents, no database server required" description = "Local-first agentic RAG with citations - hybrid search, reranking and multimodal retrieval over your own documents, no database server required"
version = "0.79.0" version = "0.82.1"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" } license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" } readme = { file = "README.md", content-type = "text/markdown" }
@ -37,7 +37,7 @@ classifiers = [
] ]
dependencies = [ dependencies = [
"haiku.rag-slim[docling,voyageai,cohere,zeroentropy,tui,cross-encoder,jina]==0.79.0", "haiku.rag-slim[docling,voyageai,cohere,zeroentropy,tui,cross-encoder,jina]==0.82.1",
] ]
[project.urls] [project.urls]
@ -52,9 +52,9 @@ haiku-rag = "haiku.rag.cli:cli"
[project.optional-dependencies] [project.optional-dependencies]
tui = ["textual>=8.2.4"] tui = ["textual>=8.2.4"]
s3 = ["haiku.rag-slim[s3]==0.79.0"] s3 = ["haiku.rag-slim[s3]==0.82.1"]
cross-encoder = ["haiku.rag-slim[cross-encoder]==0.79.0"] cross-encoder = ["haiku.rag-slim[cross-encoder]==0.82.1"]
ingester = ["haiku.rag-slim[ingester]==0.79.0"] ingester = ["haiku.rag-slim[ingester]==0.82.1"]
[build-system] [build-system]
requires = ["hatchling"] requires = ["hatchling"]

View file

@ -2,7 +2,8 @@
""" """
Version bumping script for haiku.rag workspace. Version bumping script for haiku.rag workspace.
Updates version in all pyproject.toml files and CHANGELOG.md. Updates version in all pyproject.toml files, both plugin manifests, and
CHANGELOG.md.
""" """
import re import re
@ -54,6 +55,19 @@ def update_example_dependencies(file_path: Path, new_version: str) -> None:
print(f"✓ Updated example dependencies in {file_path.relative_to(Path.cwd())}") print(f"✓ Updated example dependencies in {file_path.relative_to(Path.cwd())}")
def update_plugin_version(file_path: Path, new_version: str) -> None:
"""Update the version in a plugin manifest."""
content = file_path.read_text()
updated = re.sub(
r'^(\s*"version": )"[^"]+"',
rf'\1"{new_version}"',
content,
flags=re.MULTILINE,
)
file_path.write_text(updated)
print(f"✓ Updated {file_path.relative_to(Path.cwd())}")
def update_changelog(changelog_path: Path, new_version: str) -> None: def update_changelog(changelog_path: Path, new_version: str) -> None:
"""Update CHANGELOG.md with new version.""" """Update CHANGELOG.md with new version."""
content = changelog_path.read_text() content = changelog_path.read_text()
@ -122,10 +136,16 @@ def main():
root / "app" / "backend" / "pyproject.toml", root / "app" / "backend" / "pyproject.toml",
] ]
plugin_files = [
root / "plugins" / "haiku-rag" / ".claude-plugin" / "plugin.json",
root / "plugins" / "haiku-rag" / ".codex-plugin" / "plugin.json",
]
changelog_file = root / "CHANGELOG.md" changelog_file = root / "CHANGELOG.md"
# Check all files exist # Check all files exist
for file in pyproject_files + example_pyproject_files + [changelog_file]: for file in (
pyproject_files + example_pyproject_files + plugin_files + [changelog_file]
):
if not file.exists(): if not file.exists():
print(f"Error: {file} not found") print(f"Error: {file} not found")
sys.exit(1) sys.exit(1)
@ -155,6 +175,9 @@ def main():
for file in example_pyproject_files: for file in example_pyproject_files:
update_example_dependencies(file, new_version) update_example_dependencies(file, new_version)
for file in plugin_files:
update_plugin_version(file, new_version)
# Update CHANGELOG.md # Update CHANGELOG.md
update_changelog(changelog_file, new_version) update_changelog(changelog_file, new_version)

View file

@ -98,14 +98,10 @@ def _placed(capability) -> "Path | None":
return ref.db_path return ref.db_path
def test_capability_factories_resolve_environment_and_defaults( def test_capability_factories_resolve_defaults(temp_db_path, monkeypatch):
temp_db_path, monkeypatch """The configuration places the database; the environment plays no part."""
):
config = AppConfig() config = AppConfig()
monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path)) monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path))
assert _placed(create_rag(config=config)) == temp_db_path
monkeypatch.delenv("HAIKU_RAG_DB")
assert _placed(create_rag(config=config)) == ( assert _placed(create_rag(config=config)) == (
config.storage.data_dir / "haiku.rag.lancedb" config.storage.data_dir / "haiku.rag.lancedb"
) )
@ -124,30 +120,28 @@ class TestACapabilityFollowsTheConfiguredLocation:
"""A capability nobody handed a client opens one for itself, at the """A capability nobody handed a client opens one for itself, at the
database the configuration places.""" database the configuration places."""
def _config(self, tmp_path, uri: str) -> AppConfig: def _config(self, tmp_path, location: str) -> AppConfig:
from haiku.rag.config.models import LanceDBConfig, StorageConfig from haiku.rag.config.models import LanceDBConfig, StorageConfig
return AppConfig( return AppConfig(
lancedb=LanceDBConfig(uri=uri), lancedb=LanceDBConfig(databases={"notes": location}),
storage=StorageConfig(data_dir=tmp_path / "elsewhere"), storage=StorageConfig(data_dir=tmp_path / "elsewhere"),
) )
def test_a_configured_uri_is_left_to_the_client(self, tmp_path): def test_a_configured_location_is_the_capability_scope(self, tmp_path):
"""A path overrides a configured location, so the capability passes
None and the client resolves the configured URI."""
located = tmp_path / "notes.lancedb" located = tmp_path / "notes.lancedb"
for factory in (create_rag, create_analysis): for factory in (create_rag, create_analysis):
[local] = factory( [local] = factory(
config=self._config(tmp_path, str(located)) config=self._config(tmp_path, str(located))
).scope.databases ).scope.databases
assert local == DatabaseRef.configured(None, str(located)) assert local == DatabaseRef("notes", located)
remote = self._config(tmp_path, "s3://bucket/one.lancedb") remote = self._config(tmp_path, "s3://bucket/one.lancedb")
[ref] = factory(config=remote).scope.databases [ref] = factory(config=remote).scope.databases
assert ref == DatabaseRef(None, "s3://bucket/one.lancedb", None) assert ref == DatabaseRef("notes", "s3://bucket/one.lancedb")
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_it_opens_the_database_the_uri_places(self, tmp_path): async def test_it_opens_the_database_the_configuration_places(self, tmp_path):
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
located = tmp_path / "notes.lancedb" located = tmp_path / "notes.lancedb"
@ -162,19 +156,15 @@ class TestACapabilityFollowsTheConfiguredLocation:
finally: finally:
await capability._close() await capability._close()
def test_an_explicit_path_still_overrides_the_configured_uri(self, tmp_path): def test_a_path_beside_the_configured_placement_is_refused(self, tmp_path):
from haiku.rag.store.exceptions import AmbiguousDatabaseError
config = self._config(tmp_path, str(tmp_path / "notes.lancedb")) config = self._config(tmp_path, str(tmp_path / "notes.lancedb"))
chosen = tmp_path / "chosen.lancedb" chosen = tmp_path / "chosen.lancedb"
assert _placed(create_rag(db_path=chosen, config=config)) == chosen for factory in (create_rag, create_analysis):
with pytest.raises(AmbiguousDatabaseError, match="notes"):
def test_the_environment_still_overrides_the_configured_uri( factory(db_path=chosen, config=config)
self, tmp_path, monkeypatch
):
config = self._config(tmp_path, "s3://bucket/one.lancedb")
monkeypatch.setenv("HAIKU_RAG_DB", str(tmp_path / "from-env.lancedb"))
assert _placed(create_rag(config=config)) == tmp_path / "from-env.lancedb"
@pytest.mark.asyncio @pytest.mark.asyncio
@ -206,15 +196,15 @@ def test_domain_preamble_is_added_to_capability_instructions(temp_db_path):
def _single_database_client() -> AsyncMock: def _single_database_client() -> AsyncMock:
"""A stand-in for a client covering one unnamed database. """A stand-in for a client covering one database.
`covers_multiple`, `source` and `clients_covering` answer as one unnamed `covers_multiple`, `source` and `clients_covering` answer as one database
database does; a bare AsyncMock answers every attribute with a truthy Mock. does; a bare AsyncMock answers every attribute with a truthy Mock.
""" """
client = AsyncMock() client = AsyncMock()
client.covers_multiple = False client.covers_multiple = False
client.source_names = () client.source_names = ("test",)
client.source = None client.source = "test"
client.clients_covering.return_value = [client] client.clients_covering.return_value = [client]
return client return client
@ -412,7 +402,7 @@ async def test_a_spent_search_budget_fails_the_tool(temp_db_path):
capability.state = RAGState() capability.state = RAGState()
with pytest.raises(ToolFailed, match="Search limit reached"): with pytest.raises(ToolFailed, match="Search limit reached"):
await capability._search("anything", None) await capability._search("anything", None, 1)
def _stub_client(*batches: list[SearchResult]) -> AsyncMock: def _stub_client(*batches: list[SearchResult]) -> AsyncMock:
@ -452,7 +442,7 @@ async def _labels_of_search(temp_db_path, *sources: str) -> list[str]:
client.source_names = sources client.source_names = sources
with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=client)): with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=client)):
returned = await capability._search("cats", None) returned = await capability._search("cats", None, 1)
assert isinstance(returned, ToolReturn) assert isinstance(returned, ToolReturn)
assert returned.content is not None assert returned.content is not None
@ -483,7 +473,9 @@ async def test_a_fruitless_search_says_so(temp_db_path):
capability.state = RAGState() capability.state = RAGState()
capability.borrowed_rag = _stub_client([]) capability.borrowed_rag = _stub_client([])
assert await capability._search("nothing about this", None) == "No results found." assert (
await capability._search("nothing about this", None, 1) == "No results found."
)
@pytest.mark.asyncio @pytest.mark.asyncio
@ -500,8 +492,8 @@ async def test_a_narrower_repeat_keeps_what_the_wider_search_returned(temp_db_pa
[SearchResult(content="first", score=1.0, chunk_id="chunk-1")], [SearchResult(content="first", score=1.0, chunk_id="chunk-1")],
) )
await capability._search("Figure 3-1", 20) await capability._search("Figure 3-1", 20, 1)
await capability._search("Figure 3-1", None) await capability._search("Figure 3-1", None, 2)
stored = capability.state.searches["Figure 3-1"] stored = capability.state.searches["Figure 3-1"]
assert [result.chunk_id for result in stored] == [ assert [result.chunk_id for result in stored] == [
@ -525,8 +517,8 @@ async def test_two_databases_holding_one_chunk_id_both_survive(temp_db_path):
], ],
) )
await capability._search("cats", 20) await capability._search("cats", 20, 1)
await capability._search("cats", None) await capability._search("cats", None, 2)
stored = capability.state.searches["cats"] stored = capability.state.searches["cats"]
assert [(r.source, r.chunk_id) for r in stored] == [ assert [(r.source, r.chunk_id) for r in stored] == [
@ -1105,7 +1097,7 @@ def _record(deps: Deps, namespace: str) -> CapabilityEvidenceRecord:
return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"]) return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"])
async def _stub_search(self, query: str, _limit: int | None) -> str: async def _stub_search(self, query: str, _limit: int | None, _run_step: int) -> str:
"""Record a result the way the real search does, so citing resolves.""" """Record a result the way the real search does, so citing resolves."""
cast(Any, self.state).searches[query] = [ cast(Any, self.state).searches[query] = [
SearchResult(content="evidence", score=1.0, chunk_id="chunk-1") SearchResult(content="evidence", score=1.0, chunk_id="chunk-1")
@ -1689,22 +1681,18 @@ class TestMultipleCollectionsInstructions:
(create_rag, rag_text), (create_rag, rag_text),
(create_analysis, analysis_text), (create_analysis, analysis_text),
): ):
for config in (AppConfig(), self._config(alpha="/a.lancedb")): one_at_a_path = factory(db_path=Path("/tmp/x.lancedb"), config=AppConfig())
capability = factory(db_path=Path("/tmp/x.lancedb"), config=config) assert one_at_a_path.instruction_text == baseline()
assert capability.instruction_text == baseline() one_configured = factory(config=self._config(alpha="/a.lancedb"))
assert one_configured.instruction_text == baseline()
def test_an_explicit_path_opens_one_database(self): def test_a_path_beside_a_configured_set_is_refused(self):
"""A path names one database, whatever the configuration names.""" from haiku.rag.store.exceptions import AmbiguousDatabaseError
from haiku.rag.capabilities.analysis import instructions as analysis_text
from haiku.rag.capabilities.rag import instructions as rag_text
config = self._config(alpha="/a.lancedb", beta="/b.lancedb") config = self._config(alpha="/a.lancedb", beta="/b.lancedb")
for factory, baseline in ( for factory in (create_rag, create_analysis):
(create_rag, rag_text), with pytest.raises(AmbiguousDatabaseError, match="alpha, beta"):
(create_analysis, analysis_text), factory(db_path=Path("/tmp/one.lancedb"), config=config)
):
capability = factory(db_path=Path("/tmp/one.lancedb"), config=config)
assert capability.instruction_text == baseline()
def test_a_lent_client_covering_one_database_is_instructed_as_before(self): def test_a_lent_client_covering_one_database_is_instructed_as_before(self):
from haiku.rag.capabilities.analysis import instructions as analysis_text from haiku.rag.capabilities.analysis import instructions as analysis_text

View file

@ -29,7 +29,7 @@ class Deps:
state: dict[str, Any] = field(default_factory=dict) state: dict[str, Any] = field(default_factory=dict)
async def stub_search(self, query: str, _limit: int | None) -> str: async def stub_search(self, query: str, _limit: int | None, _run_step: int) -> str:
cast(Any, self.state).searches[query] = [ cast(Any, self.state).searches[query] = [
SearchResult(content="evidence", score=1.0, chunk_id="chunk-1") SearchResult(content="evidence", score=1.0, chunk_id="chunk-1")
] ]

View file

@ -437,7 +437,9 @@ REAL_PNG = base64.b64decode(
) )
async def _search_with_a_picture(self, query: str, _limit: int | None) -> str: async def _search_with_a_picture(
self, query: str, _limit: int | None, _run_step: int
) -> str:
"""Record a result carrying a page image, the way a real search does.""" """Record a result carrying a page image, the way a real search does."""
cast(Any, self.state).searches[query] = [ cast(Any, self.state).searches[query] = [
SearchResult( SearchResult(
@ -515,6 +517,89 @@ async def test_a_picture_that_will_not_decode_emits_neither_image_nor_label(
assert texts_of(wire[-1]) == [] assert texts_of(wire[-1]) == []
def _burst_result() -> SearchResult:
return SearchResult(
content="evidence",
score=1.0,
chunk_id="chunk-1",
document_id="doc-1",
source="main",
doc_item_refs=["#/pictures/0"],
image_data={"#/pictures/0": base64.b64encode(REAL_PNG).decode()},
)
async def _fanout_question_then_another(temp_db_path, cite: bool) -> list[list[Any]]:
"""Question 1 fans out over one picture chunk; question 2 follows compacted."""
rag = create_rag(
db_path=temp_db_path, config=AppConfig(), defer_loading=False, vision=True
)
client = AsyncMock()
client.search.side_effect = [[_burst_result()], [_burst_result()]]
client.expand_context.side_effect = lambda results: results
client.source_names = ["main"]
rag.borrowed_rag = client
citing = (
[[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-3")]]
if cite
else []
)
calls = iter(
[
[
ToolCallPart("rag_search", {"query": "figure"}, "call-1"),
ToolCallPart("rag_search", {"query": "the figure"}, "call-2"),
],
*citing,
[TextPart("first answer")],
[TextPart("second answer")],
]
)
wire: list[list[Any]] = []
async def model(messages, _info):
wire.append(list(messages))
return ModelResponse(parts=next(calls))
agent = Agent(
FunctionModel(model),
deps_type=Deps,
capabilities=[rag, create_compaction()],
)
deps = Deps()
with patch.object(
RAGCapability, "get_picture_bytes", AsyncMock(return_value=REAL_PNG)
):
first = await agent.run("what does the figure show?", deps=deps)
await agent.run(
"and what else?", deps=deps, message_history=first.all_messages()
)
return wire
@pytest.mark.asyncio
async def test_a_burst_deduplicated_picture_survives_compaction_when_cited(
temp_db_path,
):
"""Dedup attaches the picture once in its own question; the capsule re-fetches
it for the next. Neither pass may leave the model without it."""
wire = await _fanout_question_then_another(temp_db_path, cite=True)
assert [picture.data for picture in images_of(wire[1])] == [REAL_PNG]
assert [picture.data for picture in images_of(wire[-1])] == [REAL_PNG]
@pytest.mark.asyncio
async def test_a_burst_deduplicated_picture_is_dropped_by_compaction_uncited(
temp_db_path,
):
wire = await _fanout_question_then_another(temp_db_path, cite=False)
assert [picture.data for picture in images_of(wire[1])] == [REAL_PNG]
assert images_of(wire[-1]) == []
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_the_capsule_is_built_once_per_request_and_again_for_the_next( async def test_the_capsule_is_built_once_per_request_and_again_for_the_next(
temp_db_path, temp_db_path,

View file

@ -0,0 +1,329 @@
import base64
from dataclasses import dataclass, field
from io import BytesIO
from typing import Any
from unittest.mock import AsyncMock
import pytest
from PIL import Image as PILImage
from pydantic_ai import Agent
from pydantic_ai.messages import (
BinaryContent,
ModelResponse,
TextPart,
ToolCallPart,
ToolReturnPart,
)
from pydantic_ai.models.function import FunctionModel
from pydantic_ai.run import AgentRunResult
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
from haiku.rag.capabilities.rag import RAGState
from haiku.rag.capabilities.rag import create_capability as create_rag
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
@dataclass
class Deps:
state: dict[str, Any] = field(default_factory=dict)
def burst_model(bursts: list[list[str]]) -> FunctionModel:
"""Emit one `rag_search` call per query in each burst, then answer."""
responses = 0
def model_function(_messages, _info) -> ModelResponse:
nonlocal responses
responses += 1
if responses <= len(bursts):
return ModelResponse(
parts=[
ToolCallPart("rag_search", {"query": query})
for query in bursts[responses - 1]
]
)
return ModelResponse(parts=[TextPart("done")])
return FunctionModel(model_function)
def burst_agent(
bursts: list[list[str]], db_path, max_searches: int
) -> Agent[Deps, str]:
config = AppConfig()
config.qa.max_searches = max_searches
return Agent(
burst_model(bursts),
deps_type=Deps,
capabilities=[create_rag(db_path=db_path, config=config, defer_loading=False)],
)
def search_returns(result: AgentRunResult[Any]) -> list[ToolReturnPart]:
return [
part
for message in result.all_messages()
for part in message.parts
if isinstance(part, ToolReturnPart) and part.tool_name == "rag_search"
]
def outcomes(result: AgentRunResult[Any]) -> list[str]:
return [
"failed" if part.outcome == "failed" else "ok"
for part in search_returns(result)
]
@pytest.mark.asyncio
async def test_a_burst_in_one_response_consumes_one_unit(rag_db):
"""Three searches emitted together cost one unit and run in emission order."""
agent = burst_agent([["ai", "machine learning", "deep learning"]], rag_db, 1)
result = await agent.run("question", deps=Deps())
assert outcomes(result) == ["ok", "ok", "ok"]
calls = [
part
for message in result.all_messages()
for part in message.parts
if isinstance(part, ToolCallPart) and part.tool_name == "rag_search"
]
assert [part.tool_call_id for part in search_returns(result)] == [
part.tool_call_id for part in calls
]
@pytest.mark.asyncio
async def test_sequential_searches_pay_one_unit_each(rag_db):
agent = burst_agent([["ai"], ["machine learning"]], rag_db, 1)
result = await agent.run("question", deps=Deps())
assert outcomes(result) == ["ok", "failed"]
assert "Search limit reached" in str(search_returns(result)[1].content)
@pytest.mark.asyncio
async def test_max_searches_zero_fails_every_sibling(rag_db):
agent = burst_agent([["ai", "machine learning", "deep learning"]], rag_db, 0)
result = await agent.run("question", deps=Deps())
assert outcomes(result) == ["failed", "failed", "failed"]
@pytest.mark.asyncio
async def test_a_rejected_round_fails_all_its_siblings(rag_db):
agent = burst_agent([["ai"], ["ml", "deep learning", "supervised"]], rag_db, 1)
result = await agent.run("question", deps=Deps())
assert outcomes(result) == ["ok", "failed", "failed", "failed"]
@pytest.mark.asyncio
async def test_a_sibling_past_the_allowance_pays_its_own_unit(rag_db):
burst = [["ai", "machine learning", "deep learning", "supervised learning"]]
within = await burst_agent(burst, rag_db, 2).run("question", deps=Deps())
over = await burst_agent(burst, rag_db, 1).run("question", deps=Deps())
assert outcomes(within) == ["ok", "ok", "ok", "ok"]
assert outcomes(over) == ["ok", "ok", "ok", "failed"]
@pytest.mark.asyncio
async def test_unit_tracking_resets_between_runs(rag_db):
"""A second run's opening burst prices like a first run's."""
def model_function(messages, _info) -> ModelResponse:
if any(isinstance(part, ToolReturnPart) for part in messages[-1].parts):
return ModelResponse(parts=[TextPart("done")])
return ModelResponse(
parts=[
ToolCallPart("rag_search", {"query": query})
for query in ["ai", "machine learning", "deep learning"]
]
)
config = AppConfig()
config.qa.max_searches = 1
agent = Agent(
FunctionModel(model_function),
deps_type=Deps,
capabilities=[create_rag(db_path=rag_db, config=config, defer_loading=False)],
)
deps = Deps()
first = await agent.run("question", deps=deps)
second = await agent.run("another", deps=deps, message_history=first.all_messages())
assert outcomes(first) == ["ok", "ok", "ok"]
assert outcomes(second)[-3:] == ["ok", "ok", "ok"]
def _png() -> str:
buffer = BytesIO()
PILImage.new("RGB", (4, 4), "red").save(buffer, format="PNG")
return base64.b64encode(buffer.getvalue()).decode()
def make_result(**overrides: Any) -> SearchResult:
fields: dict[str, Any] = {
"content": "body",
"score": 0.9,
"source": "main",
"chunk_id": "c1",
"document_id": "d1",
"image_data": {"#/pictures/0": _png()},
}
fields.update(overrides)
return SearchResult(**fields)
def stub_client(
*batches: list[SearchResult], sources: list[str] | None = None
) -> AsyncMock:
client = AsyncMock()
client.search.side_effect = list(batches)
client.expand_context.side_effect = lambda results: results
client.source_names = sources or ["main"]
return client
def dedup_capability(client: AsyncMock, temp_db_path, *, vision: bool = True):
capability = create_rag(db_path=temp_db_path, config=AppConfig(), vision=vision)
capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
capability.borrowed_rag = client
return capability
def images_of(returned: Any) -> list[BinaryContent]:
if isinstance(returned, str):
return []
return [item for item in returned.content if isinstance(item, BinaryContent)]
def text_of(returned: Any) -> str:
return returned if isinstance(returned, str) else returned.return_value
@pytest.mark.asyncio
async def test_a_duplicate_sibling_is_elided_and_stays_citable(temp_db_path):
duplicate, novel = make_result(), make_result(chunk_id="c2", content="novel")
client = stub_client([make_result()], [duplicate, novel])
capability = dedup_capability(client, temp_db_path)
first = await capability._search("q", None, 1)
second = await capability._search("q rephrased", None, 1)
assert len(images_of(first)) == 1
assert images_of(second) == []
text = text_of(second)
assert "Also matched, shown above: [c1] [rank 1 of 2]" in text
assert "body" not in text
assert "[rank 2 of 2]" in text and "novel" in text
assert [r.chunk_id for r in capability.state.searches["q rephrased"]] == [
"c1",
"c2",
]
assert await capability._cite(["c1"]) == "Registered 1 citation(s)."
@pytest.mark.asyncio
async def test_a_new_run_step_formats_shown_results_in_full(temp_db_path):
client = stub_client([make_result()], [make_result()])
capability = dedup_capability(client, temp_db_path)
await capability._search("q", None, 1)
second = await capability._search("q again", None, 2)
assert "body" in text_of(second)
assert len(images_of(second)) == 1
@pytest.mark.asyncio
async def test_same_chunk_id_from_another_collection_is_not_elided(temp_db_path):
client = stub_client(
[make_result(source="alpha")],
[make_result(source="beta")],
sources=["alpha", "beta"],
)
capability = dedup_capability(client, temp_db_path)
await capability._search("q", None, 1)
second = await capability._search("q rephrased", None, 1)
assert "body" in text_of(second)
@pytest.mark.asyncio
async def test_same_anchor_with_new_evidence_formats_in_full(temp_db_path):
shared, extra = _png(), _png()
client = stub_client(
[make_result(content="c1 with c2", image_data={"#/pictures/1": shared})],
[
make_result(
content="c1 with c3",
image_data={"#/pictures/1": shared, "#/pictures/3": extra},
)
],
)
capability = dedup_capability(client, temp_db_path)
await capability._search("q", None, 1)
second = await capability._search("q rephrased", None, 1)
assert "c1 with c3" in text_of(second)
assert len(images_of(second)) == 1
labels = [item for item in second.content if isinstance(item, str)]
assert any("#/pictures/3" in label for label in labels)
@pytest.mark.parametrize(
("overrides", "elided"),
[
({"score": 0.1}, True),
({"content": "different"}, False),
({"document_title": "Other"}, False),
({"headings": ["Heading"]}, False),
({"labels": ["table"]}, False),
({"picture_captions": {"#/pictures/0": "A caption"}}, False),
({"image_data": {"#/pictures/9": _png()}}, False),
],
)
@pytest.mark.asyncio
async def test_equivalence_follows_the_rendered_evidence(
temp_db_path, overrides: dict[str, Any], elided: bool
):
"""Any rendered field or picture identity defeats elision; score alone does not."""
client = stub_client([make_result()], [make_result(**overrides)])
capability = dedup_capability(client, temp_db_path)
await capability._search("q", None, 1)
second = await capability._search("q rephrased", None, 1)
assert ("Also matched, shown above" in text_of(second)) is elided
@pytest.mark.asyncio
async def test_a_failed_sibling_commits_nothing(temp_db_path):
client = stub_client(
[make_result(image_data={"#/pictures/0": "AAA"})],
[make_result()],
)
capability = dedup_capability(client, temp_db_path)
evidence_before = capability.state.evidence.model_dump()
with pytest.raises(Exception):
await capability._search("q", None, 1)
assert capability.state.searches == {}
assert capability.state.evidence.model_dump() == evidence_before
second = await capability._search("q rephrased", None, 1)
assert "body" in text_of(second)
assert len(images_of(second)) == 1

View file

@ -84,17 +84,18 @@ def test_chat_capabilities_read_the_named_database(tmp_path, monkeypatch):
from haiku.rag.chat import run_chat from haiku.rag.chat import run_chat
run_chat(scope=DatabaseScope.resolve(config, database_name="b")) run_chat(scope=DatabaseScope.resolve(config, database_name="b"))
named_scope = chat_app.call_args.kwargs["scope"]
[named] = chat_app.call_args.kwargs["capabilities"] [named] = chat_app.call_args.kwargs["capabilities"]
run_chat(scope=DatabaseScope.resolve(config)) run_chat(scope=DatabaseScope.resolve(config))
covering_scope = chat_app.call_args.kwargs["scope"]
[covering] = chat_app.call_args.kwargs["capabilities"] [covering] = chat_app.call_args.kwargs["capabilities"]
# The chat lends its own client, so this scope is the fallback: it places # The app opens the scope it is handed and lends that client to the
# the named database alone. # capabilities, which keep the configuration as the caller named it.
[placed] = named.scope.databases assert named_scope.names == ("b",)
assert placed.db_path == tmp_path / "b.lancedb" assert covering_scope.names == ("a", "b")
assert named.config.lancedb.databases == {} assert set(named.config.lancedb.databases) == {"a", "b"}
assert covering.scope.names == ("a", "b")
assert set(covering.config.lancedb.databases) == {"a", "b"} assert set(covering.config.lancedb.databases) == {"a", "b"}
@ -158,8 +159,8 @@ def _make_mock_client():
# Covers one database; a bare AsyncMock answers `covers_multiple` with a # Covers one database; a bare AsyncMock answers `covers_multiple` with a
# truthy Mock. # truthy Mock.
mock_client.covers_multiple = False mock_client.covers_multiple = False
mock_client.source_names = () mock_client.source_names = ("test",)
mock_client.source = None mock_client.source = "test"
return mock_client return mock_client
@ -461,9 +462,9 @@ async def test_document_filter_updates_rag_state(temp_db_path: Path):
): ):
async with app.run_test(): async with app.run_test():
# The selection is document ids, so a repeated title cannot widen it. # The selection is document ids, so a repeated title cannot widen it.
selected = [ selected: list[tuple[str | None, str]] = [
(None, "6f1c2d4e-0000-4000-8000-000000000001"), ("test", "6f1c2d4e-0000-4000-8000-000000000001"),
(None, "6f1c2d4e-0000-4000-8000-000000000002"), ("test", "6f1c2d4e-0000-4000-8000-000000000002"),
] ]
app.on_document_filter_modal_filter_changed( app.on_document_filter_modal_filter_changed(
DocumentFilterModal.FilterChanged(selected) DocumentFilterModal.FilterChanged(selected)
@ -477,7 +478,7 @@ async def test_document_filter_updates_rag_state(temp_db_path: Path):
assert rag_state.document_filter == expected_filter assert rag_state.document_filter == expected_filter
assert rag_state.document_filter is not None assert rag_state.document_filter is not None
assert "LIKE" not in rag_state.document_filter assert "LIKE" not in rag_state.document_filter
# An unnamed database leaves the question unscoped by source. # One database leaves the question unscoped by source.
assert rag_state.sources is None assert rag_state.sources is None
# The state snapshot should also reflect the change # The state snapshot should also reflect the change
@ -486,12 +487,15 @@ async def test_document_filter_updates_rag_state(temp_db_path: Path):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_document_filter_narrows_sources_to_the_selection(temp_db_path: Path): async def test_document_filter_narrows_sources_to_the_selection(temp_db_path: Path):
"""The filter carries ids, and `sources` restricts the question to the """Over a set, the filter carries ids and `sources` restricts the question
databases the selection names.""" to the databases the selection names."""
from haiku.rag.chat.app import RAG_STATE_NAMESPACE from haiku.rag.chat.app import RAG_STATE_NAMESPACE
from haiku.rag.chat.widgets.document_filter_modal import DocumentFilterModal from haiku.rag.chat.widgets.document_filter_modal import DocumentFilterModal
app, mock_client = _make_app_with_state(temp_db_path) app, mock_client = _make_app_with_state(temp_db_path)
mock_client.covers_multiple = True
mock_client.source_names = ("alpha", "beta")
mock_client.source = None
with ( with (
patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag, patch("haiku.rag.chat.app.HaikuRAG") as _stub_rag,
@ -537,7 +541,7 @@ async def test_document_filter_cleared_when_empty(temp_db_path: Path):
async with app.run_test(): async with app.run_test():
# First set a filter # First set a filter
app.on_document_filter_modal_filter_changed( app.on_document_filter_modal_filter_changed(
DocumentFilterModal.FilterChanged([(None, "AI Overview")]) DocumentFilterModal.FilterChanged([("test", "AI Overview")])
) )
rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE]) rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE])
assert rag_state.document_filter is not None assert rag_state.document_filter is not None
@ -684,6 +688,37 @@ class TestLendingTheClient:
assert borrowed == [client] * len(app._capabilities) assert borrowed == [client] * len(app._capabilities)
assert borrowed assert borrowed
@pytest.mark.asyncio
async def test_mounting_gives_every_capability_the_apps_scope(self, tmp_path):
"""A capability built over the configured set covers what the chat
selected once mounted: the analysis sandbox is built over that scope."""
from haiku.rag.chat.app import ChatApp
from haiku.rag.client.scope import DatabaseScope
from haiku.rag.config.models import AppConfig, LanceDBConfig
config = AppConfig(
lancedb=LanceDBConfig(
databases={
"a": str(tmp_path / "a.lancedb"),
"b": str(tmp_path / "b.lancedb"),
}
)
)
selected = DatabaseScope.resolve(config, database_name="b")
capability = create_capability(config=config)
assert capability.scope.covers_multiple
client = _make_mock_client()
app = ChatApp(scope=selected, capabilities=[capability], read_only=True)
with (
patch("haiku.rag.chat.app.HaikuRAG") as stub_rag,
_covering_returns(stub_rag, client),
):
async with app.run_test():
pass
assert capability.scope == selected
class TestDocumentSelectionIdentity: class TestDocumentSelectionIdentity:
"""Two documents can share a title, within a corpus and across databases, so """Two documents can share a title, within a corpus and across databases, so
@ -785,11 +820,23 @@ class TestDocumentSelectionIdentity:
) )
] ]
((label, source, doc_id),) = _labelled(docs) ((label, source, doc_id),) = _labelled(docs, name_database=True)
box = DocumentCheckbox(label, source, doc_id, value=False) box = DocumentCheckbox(label, source, doc_id, value=False)
assert str(box.label) == "Report [/red] (alpha [/x])" assert str(box.label) == "Report [/red] (alpha [/x])"
def test_one_database_is_not_named_on_its_labels(self):
"""A single database names every document alike, so the label says
nothing a title does not."""
from haiku.rag.chat.widgets.document_filter_modal import _labelled
from haiku.rag.store.models.document import Document
docs = [Document(id="id-one", content="", title="Report", source="test")]
((label, _, _),) = _labelled(docs)
assert label == "Report"
def test_a_citation_title_that_looks_like_markup_is_text(): def test_a_citation_title_that_looks_like_markup_is_text():
from rich.text import Text from rich.text import Text
@ -1014,17 +1061,20 @@ class TestKeepingSelectionsReachable:
from haiku.rag.store.models.document import Document from haiku.rag.store.models.document import Document
picked = [ picked = [
Document(id=f"sel-{i:04d}", content="", title=f"Selected {i:04d}") Document(
id=f"sel-{i:04d}", content="", title=f"Selected {i:04d}", source="test"
)
for i in range(DOCUMENT_PAGE + 20) for i in range(DOCUMENT_PAGE + 20)
] ]
by_id = {d.id: d for d in picked} by_id = {d.id: d for d in picked}
matched = [ matched = [
Document(id=f"hit-{i}", content="", title=f"Hit {i}") for i in range(5) Document(id=f"hit-{i}", content="", title=f"Hit {i}", source="test")
for i in range(5)
] ]
client = AsyncMock() client = AsyncMock()
client.covers_multiple = False client.covers_multiple = False
client.source_names = () client.source_names = ("test",)
client.count_documents.return_value = 5 client.count_documents.return_value = 5
async def listing(limit=None, offset=0, filter=None): async def listing(limit=None, offset=0, filter=None):
@ -1036,7 +1086,7 @@ class TestKeepingSelectionsReachable:
client.list_documents.side_effect = listing client.list_documents.side_effect = listing
modal = DocumentFilterModal( modal = DocumentFilterModal(
client=client, selected=[(None, d.id or "") for d in picked] client=client, selected=[("test", d.id or "") for d in picked]
) )
app, _ = _make_app(temp_db_path, client) app, _ = _make_app(temp_db_path, client)
with ( with (
@ -1083,14 +1133,16 @@ class TestKeepingSelectionsReachable:
from haiku.rag.store.models.document import Document from haiku.rag.store.models.document import Document
picked = [ picked = [
Document(id=f"sel-{i:04d}", content="", title=f"Selected {i:04d}") Document(
id=f"sel-{i:04d}", content="", title=f"Selected {i:04d}", source="test"
)
for i in range(DOCUMENT_PAGE + 1) for i in range(DOCUMENT_PAGE + 1)
] ]
by_id = {d.id: d for d in picked} by_id = {d.id: d for d in picked}
client = AsyncMock() client = AsyncMock()
client.covers_multiple = False client.covers_multiple = False
client.source_names = () client.source_names = ("test",)
client.count_documents.return_value = 0 client.count_documents.return_value = 0
async def listing(limit=None, offset=0, filter=None): async def listing(limit=None, offset=0, filter=None):
@ -1102,7 +1154,7 @@ class TestKeepingSelectionsReachable:
client.list_documents.side_effect = listing client.list_documents.side_effect = listing
modal = DocumentFilterModal( modal = DocumentFilterModal(
client=client, selected=[(None, d.id or "") for d in picked] client=client, selected=[("test", d.id or "") for d in picked]
) )
app, _ = _make_app(temp_db_path, client) app, _ = _make_app(temp_db_path, client)
with ( with (
@ -1134,7 +1186,9 @@ class TestKeepingSelectionsReachable:
# The row is gone from the listing, not merely unchecked. # The row is gone from the listing, not merely unchecked.
assert "sel-0200" not in remaining assert "sel-0200" not in remaining
assert len(remaining) == DOCUMENT_PAGE assert len(remaining) == DOCUMENT_PAGE
assert modal._selected == {(None, d.id) for d in picked} - {(None, "sel-0200")} assert modal._selected == {("test", d.id) for d in picked} - {
("test", "sel-0200")
}
# The page it was on no longer exists, so the modal does not report it. # The page it was on no longer exists, so the modal does not report it.
assert modal._page == 0 assert modal._page == 0
assert "page" not in footer assert "page" not in footer
@ -1154,10 +1208,10 @@ class TestKeepingSelectionsReachable:
client = AsyncMock() client = AsyncMock()
client.covers_multiple = False client.covers_multiple = False
client.source_names = () client.source_names = ("test",)
client.count_documents.return_value = DOCUMENT_PAGE * 2 client.count_documents.return_value = DOCUMENT_PAGE * 2
client.list_documents.return_value = [ client.list_documents.return_value = [
Document(id="d1", content="", title="One") Document(id="d1", content="", title="One", source="test")
] ]
modal = DocumentFilterModal(client=client) modal = DocumentFilterModal(client=client)
@ -1187,7 +1241,7 @@ class TestKeepingSelectionsReachable:
client = AsyncMock() client = AsyncMock()
client.covers_multiple = False client.covers_multiple = False
client.source_names = () client.source_names = ("test",)
client.list_documents.return_value = [] client.list_documents.return_value = []
client.count_documents.return_value = 0 client.count_documents.return_value = 0
@ -1222,10 +1276,10 @@ class TestKeepingSelectionsReachable:
client = AsyncMock() client = AsyncMock()
client.covers_multiple = False client.covers_multiple = False
client.source_names = () client.source_names = ("test",)
client.list_documents.return_value = [ client.list_documents.return_value = [
Document(id="id-one", content="", title="Capital region"), Document(id="id-one", content="", title="Capital region", source="test"),
Document(id="id-two", content="", title="Nobel laureates"), Document(id="id-two", content="", title="Nobel laureates", source="test"),
] ]
client.count_documents.return_value = DOCUMENT_PAGE * 2 client.count_documents.return_value = DOCUMENT_PAGE * 2
@ -1328,10 +1382,10 @@ class TestDocumentSearchFilter:
client = AsyncMock() client = AsyncMock()
client.covers_multiple = False client.covers_multiple = False
client.source_names = () client.source_names = ("test",)
client.list_documents.return_value = [ client.list_documents.return_value = [
Document(id="id-one", content="", title="Capital region"), Document(id="id-one", content="", title="Capital region", source="test"),
Document(id="id-two", content="", title="Nobel laureates"), Document(id="id-two", content="", title="Nobel laureates", source="test"),
] ]
client.count_documents.return_value = 2 client.count_documents.return_value = 2
@ -1347,7 +1401,9 @@ class TestDocumentSearchFilter:
assert len(list(modal.query(DocumentCheckbox))) == 2 assert len(list(modal.query(DocumentCheckbox))) == 2
client.list_documents.return_value = [ client.list_documents.return_value = [
Document(id="id-two", content="", title="Nobel laureates"), Document(
id="id-two", content="", title="Nobel laureates", source="test"
),
] ]
client.count_documents.return_value = 1 client.count_documents.return_value = 1
await modal.on_input_submitted(Input.Submitted(Input(), "Nobel")) await modal.on_input_submitted(Input.Submitted(Input(), "Nobel"))

View file

@ -21,6 +21,12 @@ embeddings:
""") """)
os.environ["HAIKU_RAG_CONFIG_PATH"] = str(_test_config_path) os.environ["HAIKU_RAG_CONFIG_PATH"] = str(_test_config_path)
# telemetry.configure() passes send_to_logfire="if-token-present" explicitly,
# which beats LOGFIRE_SEND_TO_LOGFIRE, so no token must resolve: drop the
# environment variable and point credentials discovery at an empty directory.
os.environ.pop("LOGFIRE_TOKEN", None)
os.environ["LOGFIRE_CREDENTIALS_DIR"] = tempfile.mkdtemp()
import pydantic_ai.models # noqa: E402 import pydantic_ai.models # noqa: E402
import pytest # noqa: E402 import pytest # noqa: E402
import yaml # noqa: E402 import yaml # noqa: E402
@ -102,7 +108,7 @@ def temp_yaml_config(tmp_path, monkeypatch):
"vector_dim": 2560, "vector_dim": 2560,
} }
}, },
"qa": {"model": {"provider": "ollama", "name": "gpt-oss"}}, "qa": {"model": {"provider": "ollama", "name": "qwen3.8"}},
} }
with open(config_file, "w") as f: with open(config_file, "w") as f:

View file

@ -838,11 +838,11 @@ async def test_database_503_when_no_database_is_configured(state):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_the_report_follows_a_configured_uri(tmp_path, jobs, sync): async def test_the_report_follows_a_configured_uri(tmp_path, jobs, sync):
"""A configured `lancedb.uri` places the database, so the report opens that """A configured location places the database, so the report opens that
and not the local default.""" and not the local default."""
db_path = tmp_path / "configured.lancedb" db_path = tmp_path / "configured.lancedb"
await _seed_lancedb(db_path) await _seed_lancedb(db_path)
config = AppConfig(lancedb=LanceDBConfig(uri=str(db_path))) config = AppConfig(lancedb=LanceDBConfig(databases={"configured": str(db_path)}))
state = APIState( state = APIState(
config=config, config=config,
job_repo=jobs, job_repo=jobs,

View file

@ -504,31 +504,58 @@ def test_cli_entry_point_exits_on_store_state_errors(monkeypatch, capsys, error)
class TestPlacingTheIngesterDatabase: class TestPlacingTheIngesterDatabase:
"""The ingester writes wherever the configuration places the database, and """The ingester writes wherever the configuration places the database, and
resolves that once. A path is an explicit override of a configured resolves that once. `--db PATH` is the operator's explicit override and
`lancedb.uri`, so no local default stands in for one.""" constructs the scope directly."""
@staticmethod @staticmethod
def _app(config: AppConfig, db_path=None): def _app(config: AppConfig, db_path=None):
from haiku.rag.ingester.app import IngesterApp from haiku.rag.ingester.app import IngesterApp
from haiku.rag.ingester.cli import _scope_for
return IngesterApp(config=config, db_path=db_path) return IngesterApp(config=config, scope=_scope_for(db_path))
def test_a_configured_uri_becomes_the_scope(self, tmp_path): def test_a_configured_location_becomes_the_scope(self, tmp_path):
config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/prod.lancedb")) config = AppConfig(
lancedb=LanceDBConfig(databases={"prod": "s3://bucket/prod.lancedb"})
)
[ref] = self._app(config)._scope.databases [ref] = self._app(config)._scope.databases
assert ref.uri == "s3://bucket/prod.lancedb" assert ref.name == "prod"
assert ref.location == "s3://bucket/prod.lancedb"
assert ref.db_path is None assert ref.db_path is None
def test_an_override_names_the_database(self, tmp_path): def test_an_override_names_the_database_by_its_stem(self, tmp_path):
config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/prod.lancedb")) config = AppConfig(
lancedb=LanceDBConfig(databases={"prod": "s3://bucket/prod.lancedb"})
)
override = tmp_path / "local.lancedb" override = tmp_path / "local.lancedb"
[ref] = self._app(config, override)._scope.databases [ref] = self._app(config, override)._scope.databases
assert ref.db_path == override assert ref.name == "local"
assert ref.uri == "" assert ref.location == override
def test_the_override_is_the_cli_s_alone(self, tmp_path):
"""`IngesterApp` takes a resolved scope, so a Python caller has no path
to slip past the configuration; only the CLI constructs one."""
from haiku.rag.client.scope import DatabaseScope
from haiku.rag.ingester.cli import _scope_for
assert _scope_for(None) is None
assert _scope_for(tmp_path / "local.lancedb") == DatabaseScope.at(
tmp_path / "local.lancedb"
)
def test_a_path_without_a_stem_is_a_usage_error(self):
from pathlib import Path
import typer
from haiku.rag.ingester.cli import _scope_for
with pytest.raises(typer.BadParameter, match="no name"):
_scope_for(Path("/"))
def test_one_configured_database_is_accepted(self, tmp_path): def test_one_configured_database_is_accepted(self, tmp_path):
"""A one-entry mapping names which database to write.""" """A one-entry mapping names which database to write."""
@ -567,6 +594,10 @@ class TestPlacingTheIngesterDatabase:
f" a: {tmp_path / 'a.lancedb'}\n" f" a: {tmp_path / 'a.lancedb'}\n"
f" b: {tmp_path / 'b.lancedb'}\n" f" b: {tmp_path / 'b.lancedb'}\n"
) )
import haiku.rag.config as config_module
# The CLI caches the loaded configuration process-wide.
monkeypatch.setattr(config_module, "_config", None)
monkeypatch.setattr(sys, "argv", ["haiku-ingester", "run-batch"]) monkeypatch.setattr(sys, "argv", ["haiku-ingester", "run-batch"])
monkeypatch.setenv("HAIKU_RAG_CONFIG_PATH", str(config_file)) monkeypatch.setenv("HAIKU_RAG_CONFIG_PATH", str(config_file))

View file

@ -14,6 +14,7 @@ import pytest
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.client.exceptions import UnsupportedSourceError from haiku.rag.client.exceptions import UnsupportedSourceError
from haiku.rag.client.scope import DatabaseScope
from haiku.rag.config import ( from haiku.rag.config import (
APIConfig, APIConfig,
AppConfig, AppConfig,
@ -121,7 +122,7 @@ async def test_run_batch_drains_upserts(tmp_path, use_client):
use_client(client) use_client(client)
report = await IngesterApp( report = await IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb" config=_config(tmp_path), scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch() ).run_batch()
assert report.succeeded == 2 assert report.succeeded == 2
@ -146,7 +147,7 @@ async def test_run_batch_reports_progress(tmp_path, use_client):
progress = [] progress = []
report = await IngesterApp( report = await IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb" config=_config(tmp_path), scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch(progress_callback=progress.append) ).run_batch(progress_callback=progress.append)
assert report.succeeded == 2 assert report.succeeded == 2
@ -171,7 +172,9 @@ async def test_run_batch_prunes_orphans(tmp_path, use_client):
db_path = tmp_path / "db.lancedb" db_path = tmp_path / "db.lancedb"
# First batch ingests both files and records sync_state for each. # First batch ingests both files and records sync_state for each.
first = await IngesterApp(config=config, db_path=db_path).run_batch() first = await IngesterApp(
config=config, scope=DatabaseScope.at(db_path)
).run_batch()
assert first.succeeded == 2 assert first.succeeded == 2
client.delete_document.assert_not_awaited() client.delete_document.assert_not_awaited()
@ -179,7 +182,9 @@ async def test_run_batch_prunes_orphans(tmp_path, use_client):
# sync_state but not on disk -> enqueues a DELETE for it. # sync_state but not on disk -> enqueues a DELETE for it.
(tmp_path / "b.md").unlink() (tmp_path / "b.md").unlink()
second = await IngesterApp(config=config, db_path=db_path).run_batch() second = await IngesterApp(
config=config, scope=DatabaseScope.at(db_path)
).run_batch()
# a.md is unchanged (same mtime) so it's not re-ingested; only the orphan # a.md is unchanged (same mtime) so it's not re-ingested; only the orphan
# delete runs. # delete runs.
@ -198,7 +203,7 @@ async def test_run_batch_reports_dead_on_permanent_failure(tmp_path, use_client)
use_client(client) use_client(client)
report = await IngesterApp( report = await IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb" config=_config(tmp_path), scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch() ).run_batch()
assert report.succeeded == 0 assert report.succeeded == 0
@ -217,7 +222,9 @@ async def test_run_batch_recovered_doc_is_not_counted_as_dead(tmp_path, use_clie
failing = _mock_client() failing = _mock_client()
failing.create_document_from_source.side_effect = UnsupportedSourceError("nope") failing.create_document_from_source.side_effect = UnsupportedSourceError("nope")
use_client(failing) use_client(failing)
first = await IngesterApp(config=config, db_path=db_path).run_batch() first = await IngesterApp(
config=config, scope=DatabaseScope.at(db_path)
).run_batch()
assert first.dead == 1 assert first.dead == 1
healthy = _mock_client() healthy = _mock_client()
@ -226,7 +233,9 @@ async def test_run_batch_recovered_doc_is_not_counted_as_dead(tmp_path, use_clie
# records the revision in sync_state, so a plain re-run no longer retries an # records the revision in sync_state, so a plain re-run no longer retries an
# unchanged file — recovery needs the content (mtime) to change. # unchanged file — recovery needs the content (mtime) to change.
(tmp_path / "a.md").write_text("hello again") (tmp_path / "a.md").write_text("hello again")
second = await IngesterApp(config=config, db_path=db_path).run_batch() second = await IngesterApp(
config=config, scope=DatabaseScope.at(db_path)
).run_batch()
assert second.dead == 0 assert second.dead == 0
assert second.succeeded == 1 assert second.succeeded == 1
@ -248,7 +257,7 @@ async def test_run_batch_reports_failed_sweep(
with caplog.at_level("ERROR", logger="haiku.rag.ingester.pollers.base"): with caplog.at_level("ERROR", logger="haiku.rag.ingester.pollers.base"):
report = await IngesterApp( report = await IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb" config=_config(tmp_path), scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch() ).run_batch()
assert report.failed_sweeps == ["local"] assert report.failed_sweeps == ["local"]
@ -264,7 +273,7 @@ async def test_run_batch_empty_source_returns_immediately(tmp_path, use_client):
report = await asyncio.wait_for( report = await asyncio.wait_for(
IngesterApp( IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb" config=_config(tmp_path), scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch(), ).run_batch(),
timeout=5.0, timeout=5.0,
) )
@ -287,7 +296,9 @@ async def test_run_batch_dry_run_reports_manifest_without_mutating_queue(tmp_pat
finally: finally:
await engine.dispose() await engine.dispose()
report = await IngesterApp(config=config, db_path=db_path).run_batch_dry_run() report = await IngesterApp(
config=config, scope=DatabaseScope.at(db_path)
).run_batch_dry_run()
assert report.failed_sweeps == [] assert report.failed_sweeps == []
assert report.manifest.version == 1 assert report.manifest.version == 1
@ -322,7 +333,7 @@ async def test_run_batch_from_manifest_drains_changes_without_sweeping(
monkeypatch.setattr(PollerManager, "sweep_all", sweep_all) monkeypatch.setattr(PollerManager, "sweep_all", sweep_all)
report = await IngesterApp( report = await IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb" config=_config(tmp_path), scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch_from_manifest( ).run_batch_from_manifest(
_manifest( _manifest(
BatchChange( BatchChange(
@ -350,7 +361,7 @@ async def test_run_batch_from_manifest_rejects_stale_upsert_revision(
use_client(client) use_client(client)
report = await IngesterApp( report = await IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb" config=_config(tmp_path), scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch_from_manifest( ).run_batch_from_manifest(
_manifest( _manifest(
BatchChange( BatchChange(
@ -378,7 +389,7 @@ async def test_run_batch_from_manifest_delete_uses_manifest_even_if_file_reappea
use_client(client) use_client(client)
report = await IngesterApp( report = await IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb" config=_config(tmp_path), scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch_from_manifest( ).run_batch_from_manifest(
_manifest( _manifest(
BatchChange( BatchChange(
@ -431,7 +442,7 @@ async def test_run_batch_from_manifest_resumes_same_manifest_work(tmp_path, use_
await engine.dispose() await engine.dispose()
report = await IngesterApp( report = await IngesterApp(
config=config, db_path=tmp_path / "db.lancedb" config=config, scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch_from_manifest(manifest) ).run_batch_from_manifest(manifest)
assert report.succeeded == 1 assert report.succeeded == 1
@ -456,7 +467,7 @@ async def test_run_batch_from_manifest_rejects_non_manifest_pending_work(
with pytest.raises(ValueError, match="non-manifest pending work"): with pytest.raises(ValueError, match="non-manifest pending work"):
await IngesterApp( await IngesterApp(
config=config, db_path=tmp_path / "db.lancedb" config=config, scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch_from_manifest( ).run_batch_from_manifest(
_manifest( _manifest(
BatchChange( BatchChange(
@ -505,7 +516,7 @@ async def test_run_batch_from_manifest_rejects_different_manifest_pending_work(
with pytest.raises(ValueError, match="non-manifest pending work"): with pytest.raises(ValueError, match="non-manifest pending work"):
await IngesterApp( await IngesterApp(
config=config, db_path=tmp_path / "db.lancedb" config=config, scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch_from_manifest(manifest) ).run_batch_from_manifest(manifest)
@ -526,7 +537,7 @@ async def test_run_batch_from_manifest_rejects_unrelated_pending_work(
with pytest.raises(ValueError, match="non-manifest pending work"): with pytest.raises(ValueError, match="non-manifest pending work"):
await IngesterApp( await IngesterApp(
config=config, db_path=tmp_path / "db.lancedb" config=config, scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch_from_manifest( ).run_batch_from_manifest(
_manifest( _manifest(
BatchChange( BatchChange(
@ -554,7 +565,7 @@ async def test_run_batch_from_manifest_rejects_duplicate_changes(tmp_path, use_c
with pytest.raises(ValueError, match="duplicate change"): with pytest.raises(ValueError, match="duplicate change"):
await IngesterApp( await IngesterApp(
config=_config(tmp_path), db_path=tmp_path / "db.lancedb" config=_config(tmp_path), scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch_from_manifest(_manifest(change, change)) ).run_batch_from_manifest(_manifest(change, change))
@ -583,7 +594,9 @@ async def test_run_batch_aborts_when_all_workers_die(
with caplog.at_level("ERROR", logger="haiku.rag.ingester.app"): with caplog.at_level("ERROR", logger="haiku.rag.ingester.app"):
report = await asyncio.wait_for( report = await asyncio.wait_for(
IngesterApp(config=config, db_path=tmp_path / "db.lancedb").run_batch(), IngesterApp(
config=config, scope=DatabaseScope.at(tmp_path / "db.lancedb")
).run_batch(),
timeout=10.0, timeout=10.0,
) )
@ -608,7 +621,7 @@ async def test_serve_starts_workers_pollers_and_shuts_down(tmp_path, use_client,
use_client(_mock_client()) use_client(_mock_client())
config = _config(tmp_path) config = _config(tmp_path)
config.ingester.api = APIConfig(enabled=api, host="127.0.0.1", port=0) config.ingester.api = APIConfig(enabled=api, host="127.0.0.1", port=0)
app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb") app = IngesterApp(config=config, scope=DatabaseScope.at(tmp_path / "db.lancedb"))
task = asyncio.create_task(app.serve(api=api)) task = asyncio.create_task(app.serve(api=api))
try: try:
@ -653,7 +666,7 @@ async def test_stop_pool_warns_when_shutdown_grace_elapses(tmp_path, caplog):
"""When a worker doesn't stop within the shutdown grace, _stop_pool logs a """When a worker doesn't stop within the shutdown grace, _stop_pool logs a
warning and still drains any pending cancel-cleanup releases.""" warning and still drains any pending cancel-cleanup releases."""
config = _config(tmp_path, shutdown_grace_s=0.01) config = _config(tmp_path, shutdown_grace_s=0.01)
app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb") app = IngesterApp(config=config, scope=DatabaseScope.at(tmp_path / "db.lancedb"))
pool = _SlowPool() pool = _SlowPool()
app._pool = pool app._pool = pool
@ -694,7 +707,9 @@ async def test_run_batch_closes_sources_after_pool_stops(
): ):
(tmp_path / "a.md").write_text("hello") (tmp_path / "a.md").write_text("hello")
use_client(_mock_client()) use_client(_mock_client())
app = IngesterApp(config=_config(tmp_path), db_path=tmp_path / "db.lancedb") app = IngesterApp(
config=_config(tmp_path), scope=DatabaseScope.at(tmp_path / "db.lancedb")
)
order = _record_close_order(monkeypatch) order = _record_close_order(monkeypatch)
await app.run_batch() await app.run_batch()
@ -707,7 +722,7 @@ async def test_serve_closes_sources_after_pool_stops(tmp_path, use_client, monke
use_client(_mock_client()) use_client(_mock_client())
config = _config(tmp_path) config = _config(tmp_path)
config.ingester.api = APIConfig(enabled=False) config.ingester.api = APIConfig(enabled=False)
app = IngesterApp(config=config, db_path=tmp_path / "db.lancedb") app = IngesterApp(config=config, scope=DatabaseScope.at(tmp_path / "db.lancedb"))
order = _record_close_order(monkeypatch) order = _record_close_order(monkeypatch)
task = asyncio.create_task(app.serve(api=False)) task = asyncio.create_task(app.serve(api=False))

View file

@ -1054,16 +1054,23 @@ async def test_breaker_isolates_sources(client, jobs, sync):
for _ in range(10): for _ in range(10):
pool._breaker_for("bad").record_failure() pool._breaker_for("bad").record_failure()
async def _good_jobs_drained():
while True:
done = await jobs.list_jobs(status=JobStatus.SUCCEEDED, limit=50)
if len(done) == 3:
return done
await asyncio.sleep(0.02)
await pool.start() await pool.start()
try: try:
await asyncio.sleep(0.2) succeeded = await asyncio.wait_for(_good_jobs_drained(), timeout=5.0)
succeeded = await jobs.list_jobs(status=JobStatus.SUCCEEDED, limit=50)
queued = await jobs.list_jobs(status=JobStatus.QUEUED, limit=50) queued = await jobs.list_jobs(status=JobStatus.QUEUED, limit=50)
finally: finally:
await pool.stop() await pool.stop()
assert {j.uri for j in succeeded} == {"g0", "g1", "g2"} assert {j.uri for j in succeeded} == {"g0", "g1", "g2"}
assert {j.uri for j in queued} == {"b0", "b1", "b2"} assert {j.uri for j in queued} == {"b0", "b1", "b2"}
assert [j.attempts for j in queued] == [0, 0, 0]
@pytest.mark.asyncio @pytest.mark.asyncio

View file

@ -40,8 +40,8 @@ async def _restore_embedder(config, name, *, provider=None, model_name=None):
import lancedb import lancedb
_, db_path = locate_database(config.lancedb.databases[name]) db_path = locate_database(config.lancedb.databases[name])
assert db_path is not None assert not isinstance(db_path, str)
db = await lancedb.connect_async(str(db_path.resolve())) db = await lancedb.connect_async(str(db_path.resolve()))
table = await db.open_table("settings") table = await db.open_table("settings")
rows = ( rows = (

View file

@ -7,7 +7,6 @@ import pytest
from haiku.rag.capabilities._tools import search_corpus from haiku.rag.capabilities._tools import search_corpus
from haiku.rag.capabilities.rag import RAGState, create_capability from haiku.rag.capabilities.rag import RAGState, create_capability
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.client.scope import DatabaseScope
from haiku.rag.client.session import FederatedSession from haiku.rag.client.session import FederatedSession
from haiku.rag.sandbox import AnalysisContext, Sandbox from haiku.rag.sandbox import AnalysisContext, Sandbox
from haiku.rag.store.exceptions import UnknownDatabaseError from haiku.rag.store.exceptions import UnknownDatabaseError
@ -30,7 +29,7 @@ class TestAskAcrossDatabases:
capability = create_capability(config=config, rag=rag, defer_loading=False) capability = create_capability(config=config, rag=rag, defer_loading=False)
capability.state = RAGState(sources=["alpha"]) capability.state = RAGState(sources=["alpha"])
formatted = await capability._search("cats", limit=10) formatted = await capability._search("cats", 10, 1)
assert isinstance(formatted, str) assert isinstance(formatted, str)
assert "alpha" in formatted assert "alpha" in formatted
@ -47,7 +46,7 @@ class TestAskAcrossDatabases:
capability = create_capability(config=config, rag=rag, defer_loading=False) capability = create_capability(config=config, rag=rag, defer_loading=False)
capability.state = RAGState() capability.state = RAGState()
formatted = await capability._search("cats", limit=10) formatted = await capability._search("cats", 10, 1)
assert isinstance(formatted, str) assert isinstance(formatted, str)
assert "alpha document" in formatted assert "alpha document" in formatted
@ -72,7 +71,7 @@ class TestStandaloneCapabilities:
assert capability.scope.names == ("alpha", "beta") assert capability.scope.names == ("alpha", "beta")
run = await capability.for_run(make_context(Deps())) run = await capability.for_run(make_context(Deps()))
try: try:
formatted = await run._search("cats", limit=10) formatted = await run._search("cats", 10, 1)
finally: finally:
await run._close() await run._close()
@ -135,7 +134,7 @@ class TestAnalyzeAcrossDatabases:
capability = create_analysis(config=config, rag=rag, defer_loading=False) capability = create_analysis(config=config, rag=rag, defer_loading=False)
capability.state = AnalysisState(sources=["alpha"]) capability.state = AnalysisState(sources=["alpha"])
formatted = await capability._search("cats", limit=10) formatted = await capability._search("cats", 10, 1)
sandbox = await capability._ensure_sandbox() sandbox = await capability._ensure_sandbox()
await capability._close() await capability._close()
@ -163,8 +162,9 @@ class TestCollectionIdentityForTheModel:
assert "Collection" not in result.format_for_agent() assert "Collection" not in result.format_for_agent()
def test_an_unnamed_collection_is_never_mentioned(self): def test_a_hand_built_result_without_a_source_is_never_labelled(self):
"""Nothing to name, whatever the caller asked for.""" """A result built by hand carries no source to name, whatever the
caller asked for."""
result = SearchResult(content="body", score=0.9, chunk_id="c1") result = SearchResult(content="body", score=0.9, chunk_id="c1")
assert "Collection" not in result.format_for_agent(include_collection=True) assert "Collection" not in result.format_for_agent(include_collection=True)
@ -256,12 +256,9 @@ class TestLendingANamedClient:
await _seed(config, "alpha", ["alpha document about cats"]) await _seed(config, "alpha", ["alpha document about cats"])
await _seed(config, "beta", ["beta document about cats"]) await _seed(config, "beta", ["beta document about cats"])
# `run_chat` derives these for a single-database scope. # What `run_chat` builds: the capability's own scope is the set, and
scope = DatabaseScope.resolve(config, database_name="alpha") # the lent client is what narrows it.
one_config, one_path = scope.databases[0].connection(config) capability = create_capability(config=config, defer_loading=False)
capability = create_capability(
db_path=one_path, config=one_config, defer_loading=False
)
async with HaikuRAG(config=config, sources=["alpha"]) as client: async with HaikuRAG(config=config, sources=["alpha"]) as client:
# What `ChatApp.on_mount` does. # What `ChatApp.on_mount` does.
@ -305,8 +302,10 @@ class TestWhenTheModelIsToldTheCollection:
async with HaikuRAG(config=config) as rag: async with HaikuRAG(config=config) as rag:
monkeypatch.setattr(rag, "search", AsyncMock(return_value=only_alpha)) monkeypatch.setattr(rag, "search", AsyncMock(return_value=only_alpha))
spanning, _, spans = await search_corpus(rag, "cats") spanning, _, _, spans = await search_corpus(rag, "cats")
narrowed, _, narrows = await search_corpus(rag, "cats", sources=["alpha"]) narrowed, _, _, narrows = await search_corpus(
rag, "cats", sources=["alpha"]
)
assert "Collection: alpha" in spanning assert "Collection: alpha" in spanning
assert "Collection" not in narrowed assert "Collection" not in narrowed

View file

@ -299,7 +299,7 @@ class TestCitationSource:
assert citation.chunk_id == "c1" assert citation.chunk_id == "c1"
def test_a_single_database_citation_has_no_source(self): def test_a_citation_from_a_hand_built_result_has_no_source(self):
result = SearchResult( result = SearchResult(
content="body", content="body",
score=0.9, score=0.9,
@ -342,7 +342,7 @@ class TestCiteFallback:
) )
run = await capability.for_run(make_context(deps)) run = await capability.for_run(make_context(deps))
# The search returns the cats chunk, never the aardvark one. # The search returns the cats chunk, never the aardvark one.
await run._search("cats", limit=10) await run._search("cats", 10, 1)
await run._cite([aardvark.id]) await run._cite([aardvark.id])
@ -374,7 +374,7 @@ class TestCiteFallback:
state={"rag": RAGState(sources=["alpha"]).model_dump(mode="json")} state={"rag": RAGState(sources=["alpha"]).model_dump(mode="json")}
) )
run = await capability.for_run(make_context(deps)) run = await capability.for_run(make_context(deps))
await run._search("cats", limit=10) await run._search("cats", 10, 1)
with pytest.raises(ModelRetry, match="None of the supplied chunk_ids"): with pytest.raises(ModelRetry, match="None of the supplied chunk_ids"):
await run._cite([outside.id]) await run._cite([outside.id])

View file

@ -259,7 +259,8 @@ class TestLookupByIdentifier:
assert chunk is not None and chunk.content == "alpha one" assert chunk is not None and chunk.content == "alpha one"
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_an_unnamed_database_answers_to_no_name(self, temp_db_path): async def test_a_database_at_a_path_answers_to_its_stem_alone(self, temp_db_path):
stem = temp_db_path.stem
async with HaikuRAG(temp_db_path, create=True) as rag: async with HaikuRAG(temp_db_path, create=True) as rag:
docling = DoclingDocument(name="one") docling = DoclingDocument(name="one")
docling.add_text(label=DocItemLabel.TEXT, text="body") docling.add_text(label=DocItemLabel.TEXT, text="body")
@ -276,6 +277,8 @@ class TestLookupByIdentifier:
assert await rag.get_document_by_id(doc.id) is not None assert await rag.get_document_by_id(doc.id) is not None
assert await rag.get_chunk_by_id(held.id) is not None assert await rag.get_chunk_by_id(held.id) is not None
assert await rag.get_document_by_id(doc.id, stem) is not None
assert await rag.get_chunk_by_id(held.id, stem) is not None
with pytest.raises(UnknownDatabaseError): with pytest.raises(UnknownDatabaseError):
await rag.get_document_by_id(doc.id, "alpha") await rag.get_document_by_id(doc.id, "alpha")
with pytest.raises(UnknownDatabaseError): with pytest.raises(UnknownDatabaseError):
@ -459,8 +462,9 @@ class TestDocumentsNameTheirDatabase:
assert by_uri is not None and by_uri.source == "alpha" assert by_uri is not None and by_uri.source == "alpha"
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_one_database_leaves_the_source_unset(self, tmp_path, temp_db_path): async def test_one_database_at_a_path_is_named_by_its_stem(
"""Nothing names the database when there is only one to name.""" self, tmp_path, temp_db_path
):
async with HaikuRAG(temp_db_path, create=True) as rag: async with HaikuRAG(temp_db_path, create=True) as rag:
dim = get_config().embeddings.model.vector_dim dim = get_config().embeddings.model.vector_dim
doc = DoclingDocument(name="solo") doc = DoclingDocument(name="solo")
@ -472,6 +476,7 @@ class TestDocumentsNameTheirDatabase:
) )
[listed] = await rag.list_documents() [listed] = await rag.list_documents()
assert listed.source is None assert listed.source == temp_db_path.stem
assert listed.id is not None assert listed.id is not None
assert (await rag.get_document_by_id(listed.id)).source is None by_id = await rag.get_document_by_id(listed.id)
assert by_id is not None and by_id.source == temp_db_path.stem

View file

@ -475,8 +475,27 @@ class TestFailureNaming:
assert caught.value.__cause__ is None assert caught.value.__cause__ is None
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_an_unnamed_database_keeps_its_error(self, tmp_path): async def test_a_missing_default_database_names_the_remedy(self, tmp_path):
"""Nothing named it, so there is no name to report.""" """The location stays out of the message; the way to create the
database does not."""
from haiku.rag.config.models import AppConfig, StorageConfig
config = AppConfig(storage=StorageConfig(data_dir=tmp_path / "empty"))
with pytest.raises(SourceUnavailableError) as caught:
async with HaikuRAG(config=config):
pass
message = str(caught.value)
assert "database 'haiku.rag' does not exist" in message
assert "haiku-rag init" in message
assert "create=True" in message
assert str(tmp_path) not in message
assert caught.value.__cause__ is None
@pytest.mark.asyncio
async def test_a_database_given_as_a_path_keeps_its_error(self, tmp_path):
"""The caller gave the path, so the error may name it."""
with pytest.raises(FileNotFoundError): with pytest.raises(FileNotFoundError):
async with HaikuRAG(tmp_path / "nope.lancedb"): async with HaikuRAG(tmp_path / "nope.lancedb"):
pass pass

View file

@ -8,6 +8,7 @@ from haiku.rag.client.scope import DatabaseScope
from haiku.rag.config.models import AppConfig, LanceDBConfig from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.store.exceptions import ( from haiku.rag.store.exceptions import (
AmbiguousDatabaseError, AmbiguousDatabaseError,
SourceUnavailableError,
UnknownDatabaseError, UnknownDatabaseError,
) )
from haiku.rag.utils import locate_database from haiku.rag.utils import locate_database
@ -18,18 +19,13 @@ from tests.multi_db.helpers import (
class TestConfig: class TestConfig:
def test_databases_and_uri_are_mutually_exclusive(self): def test_databases_is_the_one_placement(self):
with pytest.raises(ValidationError, match="databases"):
LanceDBConfig(
uri="s3://b/one.lancedb", databases={"one": "s3://b/one.lancedb"}
)
def test_databases_alone_is_fine(self):
config = LanceDBConfig(databases={"one": "s3://b/one.lancedb"}) config = LanceDBConfig(databases={"one": "s3://b/one.lancedb"})
assert config.databases == {"one": "s3://b/one.lancedb"} assert config.databases == {"one": "s3://b/one.lancedb"}
def test_uri_alone_is_fine(self): def test_uri_is_refused_naming_the_replacement(self):
assert LanceDBConfig(uri="s3://b/one.lancedb").databases == {} with pytest.raises(ValidationError, match="lancedb.databases"):
LanceDBConfig.model_validate({"uri": "s3://b/one.lancedb"})
class TestNamingIsRequired: class TestNamingIsRequired:
@ -51,18 +47,18 @@ class TestNamingIsRequired:
class TestNamingADatabaseDirectly: class TestNamingADatabaseDirectly:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_an_explicit_db_path_wins_over_the_configured_set( async def test_a_db_path_beside_the_configured_set_is_refused(
self, tmp_path, temp_db_path self, tmp_path, temp_db_path
): ):
"""A caller that names a path means that database, not the configured """The configuration places databases; a path beside it is a second
set: the CLI resolves `--db` to one and must not fan out instead.""" placement, and the refusal names both."""
config = _config(tmp_path, ["alpha", "beta"]) config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
async with HaikuRAG(temp_db_path, config=config, create=True) as rag: with pytest.raises(AmbiguousDatabaseError, match="alpha") as raised:
assert not rag.covers_multiple async with HaikuRAG(temp_db_path, config=config, create=True):
assert rag.source is None pass
assert rag.store.db_path == temp_db_path assert str(temp_db_path) in str(raised.value)
assert not temp_db_path.exists()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_one_configured_database_is_opened_by_name(self, tmp_path): async def test_one_configured_database_is_opened_by_name(self, tmp_path):
@ -79,42 +75,43 @@ class TestNamingADatabaseDirectly:
class TestOneConfiguredLocation: class TestOneConfiguredLocation:
"""`lancedb.uri` places one unnamed database, at a URI or at a local path.""" """One entry in `lancedb.databases` places one named database, at a URI or
at a local path."""
def _config(self, location) -> AppConfig: def _config(self, location) -> AppConfig:
return AppConfig(lancedb=LanceDBConfig(uri=str(location))) return AppConfig(lancedb=LanceDBConfig(databases={"notes": str(location)}))
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_a_local_uri_opens_the_configured_database(self, tmp_path): async def test_a_local_location_opens_the_configured_database(self, tmp_path):
located = tmp_path / "notes.lancedb" located = tmp_path / "notes.lancedb"
config = self._config(located) config = self._config(located)
async with HaikuRAG(config=config, create=True) as rag: async with HaikuRAG(config=config, create=True) as rag:
assert rag.store.db_path == located assert rag.store.db_path == located
# It places a database without naming one: only `lancedb.databases` assert rag.source == "notes"
# assigns the name results and citations carry.
assert rag.source is None
assert located.exists() assert located.exists()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_an_explicit_path_overrides_a_local_uri(self, tmp_path): async def test_a_path_beside_the_configured_database_is_refused(self, tmp_path):
"""`--db` overrides the configured location for one invocation."""
config = self._config(tmp_path / "configured.lancedb") config = self._config(tmp_path / "configured.lancedb")
chosen = tmp_path / "chosen.lancedb" chosen = tmp_path / "chosen.lancedb"
async with HaikuRAG(chosen, config=config, create=True) as rag: with pytest.raises(AmbiguousDatabaseError, match="notes"):
assert rag.store.db_path == chosen async with HaikuRAG(chosen, config=config, create=True):
assert chosen.exists() pass
assert not chosen.exists()
assert not (tmp_path / "configured.lancedb").exists() assert not (tmp_path / "configured.lancedb").exists()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_a_local_uri_that_does_not_exist_is_refused(self, tmp_path): async def test_a_local_location_that_does_not_exist_is_refused(self, tmp_path):
"""A schemeless location is a local path and must exist.""" """A schemeless location is a local path and must exist. The error names
the configured database, never its location."""
config = self._config(tmp_path / "typo.lancedb") config = self._config(tmp_path / "typo.lancedb")
with pytest.raises(FileNotFoundError): with pytest.raises(SourceUnavailableError, match="notes") as caught:
async with HaikuRAG(config=config): async with HaikuRAG(config=config):
pass pass
assert "typo.lancedb" not in str(caught.value)
assert not (tmp_path / "typo.lancedb").exists() assert not (tmp_path / "typo.lancedb").exists()
def test_a_uri_with_a_scheme_stays_a_uri(self, tmp_path): def test_a_uri_with_a_scheme_stays_a_uri(self, tmp_path):
@ -125,23 +122,54 @@ class TestOneConfiguredLocation:
config = self._config("s3://bucket/one.lancedb") config = self._config("s3://bucket/one.lancedb")
[ref] = DatabaseScope.resolve(config).databases [ref] = DatabaseScope.resolve(config).databases
one, db_path = ref.connection(config)
assert db_path is None assert ref.location == "s3://bucket/one.lancedb"
assert ConnectionMode.from_config(one) == ConnectionMode.OBJECT_STORAGE assert ConnectionMode.of(ref.location) == ConnectionMode.OBJECT_STORAGE
class TestSessionsOwnTheRef:
"""A session is built from the resolved reference and hands storage only
its location; the configuration it keeps is the one the caller named."""
@pytest.mark.asyncio
async def test_a_session_opens_the_location_with_the_undivided_config(
self, tmp_path
):
from haiku.rag.client.session import SingleDatabaseSession
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
[ref] = DatabaseScope.resolve(config, database_name="alpha").databases
session = await SingleDatabaseSession(ref, config, read_only=True).open()
try:
assert session.source == "alpha"
assert session.location == ref.location
assert session.db_path == ref.location
assert session.store.location == ref.location
assert session.store._config is config
finally:
await session.aclose()
@pytest.mark.asyncio
async def test_a_client_keeps_the_configuration_it_was_given(self, tmp_path):
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
async with HaikuRAG(config=config, sources=["alpha"]) as rag:
assert rag._config is config
assert set(rag._config.lancedb.databases) == {"alpha", "beta"}
assert rag.store.location == tmp_path / "alpha.lancedb"
class TestLocate: class TestLocate:
def test_a_scheme_is_a_uri(self): def test_a_scheme_is_a_uri(self):
assert locate_database("s3://bucket/one.lancedb") == ( assert locate_database("s3://bucket/one.lancedb") == "s3://bucket/one.lancedb"
"s3://bucket/one.lancedb",
None,
)
def test_anything_else_is_a_local_path(self): def test_anything_else_is_a_local_path(self):
uri, db_path = locate_database("/data/one.lancedb") from pathlib import Path
assert uri == ""
assert db_path is not None and str(db_path) == "/data/one.lancedb" assert locate_database("/data/one.lancedb") == Path("/data/one.lancedb")
class TestSelection: class TestSelection:
@ -216,10 +244,39 @@ class TestPlacingADatabase:
assert {r.source for r in results} == {"alpha"} assert {r.source for r in results} == {"alpha"}
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_an_unnamed_database_names_nothing(self, temp_db_path): async def test_a_database_at_a_path_is_named_by_its_stem(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag: async with HaikuRAG(temp_db_path, create=True) as rag:
assert rag.source_names == () assert rag.source_names == (temp_db_path.stem,)
assert rag.source is None assert rag.source == temp_db_path.stem
@pytest.mark.asyncio
async def test_the_default_database_is_selectable_by_name(self, tmp_path):
"""Nothing configured is the one entry `haiku.rag`, an ordinary
configured database that `sources` can name."""
from haiku.rag.config.models import StorageConfig
config = AppConfig(storage=StorageConfig(data_dir=tmp_path))
async with HaikuRAG(config=config, sources=["haiku.rag"], create=True) as rag:
assert rag.source == "haiku.rag"
assert rag.store.db_path == tmp_path / "haiku.rag.lancedb"
def test_coverage_is_known_before_the_client_opens(self, tmp_path):
"""Coverage is a fact of the resolved scope, readable before entering,
and `source_names` and `covers_multiple` agree on it."""
config = _config(tmp_path, ["alpha", "beta"])
covering = HaikuRAG(config=config)
assert covering.source_names == ("alpha", "beta")
assert covering.covers_multiple
narrowed = HaikuRAG(config=config, sources=["beta"])
assert narrowed.source_names == ("beta",)
assert not narrowed.covers_multiple
at_path = HaikuRAG(tmp_path / "other.lancedb")
assert at_path.source_names == ("other",)
assert not at_path.covers_multiple
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_the_reader_for_a_database_is_the_client_holding_it(self, tmp_path): async def test_the_reader_for_a_database_is_the_client_holding_it(self, tmp_path):
@ -280,10 +337,10 @@ class TestPlacingADatabase:
await alpha.reader_for("beta") await alpha.reader_for("beta")
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_an_unnamed_database_refuses_any_name(self, temp_db_path): async def test_a_database_at_a_path_answers_to_its_stem_alone(self, temp_db_path):
"""Nothing names it, so no name can be the one it covers."""
async with HaikuRAG(temp_db_path, create=True) as rag: async with HaikuRAG(temp_db_path, create=True) as rag:
with pytest.raises(UnknownDatabaseError, match="single unnamed database"): assert await rag.reader_for(temp_db_path.stem) is rag
with pytest.raises(UnknownDatabaseError, match=temp_db_path.stem):
await rag.reader_for("anything") await rag.reader_for("anything")
@pytest.mark.asyncio @pytest.mark.asyncio

View file

@ -71,7 +71,7 @@ class TestFederatedSearch:
class TestSingleDatabaseUnchanged: class TestSingleDatabaseUnchanged:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_source_is_unset_without_configured_databases(self, temp_db_path): async def test_source_is_the_stem_without_configured_databases(self, temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as rag: async with HaikuRAG(temp_db_path, create=True) as rag:
doc = DoclingDocument(name="one") doc = DoclingDocument(name="one")
doc.add_text(label=DocItemLabel.TEXT, text="a document about cats") doc.add_text(label=DocItemLabel.TEXT, text="a document about cats")
@ -89,7 +89,7 @@ class TestSingleDatabaseUnchanged:
results = await rag.search("cats", search_type="fts") results = await rag.search("cats", search_type="fts")
assert results assert results
assert all(r.source is None for r in results) assert all(r.source == temp_db_path.stem for r in results)
class TestOneQueryVector: class TestOneQueryVector:
@ -473,9 +473,11 @@ class TestNarrowingToOneDatabase:
assert results == [] assert results == []
class TestReciprocalRankFusion: class TestFusionWithoutAReranker:
"""Without a reranker, scores from separate indexes are not comparable, so """Without a reranker, the union is ordered by cosine similarity to the
fusion ranks by position. These pin what that produces.""" query. A search with no query vector (full-text) orders by retrieval score
instead; in both, ties resolve by within-database rank and only a tie on
both falls to configured order. These pin what that produces."""
@staticmethod @staticmethod
def _ranked(source: str, count: int, top: float) -> list[tuple[Chunk, float]]: def _ranked(source: str, count: int, top: float) -> list[tuple[Chunk, float]]:
@ -489,7 +491,7 @@ class TestReciprocalRankFusion:
second, so score order and position order disagree.""" second, so score order and position order disagree."""
return [self._ranked("a", count, 0.9), self._ranked("b", count, 0.2)] return [self._ranked("a", count, 0.9), self._ranked("b", count, 0.2)]
async def _fuse_over(self, tmp_path, per_source, limit): async def _fuse_over(self, tmp_path, per_source, limit, query_vector=None):
from haiku.rag.client.search import _fuse from haiku.rag.client.search import _fuse
config = _config(tmp_path, ["alpha", "beta"]) config = _config(tmp_path, ["alpha", "beta"])
@ -498,41 +500,217 @@ class TestReciprocalRankFusion:
async with HaikuRAG(config=config) as rag: async with HaikuRAG(config=config) as rag:
assert rag.reranker is None assert rag.reranker is None
clients = await rag.clients_for(["alpha", "beta"]) clients = await rag.clients_for(["alpha", "beta"])
fused = await _fuse(rag, clients, "cats", per_source, limit) fused = await _fuse(
rag, clients, "cats", per_source, limit, query_vector=query_vector
)
return [(owner.source, chunk.id, score) for owner, chunk, score in fused] return [(owner.source, chunk.id, score) for owner, chunk, score in fused]
@staticmethod
def _embedded(
source: str, embeddings: list[list[float]]
) -> list[tuple[Chunk, float]]:
"""A ranking whose retrieval scores descend while the embeddings are
the caller's, so cosine order and score order can be made to disagree."""
return [
(
Chunk(id=f"{source}{i}", content=f"{source} {i}", embedding=e),
0.9 - i / 100,
)
for i, e in enumerate(embeddings)
]
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_databases_interleave_by_rank(self, tmp_path): async def test_cosine_orders_the_union(self, tmp_path):
"""Each contributes its rank-1 before either contributes its rank-2.""" """With a query vector, similarity to the query decides, not the
databases' own scores or ranks."""
alpha = self._embedded("a", [[0.0, 1.0], [0.6, 0.8]])
beta = self._embedded("b", [[1.0, 0.0], [0.8, 0.6]])
fused = await self._fuse_over(
tmp_path, [alpha, beta], 10, query_vector=[1.0, 0.0]
)
assert [cid for _, cid, _ in fused] == ["b0", "b1", "a1", "a0"]
assert [round(score, 2) for _, _, score in fused] == [1.0, 0.8, 0.6, 0.0]
@pytest.mark.asyncio
async def test_cosine_ties_break_by_rank_then_configured_order(self, tmp_path):
"""Identical embeddings tie on cosine; within-database rank decides,
and equal ranks fall to configured order."""
same = [1.0, 0.0]
alpha = self._embedded("a", [same, same])
beta = self._embedded("b", [same, same])
fused = await self._fuse_over(
tmp_path, [alpha, beta], 10, query_vector=[1.0, 0.0]
)
assert [cid for _, cid, _ in fused] == ["a0", "b0", "a1", "b1"]
@pytest.mark.asyncio
async def test_a_hybrid_search_takes_the_cosine_path_end_to_end(
self, tmp_path, monkeypatch
):
"""The result scores are cosines, not retrieval scores: a fusion that
silently loses the candidate embeddings reverts to score order and
returns lancedb's hybrid values, which this pins against."""
dim = get_config().embeddings.model.vector_dim
toward = [1.0] + [0.0] * (dim - 1)
away = [0.0, 1.0] + [0.0] * (dim - 2)
config = _config(tmp_path, ["alpha", "beta"])
for name, embedding in (("alpha", away), ("beta", toward)):
async with HaikuRAG(config=config, create=True, sources=[name]) as rag:
doc = DoclingDocument(name=name)
doc.add_text(label=DocItemLabel.TEXT, text=f"{name} cats")
await rag.import_document(
doc,
[Chunk(content=f"{name} cats", embedding=embedding, order=0)],
uri=f"test://{name}",
)
async def embed_query(self, text):
return toward
monkeypatch.setattr(EmbedderWrapper, "embed_query", embed_query)
async with HaikuRAG(config=config) as rag:
results = await rag.search("cats", limit=2)
assert [r.source for r in results] == ["beta", "alpha"]
assert results[0].score == pytest.approx(1.0)
assert results[1].score == pytest.approx(0.0)
@pytest.mark.asyncio
async def test_embeddings_are_materialized_only_for_cosine_fusion(
self, tmp_path, monkeypatch, query_embedding
):
"""A reranker scores the union itself, so its 10x over-fetch must not
materialize per-chunk embeddings."""
from haiku.rag.store.repositories.chunk import ChunkRepository
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
await _seed(config, "beta", ["beta document about cats"])
asked: list[bool] = []
search = ChunkRepository.search
async def spy(self, *args, **kwargs):
asked.append(kwargs.get("with_vectors", False))
return await search(self, *args, **kwargs)
monkeypatch.setattr(ChunkRepository, "search", spy)
async with HaikuRAG(config=config) as rag:
await rag.search("cats")
assert asked == [True, True]
asked.clear()
monkeypatch.setattr(HaikuRAG, "reranker", property(lambda self: StubReranker()))
async with HaikuRAG(config=config) as rag:
await rag.search("cats")
assert asked == [False, False]
# An image query skips the reranker branch, so it takes cosine fusion
# and needs vectors even with a reranker configured.
dim = get_config().embeddings.model.vector_dim
async def embed_image(self, image):
return [0.1] * dim
monkeypatch.setattr(EmbedderWrapper, "supports_images", True)
monkeypatch.setattr(EmbedderWrapper, "embed_image", embed_image)
asked.clear()
async with HaikuRAG(config=config) as rag:
await rag.search(b"\x89PNG\r\n\x1a\n")
assert asked == [True, True]
@pytest.mark.asyncio
async def test_a_candidate_without_an_embedding_disables_the_cosine(self, tmp_path):
"""One unembedded candidate makes cosine incomparable across the union,
so the whole fusion keeps retrieval-score order."""
alpha = self._embedded("a", [[0.0, 1.0]])
beta = self._ranked("b", 1, 0.2)
fused = await self._fuse_over(
tmp_path, [alpha, beta], 10, query_vector=[1.0, 0.0]
)
assert [(cid, score) for _, cid, score in fused] == [("a0", 0.9), ("b0", 0.2)]
@pytest.mark.asyncio
async def test_the_score_orders_the_union(self, tmp_path):
"""A stronger database takes consecutive slots; breadth is not
guaranteed."""
fused = await self._fuse_over(tmp_path, self._lopsided(3), 10) fused = await self._fuse_over(tmp_path, self._lopsided(3), 10)
assert [(source, cid) for source, cid, _ in fused] == [ assert [(source, cid) for source, cid, _ in fused] == [
("alpha", "a0"), ("alpha", "a0"),
("beta", "b0"),
("alpha", "a1"), ("alpha", "a1"),
("beta", "b1"),
("alpha", "a2"), ("alpha", "a2"),
("beta", "b0"),
("beta", "b1"),
("beta", "b2"), ("beta", "b2"),
] ]
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_the_score_is_the_reciprocal_of_the_rank(self, tmp_path): async def test_the_score_is_the_retrieval_score(self, tmp_path):
"""The fused score is the candidate's own, so re-sorting downstream
(context expansion) preserves the fused order."""
fused = await self._fuse_over(tmp_path, self._lopsided(2), 10) fused = await self._fuse_over(tmp_path, self._lopsided(2), 10)
assert [score for _, _, score in fused] == [ assert [score for _, _, score in fused] == [0.9, 0.89, 0.2, 0.19]
1 / 61,
1 / 61, @pytest.mark.asyncio
1 / 62, async def test_score_ties_break_by_rank_within_the_database(self, tmp_path):
1 / 62, """Equal scores can sit at different ranks: rank depends on what the
rest of a database scored. The candidate nothing in its own database
beat wins the tie."""
per_source = [
[
(Chunk(id="a0", content="a 0"), 0.9),
(Chunk(id="a1", content="a 1"), 0.5),
],
[
(Chunk(id="b0", content="b 0"), 0.5),
(Chunk(id="b1", content="b 1"), 0.3),
],
]
fused = await self._fuse_over(tmp_path, per_source, 10)
assert [cid for _, cid, _ in fused] == ["a0", "b0", "a1", "b1"]
@pytest.mark.asyncio
async def test_the_configured_order_does_not_matter(self, tmp_path):
"""The same candidates fuse to the same list whichever database is
declared first."""
forward = await self._fuse_over(tmp_path, self._lopsided(3), 10)
(tmp_path / "swapped").mkdir()
backward = await self._fuse_over(
tmp_path / "swapped",
[self._ranked("b", 3, 0.2), self._ranked("a", 3, 0.9)],
10,
)
assert [(cid, score) for _, cid, score in forward] == [
(cid, score) for _, cid, score in backward
] ]
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_equal_scores_keep_the_configured_order(self, tmp_path): async def test_exact_ties_keep_the_configured_order(self, tmp_path):
"""Every rank ties across databases, so the tiebreak decides all of it.""" """Hybrid scores are rank-derived and tie exactly when databases agree,
fused = await self._fuse_over(tmp_path, self._lopsided(2), 10) so a genuine tie must still resolve deterministically."""
per_source = [self._ranked("a", 2, 0.9), self._ranked("b", 2, 0.9)]
fused = await self._fuse_over(tmp_path, per_source, 10)
assert [source for source, _, _ in fused] == ["alpha", "beta", "alpha", "beta"] assert [source for source, _, _ in fused] == ["alpha", "beta", "alpha", "beta"]
@pytest.mark.asyncio
async def test_rank_never_overrides_the_score(self, tmp_path):
"""A database's rank-2 with a higher score precedes another's rank-0:
allocation is content-driven, not round-robin."""
fused = await self._fuse_over(tmp_path, self._lopsided(2), 10)
assert [cid for _, cid, _ in fused] == ["a0", "a1", "b0", "b1"]
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_the_limit_cuts_the_fused_list(self, tmp_path): async def test_the_limit_cuts_the_fused_list(self, tmp_path):
"""Each database was asked for enough to fill the window on its own.""" """Each database was asked for enough to fill the window on its own."""
@ -540,8 +718,8 @@ class TestReciprocalRankFusion:
assert [(source, cid) for source, cid, _ in fused] == [ assert [(source, cid) for source, cid, _ in fused] == [
("alpha", "a0"), ("alpha", "a0"),
("beta", "b0"),
("alpha", "a1"), ("alpha", "a1"),
("alpha", "a2"),
] ]

Some files were not shown because too many files have changed in this diff Show more