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

drop private title/embedding forwarders from HaikuRAG

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-24 15:52:50 +03:00
parent 0ffca7727c
commit cd4224a727
No known key found for this signature in database
6 changed files with 49 additions and 92 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
CHANGELOG.md
View file

@ -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

26
haiku_rag_slim/haiku/rag/client/__init__.py
View file

@ -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

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

@ -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,

5
haiku_rag_slim/haiku/rag/store/engine.py
View file

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

3
tests/test_context_enhancement.py
View file

@ -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,

78
tests/test_title_generation.py
View file

@ -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 <title>) 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