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

Merge pull request #564 from ggozad/refactor/prepare-and-acquire

Browse source 
Share document preparation and HTTP acquisition
This commit is contained in:
Yiorgis Gozadinos 2026-08-20 11:06:13 +03:00 committed by GitHub
commit a896ac9eec
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
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."""