diff --git a/CHANGELOG.md b/CHANGELOG.md index a81ea630..b8c6e301 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,8 @@ ### 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. The only previously-supported direct-construction path was `HaikuRAG(db_path=None).download_models()`, which still works because it only reads config. +- **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 diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index b1264882..2b9c4e9e 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -155,36 +155,10 @@ class HaikuRAG: return await chunk(self._config, docling_document) - async def _ensure_chunks_embedded(self, chunks: list[Chunk]) -> list[Chunk]: - from haiku.rag.client.processing import ensure_chunks_embedded - - return await ensure_chunks_embedded(self._config, chunks) - # ========================================================================= # Title Generation # ========================================================================= - def _extract_structural_title( - self, docling_document: "DoclingDocument" - ) -> str | None: - from haiku.rag.client.titles import extract_structural_title - - return extract_structural_title(docling_document) - - async def _generate_title_with_llm(self, content: str) -> str | None: - from haiku.rag.client.titles import generate_title_with_llm - - return await generate_title_with_llm(self._config, content) - - async def _resolve_title( - self, - docling_document: "DoclingDocument", - content: str, - ) -> str | None: - from haiku.rag.client.titles import resolve_title - - return await resolve_title(self._config, docling_document, content) - async def generate_title(self, document: Document) -> str | None: from haiku.rag.client.titles import generate_title diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 3435f91f..1cd0861e 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -7,6 +7,8 @@ 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 @@ -29,7 +31,7 @@ async def _store_document_with_chunks( Handles versioning/rollback on failure. """ # Ensure all chunks have embeddings before storing - chunks = await client._ensure_chunks_embedded(chunks) + chunks = await ensure_chunks_embedded(client._config, chunks) # Snapshot table versions for versioned rollback (if supported) versions = await client.store.current_table_versions() @@ -77,7 +79,7 @@ async def _update_document_with_chunks( """ assert document.id is not None, "Document ID is required for update" - chunks = await client._ensure_chunks_embedded(chunks) + chunks = await ensure_chunks_embedded(client._config, chunks) versions = await client.store.current_table_versions() @@ -134,7 +136,7 @@ async def create_document( stored_content = docling_document.export_to_markdown() if title is None: - title = await client._resolve_title(docling_document, stored_content) + title = await resolve_title(client._config, docling_document, stored_content) document = Document( content=stored_content, @@ -164,7 +166,7 @@ async def import_document( """ content = docling_document.export_to_markdown() if title is None: - title = await client._resolve_title(docling_document, content) + title = await resolve_title(client._config, docling_document, content) document = Document( content=content, @@ -285,15 +287,17 @@ async def _create_document_from_file( if title is not None: existing_doc.title = title elif existing_doc.title is None: - existing_doc.title = await client._resolve_title( - docling_document, stored_content + 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 client._resolve_title(docling_document, stored_content) + title = await resolve_title( + client._config, docling_document, stored_content + ) document = Document( content=stored_content, uri=uri, @@ -379,15 +383,17 @@ async def _create_or_update_document_from_url( if title is not None: existing_doc.title = title elif existing_doc.title is None: - existing_doc.title = await client._resolve_title( - docling_document, stored_content + 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 client._resolve_title(docling_document, stored_content) + title = await resolve_title( + client._config, docling_document, stored_content + ) document = Document( content=stored_content, uri=url, diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index f4204efe..9a1569ed 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -634,11 +634,6 @@ class Store: await self.settings_table.restore(int(versions["settings"])) return True - @property - def _connection(self): - """Compatibility property for repositories expecting _connection.""" - return self - async def _checkout_tables_before(self, before: datetime) -> None: """Checkout all tables to their state at or before the given datetime. diff --git a/tests/test_context_enhancement.py b/tests/test_context_enhancement.py index 0d21cd75..aad53a6d 100644 --- a/tests/test_context_enhancement.py +++ b/tests/test_context_enhancement.py @@ -4,6 +4,7 @@ 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 @@ -13,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, diff --git a/tests/test_title_generation.py b/tests/test_title_generation.py index 7411c674..c13f736b 100644 --- a/tests/test_title_generation.py +++ b/tests/test_title_generation.py @@ -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 ) 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