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

Share document preparation and HTTP acquisition

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.

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:
+async def _prepare_and_title(
-    with tempfile.NamedTemporaryFile(
+    client: "HaikuRAG", document: Document, docling_document: "DoclingDocument"
-        mode="wb", suffix=suffix, delete=False
+) -> None:
-    ) as temp_file:
+    """Fill the document from its converted form and title it if it has none.
-        temp_file.write(body)
-        temp_file.flush()
-        return Path(temp_file.name)
 
-
+    A caller-supplied title always wins: set it on the document before calling.
-async def _write_fetch_body(body: bytes, suffix: str) -> Path:
+    Update paths that must keep an existing empty title call
-    return await asyncio.to_thread(_write_fetch_body_sync, body, suffix)
+    ``_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:
+    return await _store_document_with_chunks(client, document, chunks, docling_document)
-        document.title = await resolve_title(
-            client._config, docling_document, stored_content
-        )
-
-    return await _store_document_with_chunks(
-        client, document, embedded_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)
+    await _prepare_and_title(client, document, docling_document)
-    if title is None:
-        document.title = await resolve_title(client._config, docling_document, content)
 
     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)
+        await _prepare_and_title(client, document, item.docling_document)
-        if document.title is None:
-            document.title = await resolve_title(
-                client._config, item.docling_document, content
-            )
         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:
+        await _prepare_and_title(client, existing_doc, docling_document)
-            existing_doc.title = await resolve_title(
-                client._config, docling_document, stored_content
-            )
         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)
+    await _prepare_and_title(client, document, docling_document)
-    if document.title is None:
-        document.title = await resolve_title(
-            client._config, docling_document, stored_content
-        )
     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("*"):
+            for child in walk_files(local_path):
-                # 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
                 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
     )
 
 

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
+        # One HTTP acquisition path: the same adapter the ingester fetches with.
-        async with httpx.AsyncClient() as http:
+        from haiku.rag.ingester.sources.http import HTTPSource
-            response = await http.get(source)
-            response.raise_for_status()
 
-            content_type = response.headers.get("content-type", "").lower()
+        fetcher = HTTPSource(source_id="convert")
-            file_extension = get_extension_from_content_type_or_url(
+        try:
-                source, content_type
+            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:
+        temp_path = await _write_fetch_body(result.body, file_extension)
-                raise UnsupportedSourceError(
+        try:
-                    f"Unsupported content type/extension: {content_type}/{file_extension}"
+            doc = await _convert_file(temp_path, source_uri or source)
-                )
+            _warn_if_descriptions_missing(config, doc, source)
-
+            return doc
-            with tempfile.NamedTemporaryFile(
+        finally:
-                mode="wb", suffix=file_extension, delete=False
+            temp_path.unlink(missing_ok=True)
-            ) 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)
 
     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.

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 path in walk_files(self.root):
-        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:
             if not path.is_file():
                 continue
             if not self.filter.include_file(str(path)):

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"}

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."""