Compare commits

...

9 commits

27 changed files with 800 additions and 154 deletions

View file

@ -1,6 +1,21 @@
# Changelog # Changelog
## [Unreleased] ## [Unreleased]
## [0.48.2] - 2026-06-19
### Added
- `analysis.max_executions` (default 15): caps `execute_code` calls per analysis question. Past the cap the tool returns a notice telling the skill to answer from what it has, instead of spiralling into `request_limit`. The analysis skill sets `request_limit` to 30 as a backstop.
### Changed
- `textual-image` moved to base dependencies; `ask`/`analyze`/`visualize` render image citations without the `tui` extra.
- Drop `list_documents` and `get_document` from the default RAG skill's tool set; the skill now exposes only `search` and `cite`. The tool branches remain in `create_skill_tools` and the `skill_generator` `AVAILABLE_TOOLS` set so users can still opt in when building custom skills.
### Fixed
- Analysis-skill `execute_code` no longer crashes with LanceDB `Already borrowed` when sandboxed code reads the document VFS (`content.txt`, `items.jsonl`, `toc.json`); VFS reads run on one background event loop with a single read-only connection for the sandbox's lifetime.
## [0.48.1] - 2026-05-21 ## [0.48.1] - 2026-05-21
### Changed ### Changed
@ -1595,7 +1610,8 @@ 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.48.1...HEAD [Unreleased]: https://github.com/ggozad/haiku.rag/compare/0.48.2...HEAD
[0.48.2]: https://github.com/ggozad/haiku.rag/compare/0.48.1...0.48.2
[0.48.1]: https://github.com/ggozad/haiku.rag/compare/0.48.0...0.48.1 [0.48.1]: https://github.com/ggozad/haiku.rag/compare/0.48.0...0.48.1
[0.48.0]: https://github.com/ggozad/haiku.rag/compare/0.47.0...0.48.0 [0.48.0]: https://github.com/ggozad/haiku.rag/compare/0.47.0...0.48.0
[0.47.0]: https://github.com/ggozad/haiku.rag/compare/0.46.0...0.47.0 [0.47.0]: https://github.com/ggozad/haiku.rag/compare/0.46.0...0.47.0

View file

@ -8,7 +8,7 @@ dependencies = [
"uvicorn[standard]>=0.40.0", "uvicorn[standard]>=0.40.0",
"pydantic-ai-slim[ag-ui,anthropic,openai]>=1.81.0", "pydantic-ai-slim[ag-ui,anthropic,openai]>=1.81.0",
"python-dotenv>=1.2.1", "python-dotenv>=1.2.1",
"haiku.rag-slim>=0.48.1", "haiku.rag-slim>=0.48.2",
"logfire[pydantic-ai]>=3.17.0", "logfire[pydantic-ai]>=3.17.0",
] ]

View file

@ -52,10 +52,12 @@ analysis:
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 for code execution code_timeout: 60.0 # Max seconds for code execution
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
``` ```
- **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**: Maximum seconds for each code execution (default: 60) - **code_timeout**: Maximum seconds for each code execution (default: 60)
- **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 skill is told to answer from what it has (default: 15)
See [Analysis skill](../skills/analysis.md) for usage details. See [Analysis skill](../skills/analysis.md) for usage details.

View file

