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

Share document preparation and HTTP acquisition

Browse source 
Five call sites repeated the same post-conversion preparation: store the
Docling representation and resolve a title when none was supplied.
_prepare_and_title now owns that sequence. update_document continues to
call _prepare_document_from_docling directly because an explicit update
must preserve an existing empty title.

create_document, both content-replacement branches of update_document,
and source ingestion embedded eagerly before passing chunks to a
persistence funnel that checked them again. The funnels now own
embedding, including the checks required by import_document and
import_documents for caller-supplied chunks.

Move the document.embed span into ensure_chunks_embedded after its early
return. Every path that performs embedding is now instrumented, while
operations whose chunks are already embedded emit no span.

convert() previously used its own HTTP client and temporary-file path.
Route URL conversion through HTTPSource, matching source ingestion, and
move _write_fetch_body to processing.py so both paths share temporary
file handling without an import cycle.

Add walk_files for filesystem enumeration and use it from both
FSSource.discover and one-shot directory ingestion. Symlink escape
filtering now has one implementation.
This commit is contained in:
Yiorgis Gozadinos 2026-08-20 10:50:33 +03:00
parent 9d64a0d9f0
commit d1691d3942
No known key found for this signature in database
6 changed files with 175 additions and 114 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
CHANGELOG.md
View file

@ -9,6 +9,9 @@
### Changed
- `HaikuRAG.convert(url)` fetches through `HTTPSource`, the same adapter the ingester uses, instead of its own httpx client. `_write_fetch_body` moved from `client.documents` to `client.processing`.
- Chunk embedding is owned by the persistence funnels: `create_document`, `update_document` and source ingestion no longer embed eagerly before handing chunks to a check that would embed them anyway. The `document.embed` span moved onto `ensure_chunks_embedded`, so every path is instrumented rather than only ingest.
- One-shot directory ingestion and `FSSource.discover` share `walk_files`, so the symlink-escape guard lives in one place.
- Configuration sections reject unknown keys. A typo or a setting that has been renamed or removed now fails validation with its path (`providers.docling_serve.bogus: Extra inputs are not permitted`) instead of being silently ignored.
- `processing.converter`, `processing.chunker` and `processing.chunker_type` are constrained to their supported values, so an unsupported one fails at load rather than at first use.
- Numeric settings carry bounds: sizes, limits, dimensions, token budgets, attempt counts, breaker thresholds and `min_chunks` must be positive; retention, delays, intervals and cooldowns non-negative; `doctor.duplicates.similarity_threshold` within 0-1; `ingester.api.port` within 0-65535 (0 keeps its OS-assigned meaning); `ingester.workers.worker_count` allows 0 for an API-and-reaper-only process.

92
haiku_rag_slim/haiku/rag/client/documents.py
View file

