diff --git a/CHANGELOG.md b/CHANGELOG.md index 29d34af2..a70e556d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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. diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index d67f4f99..39b93781 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -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 ) diff --git a/haiku_rag_slim/haiku/rag/client/processing.py b/haiku_rag_slim/haiku/rag/client/processing.py index 4d541658..407b4020 100644 --- a/haiku_rag_slim/haiku/rag/client/processing.py +++ b/haiku_rag_slim/haiku/rag/client/processing.py @@ -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. diff --git a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py index 26747ecd..7c77b3d5 100644 --- a/haiku_rag_slim/haiku/rag/ingester/sources/fs.py +++ b/haiku_rag_slim/haiku/rag/ingester/sources/fs.py @@ -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)): diff --git a/tests/ingester/test_fs_source.py b/tests/ingester/test_fs_source.py index 55b4fc17..1ce4a40e 100644 --- a/tests/ingester/test_fs_source.py +++ b/tests/ingester/test_fs_source.py @@ -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"} diff --git a/tests/test_client.py b/tests/test_client.py index 75164823..e543f9c2 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -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."""