danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #351 from ggozad/feat/async-lancedb

Browse source 
refactor: native async LanceDB (non-blocking DB I/O)
This commit is contained in:
Yiorgis Gozadinos 2026-04-24 17:47:12 +03:00 committed by GitHub
commit 5e3f7fb3a6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
63 changed files with 5294 additions and 3630 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
CHANGELOG.md
View file

@ -1,9 +1,17 @@
# Changelog
## [Unreleased]
### Changed
- **Native async LanceDB**: all table I/O now uses LanceDB's async API (`connect_async`, `AsyncConnection`, `AsyncTable`). Previously, repository methods were declared `async def` but called blocking sync LanceDB under the hood, stalling the event loop on every read/write. No change to the documented `async with HaikuRAG(...) as client:` usage pattern.
- **BREAKING (internal): `HaikuRAG` must be used via `async with`.** Store initialization now happens in `__aenter__`; constructing `HaikuRAG(...)` and calling methods directly without entering the context manager no longer works.
- **BREAKING (internal): `download_models` is no longer a method on `HaikuRAG`.** It's now a module-level function: `from haiku.rag.client.downloads import download_models; async for progress in download_models(config): ...`. The CLI and in-repo consumers are updated.
- **Concurrency: background vacuum tracked as a task** on the client. `__aexit__` and `rebuild_database` now await it explicitly, preventing `CreateIndex transaction was preempted` commit conflicts when destructive operations follow a `create_document` that scheduled a background vacuum.
### Fixed
- **Chat TUI now renders citations again.** After the 0.42.1 flattening of skill state `citations` to `list[str]`, the TUI still indexed `citations[-1]` and iterated the resulting chunk-id string character-by-character, so no citations resolved through `citation_index` and the citation panel stayed empty. Fixed by iterating `state.citations` directly.
- **`search(..., filter=...)` no longer silently under-returns.** The filter path used to materialize LanceDB's top-N window, filter to matching `document_id`s in pandas, and `head(limit)`. When matching chunks lived outside that top-N window (selective filters, broad queries), the caller got fewer than `limit` results even though plenty of matching chunks existed in the index. The document filter is now pushed down into the chunk query as `document_id IN (...)` so `.limit(limit)` applies to matching chunks directly. Behavior change: searches that previously under-returned will start returning the requested count.
## [0.42.1] - 2026-04-22

42
app/backend/main.py
View file