@ -37,6 +37,7 @@ The interpreter is [pydantic-monty](https://github.com/pydantic/monty), isolated
- **Limited imports.** Only `json`, `re`, `math`, `pathlib`. - **Limited imports.** Only `json`, `re`, `math`, `pathlib`.
- **Execution timeout** (default 60s, configurable via `analysis.code_timeout`). - **Execution timeout** (default 60s, configurable via `analysis.code_timeout`).
- **Output truncation** (default 50000 chars, configurable via `analysis.max_output_chars`). - **Output truncation** (default 50000 chars, configurable via `analysis.max_output_chars`).
- **Execution budget** (default 15 calls, configurable via `analysis.max_executions`). Past the budget, `execute_code` returns a notice telling the skill to answer from what it has instead of running more code.
Variables persist between `execute_code` calls within one invocation, so the agent can build state step by step. A fresh sandbox is built per `client.analyze` call. Variables persist between `execute_code` calls within one invocation, so the agent can build state step by step. A fresh sandbox is built per `client.analyze` call.
@ -190,6 +191,7 @@ analysis:
name: claude-sonnet-4-20250514 name: claude-sonnet-4-20250514
code_timeout: 60.0 # Max seconds per code execution code_timeout: 60.0 # Max seconds per code execution
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
``` ```
When `analysis.model` is unset, the skill falls back to `qa.model`. When `analysis.model` is unset, the skill falls back to `qa.model`.

View file

@ -15,10 +15,10 @@ If the question requires *computation* over the corpus (counts, aggregates, comp
| Tool | Purpose | | Tool | Purpose |
|------|---------| |------|---------|
| `search(query, limit?)` | Hybrid search (vector + full-text) with section-aware context expansion. Returns `chunk_id`, content, `doc_item_refs`, `picture_refs`, `picture_captions`, source metadata. | | `search(query, limit?)` | Hybrid search (vector + full-text) with section-aware context expansion. Returns `chunk_id`, content, `doc_item_refs`, `picture_refs`, `picture_captions`, source metadata. |
| `list_documents()` | List all documents in the knowledge base. |
| `get_document(query)` | Fetch a document by ID, title, or URI. Partial matches work. |
| `cite(chunk_ids)` | Register chunk IDs as citations for the current answer. The agent calls this before writing the final response. | | `cite(chunk_ids)` | Register chunk IDs as citations for the current answer. The agent calls this before writing the final response. |
For corpus enumeration or full-document reads, reach for the [Analysis skill](analysis.md), which exposes `await list_documents()` and a `/documents/{id}/content.txt` virtual filesystem inside `execute_code`. Both are also available as opt-in tools when building a [custom skill](custom.md).
## State ## State
The skill manages a `RAGState` under the `"rag"` namespace: The skill manages a `RAGState` under the `"rag"` namespace:
@ -33,7 +33,7 @@ class RAGState(BaseModel):
- **citation_index** — All citations indexed by chunk ID. Accumulates across invocations so historical chunk IDs stay resolvable in UI scrollback. - **citation_index** — All citations indexed by chunk ID. Accumulates across invocations so historical chunk IDs stay resolvable in UI scrollback.
- **citations** — Chunk IDs registered via `cite` during the current invocation. Deduplicated, cleared at the start of each invocation. - **citations** — Chunk IDs registered via `cite` during the current invocation. Deduplicated, cleared at the start of each invocation.
- **document_filter** — SQL WHERE clause applied to `search` and `list_documents`. Persists across invocations. - **document_filter** — SQL WHERE clause applied to `search`. Persists across invocations.
- **searches** — Search results keyed by query string. Cleared at the start of each invocation. - **searches** — Search results keyed by query string. Cleared at the start of each invocation.
## `create_skill(db_path?, config?)` ## `create_skill(db_path?, config?)`
@ -99,7 +99,7 @@ state.document_filter = "uri LIKE '%helios/v4/%'"
result = await agent.run("What's the maintenance interval for the inverters?") result = await agent.run("What's the maintenance interval for the inverters?")
``` ```
The filter applies to every `search` and `list_documents` call for the rest of the session, including the model can't bypass it from inside. The filter applies to every `search` call for the rest of the session, and the model can't bypass it from inside.
### Combining with the analysis skill ### Combining with the analysis skill

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.48.1" version = "0.48.2"
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

@ -113,6 +113,7 @@ class AnalysisConfig(BaseModel):
model: ModelConfig | None = None model: ModelConfig | None = None
code_timeout: float = 60.0 code_timeout: float = 60.0
max_output_chars: int = 50_000 max_output_chars: int = 50_000
max_executions: int = 15
class PictureDescriptionConfig(BaseModel): class PictureDescriptionConfig(BaseModel):

View file

@ -1,9 +1,8 @@
import asyncio import asyncio
import atexit
import concurrent.futures
import json import json
import os import os
from collections.abc import Callable from collections.abc import AsyncIterator, Callable, Coroutine
from contextlib import asynccontextmanager
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal from typing import TYPE_CHECKING, Any, Literal
@ -19,6 +18,8 @@ from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentIte
if TYPE_CHECKING: if TYPE_CHECKING:
from pathlib import PurePosixPath from pathlib import PurePosixPath
from haiku.rag.client import HaikuRAG
@dataclass @dataclass
class SandboxResult: class SandboxResult:
@ -29,15 +30,6 @@ class SandboxResult:
success: bool success: bool
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
atexit.register(_executor.shutdown, wait=False)
def _run_async(coro: Any) -> Any:
"""Run an async coroutine from a sync context (CallbackFile read)."""
return _executor.submit(asyncio.run, coro).result()
def _build_toc( def _build_toc(
items: list["DocumentItem"], items: list["DocumentItem"],
chunk_index: dict[str, list[str]], chunk_index: dict[str, list[str]],
@ -127,11 +119,21 @@ class Sandbox:
sandbox = Sandbox(db_path, config, context) sandbox = Sandbox(db_path, config, context)
result = await sandbox.execute("x = await search('query')") result = await sandbox.execute("x = await search('query')")
result = await sandbox.execute("print(x[0]['content'])") # x persists result = await sandbox.execute("print(x[0]['content'])") # x persists
All database access runs on the event loop that drives ``execute()``. Monty's
file callbacks are synchronous and run on the interpreter's worker thread, so
they bridge back to that loop via ``run_coroutine_threadsafe``; the loop is
free during ``feed_run_async`` (the VM runs on the worker thread), so the
bridge does not deadlock. When a ``rag`` connection is supplied it is used
for every read, so an analysis run drives a single connection on a single
loop; otherwise each read opens an ephemeral read-only connection.
""" """
_db_path: Path _db_path: Path
_config: AppConfig _config: AppConfig
_context: AnalysisContext _context: AnalysisContext
_rag: "HaikuRAG | None"
_lock: "asyncio.Lock | None"
_search_results: "list[SearchResult]" _search_results: "list[SearchResult]"
_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]]]
@ -139,16 +141,21 @@ class Sandbox:
_toc_json_cache: dict[str, str] _toc_json_cache: dict[str, str]
_repl: MontyRepl | None _repl: MontyRepl | None
_vfs: OSAccess | None _vfs: OSAccess | None
_loop: asyncio.AbstractEventLoop | None
def __init__( def __init__(
self, self,
db_path: Path, db_path: Path,
config: AppConfig, config: AppConfig,
context: AnalysisContext, context: AnalysisContext,
rag: "HaikuRAG | None" = None,
lock: "asyncio.Lock | None" = None,
): ):
self._db_path = db_path self._db_path = db_path
self._config = config self._config = config
self._context = context self._context = context
self._rag = rag
self._lock = lock
self._search_results = [] self._search_results = []
self._doc_items = {} self._doc_items = {}
self._doc_chunk_index = {} self._doc_chunk_index = {}
@ -156,11 +163,42 @@ class Sandbox:
self._toc_json_cache = {} self._toc_json_cache = {}
self._repl = None self._repl = None
self._vfs = None self._vfs = None
self._loop = None
@asynccontextmanager
async def _connection(self) -> "AsyncIterator[HaikuRAG]":
"""Yield the shared connection (serialized by the lock), or an ephemeral
read-only one. The lock guards the whole block so a read's awaits cannot
interleave with another task's operation on the same connection."""
if self._rag is not None:
if self._lock is not None:
async with self._lock:
yield self._rag
else:
yield self._rag
return
from haiku.rag.client import HaikuRAG
async with HaikuRAG(self._db_path, config=self._config, read_only=True) as rag:
yield rag
def _run_on_loop(self, coro: Coroutine[Any, Any, Any]) -> Any:
"""Run a coroutine on the execute() loop from a synchronous callback.
Called from Monty's worker thread while ``feed_run_async`` leaves the
loop free, so scheduling onto it and blocking for the result is safe.
"""
assert self._loop is not None, (
"VFS reads happen during execute(); the loop must be captured first."
)
return asyncio.run_coroutine_threadsafe(coro, self._loop).result()
def close(self) -> None:
"""Retained for API compatibility; the sandbox owns no resources."""
return
def _build_external_functions(self) -> dict[str, Any]: def _build_external_functions(self) -> dict[str, Any]:
"""Build async external functions for the Monty interpreter.""" """Build async external functions for the Monty interpreter."""
db_path = self._db_path
config = self._config
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]]:
@ -169,9 +207,7 @@ class Sandbox:
# agent's Python can't do anything with them. The driving model # agent's Python can't do anything with them. The driving model
# gets figures through the top-level `search` tool when the # gets figures through the top-level `search` tool when the
# question is visual; in-code search is for structural work. # question is visual; in-code search is for structural work.
from haiku.rag.client import HaikuRAG async with self._connection() as rag:
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
results = await rag.search(query, limit=limit, filter=context.filter) results = await rag.search(query, limit=limit, filter=context.filter)
expanded = await rag.expand_context(results) expanded = await rag.expand_context(results)
self._search_results.extend(expanded) self._search_results.extend(expanded)
@ -198,9 +234,7 @@ class Sandbox:
return out return out
async def list_documents() -> list[dict[str, Any]]: async def list_documents() -> list[dict[str, Any]]:
from haiku.rag.client import HaikuRAG async with self._connection() as rag:
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
docs = await rag.list_documents(filter=context.filter) docs = await rag.list_documents(filter=context.filter)
return [ return [
{ {
@ -226,16 +260,12 @@ class Sandbox:
- items.jsonl: CallbackFile (lazy, bulk-cached) - items.jsonl: CallbackFile (lazy, bulk-cached)
- toc.json: CallbackFile (lazy, bulk-cached) - toc.json: CallbackFile (lazy, bulk-cached)
""" """
from haiku.rag.client import HaikuRAG
db_path = self._db_path
config = self._config
files: list[MemoryFile | CallbackFile] = [] files: list[MemoryFile | CallbackFile] = []
def _deny_write(_path: "PurePosixPath", _content: str | bytes) -> None: def _deny_write(_path: "PurePosixPath", _content: str | bytes) -> None:
raise PermissionError(f"Document files are read-only: {_path}") raise PermissionError(f"Document files are read-only: {_path}")
async with HaikuRAG(db_path, config=config, read_only=True) as rag: async with self._connection() as rag:
docs = await rag.list_documents(filter=self._context.filter) docs = await rag.list_documents(filter=self._context.filter)
doc_titles = {doc.id: doc.title for doc in docs if doc.id} doc_titles = {doc.id: doc.title for doc in docs if doc.id}
@ -249,12 +279,10 @@ class Sandbox:
return cached return cached
async def _fetch() -> list[DocumentItem]: async def _fetch() -> list[DocumentItem]:
from haiku.rag.client import HaikuRAG async with sandbox._connection() as rag:
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
return await rag.document_item_repository.get_all_items(did) return await rag.document_item_repository.get_all_items(did)
items = _run_async(_fetch()) items = sandbox._run_on_loop(_fetch())
sandbox._doc_items[did] = items sandbox._doc_items[did] = items
return items return items
@ -265,17 +293,15 @@ class Sandbox:
return cached return cached
async def _fetch() -> dict[str, list[str]]: async def _fetch() -> dict[str, list[str]]:
from haiku.rag.client import HaikuRAG async with sandbox._connection() as rag:
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
index = ( index = (
await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped( await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
[did] [did]
) )
) )
return index.get(did, {}) return index.get(did, {})
chunk_index = _run_async(_fetch()) chunk_index = sandbox._run_on_loop(_fetch())
sandbox._doc_chunk_index[did] = chunk_index sandbox._doc_chunk_index[did] = chunk_index
return chunk_index return chunk_index
@ -351,15 +377,11 @@ class Sandbox:
) -> Callable[["PurePosixPath"], str]: ) -> Callable[["PurePosixPath"], str]:
def read_content(_path: "PurePosixPath") -> str: def read_content(_path: "PurePosixPath") -> str:
async def _fetch() -> str: async def _fetch() -> str:
from haiku.rag.client import HaikuRAG async with sandbox._connection() as rag:
async with HaikuRAG(
db_path, config=config, read_only=True
) as rag:
content = await rag.document_repository.get_content(did) content = await rag.document_repository.get_content(did)
return content or "" return content or ""
return _run_async(_fetch()) return sandbox._run_on_loop(_fetch())
return read_content return read_content
@ -408,6 +430,8 @@ class Sandbox:
Variables persist across calls within the same Sandbox instance. Variables persist across calls within the same Sandbox instance.
""" """
# Monty's synchronous file callbacks bridge DB reads back to this loop.
self._loop = asyncio.get_running_loop()
repl, vfs = await self._ensure_initialized() repl, vfs = await self._ensure_initialized()
external_fns = self._build_external_functions() external_fns = self._build_external_functions()

View file

@ -1,6 +1,7 @@
import asyncio
from collections.abc import AsyncIterator from collections.abc import AsyncIterator
from contextlib import asynccontextmanager from contextlib import asynccontextmanager
from dataclasses import dataclass from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
@ -15,12 +16,18 @@ if TYPE_CHECKING:
@dataclass @dataclass
class RAGRunDeps(SkillRunDeps): class RAGRunDeps(SkillRunDeps):
rag: "HaikuRAG | None" = None rag: "HaikuRAG | None" = None
# pydantic-ai runs a turn's tool calls concurrently; LanceDB's per-connection
# state cannot take two in-flight operations at once, so every use of ``rag``
# (skill tools and the analysis sandbox) serializes through this lock. Always
# present so serialization is never accidentally skipped.
rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
search_count: int = 0 search_count: int = 0
@dataclass @dataclass
class AnalysisRunDeps(RAGRunDeps): class AnalysisRunDeps(RAGRunDeps):
sandbox: "Sandbox | None" = None sandbox: "Sandbox | None" = None
execute_count: int = 0
def _reset_invocation_state(state: Any) -> None: def _reset_invocation_state(state: Any) -> None:
@ -67,12 +74,19 @@ def make_analysis_lifespan(db_path: Path, config: AppConfig):
async with HaikuRAG(db_path, config=config, read_only=True) as rag: async with HaikuRAG(db_path, config=config, read_only=True) as rag:
deps.rag = rag deps.rag = rag
deps.search_count = 0 deps.search_count = 0
deps.sandbox = Sandbox( deps.execute_count = 0
sandbox = Sandbox(
db_path=db_path, db_path=db_path,
config=config, config=config,
context=AnalysisContext(filter=doc_filter), context=AnalysisContext(filter=doc_filter),
rag=rag,
lock=deps.rag_lock,
) )
deps.sandbox = sandbox
_reset_invocation_state(deps.state) _reset_invocation_state(deps.state)
yield try:
yield
finally:
sandbox.close()
return lifespan return lifespan

View file

@ -1,3 +1,5 @@
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@ -84,6 +86,23 @@ def _require_rag(ctx: RunContext[RAGRunDeps]) -> HaikuRAG:
return ctx.deps.rag return ctx.deps.rag
@asynccontextmanager
async def _serialized(ctx: RunContext[RAGRunDeps]) -> AsyncIterator[None]:
"""Serialize access to the shared connection through the run's lock.
pydantic-ai runs a turn's tool calls concurrently and LanceDB's
per-connection state cannot take two in-flight operations at once. The lock
is always present (``RAGRunDeps`` creates one by default); the no-op branch
is a guard for a missing deps/lock.
"""
lock = ctx.deps.rag_lock if ctx.deps is not None else None
if lock is None:
yield
else:
async with lock:
yield
def _register_citations(state: Any, citations: "list[Citation]") -> None: def _register_citations(state: Any, citations: "list[Citation]") -> None:
"""Add citations to the index and record cited chunk IDs for this invocation.""" """Add citations to the index and record cited chunk IDs for this invocation."""
next_index = len(state.citation_index) + 1 next_index = len(state.citation_index) + 1
@ -195,12 +214,13 @@ def create_skill_tools(
) )
state = _get_state(ctx, state_type) state = _get_state(ctx, state_type)
formatted, results = await skill_search( async with _serialized(ctx):
_require_rag(ctx), formatted, results = await skill_search(
query, _require_rag(ctx),
limit=limit, query,
document_filter=state.document_filter if state else None, limit=limit,
) document_filter=state.document_filter if state else None,
)
if state: if state:
state.searches[query] = results state.searches[query] = results
@ -221,10 +241,11 @@ def create_skill_tools(
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
"""List all documents in the knowledge base.""" """List all documents in the knowledge base."""
state = _get_state(ctx, state_type) state = _get_state(ctx, state_type)
return await skill_list_documents( async with _serialized(ctx):
_require_rag(ctx), return await skill_list_documents(
filter=state.document_filter if state else None, _require_rag(ctx),
) filter=state.document_filter if state else None,
)
tools["list_documents"] = list_documents tools["list_documents"] = list_documents
@ -238,11 +259,13 @@ def create_skill_tools(
Args: Args:
query: Document ID, title, or URI to look up. query: Document ID, title, or URI to look up.
""" """
return await skill_get_document(_require_rag(ctx), query) async with _serialized(ctx):
return await skill_get_document(_require_rag(ctx), query)
tools["get_document"] = get_document tools["get_document"] = get_document
if "execute_code" in tool_names: if "execute_code" in tool_names:
max_executions = config.analysis.max_executions
async def execute_code(ctx: RunContext[AnalysisRunDeps], code: str) -> str: async def execute_code(ctx: RunContext[AnalysisRunDeps], code: str) -> str:
"""Execute Python code in a sandboxed interpreter. """Execute Python code in a sandboxed interpreter.
@ -258,6 +281,13 @@ def create_skill_tools(
Args: Args:
code: Python code to execute. code: Python code to execute.
""" """
ctx.deps.execute_count += 1
if ctx.deps.execute_count > max_executions:
return (
"Code-execution limit reached. Give your final answer now "
"from what you already have; do not call execute_code again."
)
assert ctx.deps is not None and ctx.deps.sandbox is not None, ( assert ctx.deps is not None and ctx.deps.sandbox is not None, (
"AnalysisRunDeps.sandbox is not set — skill lifespan must run before execute_code." "AnalysisRunDeps.sandbox is not set — skill lifespan must run before execute_code."
) )
@ -325,24 +355,25 @@ def create_skill_tools(
] ]
if missing: if missing:
rag = _require_rag(ctx) async with _serialized(ctx):
synthetic: list[SearchResult] = [] rag = _require_rag(ctx)
doc_cache: dict[str, Any] = {} synthetic: list[SearchResult] = []
for cid in missing: doc_cache: dict[str, Any] = {}
chunk = await rag.get_chunk_by_id(cid) for cid in missing:
if chunk is None or not chunk.document_id: chunk = await rag.get_chunk_by_id(cid)
continue if chunk is None or not chunk.document_id:
did = chunk.document_id continue
if did in doc_cache: did = chunk.document_id
doc = doc_cache[did] if did in doc_cache:
else: doc = doc_cache[did]
doc = await rag.get_document_by_id(did) else:
doc_cache[did] = doc doc = await rag.get_document_by_id(did)
chunk.document_uri = doc.uri if doc else None doc_cache[did] = doc
chunk.document_title = doc.title if doc else None chunk.document_uri = doc.uri if doc else None
synthetic.append(SearchResult.from_chunk(chunk, score=1.0)) chunk.document_title = doc.title if doc else None
if synthetic: synthetic.append(SearchResult.from_chunk(chunk, score=1.0))
citations.extend(resolve_citations(missing, synthetic)) if synthetic:
citations.extend(resolve_citations(missing, synthetic))
if citations: if citations:
_register_citations(state, citations) _register_citations(state, citations)

View file

@ -97,4 +97,5 @@ def create_skill(
state_namespace=STATE_NAMESPACE, state_namespace=STATE_NAMESPACE,
deps_type=AnalysisRunDeps, deps_type=AnalysisRunDeps,
lifespan=make_analysis_lifespan(db_path, config), lifespan=make_analysis_lifespan(db_path, config),
request_limit=30,
) )

View file

@ -18,7 +18,7 @@ CRITICAL RULES:
3. When a skill returns citations, always include them in your response 3. When a skill returns citations, always include them in your response
""" """
_RAG_TOOLS = ["search", "list_documents", "get_document", "cite"] _RAG_TOOLS = ["search", "cite"]
def get_agent_preamble(config: AppConfig) -> str: def get_agent_preamble(config: AppConfig) -> str:

View file

@ -21,12 +21,6 @@ Each result includes:
When a result's Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text. Use the image directly to answer questions about figures, diagrams, charts, screenshots. When a result's Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text. Use the image directly to answer questions about figures, diagrams, charts, screenshots.
### list_documents
List available documents in the knowledge base. Use when the user wants to browse what's available.
### get_document
Retrieve a document by ID, title, or URI. Partial matches work. Use when the user wants the full content of a specific document.
### cite ### cite
Register the chunk IDs that ground your answer. Call this BEFORE writing your final answer, with the `chunk_id` values from search results that support each claim. Every answer that uses search results must be backed by `cite`. Register the chunk IDs that ground your answer. Call this BEFORE writing your final answer, with the `chunk_id` values from search results that support each claim. Every answer that uses search results must be backed by `cite`.
@ -52,12 +46,6 @@ You MUST call `cite` with at least one chunk ID before producing your final answ
- If the retrieved documents do not directly address the question, say: "I cannot find enough information in the knowledge base to answer this question." Do not guess or infer from tangentially related content. In this refusal case do **not** call `cite` — there is nothing to cite. - If the retrieved documents do not directly address the question, say: "I cannot find enough information in the knowledge base to answer this question." Do not guess or infer from tangentially related content. In this refusal case do **not** call `cite` — there is nothing to cite.
- Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `cite` tool separately to register citations. - Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `cite` tool separately to register citations.
## When the user mentions a specific document
If the user says "search in [doc]", "find in [doc]", or "answer from [doc]":
- Use `get_document` or `list_documents` first to identify the document
- Then search for the topic
## When search returns irrelevant results ## When search returns irrelevant results
If your first search returns results that clearly don't match the question: If your first search returns results that clearly don't match the question:

View file

@ -2,7 +2,7 @@
name = "haiku.rag-slim" name = "haiku.rag-slim"
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies" description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies"
version = "0.48.1" version = "0.48.2"
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" }
@ -35,6 +35,7 @@ dependencies = [
"python-dotenv>=1.2.2", "python-dotenv>=1.2.2",
"pyyaml>=6.0.3", "pyyaml>=6.0.3",
"rich>=14.3.3", "rich>=14.3.3",
"textual-image>=0.8.5",
"typer>=0.21.0,<0.22.0", "typer>=0.21.0,<0.22.0",
"watchfiles>=1.1.1", "watchfiles>=1.1.1",
"zstandard>=0.23.0; python_version<'3.14'", "zstandard>=0.23.0; python_version<'3.14'",
@ -56,7 +57,6 @@ cross-encoder = ["sentence-transformers>=3.0.0"]
# TUI (chat and inspect commands) # TUI (chat and inspect commands)
tui = [ tui = [
"textual>=8.2.4", "textual>=8.2.4",
"textual-image>=0.8.5",
"tree-sitter>=0.25.2", "tree-sitter>=0.25.2",
"tree-sitter-json>=0.24.8", "tree-sitter-json>=0.24.8",
] ]

View file

@ -2,7 +2,7 @@
name = "haiku.rag" name = "haiku.rag"
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling" description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling"
version = "0.48.1" version = "0.48.2"
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" }
@ -30,7 +30,7 @@ classifiers = [
] ]
dependencies = [ dependencies = [
"haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui,cross-encoder]==0.48.1", "haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui,cross-encoder]==0.48.2",
] ]
[project.scripts] [project.scripts]
@ -38,8 +38,8 @@ 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.48.1"] s3 = ["haiku.rag-slim[s3]==0.48.2"]
cross-encoder = ["haiku.rag-slim[cross-encoder]==0.48.1"] cross-encoder = ["haiku.rag-slim[cross-encoder]==0.48.2"]
[build-system] [build-system]
requires = ["hatchling"] requires = ["hatchling"]

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -1,3 +1,5 @@
import asyncio
import threading
from pathlib import Path from pathlib import Path
import pytest import pytest
@ -412,3 +414,207 @@ class TestSandboxVFS:
assert "1" in result.stdout assert "1" in result.stdout
assert "Public Doc" in result.stdout assert "Public Doc" in result.stdout
assert "Private Doc" not in result.stdout assert "Private Doc" not in result.stdout
class TestSandboxHeldConnection:
"""VFS reads must work while another connection to the same DB stays open.
Mirrors the analysis-skill lifespan, which keeps a read-only connection open
for the whole turn while sandboxed code reads the document VFS.
"""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_vfs_reads_while_connection_held(self, temp_db_path):
"""All three VFS readers work while another connection to the DB is open."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Foxes and dogs roam the quiet hills.",
uri="test://doc",
title="Doc",
)
assert doc.id
async with HaikuRAG(temp_db_path, config=config, read_only=True):
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
try:
result = await sb.execute(
"from pathlib import Path\n"
f"d = Path('/documents/{doc.id}')\n"
"print('foxes' in (d / 'content.txt').read_text().lower())\n"
"print(len((d / 'items.jsonl').read_text().strip().split('\\n')))\n"
"print('tree' in (d / 'toc.json').read_text())"
)
assert result.success, result.stderr
lines = result.stdout.strip().split("\n")
assert lines[0] == "True"
assert int(lines[1]) > 0
assert lines[2] == "True"
finally:
sb.close()
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_vfs_reads_use_injected_connection(self, temp_db_path):
"""An injected connection services VFS reads on the calling loop.
No dedicated background loop/thread is spawned: all DB access for the
sandbox runs on the loop driving execute(), through the one connection.
"""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Foxes and dogs roam the quiet hills.",
uri="test://doc",
title="Doc",
)
assert doc.id
async with HaikuRAG(temp_db_path, config=config, read_only=True) as rag:
sb = Sandbox(
db_path=temp_db_path,
config=config,
context=AnalysisContext(),
rag=rag,
)
result = await sb.execute(
"from pathlib import Path\n"
f"print(Path('/documents/{doc.id}/content.txt').read_text())"
)
assert result.success, result.stderr
assert "Foxes" in result.stdout
assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_vfs_read_concurrent_with_connection_read(self, temp_db_path):
"""A VFS read and a direct read on the shared connection run together.
Both serialize through the shared lock, so concurrent tasks never have
two operations in flight on the one connection at once.
"""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Foxes and dogs roam the quiet hills.",
uri="test://doc",
title="Doc",
)
assert doc.id
async with HaikuRAG(temp_db_path, config=config, read_only=True) as rag:
lock = asyncio.Lock()
sb = Sandbox(
db_path=temp_db_path,
config=config,
context=AnalysisContext(),
rag=rag,
lock=lock,
)
read_code = (
"from pathlib import Path\n"
f"print(Path('/documents/{doc.id}/content.txt').read_text())"
)
async def direct_read() -> str | None:
async with lock:
return await rag.document_repository.get_content(doc.id)
exec_result, content = await asyncio.gather(
sb.execute(read_code),
direct_read(),
)
assert exec_result.success, exec_result.stderr
assert "Foxes" in exec_result.stdout
assert content and "Foxes" in content
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_vfs_reads_repeatable_across_executes(self, temp_db_path):
"""Repeated VFS reads across execute() calls return consistent content."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Foxes and dogs.",
uri="test://doc",
title="Doc",
)
assert doc.id
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
try:
read = (
"from pathlib import Path\n"
f"print(Path('/documents/{doc.id}/content.txt').read_text())"
)
first = await sb.execute(read)
second = await sb.execute(read)
assert first.success and second.success
assert "Foxes" in first.stdout
assert first.stdout == second.stdout
assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
finally:
sb.close()
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_variables_persist_across_executes_after_vfs_read(self, temp_db_path):
"""REPL state persists across execute() calls, including after a VFS read."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Foxes and dogs.",
uri="test://doc",
title="Doc",
)
assert doc.id
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
try:
first = await sb.execute(
"from pathlib import Path\n"
f"x = len(Path('/documents/{doc.id}/content.txt').read_text())"
)
assert first.success, first.stderr
second = await sb.execute("print(x)")
assert second.success, second.stderr
assert int(second.stdout.strip()) > 0
finally:
sb.close()
@pytest.mark.asyncio
async def test_close_is_safe_without_vfs_read(self, temp_db_path):
"""close() is a no-op (and safe to call twice) when no VFS read happened."""
async with HaikuRAG(temp_db_path, create=True):
config = AppConfig()
sb = Sandbox(db_path=temp_db_path, config=config, context=AnalysisContext())
sb.close()
sb.close()
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_close_is_idempotent_after_vfs_read(self, temp_db_path):
"""close() is a safe no-op after a VFS read; no background thread lingers."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Foxes and dogs.",
uri="test://doc",
title="Doc",
)
assert doc.id
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
f"print(Path('/documents/{doc.id}/content.txt').read_text())"
)
assert result.success
assert not any(t.name == "sandbox-vfs" for t in threading.enumerate())
sb.close()
sb.close()