@ -3,7 +3,6 @@ import hashlib
import json
import logging
import mimetypes
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING
@ -11,6 +10,7 @@ from urllib.parse import quote, unquote, urlparse
from haiku.rag.client.exceptions import UnsupportedSourceError
from haiku.rag.client.processing import (
_write_fetch_body,
ensure_chunks_embedded,
get_extension_from_content_type_or_url,
)
@ -84,17 +84,20 @@ async def _prepare_document_from_docling(
)
def _write_fetch_body_sync(body: bytes, suffix: str) -> Path:
with tempfile.NamedTemporaryFile(
mode="wb", suffix=suffix, delete=False
) as temp_file:
temp_file.write(body)
temp_file.flush()
return Path(temp_file.name)
async def _prepare_and_title(
client: "HaikuRAG", document: Document, docling_document: "DoclingDocument"
) -> None:
"""Fill the document from its converted form and title it if it has none.
async def _write_fetch_body(body: bytes, suffix: str) -> Path:
return await asyncio.to_thread(_write_fetch_body_sync, body, suffix)
A caller-supplied title always wins: set it on the document before calling.
Update paths that must keep an existing empty title call
``_prepare_document_from_docling`` directly instead.
"""
stored_content = await _prepare_document_from_docling(document, docling_document)
if document.title is None:
document.title = await resolve_title(
client._config, docling_document, stored_content
)
def parent_uri_filter(parent_uri: str) -> str:
@ -223,12 +226,9 @@ async def create_document(
Converts the content, chunks it, and generates embeddings.
"""
from haiku.rag.embeddings import embed_chunks
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.embedder, client._config)
document = Document(
content="",
@ -236,16 +236,9 @@ async def create_document(
title=title,
metadata=metadata or {},
)
stored_content = await _prepare_document_from_docling(document, docling_document)
await _prepare_and_title(client, document, docling_document)
if title is None:
document.title = await resolve_title(
client._config, docling_document, stored_content
)
return await _store_document_with_chunks(
client, document, embedded_chunks, docling_document
)
return await _store_document_with_chunks(client, document, chunks, docling_document)
async def import_document(
@ -267,9 +260,7 @@ async def import_document(
title=title,
metadata=metadata or {},
)
content = await _prepare_document_from_docling(document, docling_document)
if title is None:
document.title = await resolve_title(client._config, docling_document, content)
await _prepare_and_title(client, document, docling_document)
return await _store_document_with_chunks(client, document, chunks, docling_document)
@ -346,11 +337,7 @@ async def import_documents(
title=item.title,
metadata=item.metadata or {},
)
content = await _prepare_document_from_docling(document, item.docling_document)
if document.title is None:
document.title = await resolve_title(
client._config, item.docling_document, content
)
await _prepare_and_title(client, document, item.docling_document)
prepared.append((document, item.chunks, item.docling_document))
return await _store_documents_with_chunks(client, prepared)
@ -428,7 +415,6 @@ async def _ingest_fetch_result(
extension (and thus the docling format), overriding the URI/content-type
fallback. Callers pass it when ``result.uri`` cannot yield the right
extension, e.g. embedded attachments whose name lives in a URI fragment."""
from haiku.rag.embeddings import embed_chunks
converter = get_converter(client._config)
if filename is not None:
@ -463,10 +449,6 @@ async def _ingest_fetch_result(
with logfire.span("document.chunk", uri=result.uri) as chunk_span:
chunks = await client.chunk(docling_document)
chunk_span.set_attribute("chunks_created", len(chunks))
with logfire.span("document.embed", uri=result.uri):
embedded_chunks = await embed_chunks(
chunks, client.embedder, client._config
)
finally:
if cleanup_path is not None:
cleanup_path.unlink(missing_ok=True)
@ -475,18 +457,12 @@ async def _ingest_fetch_result(
if existing_doc:
existing_doc.metadata = final_metadata
stored_content = await _prepare_document_from_docling(
existing_doc, 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
)
await _prepare_and_title(client, existing_doc, docling_document)
with logfire.span("document.store", uri=result.uri, op="update") as store_span:
updated = await _update_document_with_chunks(
client, existing_doc, embedded_chunks, docling_document
client, existing_doc, chunks, docling_document
)
store_span.set_attribute("document_id", updated.id)
await _reconcile_pdf_attachments(client, updated, result.body, depth=depth)
@ -498,14 +474,10 @@ async def _ingest_fetch_result(
title=title,
metadata=final_metadata,
)
stored_content = await _prepare_document_from_docling(document, docling_document)
if document.title is None:
document.title = await resolve_title(
client._config, docling_document, stored_content
)
await _prepare_and_title(client, document, docling_document)
with logfire.span("document.store", uri=result.uri, op="create") as store_span:
created = await _store_document_with_chunks(
client, document, embedded_chunks, docling_document
client, document, chunks, docling_document
)
store_span.set_attribute("document_id", created.id)
await _reconcile_pdf_attachments(client, created, result.body, depth=depth)
@ -697,21 +669,14 @@ async def create_document_from_source(
"produces its own document with its own auto-derived URI."
)
from haiku.rag.ingester.sources.filter import FileFilter
from haiku.rag.ingester.sources.fs import walk_files
# One-shot CLI directory ingest uses the converter's supported
# extensions but no include/ignore patterns. For pattern-based
# filtering use `haiku-ingester serve` with an FS source.
documents: list[Document] = []
filter = FileFilter()
for child in local_path.rglob("*"):
# rglob does not recurse into symlinked directories, but it does
# yield symlinked files. Skip the ones resolving outside the
# directory the caller named, as FSSource.discover does.
if child.is_symlink():
resolved = child.resolve(strict=False)
if not resolved.is_relative_to(local_path.resolve()):
continue
child = resolved
for child in walk_files(local_path):
if child.is_file() and filter.include_file(str(child)):
doc = await create_document_from_source(
client,
@ -861,7 +826,6 @@ async def update_document(
ValueError: If document not found, or if both content and
docling_document are provided.
"""
from haiku.rag.embeddings import embed_chunks
if content is not None and docling_document is not None:
raise ValueError(
@ -905,11 +869,8 @@ async def update_document(
await _prepare_document_from_docling(existing_doc, docling_document)
new_chunks = await client.chunk(docling_document)
embedded_chunks = await embed_chunks(
new_chunks, client.embedder, client._config
)
return await _update_document_with_chunks(
client, existing_doc, embedded_chunks, docling_document
client, existing_doc, new_chunks, docling_document
)
assert content is not None
@ -919,9 +880,8 @@ async def update_document(
await _prepare_document_from_docling(existing_doc, converted_docling)
new_chunks = await client.chunk(converted_docling)
embedded_chunks = await embed_chunks(new_chunks, client.embedder, client._config)
return await _update_document_with_chunks(
client, existing_doc, embedded_chunks, converted_docling
client, existing_doc, new_chunks, converted_docling
)

66
haiku_rag_slim/haiku/rag/client/processing.py
View file

@ -6,7 +6,7 @@ from pathlib import Path
from typing import TYPE_CHECKING
from urllib.parse import urlparse
import httpx
import logfire
from haiku.rag.client.exceptions import UnsupportedSourceError
from haiku.rag.config import AppConfig
@ -56,6 +56,19 @@ def _warn_if_descriptions_missing(
)
def _write_fetch_body_sync(body: bytes, suffix: str) -> Path:
with tempfile.NamedTemporaryFile(
mode="wb", suffix=suffix, delete=False
) as temp_file:
temp_file.write(body)
temp_file.flush()
return Path(temp_file.name)
async def _write_fetch_body(body: bytes, suffix: str) -> Path:
return await asyncio.to_thread(_write_fetch_body_sync, body, suffix)
async def convert(
config: AppConfig,
source: Path | str,
@ -117,35 +130,31 @@ async def convert(
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()
# One HTTP acquisition path: the same adapter the ingester fetches with.
from haiku.rag.ingester.sources.http import HTTPSource
content_type = response.headers.get("content-type", "").lower()
file_extension = get_extension_from_content_type_or_url(
source, content_type
fetcher = HTTPSource(source_id="convert")
try:
result = await fetcher.fetch(source)
finally:
await fetcher.aclose()
file_extension = get_extension_from_content_type_or_url(
source, result.content_type
)
if file_extension not in converter.supported_extensions:
raise UnsupportedSourceError(
f"Unsupported content type/extension: "
f"{result.content_type}/{file_extension}"
)
if file_extension not in converter.supported_extensions:
raise UnsupportedSourceError(
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:
effective_uri = source_uri or source
doc = await _convert_file(temp_path, effective_uri)
_warn_if_descriptions_missing(config, doc, source)
return doc
finally:
temp_path.unlink(missing_ok=True)
temp_path = await _write_fetch_body(result.body, file_extension)
try:
doc = await _convert_file(temp_path, source_uri or source)
_warn_if_descriptions_missing(config, doc, source)
return doc
finally:
temp_path.unlink(missing_ok=True)
elif parsed.scheme == "file":
# file:// URI
@ -350,7 +359,8 @@ async def ensure_chunks_embedded(
if not chunks_to_embed:
return chunks
embedded = await embed_chunks(chunks_to_embed, embedder, config)
with logfire.span("document.embed", chunks=len(chunks_to_embed)):
embedded = await embed_chunks(chunks_to_embed, embedder, config)
# embed_chunks preserves input order; fill positionally, since duplicate
# chunk texts across documents make a content-keyed lookup ambiguous.

43
haiku_rag_slim/haiku/rag/ingester/sources/fs.py
View file

@ -29,6 +29,32 @@ def _uri_to_path(uri: str) -> Path:
raise ValueError(f"Unsupported URI scheme for FSSource: {uri}")
def walk_files(root: Path) -> list[Path]:
"""Every file under ``root``, sorted, with symlink escapes dropped.
Directory symlinks are never followed. A symlinked file is resolved and kept
only when its target is inside ``root``, so a link cannot pull in a file from
outside the tree the caller named. Comparison is against the resolved root, so
a root reached through a symlink (macOS ``/tmp``) compares like any other.
"""
resolved_root = root.resolve()
candidates: list[Path] = []
for dirpath, _dirnames, filenames in os.walk(root, followlinks=False):
for filename in filenames:
path = Path(dirpath) / filename
if path.is_symlink():
try:
target = path.resolve(strict=False)
except OSError: # pragma: no cover - strict=False absorbs these
continue
if not target.is_relative_to(resolved_root):
continue
path = target
candidates.append(path)
candidates.sort()
return candidates
class FSSource:
def __init__(
self,
@ -136,22 +162,7 @@ class FSSource:
# supports/head/fetch's resolve-then-check behaviour. Out-of-root
# targets stay skipped so a stray link can't exfiltrate data the
# operator didn't intend to expose.
candidates: list[Path] = []
for dirpath, _dirnames, filenames in os.walk(self.root, followlinks=False):
for filename in filenames:
path = Path(dirpath) / filename
if path.is_symlink():
try:
resolved = path.resolve(strict=False)
except OSError: # pragma: no cover - strict=False absorbs these
continue
if not resolved.is_relative_to(self.root):
continue
path = resolved
candidates.append(path)
candidates.sort()
for path in candidates:
for path in walk_files(self.root):
if not path.is_file():
continue
if not self.filter.include_file(str(path)):

19
tests/ingester/test_fs_source.py
View file

@ -373,3 +373,22 @@ async def test_discover_skips_symlink_to_missing_in_root_target(tmp_path):
events = [e async for e in src.discover()]
assert {e.uri for e in events} == {(tmp_path / "real.md").as_uri()}
def test_walk_files_drops_links_escaping_the_root(tmp_path):
import os
from haiku.rag.ingester.sources.fs import walk_files
tree = tmp_path / "tree"
outside = tmp_path / "outside"
tree.mkdir()
outside.mkdir()
(tree / "real.txt").write_text("in tree")
(outside / "secret.txt").write_text("out of tree")
os.symlink(outside / "secret.txt", tree / "escape.txt")
os.symlink(outside, tree / "escape_dir")
found = {path.name for path in walk_files(tree)}
assert found == {"real.txt"}

66
tests/test_client.py
View file

@ -14,9 +14,9 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.client.documents import (
DocumentImport,
_prepare_document_from_docling,
_write_fetch_body,
check_source_accessible,
)
from haiku.rag.client.processing import _write_fetch_body
from haiku.rag.config import get_config
from haiku.rag.embeddings import EmbedderWrapper
from haiku.rag.ingester.sources.base import FetchResult
@ -62,17 +62,17 @@ async def test_prepare_document_from_docling_runs_off_event_loop_thread(monkeypa
@pytest.mark.asyncio
async def test_write_fetch_body_runs_off_event_loop_thread(monkeypatch):
import haiku.rag.client.documents as documents
import haiku.rag.client.processing as processing
event_loop_thread = threading.current_thread()
called_from: list[threading.Thread] = []
original = documents._write_fetch_body_sync
original = processing._write_fetch_body_sync
def spy(body, suffix):
called_from.append(threading.current_thread())
return original(body, suffix)
monkeypatch.setattr(documents, "_write_fetch_body_sync", spy)
monkeypatch.setattr(processing, "_write_fetch_body_sync", spy)
path = await _write_fetch_body(b"payload", ".bin")
try:
@ -974,6 +974,64 @@ async def test_client_import_documents_batches_embeddings(temp_db_path):
assert all(len(row["vector"]) == dim for row in rows)
async def test_single_and_batch_import_store_the_same_document(temp_db_path):
"""import_document and import_documents share preparation and persistence, so
the same input has to land as the same stored document."""
dim = get_config().embeddings.model.vector_dim
async with HaikuRAG(temp_db_path, create=True) as client:
client.store.embedder = _CountingEmbedder(dim)
single = await client.import_document(
_docling_doc("a", "Alpha document body"),
[Chunk(content="Alpha document body", order=0)],
uri="mem://single",
)
[batched] = await client.import_documents(
[
DocumentImport(
docling_document=_docling_doc("a", "Alpha document body"),
chunks=[Chunk(content="Alpha document body", order=0)],
uri="mem://batch",
)
]
)
assert single.title == batched.title
assert single.content == batched.content
for doc in (single, batched):
chunks = await client.chunk_repository.get_by_document_id(doc.id)
assert [c.content for c in chunks] == ["Alpha document body"]
# get_by_document_id does not project the vector, so read it directly.
rows = await (
client.store.chunks_table.query()
.select(["document_id", "vector"])
.to_list()
)
vectors = {row["document_id"]: row["vector"] for row in rows}
assert set(vectors) == {single.id, batched.id}
assert all(len(vector) == dim for vector in vectors.values())
single_items = await client.document_item_repository.get_item_count(single.id)
batched_items = await client.document_item_repository.get_item_count(batched.id)
assert single_items == batched_items > 0
async def test_create_document_embeds_in_one_pass(temp_db_path):
"""Embedding is owned by the persistence funnel, so an operation makes one
embedder pass — no eager embed followed by a check that could embed again."""
dim = get_config().embeddings.model.vector_dim
embedder = _CountingEmbedder(dim)
async with HaikuRAG(temp_db_path, create=True) as client:
client.store.embedder = embedder
await client.create_document("Alpha document body")
assert len(embedder.batches) == 1
async def test_client_import_documents_mixed_embeddings(temp_db_path):
"""Pre-embedded chunks keep their vectors; only the unembedded ones go
through the embedder, in one batch."""