@ -1,5 +1,7 @@
import asyncio
import logging
import os
from contextlib import asynccontextmanager
from pathlib import Path
from dotenv import find_dotenv, load_dotenv
@ -58,13 +60,23 @@ logger.info(f"QA Provider: {Config.qa.model.provider}, Model: {Config.qa.model.n
# Only HaikuRAG client is a singleton (expensive to create)
_client: HaikuRAG | None = None
_client_lock = asyncio.Lock()
def get_client() -> HaikuRAG:
"""Get or create cached client."""
async def get_client() -> HaikuRAG:
"""Get or create the cached client.
Guarded by a lock because the first request after startup can race with
itself: two concurrent callers would both pass the None check, each build
and enter a HaikuRAG, and the loser would leak its LanceDB connection.
"""
global _client
if _client is None:
_client = HaikuRAG(db_path=db_path, config=Config, create=True)
async with _client_lock:
if _client is None:
client = HaikuRAG(db_path=db_path, config=Config, create=True)
await client.__aenter__()
_client = client
return _client
@ -126,7 +138,7 @@ async def list_documents(_: Request) -> JSONResponse:
if not db_path.exists():
return JSONResponse({"documents": [], "error": "Database not found"})
client = get_client()
client = await get_client()
docs = await client.document_repository.list_all()
return JSONResponse(
{
@ -151,8 +163,8 @@ async def db_info(_: Request) -> JSONResponse:
from haiku.rag.store.engine import get_database_stats
client = get_client()
stats = get_database_stats(client.store.db)
client = await get_client()
stats = await get_database_stats(client.store.db)
return JSONResponse(
{
@ -177,7 +189,7 @@ async def visualize_chunk(request: Request) -> JSONResponse:
if not db_path.exists():
return JSONResponse({"error": "Database not found"}, status_code=404)
client = get_client()
client = await get_client()
chunk = await client.chunk_repository.get_by_id(chunk_id)
if not chunk:
@ -203,6 +215,21 @@ async def visualize_chunk(request: Request) -> JSONResponse:
)
@asynccontextmanager
async def lifespan(app: Starlette):
"""Shut down the cached HaikuRAG client cleanly on app exit.
Awaits any in-flight background vacuum tasks and closes the LanceDB
connection. Without this, vacuum tasks are cancelled abruptly and the
connection is never closed on process shutdown.
"""
yield
global _client
if _client is not None:
await _client.__aexit__(None, None, None)
_client = None
# Create Starlette app
app = Starlette(
routes=[
@ -221,6 +248,7 @@ app = Starlette(
allow_headers=["*"],
)
],
lifespan=lifespan,
)
if __name__ == "__main__":

2
app/frontend/components/Chat.tsx
View file

@ -376,7 +376,7 @@ function MessageViewWithCitations({
if (latestCitations.length > 0) {
result.push(
<CitationBlock
key={`citations-${i}`}
key={`citations-${msg.id}`}
citations={latestCitations}
/>,
);

23
app/frontend/components/CitationBlock.tsx
View file

@ -1,6 +1,6 @@
"use client";
import { useCallback, useState } from "react";
import { useCallback, useEffect, useRef, useState } from "react";
import type { Citation } from "../lib/sessionStorage";
interface CitationBlockProps {
@ -73,8 +73,20 @@ export default function CitationBlock({ citations }: CitationBlockProps) {
loading: false,
error: null,
});
// AbortController for the in-flight visualize fetch so a rapid close/reopen
// doesn't let a stale response overwrite the new request's state.
const abortRef = useRef<AbortController | null>(null);
// Abort any in-flight request on unmount.
useEffect(() => {
return () => abortRef.current?.abort();
}, []);
const fetchVisualGrounding = useCallback(async (chunkId: string) => {
abortRef.current?.abort();
const controller = new AbortController();
abortRef.current = controller;
setVisualGrounding({
isOpen: true,
chunkId,
@ -84,10 +96,12 @@ export default function CitationBlock({ citations }: CitationBlockProps) {
});
try {
const backendUrl = process.env.NEXT_PUBLIC_BACKEND_URL || "";
const response = await fetch(`${backendUrl}/api/visualize/${chunkId}`);
const response = await fetch(`/api/visualize/${chunkId}`, {
signal: controller.signal,
});
const data = await response.json();
if (controller.signal.aborted) return;
if (!response.ok) {
throw new Error(data.error || "Failed to fetch visual grounding");
}
@ -99,6 +113,7 @@ export default function CitationBlock({ citations }: CitationBlockProps) {
error: data.images?.length === 0 ? data.message : null,
}));
} catch (err) {
if (controller.signal.aborted) return;
setVisualGrounding((prev) => ({
...prev,
loading: false,
@ -108,6 +123,8 @@ export default function CitationBlock({ citations }: CitationBlockProps) {
}, []);
const closeVisualGrounding = useCallback(() => {
abortRef.current?.abort();
abortRef.current = null;
setVisualGrounding({
isOpen: false,
chunkId: null,

3
app/frontend/components/DbInfo.tsx
View file

@ -25,8 +25,7 @@ export default function DbInfo() {
const [error, setError] = useState<string | null>(null);
useEffect(() => {
const backendUrl = process.env.NEXT_PUBLIC_BACKEND_URL || "";
fetch(`${backendUrl}/api/info`)
fetch("/api/info")
.then((res) => res.json())
.then(setInfo)
.catch((err) => setError(err.message));

75
app/frontend/components/DocumentFilter.tsx
View file

@ -16,6 +16,8 @@ interface DocumentFilterProps {
onApply: (selected: string[]) => void;
}
const getDisplayName = (doc: Document) => doc.title || doc.uri || doc.id;
export default function DocumentFilter({
isOpen,
onClose,
@ -26,33 +28,40 @@ export default function DocumentFilter({
const [documents, setDocuments] = useState<Document[]>([]);
const [loading, setLoading] = useState(true);
const [searchTerm, setSearchTerm] = useState("");
const [localSelected, setLocalSelected] = useState<Set<string>>(
new Set(selected),
);
// Track selection by document id — two docs can share a title, but ids
// are unique. Display names are only used for rendering and for the
// filter string returned to the parent.
const [localSelected, setLocalSelected] = useState<Set<string>>(new Set());
// Reset local state when modal opens
// Refetch on every open so newly-added or deleted documents show up.
useEffect(() => {
if (isOpen) {
setLocalSelected(new Set(selected));
setSearchTerm("");
}
}, [isOpen, selected]);
if (!isOpen) return;
setLoading(true);
fetch("/api/documents")
.then((res) => res.json())
.then((data) => {
setDocuments(data.documents || []);
setLoading(false);
})
.catch(() => {
setLoading(false);
});
}, [isOpen]);
// Fetch documents when modal opens
// Seed local selection from the parent's display-name list once documents
// are available. Any doc whose display name is in `selected` starts checked.
useEffect(() => {
if (isOpen && documents.length === 0) {
setLoading(true);
fetch("/api/documents")
.then((res) => res.json())
.then((data) => {
setDocuments(data.documents || []);
setLoading(false);
})
.catch(() => {
setLoading(false);
});
}
}, [isOpen, documents.length]);
if (!isOpen) return;
const selectedNames = new Set(selected);
setLocalSelected(
new Set(
documents
.filter((d) => selectedNames.has(getDisplayName(d)))
.map((d) => d.id),
),
);
setSearchTerm("");
}, [isOpen, selected, documents]);
const handleKeyDown = useCallback(
(e: React.KeyboardEvent) => {
@ -63,20 +72,24 @@ export default function DocumentFilter({
[onClose],
);
const toggleDocument = (displayName: string) => {
const toggleDocument = (docId: string) => {
setLocalSelected((prev) => {
const next = new Set(prev);
if (next.has(displayName)) {
next.delete(displayName);
if (next.has(docId)) {
next.delete(docId);
} else {
next.add(displayName);
next.add(docId);
}
return next;
});
};
const handleApply = () => {
onApply(Array.from(localSelected));
const names = documents
.filter((d) => localSelected.has(d.id))
.map(getDisplayName);
// Dedupe: two selected docs sharing a title collapse to one filter term.
onApply(Array.from(new Set(names)));
onClose();
};
@ -84,8 +97,6 @@ export default function DocumentFilter({
setLocalSelected(new Set());
};
const getDisplayName = (doc: Document) => doc.title || doc.uri || doc.id;
const filteredDocuments = documents.filter((doc) => {
if (!searchTerm) return true;
const displayName = getDisplayName(doc).toLowerCase();
@ -144,8 +155,8 @@ export default function DocumentFilter({
<label key={doc.id} className="filter-item">
<input
type="checkbox"
checked={localSelected.has(displayName)}
onChange={() => toggleDocument(displayName)}
checked={localSelected.has(doc.id)}
onChange={() => toggleDocument(doc.id)}
/>
<span className="filter-item-label">{displayName}</span>
</label>

3
docs/configuration/index.md
View file

@ -171,7 +171,8 @@ custom_config = AppConfig(
)
# Pass configuration to the client
client = HaikuRAG(config=custom_config)
async with HaikuRAG(config=custom_config) as client:
...
```
If you don't pass a config, the client uses the global configuration loaded from your YAML file or defaults.

2
docs/development.md
View file

@ -102,7 +102,7 @@ Integration tests are skipped in CI but run locally when you have the required s
```bash
uv run ruff check
uv run ruff format
uv run pyright
uv run ty check
```
## Mock API Keys

100
haiku_rag_slim/haiku/rag/app.py
View file

@ -59,8 +59,8 @@ class HaikuRAGApp: # pragma: no cover
return
# Create the database
client = HaikuRAG(db_path=self.db_path, config=self.config, create=True)
client.close()
async with HaikuRAG(db_path=self.db_path, config=self.config, create=True):
pass
self.console.print(
f"[bold green]Database initialized at {self._display_path}[/bold green]"
)
@ -88,8 +88,8 @@ class HaikuRAGApp: # pragma: no cover
# Connect directly. Don't go through Store so a database that is
# missing tables (e.g. pre-migration) still reports what it can.
db = connect_lancedb(self.config, self.db_path)
stats = get_database_stats(db)
db = await connect_lancedb(self.config, self.db_path)
stats = await get_database_stats(db)
if not any(entry["exists"] for entry in stats.values()):
self.console.print(
@ -104,14 +104,10 @@ class HaikuRAGApp: # pragma: no cover
embed_model = "unknown"
vector_dim = None
if stats["settings"]["exists"]:
settings_tbl = db.open_table("settings")
settings_tbl = await db.open_table("settings")
rows = (
settings_tbl.search()
.where("id = 'settings'")
.limit(1)
.to_arrow()
.to_pylist()
)
await settings_tbl.query().where("id = 'settings'").limit(1).to_arrow()
).to_pylist()
if rows:
raw = rows[0].get("settings") or "{}"
data = json.loads(raw) if isinstance(raw, str) else (raw or {})
@ -237,50 +233,46 @@ class HaikuRAGApp: # pragma: no cover
self.console.print("[red]Database path does not exist.[/red]")
return
store = Store(
async with Store(
self.db_path,
config=self.config,
skip_validation=True,
read_only=True,
skip_migration_check=True,
before=self.before,
)
) as store:
tables = ["documents", "chunks", "settings"]
if table:
if table not in tables:
self.console.print(
f"[red]Unknown table: {table}. Must be one of: {', '.join(tables)}[/red]"
)
return
tables = [table]
tables = ["documents", "chunks", "settings"]
if table:
if table not in tables:
self.console.print(
f"[red]Unknown table: {table}. Must be one of: {', '.join(tables)}[/red]"
)
store.close()
return
tables = [table]
self.console.print("[bold]Version History[/bold]")
self.console.print("[bold]Version History[/bold]")
for table_name in tables:
versions = await store.list_table_versions(table_name)
for table_name in tables:
versions = store.list_table_versions(table_name)
# Sort by version descending (newest first)
versions = sorted(versions, key=lambda v: v["version"], reverse=True)
# Sort by version descending (newest first)
versions = sorted(versions, key=lambda v: v["version"], reverse=True)
if limit:
versions = versions[:limit]
if limit:
versions = versions[:limit]
self.console.print(f"\n[bold cyan]{table_name}[/bold cyan]")
self.console.print(f"\n[bold cyan]{table_name}[/bold cyan]")
if not versions:
self.console.print(" [dim]No versions found[/dim]")
continue
if not versions:
self.console.print(" [dim]No versions found[/dim]")
continue
for v in versions:
version_num = v["version"]
timestamp = v["timestamp"]
self.console.print(
f" [repr.attrib_name]v{version_num}[/repr.attrib_name]: {timestamp}"
)
store.close()
for v in versions:
version_num = v["version"]
timestamp = v["timestamp"]
self.console.print(
f" [repr.attrib_name]v{version_num}[/repr.attrib_name]: {timestamp}"
)
async def list_documents(self, filter: str | None = None):
async with HaikuRAG(
@ -601,7 +593,7 @@ class HaikuRAGApp: # pragma: no cover
await client.vacuum()
self.console.print("[bold green]Vacuum completed successfully.[/bold green]")
def migrate(self) -> list[str]:
async def migrate(self) -> list[str]:
"""Run pending database migrations.
Returns:
@ -609,17 +601,13 @@ class HaikuRAGApp: # pragma: no cover
"""
from haiku.rag.store.engine import Store
store = Store(
async with Store(
self.db_path,
config=self.config,
skip_validation=True,
skip_migration_check=True,
)
try:
applied = store.migrate()
return applied
finally:
store.close()
) as store:
return await store.migrate()
async def create_index(self):
"""Create vector index on the chunks table."""
@ -630,7 +618,7 @@ class HaikuRAGApp: # pragma: no cover
read_only=self.read_only,
before=self.before,
) as client:
row_count = client.store.chunks_table.count_rows()
row_count = await client.store.chunks_table.count_rows()
self.console.print(f"Chunks in database: {row_count}")
if row_count < 256:
@ -640,7 +628,7 @@ class HaikuRAGApp: # pragma: no cover
return
# Check if index already exists
indices = client.store.chunks_table.list_indices()
indices = await client.store.chunks_table.list_indices()
has_vector_index = any("vector" in str(idx).lower() for idx in indices)
if has_vector_index:
@ -650,23 +638,21 @@ class HaikuRAGApp: # pragma: no cover
else:
self.console.print("[bold]Creating vector index...[/bold]")
client.store._ensure_vector_index()
await client.store._ensure_vector_index()
self.console.print(
"[bold green]Vector index created successfully.[/bold green]"
)
async def download_models(self):
"""Download Docling, HuggingFace tokenizer, and Ollama models per config."""
from haiku.rag.client import HaikuRAG
client = HaikuRAG(db_path=None, config=self.config)
from haiku.rag.client.downloads import download_models
progress: Progress | None = None
task_id: TaskID | None = None
current_model = ""
current_digest = ""
async for event in client.download_models():
async for event in download_models(self.config):
if event.status == "start":
self.console.print(
f"[bold blue]Downloading {event.model}...[/bold blue]"

2
haiku_rag_slim/haiku/rag/cli.py
View file

@ -529,7 +529,7 @@ def migrate( # pragma: no cover
):
app = create_app(db)
try:
applied = app.migrate()
applied = asyncio.run(app.migrate())
if applied:
typer.echo(f"Applied {len(applied)} migration(s):")
for desc in applied:

1839
haiku_rag_slim/haiku/rag/client.py
View file

File diff suppressed because it is too large Load diff

390
haiku_rag_slim/haiku/rag/client/__init__.py Normal file
View file

@ -0,0 +1,390 @@
import asyncio
import hashlib
import json
import logging
import mimetypes
import tempfile
from collections.abc import AsyncGenerator
from datetime import datetime
from enum import Enum
from pathlib import Path
from typing import TYPE_CHECKING, overload
from urllib.parse import urlparse
import httpx
from haiku.rag.config import AppConfig, Config
from haiku.rag.converters import get_converter
from haiku.rag.reranking import get_reranker
from haiku.rag.store.engine import Store
from haiku.rag.store.models.chunk import Chunk, SearchResult
from haiku.rag.store.models.document import Document
from haiku.rag.store.models.document_item import extract_items
from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository
from haiku.rag.store.repositories.document_item import DocumentItemRepository
from haiku.rag.store.repositories.settings import SettingsRepository
from haiku.rag.utils import escape_sql_string
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.agents.analysis.models import AnalysisResult
from haiku.rag.agents.research.models import (
Citation,
ResearchReport,
)
logger = logging.getLogger(__name__)
class RebuildMode(Enum):
"""Mode for rebuilding the database."""
FULL = "full" # Re-convert from source, re-chunk, re-embed
RECHUNK = "rechunk" # Re-chunk from existing content, re-embed
EMBED_ONLY = "embed_only" # Keep chunks, only regenerate embeddings
TITLE_ONLY = "title_only" # Only generate titles for untitled documents
class HaikuRAG:
"""High-level haiku-rag client."""
def __init__(
self,
db_path: Path | None = None,
config: AppConfig = Config,
skip_validation: bool = False,
create: bool = False,
read_only: bool = False,
before: datetime | None = None,
):
"""Initialize the RAG client with a database path.
Args:
db_path: Path to the database file. If None, uses config.storage.data_dir.
config: Configuration to use. Defaults to global Config.
skip_validation: Whether to skip configuration validation on database load.
create: Whether to create the database if it doesn't exist.
read_only: Whether to open the database in read-only mode.
before: Query the database as it existed at this datetime.
Implies read_only=True.
"""
self._config = config
if db_path is None:
db_path = self._config.storage.data_dir / "haiku.rag.lancedb"
self._db_path = db_path
self._skip_validation = skip_validation
self._create = create
self._read_only = read_only
self._before = before
self._vacuum_tasks: set[asyncio.Task] = set()
@property
def is_read_only(self) -> bool:
"""Whether the client is in read-only mode."""
return self.store.is_read_only
async def __aenter__(self):
"""Async context manager entry — initializes store and repositories."""
self.store = Store(
self._db_path,
config=self._config,
skip_validation=self._skip_validation,
create=self._create,
read_only=self._read_only,
before=self._before,
)
# If _initialize fails mid-way (e.g. migration check raises after
# connect), close the store so we don't leak the LanceDB connection —
# __aexit__ won't run because the `async with` never entered.
try:
await self.store._initialize()
except BaseException:
self.store.close()
raise
self.document_repository = DocumentRepository(self.store)
self.chunk_repository = ChunkRepository(self.store)
self.document_item_repository = DocumentItemRepository(self.store)
return self
async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: ARG002
"""Async context manager exit."""
await self._await_vacuum_tasks()
self.close()
return False
async def _await_vacuum_tasks(self) -> None:
"""Wait for all in-flight background vacuum tasks to complete.
Each create_document / update_document can schedule its own vacuum task;
all must be awaited before tearing down the connection, not just the
most recently scheduled one.
"""
if self._vacuum_tasks:
await asyncio.gather(*self._vacuum_tasks, return_exceptions=True)
def _schedule_vacuum(self) -> None:
"""Schedule a background vacuum and track the task for later awaiting."""
task = asyncio.create_task(self.store.vacuum())
self._vacuum_tasks.add(task)
task.add_done_callback(self._vacuum_tasks.discard)
# =========================================================================
# Processing Primitives
# =========================================================================
@overload
async def convert(self, source: Path) -> "DoclingDocument": ...
@overload
async def convert(
self, source: str, *, format: str = "md"
) -> "DoclingDocument": ...
async def convert(
self, source: Path | str, *, format: str = "md"
) -> "DoclingDocument":
from haiku.rag.client.processing import convert
return await convert(self._config, source, format=format)
async def chunk(self, docling_document: "DoclingDocument") -> list[Chunk]:
from haiku.rag.client.processing import chunk
return await chunk(self._config, docling_document)
# =========================================================================
# Title Generation
# =========================================================================
async def generate_title(self, document: Document) -> str | None:
from haiku.rag.client.titles import generate_title
return await generate_title(self._config, document)
async def create_document(
self,
content: str,
uri: str | None = None,
title: str | None = None,
metadata: dict | None = None,
format: str = "md",
) -> Document:
from haiku.rag.client.documents import create_document
return await create_document(self, content, uri, title, metadata, format)
async def import_document(
self,
docling_document: "DoclingDocument",
chunks: list[Chunk],
uri: str | None = None,
title: str | None = None,
metadata: dict | None = None,
) -> Document:
from haiku.rag.client.documents import import_document
return await import_document(
self, docling_document, chunks, uri, title, metadata
)
async def create_document_from_source(
self,
source: str | Path,
title: str | None = None,
metadata: dict | None = None,
) -> Document | list[Document]:
from haiku.rag.client.documents import create_document_from_source
return await create_document_from_source(self, source, title, metadata)
async def update_document(
self,
document_id: str,
content: str | None = None,
metadata: dict | None = None,
chunks: list[Chunk] | None = None,
title: str | None = None,
docling_document: "DoclingDocument | None" = None,
) -> Document:
from haiku.rag.client.documents import update_document
return await update_document(
self,
document_id,
content,
metadata,
chunks,
title,
docling_document,
)
async def get_document_by_id(self, document_id: str) -> Document | None:
"""Get a document by its ID.
Args:
document_id: The unique identifier of the document.
Returns:
The Document instance if found, None otherwise.
"""
return await self.document_repository.get_by_id(document_id)
async def get_chunk_by_id(self, chunk_id: str) -> Chunk | None:
"""Get a chunk by its ID.
Args:
chunk_id: The unique identifier of the chunk.
Returns:
The Chunk instance if found, None otherwise.
"""
return await self.chunk_repository.get_by_id(chunk_id)
async def get_document_by_uri(self, uri: str) -> Document | None:
"""Get a document by its URI.
Args:
uri: The URI identifier of the document.
Returns:
The Document instance if found, None otherwise.
"""
return await self.document_repository.get_by_uri(uri)
async def resolve_document(self, id_or_title: str) -> Document | None:
"""Resolve a document by ID, title, or URI (in that order).
Args:
id_or_title: Document ID, title, or URI to look up.
Returns:
The Document instance if found, None otherwise.
"""
doc = await self.get_document_by_id(id_or_title)
if doc:
return doc
safe_input = escape_sql_string(id_or_title)
docs = await self.list_documents(filter=f"title = '{safe_input}'")
if docs and docs[0].id:
return await self.get_document_by_id(docs[0].id)
docs = await self.list_documents(filter=f"uri = '{safe_input}'")
if docs and docs[0].id:
return await self.get_document_by_id(docs[0].id)
return None
async def delete_document(self, document_id: str) -> bool:
"""Delete a document by its ID."""
return await self.document_repository.delete(document_id)
async def list_documents(
self,
limit: int | None = None,
offset: int | None = None,
filter: str | None = None,
include_content: bool = False,
) -> list[Document]:
"""List all documents with optional pagination and filtering.
Args:
limit: Maximum number of documents to return.
offset: Number of documents to skip.
filter: Optional SQL WHERE clause to filter documents.
include_content: Whether to load content and docling_document.
Defaults to False to avoid loading large blobs.
Returns:
List of Document instances matching the criteria.
"""
return await self.document_repository.list_all(
limit=limit, offset=offset, filter=filter, include_content=include_content
)
async def count_documents(self, filter: str | None = None) -> int:
"""Count documents with optional filtering.
Args:
filter: Optional SQL WHERE clause to filter documents.
Returns:
Number of documents matching the criteria.
"""
return await self.document_repository.count(filter=filter)
async def search(
self,
query: str,
limit: int | None = None,
search_type: str = "hybrid",
filter: str | None = None,
) -> list[SearchResult]:
from haiku.rag.client.search import search
return await search(self, query, limit, search_type, filter)
async def expand_context(
self,
search_results: list[SearchResult],
) -> list[SearchResult]:
from haiku.rag.client.search import expand_context
return await expand_context(self, search_results)
async def ask(
self,
question: str,
system_prompt: str | None = None,
filter: str | None = None,
) -> "tuple[str, list[Citation]]":
from haiku.rag.client.agents import ask
return await ask(self, question, system_prompt, filter)
async def research(
self,
question: str,
*,
filter: str | None = None,
max_iterations: int | None = None,
) -> "ResearchReport":
from haiku.rag.client.agents import research
return await research(
self, question, filter=filter, max_iterations=max_iterations
)
async def analyze(
self,
question: str,
documents: list[str] | None = None,
filter: str | None = None,
) -> "AnalysisResult":
from haiku.rag.client.agents import analyze
return await analyze(self, question, documents, filter)
async def visualize_chunk(self, chunk: Chunk) -> list:
from haiku.rag.client.search import visualize_chunk
return await visualize_chunk(self, chunk)
async def rebuild_database(
self, mode: RebuildMode = RebuildMode.FULL
) -> AsyncGenerator[str, None]:
from haiku.rag.client.rebuild import rebuild_database
async for doc_id in rebuild_database(self, mode):
yield doc_id
async def vacuum(self) -> None:
"""Optimize and clean up old versions across all tables."""
await self.store.vacuum()
def close(self):
"""Close the underlying store connection."""
self.store.close()

140
haiku_rag_slim/haiku/rag/client/agents.py Normal file
View file

@ -0,0 +1,140 @@
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from haiku.rag.agents.analysis.models import AnalysisResult
from haiku.rag.agents.research.models import Citation, ResearchReport
from haiku.rag.client import HaikuRAG
async def ask(
client: "HaikuRAG",
question: str,
system_prompt: str | None = None,
filter: str | None = None,
) -> "tuple[str, list[Citation]]":
"""Ask a question using the configured QA agent.
Args:
client: The HaikuRAG client.
question: The question to ask.
system_prompt: Optional custom system prompt for the QA agent.
filter: SQL WHERE clause to filter documents.
Returns:
Tuple of (answer text, list of resolved citations).
"""
from haiku.rag.agents.qa import get_qa_agent
qa_agent = get_qa_agent(client, config=client._config, system_prompt=system_prompt)
return await qa_agent.answer(question, filter=filter)
async def research(
client: "HaikuRAG",
question: str,
*,
filter: str | None = None,
max_iterations: int | None = None,
) -> "ResearchReport":
"""Run multi-agent research to investigate a question.
Args:
client: The HaikuRAG client.
question: The research question to investigate.
filter: SQL WHERE clause to filter documents.
max_iterations: Override max iterations (None uses config default).
Returns:
ResearchReport with structured findings.
"""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
graph = build_research_graph(config=client._config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(
context=context, config=client._config, max_iterations=max_iterations
)
state.search_filter = filter
deps = ResearchDeps(client=client)
return await graph.run(state=state, deps=deps)
async def analyze(
client: "HaikuRAG",
question: str,
documents: list[str] | None = None,
filter: str | None = None,
) -> "AnalysisResult":
"""Answer a question using the analysis agent with code execution.
The analysis agent can write and execute Python code in a sandboxed
environment to solve problems that require computation, aggregation, or
complex traversal across documents.
Args:
client: The HaikuRAG client.
question: The question to answer.
documents: Optional list of document IDs or titles to pre-load.
filter: SQL WHERE clause to filter documents during searches.
Returns:
AnalysisResult with the answer and the final consolidated program.
"""
from haiku.rag.agents.analysis import (
AnalysisContext,
AnalysisDeps,
Sandbox,
create_analysis_agent,
)
from haiku.rag.agents.analysis.models import AnalysisResult
from haiku.rag.agents.research.models import Citation
context = AnalysisContext(filter=filter)
if documents:
loaded_docs = []
for doc_ref in documents:
doc = await client.resolve_document(doc_ref)
if doc:
loaded_docs.append(doc)
context.documents = loaded_docs if loaded_docs else None
sandbox = Sandbox(
db_path=client.store.db_path,
config=client._config,
context=context,
)
deps = AnalysisDeps(
sandbox=sandbox,
context=context,
)
agent = create_analysis_agent(client._config)
result = await agent.run(question, deps=deps)
output = result.output
seen: set[str] = set()
citations: list[Citation] = []
for sr in sandbox._search_results:
if sr.chunk_id and sr.chunk_id not in seen:
seen.add(sr.chunk_id)
citations.append(
Citation(
index=len(seen),
document_id=sr.document_id or "",
chunk_id=sr.chunk_id,
document_uri=sr.document_uri or "",
document_title=sr.document_title,
page_numbers=sr.page_numbers,
headings=sr.headings,
content=sr.content,
)
)
return AnalysisResult(
answer=output.answer,
program=output.program,
citations=citations,
)

497
haiku_rag_slim/haiku/rag/client/documents.py Normal file
View file

@ -0,0 +1,497 @@
import hashlib
import mimetypes
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING
from urllib.parse import urlparse
import httpx
from haiku.rag.client.processing import ensure_chunks_embedded
from haiku.rag.client.titles import resolve_title
from haiku.rag.converters import get_converter
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
from haiku.rag.store.models.document_item import extract_items
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.client import HaikuRAG
async def _store_document_with_chunks(
client: "HaikuRAG",
document: Document,
chunks: list[Chunk],
docling_document: "DoclingDocument",
) -> Document:
"""Store a document with chunks, embedding any that lack embeddings.
Handles versioning/rollback on failure.
"""
# Ensure all chunks have embeddings before storing
chunks = await ensure_chunks_embedded(client._config, chunks)
# Snapshot table versions for versioned rollback (if supported)
versions = await client.store.current_table_versions()
# Create the document
created_doc = await client.document_repository.create(document)
try:
assert created_doc.id is not None, (
"Document ID should not be None after creation"
)
# Set document_id and order for all chunks
for order, chunk in enumerate(chunks):
chunk.document_id = created_doc.id
chunk.order = order
# Batch create all chunks in a single operation
await client.chunk_repository.create(chunks)
# Extract and store document items for context expansion
items = extract_items(created_doc.id, docling_document)
await client.document_item_repository.create_items(created_doc.id, items)
# Vacuum old versions in background (non-blocking) if auto_vacuum enabled
if client._config.storage.auto_vacuum:
client._schedule_vacuum()
return created_doc
except Exception:
# Roll back to the captured versions and re-raise
await client.store.restore_table_versions(versions)
raise
async def _update_document_with_chunks(
client: "HaikuRAG",
document: Document,
chunks: list[Chunk],
docling_document: "DoclingDocument | None" = None,
) -> Document:
"""Update a document and replace its chunks, embedding any that lack embeddings.
Handles versioning/rollback on failure. When `docling_document` is None,
existing items are preserved.
"""
assert document.id is not None, "Document ID is required for update"
chunks = await ensure_chunks_embedded(client._config, chunks)
versions = await client.store.current_table_versions()
# Delete existing chunks before writing new ones
await client.chunk_repository.delete_by_document_id(document.id)
try:
updated_doc = await client.document_repository.update(document)
assert updated_doc.id is not None
for order, chunk in enumerate(chunks):
chunk.document_id = updated_doc.id
chunk.order = order
await client.chunk_repository.create(chunks)
# Replace document items when a new DoclingDocument is provided
if docling_document is not None:
await client.document_item_repository.delete_by_document_id(updated_doc.id)
items = extract_items(updated_doc.id, docling_document)
await client.document_item_repository.create_items(updated_doc.id, items)
if client._config.storage.auto_vacuum:
client._schedule_vacuum()
return updated_doc
except Exception:
await client.store.restore_table_versions(versions)
raise
async def create_document(
client: "HaikuRAG",
content: str,
uri: str | None = None,
title: str | None = None,
metadata: dict | None = None,
format: str = "md",
) -> Document:
"""Create a new document from text content.
Converts the content, chunks it, and generates embeddings.
"""
from haiku.rag.embeddings import embed_chunks
# Convert → Chunk → Embed using primitives
converter = get_converter(client._config)
docling_document = await converter.convert_text(content, format=format)
chunks = await client.chunk(docling_document)
embedded_chunks = await embed_chunks(chunks, client._config)
# Store markdown export as content for better display/readability.
# The original is preserved in docling_document.
stored_content = docling_document.export_to_markdown()
if title is None:
title = await resolve_title(client._config, docling_document, stored_content)
document = Document(
content=stored_content,
uri=uri,
title=title,
metadata=metadata or {},
)
document.set_docling(docling_document)
return await _store_document_with_chunks(
client, document, embedded_chunks, docling_document
)
async def import_document(
client: "HaikuRAG",
docling_document: "DoclingDocument",
chunks: list[Chunk],
uri: str | None = None,
title: str | None = None,
metadata: dict | None = None,
) -> Document:
"""Import a pre-processed document with chunks.
Use this when conversion, chunking, and embedding were done externally.
Chunks without embeddings will be automatically embedded.
"""
content = docling_document.export_to_markdown()
if title is None:
title = await resolve_title(client._config, docling_document, content)
document = Document(
content=content,
uri=uri,
title=title,
metadata=metadata or {},
)
document.set_docling(docling_document)
return await _store_document_with_chunks(client, document, chunks, docling_document)
async def create_document_from_source(
client: "HaikuRAG",
source: str | Path,
title: str | None = None,
metadata: dict | None = None,
) -> Document | list[Document]:
"""Create or update document(s) from a file path, directory, or URL.
Checks if a document with the same URI already exists:
- If MD5 is unchanged, returns existing document
- If MD5 changed, updates the document
- If no document exists, creates a new one
Returns a single Document for files/URLs, a list for directories.
"""
metadata = metadata or {}
source_str = str(source)
parsed_url = urlparse(source_str)
if parsed_url.scheme in ("http", "https"):
return await _create_or_update_document_from_url(
client, source_str, title=title, metadata=metadata
)
elif parsed_url.scheme == "file":
source_path = Path(parsed_url.path)
else:
source_path = Path(source) if isinstance(source, str) else source
if source_path.is_dir():
from haiku.rag.monitor import FileFilter
documents = []
filter = FileFilter(
ignore_patterns=client._config.monitor.ignore_patterns or None,
include_patterns=client._config.monitor.include_patterns or None,
)
for path in source_path.rglob("*"):
if path.is_file() and filter.include_file(str(path)):
doc = await _create_document_from_file(
client, path, title=None, metadata=metadata
)
documents.append(doc)
return documents
return await _create_document_from_file(
client, source_path, title=title, metadata=metadata
)
async def _create_document_from_file(
client: "HaikuRAG",
source_path: Path,
title: str | None = None,
metadata: dict | None = None,
) -> Document:
"""Create or update a document from a single file path."""
from haiku.rag.embeddings import embed_chunks
metadata = metadata or {}
converter = get_converter(client._config)
if source_path.suffix.lower() not in converter.supported_extensions:
raise ValueError(f"Unsupported file extension: {source_path.suffix}")
if not source_path.exists():
raise ValueError(f"File does not exist: {source_path}")
uri = source_path.absolute().as_uri()
md5_hash = hashlib.md5(source_path.read_bytes(), usedforsecurity=False).hexdigest()
content_type, _ = mimetypes.guess_type(str(source_path))
if not content_type:
content_type = "application/octet-stream"
metadata.update({"contentType": content_type, "md5": md5_hash})
# Check if document already exists
existing_doc = await client.get_document_by_uri(uri)
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
# MD5 unchanged; update title/metadata if provided
updated = False
if title is not None and title != existing_doc.title:
existing_doc.title = title
updated = True
merged_metadata = {**(existing_doc.metadata or {}), **metadata}
if merged_metadata != existing_doc.metadata:
existing_doc.metadata = merged_metadata
updated = True
if updated:
return await client.document_repository.update(existing_doc)
return existing_doc
# Convert → Chunk → Embed
docling_document = await client.convert(source_path)
chunks = await client.chunk(docling_document)
embedded_chunks = await embed_chunks(chunks, client._config)
stored_content = docling_document.export_to_markdown()
if existing_doc:
# Update existing document and rechunk
existing_doc.content = stored_content
existing_doc.metadata = metadata
existing_doc.set_docling(docling_document)
if title is not None:
existing_doc.title = title
elif existing_doc.title is None:
existing_doc.title = await resolve_title(
client._config, docling_document, stored_content
)
return await _update_document_with_chunks(
client, existing_doc, embedded_chunks, docling_document
)
else:
if title is None:
title = await resolve_title(
client._config, docling_document, stored_content
)
document = Document(
content=stored_content,
uri=uri,
title=title,
metadata=metadata,
)
document.set_docling(docling_document)
return await _store_document_with_chunks(
client, document, embedded_chunks, docling_document
)
async def _create_or_update_document_from_url(
client: "HaikuRAG",
url: str,
title: str | None = None,
metadata: dict | None = None,
) -> Document:
"""Create or update a document from a URL by downloading and parsing the content."""
from haiku.rag.client.processing import get_extension_from_content_type_or_url
from haiku.rag.embeddings import embed_chunks
metadata = metadata or {}
converter = get_converter(client._config)
supported_extensions = converter.supported_extensions
async with httpx.AsyncClient() as http:
response = await http.get(url)
response.raise_for_status()
md5_hash = hashlib.md5(response.content).hexdigest()
content_type = response.headers.get("content-type", "").lower()
# Check if document already exists
existing_doc = await client.get_document_by_uri(url)
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
updated = False
if title is not None and title != existing_doc.title:
existing_doc.title = title
updated = True
metadata.update({"contentType": content_type, "md5": md5_hash})
merged_metadata = {**(existing_doc.metadata or {}), **metadata}
if merged_metadata != existing_doc.metadata:
existing_doc.metadata = merged_metadata
updated = True
if updated:
return await client.document_repository.update(existing_doc)
return existing_doc
file_extension = get_extension_from_content_type_or_url(url, content_type)
if file_extension not in supported_extensions:
raise ValueError(
f"Unsupported content type/extension: {content_type}/{file_extension}"
)
with tempfile.NamedTemporaryFile(
mode="wb", suffix=file_extension, delete=False
) as temp_file:
temp_file.write(response.content)
temp_file.flush()
temp_path = Path(temp_file.name)
try:
docling_document = await client.convert(temp_path)
chunks = await client.chunk(docling_document)
embedded_chunks = await embed_chunks(chunks, client._config)
finally:
temp_path.unlink(missing_ok=True)
metadata.update({"contentType": content_type, "md5": md5_hash})
stored_content = docling_document.export_to_markdown()
if existing_doc:
existing_doc.content = stored_content
existing_doc.metadata = metadata
existing_doc.set_docling(docling_document)
if title is not None:
existing_doc.title = title
elif existing_doc.title is None:
existing_doc.title = await resolve_title(
client._config, docling_document, stored_content
)
return await _update_document_with_chunks(
client, existing_doc, embedded_chunks, docling_document
)
else:
if title is None:
title = await resolve_title(
client._config, docling_document, stored_content
)
document = Document(
content=stored_content,
uri=url,
title=title,
metadata=metadata,
)
document.set_docling(docling_document)
return await _store_document_with_chunks(
client, document, embedded_chunks, docling_document
)
async def update_document(
client: "HaikuRAG",
document_id: str,
content: str | None = None,
metadata: dict | None = None,
chunks: list[Chunk] | None = None,
title: str | None = None,
docling_document: "DoclingDocument | None" = None,
) -> Document:
"""Update a document by ID.
Updates specified fields. When content or docling_document is provided, the
document is rechunked and re-embedded. Updates to only metadata or title
skip rechunking for efficiency.
Raises:
ValueError: If document not found, or if both content and
docling_document are provided.
"""
from haiku.rag.embeddings import embed_chunks
# Validate: content and docling_document are mutually exclusive
if content is not None and docling_document is not None:
raise ValueError(
"content and docling_document are mutually exclusive. "
"Provide one or the other, not both."
)
existing_doc = await client.get_document_by_id(document_id)
if existing_doc is None:
raise ValueError(f"Document with ID {document_id} not found")
if title is not None:
existing_doc.title = title
if metadata is not None:
existing_doc.metadata = metadata
# Only metadata/title update - no rechunking needed
if content is None and chunks is None and docling_document is None:
return await client.document_repository.update(existing_doc)
# Custom chunks provided - use them as-is
if chunks is not None:
if docling_document is not None:
existing_doc.content = docling_document.export_to_markdown()
existing_doc.set_docling(docling_document)
elif content is not None:
existing_doc.content = content
return await _update_document_with_chunks(
client, existing_doc, chunks, docling_document
)
# DoclingDocument provided without chunks - chunk and embed
if docling_document is not None:
existing_doc.content = docling_document.export_to_markdown()
existing_doc.set_docling(docling_document)
new_chunks = await client.chunk(docling_document)
embedded_chunks = await embed_chunks(new_chunks, client._config)
return await _update_document_with_chunks(
client, existing_doc, embedded_chunks, docling_document
)
# Content provided without chunks - convert, chunk, and embed
assert content is not None
existing_doc.content = content
converter = get_converter(client._config)
converted_docling = await converter.convert_text(existing_doc.content, format="md")
existing_doc.set_docling(converted_docling)
new_chunks = await client.chunk(converted_docling)
embedded_chunks = await embed_chunks(new_chunks, client._config)
return await _update_document_with_chunks(
client, existing_doc, embedded_chunks, converted_docling
)
def check_source_accessible(uri: str) -> bool:
"""Check if a document's source URI is accessible."""
parsed_url = urlparse(uri)
try:
if parsed_url.scheme == "file":
return Path(parsed_url.path).exists()
elif parsed_url.scheme in ("http", "https"):
return True
return False
except Exception:
return False

155
haiku_rag_slim/haiku/rag/client/downloads.py Normal file
View file

@ -0,0 +1,155 @@
import asyncio
import json
from collections.abc import AsyncGenerator
from dataclasses import dataclass
import httpx
from haiku.rag.config import AppConfig
@dataclass
class DownloadProgress:
"""Progress event for model downloads."""
model: str
status: str
completed: int = 0
total: int = 0
digest: str = ""
async def download_models(
config: AppConfig,
) -> AsyncGenerator[DownloadProgress, None]:
"""Download required models per config, yielding progress events.
Yields DownloadProgress events for:
- Docling models
- HuggingFace tokenizer
- Sentence-transformers embedder (if configured)
- HuggingFace reranker models (mxbai, jina-local)
- Ollama models
"""
# Docling models
try:
from docling.utils.model_downloader import download_models
yield DownloadProgress(model="docling", status="start")
await asyncio.to_thread(download_models)
yield DownloadProgress(model="docling", status="done")
except ImportError:
pass
# HuggingFace tokenizer
from transformers import AutoTokenizer
tokenizer_name = config.processing.chunking_tokenizer
yield DownloadProgress(model=tokenizer_name, status="start")
await asyncio.to_thread(AutoTokenizer.from_pretrained, tokenizer_name)
yield DownloadProgress(model=tokenizer_name, status="done")
# Sentence-transformers embedder
if config.embeddings.model.provider == "sentence-transformers": # pragma: no cover
try:
from sentence_transformers import ( # type: ignore[import-not-found] # ty: ignore[unresolved-import]
SentenceTransformer,
)
model_name = config.embeddings.model.name
yield DownloadProgress(model=model_name, status="start")
await asyncio.to_thread(SentenceTransformer, model_name)
yield DownloadProgress(model=model_name, status="done")
except ImportError:
pass
# HuggingFace reranker models
if config.reranking.model: # pragma: no cover
provider = config.reranking.model.provider
model_name = config.reranking.model.name
if provider == "mxbai":
try:
from mxbai_rerank import MxbaiRerankV2
yield DownloadProgress(model=model_name, status="start")
await asyncio.to_thread(
MxbaiRerankV2, model_name, disable_transformers_warnings=True
)
yield DownloadProgress(model=model_name, status="done")
except ImportError:
pass
elif provider == "jina-local":
try:
from transformers import AutoModel
yield DownloadProgress(model=model_name, status="start")
await asyncio.to_thread(
AutoModel.from_pretrained,
model_name,
trust_remote_code=True,
)
yield DownloadProgress(model=model_name, status="done")
except ImportError:
pass
# Collect Ollama models from config
required_models: set[str] = set()
if config.embeddings.model.provider == "ollama":
required_models.add(config.embeddings.model.name)
if config.qa.model.provider == "ollama":
required_models.add(config.qa.model.name)
if config.research.model.provider == "ollama":
required_models.add(config.research.model.name)
if config.reranking.model and config.reranking.model.provider == "ollama":
required_models.add(config.reranking.model.name)
pic_desc = config.processing.conversion_options.picture_description
if pic_desc.enabled and pic_desc.model.provider == "ollama":
required_models.add(pic_desc.model.name)
if (
config.processing.auto_title
and config.processing.title_model.provider == "ollama"
):
required_models.add(config.processing.title_model.name)
if not required_models:
return
base_url = config.providers.ollama.base_url
try:
async with httpx.AsyncClient(timeout=None) as client:
for model in sorted(required_models):
yield DownloadProgress(model=model, status="pulling")
async with client.stream(
"POST", f"{base_url}/api/pull", json={"model": model}
) as r:
async for line in r.aiter_lines():
if not line:
continue
try:
data = json.loads(line)
status = data.get("status", "")
digest = data.get("digest", "")
if digest and "total" in data:
yield DownloadProgress(
model=model,
status="downloading",
total=data.get("total", 0),
completed=data.get("completed", 0),
digest=digest,
)
elif status:
yield DownloadProgress(model=model, status=status)
except json.JSONDecodeError:
pass
yield DownloadProgress(model=model, status="done")
except httpx.ConnectError:
raise ConnectionError(
f"Cannot connect to Ollama at {base_url}. "
"Is Ollama running? Start it with 'ollama serve'."
)

160
haiku_rag_slim/haiku/rag/client/processing.py Normal file
View file

@ -0,0 +1,160 @@
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING
from urllib.parse import urlparse
import httpx
from haiku.rag.config import AppConfig
from haiku.rag.converters import get_converter
from haiku.rag.store.models.chunk import Chunk
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
async def convert(
config: AppConfig, source: Path | str, *, format: str = "md"
) -> "DoclingDocument":
"""Convert a file, URL, or text to DoclingDocument.
Args:
config: Application configuration.
source: One of:
- Path: Local file path to convert
- str (URL): HTTP/HTTPS URL to download and convert
- str (text): Raw text content to convert
format: The format of text content ("md", "html", or "plain").
Defaults to "md". Use "plain" for plain text without parsing.
Only used when source is raw text (not a file path or URL).
Files and URLs determine format from extension/content-type.
Returns:
DoclingDocument from the converted source.
Raises:
ValueError: If the file doesn't exist or has unsupported extension.
httpx.RequestError: If URL download fails.
"""
converter = get_converter(config)
# Path object - convert file directly
if isinstance(source, Path):
if not source.exists():
raise ValueError(f"File does not exist: {source}")
if source.suffix.lower() not in converter.supported_extensions:
raise ValueError(f"Unsupported file extension: {source.suffix}")
return await converter.convert_file(source)
# String - check if URL or text
parsed = urlparse(source)
if parsed.scheme in ("http", "https"):
# URL - download and convert
async with httpx.AsyncClient() as http:
response = await http.get(source)
response.raise_for_status()
content_type = response.headers.get("content-type", "").lower()
file_extension = get_extension_from_content_type_or_url(
source, content_type
)
if file_extension not in converter.supported_extensions:
raise ValueError(
f"Unsupported content type/extension: {content_type}/{file_extension}"
)
with tempfile.NamedTemporaryFile(
mode="wb", suffix=file_extension, delete=False
) as temp_file:
temp_file.write(response.content)
temp_file.flush()
temp_path = Path(temp_file.name)
try:
return await converter.convert_file(temp_path)
finally:
temp_path.unlink(missing_ok=True)
elif parsed.scheme == "file":
# file:// URI
file_path = Path(parsed.path)
if not file_path.exists():
raise ValueError(f"File does not exist: {file_path}")
if file_path.suffix.lower() not in converter.supported_extensions:
raise ValueError(f"Unsupported file extension: {file_path.suffix}")
return await converter.convert_file(file_path)
else:
# Treat as text content
return await converter.convert_text(source, format=format)
async def chunk(config: AppConfig, docling_document: "DoclingDocument") -> list[Chunk]:
"""Chunk a DoclingDocument into Chunks.
Returns chunks without embeddings or document_id. Each chunk's `order`
field is set to its position in the list.
"""
from haiku.rag.chunkers import get_chunker
chunker = get_chunker(config)
return await chunker.chunk(docling_document)
async def ensure_chunks_embedded(config: AppConfig, chunks: list[Chunk]) -> list[Chunk]:
"""Ensure all chunks have embeddings, embedding any that don't.
Chunks that already have embeddings are passed through unchanged; missing
embeddings are filled in in-place in the returned list (preserving order).
"""
from haiku.rag.embeddings import embed_chunks
chunks_to_embed = [c for c in chunks if c.embedding is None]
if not chunks_to_embed:
return chunks
embedded = await embed_chunks(chunks_to_embed, config)
# Build result maintaining original order
embedded_map = {(c.content, c.order): c for c in embedded}
result = []
for ch in chunks:
if ch.embedding is not None:
result.append(ch)
else:
result.append(embedded_map[(ch.content, ch.order)])
return result
def get_extension_from_content_type_or_url(url: str, content_type: str) -> str:
"""Determine file extension from HTTP Content-Type header or URL suffix.
Returns the mapped extension for known content types, falling back to the
URL path suffix, and finally `.html` for generic web content.
"""
content_type_map = {
"text/html": ".html",
"text/plain": ".txt",
"text/markdown": ".md",
"application/pdf": ".pdf",
"application/json": ".json",
"text/csv": ".csv",
"application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",
"application/vnd.openxmlformats-officedocument.presentationml.presentation": ".pptx",
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx",
}
for ct, ext in content_type_map.items():
if ct in content_type:
return ext
parsed_url = urlparse(url)
path = Path(parsed_url.path)
if path.suffix:
return path.suffix.lower()
return ".html"

332
haiku_rag_slim/haiku/rag/client/rebuild.py Normal file
View file

@ -0,0 +1,332 @@
import json
import logging
from collections.abc import AsyncGenerator
from datetime import datetime
from typing import TYPE_CHECKING
from haiku.rag.client.documents import check_source_accessible
from haiku.rag.converters import get_converter
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
from haiku.rag.store.models.document_item import extract_items
from haiku.rag.store.repositories.settings import SettingsRepository
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG, RebuildMode
logger = logging.getLogger(__name__)
_REBUILD_BATCH_SIZE = 50
async def rebuild_database(
client: "HaikuRAG", mode: "RebuildMode | None" = None
) -> AsyncGenerator[str, None]:
"""Rebuild the database with the specified mode.
Yields the ID of each document as it is processed.
"""
from haiku.rag.client import RebuildMode
if mode is None:
mode = RebuildMode.FULL
# Wait for any already-scheduled background vacuum before the destructive
# table operations at the top of RECHUNK / FULL. Rebuild drops and
# recreates tables (and creates indices); a concurrent optimize on the
# same table fails with "CreateIndex transaction was preempted" from
# lance. Note: FULL calls create_document_from_source inside its loop,
# which may schedule *new* background vacuums — those run after the
# destructive phase and are fine.
await client._await_vacuum_tasks()
# Update settings to current config
settings_repo = SettingsRepository(client.store)
await settings_repo.save_current_settings()
documents = await client.list_documents(include_content=True)
if mode == RebuildMode.TITLE_ONLY:
async for doc_id in _rebuild_title_only(client, documents):
yield doc_id
elif mode == RebuildMode.EMBED_ONLY:
async for doc_id in _rebuild_embed_only(client, documents):
yield doc_id
elif mode == RebuildMode.RECHUNK:
await client.chunk_repository.delete_all()
await client.store.recreate_embeddings_table()
async for doc_id in _rebuild_rechunk(client, documents):
yield doc_id
else: # FULL
await client.chunk_repository.delete_all()
await client.store.recreate_embeddings_table()
async for doc_id in _rebuild_full(client, documents):
yield doc_id
# Final maintenance if auto_vacuum enabled. Swallowing only so that a
# failed post-rebuild optimize doesn't mask a successful rebuild — but
# log it so the failure is visible in the output.
if client._config.storage.auto_vacuum:
try:
await client.store.vacuum()
except Exception:
logger.warning("Post-rebuild vacuum failed", exc_info=True)
async def _rebuild_title_only(
client: "HaikuRAG", documents: list[Document]
) -> AsyncGenerator[str, None]:
"""Generate titles for documents that don't have one."""
for doc in documents:
if doc.title is not None:
continue
assert doc.id is not None
try:
title = await client.generate_title(doc)
except Exception:
logger.warning(
"Failed to generate title for document %s", doc.id, exc_info=True
)
continue
if title is not None:
doc.title = title
await client.document_repository.update(doc)
yield doc.id
async def _rebuild_embed_only(
client: "HaikuRAG", documents: list[Document]
) -> AsyncGenerator[str, None]:
"""Re-embed all chunks without changing chunk boundaries."""
from haiku.rag.embeddings import contextualize
# Collect all chunks with new embeddings
all_chunk_data: list[tuple[str, dict]] = []
for doc in documents:
assert doc.id is not None
chunks = await client.chunk_repository.get_by_document_id(doc.id)
if not chunks:
continue
texts = contextualize(chunks)
embeddings = await client.chunk_repository.embedder.embed_documents(texts)
for chunk, content_fts, embedding in zip(chunks, texts, embeddings):
all_chunk_data.append(
(
doc.id,
{
"id": chunk.id,
"document_id": chunk.document_id,
"content": chunk.content,
"content_fts": content_fts,
"metadata": json.dumps(chunk.metadata),
"order": chunk.order,
"vector": embedding,
},
)
)
# Recreate chunks table (handles dimension changes)
await client.store.recreate_embeddings_table()
# Insert all chunks
if all_chunk_data:
records = [client.store.ChunkRecord(**data) for _, data in all_chunk_data]
await client.store.chunks_table.add(records)
# Yield all processed doc IDs
yielded_docs: set[str] = set()
for doc_id, _ in all_chunk_data:
if doc_id not in yielded_docs:
yielded_docs.add(doc_id)
yield doc_id
# Yield docs with no chunks
for doc in documents:
if doc.id and doc.id not in yielded_docs:
yield doc.id
async def _flush_rebuild_batch(
client: "HaikuRAG", documents: list[Document], chunks: list[Chunk]
) -> None:
"""Batch write documents and chunks during rebuild.
Performs two writes: one for all document updates (via merge_insert), one
for all chunks. Also repopulates document items from the stored docling
document. Used by RECHUNK and FULL modes after the chunks table has been
cleared.
"""
from haiku.rag.store.engine import DocumentRecord
if not documents:
return
now = datetime.now().isoformat()
# Batch update documents using merge_insert (single LanceDB version)
doc_records = []
for doc in documents:
assert doc.id is not None
doc_records.append(
DocumentRecord(
id=doc.id,
content=doc.content,
uri=doc.uri,
title=doc.title,
metadata=json.dumps(doc.metadata),
docling_document=doc.docling_document,
docling_pages=doc.docling_pages,
docling_version=doc.docling_version,
created_at=doc.created_at.isoformat() if doc.created_at else now,
updated_at=now,
)
)
await (
client.store.documents_table.merge_insert("id")
.when_matched_update_all()
.execute(doc_records)
)
# Batch create all chunks (single LanceDB version)
if chunks:
await client.chunk_repository.create(chunks)
# Repopulate document items from stored docling data
for doc in documents:
assert doc.id is not None
docling_doc = doc.get_docling_document()
if docling_doc is not None:
await client.document_item_repository.delete_by_document_id(doc.id)
items = extract_items(doc.id, docling_doc)
await client.document_item_repository.create_items(doc.id, items)
async def _rebuild_rechunk(
client: "HaikuRAG", documents: list[Document]
) -> AsyncGenerator[str, None]:
"""Re-chunk and re-embed from existing document content."""
from haiku.rag.embeddings import embed_chunks
pending_chunks: list[Chunk] = []
pending_docs: list[Document] = []
pending_doc_ids: list[str] = []
converter = get_converter(client._config)
for doc in documents:
assert doc.id is not None
# Convert stored markdown to DoclingDocument
docling_document = await converter.convert_text(doc.content, format="md")
# Chunk and embed
chunks = await client.chunk(docling_document)
embedded_chunks = await embed_chunks(chunks, client._config)
# Update document fields
doc.set_docling(docling_document)
# Prepare chunks with document_id and order
for order, chunk in enumerate(embedded_chunks):
chunk.document_id = doc.id
chunk.order = order
pending_chunks.extend(embedded_chunks)
pending_docs.append(doc)
pending_doc_ids.append(doc.id)
# Flush batch when size reached
if len(pending_docs) >= _REBUILD_BATCH_SIZE:
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
for doc_id in pending_doc_ids:
yield doc_id
pending_chunks = []
pending_docs = []
pending_doc_ids = []
# Flush remaining
if pending_docs:
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
for doc_id in pending_doc_ids:
yield doc_id
async def _rebuild_full(
client: "HaikuRAG", documents: list[Document]
) -> AsyncGenerator[str, None]:
"""Full rebuild: re-convert from source, re-chunk, re-embed."""
from haiku.rag.embeddings import embed_chunks
pending_chunks: list[Chunk] = []
pending_docs: list[Document] = []
pending_doc_ids: list[str] = []
converter = get_converter(client._config)
for doc in documents:
assert doc.id is not None
# Try to rebuild from source if available
if doc.uri and check_source_accessible(doc.uri):
try:
# Flush pending batch before source rebuild (creates new doc)
if pending_docs:
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
for doc_id in pending_doc_ids:
yield doc_id
pending_chunks = []
pending_docs = []
pending_doc_ids = []
await client.delete_document(doc.id)
new_doc = await client.create_document_from_source(
source=doc.uri, metadata=doc.metadata or {}
)
assert isinstance(new_doc, Document)
assert new_doc.id is not None
yield new_doc.id
continue
except Exception as e:
logger.error(
"Error recreating document from source %s: %s",
doc.uri,
e,
)
continue
# Fallback: rebuild from stored content
if doc.uri:
logger.warning("Source missing for %s, re-embedding from content", doc.uri)
docling_document = await converter.convert_text(doc.content, format="md")
chunks = await client.chunk(docling_document)
embedded_chunks = await embed_chunks(chunks, client._config)
doc.set_docling(docling_document)
# Prepare chunks with document_id and order
for order, chunk in enumerate(embedded_chunks):
chunk.document_id = doc.id
chunk.order = order
pending_chunks.extend(embedded_chunks)
pending_docs.append(doc)
pending_doc_ids.append(doc.id)
# Flush batch when size reached
if len(pending_docs) >= _REBUILD_BATCH_SIZE:
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
for doc_id in pending_doc_ids:
yield doc_id
pending_chunks = []
pending_docs = []
pending_doc_ids = []
# Flush remaining
if pending_docs:
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
for doc_id in pending_doc_ids:
yield doc_id

196
haiku_rag_slim/haiku/rag/client/search.py Normal file
View file

@ -0,0 +1,196 @@
from typing import TYPE_CHECKING
from haiku.rag.reranking import get_reranker
from haiku.rag.store.models.chunk import Chunk, SearchResult
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
async def search(
client: "HaikuRAG",
query: str,
limit: int | None = None,
search_type: str = "hybrid",
filter: str | None = None,
) -> list[SearchResult]:
"""Search for relevant chunks with optional reranking.
Args:
client: The HaikuRAG client (provides config + chunk repository).
query: The search query string.
limit: Maximum number of results to return. Defaults to config.search.limit.
search_type: Type of search - "vector", "fts", or "hybrid" (default).
filter: Optional SQL WHERE clause to filter documents before searching chunks.
Returns:
List of SearchResult objects ordered by relevance.
"""
if limit is None:
limit = client._config.search.limit
reranker = get_reranker(config=client._config)
if reranker is None:
chunk_results = await client.chunk_repository.search(
query, limit, search_type, filter
)
else:
search_limit = limit * 10
raw_results = await client.chunk_repository.search(
query, search_limit, search_type, filter
)
chunks = [chunk for chunk, _ in raw_results]
chunk_results = await reranker.rerank(query, chunks, top_n=limit)
return [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results]
async def expand_context(
client: "HaikuRAG",
search_results: list[SearchResult],
) -> list[SearchResult]:
"""Expand search results with surrounding content from the document.
Uses the document_items table for section-bounded expansion.
See haiku.rag.context for the algorithm description.
Results without doc_item_refs pass through unexpanded. This happens when
chunks were created without docling metadata (e.g., custom chunks passed
to import_document).
"""
from haiku.rag.context import expand_with_items
max_chars = client._config.search.max_context_chars
# Group by document_id for efficient processing
document_groups: dict[str | None, list[SearchResult]] = {}
for result in search_results:
doc_id = result.document_id
if doc_id not in document_groups:
document_groups[doc_id] = []
document_groups[doc_id].append(result)
expanded_results = []
for doc_id, doc_results in document_groups.items():
if doc_id is None:
expanded_results.extend(doc_results)
continue
has_refs = any(r.doc_item_refs for r in doc_results)
if not has_refs:
expanded_results.extend(doc_results)
continue
expanded = await expand_with_items(
client.document_item_repository,
doc_id,
doc_results,
max_chars,
)
expanded_results.extend(expanded)
expanded_results.sort(key=lambda r: r.score, reverse=True)
return expanded_results
async def visualize_chunk(client: "HaikuRAG", chunk: Chunk) -> list:
"""Render page images with bounding box highlights for a chunk.
Expands the chunk's context to find the full section, then resolves
bounding boxes from all items in the expanded range. This ensures
visualization covers all pages the expanded content spans.
Returns a list of PIL Image objects, one per page with bounding boxes.
Empty list if no bounding boxes or page images available.
"""
from copy import deepcopy
from PIL import ImageDraw
from haiku.rag.store.models.chunk import ChunkMetadata
if not chunk.document_id:
return []
doc = await client.document_repository.get_docling_data(chunk.document_id)
if not doc:
return []
docling_doc = doc.get_docling_document()
if not docling_doc:
return []
# Expand context to get all doc_item_refs in the section
chunk_meta = chunk.get_chunk_metadata()
if chunk_meta.doc_item_refs:
search_result = SearchResult(
content=chunk.content,
score=1.0,
chunk_id=chunk.id,
document_id=chunk.document_id,
doc_item_refs=chunk_meta.doc_item_refs,
page_numbers=chunk_meta.page_numbers,
)
expanded = await expand_context(client, [search_result])
refs = expanded[0].doc_item_refs if expanded else chunk_meta.doc_item_refs
meta = ChunkMetadata(doc_item_refs=refs)
else:
meta = chunk_meta
bounding_boxes = meta.resolve_bounding_boxes(docling_doc)
if not bounding_boxes:
return []
# Group bounding boxes by page
boxes_by_page: dict[int, list] = {}
for bbox in bounding_boxes:
if bbox.page_no not in boxes_by_page:
boxes_by_page[bbox.page_no] = []
boxes_by_page[bbox.page_no].append(bbox)
# Load only the needed page images
pages_doc = await client.document_repository.get_pages_data(chunk.document_id)
if not pages_doc:
return []
page_images = pages_doc.get_page_images(list(boxes_by_page.keys()))
# Render each page with its bounding boxes
images = []
for page_no in sorted(boxes_by_page.keys()):
if page_no not in page_images:
continue
page = page_images[page_no]
if page.image is None or page.image.pil_image is None:
continue
pil_image = page.image.pil_image
page_height = page.size.height
# Scale factor: image pixels vs document coordinates
scale_x = pil_image.width / page.size.width
scale_y = pil_image.height / page.size.height
image = deepcopy(pil_image)
draw = ImageDraw.Draw(image, "RGBA")
for bbox in boxes_by_page[page_no]:
# Document coords are bottom-left origin; PIL uses top-left
x0 = bbox.left * scale_x
y0 = (page_height - bbox.top) * scale_y
x1 = bbox.right * scale_x
y1 = (page_height - bbox.bottom) * scale_y
if y0 > y1:
y0, y1 = y1, y0
fill_color = (255, 255, 0, 40) # Yellow with transparency
outline_color = (255, 165, 0, 100) # Orange outline
draw.rectangle([(x0, y0), (x1, y1)], fill=fill_color, outline=None)
draw.rectangle([(x0, y0), (x1, y1)], outline=outline_color, width=1)
images.append(image)
return images

105
haiku_rag_slim/haiku/rag/client/titles.py Normal file
View file

@ -0,0 +1,105 @@
import logging
from typing import TYPE_CHECKING
from haiku.rag.config import AppConfig
from haiku.rag.store.models.document import Document
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
logger = logging.getLogger(__name__)
def extract_structural_title(docling_document: "DoclingDocument") -> str | None:
"""Extract a title from DoclingDocument structural metadata.
Priority: FURNITURE TITLE > BODY TITLE > first SECTION_HEADER.
"""
from docling_core.types.doc.document import ContentLayer
from docling_core.types.doc.labels import DocItemLabel
furniture_title = None
body_title = None
first_section_header = None
for item in docling_document.texts:
if item.label == DocItemLabel.TITLE:
text = item.text.strip()
if not text:
continue
if item.content_layer == ContentLayer.FURNITURE:
furniture_title = text
elif body_title is None:
body_title = text
elif item.label == DocItemLabel.SECTION_HEADER and first_section_header is None:
text = item.text.strip()
if text:
first_section_header = text
return furniture_title or body_title or first_section_header
async def generate_title_with_llm(config: AppConfig, content: str) -> str | None:
"""Generate a title using LLM from document content."""
from pydantic_ai import Agent
from haiku.rag.utils import get_model
truncated = content[:2000]
model = get_model(config.processing.title_model, config)
agent: Agent[None, str] = Agent(
model=model,
output_type=str,
instructions=(
"Generate a concise, descriptive title for the following document. "
"The title should be at most 10 words. "
"Return ONLY the title text, nothing else."
),
)
result = await agent.run(truncated)
title = result.output.strip()
return title if title else None
async def resolve_title(
config: AppConfig,
docling_document: "DoclingDocument",
content: str,
) -> str | None:
"""Auto-generate a title from document structure or LLM.
Returns None if auto_title is disabled or generation fails.
"""
if not config.processing.auto_title:
return None
structural = extract_structural_title(docling_document)
if structural:
return structural
try:
return await generate_title_with_llm(config, content)
except Exception:
logger.warning("LLM title generation failed during ingestion", exc_info=True)
return None
async def generate_title(config: AppConfig, document: Document) -> str | None:
"""Generate a title for a document.
Attempts structural extraction from the stored DoclingDocument, then falls
back to LLM generation. Bypasses the auto_title config since this is an
explicit call.
Does NOT update the document — caller decides.
"""
docling_doc = document.get_docling_document()
content = document.content or ""
if docling_doc is not None:
structural = extract_structural_title(docling_doc)
if structural:
return structural
return await generate_title_with_llm(config, content)

10
haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py
View file

@ -84,8 +84,8 @@ class InfoModal(ModalScreen):
# Connect to get table info
config = self.client.store._config
try:
db = connect_lancedb(config, self.db_path)
stats = get_database_stats(db)
db = await connect_lancedb(config, self.db_path)
stats = await get_database_stats(db)
except Exception as e:
lines.append(f"[red]Failed to open database: {e}[/red]")
self._content_widget.update("\n".join(lines))
@ -101,8 +101,10 @@ class InfoModal(ModalScreen):
vector_dim: int | None = None
if stats["settings"]["exists"]:
settings_tbl = db.open_table("settings")
arrow = settings_tbl.search().where("id = 'settings'").limit(1).to_arrow()
settings_tbl = await db.open_table("settings")
arrow = await (
settings_tbl.query().where("id = 'settings'").limit(1).to_arrow()
)
rows = arrow.to_pylist() if arrow is not None else []
if rows:
raw = rows[0].get("settings") or "{}"

313
haiku_rag_slim/haiku/rag/store/engine.py
View file

@ -5,11 +5,12 @@ from datetime import datetime, timedelta
from enum import Enum
from importlib import metadata
from pathlib import Path
from typing import Any
from typing import TYPE_CHECKING, Any, cast
from uuid import uuid4
import lancedb
import pyarrow as pa
from lancedb.index import FTS, BTree, IvfPq
from lancedb.pydantic import LanceModel, Vector
from pydantic import Field
@ -17,9 +18,25 @@ from haiku.rag.config import AppConfig, Config
from haiku.rag.embeddings import get_embedder
from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError
if TYPE_CHECKING:
from lancedb.query import AsyncQueryBase
logger = logging.getLogger(__name__)
async def query_to_pydantic[T: LanceModel](
query: "AsyncQueryBase", model: type[T]
) -> list[T]:
"""Typed wrapper around AsyncQueryBase.to_pydantic.
The upstream stub annotates `.to_pydantic()` as returning `list[LanceModel]`
regardless of the concrete model passed in. This helper narrows the return
type to the concrete model so attribute access on the results type-checks
at call sites without needing per-line cast / ignore comments.
"""
return cast("list[T]", await query.to_pydantic(model))
class ConnectionMode(Enum):
LOCAL = "local"
CLOUD = "cloud"
@ -35,10 +52,12 @@ class ConnectionMode(Enum):
return ConnectionMode.OBJECT_STORAGE
def connect_lancedb(config: AppConfig, db_path: Path | None = None):
async def connect_lancedb(
config: AppConfig, db_path: Path | None = None
) -> lancedb.AsyncConnection:
mode = ConnectionMode.from_config(config)
if mode == ConnectionMode.CLOUD:
return lancedb.connect(
return await lancedb.connect_async(
uri=config.lancedb.uri,
api_key=config.lancedb.api_key,
region=config.lancedb.region,
@ -47,11 +66,11 @@ def connect_lancedb(config: AppConfig, db_path: Path | None = None):
kwargs: dict[str, Any] = {"uri": config.lancedb.uri}
if config.lancedb.storage_options:
kwargs["storage_options"] = config.lancedb.storage_options
return lancedb.connect(**kwargs)
return await lancedb.connect_async(**kwargs)
else:
if db_path is None:
raise ValueError("No lancedb.uri configured and no db_path provided")
return lancedb.connect(db_path)
return await lancedb.connect_async(db_path)
class DocumentRecord(LanceModel):
@ -89,19 +108,26 @@ def get_documents_arrow_schema() -> pa.Schema:
return pa.schema(fields)
def create_chunk_model(vector_dim: int):
"""Create a ChunkRecord model with the specified vector dimension.
This creates a model with proper vector typing for LanceDB.
class ChunkRecordBase(LanceModel):
"""Static base for ChunkRecord — declares the fields so attribute access
and constructor calls type-check. The concrete `vector` field is overridden
by create_chunk_model() with a Vector(dim) whose fixed-size-list dimension
is only known at runtime.
"""
class ChunkRecord(LanceModel):
id: str = Field(default_factory=lambda: str(uuid4()))
document_id: str
content: str
content_fts: str = Field(default="")
metadata: str = Field(default="{}")
order: int = Field(default=0)
id: str = Field(default_factory=lambda: str(uuid4()))
document_id: str
content: str
content_fts: str = Field(default="")
metadata: str = Field(default="{}")
order: int = Field(default=0)
vector: list[float] = Field(default_factory=list)
def create_chunk_model(vector_dim: int) -> type[ChunkRecordBase]:
"""Create a ChunkRecord model with the specified vector dimension."""
class ChunkRecord(ChunkRecordBase):
vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim) # type: ignore
return ChunkRecord
@ -124,7 +150,7 @@ class SettingsRecord(LanceModel):
REQUIRED_TABLES: tuple[str, ...] = ("documents", "chunks", "document_items", "settings")
def get_database_stats(db: lancedb.DBConnection) -> dict:
async def get_database_stats(db: lancedb.AsyncConnection) -> dict:
"""Collect stats for every haiku.rag table on the connection.
Missing tables return ``{"exists": False}``. Present tables include
@ -132,7 +158,7 @@ def get_database_stats(db: lancedb.DBConnection) -> dict:
entry additionally reports vector index status and, when an index
exists, ``num_indexed_rows`` and ``num_unindexed_rows``.
"""
existing = set(db.list_tables().tables)
existing = set((await db.list_tables()).tables)
stats: dict = {}
tables: dict = {}
@ -140,24 +166,24 @@ def get_database_stats(db: lancedb.DBConnection) -> dict:
if name not in existing:
stats[name] = {"exists": False}
continue
tbl = db.open_table(name)
tbl = await db.open_table(name)
tables[name] = tbl
# lancedb's .stats() stub claims TableStatistics but returns a plain dict at runtime.
tbl_stats: dict = tbl.stats() # ty: ignore[invalid-assignment]
tbl_stats: dict = await tbl.stats() # type: ignore[assignment] # ty: ignore[invalid-assignment]
stats[name] = {
"exists": True,
"num_rows": tbl_stats.get("num_rows", 0),
"total_bytes": tbl_stats.get("total_bytes", 0),
"num_versions": len(list(tbl.list_versions())),
"num_versions": len(await tbl.list_versions()),
}
if stats["chunks"]["exists"]:
chunks_tbl = tables["chunks"]
indices = chunks_tbl.list_indices()
indices = await chunks_tbl.list_indices()
has_vector_index = any("vector" in str(idx).lower() for idx in indices)
stats["chunks"]["has_vector_index"] = has_vector_index
if has_vector_index:
index_stats = chunks_tbl.index_stats("vector_idx")
index_stats = await chunks_tbl.index_stats("vector_idx")
if index_stats is not None:
stats["chunks"]["num_indexed_rows"] = index_stats.num_indexed_rows
stats["chunks"]["num_unindexed_rows"] = index_stats.num_unindexed_rows
@ -181,10 +207,13 @@ class Store:
self._before = before
# Time-travel mode is always read-only
self._read_only = read_only or (before is not None)
self._create = create
self._skip_validation = skip_validation
self._skip_migration_check = skip_migration_check
self._vacuum_lock = asyncio.Lock()
self._is_new_db = False
# Check if database exists (for local filesystem only)
is_new_db = False
if self._connection_mode == ConnectionMode.LOCAL:
if not db_path.exists():
if not create:
@ -192,17 +221,27 @@ class Store:
f"Database does not exist at {self.db_path.absolute()}. "
"Use 'haiku-rag init' to create a new database."
)
is_new_db = True
self._is_new_db = True
# Ensure parent directories exist for new databases
if not db_path.parent.exists():
Path.mkdir(db_path.parent, parents=True)
# Connect to LanceDB
self.db = connect_lancedb(self._config, db_path)
# Create embedder (sync — no LanceDB needed)
self.embedder = get_embedder(config=self._config)
# For remote stores, detect new DB by checking if tables exist
if not is_new_db and self._connection_mode != ConnectionMode.LOCAL:
existing_tables = self.db.list_tables().tables
async def _initialize(self):
"""Perform async initialization: connect to LanceDB, init tables, validate."""
# Connect to LanceDB
self.db: lancedb.AsyncConnection = await connect_lancedb(
self._config, self.db_path
)
# For remote stores (and as a safety net for local paths that exist but
# have no tables — e.g. a previously failed init), detect new DB by
# checking whether any tables exist.
is_new_db = self._is_new_db
if not is_new_db:
existing_tables = (await self.db.list_tables()).tables
if not existing_tables:
is_new_db = True
@ -210,57 +249,68 @@ class Store:
# that can read existing chunks. For new databases, use config's dimension.
stored_vector_dim = None
if not is_new_db:
stored_vector_dim = self._get_stored_vector_dim()
# Create embedder with config's dimension (for generating new embeddings)
self.embedder = get_embedder(config=self._config)
stored_vector_dim = await self._get_stored_vector_dim()
# Create ChunkRecord with stored dimension (for reading) or config dimension (for new DB)
chunk_vector_dim = stored_vector_dim or self.embedder._vector_dim
self.ChunkRecord = create_chunk_model(chunk_vector_dim)
self.ChunkRecord: type[ChunkRecordBase] = create_chunk_model(chunk_vector_dim)
# Initialize tables (creates them if they don't exist)
self._init_tables()
await self._init_tables()
# Checkout tables to historical state if before is specified
if before is not None:
self._checkout_tables_before(before)
if self._before is not None:
await self._checkout_tables_before(self._before)
# Set version for new databases, check migrations for existing ones
if is_new_db:
if not self._read_only:
self._set_initial_version()
elif not skip_migration_check:
self._check_migrations()
await self._set_initial_version()
elif not self._skip_migration_check:
await self._check_migrations()
# Validate config compatibility after connection is established
if not skip_validation:
self._validate_configuration()
if not self._skip_validation:
await self._validate_configuration()
async def __aenter__(self):
# If _initialize connects to LanceDB but then fails (e.g. migration
# check, config validation), close the connection so it doesn't
# leak — __aexit__ won't run because the `async with` never entered.
try:
await self._initialize()
except BaseException:
self.close()
raise
return self
async def __aexit__(self, exc_type, exc_val, exc_tb): # noqa: ARG002
self.close()
return False
@property
def is_read_only(self) -> bool:
"""Whether the store is in read-only mode."""
return self._read_only
def _get_stored_vector_dim(self) -> int | None:
async def _get_stored_vector_dim(self) -> int | None:
"""Read the stored vector dimension from the settings table.
Returns:
The stored vector dimension, or None if not found.
"""
try:
existing_tables = self.db.list_tables().tables
existing_tables = (await self.db.list_tables()).tables
if "settings" not in existing_tables:
return None
settings_table = self.db.open_table("settings")
settings_table = await self.db.open_table("settings")
rows = (
settings_table.search()
await settings_table.query()
.where("id = 'settings'")
.limit(1)
.to_arrow()
.to_pylist()
)
).to_pylist()
if not rows or not rows[0].get("settings"):
return None
@ -312,7 +362,7 @@ class Store:
self.document_items_table,
self.settings_table,
]:
table.optimize(cleanup_older_than=retention)
await table.optimize(cleanup_older_than=retention)
except (RuntimeError, OSError) as e:
# Handle resource errors gracefully
logger.debug(f"Vacuum skipped due to resource constraints: {e}")
@ -321,7 +371,7 @@ class Store:
def _connection_mode(self) -> ConnectionMode:
return ConnectionMode.from_config(self._config)
def _ensure_vector_index(self) -> None:
async def _ensure_vector_index(self) -> None:
"""Create or rebuild vector index on chunks table.
Cloud deployments auto-create indexes, so we skip for those.
@ -334,7 +384,7 @@ class Store:
try:
# Check if table has enough data (indexes require training data)
row_count = self.chunks_table.count_rows()
row_count = await self.chunks_table.count_rows()
if row_count < 256:
logger.debug(
f"Skipping vector index creation: need at least 256 rows, have {row_count}"
@ -343,30 +393,34 @@ class Store:
# Create or replace index (replace=True is the default)
logger.info("Creating vector index on chunks table...")
self.chunks_table.create_index(
metric=self._config.search.vector_index_metric,
index_type="IVF_PQ",
replace=True, # Explicit: replace existing index
await self.chunks_table.create_index(
"vector",
config=IvfPq(
distance_type=self._config.search.vector_index_metric,
),
replace=True,
)
# Wait for index creation to complete
# Index name is column_name + "_idx"
self.chunks_table.wait_for_index(["vector_idx"], timeout=timedelta(hours=1))
await self.chunks_table.wait_for_index(
["vector_idx"], timeout=timedelta(hours=1)
)
logger.info("Vector index created successfully")
except Exception as e:
logger.warning(f"Could not create vector index: {e}")
def _validate_configuration(self) -> None:
async def _validate_configuration(self) -> None:
"""Validate that the configuration is compatible with the database."""
from haiku.rag.store.repositories.settings import SettingsRepository
settings_repo = SettingsRepository(self)
settings_repo.validate_config_compatibility()
await settings_repo.validate_config_compatibility()
def _init_tables(self):
async def _init_tables(self):
"""Initialize database tables (create if they don't exist)."""
existing_tables = self.db.list_tables().tables
existing_tables = (await self.db.list_tables()).tables
missing_tables = set(REQUIRED_TABLES) - set(existing_tables)
if missing_tables and self._read_only:
@ -377,57 +431,61 @@ class Store:
# Create or open documents table
if "documents" in existing_tables:
self.documents_table = self.db.open_table("documents")
self.documents_table = await self.db.open_table("documents")
else:
self.documents_table = self.db.create_table(
self.documents_table = await self.db.create_table(
"documents", schema=get_documents_arrow_schema()
)
# Create or open chunks table
if "chunks" in existing_tables:
self.chunks_table = self.db.open_table("chunks")
self.chunks_table = await self.db.open_table("chunks")
else:
self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord)
self.chunks_table = await self.db.create_table(
"chunks", schema=self.ChunkRecord
)
# Create FTS index on content_fts (contextualized content) for better search
self.chunks_table.create_fts_index(
"content_fts", replace=True, with_position=True, remove_stop_words=False
await self.chunks_table.create_index(
"content_fts",
config=FTS(with_position=True, remove_stop_words=False),
replace=True,
)
# Create or open document_items table
if "document_items" in existing_tables:
self.document_items_table = self.db.open_table("document_items")
self.document_items_table = await self.db.open_table("document_items")
else:
self.document_items_table = self.db.create_table(
self.document_items_table = await self.db.create_table(
"document_items", schema=DocumentItemRecord
)
self.document_items_table.create_scalar_index(
"document_id", index_type="BTREE", replace=True
await self.document_items_table.create_index(
"document_id", config=BTree(), replace=True
)
self.document_items_table.create_scalar_index(
"position", index_type="BTREE", replace=True
await self.document_items_table.create_index(
"position", config=BTree(), replace=True
)
self.document_items_table.create_scalar_index(
"self_ref", index_type="BTREE", replace=True
await self.document_items_table.create_index(
"self_ref", config=BTree(), replace=True
)
# Create or open settings table
if "settings" in existing_tables:
self.settings_table = self.db.open_table("settings")
self.settings_table = await self.db.open_table("settings")
else:
self.settings_table = self.db.create_table(
self.settings_table = await self.db.create_table(
"settings", schema=SettingsRecord
)
# Save current settings to the new database
settings_data = self._config.model_dump(mode="json")
self.settings_table.add(
await self.settings_table.add(
[SettingsRecord(id="settings", settings=json.dumps(settings_data))]
)
def _set_initial_version(self):
async def _set_initial_version(self):
"""Set the initial version for a new database."""
self.set_haiku_version(metadata.version("haiku.rag-slim"))
await self.set_haiku_version(metadata.version("haiku.rag-slim"))
def _check_migrations(self) -> None:
async def _check_migrations(self) -> None:
"""Check if migrations are pending and error or update version accordingly.
Raises:
@ -436,7 +494,7 @@ class Store:
from haiku.rag.store.upgrades import get_pending_upgrades
current_version = metadata.version("haiku.rag-slim")
db_version = self.get_haiku_version()
db_version = await self.get_haiku_version()
pending = get_pending_upgrades(db_version)
@ -450,9 +508,9 @@ class Store:
# No pending migrations - update version silently if needed (writable only)
if not self._read_only and db_version != current_version:
self.set_haiku_version(current_version)
await self.set_haiku_version(current_version)
def migrate(self) -> list[str]:
async def migrate(self) -> list[str]:
"""Run pending database migrations.
Returns:
@ -465,21 +523,21 @@ class Store:
from haiku.rag.store.upgrades import run_pending_upgrades
db_version = self.get_haiku_version()
db_version = await self.get_haiku_version()
current_version = metadata.version("haiku.rag-slim")
applied = run_pending_upgrades(self, db_version)
applied = await run_pending_upgrades(self, db_version)
# Update version after successful migration
if applied or db_version != current_version:
self.set_haiku_version(current_version)
await self.set_haiku_version(current_version)
return applied
def get_haiku_version(self) -> str:
async def get_haiku_version(self) -> str:
"""Returns the user version stored in settings."""
settings_records = list(
self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
settings_records = await query_to_pydantic(
self.settings_table.query().limit(1), SettingsRecord
)
if settings_records:
settings = (
@ -490,15 +548,15 @@ class Store:
return settings.get("version", "0.0.0")
return "0.0.0"
def set_haiku_version(self, version: str) -> None:
async def set_haiku_version(self, version: str) -> None:
"""Updates the user version in settings.
Raises:
ReadOnlyError: If the store is in read-only mode.
"""
self._assert_writable()
settings_records = list(
self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
settings_records = await query_to_pydantic(
self.settings_table.query().limit(1), SettingsRecord
)
if settings_records:
# Only write if version actually changes to avoid creating new table versions
@ -509,73 +567,74 @@ class Store:
)
if current.get("version") != version:
current["version"] = version
self.settings_table.update(
await self.settings_table.update(
{"settings": json.dumps(current)},
where="id = 'settings'",
values={"settings": json.dumps(current)},
)
else:
# Create new settings record
settings_data = Config.model_dump(mode="json")
settings_data["version"] = version
self.settings_table.add(
await self.settings_table.add(
[SettingsRecord(id="settings", settings=json.dumps(settings_data))]
)
def recreate_embeddings_table(self) -> None:
async def recreate_embeddings_table(self) -> None:
"""Recreate the chunks table with current vector dimensions.
Raises:
ReadOnlyError: If the store is in read-only mode.
"""
self._assert_writable()
# Drop and recreate chunks table
try:
self.db.drop_table("chunks")
except Exception:
pass
# Drop and recreate chunks table. Check existence first rather than
# catching-and-swallowing drop_table's errors — a catch-all would
# hide real failures (permissions, storage-backend errors) and then
# the subsequent create_table would fail confusingly.
if "chunks" in (await self.db.list_tables()).tables:
await self.db.drop_table("chunks")
# Update the ChunkRecord model with new vector dimension
self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord)
self.chunks_table = await self.db.create_table(
"chunks", schema=self.ChunkRecord
)
# Create FTS index on content_fts (contextualized content) for better search
self.chunks_table.create_fts_index(
"content_fts", replace=True, with_position=True, remove_stop_words=False
await self.chunks_table.create_index(
"content_fts",
config=FTS(with_position=True, remove_stop_words=False),
replace=True,
)
def close(self):
"""Close the database connection."""
# LanceDB connections are automatically managed
pass
# AsyncConnection.close() is synchronous
if hasattr(self, "db"):
self.db.close()
def current_table_versions(self) -> dict[str, int]:
async def current_table_versions(self) -> dict[str, int]:
"""Capture current versions of key tables for rollback using LanceDB's API."""
return {
"documents": int(self.documents_table.version),
"chunks": int(self.chunks_table.version),
"document_items": int(self.document_items_table.version),
"settings": int(self.settings_table.version),
"documents": await self.documents_table.version(),
"chunks": await self.chunks_table.version(),
"document_items": await self.document_items_table.version(),
"settings": await self.settings_table.version(),
}
def restore_table_versions(self, versions: dict[str, int]) -> bool:
async def restore_table_versions(self, versions: dict[str, int]) -> bool:
"""Restore tables to the provided versions using LanceDB's API.
Raises:
ReadOnlyError: If the store is in read-only mode.
"""
self._assert_writable()
self.documents_table.restore(int(versions["documents"]))
self.chunks_table.restore(int(versions["chunks"]))
self.document_items_table.restore(int(versions["document_items"]))
self.settings_table.restore(int(versions["settings"]))
await self.documents_table.restore(int(versions["documents"]))
await self.chunks_table.restore(int(versions["chunks"]))
await self.document_items_table.restore(int(versions["document_items"]))
await self.settings_table.restore(int(versions["settings"]))
return True
@property
def _connection(self):
"""Compatibility property for repositories expecting _connection."""
return self
def _checkout_tables_before(self, before: datetime) -> None:
async def _checkout_tables_before(self, before: datetime) -> None:
"""Checkout all tables to their state at or before the given datetime.
Args:
@ -601,7 +660,7 @@ class Store:
]
for table_name, table in tables:
versions = table.list_versions()
versions = await table.list_versions()
# Find the latest version at or before the target datetime
# Versions are sorted by version number, not timestamp, so we need to check all
best_version = None
@ -634,9 +693,9 @@ class Store:
)
# Checkout to the found version
table.checkout(best_version)
await table.checkout(best_version)
def list_table_versions(self, table_name: str) -> list[dict[str, Any]]:
async def list_table_versions(self, table_name: str) -> list[dict[str, Any]]:
"""List version history for a table.
Args:
@ -655,4 +714,4 @@ class Store:
if table is None:
raise ValueError(f"Unknown table: {table_name}")
return list(table.list_versions())
return list(await table.list_versions())

152
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -1,19 +1,16 @@
import json
import logging
from typing import TYPE_CHECKING, cast
from typing import TYPE_CHECKING
from uuid import uuid4
if TYPE_CHECKING:
import pandas as pd
from lancedb.query import (
LanceHybridQueryBuilder,
LanceQueryBuilder,
LanceVectorQueryBuilder,
)
from lancedb.query import AsyncQueryBase
from lancedb.index import FTS
from lancedb.rerankers import RRFReranker
from haiku.rag.store.engine import Store
from haiku.rag.store.engine import Store, query_to_pydantic
from haiku.rag.store.models.chunk import Chunk
logger = logging.getLogger(__name__)
@ -26,11 +23,13 @@ class ChunkRepository:
self.store = store
self.embedder = store.embedder
def _ensure_fts_index(self) -> None:
async def _ensure_fts_index(self) -> None:
"""Ensure FTS index exists on the content_fts column."""
try:
self.store.chunks_table.create_fts_index(
"content_fts", replace=True, with_position=True, remove_stop_words=False
await self.store.chunks_table.create_index(
"content_fts",
config=FTS(with_position=True, remove_stop_words=False),
replace=True,
)
except Exception as e:
# Log the error but don't fail - FTS might already exist
@ -69,7 +68,7 @@ class ChunkRepository:
vector=entity.embedding,
)
self.store.chunks_table.add([chunk_record])
await self.store.chunks_table.add([chunk_record])
entity.id = chunk_id
return entity
@ -90,6 +89,7 @@ class ChunkRepository:
chunk_id = str(uuid4())
assert chunk.document_id is not None
assert chunk.embedding is not None
chunk_record = self.store.ChunkRecord(
id=chunk_id,
document_id=chunk.document_id,
@ -105,17 +105,15 @@ class ChunkRepository:
chunk.id = chunk_id
# Single batch insert for all chunks
self.store.chunks_table.add(chunk_records)
await self.store.chunks_table.add(chunk_records)
return chunks
async def get_by_id(self, entity_id: str) -> Chunk | None:
"""Get a chunk by its ID."""
results = list(
self.store.chunks_table.search()
.where(f"id = '{entity_id}'")
.limit(1)
.to_pydantic(self.store.ChunkRecord)
results = await query_to_pydantic(
self.store.chunks_table.query().where(f"id = '{entity_id}'").limit(1),
self.store.ChunkRecord,
)
if not results:
@ -140,9 +138,8 @@ class ChunkRepository:
assert entity.id, "Chunk ID is required for update"
assert entity.embedding is not None, "Chunk must have an embedding"
self.store.chunks_table.update(
where=f"id = '{entity.id}'",
values={
await self.store.chunks_table.update(
{
"document_id": entity.document_id,
"content": entity.content,
"content_fts": self._contextualize_content(entity),
@ -152,6 +149,7 @@ class ChunkRepository:
"order": int(entity.order),
"vector": entity.embedding,
},
where=f"id = '{entity.id}'",
)
return entity
@ -162,21 +160,21 @@ class ChunkRepository:
if chunk is None:
return False
self.store.chunks_table.delete(f"id = '{entity_id}'")
await self.store.chunks_table.delete(f"id = '{entity_id}'")
return True
async def list_all(
self, limit: int | None = None, offset: int | None = None
) -> list[Chunk]:
"""List all chunks with optional pagination."""
query = self.store.chunks_table.search()
query = self.store.chunks_table.query()
if offset is not None:
query = query.offset(offset)
if limit is not None:
query = query.limit(limit)
results = list(query.to_pydantic(self.store.ChunkRecord))
results = await query_to_pydantic(query, self.store.ChunkRecord)
chunks: list[Chunk] = []
for rec in results:
@ -196,13 +194,15 @@ class ChunkRepository:
"""Delete all chunks from the database."""
self.store._assert_writable()
# Drop and recreate table to clear all data
self.store.db.drop_table("chunks")
self.store.chunks_table = self.store.db.create_table(
await self.store.db.drop_table("chunks")
self.store.chunks_table = await self.store.db.create_table(
"chunks", schema=self.store.ChunkRecord
)
# Create FTS index on content_fts (contextualized content) for better search
self.store.chunks_table.create_fts_index(
"content_fts", replace=True, with_position=True, remove_stop_words=False
await self.store.chunks_table.create_index(
"content_fts",
config=FTS(with_position=True, remove_stop_words=False),
replace=True,
)
async def delete_by_document_id(self, document_id: str) -> bool:
@ -213,7 +213,7 @@ class ChunkRepository:
if not chunks:
return False
self.store.chunks_table.delete(f"document_id = '{document_id}'")
await self.store.chunks_table.delete(f"document_id = '{document_id}'")
return True
async def search(
@ -236,63 +236,51 @@ class ChunkRepository:
"""
if not query.strip():
return []
filtered_doc_ids = None
chunk_filter: str | None = None
if filter:
# We perform filtering as a two-step process, first filtering documents, then
# filtering chunks based on those document IDs.
# This is because LanceDB does not support joins directly in search queries.
docs_df = (
self.store.documents_table.search()
# Translate the document-level filter into a chunk-level
# document_id IN (...) clause so LanceDB can combine it with
# limit. The previous two-step pattern (materialize top-N,
# filter in pandas, head(limit)) silently under-returned
# whenever the top-N window lacked `limit` matching chunks.
docs_df = await (
self.store.documents_table.query()
.select(["id"])
.where(filter)
.to_pandas()
)
# Early exit if no documents match the filter
if docs_df.empty:
return []
# Keep as pandas Series for efficient vectorized operations
filtered_doc_ids = docs_df["id"]
id_list = ", ".join(f"'{d}'" for d in docs_df["id"])
chunk_filter = f"document_id IN ({id_list})"
# Prepare search query based on search type
if search_type == "vector":
query_embedding = await self.embedder.embed_query(query)
vector_query = cast(
"LanceVectorQueryBuilder",
self.store.chunks_table.search(
query_embedding, query_type="vector", vector_column_name="vector"
),
results = (
self.store.chunks_table.query()
.nearest_to(query_embedding)
.column("vector")
.refine_factor(self.store._config.search.vector_refine_factor)
)
results = vector_query.refine_factor(
self.store._config.search.vector_refine_factor
)
elif search_type == "fts":
results = self.store.chunks_table.search(query, query_type="fts")
results = self.store.chunks_table.query().nearest_to_text(
query, columns="content_fts"
)
else: # hybrid (default)
query_embedding = await self.embedder.embed_query(query)
# Create RRF reranker
reranker = RRFReranker()
# Perform native hybrid search with RRF reranking
hybrid_query = cast(
"LanceHybridQueryBuilder",
self.store.chunks_table.search(query_type="hybrid")
.vector(query_embedding)
.text(query),
results = (
self.store.chunks_table.query()
.nearest_to(query_embedding)
.column("vector")
.nearest_to_text(query, columns="content_fts")
.refine_factor(self.store._config.search.vector_refine_factor)
.rerank(reranker)
)
results = hybrid_query.refine_factor(
self.store._config.search.vector_refine_factor
).rerank(reranker)
# Apply filtering if needed (common for all search types)
if filtered_doc_ids is not None:
chunks_df = results.to_pandas()
filtered_chunks_df = chunks_df.loc[
chunks_df["document_id"].isin(filtered_doc_ids)
].head(limit)
return await self._process_search_results(filtered_chunks_df)
# No filtering needed, apply limit and return
if chunk_filter is not None:
results = results.where(chunk_filter)
results = results.limit(limit)
return await self._process_search_results(results)
@ -312,18 +300,18 @@ class ChunkRepository:
Returns:
List of chunks ordered by their order field.
"""
query = self.store.chunks_table.search().where(f"document_id = '{document_id}'")
query = self.store.chunks_table.query().where(f"document_id = '{document_id}'")
if offset is not None:
query = query.offset(offset)
if limit is not None:
query = query.limit(limit)
results = list(query.to_pydantic(self.store.ChunkRecord))
results = await query_to_pydantic(query, self.store.ChunkRecord)
# Get document info (only metadata columns, skip content/docling blobs)
doc_rows = list(
self.store.documents_table.search()
doc_rows = await (
self.store.documents_table.query()
.select(["id", "uri", "title", "metadata"])
.where(f"id = '{document_id}'")
.limit(1)
@ -355,8 +343,8 @@ class ChunkRepository:
async def count_by_document_id(self, document_id: str) -> int:
"""Count the number of chunks for a specific document."""
df = (
self.store.chunks_table.search()
df = await (
self.store.chunks_table.query()
.select(["id"])
.where(f"document_id = '{document_id}'")
.to_pandas()
@ -381,10 +369,8 @@ class ChunkRepository:
f" AND `order` >= {min_order}"
f" AND `order` <= {max_order}"
)
results = list(
self.store.chunks_table.search()
.where(where)
.to_pydantic(self.store.ChunkRecord)
results = await query_to_pydantic(
self.store.chunks_table.query().where(where), self.store.ChunkRecord
)
return [
Chunk(
@ -398,12 +384,12 @@ class ChunkRepository:
]
async def _process_search_results(
self, query_result: "pd.DataFrame | LanceQueryBuilder"
self, query_result: "pd.DataFrame | AsyncQueryBase"
) -> list[tuple[Chunk, float]]:
"""Process search results into chunks with document info and scores.
Args:
query_result: Either a pandas DataFrame or a LanceDB query result
query_result: Either a pandas DataFrame or a LanceDB async query result
"""
import pandas as pd
@ -426,7 +412,7 @@ class ChunkRepository:
df = query_result
else:
# Convert LanceDB query result to DataFrame
df = query_result.to_pandas()
df = await query_result.to_pandas()
# Extract scores
scores = extract_scores(df)
@ -452,8 +438,8 @@ class ChunkRepository:
if document_ids:
id_list = "', '".join(document_ids)
where_clause = f"id IN ('{id_list}')"
doc_rows = list(
self.store.documents_table.search()
doc_rows = await (
self.store.documents_table.query()
.select(["id", "uri", "title", "metadata"])
.where(where_clause)
.to_list()

79
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

@ -2,7 +2,14 @@ import json
from datetime import datetime
from uuid import uuid4
from haiku.rag.store.engine import DocumentRecord, Store, get_documents_arrow_schema
from lancedb.index import BTree
from haiku.rag.store.engine import (
DocumentRecord,
Store,
get_documents_arrow_schema,
query_to_pydantic,
)
from haiku.rag.store.models.document import Document
from haiku.rag.utils import escape_sql_string
@ -78,7 +85,7 @@ class DocumentRepository:
)
# Add to table
self.store.documents_table.add([doc_record])
await self.store.documents_table.add([doc_record])
entity.id = doc_id
entity.created_at = datetime.fromisoformat(now)
@ -88,11 +95,9 @@ class DocumentRepository:
async def get_by_id(self, entity_id: str) -> Document | None:
"""Get a document by its ID."""
safe_id = escape_sql_string(entity_id)
results = list(
self.store.documents_table.search()
.where(f"id = '{safe_id}'")
.limit(1)
.to_pydantic(DocumentRecord)
results = await query_to_pydantic(
self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1),
DocumentRecord,
)
if not results:
@ -103,8 +108,8 @@ class DocumentRepository:
async def get_content(self, entity_id: str) -> str | None:
"""Get only the text content of a document (skips docling blobs)."""
safe_id = escape_sql_string(entity_id)
results = list(
self.store.documents_table.search()
results = await (
self.store.documents_table.query()
.select(["content"])
.where(f"id = '{safe_id}'")
.limit(1)
@ -119,8 +124,8 @@ class DocumentRepository:
async def get_docling_data(self, entity_id: str) -> Document | None:
"""Get a document with only docling data loaded (skips content blob)."""
safe_id = escape_sql_string(entity_id)
results = list(
self.store.documents_table.search()
results = await (
self.store.documents_table.query()
.select(self._DOCLING_COLUMNS)
.where(f"id = '{safe_id}'")
.limit(1)
@ -141,8 +146,8 @@ class DocumentRepository:
async def get_pages_data(self, entity_id: str) -> Document | None:
"""Get a document with only page image data loaded."""
safe_id = escape_sql_string(entity_id)
results = list(
self.store.documents_table.search()
results = await (
self.store.documents_table.query()
.select(["id", "docling_pages"])
.where(f"id = '{safe_id}'")
.limit(1)
@ -171,9 +176,8 @@ class DocumentRepository:
# Update the record
safe_id = escape_sql_string(entity.id)
self.store.documents_table.update(
where=f"id = '{safe_id}'",
values={
await self.store.documents_table.update(
{
"content": entity.content,
"uri": entity.uri,
"title": entity.title,
@ -183,6 +187,7 @@ class DocumentRepository:
"docling_version": entity.docling_version,
"updated_at": now,
},
where=f"id = '{safe_id}'",
)
return entity
@ -202,7 +207,7 @@ class DocumentRepository:
# Delete the document
safe_id = escape_sql_string(entity_id)
self.store.documents_table.delete(f"id = '{safe_id}'")
await self.store.documents_table.delete(f"id = '{safe_id}'")
return True
_LISTING_COLUMNS = ["id", "title", "uri", "metadata", "created_at", "updated_at"]
@ -226,7 +231,7 @@ class DocumentRepository:
Returns:
List of Document instances matching the criteria.
"""
query = self.store.documents_table.search()
query = self.store.documents_table.query()
if not include_content:
query = query.select(self._LISTING_COLUMNS)
@ -238,7 +243,7 @@ class DocumentRepository:
query = query.limit(limit)
if include_content:
results = list(query.to_pydantic(DocumentRecord))
results = await query_to_pydantic(query, DocumentRecord)
return [self._record_to_document(doc) for doc in results]
return [
@ -255,7 +260,7 @@ class DocumentRepository:
if row.get("updated_at")
else datetime.now(),
)
for row in query.to_list()
for row in await query.to_list()
]
async def count(self, filter: str | None = None) -> int:
@ -267,16 +272,14 @@ class DocumentRepository:
Returns:
Number of documents matching the criteria.
"""
return self.store.documents_table.count_rows(filter=filter)
return await self.store.documents_table.count_rows(filter=filter)
async def get_by_uri(self, uri: str) -> Document | None:
"""Get a document by its URI."""
escaped_uri = escape_sql_string(uri)
results = list(
self.store.documents_table.search()
.where(f"uri = '{escaped_uri}'")
.limit(1)
.to_pydantic(DocumentRecord)
results = await query_to_pydantic(
self.store.documents_table.query().where(f"uri = '{escaped_uri}'").limit(1),
DocumentRecord,
)
if not results:
@ -291,29 +294,29 @@ class DocumentRepository:
# Delete all chunks and items first
await self.chunk_repository.delete_all()
self.store.db.drop_table("document_items")
self.store.document_items_table = self.store.db.create_table(
await self.store.db.drop_table("document_items")
self.store.document_items_table = await self.store.db.create_table(
"document_items", schema=DocumentItemRecord
)
self.store.document_items_table.create_scalar_index(
"document_id", index_type="BTREE", replace=True
await self.store.document_items_table.create_index(
"document_id", config=BTree(), replace=True
)
self.store.document_items_table.create_scalar_index(
"position", index_type="BTREE", replace=True
await self.store.document_items_table.create_index(
"position", config=BTree(), replace=True
)
self.store.document_items_table.create_scalar_index(
"self_ref", index_type="BTREE", replace=True
await self.store.document_items_table.create_index(
"self_ref", config=BTree(), replace=True
)
# Get count before deletion
count = len(
list(
self.store.documents_table.search().limit(1).to_pydantic(DocumentRecord)
await query_to_pydantic(
self.store.documents_table.query().limit(1), DocumentRecord
)
)
if count > 0:
# Drop and recreate table to clear all data
self.store.db.drop_table("documents")
self.store.documents_table = self.store.db.create_table(
await self.store.db.drop_table("documents")
self.store.documents_table = await self.store.db.create_table(
"documents", schema=get_documents_arrow_schema()
)

22
haiku_rag_slim/haiku/rag/store/repositories/document_item.py
View file

@ -38,13 +38,13 @@ class DocumentItemRepository:
)
for item in items
]
self.store.document_items_table.add(records)
await self.store.document_items_table.add(records)
async def get_all_items(self, document_id: str) -> list[DocumentItem]:
"""Get all items for a document, sorted by position."""
safe_id = escape_sql_string(document_id)
rows = (
self.store.document_items_table.search()
rows = await (
self.store.document_items_table.query()
.where(f"document_id = '{safe_id}'")
.to_list()
)
@ -64,11 +64,11 @@ class DocumentItemRepository:
Returns:
Dict mapping document_id to sorted list of DocumentItem.
"""
query = self.store.document_items_table.search()
query = self.store.document_items_table.query()
if document_ids is not None:
safe_ids = ", ".join(f"'{escape_sql_string(did)}'" for did in document_ids)
query = query.where(f"document_id IN ({safe_ids})")
rows = query.to_list()
rows = await query.to_list()
grouped: dict[str, list[DocumentItem]] = {}
for row in rows:
@ -83,8 +83,8 @@ class DocumentItemRepository:
) -> list[DocumentItem]:
"""Get items for a document within a position range (inclusive)."""
safe_id = escape_sql_string(document_id)
rows = (
self.store.document_items_table.search()
rows = await (
self.store.document_items_table.query()
.where(
f"document_id = '{safe_id}' "
f"AND position >= {start} AND position <= {end}"
@ -102,8 +102,8 @@ class DocumentItemRepository:
safe_id = escape_sql_string(document_id)
refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs)
rows = (
self.store.document_items_table.search()
rows = await (
self.store.document_items_table.query()
.select(["self_ref", "position"])
.where(f"document_id = '{safe_id}' AND self_ref IN ({refs_sql})")
.to_list()
@ -113,7 +113,7 @@ class DocumentItemRepository:
async def get_item_count(self, document_id: str) -> int:
"""Count items for a document."""
safe_id = escape_sql_string(document_id)
return self.store.document_items_table.count_rows(
return await self.store.document_items_table.count_rows(
filter=f"document_id = '{safe_id}'"
)
@ -121,4 +121,4 @@ class DocumentItemRepository:
"""Delete all items for a document."""
self.store._assert_writable()
safe_id = escape_sql_string(document_id)
self.store.document_items_table.delete(f"document_id = '{safe_id}'")
await self.store.document_items_table.delete(f"document_id = '{safe_id}'")

54
haiku_rag_slim/haiku/rag/store/repositories/settings.py
View file

@ -1,6 +1,6 @@
import json
from haiku.rag.store.engine import SettingsRecord, Store
from haiku.rag.store.engine import SettingsRecord, Store, query_to_pydantic
class ConfigMismatchError(Exception):
@ -18,16 +18,14 @@ class SettingsRepository:
async def create(self, entity: dict) -> dict:
"""Create settings in the database."""
settings_record = SettingsRecord(id="settings", settings=json.dumps(entity))
self.store.settings_table.add([settings_record])
await self.store.settings_table.add([settings_record])
return entity
async def get_by_id(self, entity_id: str) -> dict | None:
"""Get settings by ID."""
results = list(
self.store.settings_table.search()
.where(f"id = '{entity_id}'")
.limit(1)
.to_pydantic(SettingsRecord)
results = await query_to_pydantic(
self.store.settings_table.query().where(f"id = '{entity_id}'").limit(1),
SettingsRecord,
)
if not results:
@ -37,32 +35,32 @@ class SettingsRepository:
async def update(self, entity: dict) -> dict:
"""Update existing settings."""
self.store.settings_table.update(
where="id = 'settings'", values={"settings": json.dumps(entity)}
await self.store.settings_table.update(
{"settings": json.dumps(entity)}, where="id = 'settings'"
)
return entity
async def delete(self, entity_id: str) -> bool:
"""Delete settings by ID."""
self.store.settings_table.delete(f"id = '{entity_id}'")
await self.store.settings_table.delete(f"id = '{entity_id}'")
return True
async def list_all(
self, limit: int | None = None, offset: int | None = None
) -> list[dict]:
"""List all settings."""
results = list(self.store.settings_table.search().to_pydantic(SettingsRecord))
results = await query_to_pydantic(
self.store.settings_table.query(), SettingsRecord
)
return [
json.loads(record.settings) if record.settings else {} for record in results
]
def get_current_settings(self) -> dict:
async def get_current_settings(self) -> dict:
"""Get the current settings."""
results = list(
self.store.settings_table.search()
.where("id = 'settings'")
.limit(1)
.to_pydantic(SettingsRecord)
results = await query_to_pydantic(
self.store.settings_table.query().where("id = 'settings'").limit(1),
SettingsRecord,
)
if not results:
@ -70,17 +68,15 @@ class SettingsRepository:
return json.loads(results[0].settings) if results[0].settings else {}
def save_current_settings(self) -> None:
async def save_current_settings(self) -> None:
"""Save the current configuration to the database."""
self.store._assert_writable()
current_config = self.store._config.model_dump(mode="json")
# Check if settings exist
existing = list(
self.store.settings_table.search()
.where("id = 'settings'")
.limit(1)
.to_pydantic(SettingsRecord)
existing = await query_to_pydantic(
self.store.settings_table.query().where("id = 'settings'").limit(1),
SettingsRecord,
)
if existing:
@ -91,24 +87,24 @@ class SettingsRepository:
# Update existing settings
if existing_settings != current_config:
self.store.settings_table.update(
await self.store.settings_table.update(
{"settings": json.dumps(current_config)},
where="id = 'settings'",
values={"settings": json.dumps(current_config)},
)
else:
# Create new settings
settings_record = SettingsRecord(
id="settings", settings=json.dumps(current_config)
)
self.store.settings_table.add([settings_record])
await self.store.settings_table.add([settings_record])
def validate_config_compatibility(self) -> None:
async def validate_config_compatibility(self) -> None:
"""Validate that the current configuration is compatible with stored settings."""
stored_settings = self.get_current_settings()
stored_settings = await self.get_current_settings()
# If no stored settings, this is a new database - save current config and return
if not stored_settings:
self.save_current_settings()
await self.save_current_settings()
return
current_config = self.store._config.model_dump(mode="json")

10
haiku_rag_slim/haiku/rag/store/upgrades/__init__.py
View file

@ -1,7 +1,7 @@
import logging
from collections.abc import Callable
from collections.abc import Callable, Coroutine
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any
from packaging.version import Version, parse
@ -16,7 +16,7 @@ class Upgrade:
"""Represents a database upgrade step."""
version: str
apply: Callable[["Store"], None]
apply: Callable[["Store"], Coroutine[Any, Any, None]]
description: str = ""
@ -36,7 +36,7 @@ def get_pending_upgrades(from_version: str) -> list[Upgrade]:
return [s for s in sorted_steps if v_from < parse(s.version)]
def run_pending_upgrades(store: "Store", from_version: str) -> list[str]:
async def run_pending_upgrades(store: "Store", from_version: str) -> list[str]:
"""Run upgrades where from_version < step.version.
Returns:
@ -58,7 +58,7 @@ def run_pending_upgrades(store: "Store", from_version: str) -> list[str]:
idx,
len(applicable),
)
step.apply(store)
await step.apply(store)
logger.info("Completed upgrade %s", step.version)
applied.append(
f"{step.version}: {step.description}" if step.description else step.version

12
haiku_rag_slim/haiku/rag/store/upgrades/v0_20_0.py
View file

@ -7,12 +7,12 @@ from haiku.rag.store.engine import Store
from haiku.rag.store.upgrades import Upgrade
def _apply_add_docling_document_columns(store: Store) -> None: # pragma: no cover
async def _apply_add_docling_document_columns(store: Store) -> None: # pragma: no cover
"""Add 'docling_document_json' and 'docling_version' columns to documents table."""
# Read existing rows using Arrow for schema-agnostic access
try:
docs_arrow = store.documents_table.search().to_arrow()
docs_arrow = await store.documents_table.query().to_arrow()
rows = docs_arrow.to_pylist()
except Exception:
rows = []
@ -30,11 +30,13 @@ def _apply_add_docling_document_columns(store: Store) -> None: # pragma: no cov
# Drop and recreate documents table with the new schema
try:
store.db.drop_table("documents")
await store.db.drop_table("documents")
except Exception:
pass
store.documents_table = store.db.create_table("documents", schema=DocumentRecordV3)
store.documents_table = await store.db.create_table(
"documents", schema=DocumentRecordV3
)
# Reinsert previous rows with new columns as None
if rows:
@ -58,7 +60,7 @@ def _apply_add_docling_document_columns(store: Store) -> None: # pragma: no cov
)
)
store.documents_table.add(backfilled)
await store.documents_table.add(backfilled)
upgrade_add_docling_document = Upgrade(

19
haiku_rag_slim/haiku/rag/store/upgrades/v0_23_1.py
View file

@ -1,5 +1,6 @@
import json
from lancedb.index import FTS
from lancedb.pydantic import LanceModel, Vector
from pydantic import Field
@ -7,11 +8,11 @@ from haiku.rag.store.engine import Store
from haiku.rag.store.upgrades import Upgrade
def _apply_add_content_fts(store: Store) -> None: # pragma: no cover
async def _apply_add_content_fts(store: Store) -> None: # pragma: no cover
"""Add content_fts column with contextualized content for better FTS."""
# Read existing chunks
try:
chunks_arrow = store.chunks_table.search().to_arrow()
chunks_arrow = await store.chunks_table.query().to_arrow()
rows = chunks_arrow.to_pylist()
except Exception:
return
@ -38,11 +39,11 @@ def _apply_add_content_fts(store: Store) -> None: # pragma: no cover
# Drop and recreate table with new schema
try:
store.db.drop_table("chunks")
await store.db.drop_table("chunks")
except Exception:
pass
store.chunks_table = store.db.create_table("chunks", schema=ChunkRecord)
store.chunks_table = await store.db.create_table("chunks", schema=ChunkRecord)
# Populate content_fts with contextualized content
new_records: list[ChunkRecord] = []
@ -79,17 +80,19 @@ def _apply_add_content_fts(store: Store) -> None: # pragma: no cover
)
if new_records:
store.chunks_table.add(new_records)
await store.chunks_table.add(new_records)
# Drop old FTS index on content column if it exists
try:
store.chunks_table.drop_index("content_idx")
await store.chunks_table.drop_index("content_idx")
except Exception:
pass
# Create FTS index on content_fts
store.chunks_table.create_fts_index(
"content_fts", replace=True, with_position=True, remove_stop_words=False
await store.chunks_table.create_index(
"content_fts",
config=FTS(with_position=True, remove_stop_words=False),
replace=True,
)

87
haiku_rag_slim/haiku/rag/store/upgrades/v0_25_0.py
View file

@ -15,7 +15,7 @@ logger = logging.getLogger(__name__)
BATCH_SIZE = 10
def _apply_compress_docling_document(store: Store) -> None: # pragma: no cover
async def _apply_compress_docling_document(store: Store) -> None: # pragma: no cover
"""Migrate docling_document_json (str) to docling_document (compressed bytes)."""
class DocumentRecordV4(LanceModel):
@ -76,36 +76,33 @@ def _apply_compress_docling_document(store: Store) -> None: # pragma: no cover
# First pass: collect document IDs to process
try:
ids = [
row["id"]
for row in store.documents_table.search()
.select(["id"])
.to_arrow()
.to_pylist()
]
ids = (
await store.documents_table.query().select(["id"]).to_arrow()
).to_pylist()
ids = [row["id"] for row in ids]
except Exception:
ids = []
if not ids:
# Check if there's a staging table from a failed migration to recover from
if "documents_v4_staging" in store.db.list_tables().tables:
staging_table = store.db.open_table("documents_v4_staging")
staging_ids = [
row["id"]
for row in staging_table.search().select(["id"]).to_arrow().to_pylist()
]
if "documents_v4_staging" in (await store.db.list_tables()).tables:
staging_table = await store.db.open_table("documents_v4_staging")
staging_ids = (
await staging_table.query().select(["id"]).to_arrow()
).to_pylist()
staging_ids = [row["id"] for row in staging_ids]
if staging_ids:
logger.info(
"Recovering %d documents from failed migration", len(staging_ids)
)
# Create new documents table and copy from staging
store.documents_table = None
if "documents" in store.db.list_tables().tables:
store.db.drop_table("documents")
store.documents_table = store.db.create_table(
if "documents" in (await store.db.list_tables()).tables:
await store.db.drop_table("documents")
store.documents_table = await store.db.create_table(
"documents", schema=get_documents_arrow_schema_v4()
)
# Copy data from staging (reuse the copy logic below by jumping there)
# Copy data from staging
total_batches = (len(staging_ids) + BATCH_SIZE - 1) // BATCH_SIZE
for batch_num, i in enumerate(
range(0, len(staging_ids), BATCH_SIZE), 1
@ -113,11 +110,10 @@ def _apply_compress_docling_document(store: Store) -> None: # pragma: no cover
batch_ids = staging_ids[i : i + BATCH_SIZE]
id_list = ", ".join(f"'{id}'" for id in batch_ids)
batch = (
staging_table.search()
await staging_table.query()
.where(f"id IN ({id_list})")
.to_arrow()
.to_pylist()
)
).to_pylist()
records = [
DocumentRecordV4(
id=row["id"],
@ -133,26 +129,26 @@ def _apply_compress_docling_document(store: Store) -> None: # pragma: no cover
for row in batch
]
if records:
store.documents_table.add(records)
await store.documents_table.add(records)
logger.info("Recovered batch %d/%d", batch_num, total_batches)
# Cleanup staging
store.db.drop_table("documents_v4_staging")
await store.db.drop_table("documents_v4_staging")
logger.info("Recovery complete")
return
# No documents and no staging to recover, just recreate table with new schema
store.documents_table = None
if "documents" in store.db.list_tables().tables:
store.db.drop_table("documents")
store.documents_table = store.db.create_table(
if "documents" in (await store.db.list_tables()).tables:
await store.db.drop_table("documents")
store.documents_table = await store.db.create_table(
"documents", schema=get_documents_arrow_schema_v4()
)
return
# Create staging table with new schema
if "documents_v4_staging" in store.db.list_tables().tables:
store.db.drop_table("documents_v4_staging")
staging_table = store.db.create_table(
if "documents_v4_staging" in (await store.db.list_tables()).tables:
await store.db.drop_table("documents_v4_staging")
staging_table = await store.db.create_table(
"documents_v4_staging", schema=get_documents_arrow_schema_v4()
)
@ -166,15 +162,12 @@ def _apply_compress_docling_document(store: Store) -> None: # pragma: no cover
id_list = ", ".join(f"'{id}'" for id in batch_ids)
batch = (
store.documents_table.search()
.where(f"id IN ({id_list})")
.to_arrow()
.to_pylist()
)
await store.documents_table.query().where(f"id IN ({id_list})").to_arrow()
).to_pylist()
migrated_batch = [migrate_row(row) for row in batch]
if migrated_batch:
staging_table.add(migrated_batch)
await staging_table.add(migrated_batch)
logger.info(
"Compressed batch %d/%d (%d documents)",
@ -185,17 +178,15 @@ def _apply_compress_docling_document(store: Store) -> None: # pragma: no cover
# Replace old table with staging table
store.documents_table = None
if "documents" in store.db.list_tables().tables:
store.db.drop_table("documents")
store.documents_table = store.db.create_table(
if "documents" in (await store.db.list_tables()).tables:
await store.db.drop_table("documents")
store.documents_table = await store.db.create_table(
"documents", schema=get_documents_arrow_schema_v4()
)
# Copy from staging to final table in batches
staging_ids = [
row["id"]
for row in staging_table.search().select(["id"]).to_arrow().to_pylist()
]
staging_ids = (await staging_table.query().select(["id"]).to_arrow()).to_pylist()
staging_ids = [row["id"] for row in staging_ids]
logger.info("Copying %d documents to new table", len(staging_ids))
@ -204,8 +195,8 @@ def _apply_compress_docling_document(store: Store) -> None: # pragma: no cover
id_list = ", ".join(f"'{id}'" for id in batch_ids)
batch = (
staging_table.search().where(f"id IN ({id_list})").to_arrow().to_pylist()
)
await staging_table.query().where(f"id IN ({id_list})").to_arrow()
).to_pylist()
records = [
DocumentRecordV4(
id=row["id"],
@ -221,18 +212,18 @@ def _apply_compress_docling_document(store: Store) -> None: # pragma: no cover
for row in batch
]
if records:
store.documents_table.add(records)
await store.documents_table.add(records)
logger.info("Copied batch %d/%d", batch_num, total_batches)
# Cleanup staging table
if "documents_v4_staging" in store.db.list_tables().tables:
store.db.drop_table("documents_v4_staging")
if "documents_v4_staging" in (await store.db.list_tables()).tables:
await store.db.drop_table("documents_v4_staging")
# Vacuum all tables (destructive migration, no history preserved)
logger.info("Vacuuming database")
for table in [store.documents_table, store.chunks_table, store.settings_table]:
try:
table.optimize(cleanup_older_than=timedelta(seconds=0))
await table.optimize(cleanup_older_than=timedelta(seconds=0))
except Exception:
pass

85
haiku_rag_slim/haiku/rag/store/upgrades/v0_38_0.py
View file

@ -16,7 +16,7 @@ logger = logging.getLogger(__name__)
BATCH_SIZE = 5
def _apply_split_pages_zstd(store: Store) -> None: # pragma: no cover
async def _apply_split_pages_zstd(store: Store) -> None: # pragma: no cover
"""Split docling_document into structure + pages and re-compress with zstd."""
class DocumentRecordV5(LanceModel):
@ -99,32 +99,29 @@ def _apply_split_pages_zstd(store: Store) -> None: # pragma: no cover
# First pass: collect document IDs to process
try:
ids = [
row["id"]
for row in store.documents_table.search()
.select(["id"])
.to_arrow()
.to_pylist()
]
ids = (
await store.documents_table.query().select(["id"]).to_arrow()
).to_pylist()
ids = [row["id"] for row in ids]
except (pa.ArrowInvalid, pa.ArrowNotImplementedError, OSError):
ids = []
if not ids:
# Check for staging table from a failed migration
if staging_name in store.db.list_tables().tables:
staging_table = store.db.open_table(staging_name)
staging_ids = [
row["id"]
for row in staging_table.search().select(["id"]).to_arrow().to_pylist()
]
if staging_name in (await store.db.list_tables()).tables:
staging_table = await store.db.open_table(staging_name)
staging_ids = (
await staging_table.query().select(["id"]).to_arrow()
).to_pylist()
staging_ids = [row["id"] for row in staging_ids]
if staging_ids:
logger.info(
"Recovering %d documents from failed migration", len(staging_ids)
)
store.documents_table = None
if "documents" in store.db.list_tables().tables:
store.db.drop_table("documents")
store.documents_table = store.db.create_table(
if "documents" in (await store.db.list_tables()).tables:
await store.db.drop_table("documents")
store.documents_table = await store.db.create_table(
"documents", schema=get_documents_arrow_schema_v5()
)
total_batches = (len(staging_ids) + BATCH_SIZE - 1) // BATCH_SIZE
@ -134,32 +131,31 @@ def _apply_split_pages_zstd(store: Store) -> None: # pragma: no cover
batch_ids = staging_ids[i : i + BATCH_SIZE]
id_list = ", ".join(f"'{doc_id}'" for doc_id in batch_ids)
batch = (
staging_table.search()
await staging_table.query()
.where(f"id IN ({id_list})")
.to_arrow()
.to_pylist()
)
).to_pylist()
records = [copy_staging_row(row) for row in batch]
if records:
store.documents_table.add(records)
await store.documents_table.add(records)
logger.info("Recovered batch %d/%d", batch_num, total_batches)
store.db.drop_table(staging_name)
await store.db.drop_table(staging_name)
logger.info("Recovery complete")
return
# No documents and no staging — recreate table with new schema
store.documents_table = None
if "documents" in store.db.list_tables().tables:
store.db.drop_table("documents")
store.documents_table = store.db.create_table(
if "documents" in (await store.db.list_tables()).tables:
await store.db.drop_table("documents")
store.documents_table = await store.db.create_table(
"documents", schema=get_documents_arrow_schema_v5()
)
return
# Create staging table with new schema
if staging_name in store.db.list_tables().tables:
store.db.drop_table(staging_name)
staging_table = store.db.create_table(
if staging_name in (await store.db.list_tables()).tables:
await store.db.drop_table(staging_name)
staging_table = await store.db.create_table(
staging_name, schema=get_documents_arrow_schema_v5()
)
@ -177,15 +173,12 @@ def _apply_split_pages_zstd(store: Store) -> None: # pragma: no cover
id_list = ", ".join(f"'{doc_id}'" for doc_id in batch_ids)
batch = (
store.documents_table.search()
.where(f"id IN ({id_list})")
.to_arrow()
.to_pylist()
)
await store.documents_table.query().where(f"id IN ({id_list})").to_arrow()
).to_pylist()
migrated_batch = [migrate_row(row) for row in batch]
if migrated_batch:
staging_table.add(migrated_batch)
await staging_table.add(migrated_batch)
logger.info(
"Migrated batch %d/%d (%d documents)",
@ -196,17 +189,15 @@ def _apply_split_pages_zstd(store: Store) -> None: # pragma: no cover
# Replace old table with staging table
store.documents_table = None
if "documents" in store.db.list_tables().tables:
store.db.drop_table("documents")
store.documents_table = store.db.create_table(
if "documents" in (await store.db.list_tables()).tables:
await store.db.drop_table("documents")
store.documents_table = await store.db.create_table(
"documents", schema=get_documents_arrow_schema_v5()
)
# Copy from staging to final table in batches
staging_ids = [
row["id"]
for row in staging_table.search().select(["id"]).to_arrow().to_pylist()
]
staging_ids = (await staging_table.query().select(["id"]).to_arrow()).to_pylist()
staging_ids = [row["id"] for row in staging_ids]
logger.info("Copying %d documents to new table", len(staging_ids))
@ -215,22 +206,22 @@ def _apply_split_pages_zstd(store: Store) -> None: # pragma: no cover
id_list = ", ".join(f"'{doc_id}'" for doc_id in batch_ids)
batch = (
staging_table.search().where(f"id IN ({id_list})").to_arrow().to_pylist()
)
await staging_table.query().where(f"id IN ({id_list})").to_arrow()
).to_pylist()
records = [copy_staging_row(row) for row in batch]
if records:
store.documents_table.add(records)
await store.documents_table.add(records)
logger.info("Copied batch %d/%d", batch_num, total_batches)
# Cleanup staging table
if staging_name in store.db.list_tables().tables:
store.db.drop_table(staging_name)
if staging_name in (await store.db.list_tables()).tables:
await store.db.drop_table(staging_name)
# Vacuum all tables
logger.info("Vacuuming database")
for table in [store.documents_table, store.chunks_table, store.settings_table]:
try:
table.optimize(cleanup_older_than=timedelta(seconds=0))
await table.optimize(cleanup_older_than=timedelta(seconds=0))
except Exception:
pass

14
haiku_rag_slim/haiku/rag/store/upgrades/v0_40_0.py
View file

@ -8,7 +8,7 @@ from haiku.rag.utils import escape_sql_string
logger = logging.getLogger(__name__)
def _apply_populate_document_items(store: Store) -> None: # pragma: no cover
async def _apply_populate_document_items(store: Store) -> None: # pragma: no cover
"""Populate document_items table from existing docling documents."""
from docling_core.types.doc.document import DoclingDocument
@ -16,10 +16,8 @@ def _apply_populate_document_items(store: Store) -> None: # pragma: no cover
from haiku.rag.store.models.document_item import extract_items
# Get all document IDs that have docling data
ids = [
row["id"]
for row in store.documents_table.search().select(["id"]).to_arrow().to_pylist()
]
ids = (await store.documents_table.query().select(["id"]).to_arrow()).to_pylist()
ids = [row["id"] for row in ids]
if not ids:
logger.info("No documents to migrate")
@ -33,8 +31,8 @@ def _apply_populate_document_items(store: Store) -> None: # pragma: no cover
for idx, doc_id in enumerate(ids, 1):
# Load only docling data
safe_id = escape_sql_string(doc_id)
rows = (
store.documents_table.search()
rows = await (
store.documents_table.query()
.select(["id", "docling_document"])
.where(f"id = '{safe_id}'")
.limit(1)
@ -68,7 +66,7 @@ def _apply_populate_document_items(store: Store) -> None: # pragma: no cover
)
for item in items
]
store.document_items_table.add(records)
await store.document_items_table.add(records)
migrated += 1
if idx % 10 == 0 or idx == total:

67
tests/agents/qa/test_qa.py
View file

@ -15,56 +15,55 @@ def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_qa")
def test_get_qa_agent_factory(temp_db_path):
@pytest.mark.asyncio
async def test_get_qa_agent_factory(temp_db_path):
"""Test get_qa_agent factory function creates a properly configured agent."""
from haiku.rag.agents.qa import get_qa_agent
client = HaikuRAG(temp_db_path, create=True)
agent = get_qa_agent(client, Config)
async with HaikuRAG(temp_db_path, create=True) as client:
agent = get_qa_agent(client, Config)
assert agent is not None
assert isinstance(agent, QuestionAnswerAgent)
# Verify internal client is set correctly
assert agent._client is client
client.close()
assert agent is not None
assert isinstance(agent, QuestionAnswerAgent)
# Verify internal client is set correctly
assert agent._client is client
def test_get_qa_agent_with_custom_prompt(temp_db_path):
@pytest.mark.asyncio
async def test_get_qa_agent_with_custom_prompt(temp_db_path):
"""Test get_qa_agent factory with custom system prompt."""
from haiku.rag.agents.qa import get_qa_agent
client = HaikuRAG(temp_db_path, create=True)
custom_prompt = "You are a custom QA assistant."
agent = get_qa_agent(client, Config, system_prompt=custom_prompt)
async with HaikuRAG(temp_db_path, create=True) as client:
custom_prompt = "You are a custom QA assistant."
agent = get_qa_agent(client, Config, system_prompt=custom_prompt)
assert agent is not None
assert isinstance(agent, QuestionAnswerAgent)
assert agent._system_prompt == custom_prompt
client.close()
assert agent is not None
assert isinstance(agent, QuestionAnswerAgent)
assert agent._system_prompt == custom_prompt
@pytest.mark.vcr()
async def test_qa_ollama(allow_model_requests, qa_corpus: Dataset, temp_db_path):
"""Test Ollama QA with LLM judge (VCR recorded)."""
client = HaikuRAG(temp_db_path, create=True)
qa = QuestionAnswerAgent(
client, ModelConfig(provider="ollama", name="gpt-oss", enable_thinking=True)
)
llm_judge = LLMJudge()
async with HaikuRAG(temp_db_path, create=True) as client:
qa = QuestionAnswerAgent(
client,
ModelConfig(provider="ollama", name="gpt-oss", enable_thinking=True),
)
llm_judge = LLMJudge()
doc = qa_corpus[1]
await client.create_document(
content=doc["document_extracted"], uri=doc["document_id"]
)
doc = qa_corpus[1]
await client.create_document(
content=doc["document_extracted"], uri=doc["document_id"]
)
question = doc["question"]
expected_answer = doc["answer"]
question = doc["question"]
expected_answer = doc["answer"]
answer, _ = await qa.answer(question)
is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer)
answer, _ = await qa.answer(question)
is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer)
assert is_equivalent, (
f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}"
)
assert is_equivalent, (
f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}"
)

34
tests/agents/research/test_research_graph.py
View file

@ -21,28 +21,26 @@ async def test_graph_end_to_end(allow_model_requests, temp_db_path, qa_corpus):
"""Test research graph with real LLM calls recorded via VCR."""
graph = build_research_graph()
client = HaikuRAG(temp_db_path, create=True)
doc = qa_corpus[0]
await client.create_document(
content=doc["document_extracted"], uri=doc["document_id"]
)
async with HaikuRAG(temp_db_path, create=True) as client:
doc = qa_corpus[0]
await client.create_document(
content=doc["document_extracted"], uri=doc["document_id"]
)
state = ResearchState(
context=ResearchContext(original_question=doc["question"]),
max_iterations=1,
max_concurrency=1,
)
state = ResearchState(
context=ResearchContext(original_question=doc["question"]),
max_iterations=1,
max_concurrency=1,
)
deps = ResearchDeps(client=client)
deps = ResearchDeps(client=client)
result = await graph.run(state=state, deps=deps)
result = await graph.run(state=state, deps=deps)
assert result is not None
assert isinstance(result, ResearchReport)
assert result.title
assert result.executive_summary
client.close()
assert result is not None
assert isinstance(result, ResearchReport)
assert result.title
assert result.executive_summary
def test_iterative_plan_result_model():

25
tests/agents/research/test_search_filter.py
View file

@ -16,21 +16,18 @@ def vcr_cassette_dir():
@pytest.fixture
async def client_with_docs(temp_db_path):
"""Create a client with two distinct documents."""
client = HaikuRAG(temp_db_path, create=True)
async with HaikuRAG(temp_db_path, create=True) as client:
# Add two documents with distinct content
doc1 = await client.create_document(
"Document about cats: Cats are small furry mammals that purr.",
title="Cat Facts",
)
doc2 = await client.create_document(
"Document about dogs: Dogs are loyal companions that bark.",
title="Dog Facts",
)
# Add two documents with distinct content
doc1 = await client.create_document(
"Document about cats: Cats are small furry mammals that purr.",
title="Cat Facts",
)
doc2 = await client.create_document(
"Document about dogs: Dogs are loyal companions that bark.",
title="Dog Facts",
)
yield client, doc1.id, doc2.id
client.close()
yield client, doc1.id, doc2.id
@pytest.mark.vcr()

534
tests/cassettes/test_filter/test_search_with_filter_returns_full_limit.yaml Normal file
View file

File diff suppressed because one or more lines are too long

242
tests/cassettes/test_rebuild/test_rebuild_batch_size_flush.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_rebuild/test_rebuild_full_source_failure_is_logged_and_skipped.yaml Normal file
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_rebuild/test_rebuild_full_with_accessible_source.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_search/test_fts_search_targets_content_fts_column.yaml Normal file
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_versioning/test_aexit_awaits_all_background_vacuums.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_versioning/test_aexit_awaits_background_vacuum.yaml Normal file
View file

File diff suppressed because one or more lines are too long

5
tests/conftest.py
View file

@ -47,13 +47,12 @@ def qa_corpus() -> "Dataset":
@pytest.fixture
def temp_db_path():
def temp_db_path(tmp_path):
"""Create a temporary database path for testing.
Note: Tests that need a database should use HaikuRAG with create=True.
"""
with tempfile.TemporaryDirectory() as temp_dir:
yield Path(temp_dir) / "test.lancedb"
return tmp_path / "test.lancedb"
@pytest.fixture

108
tests/store/test_document_items.py
View file

@ -1,6 +1,10 @@
import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.client.documents import (
_store_document_with_chunks,
_update_document_with_chunks,
)
from haiku.rag.store.engine import Store
from haiku.rag.store.models.document_item import (
DocumentItem,
@ -217,7 +221,7 @@ class TestDocumentItemPopulation:
# Use _store_document_with_chunks directly with empty chunks
# to avoid needing embeddings
created = await rag._store_document_with_chunks(document, [], docling_doc)
created = await _store_document_with_chunks(rag, document, [], docling_doc)
assert created.id is not None
count = await rag.document_item_repository.get_item_count(created.id)
@ -245,7 +249,7 @@ class TestDocumentItemPopulation:
uri="test://doc",
)
document.set_docling(docling_doc)
created = await rag._store_document_with_chunks(document, [], docling_doc)
created = await _store_document_with_chunks(rag, document, [], docling_doc)
assert created.id is not None
assert await rag.document_item_repository.get_item_count(created.id) == 6
@ -254,7 +258,7 @@ class TestDocumentItemPopulation:
new_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Only one item now.")
created.set_docling(new_doc)
await rag._update_document_with_chunks(created, [], new_doc)
await _update_document_with_chunks(rag, created, [], new_doc)
assert await rag.document_item_repository.get_item_count(created.id) == 1
async def test_delete_document_cascades_items(self, temp_db_path):
@ -269,7 +273,7 @@ class TestDocumentItemPopulation:
uri="test://doc",
)
document.set_docling(docling_doc)
created = await rag._store_document_with_chunks(document, [], docling_doc)
created = await _store_document_with_chunks(rag, document, [], docling_doc)
assert created.id is not None
assert await rag.document_item_repository.get_item_count(created.id) == 6
@ -277,8 +281,9 @@ class TestDocumentItemPopulation:
assert await rag.document_item_repository.get_item_count(created.id) == 0
@pytest.mark.asyncio
class TestDocumentItemMigration:
def test_migration_populates_items_for_existing_documents(self, temp_db_path):
async def test_migration_populates_items_for_existing_documents(self, temp_db_path):
"""Test that the v0.40.0 migration populates items for pre-existing documents."""
from haiku.rag.store.compression import compress_docling_split
from haiku.rag.store.engine import DocumentRecord
@ -288,64 +293,59 @@ class TestDocumentItemMigration:
structure, pages = compress_docling_split(json_str)
# Create a database at a pre-migration version with a document
store = Store(temp_db_path, create=True, skip_migration_check=True)
store.set_haiku_version("0.39.0")
doc_record = DocumentRecord(
id="test-doc-1",
content="test content",
uri="test://doc",
docling_document=structure,
docling_pages=pages,
docling_version=docling_doc.version,
)
store.documents_table.add([doc_record])
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
await store.set_haiku_version("0.39.0")
doc_record = DocumentRecord(
id="test-doc-1",
content="test content",
uri="test://doc",
docling_document=structure,
docling_pages=pages,
docling_version=docling_doc.version,
)
await store.documents_table.add([doc_record])
# Verify no items exist yet
assert store.document_items_table.count_rows() == 0
store.close()
# Verify no items exist yet
assert await store.document_items_table.count_rows() == 0
# Re-open with skip_migration_check and run migration
store = Store(temp_db_path, skip_migration_check=True)
applied = store.migrate()
async with Store(temp_db_path, skip_migration_check=True) as store:
applied = await store.migrate()
# Should have applied the v0.40.0 migration
assert any("document_items" in desc for desc in applied)
# Should have applied the v0.40.0 migration
assert any("document_items" in desc for desc in applied)
# Items should now exist
item_count = store.document_items_table.count_rows(
filter="document_id = 'test-doc-1'"
)
assert item_count == 6
# Items should now exist
item_count = await store.document_items_table.count_rows(
filter="document_id = 'test-doc-1'"
)
assert item_count == 6
# Verify item content
items = (
store.document_items_table.search()
.where("document_id = 'test-doc-1'")
.to_list()
)
labels = {row["label"] for row in items}
assert "section_header" in labels
assert "paragraph" in labels
assert "table" in labels
# Verify item content
items = await (
store.document_items_table.query()
.where("document_id = 'test-doc-1'")
.to_list()
)
labels = {row["label"] for row in items}
assert "section_header" in labels
assert "paragraph" in labels
assert "table" in labels
store.close()
def test_migration_skips_documents_without_docling(self, temp_db_path):
async def test_migration_skips_documents_without_docling(self, temp_db_path):
"""Test that migration handles documents without docling data."""
from haiku.rag.store.engine import DocumentRecord
store = Store(temp_db_path, create=True, skip_migration_check=True)
store.set_haiku_version("0.39.0")
doc_record = DocumentRecord(
id="no-docling",
content="plain text document",
)
store.documents_table.add([doc_record])
store.close()
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
await store.set_haiku_version("0.39.0")
doc_record = DocumentRecord(
id="no-docling",
content="plain text document",
)
await store.documents_table.add([doc_record])
store = Store(temp_db_path, skip_migration_check=True)
store.migrate()
async with Store(temp_db_path, skip_migration_check=True) as store:
await store.migrate()
# No items should have been created
assert store.document_items_table.count_rows() == 0
store.close()
# No items should have been created
assert await store.document_items_table.count_rows() == 0

120
tests/store/test_engine.py
View file

@ -5,25 +5,24 @@ from haiku.rag.store.engine import get_database_stats
class TestGetDatabaseStats:
def test_empty_database_stats(self, temp_db_path):
@pytest.mark.asyncio
async def test_empty_database_stats(self, temp_db_path):
"""get_database_stats() on a fresh database reports zero rows and no vector index."""
store = Store(temp_db_path, create=True)
async with Store(temp_db_path, create=True) as store:
stats = await get_database_stats(store.db)
stats = get_database_stats(store.db)
for name in ("documents", "chunks", "document_items", "settings"):
assert stats[name]["exists"] is True
assert stats[name]["num_rows"] >= 0
assert stats[name]["total_bytes"] >= 0
assert stats[name]["num_versions"] >= 1
for name in ("documents", "chunks", "document_items", "settings"):
assert stats[name]["exists"] is True
assert stats[name]["num_rows"] >= 0
assert stats[name]["total_bytes"] >= 0
assert stats[name]["num_versions"] >= 1
assert stats["documents"]["num_rows"] == 0
assert stats["chunks"]["num_rows"] == 0
assert stats["chunks"]["has_vector_index"] is False
assert stats["documents"]["num_rows"] == 0
assert stats["chunks"]["num_rows"] == 0
assert stats["chunks"]["has_vector_index"] is False
store.close()
def test_missing_tables_report_absent(self, temp_db_path):
@pytest.mark.asyncio
async def test_missing_tables_report_absent(self, temp_db_path):
"""Tables that don't exist on the connection are reported as absent."""
import lancedb
from lancedb.pydantic import LanceModel
@ -33,10 +32,10 @@ class TestGetDatabaseStats:
id: str = Field(default="settings")
settings: str = Field(default="{}")
db = lancedb.connect(temp_db_path)
db.create_table("settings", schema=SettingsRecord)
db = await lancedb.connect_async(temp_db_path)
await db.create_table("settings", schema=SettingsRecord)
stats = get_database_stats(db)
stats = await get_database_stats(db)
assert stats["settings"]["exists"] is True
assert stats["documents"] == {"exists": False}
@ -50,54 +49,57 @@ class TestGetDatabaseStats:
from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store)
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store)
doc = await doc_repo.create(Document(content="hello world"))
assert doc.id is not None
doc = await doc_repo.create(Document(content="hello world"))
assert doc.id is not None
await chunk_repo.create(
Chunk(
content="hello world",
document_id=doc.id,
embedding=[0.0] * store.embedder._vector_dim,
await chunk_repo.create(
Chunk(
content="hello world",
document_id=doc.id,
embedding=[0.0] * store.embedder._vector_dim,
)
)
)
stats = get_database_stats(store.db)
assert stats["documents"]["num_rows"] == 1
assert stats["chunks"]["num_rows"] == 1
store.close()
stats = await get_database_stats(store.db)
assert stats["documents"]["num_rows"] == 1
assert stats["chunks"]["num_rows"] == 1
def test_stats_with_vector_index(self, temp_db_path):
@pytest.mark.asyncio
async def test_stats_with_vector_index(self, temp_db_path):
"""get_database_stats() reports vector index details once an index exists."""
from datetime import timedelta
store = Store(temp_db_path, create=True)
dim = store.embedder._vector_dim
from lancedb.index import IvfPq
# Need >=256 rows for IVF_PQ training.
rows = [
{
"id": f"chunk-{i}",
"document_id": "doc-1",
"content": f"content {i}",
"content_fts": "",
"metadata": "{}",
"order": i,
"vector": [float(i % 7) + 0.01 * j for j in range(dim)],
}
for i in range(256)
]
store.chunks_table.add(rows)
store.chunks_table.create_index(
metric="cosine", index_type="IVF_PQ", replace=True
)
store.chunks_table.wait_for_index(["vector_idx"], timeout=timedelta(minutes=1))
async with Store(temp_db_path, create=True) as store:
dim = store.embedder._vector_dim
stats = get_database_stats(store.db)
assert stats["chunks"]["has_vector_index"] is True
assert stats["chunks"]["num_indexed_rows"] >= 0
assert "num_unindexed_rows" in stats["chunks"]
store.close()
# Need >=256 rows for IVF_PQ training.
rows = [
{
"id": f"chunk-{i}",
"document_id": "doc-1",
"content": f"content {i}",
"content_fts": "",
"metadata": "{}",
"order": i,
"vector": [float(i % 7) + 0.01 * j for j in range(dim)],
}
for i in range(256)
]
await store.chunks_table.add(rows)
await store.chunks_table.create_index(
"vector", config=IvfPq(distance_type="cosine"), replace=True
)
await store.chunks_table.wait_for_index(
["vector_idx"], timeout=timedelta(minutes=1)
)
stats = await get_database_stats(store.db)
assert stats["chunks"]["has_vector_index"] is True
assert stats["chunks"]["num_indexed_rows"] >= 0
assert "num_unindexed_rows" in stats["chunks"]

198
tests/store/test_migrations.py
View file

@ -19,150 +19,160 @@ class TestMigrationRequiredError:
class TestMigrationCheck:
def test_new_database_sets_version(self, temp_db_path):
@pytest.mark.asyncio
async def test_new_database_sets_version(self, temp_db_path):
"""New database should set the current package version."""
store = Store(temp_db_path, create=True)
version = store.get_haiku_version()
expected = metadata.version("haiku.rag-slim")
assert version == expected
store.close()
async with Store(temp_db_path, create=True) as store:
version = await store.get_haiku_version()
expected = metadata.version("haiku.rag-slim")
assert version == expected
def test_existing_database_same_version_no_error(self, temp_db_path):
@pytest.mark.asyncio
async def test_existing_database_same_version_no_error(self, temp_db_path):
"""Opening a database with the same version should not error."""
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
# Re-open - should work without error
store = Store(temp_db_path)
store.close()
async with Store(temp_db_path):
pass
def test_version_bump_without_pending_migrations_updates_silently(
@pytest.mark.asyncio
async def test_version_bump_without_pending_migrations_updates_silently(
self, temp_db_path
):
"""When version is outdated but no migrations pending, update version silently."""
store = Store(temp_db_path, create=True)
# Set an older version that has no pending migrations
# (newer than all current upgrade steps)
store.set_haiku_version("100.0.0")
store.close()
async with Store(temp_db_path, create=True) as store:
# Set an older version that has no pending migrations
# (newer than all current upgrade steps)
await store.set_haiku_version("100.0.0")
# Re-open - should update version silently, no error
store = Store(temp_db_path)
# Version should now be current
version = store.get_haiku_version()
expected = metadata.version("haiku.rag-slim")
assert version == expected
store.close()
async with Store(temp_db_path) as store:
# Version should now be current
version = await store.get_haiku_version()
expected = metadata.version("haiku.rag-slim")
assert version == expected
def test_pending_migrations_raises_error(self, temp_db_path):
@pytest.mark.asyncio
async def test_pending_migrations_raises_error(self, temp_db_path):
"""When actual migrations are pending, should raise MigrationRequiredError."""
store = Store(temp_db_path, create=True)
# Set version to before the first upgrade step
store.set_haiku_version("0.19.0")
store.close()
async with Store(temp_db_path, create=True) as store:
# Set version to before the first upgrade step
await store.set_haiku_version("0.19.0")
# Re-open should raise
with pytest.raises(MigrationRequiredError) as exc_info:
Store(temp_db_path)
async with Store(temp_db_path) as store:
pass
assert "migrate" in str(exc_info.value).lower()
def test_pending_migrations_read_only_raises_error(self, temp_db_path):
@pytest.mark.asyncio
async def test_pending_migrations_raises_error_with_create_flag(self, temp_db_path):
"""Opening an existing DB with create=True must still check migrations."""
async with Store(temp_db_path, create=True) as store:
await store.set_haiku_version("0.19.0")
# create=True is idempotent — must not mark an existing populated DB as new
with pytest.raises(MigrationRequiredError):
async with Store(temp_db_path, create=True) as store:
pass
@pytest.mark.asyncio
async def test_pending_migrations_read_only_raises_error(self, temp_db_path):
"""Read-only mode with pending migrations should still raise."""
store = Store(temp_db_path, create=True)
store.set_haiku_version("0.19.0")
store.close()
async with Store(temp_db_path, create=True) as store:
await store.set_haiku_version("0.19.0")
with pytest.raises(MigrationRequiredError):
Store(temp_db_path, read_only=True)
async with Store(temp_db_path, read_only=True) as store:
pass
def test_read_only_version_bump_without_migrations_ok(self, temp_db_path):
@pytest.mark.asyncio
async def test_read_only_version_bump_without_migrations_ok(self, temp_db_path):
"""Read-only mode with version bump but no migrations should work."""
store = Store(temp_db_path, create=True)
# Set a version newer than all upgrade steps
store.set_haiku_version("100.0.0")
store.close()
async with Store(temp_db_path, create=True) as store:
# Set a version newer than all upgrade steps
await store.set_haiku_version("100.0.0")
# Read-only open should work (version not updated, but no error)
store = Store(temp_db_path, read_only=True)
# Version should stay at the old value (can't update in read-only)
assert store.get_haiku_version() == "100.0.0"
store.close()
async with Store(temp_db_path, read_only=True) as store:
# Version should stay at the old value (can't update in read-only)
assert await store.get_haiku_version() == "100.0.0"
def test_skip_migration_check_bypasses_error(self, temp_db_path):
@pytest.mark.asyncio
async def test_skip_migration_check_bypasses_error(self, temp_db_path):
"""skip_migration_check=True should bypass migration error."""
store = Store(temp_db_path, create=True)
store.set_haiku_version("0.19.0")
store.close()
async with Store(temp_db_path, create=True) as store:
await store.set_haiku_version("0.19.0")
# Open with skip_migration_check should work
store = Store(temp_db_path, skip_migration_check=True)
# Version should remain old (no auto-migration)
assert store.get_haiku_version() == "0.19.0"
store.close()
async with Store(temp_db_path, skip_migration_check=True) as store:
# Version should remain old (no auto-migration)
assert await store.get_haiku_version() == "0.19.0"
class TestMigrateMethod:
def test_migrate_applies_pending_upgrades(self, temp_db_path):
@pytest.mark.asyncio
async def test_migrate_applies_pending_upgrades(self, temp_db_path):
"""Store.migrate() should apply pending upgrades and update version."""
store = Store(temp_db_path, create=True)
store.set_haiku_version("0.19.0")
store.close()
async with Store(temp_db_path, create=True) as store:
await store.set_haiku_version("0.19.0")
# Open with skip_migration_check to avoid error
store = Store(temp_db_path, skip_migration_check=True)
old_version = store.get_haiku_version()
assert old_version == "0.19.0"
async with Store(temp_db_path, skip_migration_check=True) as store:
old_version = await store.get_haiku_version()
assert old_version == "0.19.0"
# Run migration
applied = store.migrate()
# Run migration
applied = await store.migrate()
# Should have applied migrations
assert len(applied) > 0
# Should have applied migrations
assert len(applied) > 0
# Version should be updated
new_version = store.get_haiku_version()
expected = metadata.version("haiku.rag-slim")
assert new_version == expected
store.close()
# Version should be updated
new_version = await store.get_haiku_version()
expected = metadata.version("haiku.rag-slim")
assert new_version == expected
def test_migrate_returns_applied_upgrades(self, temp_db_path):
@pytest.mark.asyncio
async def test_migrate_returns_applied_upgrades(self, temp_db_path):
"""Store.migrate() should return list of applied upgrade descriptions."""
store = Store(temp_db_path, create=True)
store.set_haiku_version("0.19.0")
store.close()
async with Store(temp_db_path, create=True) as store:
await store.set_haiku_version("0.19.0")
store = Store(temp_db_path, skip_migration_check=True)
applied = store.migrate()
async with Store(temp_db_path, skip_migration_check=True) as store:
applied = await store.migrate()
# Should return descriptions of applied upgrades
assert isinstance(applied, list)
for item in applied:
assert isinstance(item, str)
store.close()
# Should return descriptions of applied upgrades
assert isinstance(applied, list)
for item in applied:
assert isinstance(item, str)
def test_migrate_with_no_pending_returns_empty(self, temp_db_path):
@pytest.mark.asyncio
async def test_migrate_with_no_pending_returns_empty(self, temp_db_path):
"""Store.migrate() with no pending migrations returns empty list."""
store = Store(temp_db_path, create=True)
# Already at current version
store.close()
async with Store(temp_db_path, create=True) as store:
# Already at current version
pass
store = Store(temp_db_path, skip_migration_check=True)
applied = store.migrate()
assert applied == []
store.close()
async with Store(temp_db_path, skip_migration_check=True) as store:
applied = await store.migrate()
assert applied == []
def test_migrate_raises_read_only_error(self, temp_db_path):
@pytest.mark.asyncio
async def test_migrate_raises_read_only_error(self, temp_db_path):
"""Store.migrate() should raise ReadOnlyError in read-only mode."""
from haiku.rag.store.exceptions import ReadOnlyError
store = Store(temp_db_path, create=True)
store.set_haiku_version("0.19.0")
store.close()
async with Store(temp_db_path, create=True) as store:
await store.set_haiku_version("0.19.0")
store = Store(temp_db_path, skip_migration_check=True, read_only=True)
with pytest.raises(ReadOnlyError):
store.migrate()
store.close()
async with Store(
temp_db_path, skip_migration_check=True, read_only=True
) as store:
with pytest.raises(ReadOnlyError):
await store.migrate()
class TestGetPendingUpgrades:

268
tests/store/test_read_only.py
View file

@ -28,7 +28,8 @@ class TestReadOnlyError:
class TestStoreReadOnly:
def test_store_read_only_raises_on_empty_directory(self, tmp_path):
@pytest.mark.asyncio
async def test_store_read_only_raises_on_empty_directory(self, tmp_path):
"""Opening an empty directory in read-only mode raises ReadOnlyError."""
empty_dir = tmp_path / "empty_db"
empty_dir.mkdir()
@ -36,155 +37,148 @@ class TestStoreReadOnly:
with pytest.raises(
ReadOnlyError, match="Cannot create tables in read-only mode"
):
Store(
async with Store(
empty_dir,
read_only=True,
skip_validation=True,
skip_migration_check=True,
)
):
pass
def test_store_default_is_not_read_only(self, temp_db_path):
@pytest.mark.asyncio
async def test_store_default_is_not_read_only(self, temp_db_path):
"""Store defaults to not read-only."""
store = Store(temp_db_path, create=True)
assert store.is_read_only is False
store.close()
async with Store(temp_db_path, create=True) as store:
assert store.is_read_only is False
def test_store_can_be_created_read_only(self, temp_db_path):
@pytest.mark.asyncio
async def test_store_can_be_created_read_only(self, temp_db_path):
"""Store can be created with read_only=True."""
# First create a normal store to initialize the database
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
# Now open in read-only mode
store = Store(temp_db_path, read_only=True)
assert store.is_read_only is True
store.close()
async with Store(temp_db_path, read_only=True) as store:
assert store.is_read_only is True
def test_assert_writable_raises_when_read_only(self, temp_db_path):
@pytest.mark.asyncio
async def test_assert_writable_raises_when_read_only(self, temp_db_path):
"""_assert_writable() raises ReadOnlyError when read_only=True."""
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
store = Store(temp_db_path, read_only=True)
with pytest.raises(ReadOnlyError):
store._assert_writable()
store.close()
async with Store(temp_db_path, read_only=True) as store:
with pytest.raises(ReadOnlyError):
store._assert_writable()
def test_assert_writable_passes_when_not_read_only(self, temp_db_path):
@pytest.mark.asyncio
async def test_assert_writable_passes_when_not_read_only(self, temp_db_path):
"""_assert_writable() does not raise when read_only=False."""
store = Store(temp_db_path, create=True)
store._assert_writable() # Should not raise
store.close()
async with Store(temp_db_path, create=True) as store:
store._assert_writable() # Should not raise
@pytest.mark.asyncio
async def test_vacuum_raises_when_read_only(self, temp_db_path):
"""vacuum() raises ReadOnlyError when read_only=True."""
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
store = Store(temp_db_path, read_only=True)
with pytest.raises(ReadOnlyError):
await store.vacuum()
store.close()
async with Store(temp_db_path, read_only=True) as store:
with pytest.raises(ReadOnlyError):
await store.vacuum()
def test_set_haiku_version_raises_when_read_only(self, temp_db_path):
@pytest.mark.asyncio
async def test_set_haiku_version_raises_when_read_only(self, temp_db_path):
"""set_haiku_version() raises ReadOnlyError when read_only=True."""
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
store = Store(temp_db_path, read_only=True)
with pytest.raises(ReadOnlyError):
store.set_haiku_version("1.0.0")
store.close()
async with Store(temp_db_path, read_only=True) as store:
with pytest.raises(ReadOnlyError):
await store.set_haiku_version("1.0.0")
def test_recreate_embeddings_table_raises_when_read_only(self, temp_db_path):
@pytest.mark.asyncio
async def test_recreate_embeddings_table_raises_when_read_only(self, temp_db_path):
"""recreate_embeddings_table() raises ReadOnlyError when read_only=True."""
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
store = Store(temp_db_path, read_only=True)
with pytest.raises(ReadOnlyError):
store.recreate_embeddings_table()
store.close()
async with Store(temp_db_path, read_only=True) as store:
with pytest.raises(ReadOnlyError):
await store.recreate_embeddings_table()
def test_restore_table_versions_raises_when_read_only(self, temp_db_path):
@pytest.mark.asyncio
async def test_restore_table_versions_raises_when_read_only(self, temp_db_path):
"""restore_table_versions() raises ReadOnlyError when read_only=True."""
store = Store(temp_db_path, create=True)
versions = store.current_table_versions()
store.close()
async with Store(temp_db_path, create=True) as store:
versions = await store.current_table_versions()
store = Store(temp_db_path, read_only=True)
with pytest.raises(ReadOnlyError):
store.restore_table_versions(versions)
store.close()
async with Store(temp_db_path, read_only=True) as store:
with pytest.raises(ReadOnlyError):
await store.restore_table_versions(versions)
class TestDocumentRepositoryReadOnly:
@pytest.mark.asyncio
async def test_create_raises_when_read_only(self, temp_db_path):
"""DocumentRepository.create() raises ReadOnlyError when read_only=True."""
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
store = Store(temp_db_path, read_only=True)
repo = DocumentRepository(store)
doc = Document(content="test content")
async with Store(temp_db_path, read_only=True) as store:
repo = DocumentRepository(store)
doc = Document(content="test content")
with pytest.raises(ReadOnlyError):
await repo.create(doc)
store.close()
with pytest.raises(ReadOnlyError):
await repo.create(doc)
@pytest.mark.asyncio
async def test_update_raises_when_read_only(self, temp_db_path):
"""DocumentRepository.update() raises ReadOnlyError when read_only=True."""
# First create a document
store = Store(temp_db_path, create=True)
repo = DocumentRepository(store)
doc = Document(content="test content")
created_doc = await repo.create(doc)
store.close()
async with Store(temp_db_path, create=True) as store:
repo = DocumentRepository(store)
doc = Document(content="test content")
created_doc = await repo.create(doc)
# Try to update in read-only mode
store = Store(temp_db_path, read_only=True)
repo = DocumentRepository(store)
created_doc.content = "updated content"
async with Store(temp_db_path, read_only=True) as store:
repo = DocumentRepository(store)
created_doc.content = "updated content"
with pytest.raises(ReadOnlyError):
await repo.update(created_doc)
store.close()
with pytest.raises(ReadOnlyError):
await repo.update(created_doc)
@pytest.mark.asyncio
async def test_delete_raises_when_read_only(self, temp_db_path):
"""DocumentRepository.delete() raises ReadOnlyError when read_only=True."""
# First create a document
store = Store(temp_db_path, create=True)
repo = DocumentRepository(store)
doc = Document(content="test content")
created_doc = await repo.create(doc)
assert created_doc.id is not None
doc_id = created_doc.id
store.close()
async with Store(temp_db_path, create=True) as store:
repo = DocumentRepository(store)
doc = Document(content="test content")
created_doc = await repo.create(doc)
assert created_doc.id is not None
doc_id = created_doc.id
# Try to delete in read-only mode
store = Store(temp_db_path, read_only=True)
repo = DocumentRepository(store)
async with Store(temp_db_path, read_only=True) as store:
repo = DocumentRepository(store)
with pytest.raises(ReadOnlyError):
await repo.delete(doc_id)
store.close()
with pytest.raises(ReadOnlyError):
await repo.delete(doc_id)
@pytest.mark.asyncio
async def test_delete_all_raises_when_read_only(self, temp_db_path):
"""DocumentRepository.delete_all() raises ReadOnlyError when read_only=True."""
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
store = Store(temp_db_path, read_only=True)
repo = DocumentRepository(store)
async with Store(temp_db_path, read_only=True) as store:
repo = DocumentRepository(store)
with pytest.raises(ReadOnlyError):
await repo.delete_all()
store.close()
with pytest.raises(ReadOnlyError):
await repo.delete_all()
class TestChunkRepositoryReadOnly:
@ -192,86 +186,82 @@ class TestChunkRepositoryReadOnly:
async def test_create_raises_when_read_only(self, temp_db_path):
"""ChunkRepository.create() raises ReadOnlyError when read_only=True."""
# First create a document to have a valid document_id
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
doc = Document(content="test content")
created_doc = await doc_repo.create(doc)
store.close()
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
doc = Document(content="test content")
created_doc = await doc_repo.create(doc)
store = Store(temp_db_path, read_only=True)
repo = ChunkRepository(store)
chunk = Chunk(
content="test chunk",
document_id=created_doc.id,
embedding=[0.0] * store.embedder._vector_dim,
)
async with Store(temp_db_path, read_only=True) as store:
repo = ChunkRepository(store)
chunk = Chunk(
content="test chunk",
document_id=created_doc.id,
embedding=[0.0] * store.embedder._vector_dim,
)
with pytest.raises(ReadOnlyError):
await repo.create(chunk)
store.close()
with pytest.raises(ReadOnlyError):
await repo.create(chunk)
@pytest.mark.asyncio
async def test_delete_by_document_id_raises_when_read_only(self, temp_db_path):
"""ChunkRepository.delete_by_document_id() raises ReadOnlyError when read_only=True."""
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
store = Store(temp_db_path, read_only=True)
repo = ChunkRepository(store)
async with Store(temp_db_path, read_only=True) as store:
repo = ChunkRepository(store)
with pytest.raises(ReadOnlyError):
await repo.delete_by_document_id("some-id")
store.close()
with pytest.raises(ReadOnlyError):
await repo.delete_by_document_id("some-id")
@pytest.mark.asyncio
async def test_delete_all_raises_when_read_only(self, temp_db_path):
"""ChunkRepository.delete_all() raises ReadOnlyError when read_only=True."""
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
store = Store(temp_db_path, read_only=True)
repo = ChunkRepository(store)
async with Store(temp_db_path, read_only=True) as store:
repo = ChunkRepository(store)
with pytest.raises(ReadOnlyError):
await repo.delete_all()
store.close()
with pytest.raises(ReadOnlyError):
await repo.delete_all()
class TestSettingsRepositoryReadOnly:
def test_save_current_settings_raises_when_read_only(self, temp_db_path):
@pytest.mark.asyncio
async def test_save_current_settings_raises_when_read_only(self, temp_db_path):
"""SettingsRepository.save_current_settings() raises ReadOnlyError when read_only=True."""
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
store = Store(temp_db_path, read_only=True)
repo = SettingsRepository(store)
async with Store(temp_db_path, read_only=True) as store:
repo = SettingsRepository(store)
with pytest.raises(ReadOnlyError):
repo.save_current_settings()
store.close()
with pytest.raises(ReadOnlyError):
await repo.save_current_settings()
class TestClientReadOnly:
def test_client_default_is_not_read_only(self, temp_db_path):
@pytest.mark.asyncio
async def test_client_default_is_not_read_only(self, temp_db_path):
"""Client defaults to not read-only."""
client = HaikuRAG(temp_db_path, create=True)
assert client.is_read_only is False
client.close()
async with HaikuRAG(temp_db_path, create=True) as client:
assert client.is_read_only is False
def test_client_can_be_created_read_only(self, temp_db_path):
@pytest.mark.asyncio
async def test_client_can_be_created_read_only(self, temp_db_path):
"""Client can be created with read_only=True."""
client = HaikuRAG(temp_db_path, create=True)
client.close()
async with HaikuRAG(temp_db_path, create=True):
pass
client = HaikuRAG(temp_db_path, read_only=True)
assert client.is_read_only is True
client.close()
async with HaikuRAG(temp_db_path, read_only=True) as client:
assert client.is_read_only is True
@pytest.mark.vcr()
async def test_client_create_document_raises_when_read_only(self, temp_db_path):
"""Client.create_document() raises ReadOnlyError when read_only=True."""
client = HaikuRAG(temp_db_path, create=True)
client.close()
async with HaikuRAG(temp_db_path, create=True):
pass
async with HaikuRAG(temp_db_path, read_only=True) as client:
with pytest.raises(ReadOnlyError):

124
tests/store/test_time_travel.py
View file

@ -9,104 +9,88 @@ from haiku.rag.store.repositories.document import DocumentRepository
class TestStoreTimeTravel:
def test_store_with_before_is_read_only(self, temp_db_path):
@pytest.mark.asyncio
async def test_store_with_before_is_read_only(self, temp_db_path):
"""Store with before parameter is automatically read-only."""
# Create a store first
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
# Open with before - should be read-only
before = datetime.now(UTC) + timedelta(hours=1)
store = Store(temp_db_path, before=before)
assert store.is_read_only is True
store.close()
async with Store(temp_db_path, before=before) as store:
assert store.is_read_only is True
def test_store_before_raises_on_write(self, temp_db_path):
@pytest.mark.asyncio
async def test_store_before_raises_on_write(self, temp_db_path):
"""Store with before parameter raises on write operations."""
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
before = datetime.now(UTC) + timedelta(hours=1)
store = Store(temp_db_path, before=before)
with pytest.raises(ReadOnlyError):
store._assert_writable()
store.close()
async with Store(temp_db_path, before=before) as store:
with pytest.raises(ReadOnlyError):
store._assert_writable()
@pytest.mark.asyncio
async def test_store_before_checks_out_historical_state(self, temp_db_path):
"""Store with before parameter checks out tables to historical state."""
# Create store and add a document
store = Store(temp_db_path, create=True)
repo = DocumentRepository(store)
await repo.create(Document(content="First document"))
async with Store(temp_db_path, create=True) as store:
repo = DocumentRepository(store)
await repo.create(Document(content="First document"))
# Get the version timestamp after first document
versions_after_first = store.list_table_versions("documents")
# Find the latest version timestamp
latest_version = max(versions_after_first, key=lambda v: v["version"])
time_after_first = latest_version["timestamp"]
versions_after_first = await store.list_table_versions("documents")
latest_version = max(versions_after_first, key=lambda v: v["version"])
time_after_first = latest_version["timestamp"]
# Wait a bit to ensure the next write gets a distinct timestamp
await asyncio.sleep(0.5)
await asyncio.sleep(0.5)
# Add second document
await repo.create(Document(content="Second document"))
await repo.create(Document(content="Second document"))
# Verify we have more versions now
versions_after_second = store.list_table_versions("documents")
assert len(versions_after_second) > len(versions_after_first)
versions_after_second = await store.list_table_versions("documents")
assert len(versions_after_second) > len(versions_after_first)
store.close()
async with Store(temp_db_path, before=time_after_first) as store:
repo = DocumentRepository(store)
# Open at historical state (using the timestamp from after first write)
store = Store(temp_db_path, before=time_after_first)
repo = DocumentRepository(store)
docs = await repo.list_all(include_content=True)
assert len(docs) == 1
assert docs[0].content == "First document"
# Should only see first document
docs = await repo.list_all(include_content=True)
assert len(docs) == 1
assert docs[0].content == "First document"
store.close()
async with Store(temp_db_path) as store:
repo = DocumentRepository(store)
# Open at current state
store = Store(temp_db_path)
repo = DocumentRepository(store)
docs = await repo.list_all()
assert len(docs) == 2
# Should see both documents
docs = await repo.list_all()
assert len(docs) == 2
store.close()
def test_store_before_no_version_raises(self, temp_db_path):
@pytest.mark.asyncio
async def test_store_before_no_version_raises(self, temp_db_path):
"""Store with before datetime before any version raises ValueError."""
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
# Try to open before the database was created
before = datetime(2000, 1, 1, tzinfo=UTC)
with pytest.raises(ValueError) as exc_info:
Store(temp_db_path, before=before)
async with Store(temp_db_path, before=before):
pass
assert "No data exists before" in str(exc_info.value)
def test_current_table_versions_returns_versions(self, temp_db_path):
@pytest.mark.asyncio
async def test_current_table_versions_returns_versions(self, temp_db_path):
"""current_table_versions returns dict of table versions."""
store = Store(temp_db_path, create=True)
versions = store.current_table_versions()
async with Store(temp_db_path, create=True) as store:
versions = await store.current_table_versions()
assert "documents" in versions
assert "chunks" in versions
assert "settings" in versions
assert all(isinstance(v, int) for v in versions.values())
store.close()
assert "documents" in versions
assert "chunks" in versions
assert "settings" in versions
assert all(isinstance(v, int) for v in versions.values())
def test_list_table_versions_returns_history(self, temp_db_path):
@pytest.mark.asyncio
async def test_list_table_versions_returns_history(self, temp_db_path):
"""list_table_versions returns version history for a table."""
store = Store(temp_db_path, create=True)
versions = store.list_table_versions("documents")
async with Store(temp_db_path, create=True) as store:
versions = await store.list_table_versions("documents")
assert len(versions) >= 1
for v in versions:
assert "version" in v
assert "timestamp" in v
store.close()
assert len(versions) >= 1
for v in versions:
assert "version" in v
assert "timestamp" in v

74
tests/test_chunk.py
View file

@ -9,42 +9,40 @@ from haiku.rag.store.models.chunk import Chunk, ChunkMetadata, SearchResult
@pytest.mark.vcr()
async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path):
"""Test ChunkRepository operations."""
# Create client
client = HaikuRAG(db_path=temp_db_path, config=Config, create=True)
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Get the first document from the corpus
first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"]
# Get the first document from the corpus
first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"]
# Create a document first with chunks
created_document = await client.create_document(
content=document_text, metadata={"source": "test"}
)
assert created_document.id is not None
# Create a document first with chunks
created_document = await client.create_document(
content=document_text, metadata={"source": "test"}
)
assert created_document.id is not None
# Test getting chunks by document ID
chunks = await client.chunk_repository.get_by_document_id(created_document.id)
assert len(chunks) > 0
assert all(chunk.document_id == created_document.id for chunk in chunks)
# Test getting chunks by document ID
chunks = await client.chunk_repository.get_by_document_id(created_document.id)
assert len(chunks) > 0
assert all(chunk.document_id == created_document.id for chunk in chunks)
# Test chunk search
results = await client.chunk_repository.search(
"election", limit=2, search_type="vector"
)
assert len(results) <= 2
assert all(hasattr(chunk, "content") for chunk, _ in results)
# Test chunk search
results = await client.chunk_repository.search(
"election", limit=2, search_type="vector"
)
assert len(results) <= 2
assert all(hasattr(chunk, "content") for chunk, _ in results)
# Test deleting chunks by document ID
deleted = await client.chunk_repository.delete_by_document_id(
created_document.id
)
assert deleted is True
# Test deleting chunks by document ID
deleted = await client.chunk_repository.delete_by_document_id(created_document.id)
assert deleted is True
# Verify chunks are gone
chunks_after_delete = await client.chunk_repository.get_by_document_id(
created_document.id
)
assert len(chunks_after_delete) == 0
client.close()
# Verify chunks are gone
chunks_after_delete = await client.chunk_repository.get_by_document_id(
created_document.id
)
assert len(chunks_after_delete) == 0
@pytest.mark.vcr()
@ -409,13 +407,12 @@ async def test_chunk_content_fts_populated(temp_db_path):
await client.chunk_repository.create(chunk)
# Read the raw record from the database
records = list(
client.store.chunks_table.search()
records = (
await client.store.chunks_table.query()
.where(f"id = '{chunk.id}'")
.limit(1)
.to_arrow()
.to_pylist()
)
).to_pylist()
assert len(records) == 1
record = records[0]
@ -453,13 +450,12 @@ async def test_chunk_content_fts_without_headings(temp_db_path):
await client.chunk_repository.create(chunk)
# Read the raw record from the database
records = list(
client.store.chunks_table.search()
records = (
await client.store.chunks_table.query()
.where(f"id = '{chunk.id}'")
.limit(1)
.to_arrow()
.to_pylist()
)
).to_pylist()
assert len(records) == 1
record = records[0]

97
tests/test_client.py
View file

@ -481,49 +481,35 @@ async def test_client_create_document_from_url_http_error(temp_db_path):
)
@pytest.mark.vcr()
async def test_get_extension_from_content_type_or_url(temp_db_path):
"""Test the helper method for determining file extensions."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Test content type mappings
assert (
client._get_extension_from_content_type_or_url("", "text/html") == ".html"
)
assert (
client._get_extension_from_content_type_or_url("", "application/pdf")
== ".pdf"
)
assert (
client._get_extension_from_content_type_or_url("", "text/plain") == ".txt"
)
def test_get_extension_from_content_type_or_url():
"""Test the helper function for determining file extensions."""
from haiku.rag.client.processing import get_extension_from_content_type_or_url
# Test URL extension detection
assert (
client._get_extension_from_content_type_or_url(
"https://example.com/doc.pdf", ""
)
== ".pdf"
)
assert (
client._get_extension_from_content_type_or_url(
"https://example.com/data.json", ""
)
== ".json"
)
# Content type mappings
assert get_extension_from_content_type_or_url("", "text/html") == ".html"
assert get_extension_from_content_type_or_url("", "application/pdf") == ".pdf"
assert get_extension_from_content_type_or_url("", "text/plain") == ".txt"
# Test default fallback
assert (
client._get_extension_from_content_type_or_url("https://example.com/", "")
== ".html"
)
# URL extension detection
assert (
get_extension_from_content_type_or_url("https://example.com/doc.pdf", "")
== ".pdf"
)
assert (
get_extension_from_content_type_or_url("https://example.com/data.json", "")
== ".json"
)
# Test content type priority over URL extension
assert (
client._get_extension_from_content_type_or_url(
"https://example.com/file.txt", "application/pdf"
)
== ".pdf"
# Default fallback
assert get_extension_from_content_type_or_url("https://example.com/", "") == ".html"
# Content type priority over URL extension
assert (
get_extension_from_content_type_or_url(
"https://example.com/file.txt", "application/pdf"
)
== ".pdf"
)
@pytest.mark.vcr()
@ -1286,6 +1272,39 @@ async def test_client_convert_file_not_found(temp_db_path):
await client.convert(Path("/nonexistent/path/file.txt"))
async def test_client_convert_from_url(temp_db_path):
"""convert() with an http(s) URL downloads to a tempfile and converts."""
from docling_core.types.doc.document import DoclingDocument
async with HaikuRAG(temp_db_path, create=True) as client:
mock_response = AsyncMock()
mock_response.content = (
b"<html><body><p>URL convert path content.</p></body></html>"
)
mock_response.headers = {"content-type": "text/html"}
mock_response.raise_for_status = AsyncMock()
with patch("httpx.AsyncClient.get", return_value=mock_response):
docling_doc = await client.convert("https://example.com/page.html")
assert isinstance(docling_doc, DoclingDocument)
markdown = docling_doc.export_to_markdown()
assert "URL convert path content" in markdown
async def test_client_convert_from_url_unsupported_content_type(temp_db_path):
"""convert() rejects URLs whose content type isn't supported by the converter."""
async with HaikuRAG(temp_db_path, create=True) as client:
mock_response = AsyncMock()
mock_response.content = b"\x00\x01\x02binary"
mock_response.headers = {"content-type": "application/octet-stream"}
mock_response.raise_for_status = AsyncMock()
with patch("httpx.AsyncClient.get", return_value=mock_response):
with pytest.raises(ValueError, match="Unsupported content type"):
await client.convert("https://example.com/blob.bin")
@pytest.mark.vcr()
async def test_client_convert_unsupported_extension(temp_db_path):
"""Test convert() raises ValueError for unsupported file extension."""

5
tests/test_context.py
View file

@ -1,5 +1,6 @@
import pytest
from haiku.rag.client.documents import _store_document_with_chunks
from haiku.rag.context import (
_expand_outward,
_find_expansion_range,
@ -249,7 +250,8 @@ class TestExpandWithItems:
from haiku.rag.store.models.document import Document
async with HaikuRAG(temp_db_path, create=True) as rag:
doc = await rag._store_document_with_chunks(
doc = await _store_document_with_chunks(
rag,
Document(content="test"),
[],
__import__(
@ -262,6 +264,7 @@ class TestExpandWithItems:
document_id=doc.id,
doc_item_refs=["#/texts/999999"],
)
assert doc.id is not None
expanded = await expand_with_items(
rag.document_item_repository, doc.id, [result], 5000
)

6
tests/test_context_enhancement.py
View file

@ -3,6 +3,8 @@ from docling_core.types.doc.document import DoclingDocument, TableData
from docling_core.types.doc.labels import DocItemLabel
from haiku.rag.client import HaikuRAG
from haiku.rag.client.documents import _store_document_with_chunks
from haiku.rag.client.processing import ensure_chunks_embedded
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import SearchResult
@ -12,7 +14,7 @@ async def create_document_with_docling(
):
"""Helper to create a document from a DoclingDocument using import_document."""
chunks = await client.chunk(docling_doc)
embedded_chunks = await client._ensure_chunks_embedded(chunks)
embedded_chunks = await ensure_chunks_embedded(client._config, chunks)
return await client.import_document(
docling_document=docling_doc,
chunks=embedded_chunks,
@ -314,7 +316,7 @@ async def test_expand_context_single_item_document(temp_db_path):
async with HaikuRAG(temp_db_path, create=True) as client:
document = Document(content="Simple test content")
document.set_docling(docling_doc)
doc = await client._store_document_with_chunks(document, [], docling_doc)
doc = await _store_document_with_chunks(client, document, [], docling_doc)
assert doc.id is not None
# Create a search result with a doc_item_ref pointing to the item

52
tests/test_database_autocreate.py
View file

@ -1,51 +1,45 @@
import tempfile
from pathlib import Path
import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig
def test_database_not_created_without_create_flag():
async def test_database_not_created_without_create_flag(tmp_path):
"""Test that database is not created without create=True."""
with tempfile.TemporaryDirectory() as tmpdir:
db_path = Path(tmpdir) / "test.lancedb"
db_path = tmp_path / "test.lancedb"
config = AppConfig()
config = AppConfig()
with pytest.raises(FileNotFoundError, match="Database does not exist"):
HaikuRAG(db_path=db_path, config=config)
with pytest.raises(FileNotFoundError, match="Database does not exist"):
async with HaikuRAG(db_path=db_path, config=config):
pass
def test_database_created_with_create_flag():
async def test_database_created_with_create_flag(tmp_path):
"""Test that database is created with create=True."""
with tempfile.TemporaryDirectory() as tmpdir:
db_path = Path(tmpdir) / "test.lancedb"
db_path = tmp_path / "test.lancedb"
config = AppConfig()
config = AppConfig()
client = HaikuRAG(db_path=db_path, config=config, create=True)
async with HaikuRAG(db_path=db_path, config=config, create=True):
assert db_path.exists()
client.close()
@pytest.mark.vcr()
async def test_operations_work_after_database_created():
async def test_operations_work_after_database_created(tmp_path):
"""Test that operations work after DB is created."""
with tempfile.TemporaryDirectory() as tmpdir:
db_path = Path(tmpdir) / "test.lancedb"
db_path = tmp_path / "test.lancedb"
config = AppConfig()
config = AppConfig()
# First, create DB with create=True and add document
async with HaikuRAG(db_path=db_path, config=config, create=True) as client:
await client.create_document("Test content", uri="test://doc1")
# First, create DB with create=True and add document
async with HaikuRAG(db_path=db_path, config=config, create=True) as client:
await client.create_document("Test content", uri="test://doc1")
# Re-open without create flag and verify we can read the document
async with HaikuRAG(db_path=db_path, config=config) as client:
docs = await client.list_documents()
assert len(docs) == 1
doc = await client.get_document_by_id(docs[0].id)
assert doc is not None
assert doc.content == "Test content"
# Re-open without create flag and verify we can read the document
async with HaikuRAG(db_path=db_path, config=config) as client:
docs = await client.list_documents()
assert len(docs) == 1
doc = await client.get_document_by_id(docs[0].id)
assert doc is not None
assert doc.content == "Test content"

277
tests/test_document.py
View file

@ -11,27 +11,25 @@ async def test_document_list_excludes_content_by_default(
qa_corpus: Dataset, temp_db_path
):
"""list_all excludes content and docling_document by default."""
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.txt",
title="Test Document",
metadata={"key": "value"},
)
created = await doc_repo.create(doc)
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.txt",
title="Test Document",
metadata={"key": "value"},
)
created = await doc_repo.create(doc)
docs = await doc_repo.list_all()
assert len(docs) == 1
assert docs[0].id == created.id
assert docs[0].title == "Test Document"
assert docs[0].uri == "https://example.com/doc.txt"
assert docs[0].metadata == {"key": "value"}
assert docs[0].content == ""
assert docs[0].docling_document is None
store.close()
docs = await doc_repo.list_all()
assert len(docs) == 1
assert docs[0].id == created.id
assert docs[0].title == "Test Document"
assert docs[0].uri == "https://example.com/doc.txt"
assert docs[0].metadata == {"key": "value"}
assert docs[0].content == ""
assert docs[0].docling_document is None
@pytest.mark.asyncio
@ -39,62 +37,61 @@ async def test_document_list_includes_content_when_requested(
qa_corpus: Dataset, temp_db_path
):
"""list_all returns content when include_content=True."""
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
content = qa_corpus[0]["document_extracted"]
doc = Document(content=content, uri="https://example.com/doc.txt")
created = await doc_repo.create(doc)
content = qa_corpus[0]["document_extracted"]
doc = Document(content=content, uri="https://example.com/doc.txt")
created = await doc_repo.create(doc)
docs = await doc_repo.list_all(include_content=True)
assert len(docs) == 1
assert docs[0].id == created.id
assert docs[0].content == content
store.close()
docs = await doc_repo.list_all(include_content=True)
assert len(docs) == 1
assert docs[0].id == created.id
assert docs[0].content == content
@pytest.mark.asyncio
async def test_document_list_with_filter(qa_corpus: Dataset, temp_db_path):
"""Test listing documents with filter clause."""
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"]
first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"]
doc1 = Document(
content=document_text,
uri="https://example.com/doc1.txt",
metadata={"source": "test", "category": "A"},
)
doc2 = Document(
content=document_text,
uri="https://arxiv.org/paper.pdf",
metadata={"source": "test", "category": "B"},
)
doc3 = Document(
content=document_text,
uri="https://example.com/doc3.txt",
metadata={"source": "test", "category": "A"},
)
doc1 = Document(
content=document_text,
uri="https://example.com/doc1.txt",
metadata={"source": "test", "category": "A"},
)
doc2 = Document(
content=document_text,
uri="https://arxiv.org/paper.pdf",
metadata={"source": "test", "category": "B"},
)
doc3 = Document(
content=document_text,
uri="https://example.com/doc3.txt",
metadata={"source": "test", "category": "A"},
)
created_doc1 = await doc_repo.create(doc1)
created_doc2 = await doc_repo.create(doc2)
created_doc3 = await doc_repo.create(doc3)
created_doc1 = await doc_repo.create(doc1)
created_doc2 = await doc_repo.create(doc2)
created_doc3 = await doc_repo.create(doc3)
all_documents = await doc_repo.list_all()
assert len(all_documents) == 3
all_documents = await doc_repo.list_all()
assert len(all_documents) == 3
arxiv_documents = await doc_repo.list_all(filter="uri LIKE '%arxiv%'")
assert len(arxiv_documents) == 1
assert arxiv_documents[0].id == created_doc2.id
arxiv_documents = await doc_repo.list_all(filter="uri LIKE '%arxiv%'")
assert len(arxiv_documents) == 1
assert arxiv_documents[0].id == created_doc2.id
example_documents = await doc_repo.list_all(filter="uri LIKE '%example.com%'")
assert len(example_documents) == 2
assert {doc.id for doc in example_documents} == {created_doc1.id, created_doc3.id}
store.close()
example_documents = await doc_repo.list_all(filter="uri LIKE '%example.com%'")
assert len(example_documents) == 2
assert {doc.id for doc in example_documents} == {
created_doc1.id,
created_doc3.id,
}
def test_document_get_docling_document():
@ -239,45 +236,43 @@ async def test_get_docling_data_loads_only_docling_columns(
from haiku.rag.store.compression import compress_json
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
doc_json = {
"name": "test_doc",
"texts": [],
"tables": [],
"pictures": [],
"groups": [],
"body": {"self_ref": "#/body", "children": []},
"furniture": {"self_ref": "#/furniture", "children": []},
}
compressed = compress_json(json.dumps(doc_json))
doc_json = {
"name": "test_doc",
"texts": [],
"tables": [],
"pictures": [],
"groups": [],
"body": {"self_ref": "#/body", "children": []},
"furniture": {"self_ref": "#/furniture", "children": []},
}
compressed = compress_json(json.dumps(doc_json))
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.txt",
docling_document=compressed,
docling_version="2.1.0",
)
created = await doc_repo.create(doc)
assert created.id is not None
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.txt",
docling_document=compressed,
docling_version="2.1.0",
)
created = await doc_repo.create(doc)
assert created.id is not None
result = await doc_repo.get_docling_data(created.id)
assert result is not None
assert result.id == created.id
assert result.content == ""
assert result.docling_document == compressed
assert result.docling_version == "2.1.0"
result = await doc_repo.get_docling_data(created.id)
assert result is not None
assert result.id == created.id
assert result.content == ""
assert result.docling_document == compressed
assert result.docling_version == "2.1.0"
# Verify docling document can be parsed
docling_doc = result.get_docling_document()
assert docling_doc is not None
assert docling_doc.name == "test_doc"
# Verify docling document can be parsed
docling_doc = result.get_docling_document()
assert docling_doc is not None
assert docling_doc.name == "test_doc"
# Non-existent ID returns None
assert await doc_repo.get_docling_data("nonexistent-id") is None
store.close()
# Non-existent ID returns None
assert await doc_repo.get_docling_data("nonexistent-id") is None
@pytest.mark.asyncio
@ -291,27 +286,25 @@ async def test_get_pages_data_loads_only_pages_column(qa_corpus: Dataset, temp_d
json.dumps({"1": {"size": {"width": 612, "height": 792}, "page_no": 1}})
)
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.txt",
docling_pages=pages_blob,
)
created = await doc_repo.create(doc)
assert created.id is not None
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.txt",
docling_pages=pages_blob,
)
created = await doc_repo.create(doc)
assert created.id is not None
result = await doc_repo.get_pages_data(created.id)
assert result is not None
assert result.id == created.id
assert result.content == ""
assert result.docling_pages == pages_blob
result = await doc_repo.get_pages_data(created.id)
assert result is not None
assert result.id == created.id
assert result.content == ""
assert result.docling_pages == pages_blob
# Non-existent ID returns None
assert await doc_repo.get_pages_data("nonexistent-id") is None
store.close()
# Non-existent ID returns None
assert await doc_repo.get_pages_data("nonexistent-id") is None
@pytest.mark.asyncio
@ -319,22 +312,20 @@ async def test_get_pages_data_none_for_markdown_document(
qa_corpus: Dataset, temp_db_path
):
"""Markdown documents have no page images — get_pages_data returns None pages."""
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.md",
)
created = await doc_repo.create(doc)
assert created.id is not None
doc = Document(
content=qa_corpus[0]["document_extracted"],
uri="https://example.com/doc.md",
)
created = await doc_repo.create(doc)
assert created.id is not None
result = await doc_repo.get_pages_data(created.id)
assert result is not None
assert result.id == created.id
assert result.docling_pages is None
store.close()
result = await doc_repo.get_pages_data(created.id)
assert result is not None
assert result.id == created.id
assert result.docling_pages is None
@pytest.mark.asyncio
@ -342,23 +333,21 @@ async def test_document_get_by_uri_with_special_characters(
qa_corpus: Dataset, temp_db_path
):
"""Test get_by_uri handles URIs with special characters like single quotes."""
store = Store(temp_db_path, create=True)
doc_repo = DocumentRepository(store)
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"]
first_doc = qa_corpus[0]
document_text = first_doc["document_extracted"]
doc_with_quote = Document(
content=document_text,
uri="Hamish and Andy's Gap Year",
metadata={"source": "test"},
)
doc_with_quote = Document(
content=document_text,
uri="Hamish and Andy's Gap Year",
metadata={"source": "test"},
)
created_doc = await doc_repo.create(doc_with_quote)
created_doc = await doc_repo.create(doc_with_quote)
retrieved = await doc_repo.get_by_uri("Hamish and Andy's Gap Year")
assert retrieved is not None
assert retrieved.id == created_doc.id
assert retrieved.uri == "Hamish and Andy's Gap Year"
store.close()
retrieved = await doc_repo.get_by_uri("Hamish and Andy's Gap Year")
assert retrieved is not None
assert retrieved.id == created_doc.id
assert retrieved.uri == "Hamish and Andy's Gap Year"

124
tests/test_download_models.py
View file

@ -4,13 +4,14 @@ from unittest.mock import AsyncMock, patch
import httpx
import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.client.downloads import download_models
from haiku.rag.config import Config
@pytest.fixture
def mock_to_thread():
"""Patch asyncio.to_thread to skip docling/tokenizer downloads."""
with patch("haiku.rag.client.asyncio.to_thread", new_callable=AsyncMock):
with patch("haiku.rag.client.downloads.asyncio.to_thread", new_callable=AsyncMock):
yield
@ -22,79 +23,77 @@ async def _mock_httpx_client(stream_fn):
yield mock_client
async def test_download_models_ollama_connect_error(temp_db_path, mock_to_thread):
async def test_download_models_ollama_connect_error(mock_to_thread):
"""When Ollama is not running, download_models raises ConnectionError."""
async with HaikuRAG(temp_db_path, create=True) as client:
@asynccontextmanager
async def failing_stream(method, url, **kwargs):
raise httpx.ConnectError("All connection attempts failed")
yield # unreachable, but needed for generator syntax
@asynccontextmanager
async def failing_stream(method, url, **kwargs):
raise httpx.ConnectError("All connection attempts failed")
yield # unreachable, but needed for generator syntax
with patch(
"haiku.rag.client.httpx.AsyncClient",
return_value=_mock_httpx_client(failing_stream),
):
with pytest.raises(
ConnectionError, match="Cannot connect to Ollama"
) as exc_info:
async for _ in client.download_models():
pass
with patch(
"haiku.rag.client.downloads.httpx.AsyncClient",
return_value=_mock_httpx_client(failing_stream),
):
with pytest.raises(
ConnectionError, match="Cannot connect to Ollama"
) as exc_info:
async for _ in download_models(Config):
pass
assert "ollama serve" in str(exc_info.value)
assert "ollama serve" in str(exc_info.value)
async def test_download_models_ollama_pulls_models(temp_db_path, mock_to_thread):
async def test_download_models_ollama_pulls_models(mock_to_thread):
"""download_models yields correct progress events for Ollama model pulls."""
async with HaikuRAG(temp_db_path, create=True) as client:
stream_lines = [
'{"status": "pulling manifest"}',
"",
'{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 500}',
'{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 1000}',
"not valid json",
'{"status": "verifying sha256 digest"}',
'{"status": "writing manifest"}',
'{"status": "success"}',
]
stream_lines = [
'{"status": "pulling manifest"}',
"",
'{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 500}',
'{"status": "downloading", "digest": "sha256:abc", "total": 1000, "completed": 1000}',
"not valid json",
'{"status": "verifying sha256 digest"}',
'{"status": "writing manifest"}',
'{"status": "success"}',
]
@asynccontextmanager
async def mock_stream(method, url, **kwargs):
mock_resp = AsyncMock()
@asynccontextmanager
async def mock_stream(method, url, **kwargs):
mock_resp = AsyncMock()
async def aiter_lines():
for line in stream_lines:
yield line
async def aiter_lines():
for line in stream_lines:
yield line
mock_resp.aiter_lines = aiter_lines
yield mock_resp
mock_resp.aiter_lines = aiter_lines
yield mock_resp
with patch(
"haiku.rag.client.httpx.AsyncClient",
return_value=_mock_httpx_client(mock_stream),
):
events = []
async for progress in client.download_models():
events.append(progress)
with patch(
"haiku.rag.client.downloads.httpx.AsyncClient",
return_value=_mock_httpx_client(mock_stream),
):
events = []
async for progress in download_models(Config):
events.append(progress)
# Default config has embeddings=qwen3-embedding:4b, qa/research=gpt-oss
ollama_models = {"gpt-oss", "qwen3-embedding:4b"}
ollama_events = [e for e in events if e.model in ollama_models]
pulling_events = [e for e in ollama_events if e.status == "pulling"]
done_events = [e for e in ollama_events if e.status == "done"]
download_events = [e for e in ollama_events if e.status == "downloading"]
# Default config has embeddings=qwen3-embedding:4b, qa/research=gpt-oss
ollama_models = {"gpt-oss", "qwen3-embedding:4b"}
ollama_events = [e for e in events if e.model in ollama_models]
pulling_events = [e for e in ollama_events if e.status == "pulling"]
done_events = [e for e in ollama_events if e.status == "done"]
download_events = [e for e in ollama_events if e.status == "downloading"]
assert len(pulling_events) == 2
assert len(done_events) == 2
assert len(download_events) > 0
assert len(pulling_events) == 2
assert len(done_events) == 2
assert len(download_events) > 0
for de in download_events:
assert de.digest == "sha256:abc"
assert de.total == 1000
assert de.completed > 0
for de in download_events:
assert de.digest == "sha256:abc"
assert de.total == 1000
assert de.completed > 0
async def test_download_models_no_ollama_models(temp_db_path, mock_to_thread):
async def test_download_models_no_ollama_models(mock_to_thread):
"""When no Ollama models are configured, no Ollama pull events are yielded."""
from haiku.rag.config import AppConfig
@ -103,10 +102,9 @@ async def test_download_models_no_ollama_models(temp_db_path, mock_to_thread):
config.qa.model.provider = "openai"
config.research.model.provider = "openai"
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
events = []
async for progress in client.download_models():
events.append(progress)
events = []
async for progress in download_models(config):
events.append(progress)
models = {e.model for e in events}
assert "qwen3-embedding:4b" not in models

42
tests/test_filter.py
View file

@ -179,3 +179,45 @@ async def test_search_filter_with_all_search_types(temp_db_path):
for result in results:
assert result.document_uri is not None
assert "other.com" in result.document_uri
@pytest.mark.vcr()
async def test_search_with_filter_returns_full_limit(temp_db_path):
"""Regression: filter + limit must return up to `limit` matching chunks
even when non-matching chunks would dominate the top-N window.
Previously the filter path materialized LanceDB's default top-N window
(~10), filtered to matching document_ids in pandas, then took `head(limit)`.
If the top-N window was dominated by non-matching chunks, the caller got
silently fewer results than requested — even when plenty of matching
chunks existed further down the ranking. This test puts the target
document behind many distractor documents and asserts we still get the
requested count back.
"""
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
for i in range(12):
await client.create_document(
content=(
"machine learning neural network deep learning model "
"machine learning neural network deep learning model "
"machine learning neural network deep learning model"
),
uri=f"https://distractor.com/doc{i}.html",
title=f"Distractor {i}",
)
await client.create_document(
content="one passing mention of machine learning here",
uri="https://target.com/one.html",
title="Target One",
)
results = await client.search(
"machine learning",
limit=5,
search_type="fts",
filter="uri LIKE '%target.com%'",
)
assert len(results) == 1
assert results[0].document_uri == "https://target.com/one.html"

87
tests/test_info.py
View file

@ -1,5 +1,5 @@
import json
from unittest.mock import patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -15,7 +15,7 @@ async def test_app_info_outputs(temp_db_path, capsys):
from lancedb.pydantic import LanceModel, Vector
from pydantic import Field
db = lancedb.connect(temp_db_path)
db = await lancedb.connect_async(temp_db_path)
class SettingsRecord(LanceModel):
id: str = Field(default="settings")
@ -31,13 +31,13 @@ async def test_app_info_outputs(temp_db_path, capsys):
content: str
vector: Vector(3) # type: ignore
settings_tbl = db.create_table("settings", schema=SettingsRecord)
docs_tbl = db.create_table("documents", schema=DocumentRecord)
chunks_tbl = db.create_table("chunks", schema=ChunkRecord)
db.create_table("document_items", schema=DocumentItemRecord)
settings_tbl = await db.create_table("settings", schema=SettingsRecord)
docs_tbl = await db.create_table("documents", schema=DocumentRecord)
chunks_tbl = await db.create_table("chunks", schema=ChunkRecord)
await db.create_table("document_items", schema=DocumentItemRecord)
# Insert one of each - using the new config format
settings_tbl.add(
await settings_tbl.add(
[
SettingsRecord(
id="settings",
@ -56,8 +56,8 @@ async def test_app_info_outputs(temp_db_path, capsys):
)
]
)
docs_tbl.add([DocumentRecord(id="doc-1", content="hello")])
chunks_tbl.add(
await docs_tbl.add([DocumentRecord(id="doc-1", content="hello")])
await chunks_tbl.add(
[ChunkRecord(id="c1", document_id="doc-1", content="c", vector=[0.1, 0.2, 0.3])]
)
@ -68,7 +68,9 @@ async def test_app_info_outputs(temp_db_path, capsys):
# Validate expected content substrings
# Note: Rich console may wrap long paths to new lines, so check separately
assert "path:" in out
assert str(temp_db_path) in out
# Rich may wrap long paths across lines — check with newlines stripped
out_no_wrap = out.replace("\n", "")
assert str(temp_db_path) in out_no_wrap
assert "haiku.rag version (db):" in out
assert "embeddings: openai/text-embedding-3-small (dim: 3)" in out
assert "documents: 1" in out
@ -93,10 +95,11 @@ async def test_app_info_outputs(temp_db_path, capsys):
async def test_app_info_with_vector_index(temp_db_path, capsys):
# Build a database with enough chunks to create a vector index
import lancedb
from lancedb.index import IvfPq
from lancedb.pydantic import LanceModel, Vector
from pydantic import Field
db = lancedb.connect(temp_db_path)
db = await lancedb.connect_async(temp_db_path)
class SettingsRecord(LanceModel):
id: str = Field(default="settings")
@ -112,13 +115,13 @@ async def test_app_info_with_vector_index(temp_db_path, capsys):
content: str
vector: Vector(3) # type: ignore
settings_tbl = db.create_table("settings", schema=SettingsRecord)
docs_tbl = db.create_table("documents", schema=DocumentRecord)
chunks_tbl = db.create_table("chunks", schema=ChunkRecord)
db.create_table("document_items", schema=DocumentItemRecord)
settings_tbl = await db.create_table("settings", schema=SettingsRecord)
docs_tbl = await db.create_table("documents", schema=DocumentRecord)
chunks_tbl = await db.create_table("chunks", schema=ChunkRecord)
await db.create_table("document_items", schema=DocumentItemRecord)
# Insert settings
settings_tbl.add(
await settings_tbl.add(
[
SettingsRecord(
id="settings",
@ -128,7 +131,7 @@ async def test_app_info_with_vector_index(temp_db_path, capsys):
)
# Insert document
docs_tbl.add([DocumentRecord(id="doc-1", content="test")])
await docs_tbl.add([DocumentRecord(id="doc-1", content="test")])
# Insert 512 chunks to allow index creation (PQ needs more than 256 for training)
chunks = [
@ -140,10 +143,10 @@ async def test_app_info_with_vector_index(temp_db_path, capsys):
)
for i in range(512)
]
chunks_tbl.add(chunks)
await chunks_tbl.add(chunks)
# Create vector index
chunks_tbl.create_index(metric="cosine", index_type="IVF_PQ")
await chunks_tbl.create_index("vector", config=IvfPq(distance_type="cosine"))
app = HaikuRAGApp(db_path=temp_db_path)
await app.info()
@ -172,9 +175,14 @@ async def test_app_info_uses_connect_lancedb_for_remote(tmp_path):
)
app = HaikuRAGApp(db_path=nonexistent, config=config)
with patch("haiku.rag.store.engine.connect_lancedb") as mock_connect:
with patch(
"haiku.rag.store.engine.connect_lancedb", new_callable=AsyncMock
) as mock_connect:
# Empty DB triggers the early-return path - enough to prove connect_lancedb was used
mock_connect.return_value.list_tables.return_value.tables = []
mock_db = mock_connect.return_value
mock_list_result = MagicMock()
mock_list_result.tables = []
mock_db.list_tables = AsyncMock(return_value=mock_list_result)
await app.info()
mock_connect.assert_called_once_with(config, nonexistent)
@ -188,7 +196,7 @@ async def test_app_info_with_missing_document_items_table(temp_db_path, capsys):
from lancedb.pydantic import LanceModel, Vector
from pydantic import Field
db = lancedb.connect(temp_db_path)
db = await lancedb.connect_async(temp_db_path)
class SettingsRecord(LanceModel):
id: str = Field(default="settings")
@ -204,12 +212,12 @@ async def test_app_info_with_missing_document_items_table(temp_db_path, capsys):
content: str
vector: Vector(3) # type: ignore
settings_tbl = db.create_table("settings", schema=SettingsRecord)
docs_tbl = db.create_table("documents", schema=DocumentRecord)
chunks_tbl = db.create_table("chunks", schema=ChunkRecord)
settings_tbl = await db.create_table("settings", schema=SettingsRecord)
docs_tbl = await db.create_table("documents", schema=DocumentRecord)
chunks_tbl = await db.create_table("chunks", schema=ChunkRecord)
# Intentionally omit document_items (added in 0.40.0)
settings_tbl.add(
await settings_tbl.add(
[
SettingsRecord(
id="settings",
@ -228,8 +236,8 @@ async def test_app_info_with_missing_document_items_table(temp_db_path, capsys):
)
]
)
docs_tbl.add([DocumentRecord(id="doc-1", content="hello")])
chunks_tbl.add(
await docs_tbl.add([DocumentRecord(id="doc-1", content="hello")])
await chunks_tbl.add(
[ChunkRecord(id="c1", document_id="doc-1", content="c", vector=[0.1, 0.2, 0.3])]
)
@ -261,7 +269,7 @@ async def test_app_info_reports_up_to_date(temp_db_path, capsys):
from lancedb.pydantic import LanceModel, Vector
from pydantic import Field
db = lancedb.connect(temp_db_path)
db = await lancedb.connect_async(temp_db_path)
class SettingsRecord(LanceModel):
id: str = Field(default="settings")
@ -277,13 +285,13 @@ async def test_app_info_reports_up_to_date(temp_db_path, capsys):
content: str
vector: Vector(3) # type: ignore
settings_tbl = db.create_table("settings", schema=SettingsRecord)
db.create_table("documents", schema=DocumentRecord)
db.create_table("chunks", schema=ChunkRecord)
db.create_table("document_items", schema=DocumentItemRecord)
settings_tbl = await db.create_table("settings", schema=SettingsRecord)
await db.create_table("documents", schema=DocumentRecord)
await db.create_table("chunks", schema=ChunkRecord)
await db.create_table("document_items", schema=DocumentItemRecord)
current_version = metadata.version("haiku.rag-slim")
settings_tbl.add(
await settings_tbl.add(
[
SettingsRecord(
id="settings",
@ -324,6 +332,9 @@ async def test_app_init_skips_exists_check_for_remote(tmp_path):
app = HaikuRAGApp(db_path=nonexistent, config=config)
with patch("haiku.rag.app.HaikuRAG") as mock_client_cls:
mock_client = AsyncMock()
mock_client_cls.return_value.__aenter__ = AsyncMock(return_value=mock_client)
mock_client_cls.return_value.__aexit__ = AsyncMock(return_value=False)
await app.init()
# Should have called HaikuRAG to create, not returned early
mock_client_cls.assert_called_once()
@ -342,9 +353,9 @@ async def test_app_history_skips_exists_check_for_remote(tmp_path):
app = HaikuRAGApp(db_path=nonexistent, config=config)
with patch("haiku.rag.store.engine.Store") as mock_store_cls:
mock_store = mock_store_cls.return_value
mock_store.documents_table.list_versions.return_value = []
mock_store.chunks_table.list_versions.return_value = []
mock_store.settings_table.list_versions.return_value = []
mock_store = AsyncMock()
mock_store.list_table_versions = AsyncMock(return_value=[])
mock_store_cls.return_value.__aenter__ = AsyncMock(return_value=mock_store)
mock_store_cls.return_value.__aexit__ = AsyncMock(return_value=False)
await app.history()
mock_store_cls.assert_called_once()

260
tests/test_lancedb_connection.py
View file

@ -1,4 +1,4 @@
from unittest.mock import patch
from unittest.mock import AsyncMock, patch
import pytest
@ -42,25 +42,32 @@ class TestConnectionMode:
class TestConnectLancedb:
def test_local_passes_db_path(self, temp_db_path):
@pytest.mark.asyncio
async def test_local_passes_db_path(self, temp_db_path):
config = AppConfig(lancedb=LanceDBConfig(uri=""))
with patch("haiku.rag.store.engine.lancedb.connect") as mock_connect:
connect_lancedb(config, db_path=temp_db_path)
with patch(
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config, db_path=temp_db_path)
mock_connect.assert_called_once_with(temp_db_path)
def test_cloud_passes_uri_api_key_region(self):
@pytest.mark.asyncio
async def test_cloud_passes_uri_api_key_region(self):
config = AppConfig(
lancedb=LanceDBConfig(
uri="db://my-database", api_key="test-key", region="us-west-2"
)
)
with patch("haiku.rag.store.engine.lancedb.connect") as mock_connect:
connect_lancedb(config)
with patch(
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config)
mock_connect.assert_called_once_with(
uri="db://my-database", api_key="test-key", region="us-west-2"
)
def test_object_storage_passes_uri_and_storage_options(self):
@pytest.mark.asyncio
async def test_object_storage_passes_uri_and_storage_options(self):
config = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path",
@ -70,8 +77,10 @@ class TestConnectLancedb:
},
)
)
with patch("haiku.rag.store.engine.lancedb.connect") as mock_connect:
connect_lancedb(config)
with patch(
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config)
mock_connect.assert_called_once_with(
uri="s3://bucket/path",
storage_options={
@ -80,120 +89,136 @@ class TestConnectLancedb:
},
)
def test_object_storage_without_storage_options(self):
@pytest.mark.asyncio
async def test_object_storage_without_storage_options(self):
config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
with patch("haiku.rag.store.engine.lancedb.connect") as mock_connect:
connect_lancedb(config)
with patch(
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
) as mock_connect:
await connect_lancedb(config)
mock_connect.assert_called_once_with(uri="s3://bucket/path")
def test_local_without_db_path_raises(self):
@pytest.mark.asyncio
async def test_local_without_db_path_raises(self):
config = AppConfig(lancedb=LanceDBConfig(uri=""))
with pytest.raises(
ValueError, match="No lancedb.uri configured and no db_path provided"
):
connect_lancedb(config)
await connect_lancedb(config)
class TestStoreConnectionMode:
def test_store_connection_mode_local(self, temp_db_path):
store = Store(temp_db_path, create=True)
assert store._connection_mode == ConnectionMode.LOCAL
store.close()
@pytest.mark.asyncio
async def test_store_connection_mode_local(self, temp_db_path):
async with Store(temp_db_path, create=True) as store:
assert store._connection_mode == ConnectionMode.LOCAL
def test_store_connection_mode_cloud(self, temp_db_path):
store = Store(temp_db_path, create=True)
with (
patch.object(Config.lancedb, "uri", "db://test-database"),
patch.object(Config.lancedb, "api_key", "test-api-key"),
patch.object(Config.lancedb, "region", "us-east-1"),
):
assert store._connection_mode == ConnectionMode.CLOUD
store.close()
@pytest.mark.asyncio
async def test_store_connection_mode_cloud(self, temp_db_path):
async with Store(temp_db_path, create=True) as store:
with (
patch.object(Config.lancedb, "uri", "db://test-database"),
patch.object(Config.lancedb, "api_key", "test-api-key"),
patch.object(Config.lancedb, "region", "us-east-1"),
):
assert store._connection_mode == ConnectionMode.CLOUD
def test_store_connection_mode_object_storage(self, temp_db_path):
store = Store(temp_db_path, create=True)
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
assert store._connection_mode == ConnectionMode.OBJECT_STORAGE
store.close()
@pytest.mark.asyncio
async def test_store_connection_mode_object_storage(self, temp_db_path):
async with Store(temp_db_path, create=True) as store:
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
assert store._connection_mode == ConnectionMode.OBJECT_STORAGE
class TestVacuumByConnectionMode:
@pytest.mark.asyncio
async def test_cloud_skips_vacuum(self, temp_db_path):
store = Store(temp_db_path, create=True)
with (
patch.object(Config.lancedb, "uri", "db://test-database"),
patch.object(Config.lancedb, "api_key", "test-api-key"),
patch.object(Config.lancedb, "region", "us-east-1"),
):
with patch.object(store.chunks_table, "optimize") as mock_optimize:
await store.vacuum()
mock_optimize.assert_not_called()
store.close()
async with Store(temp_db_path, create=True) as store:
with (
patch.object(Config.lancedb, "uri", "db://test-database"),
patch.object(Config.lancedb, "api_key", "test-api-key"),
patch.object(Config.lancedb, "region", "us-east-1"),
):
with patch.object(
store.chunks_table, "optimize", new_callable=AsyncMock
) as mock_optimize:
await store.vacuum()
mock_optimize.assert_not_called()
@pytest.mark.asyncio
async def test_object_storage_runs_vacuum(self, temp_db_path):
store = Store(temp_db_path, create=True)
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
with patch.object(store.chunks_table, "optimize") as mock_optimize:
await store.vacuum()
mock_optimize.assert_called()
store.close()
async with Store(temp_db_path, create=True) as store:
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
with patch.object(
store.chunks_table, "optimize", new_callable=AsyncMock
) as mock_optimize:
await store.vacuum()
mock_optimize.assert_called()
@pytest.mark.asyncio
async def test_local_runs_vacuum(self, temp_db_path):
store = Store(temp_db_path, create=True)
with patch.object(Config.lancedb, "uri", ""):
with patch.object(store.chunks_table, "optimize") as mock_optimize:
await store.vacuum()
mock_optimize.assert_called()
store.close()
async with Store(temp_db_path, create=True) as store:
with patch.object(Config.lancedb, "uri", ""):
with patch.object(
store.chunks_table, "optimize", new_callable=AsyncMock
) as mock_optimize:
await store.vacuum()
mock_optimize.assert_called()
class TestVectorIndexByConnectionMode:
def test_cloud_skips_index_creation(self, temp_db_path):
store = Store(temp_db_path, create=True)
with (
patch.object(Config.lancedb, "uri", "db://test-database"),
patch.object(Config.lancedb, "api_key", "test-api-key"),
patch.object(Config.lancedb, "region", "us-east-1"),
):
with patch.object(store.chunks_table, "count_rows") as mock_count:
store._ensure_vector_index()
mock_count.assert_not_called()
store.close()
@pytest.mark.asyncio
async def test_cloud_skips_index_creation(self, temp_db_path):
async with Store(temp_db_path, create=True) as store:
with (
patch.object(Config.lancedb, "uri", "db://test-database"),
patch.object(Config.lancedb, "api_key", "test-api-key"),
patch.object(Config.lancedb, "region", "us-east-1"),
):
with patch.object(
store.chunks_table, "count_rows", new_callable=AsyncMock
) as mock_count:
await store._ensure_vector_index()
mock_count.assert_not_called()
def test_object_storage_runs_index_creation(self, temp_db_path):
store = Store(temp_db_path, create=True)
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
with patch.object(
store.chunks_table, "count_rows", return_value=0
) as mock_count:
store._ensure_vector_index()
mock_count.assert_called()
store.close()
@pytest.mark.asyncio
async def test_object_storage_runs_index_creation(self, temp_db_path):
async with Store(temp_db_path, create=True) as store:
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
with patch.object(
store.chunks_table,
"count_rows",
new_callable=AsyncMock,
return_value=0,
) as mock_count:
await store._ensure_vector_index()
mock_count.assert_called()
class TestStoreSkipsPathValidationForRemote:
def test_skips_path_check_for_cloud(self, tmp_path):
@pytest.mark.asyncio
async def test_skips_path_check_for_cloud(self, tmp_path):
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
config = AppConfig(
lancedb=LanceDBConfig(
uri="db://test-database", api_key="key", region="us-east-1"
)
)
with patch("haiku.rag.store.engine.lancedb.connect"):
with patch.object(Store, "_init_tables"):
store = Store(
with patch(
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
):
with patch.object(Store, "_init_tables", new_callable=AsyncMock):
async with Store(
nonexistent,
config=config,
create=True,
skip_validation=True,
skip_migration_check=True,
)
store.close()
) as store:
assert store is not None
def test_skips_path_check_for_object_storage(self, tmp_path):
@pytest.mark.asyncio
async def test_skips_path_check_for_object_storage(self, tmp_path):
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
config = AppConfig(
lancedb=LanceDBConfig(
@ -201,13 +226,74 @@ class TestStoreSkipsPathValidationForRemote:
storage_options={"endpoint": "http://localhost:9000"},
)
)
with patch("haiku.rag.store.engine.lancedb.connect"):
with patch.object(Store, "_init_tables"):
store = Store(
with patch(
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
):
with patch.object(Store, "_init_tables", new_callable=AsyncMock):
async with Store(
nonexistent,
config=config,
create=True,
skip_validation=True,
skip_migration_check=True,
)
store.close()
) as store:
assert store is not None
class TestInitFailureCleanup:
@pytest.mark.asyncio
async def test_store_aenter_closes_connection_on_init_failure(
self, temp_db_path, monkeypatch
):
"""If _initialize raises after connect, __aenter__ must close the
AsyncConnection so it doesn't leak (no __aexit__ runs in that case)."""
mock_conn = AsyncMock()
mock_conn.close = lambda: mock_conn.close_calls.append(True) # type: ignore[attr-defined]
mock_conn.close_calls = [] # type: ignore[attr-defined]
async def fake_connect(*args, **kwargs):
return mock_conn
async def failing_init_tables(self):
raise RuntimeError("simulated table init failure")
monkeypatch.setattr("haiku.rag.store.engine.connect_lancedb", fake_connect)
monkeypatch.setattr(Store, "_init_tables", failing_init_tables)
with pytest.raises(RuntimeError, match="simulated table init failure"):
async with Store(temp_db_path, create=True) as store:
assert store is not None
assert mock_conn.close_calls == [True], (
"AsyncConnection.close() was not called on init failure"
)
@pytest.mark.asyncio
async def test_client_aenter_closes_store_on_init_failure(
self, temp_db_path, monkeypatch
):
"""HaikuRAG.__aenter__ must close the store if _initialize fails."""
from haiku.rag.client import HaikuRAG
close_calls: list[bool] = []
original_close = Store.close
def tracking_close(self):
close_calls.append(True)
original_close(self)
async def failing_init(self):
# Set db so close() has something to close
self.db = AsyncMock()
self.db.close = lambda: None
raise RuntimeError("simulated initialize failure")
monkeypatch.setattr(Store, "_initialize", failing_init)
monkeypatch.setattr(Store, "close", tracking_close)
with pytest.raises(RuntimeError, match="simulated initialize failure"):
async with HaikuRAG(temp_db_path, create=True):
pass
assert close_calls, "Store.close() was not called when _initialize raised"

194
tests/test_rebuild.py
View file

@ -1,3 +1,5 @@
import tempfile
from pathlib import Path
from typing import TypedDict
import pytest
@ -89,8 +91,8 @@ async def test_rebuild_embed_only_skips_unchanged(qa_corpus: Dataset, temp_db_pa
assert doc.id is not None
# Get embeddings before rebuild
records_before = list(
client.store.chunks_table.search()
records_before = await (
client.store.chunks_table.query()
.where(f"document_id = '{doc.id}'")
.to_pydantic(client.store.ChunkRecord)
)
@ -104,8 +106,8 @@ async def test_rebuild_embed_only_skips_unchanged(qa_corpus: Dataset, temp_db_pa
assert doc.id in processed_ids
# Get embeddings after rebuild
records_after = list(
client.store.chunks_table.search()
records_after = await (
client.store.chunks_table.query()
.where(f"document_id = '{doc.id}'")
.to_pydantic(client.store.ChunkRecord)
)
@ -156,7 +158,7 @@ async def test_rebuild_embed_only_with_changed_vector_dim(
]
# Step 2: Manually recreate chunks table with 4096-dim vectors (simulating old DB)
db = lancedb.connect(temp_db_path)
db = await lancedb.connect_async(temp_db_path)
class ChunkRecord4096(LanceModel):
id: str
@ -167,8 +169,8 @@ async def test_rebuild_embed_only_with_changed_vector_dim(
order: int = Field(default=0)
vector: Vector(4096) = Field(default_factory=lambda: [0.0] * 4096) # type: ignore
db.drop_table("chunks")
chunks_table = db.create_table("chunks", schema=ChunkRecord4096)
await db.drop_table("chunks")
chunks_table = await db.create_table("chunks", schema=ChunkRecord4096)
# Insert chunks with 4096-dim fake vectors
records_4096 = [
@ -183,19 +185,20 @@ async def test_rebuild_embed_only_with_changed_vector_dim(
)
for c in chunk_data
]
chunks_table.add(records_4096)
await chunks_table.add(records_4096)
# Update settings to reflect the 4096-dim model used
settings_table = db.open_table("settings")
settings_table = await db.open_table("settings")
rows = (
settings_table.search().where("id = 'settings'").limit(1).to_arrow().to_pylist()
)
await settings_table.query().where("id = 'settings'").limit(1).to_arrow()
).to_pylist()
settings = json.loads(rows[0]["settings"])
settings["embeddings"]["model"]["vector_dim"] = 4096
settings["embeddings"]["model"]["name"] = "qwen3-embedding:8b"
settings_table.update(
where="id = 'settings'", values={"settings": json.dumps(settings)}
await settings_table.update(
{"settings": json.dumps(settings)}, where="id = 'settings'"
)
db.close()
# Step 3: Open with skip_validation (different config) and run embed-only rebuild
# This should work: Store should use stored vector_dim for reading,
@ -213,11 +216,10 @@ async def test_rebuild_embed_only_with_changed_vector_dim(
# Check that embeddings in DB are now 2560-dim
raw_chunks = (
client.store.chunks_table.search()
await client.store.chunks_table.query()
.where(f"document_id = '{doc.id}'")
.to_arrow()
.to_pylist()
)
).to_pylist()
for raw_chunk in raw_chunks:
assert len(raw_chunk["vector"]) == 2560
@ -263,3 +265,163 @@ async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path):
# Chunk IDs should change (chunks are recreated)
assert chunk_ids_before.isdisjoint(chunk_ids_after)
@pytest.mark.vcr()
async def test_rebuild_full_with_accessible_source(temp_db_path):
"""FULL rebuild re-ingests from source when the URI is accessible.
Covers the main path in _rebuild_full (source-accessible branch): the
document is deleted and re-created from its URI, producing a new ID.
"""
async with HaikuRAG(temp_db_path, create=True) as client:
with tempfile.TemporaryDirectory() as temp_dir:
source_path = Path(temp_dir) / "source.txt"
source_path.write_text("Fresh content from an accessible file source.")
original = await client.create_document_from_source(source=source_path)
assert not isinstance(original, list)
assert original.id is not None
original_id = original.id
processed_ids = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)
]
# Original doc was deleted and a new one created; the old ID
# must not appear, and exactly one new ID must have been yielded.
assert original_id not in processed_ids
assert len(processed_ids) == 1
new_doc = await client.get_document_by_id(processed_ids[0])
assert new_doc is not None
assert new_doc.uri == source_path.as_uri()
assert "Fresh content" in new_doc.content
async def test_rebuild_title_only_handles_llm_failure(temp_db_path, monkeypatch):
"""TITLE_ONLY: a failure on one document does not abort the generator.
The first document raises during title generation (simulated LLM error);
the second succeeds. Rebuild must log-and-skip the failure, yield only
the successful document, and persist its new title.
"""
from haiku.rag.store.models.document import Document
async with HaikuRAG(temp_db_path, create=True) as client:
# Skip embedding — TITLE_ONLY only touches documents.
doc1 = await client.document_repository.create(
Document(content="doc one body", metadata={})
)
doc2 = await client.document_repository.create(
Document(content="doc two body", metadata={})
)
assert doc1.id is not None and doc2.id is not None
async def fake_generate_title(doc):
if doc.id == doc1.id:
raise RuntimeError("simulated LLM failure")
return "Second Title"
monkeypatch.setattr(client, "generate_title", fake_generate_title)
processed_ids = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY)
]
assert processed_ids == [doc2.id]
refreshed = await client.get_document_by_id(doc2.id)
assert refreshed is not None
assert refreshed.title == "Second Title"
untouched = await client.get_document_by_id(doc1.id)
assert untouched is not None
assert untouched.title is None
@pytest.mark.vcr()
async def test_rebuild_full_source_failure_is_logged_and_skipped(
temp_db_path, monkeypatch
):
"""FULL rebuild logs-and-continues when re-ingesting from source fails.
Covers _rebuild_full's `except Exception` branch: when
create_document_from_source raises, the doc is skipped (no yield) and
the error is logged. Regression guard against silent failures.
"""
import logging
from haiku.rag.client import rebuild as rebuild_module
async with HaikuRAG(temp_db_path, create=True) as client:
with tempfile.TemporaryDirectory() as temp_dir:
source_path = Path(temp_dir) / "source.txt"
source_path.write_text("Content that will vanish by rebuild time.")
original = await client.create_document_from_source(source=source_path)
assert not isinstance(original, list)
assert original.id is not None
# Force the source rebuild branch to raise.
async def failing_create(*args, **kwargs):
raise RuntimeError("simulated ingestion failure")
monkeypatch.setattr(client, "create_document_from_source", failing_create)
# Attach directly to the rebuild module's logger rather than
# relying on caplog — `haiku.rag.logging.get_logger()` (invoked
# by other tests) sets `propagate=False` on the `haiku.rag`
# logger, which breaks caplog under xdist ordering.
records: list[logging.LogRecord] = []
class _ListHandler(logging.Handler):
def emit(self, record: logging.LogRecord) -> None:
records.append(record)
handler = _ListHandler(level=logging.ERROR)
rebuild_module.logger.addHandler(handler)
try:
processed_ids = [
doc_id
async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)
]
finally:
rebuild_module.logger.removeHandler(handler)
assert processed_ids == []
assert any(
"Error recreating document from source" in rec.getMessage()
for rec in records
)
@pytest.mark.vcr()
async def test_rebuild_batch_size_flush(temp_db_path, monkeypatch):
"""RECHUNK flushes in batches and yields every document.
Forces a tiny batch size so three docs trigger at least one mid-loop
flush plus the final flush. Regression guard for the batched-write path
in _rebuild_rechunk.
"""
from haiku.rag.client import rebuild as rebuild_module
monkeypatch.setattr(rebuild_module, "_REBUILD_BATCH_SIZE", 2)
async with HaikuRAG(temp_db_path, create=True) as client:
ids: list[str] = []
for i in range(3):
doc = await client.create_document(content=f"batch flush doc {i}")
assert doc.id is not None
ids.append(doc.id)
processed = [
doc_id async for doc_id in client.rebuild_database(mode=RebuildMode.RECHUNK)
]
assert sorted(processed) == sorted(ids)
for doc_id in ids:
chunks = await client.chunk_repository.get_by_document_id(doc_id)
assert len(chunks) > 0

32
tests/test_s3_integration.py
View file

@ -51,37 +51,35 @@ def _make_config() -> AppConfig:
)
def test_store_connect_and_create(tmp_path):
@pytest.mark.asyncio
async def test_store_connect_and_create(tmp_path):
from haiku.rag.store.engine import get_database_stats
config = _make_config()
store = Store(tmp_path / "unused", config=config, create=True)
stats = get_database_stats(store.db)
assert stats["documents"]["exists"]
assert stats["chunks"]["exists"]
store.close()
async with Store(tmp_path / "unused", config=config, create=True) as store:
stats = await get_database_stats(store.db)
assert stats["documents"]["exists"]
assert stats["chunks"]["exists"]
@pytest.mark.asyncio
async def test_store_vacuum(tmp_path):
config = _make_config()
store = Store(tmp_path / "unused", config=config, create=True)
await store.vacuum()
store.close()
async with Store(tmp_path / "unused", config=config, create=True) as store:
await store.vacuum()
def test_store_add_document(tmp_path):
@pytest.mark.asyncio
async def test_store_add_document(tmp_path):
from haiku.rag.store.engine import DocumentRecord, get_database_stats
config = _make_config()
store = Store(tmp_path / "unused", config=config, create=True)
async with Store(tmp_path / "unused", config=config, create=True) as store:
doc = DocumentRecord(content="The quick brown fox jumps over the lazy dog.")
await store.documents_table.add([doc])
doc = DocumentRecord(content="The quick brown fox jumps over the lazy dog.")
store.documents_table.add([doc])
stats = get_database_stats(store.db)
assert stats["documents"]["num_rows"] == 1
store.close()
stats = await get_database_stats(store.db)
assert stats["documents"]["num_rows"] == 1
@pytest.mark.asyncio

356
tests/test_search.py
View file

@ -9,187 +9,179 @@ from haiku.rag.store.models import SearchResult
@pytest.mark.vcr()
async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
"""Test that documents can be found by searching with their associated questions."""
# Create client
client = HaikuRAG(db_path=temp_db_path, config=Config, create=True)
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Load unique documents (limited to 10)
seen_documents = set()
documents = []
# Load unique documents (limited to 10)
seen_documents = set()
documents = []
for doc_data in qa_corpus:
if len(seen_documents) >= 10:
break
document_text = doc_data["document_extracted"]
document_id = doc_data.get("document_id", "")
for doc_data in qa_corpus:
if len(seen_documents) >= 10:
break
document_text = doc_data["document_extracted"]
document_id = doc_data.get("document_id", "")
if document_id in seen_documents:
continue
seen_documents.add(document_id)
if document_id in seen_documents:
continue
seen_documents.add(document_id)
# Create the document with chunks and embeddings
created_document = await client.create_document(content=document_text)
documents.append((created_document, doc_data))
# Create the document with chunks and embeddings
created_document = await client.create_document(content=document_text)
documents.append((created_document, doc_data))
# Test with first few unique documents
# Test with first few unique documents
for target_document, doc_data in documents:
question = doc_data["question"]
for target_document, doc_data in documents:
question = doc_data["question"]
# Test vector search (limit=10 to accommodate different embedding models)
vector_results = await client.chunk_repository.search(
question, limit=10, search_type="vector"
)
target_document_ids = {chunk.document_id for chunk, _ in vector_results}
assert target_document.id in target_document_ids
# Test vector search (limit=10 to accommodate different embedding models)
vector_results = await client.chunk_repository.search(
question, limit=10, search_type="vector"
)
target_document_ids = {chunk.document_id for chunk, _ in vector_results}
assert target_document.id in target_document_ids
# Test FTS search
fts_results = await client.chunk_repository.search(
question, limit=10, search_type="fts"
)
target_document_ids = {chunk.document_id for chunk, _ in fts_results}
assert target_document.id in target_document_ids
# Test FTS search
fts_results = await client.chunk_repository.search(
question, limit=10, search_type="fts"
)
target_document_ids = {chunk.document_id for chunk, _ in fts_results}
assert target_document.id in target_document_ids
# Test hybrid search
hybrid_results = await client.chunk_repository.search(
question, limit=10, search_type="hybrid"
)
target_document_ids = {chunk.document_id for chunk, _ in hybrid_results}
assert target_document.id in target_document_ids
client.close()
# Test hybrid search
hybrid_results = await client.chunk_repository.search(
question, limit=10, search_type="hybrid"
)
target_document_ids = {chunk.document_id for chunk, _ in hybrid_results}
assert target_document.id in target_document_ids
@pytest.mark.vcr()
async def test_search_chunk_includes_document_provenance(temp_db_path):
"""Test that raw chunk search results include document URI, metadata, and ID."""
client = HaikuRAG(db_path=temp_db_path, config=Config, create=True)
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Create a document with URI and metadata but no title
created_document = await client.create_document(
content="This is a test document with some content for searching.",
uri="https://example.com/test.html",
metadata={"title": "Test Document", "author": "Test Author"},
)
# Create a document with URI and metadata but no title
created_document = await client.create_document(
content="This is a test document with some content for searching.",
uri="https://example.com/test.html",
metadata={"title": "Test Document", "author": "Test Author"},
)
# Search for chunks
results = await client.chunk_repository.search(
"test document", limit=1, search_type="hybrid"
)
# Search for chunks
results = await client.chunk_repository.search(
"test document", limit=1, search_type="hybrid"
)
assert len(results) > 0
chunk, score = results[0]
assert len(results) > 0
chunk, score = results[0]
# Test that score is valid
assert isinstance(score, int | float), (
f"Score should be numeric, got {type(score)}"
)
assert score >= 0, f"Score should be non-negative, got {score}"
# Test that score is valid
assert isinstance(score, int | float), f"Score should be numeric, got {type(score)}"
assert score >= 0, f"Score should be non-negative, got {score}"
# Verify the chunk includes document information
assert chunk.document_uri == "https://example.com/test.html"
assert chunk.document_meta == {"title": "Test Document", "author": "Test Author"}
assert chunk.document_id == created_document.id
assert chunk.document_title is None
client.close()
# Verify the chunk includes document information
assert chunk.document_uri == "https://example.com/test.html"
assert chunk.document_meta == {
"title": "Test Document",
"author": "Test Author",
}
assert chunk.document_id == created_document.id
assert chunk.document_title is None
@pytest.mark.vcr()
async def test_search_score_types(temp_db_path):
"""Test that different search types return appropriate score ranges."""
client = HaikuRAG(db_path=temp_db_path, config=Config, create=True)
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Create multiple documents with different content
documents_content = [
"Machine learning algorithms are powerful tools for data analysis and pattern recognition.",
"Deep learning neural networks can process complex datasets and identify hidden patterns.",
"Natural language processing enables computers to understand and generate human text.",
"Computer vision systems can interpret and analyze visual information from images.",
]
# Create multiple documents with different content
documents_content = [
"Machine learning algorithms are powerful tools for data analysis and pattern recognition.",
"Deep learning neural networks can process complex datasets and identify hidden patterns.",
"Natural language processing enables computers to understand and generate human text.",
"Computer vision systems can interpret and analyze visual information from images.",
]
for content in documents_content:
await client.create_document(content=content)
for content in documents_content:
await client.create_document(content=content)
query = "machine learning"
query = "machine learning"
# Test vector search scores (should be converted from distances)
vector_results = await client.chunk_repository.search(
query, limit=3, search_type="vector"
)
assert len(vector_results) > 0
vector_scores = [score for _, score in vector_results]
# Test vector search scores (should be converted from distances)
vector_results = await client.chunk_repository.search(
query, limit=3, search_type="vector"
)
assert len(vector_results) > 0
vector_scores = [score for _, score in vector_results]
# Test FTS search scores (should be native LanceDB FTS scores)
fts_results = await client.chunk_repository.search(
query, limit=3, search_type="fts"
)
assert len(fts_results) > 0
fts_scores = [score for _, score in fts_results]
# Test FTS search scores (should be native LanceDB FTS scores)
fts_results = await client.chunk_repository.search(
query, limit=3, search_type="fts"
)
assert len(fts_results) > 0
fts_scores = [score for _, score in fts_results]
# Test hybrid search scores (should be native LanceDB relevance scores)
hybrid_results = await client.chunk_repository.search(
query, limit=3, search_type="hybrid"
)
assert len(hybrid_results) > 0
hybrid_scores = [score for _, score in hybrid_results]
# Test hybrid search scores (should be native LanceDB relevance scores)
hybrid_results = await client.chunk_repository.search(
query, limit=3, search_type="hybrid"
)
assert len(hybrid_results) > 0
hybrid_scores = [score for _, score in hybrid_results]
# All scores should be numeric and non-negative
for scores, search_type in [
(vector_scores, "vector"),
(fts_scores, "fts"),
(hybrid_scores, "hybrid"),
]:
for score in scores:
assert isinstance(score, int | float), (
f"{search_type} score should be numeric"
)
assert score >= 0, f"{search_type} score should be non-negative"
# All scores should be numeric and non-negative
for scores, search_type in [
(vector_scores, "vector"),
(fts_scores, "fts"),
(hybrid_scores, "hybrid"),
]:
for score in scores:
assert isinstance(score, int | float), (
f"{search_type} score should be numeric"
)
assert score >= 0, f"{search_type} score should be non-negative"
# Vector scores should typically be small (0-1 range due to distance conversion)
assert all(0 <= score <= 1 for score in vector_scores), (
"Vector scores should be in 0-1 range"
)
# Vector scores should typically be small (0-1 range due to distance conversion)
assert all(0 <= score <= 1 for score in vector_scores), (
"Vector scores should be in 0-1 range"
)
# Scores should be sorted in descending order (most relevant first)
for scores, search_type in [
(vector_scores, "vector"),
(fts_scores, "fts"),
(hybrid_scores, "hybrid"),
]:
for i in range(len(scores) - 1):
assert scores[i] >= scores[i + 1], (
f"{search_type} results should be sorted by score descending"
)
client.close()
# Scores should be sorted in descending order (most relevant first)
for scores, search_type in [
(vector_scores, "vector"),
(fts_scores, "fts"),
(hybrid_scores, "hybrid"),
]:
for i in range(len(scores) - 1):
assert scores[i] >= scores[i + 1], (
f"{search_type} results should be sorted by score descending"
)
@pytest.mark.vcr()
async def test_search_returns_search_result(temp_db_path):
"""Test that client.search() returns SearchResult with provenance info."""
client = HaikuRAG(db_path=temp_db_path, config=Config, create=True)
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
await client.create_document(
content="Machine learning models can classify images with high accuracy.",
uri="https://example.com/ml.html",
title="ML Guide",
)
await client.create_document(
content="Machine learning models can classify images with high accuracy.",
uri="https://example.com/ml.html",
title="ML Guide",
)
results = await client.search("machine learning", limit=3)
results = await client.search("machine learning", limit=3)
assert len(results) > 0
result = results[0]
assert isinstance(result, SearchResult)
assert result.content
assert result.score > 0
assert result.document_uri == "https://example.com/ml.html"
assert result.document_title == "ML Guide"
assert result.chunk_id is not None
assert result.document_id is not None
# page_numbers and headings come from chunk metadata
assert isinstance(result.page_numbers, list)
assert isinstance(result.labels, list)
assert len(result.labels) > 0
client.close()
assert len(results) > 0
result = results[0]
assert isinstance(result, SearchResult)
assert result.content
assert result.score > 0
assert result.document_uri == "https://example.com/ml.html"
assert result.document_title == "ML Guide"
assert result.chunk_id is not None
assert result.document_id is not None
# page_numbers and headings come from chunk metadata
assert isinstance(result.page_numbers, list)
assert isinstance(result.labels, list)
assert len(result.labels) > 0
@pytest.mark.vcr()
@ -197,30 +189,27 @@ async def test_search_graceful_degradation(temp_db_path):
"""Test search works when docling data is unavailable."""
from haiku.rag.store.models import Chunk
client = HaikuRAG(db_path=temp_db_path, config=Config, create=True)
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Import document with custom chunks (no docling document)
custom_chunks = [
Chunk(content="Custom chunk without docling metadata", metadata={}),
]
docling_doc = await client.convert("Document with custom chunks")
await client.import_document(
docling_document=docling_doc,
chunks=custom_chunks,
uri="https://example.com/custom.html",
)
# Import document with custom chunks (no docling document)
custom_chunks = [
Chunk(content="Custom chunk without docling metadata", metadata={}),
]
docling_doc = await client.convert("Document with custom chunks")
await client.import_document(
docling_document=docling_doc,
chunks=custom_chunks,
uri="https://example.com/custom.html",
)
results = await client.search("custom chunk", limit=3)
results = await client.search("custom chunk", limit=3)
assert len(results) > 0
result = results[0]
assert isinstance(result, SearchResult)
assert result.content
# Metadata defaults should still work
assert result.page_numbers == []
assert result.labels == []
client.close()
assert len(results) > 0
result = results[0]
assert isinstance(result, SearchResult)
assert result.content
# Metadata defaults should still work
assert result.page_numbers == []
assert result.labels == []
@pytest.mark.vcr()
@ -252,6 +241,41 @@ async def test_search_result_format_includes_metadata(temp_db_path):
assert "machine learning" in formatted.lower()
@pytest.mark.vcr()
async def test_fts_search_targets_content_fts_column(temp_db_path):
"""FTS search must target the content_fts column (where the FTS index
lives and where contextualized heading prefixes end up) — not the raw
content column. Regression guard against upstream default-column changes
in LanceDB's nearest_to_text().
"""
from haiku.rag.store.models.chunk import Chunk
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content="seed", uri="test://doc")
assert doc.id is not None
# Heading-only term — contextualization will prepend headings to the
# body when populating content_fts, so this word ends up ONLY in
# content_fts, not in the content column.
heading_only_term = "zxqvjfoowizardry"
chunk = Chunk(
content="unrelated body text",
document_id=doc.id,
metadata={"headings": [heading_only_term]},
embedding=[0.0] * client.store.embedder._vector_dim,
)
await client.chunk_repository.create(chunk)
# FTS on the heading-only word must match via content_fts.
results = await client.chunk_repository.search(
heading_only_term, limit=5, search_type="fts"
)
assert any(c.content == "unrelated body text" for c, _ in results), (
"FTS did not match a heading-only term — nearest_to_text is not "
"targeting the content_fts column"
)
def test_search_result_primary_label_prioritizes_structural_types():
"""Test _get_primary_label prioritizes structural labels correctly."""
# Table should be prioritized

147
tests/test_settings.py
View file

@ -4,43 +4,42 @@ from haiku.rag.config import AppConfig, Config
from haiku.rag.store.repositories.settings import ConfigMismatchError
def test_settings_table_populated_on_store_init(temp_db_path):
@pytest.mark.asyncio
async def test_settings_table_populated_on_store_init(temp_db_path):
"""Test that settings table is populated with current config when store is initialized."""
from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.settings import SettingsRepository
store = Store(temp_db_path, create=True)
settings_repo = SettingsRepository(store)
async with Store(temp_db_path, create=True) as store:
settings_repo = SettingsRepository(store)
db_settings = settings_repo.get_current_settings()
config_dict = Config.model_dump(mode="json")
db_settings = await settings_repo.get_current_settings()
config_dict = Config.model_dump(mode="json")
# Remove version from db_settings since it's added automatically
db_settings_without_version = {
k: v for k, v in db_settings.items() if k != "version"
}
assert db_settings_without_version == config_dict
store.close()
# Remove version from db_settings since it's added automatically
db_settings_without_version = {
k: v for k, v in db_settings.items() if k != "version"
}
assert db_settings_without_version == config_dict
def test_settings_save_and_retrieve(temp_db_path):
@pytest.mark.asyncio
async def test_settings_save_and_retrieve(temp_db_path):
"""Test saving and retrieving settings after config change."""
from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.settings import SettingsRepository
store = Store(temp_db_path, create=True)
settings_repo = SettingsRepository(store)
async with Store(temp_db_path, create=True) as store:
settings_repo = SettingsRepository(store)
original_chunk_size = Config.processing.chunk_size
Config.processing.chunk_size = 2 * original_chunk_size
original_chunk_size = Config.processing.chunk_size
Config.processing.chunk_size = 2 * original_chunk_size
settings_repo.save_current_settings()
retrieved_settings = settings_repo.get_current_settings()
assert retrieved_settings["processing"]["chunk_size"] == 2 * original_chunk_size
await settings_repo.save_current_settings()
retrieved_settings = await settings_repo.get_current_settings()
assert retrieved_settings["processing"]["chunk_size"] == 2 * original_chunk_size
Config.processing.chunk_size = original_chunk_size
store.close()
Config.processing.chunk_size = original_chunk_size
def test_monitor_filter_patterns_config():
@ -56,103 +55,111 @@ def test_monitor_filter_patterns_config():
class TestValidateConfigCompatibility:
"""Tests for validate_config_compatibility method."""
def test_empty_settings_saves_config(self, temp_db_path):
@pytest.mark.asyncio
async def test_empty_settings_saves_config(self, temp_db_path):
"""When settings row is missing, validation saves current config."""
from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.settings import SettingsRepository
store = Store(temp_db_path, create=True, skip_validation=True)
settings_repo = SettingsRepository(store)
async with Store(temp_db_path, create=True, skip_validation=True) as store:
settings_repo = SettingsRepository(store)
# Clear settings to simulate empty state
store.settings_table.delete("id = 'settings'")
assert settings_repo.get_current_settings() == {}
# Clear settings to simulate empty state
await store.settings_table.delete("id = 'settings'")
assert await settings_repo.get_current_settings() == {}
# Validation should save settings
settings_repo.validate_config_compatibility()
# Validation should save settings
await settings_repo.validate_config_compatibility()
# Now settings should exist
saved = settings_repo.get_current_settings()
assert saved.get("embeddings", {}).get("model", {}).get("provider") is not None
store.close()
# Now settings should exist
saved = await settings_repo.get_current_settings()
assert (
saved.get("embeddings", {}).get("model", {}).get("provider") is not None
)
def test_compatible_config_no_error(self, temp_db_path):
@pytest.mark.asyncio
async def test_compatible_config_no_error(self, temp_db_path):
"""Compatible config does not raise error."""
from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.settings import SettingsRepository
store = Store(temp_db_path, create=True)
settings_repo = SettingsRepository(store)
async with Store(temp_db_path, create=True) as store:
settings_repo = SettingsRepository(store)
# Should not raise - same config
settings_repo.validate_config_compatibility()
store.close()
# Should not raise - same config
await settings_repo.validate_config_compatibility()
def test_provider_mismatch_raises_error(self, temp_db_path):
@pytest.mark.asyncio
async def test_provider_mismatch_raises_error(self, temp_db_path):
"""Different embedding provider raises ConfigMismatchError."""
from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.settings import SettingsRepository
# Create store with default config (ollama)
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
# Create new config with different provider
new_config = AppConfig()
new_config.embeddings.model.provider = "openai"
store2 = Store(temp_db_path, config=new_config, skip_validation=True)
settings_repo = SettingsRepository(store2)
async with Store(
temp_db_path, config=new_config, skip_validation=True
) as store2:
settings_repo = SettingsRepository(store2)
with pytest.raises(ConfigMismatchError) as exc_info:
settings_repo.validate_config_compatibility()
with pytest.raises(ConfigMismatchError) as exc_info:
await settings_repo.validate_config_compatibility()
assert "embedding provider" in str(exc_info.value)
assert "ollama" in str(exc_info.value)
assert "openai" in str(exc_info.value)
store2.close()
assert "embedding provider" in str(exc_info.value)
assert "ollama" in str(exc_info.value)
assert "openai" in str(exc_info.value)
def test_model_mismatch_raises_error(self, temp_db_path):
@pytest.mark.asyncio
async def test_model_mismatch_raises_error(self, temp_db_path):
"""Different embedding model raises ConfigMismatchError."""
from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.settings import SettingsRepository
# Create store with default config
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
# Create new config with different model
new_config = AppConfig()
new_config.embeddings.model.name = "different-model"
store2 = Store(temp_db_path, config=new_config, skip_validation=True)
settings_repo = SettingsRepository(store2)
async with Store(
temp_db_path, config=new_config, skip_validation=True
) as store2:
settings_repo = SettingsRepository(store2)
with pytest.raises(ConfigMismatchError) as exc_info:
settings_repo.validate_config_compatibility()
with pytest.raises(ConfigMismatchError) as exc_info:
await settings_repo.validate_config_compatibility()
assert "embedding model" in str(exc_info.value)
store2.close()
assert "embedding model" in str(exc_info.value)
def test_vector_dim_mismatch_raises_error(self, temp_db_path):
@pytest.mark.asyncio
async def test_vector_dim_mismatch_raises_error(self, temp_db_path):
"""Different vector dimension raises ConfigMismatchError."""
from haiku.rag.store.engine import Store
from haiku.rag.store.repositories.settings import SettingsRepository
# Create store with default config
store = Store(temp_db_path, create=True)
store.close()
async with Store(temp_db_path, create=True):
pass
# Create new config with different vector dimension
new_config = AppConfig()
new_config.embeddings.model.vector_dim = 9999
store2 = Store(temp_db_path, config=new_config, skip_validation=True)
settings_repo = SettingsRepository(store2)
async with Store(
temp_db_path, config=new_config, skip_validation=True
) as store2:
settings_repo = SettingsRepository(store2)
with pytest.raises(ConfigMismatchError) as exc_info:
settings_repo.validate_config_compatibility()
with pytest.raises(ConfigMismatchError) as exc_info:
await settings_repo.validate_config_compatibility()
assert "vector dimension" in str(exc_info.value)
assert "9999" in str(exc_info.value)
store2.close()
assert "vector dimension" in str(exc_info.value)
assert "9999" in str(exc_info.value)

78
tests/test_title_generation.py
View file

@ -5,6 +5,7 @@ from docling_core.types.doc.document import ContentLayer, DoclingDocument
from docling_core.types.doc.labels import DocItemLabel
from haiku.rag.client import HaikuRAG
from haiku.rag.client.titles import extract_structural_title, resolve_title
from haiku.rag.config import AppConfig
from haiku.rag.config.models import ProcessingConfig
from haiku.rag.embeddings import EmbedderWrapper
@ -35,11 +36,7 @@ def mock_embedder(monkeypatch):
class TestExtractStructuralTitle:
def _make_client(self, tmp_path):
config = AppConfig(processing=ProcessingConfig(auto_title=True))
return HaikuRAG(tmp_path / "test.lancedb", config=config, create=True)
def test_furniture_title(self, tmp_path):
def test_furniture_title(self):
"""TITLE on FURNITURE layer (HTML <title>) is extracted."""
doc = DoclingDocument(name="test")
doc.add_text(
@ -49,11 +46,9 @@ class TestExtractStructuralTitle:
)
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Body text")
client = self._make_client(tmp_path)
result = client._extract_structural_title(doc)
assert result == "Website Page Title"
assert extract_structural_title(doc) == "Website Page Title"
def test_body_title(self, tmp_path):
def test_body_title(self):
"""TITLE on BODY layer (h1, PDF title) is extracted."""
doc = DoclingDocument(name="test")
doc.add_text(
@ -63,31 +58,25 @@ class TestExtractStructuralTitle:
)
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Body text")
client = self._make_client(tmp_path)
result = client._extract_structural_title(doc)
assert result == "Document Heading"
assert extract_structural_title(doc) == "Document Heading"
def test_section_header_fallback(self, tmp_path):
def test_section_header_fallback(self):
"""First SECTION_HEADER is used when no TITLE exists."""
doc = DoclingDocument(name="test")
doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Introduction")
doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Background")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Body text")
client = self._make_client(tmp_path)
result = client._extract_structural_title(doc)
assert result == "Introduction"
assert extract_structural_title(doc) == "Introduction"
def test_no_title_or_headers(self, tmp_path):
def test_no_title_or_headers(self):
"""Returns None when no TITLE or SECTION_HEADER exists."""
doc = DoclingDocument(name="test")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Just a paragraph")
client = self._make_client(tmp_path)
result = client._extract_structural_title(doc)
assert result is None
assert extract_structural_title(doc) is None
def test_furniture_title_preferred_over_body_title(self, tmp_path):
def test_furniture_title_preferred_over_body_title(self):
"""FURNITURE TITLE takes priority over BODY TITLE."""
doc = DoclingDocument(name="test")
doc.add_text(
@ -101,11 +90,9 @@ class TestExtractStructuralTitle:
content_layer=ContentLayer.FURNITURE,
)
client = self._make_client(tmp_path)
result = client._extract_structural_title(doc)
assert result == "HTML Page Title"
assert extract_structural_title(doc) == "HTML Page Title"
def test_whitespace_stripped(self, tmp_path):
def test_whitespace_stripped(self):
"""Whitespace is stripped from extracted titles."""
doc = DoclingDocument(name="test")
doc.add_text(
@ -114,11 +101,9 @@ class TestExtractStructuralTitle:
content_layer=ContentLayer.BODY,
)
client = self._make_client(tmp_path)
result = client._extract_structural_title(doc)
assert result == "Padded Title"
assert extract_structural_title(doc) == "Padded Title"
def test_empty_title_text_skipped(self, tmp_path):
def test_empty_title_text_skipped(self):
"""Empty or whitespace-only TITLE text is skipped."""
doc = DoclingDocument(name="test")
doc.add_text(
@ -128,54 +113,49 @@ class TestExtractStructuralTitle:
)
doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Actual Heading")
client = self._make_client(tmp_path)
result = client._extract_structural_title(doc)
assert result == "Actual Heading"
assert extract_structural_title(doc) == "Actual Heading"
# =========================================================================
# _resolve_title
# resolve_title
# =========================================================================
class TestResolveTitle:
def _make_client(self, tmp_path, auto_title=True):
config = AppConfig(processing=ProcessingConfig(auto_title=auto_title))
return HaikuRAG(tmp_path / "test.lancedb", config=config, create=True)
@pytest.mark.asyncio
async def test_auto_title_disabled_returns_none(self, tmp_path):
async def test_auto_title_disabled_returns_none(self):
"""When auto_title is False, returns None (no title generation)."""
doc = DoclingDocument(name="test")
doc.add_text(label=DocItemLabel.TITLE, text="Structural Title")
client = self._make_client(tmp_path, auto_title=False)
result = await client._resolve_title(doc, "some content")
config = AppConfig(processing=ProcessingConfig(auto_title=False))
result = await resolve_title(config, doc, "some content")
assert result is None
@pytest.mark.asyncio
async def test_structural_title_extracted(self, tmp_path):
async def test_structural_title_extracted(self):
"""Structural title is extracted when auto_title is enabled."""
doc = DoclingDocument(name="test")
doc.add_text(label=DocItemLabel.TITLE, text="Auto Extracted Title")
client = self._make_client(tmp_path)
result = await client._resolve_title(doc, "some content")
config = AppConfig(processing=ProcessingConfig(auto_title=True))
result = await resolve_title(config, doc, "some content")
assert result == "Auto Extracted Title"
@pytest.mark.asyncio
async def test_llm_failure_returns_none(self, tmp_path, monkeypatch):
async def test_llm_failure_returns_none(self, monkeypatch):
"""LLM failure during ingestion returns None instead of raising."""
doc = DoclingDocument(name="test")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Just text")
client = self._make_client(tmp_path)
async def exploding_llm(self, content):
async def exploding_llm(config, content):
raise RuntimeError("LLM is down")
monkeypatch.setattr(HaikuRAG, "_generate_title_with_llm", exploding_llm)
result = await client._resolve_title(doc, "some content")
monkeypatch.setattr(
"haiku.rag.client.titles.generate_title_with_llm", exploding_llm
)
config = AppConfig(processing=ProcessingConfig(auto_title=True))
result = await resolve_title(config, doc, "some content")
assert result is None

147
tests/test_versioning.py
View file

@ -1,3 +1,5 @@
import asyncio
import pytest
from haiku.rag.client import HaikuRAG
@ -60,7 +62,7 @@ async def test_version_rollback_on_update_failure(temp_db_path):
assert len(original_chunks) > 0
def test_new_database_does_not_run_upgrades(monkeypatch, temp_db_path):
async def test_new_database_does_not_run_upgrades(monkeypatch, temp_db_path):
def fail_if_called(*_args, **_kwargs):
raise AssertionError("run_pending_upgrades should not be called for new DB")
@ -69,11 +71,13 @@ def test_new_database_does_not_run_upgrades(monkeypatch, temp_db_path):
fail_if_called,
)
Store(temp_db_path, create=True)
async with Store(temp_db_path, create=True):
pass
def test_existing_database_checks_migrations(monkeypatch, temp_db_path):
Store(temp_db_path, create=True)
async def test_existing_database_checks_migrations(monkeypatch, temp_db_path):
async with Store(temp_db_path, create=True):
pass
from haiku.rag.store import upgrades
@ -90,25 +94,33 @@ def test_existing_database_checks_migrations(monkeypatch, temp_db_path):
)
# Opening an existing database should check for pending migrations
Store(temp_db_path)
async with Store(temp_db_path):
pass
assert called["value"]
async def _wait_for_background_vacuum(client):
"""Wait for any in-flight background vacuum tasks to complete."""
await client._await_vacuum_tasks()
@pytest.mark.vcr()
async def test_vacuum_with_retention_threshold(temp_db_path):
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
# Create first document
await client.create_document(content="First document")
await _wait_for_background_vacuum(client)
# Create second document
await client.create_document(content="Second document")
await _wait_for_background_vacuum(client)
store = client.store
# Get initial version counts (should have multiple versions from creates)
initial_doc_versions = len(list(store.documents_table.list_versions()))
initial_chunk_versions = len(list(store.chunks_table.list_versions()))
initial_doc_versions = len(await store.documents_table.list_versions())
initial_chunk_versions = len(await store.chunks_table.list_versions())
assert initial_doc_versions > 1, "Should have multiple document table versions"
assert initial_chunk_versions > 1, "Should have multiple chunk table versions"
@ -117,8 +129,8 @@ async def test_vacuum_with_retention_threshold(temp_db_path):
# Note: vacuum may create new versions even when not cleaning up old ones
await store.vacuum()
after_default_doc_versions = len(list(store.documents_table.list_versions()))
after_default_chunk_versions = len(list(store.chunks_table.list_versions()))
after_default_doc_versions = len(await store.documents_table.list_versions())
after_default_chunk_versions = len(await store.chunks_table.list_versions())
# After vacuum with retention, version count should stay the same or increase
# (optimize may create new versions) but not decrease
@ -132,8 +144,8 @@ async def test_vacuum_with_retention_threshold(temp_db_path):
# Vacuum with 0 threshold - should significantly reduce versions
await store.vacuum(retention_seconds=0)
after_zero_doc_versions = len(list(store.documents_table.list_versions()))
after_zero_chunk_versions = len(list(store.chunks_table.list_versions()))
after_zero_doc_versions = len(await store.documents_table.list_versions())
after_zero_chunk_versions = len(await store.chunks_table.list_versions())
# After aggressive vacuum, should have minimal versions (1-2)
# Note: optimize operation may create a version after cleanup
@ -168,16 +180,96 @@ async def test_vacuum_completes_before_context_exit(temp_db_path, monkeypatch):
await client.create_document(content=f"Test document {i}")
# After context exit, automatic vacuum should have kept versions minimal
store = Store(temp_db_path, create=True)
final_versions = len(list(store.documents_table.list_versions()))
async with Store(temp_db_path, create=True) as store:
final_versions = len(await store.documents_table.list_versions())
# With retention_seconds=0, vacuum aggressively cleans up between operations
# Should have very few versions remaining (1-2)
assert final_versions <= 2, (
f"Aggressive vacuum should keep minimal versions, got {final_versions}"
# With retention_seconds=0, vacuum aggressively cleans up between operations
# Should have very few versions remaining (1-2)
assert final_versions <= 2, (
f"Aggressive vacuum should keep minimal versions, got {final_versions}"
)
assert final_versions >= 1, "Should have at least one version remaining"
@pytest.mark.vcr()
async def test_aexit_awaits_background_vacuum(temp_db_path, monkeypatch):
"""__aexit__ must await any in-flight background vacuum, not just release the lock.
Background vacuum runs as an asyncio task; the event loop may not have scheduled
it yet when __aexit__ runs. Simply acquiring the vacuum lock (which is free until
the task actually starts) would let close() proceed before vacuum runs.
"""
from haiku.rag.config import Config
monkeypatch.setattr(Config.storage, "auto_vacuum", True)
vacuum_started = asyncio.Event()
vacuum_completed = asyncio.Event()
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
original_vacuum = client.store.vacuum
async def instrumented_vacuum(*args, **kwargs):
vacuum_started.set()
# Delay so __aexit__ would see an unstarted/incomplete task if it
# relied on the lock rather than awaiting the task directly.
await asyncio.sleep(0.05)
await original_vacuum(*args, **kwargs)
vacuum_completed.set()
client.store.vacuum = instrumented_vacuum
await client.create_document(content="triggers background vacuum")
assert vacuum_started.is_set(), "Background vacuum task never ran"
assert vacuum_completed.is_set(), "__aexit__ exited before vacuum finished"
@pytest.mark.vcr()
async def test_aexit_awaits_all_background_vacuums(temp_db_path, monkeypatch):
"""Multiple create_document calls schedule multiple vacuum tasks; __aexit__
must await all of them, not just the last-scheduled one.
Scenario: Task A acquires the vacuum lock and is slow. Task B is scheduled
while Task A still holds the lock — Task B sees the lock held and returns
immediately. If the client only tracks the most recently scheduled task,
__aexit__ awaits the fast no-op B and closes the connection while Task A
is still running.
"""
from haiku.rag.config import Config
monkeypatch.setattr(Config.storage, "auto_vacuum", True)
first_vacuum_completed = asyncio.Event()
async with HaikuRAG(db_path=temp_db_path, create=True) as client:
call_count = 0
async def slow_vacuum(*_args, **_kwargs):
nonlocal call_count
call_count += 1
my_num = call_count
# Mimic the real vacuum's skip-if-running behavior.
if client.store._vacuum_lock.locked():
return
async with client.store._vacuum_lock:
if my_num == 1:
# Hold the lock longer than any other operation in the
# test so Task A cannot finish incidentally. __aexit__
# must explicitly wait for this task.
await asyncio.sleep(2.0)
first_vacuum_completed.set()
client.store.vacuum = slow_vacuum
await client.create_document(content="triggers first vacuum")
# Let Task A start and acquire the vacuum lock before scheduling B.
await asyncio.sleep(0.02)
await client.create_document(content="triggers second vacuum")
assert first_vacuum_completed.is_set(), (
"__aexit__ returned before the first vacuum task finished"
)
assert final_versions >= 1, "Should have at least one version remaining"
store.close()
@pytest.mark.vcr()
@ -194,8 +286,8 @@ async def test_auto_vacuum_disabled_skips_vacuum(temp_db_path, monkeypatch):
await client.create_document(content=f"Test document {i}")
# Count versions - should accumulate without vacuum
doc_versions = len(list(client.store.documents_table.list_versions()))
chunk_versions = len(list(client.store.chunks_table.list_versions()))
doc_versions = len(await client.store.documents_table.list_versions())
chunk_versions = len(await client.store.chunks_table.list_versions())
# Without auto-vacuum, versions should accumulate (more than 3 from creates)
assert doc_versions >= 3, (
@ -221,11 +313,10 @@ async def test_auto_vacuum_enabled_triggers_vacuum(temp_db_path, monkeypatch):
await client.create_document(content=f"Test document {i}")
# After context exit, vacuum should have cleaned up
store = Store(temp_db_path, create=True)
final_versions = len(list(store.documents_table.list_versions()))
async with Store(temp_db_path, create=True) as store:
final_versions = len(await store.documents_table.list_versions())
# With auto_vacuum=True and retention=0, should have minimal versions
assert final_versions <= 2, (
f"With auto-vacuum enabled, should have minimal versions, got {final_versions}"
)
store.close()
# With auto_vacuum=True and retention=0, should have minimal versions
assert final_versions <= 2, (
f"With auto-vacuum enabled, should have minimal versions, got {final_versions}"
)