View file

@ -7,6 +7,7 @@ list of siblings; HTML/markdown corpora with real heading hierarchy get a
nested tree. Items with no section_header at all produce `tree: []`. nested tree. Items with no section_header at all produce `tree: []`.
""" """
import asyncio
import json import json
from pathlib import PurePosixPath from pathlib import PurePosixPath
@ -56,15 +57,21 @@ def _header(
) )
async def _read_toc(sandbox: Sandbox, doc_id: str) -> dict: async def _read_vfs_text(sandbox: Sandbox, path: str) -> str:
"""Read a VFS file the way execute() does: the synchronous reader runs on a
worker thread and bridges DB access back to the (free) calling loop."""
vfs = await sandbox._build_vfs() vfs = await sandbox._build_vfs()
raw = vfs.path_read_text(PurePosixPath(f"/documents/{doc_id}/toc.json")) sandbox._loop = asyncio.get_running_loop()
return await asyncio.to_thread(vfs.path_read_text, PurePosixPath(path))
async def _read_toc(sandbox: Sandbox, doc_id: str) -> dict:
raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/toc.json")
return json.loads(raw) return json.loads(raw)
async def _read_items_jsonl(sandbox: Sandbox, doc_id: str) -> list[dict]: async def _read_items_jsonl(sandbox: Sandbox, doc_id: str) -> list[dict]:
vfs = await sandbox._build_vfs() raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/items.jsonl")
raw = vfs.path_read_text(PurePosixPath(f"/documents/{doc_id}/items.jsonl"))
return [json.loads(line) for line in raw.strip().splitlines()] if raw else [] return [json.loads(line) for line in raw.strip().splitlines()] if raw else []

View file

@ -55,6 +55,14 @@ class TestAnalysisSkillCreation:
assert skill.metadata.description assert skill.metadata.description
assert skill.instructions assert skill.instructions
def test_create_skill_sets_request_limit_backstop(
self, test_app_config, temp_db_path
):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.request_limit == 30
def test_create_skill_has_expected_tools(self, test_app_config, temp_db_path): def test_create_skill_has_expected_tools(self, test_app_config, temp_db_path):
from haiku.rag.skills.analysis import create_skill from haiku.rag.skills.analysis import create_skill
@ -148,6 +156,23 @@ class TestExecuteCodeTool:
assert "ZeroDivisionError" in result assert "ZeroDivisionError" in result
assert state.executions[0].success is False assert state.executions[0].success is False
async def test_execute_code_rate_limited(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill
config = AppConfig()
config.analysis.max_executions = 2
skill = create_skill(db_path=rag_db, config=config)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state, sandbox=sandbox_factory())
await execute_code(ctx, code="print('first')")
await execute_code(ctx, code="print('second')")
result = await execute_code(ctx, code="print('third')")
assert "limit reached" in result.lower()
assert ctx.deps.execute_count == 3
assert len(state.executions) == 2
async def test_execute_code_applies_document_filter(self, rag_db, sandbox_factory): async def test_execute_code_applies_document_filter(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill from haiku.rag.skills.analysis import create_skill
@ -295,6 +320,7 @@ class TestAnalysisLifespan:
assert deps.rag is not None assert deps.rag is not None
assert deps.rag.is_read_only assert deps.rag.is_read_only
assert deps.search_count == 0 assert deps.search_count == 0
assert deps.execute_count == 0
assert isinstance(deps.sandbox, Sandbox) assert isinstance(deps.sandbox, Sandbox)
docs = await deps.rag.list_documents() docs = await deps.rag.list_documents()
assert len(docs) == 2 assert len(docs) == 2
@ -311,6 +337,16 @@ class TestAnalysisLifespan:
assert deps.sandbox is not None assert deps.sandbox is not None
assert deps.sandbox._context.filter == "title = 'AI Overview'" assert deps.sandbox._context.filter == "title = 'AI Overview'"
async def test_lifespan_resets_counts_per_invocation(self, rag_db):
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
config = AppConfig()
lifespan = make_analysis_lifespan(rag_db, config)
deps = AnalysisRunDeps(search_count=7, execute_count=42)
async with lifespan(deps):
assert deps.search_count == 0
assert deps.execute_count == 0
async def test_skill_has_lifespan_and_deps_type( async def test_skill_has_lifespan_and_deps_type(
self, test_app_config, temp_db_path self, test_app_config, temp_db_path
): ):

View file

@ -1,3 +1,5 @@
import asyncio
import pytest import pytest
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
@ -112,7 +114,7 @@ class TestRAGSkillCreation:
skill = create_skill(config=test_app_config, db_path=temp_db_path) skill = create_skill(config=test_app_config, db_path=temp_db_path)
tool_names = {getattr(t, "__name__") for t in skill.tools if callable(t)} tool_names = {getattr(t, "__name__") for t in skill.tools if callable(t)}
assert tool_names == {"search", "list_documents", "get_document", "cite"} assert tool_names == {"search", "cite"}
def test_create_skill_has_state(self, test_app_config, temp_db_path): def test_create_skill_has_state(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import RAGState, create_skill from haiku.rag.skills.rag import RAGState, create_skill
@ -169,6 +171,38 @@ class TestSearchTool:
assert isinstance(result, str) assert isinstance(result, str)
assert len(result) > 0 assert len(result) > 0
async def test_concurrent_searches_serialize_on_lock(
self, rag_db, rag_client, monkeypatch
):
"""Tool calls in a turn run concurrently; the shared lock keeps only one
connection operation in flight (pydantic-ai borrow guard)."""
from haiku.rag.skills import _tools
from haiku.rag.skills.rag import create_skill
original = _tools.skill_search
inflight = {"now": 0, "max": 0}
async def tracking_search(*args, **kwargs):
inflight["now"] += 1
inflight["max"] = max(inflight["max"], inflight["now"])
try:
await asyncio.sleep(0.05) # widen the window an overlap would use
return await original(*args, **kwargs)
finally:
inflight["now"] -= 1
monkeypatch.setattr(_tools, "skill_search", tracking_search)
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
ctx = _make_ctx(rag=rag_client) # one ctx -> one shared lock, as in a turn
results = await asyncio.gather(
search(ctx, query="artificial intelligence"),
search(ctx, query="machine learning"),
)
assert all(isinstance(r, str) and r for r in results)
assert inflight["max"] == 1
async def test_search_updates_state(self, rag_db, rag_client): async def test_search_updates_state(self, rag_db, rag_client):
from haiku.rag.skills.rag import RAGState, create_skill from haiku.rag.skills.rag import RAGState, create_skill
@ -220,52 +254,6 @@ class TestSearchTool:
assert len(state.searches) == 2 assert len(state.searches) == 2
class TestListDocumentsTool:
async def test_list_documents_returns_results(self, rag_db, rag_client):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
list_docs = _get_tool(skill, "list_documents")
ctx = _make_ctx(rag=rag_client)
results = await list_docs(ctx)
assert isinstance(results, list)
assert len(results) == 2
async def test_list_documents_applies_document_filter_from_state(
self, rag_db, rag_client
):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
list_docs = _get_tool(skill, "list_documents")
state = RAGState(document_filter="title = 'AI Overview'")
ctx = _make_ctx(state, rag=rag_client)
results = await list_docs(ctx)
assert len(results) == 1
assert results[0]["title"] == "AI Overview"
class TestGetDocumentTool:
async def test_get_document_by_title(self, rag_db, rag_client):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
get_doc = _get_tool(skill, "get_document")
ctx = _make_ctx(rag=rag_client)
result = await get_doc(ctx, query="AI Overview")
assert result is not None
assert result["title"] == "AI Overview"
async def test_get_document_not_found(self, rag_db, rag_client):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
get_doc = _get_tool(skill, "get_document")
ctx = _make_ctx(rag=rag_client)
result = await get_doc(ctx, query="nonexistent document xyz")
assert result is None
class TestCiteTool: class TestCiteTool:
async def test_cite_registers_citations(self, rag_db, rag_client): async def test_cite_registers_citations(self, rag_db, rag_client):
from haiku.rag.skills.rag import RAGState, create_skill from haiku.rag.skills.rag import RAGState, create_skill

View file

@ -24,6 +24,7 @@ from haiku.rag.skills._deps import RAGRunDeps
from haiku.rag.skills._tools import create_skill_tools from haiku.rag.skills._tools import create_skill_tools
from haiku.rag.skills.rag import RAGState from haiku.rag.skills.rag import RAGState
from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document import Document
def _make_png(color: str = "red") -> bytes: def _make_png(color: str = "red") -> bytes:
@ -239,3 +240,80 @@ async def test_skill_search_keeps_same_self_ref_from_different_documents():
payloads = {part.data for part in result.content} # type: ignore[attr-defined] payloads = {part.data for part in result.content} # type: ignore[attr-defined]
assert PICTURE_BYTES in payloads assert PICTURE_BYTES in payloads
assert other_bytes in payloads assert other_bytes in payloads
def _build_tool(config: AppConfig, name: str):
tools = create_skill_tools(
db_path=Path("/tmp/unused.lancedb"),
config=config,
state_type=RAGState,
tool_names=[name],
model=config.qa.model,
)
return tools[name]
@pytest.mark.asyncio
async def test_list_documents_tool_returns_shaped_dicts():
config = AppConfig()
list_documents = _build_tool(config, "list_documents")
rag = AsyncMock()
rag.list_documents = AsyncMock(
return_value=[
Document(id="d1", content="x", title="AI", uri="test://ai"),
Document(id="d2", content="y", title="ML", uri="test://ml"),
]
)
ctx = _make_ctx(rag, RAGState())
results = await list_documents(ctx)
assert [r["title"] for r in results] == ["AI", "ML"]
assert all(
set(r.keys()) == {"id", "title", "uri", "metadata", "created_at", "updated_at"}
for r in results
)
@pytest.mark.asyncio
async def test_list_documents_tool_forwards_document_filter_from_state():
config = AppConfig()
list_documents = _build_tool(config, "list_documents")
rag = AsyncMock()
rag.list_documents = AsyncMock(return_value=[])
state = RAGState(document_filter="title = 'AI Overview'")
ctx = _make_ctx(rag, state)
await list_documents(ctx)
rag.list_documents.assert_awaited_once_with(filter="title = 'AI Overview'")
@pytest.mark.asyncio
async def test_get_document_tool_returns_shaped_dict():
config = AppConfig()
get_document = _build_tool(config, "get_document")
rag = AsyncMock()
rag.resolve_document = AsyncMock(
return_value=Document(id="d1", content="full text", title="AI", uri="test://ai")
)
ctx = _make_ctx(rag, RAGState())
result = await get_document(ctx, "AI")
assert result is not None
assert result["content"] == "full text"
assert result["title"] == "AI"
@pytest.mark.asyncio
async def test_get_document_tool_returns_none_when_missing():
config = AppConfig()
get_document = _build_tool(config, "get_document")
rag = AsyncMock()
rag.resolve_document = AsyncMock(return_value=None)
ctx = _make_ctx(rag, RAGState())
result = await get_document(ctx, "nonexistent")
assert result is None

10
uv.lock
View file

@ -1432,7 +1432,7 @@ wheels = [
[[package]] [[package]]
name = "haiku-rag" name = "haiku-rag"
version = "0.48.1" version = "0.48.2"
source = { editable = "." } source = { editable = "." }
dependencies = [ dependencies = [
{ name = "haiku-rag-slim", extra = ["cohere", "cross-encoder", "docling", "mxbai", "tui", "voyageai", "zeroentropy"] }, { name = "haiku-rag-slim", extra = ["cohere", "cross-encoder", "docling", "mxbai", "tui", "voyageai", "zeroentropy"] },
@ -1497,7 +1497,7 @@ dev = [
[[package]] [[package]]
name = "haiku-rag-evals" name = "haiku-rag-evals"
version = "0.48.1" version = "0.48.2"
source = { editable = "evaluations" } source = { editable = "evaluations" }
dependencies = [ dependencies = [
{ name = "datasets" }, { name = "datasets" },
@ -1520,7 +1520,7 @@ requires-dist = [
[[package]] [[package]]
name = "haiku-rag-slim" name = "haiku-rag-slim"
version = "0.48.1" version = "0.48.2"
source = { editable = "haiku_rag_slim" } source = { editable = "haiku_rag_slim" }
dependencies = [ dependencies = [
{ name = "docling-core" }, { name = "docling-core" },
@ -1536,6 +1536,7 @@ dependencies = [
{ name = "python-dotenv" }, { name = "python-dotenv" },
{ name = "pyyaml" }, { name = "pyyaml" },
{ name = "rich" }, { name = "rich" },
{ name = "textual-image" },
{ name = "typer" }, { name = "typer" },
{ name = "watchfiles" }, { name = "watchfiles" },
{ name = "zstandard", marker = "python_full_version < '3.14'" }, { name = "zstandard", marker = "python_full_version < '3.14'" },
@ -1580,7 +1581,6 @@ s3 = [
] ]
tui = [ tui = [
{ name = "textual" }, { name = "textual" },
{ name = "textual-image" },
{ name = "tree-sitter" }, { name = "tree-sitter" },
{ name = "tree-sitter-json" }, { name = "tree-sitter-json" },
] ]
@ -1623,7 +1623,7 @@ requires-dist = [
{ name = "rich", specifier = ">=14.3.3" }, { name = "rich", specifier = ">=14.3.3" },
{ name = "sentence-transformers", marker = "extra == 'cross-encoder'", specifier = ">=3.0.0" }, { name = "sentence-transformers", marker = "extra == 'cross-encoder'", specifier = ">=3.0.0" },
{ name = "textual", marker = "extra == 'tui'", specifier = ">=8.2.4" }, { name = "textual", marker = "extra == 'tui'", specifier = ">=8.2.4" },
{ name = "textual-image", marker = "extra == 'tui'", specifier = ">=0.8.5" }, { name = "textual-image", specifier = ">=0.8.5" },
{ name = "torch", marker = "extra == 'jina'", specifier = ">=2.0.0" }, { name = "torch", marker = "extra == 'jina'", specifier = ">=2.0.0" },
{ name = "transformers", marker = "extra == 'jina'", specifier = ">=4.40.0" }, { name = "transformers", marker = "extra == 'jina'", specifier = ">=4.40.0" },
{ name = "transformers", marker = "extra == 'mxbai'", specifier = ">=4.49.0,<5.0.0" }, { name = "transformers", marker = "extra == 'mxbai'", specifier = ">=4.49.0,<5.0.0" },