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

Merge pull request #288 from ggozad/feat/title-generation

Browse source 
Automatic title generation for documents
This commit is contained in:
Yiorgis Gozadinos 2026-02-26 11:21:05 +02:00 committed by GitHub
commit 3111e54718
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 661 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
CHANGELOG.md
View file

@ -1,6 +1,13 @@
# Changelog
## [Unreleased]
### Added
- **Automatic title generation**: Documents can now have titles auto-generated during ingestion via `processing.auto_title: true`. Uses two-tier extraction: structural metadata from DoclingDocument (HTML `<title>`, h1, section headers) first, with LLM fallback via configurable `processing.title_model`
- **`generate_title()`**: Public method on `HaikuRAG` to generate a title for an existing document on demand
- **`rebuild --title-only`**: New rebuild mode that generates titles only for untitled documents without re-chunking or re-embedding
- **`add --title`**: CLI option to set a title when adding text documents
## [0.32.0] - 2026-02-24
### Changed

7
docs/cli.md
View file

@ -50,6 +50,9 @@ From text:
```bash
haiku-rag add "Your document content here"
# Set a title
haiku-rag add "Your document content here" --title "My Document"
# Attach metadata (repeat --meta for multiple entries)
haiku-rag add "Your document content here" --meta author=alice --meta topic=notes
```
@ -391,6 +394,9 @@ haiku-rag rebuild --rechunk
# Only regenerate embeddings (fastest, keeps existing chunks)
haiku-rag rebuild --embed-only
# Only generate titles for untitled documents
haiku-rag rebuild --title-only
```
**Rebuild modes:**
@ -400,6 +406,7 @@ haiku-rag rebuild --embed-only
| Full | (default) | Changed converter, source files updated |
| Rechunk | `--rechunk` | Changed chunking strategy or chunk size |
| Embed only | `--embed-only` | Changed embedding model or vector dimensions |
| Title only | `--title-only` | Generate titles for documents without one |
### Download Models

5
docs/configuration/index.md
View file

@ -117,6 +117,11 @@ processing:
chunking_tokenizer: "Qwen/Qwen3-Embedding-0.6B"
chunking_merge_peers: true
chunking_use_markdown_tables: false
auto_title: false # Auto-generate titles on ingestion
title_model:
provider: ollama
name: gpt-oss
enable_thinking: false
conversion_options:
do_ocr: true
force_ocr: false

31
docs/configuration/processing.md
View file

@ -21,6 +21,13 @@ processing:
chunking_merge_peers: true # Merge undersized successive chunks
chunking_use_markdown_tables: false # Use markdown tables vs narrative format
# Automatic title generation
auto_title: false # Auto-generate titles on ingestion
title_model: # LLM for title generation (fallback)
provider: ollama
name: gpt-oss
enable_thinking: false
# Conversion options (works with both local and remote converters)
conversion_options:
# OCR settings
@ -201,6 +208,30 @@ picture_description:
See [VLM Picture Description with docling-serve](../remote-processing.md#vlm-picture-description-with-docling-serve) for a complete example.
### Automatic Title Generation
Enable automatic title generation during document ingestion:
```yaml
processing:
auto_title: true
title_model:
provider: ollama
name: gpt-oss
enable_thinking: false
```
When `auto_title` is enabled, haiku.rag attempts to extract a title for each document during ingestion using a two-tier approach:
1. **Structural extraction** (free, no model calls): Scans the DoclingDocument for semantic labels — HTML `<title>` tags, `<h1>` headings, PDF title blocks, and section headers
2. **LLM fallback**: When no structural title is found (e.g., plain text), generates a title using the configured `title_model`
Priority order: HTML `<title>` (furniture layer) → h1/PDF title (body layer) → first section header → LLM generation.
Explicit titles passed via `title=` parameter always take precedence and are never overridden. When updating documents, existing titles are preserved — auto-generation only applies to untitled documents.
To generate titles for existing untitled documents, use [`rebuild --title-only`](../cli.md#rebuild-database).
### Local vs Remote Processing
**Local processing** (default):

22
docs/python.md
View file

@ -229,6 +229,28 @@ async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY):
- `RebuildMode.FULL` - Re-convert from source files, re-chunk, re-embed (default)
- `RebuildMode.RECHUNK` - Re-chunk from existing document content, re-embed
- `RebuildMode.EMBED_ONLY` - Keep existing chunks, only regenerate embeddings
- `RebuildMode.TITLE_ONLY` - Generate titles for untitled documents (no re-chunking or re-embedding)
### Generating Titles
Generate a title for an existing document on demand:
```python
title = await client.generate_title(doc)
if title:
await client.update_document(document_id=doc.id, title=title)
```
Uses the same two-tier approach as automatic ingestion: structural extraction from DoclingDocument metadata first, with LLM fallback via `processing.title_model`. Unlike ingestion, this method does not catch exceptions — if the LLM call fails, the error propagates.
To batch-generate titles for all untitled documents, use `RebuildMode.TITLE_ONLY`:
```python
async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY):
print(f"Generated title for {doc_id}")
```
See [Automatic Title Generation](configuration/processing.md#automatic-title-generation) for configuration details.
## Maintenance

9
haiku_rag_slim/haiku/rag/app.py
View file

@ -261,14 +261,18 @@ class HaikuRAGApp: # pragma: no cover
for doc in documents:
self._rich_print_document(doc, truncate=True)
async def add_document_from_text(self, text: str, metadata: dict | None = None):
async def add_document_from_text(
self, text: str, title: str | None = None, metadata: dict | None = None
):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
doc = await self.client.create_document(text, metadata=metadata)
doc = await self.client.create_document(
text, title=title, metadata=metadata
)
self._rich_print_document(doc, truncate=True)
self.console.print(
f"[bold green]Document {doc.id} added successfully.[/bold green]"
@ -534,6 +538,7 @@ class HaikuRAGApp: # pragma: no cover
RebuildMode.FULL: "full rebuild",
RebuildMode.RECHUNK: "rechunk",
RebuildMode.EMBED_ONLY: "embed only",
RebuildMode.TITLE_ONLY: "title only",
}[mode]
self.console.print(

23
haiku_rag_slim/haiku/rag/cli.py
View file

@ -200,6 +200,11 @@ def add_document_text( # pragma: no cover
text: str = typer.Argument(
help="The text content of the document to add",
),
title: str | None = typer.Option(
None,
"--title",
help="Optional title for the document",
),
meta: list[str] | None = typer.Option(
None,
"--meta",
@ -214,7 +219,9 @@ def add_document_text( # pragma: no cover
):
app = create_app(db)
metadata = _parse_meta_options(meta)
asyncio.run(app.add_document_from_text(text=text, metadata=metadata or None))
asyncio.run(
app.add_document_from_text(text=text, title=title, metadata=metadata or None)
)
@_cli.command("add-src", help="Add a document from a file path, directory, or URL")
@ -469,17 +476,27 @@ def rebuild(
"--rechunk",
help="Re-chunk from existing content without accessing source files",
),
title_only: bool = typer.Option(
False,
"--title-only",
help="Only generate titles for documents without one",
),
):
from haiku.rag.client import RebuildMode
if embed_only and rechunk:
typer.echo("Error: --embed-only and --rechunk are mutually exclusive")
exclusive = sum([embed_only, rechunk, title_only])
if exclusive > 1:
typer.echo(
"Error: --embed-only, --rechunk, and --title-only are mutually exclusive"
)
raise typer.Exit(1)
if embed_only: # pragma: no cover
mode = RebuildMode.EMBED_ONLY
elif rechunk: # pragma: no cover
mode = RebuildMode.RECHUNK
elif title_only: # pragma: no cover
mode = RebuildMode.TITLE_ONLY
else: # pragma: no cover
mode = RebuildMode.FULL

167
haiku_rag_slim/haiku/rag/client.py
View file

@ -46,6 +46,7 @@ class RebuildMode(Enum):
FULL = "full" # Re-convert from source, re-chunk, re-embed
RECHUNK = "rechunk" # Re-chunk from existing content, re-embed
EMBED_ONLY = "embed_only" # Keep chunks, only regenerate embeddings
TITLE_ONLY = "title_only" # Only generate titles for untitled documents
@dataclass
@ -248,6 +249,108 @@ class HaikuRAG:
return result
# =========================================================================
# Title Generation
# =========================================================================
def _extract_structural_title(
self, docling_document: "DoclingDocument"
) -> str | None:
"""Extract a title from DoclingDocument structural metadata.
Priority: FURNITURE TITLE > BODY TITLE > first SECTION_HEADER.
"""
from docling_core.types.doc.document import ContentLayer
from docling_core.types.doc.labels import DocItemLabel
furniture_title = None
body_title = None
first_section_header = None
for item in docling_document.texts:
if item.label == DocItemLabel.TITLE:
text = item.text.strip()
if not text:
continue
if item.content_layer == ContentLayer.FURNITURE:
furniture_title = text
elif body_title is None:
body_title = text
elif (
item.label == DocItemLabel.SECTION_HEADER
and first_section_header is None
):
text = item.text.strip()
if text:
first_section_header = text
return furniture_title or body_title or first_section_header
async def _generate_title_with_llm(self, content: str) -> str | None:
"""Generate a title using LLM from document content."""
from pydantic_ai import Agent
from haiku.rag.utils import get_model
truncated = content[:2000]
model = get_model(self._config.processing.title_model, self._config)
agent: Agent[None, str] = Agent(
model=model,
output_type=str,
instructions=(
"Generate a concise, descriptive title for the following document. "
"The title should be at most 10 words. "
"Return ONLY the title text, nothing else."
),
)
result = await agent.run(truncated)
title = result.output.strip()
return title if title else None
async def _resolve_title(
self,
docling_document: "DoclingDocument",
content: str,
) -> str | None:
"""Auto-generate a title from document structure or LLM.
Returns None if auto_title is disabled or generation fails.
"""
if not self._config.processing.auto_title:
return None
structural = self._extract_structural_title(docling_document)
if structural:
return structural
try:
return await self._generate_title_with_llm(content)
except Exception:
logger.warning(
"LLM title generation failed during ingestion", exc_info=True
)
return None
async def generate_title(self, document: Document) -> str | None:
"""Generate a title for a document.
Attempts structural extraction from the stored DoclingDocument,
then falls back to LLM generation. Bypasses the auto_title config
since this is an explicit call.
Does NOT update the document — caller decides.
"""
docling_doc = document.get_docling_document()
content = document.content or ""
if docling_doc is not None:
structural = self._extract_structural_title(docling_doc)
if structural:
return structural
return await self._generate_title_with_llm(content)
async def _store_document_with_chunks(
self,
document: Document,
@ -383,6 +486,9 @@ class HaikuRAG:
# The original content is preserved in docling_document
stored_content = docling_document.export_to_markdown()
if title is None:
title = await self._resolve_title(docling_document, stored_content)
# Create document model
document = Document(
content=stored_content,
@ -420,8 +526,12 @@ class HaikuRAG:
Returns:
The created Document instance.
"""
content = docling_document.export_to_markdown()
if title is None:
title = await self._resolve_title(docling_document, content)
document = Document(
content=docling_document.export_to_markdown(),
content=content,
uri=uri,
title=title,
metadata=metadata or {},
@ -556,9 +666,11 @@ class HaikuRAG:
chunks = await self.chunk(docling_document)
embedded_chunks = await embed_chunks(chunks, self._config)
stored_content = docling_document.export_to_markdown()
if existing_doc:
# Update existing document and rechunk
existing_doc.content = docling_document.export_to_markdown()
existing_doc.content = stored_content
existing_doc.metadata = metadata
existing_doc.docling_document = compress_json(
docling_document.model_dump_json()
@ -566,13 +678,19 @@ class HaikuRAG:
existing_doc.docling_version = docling_document.version
if title is not None:
existing_doc.title = title
elif existing_doc.title is None:
existing_doc.title = await self._resolve_title(
docling_document, stored_content
)
return await self._update_document_with_chunks(
existing_doc, embedded_chunks
)
else:
# Create new document
if title is None:
title = await self._resolve_title(docling_document, stored_content)
document = Document(
content=docling_document.export_to_markdown(),
content=stored_content,
uri=uri,
title=title,
metadata=metadata,
@ -665,9 +783,11 @@ class HaikuRAG:
# Merge metadata with contentType and md5
metadata.update({"contentType": content_type, "md5": md5_hash})
stored_content = docling_document.export_to_markdown()
if existing_doc:
# Update existing document and rechunk
existing_doc.content = docling_document.export_to_markdown()
existing_doc.content = stored_content
existing_doc.metadata = metadata
existing_doc.docling_document = compress_json(
docling_document.model_dump_json()
@ -675,13 +795,19 @@ class HaikuRAG:
existing_doc.docling_version = docling_document.version
if title is not None:
existing_doc.title = title
elif existing_doc.title is None:
existing_doc.title = await self._resolve_title(
docling_document, stored_content
)
return await self._update_document_with_chunks(
existing_doc, embedded_chunks
)
else:
# Create new document
if title is None:
title = await self._resolve_title(docling_document, stored_content)
document = Document(
content=docling_document.export_to_markdown(),
content=stored_content,
uri=url,
title=title,
metadata=metadata,
@ -1519,6 +1645,7 @@ class HaikuRAG:
- FULL: Re-convert from source files, re-chunk, re-embed (default)
- RECHUNK: Re-chunk from existing content, re-embed (no source access)
- EMBED_ONLY: Keep existing chunks, only regenerate embeddings
- TITLE_ONLY: Only generate titles for untitled documents
Yields:
The ID of the document currently being processed.
@ -1529,7 +1656,10 @@ class HaikuRAG:
documents = await self.list_documents(include_content=True)
if mode == RebuildMode.EMBED_ONLY:
if mode == RebuildMode.TITLE_ONLY:
async for doc_id in self._rebuild_title_only(documents):
yield doc_id
elif mode == RebuildMode.EMBED_ONLY:
async for doc_id in self._rebuild_embed_only(documents):
yield doc_id
elif mode == RebuildMode.RECHUNK:
@ -1550,6 +1680,26 @@ class HaikuRAG:
except Exception:
pass
async def _rebuild_title_only(
self, documents: list[Document]
) -> AsyncGenerator[str, None]:
"""Generate titles for documents that don't have one."""
for doc in documents:
if doc.title is not None:
continue
assert doc.id is not None
try:
title = await self.generate_title(doc)
except Exception:
logger.warning(
"Failed to generate title for document %s", doc.id, exc_info=True
)
continue
if title is not None:
doc.title = title
await self.document_repository.update(doc)
yield doc.id
async def _rebuild_embed_only(
self, documents: list[Document]
) -> AsyncGenerator[str, None]:
@ -1884,6 +2034,11 @@ class HaikuRAG:
pic_desc = self._config.processing.conversion_options.picture_description
if pic_desc.enabled and pic_desc.model.provider == "ollama":
required_models.add(pic_desc.model.name)
if (
self._config.processing.auto_title
and self._config.processing.title_model.provider == "ollama"
):
required_models.add(self._config.processing.title_model.name)
if not required_models:
return

8
haiku_rag_slim/haiku/rag/config/models.py
View file

@ -156,6 +156,14 @@ class ProcessingConfig(BaseModel):
chunking_merge_peers: bool = True
chunking_use_markdown_tables: bool = False
conversion_options: ConversionOptions = Field(default_factory=ConversionOptions)
auto_title: bool = False
title_model: ModelConfig = Field(
default_factory=lambda: ModelConfig(
provider="ollama",
name="gpt-oss",
enable_thinking=False,
)
)
class SearchConfig(BaseModel):

393
tests/test_title_generation.py Normal file
View file

@ -0,0 +1,393 @@
import random
import pytest
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.config import AppConfig
from haiku.rag.config.models import ProcessingConfig
from haiku.rag.embeddings import EmbedderWrapper
@pytest.fixture(autouse=True)
def mock_embedder(monkeypatch):
"""Monkeypatch the embedder to return deterministic vectors."""
async def fake_embed_query(self, text):
random.seed(hash(text) % (2**32))
return [random.random() for _ in range(2560)]
async def fake_embed_documents(self, texts):
result = []
for t in texts:
random.seed(hash(t) % (2**32))
result.append([random.random() for _ in range(2560)])
return result
monkeypatch.setattr(EmbedderWrapper, "embed_query", fake_embed_query)
monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents)
# =========================================================================
# Structural title extraction
# =========================================================================
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):
"""TITLE on FURNITURE layer (HTML <title>) is extracted."""
doc = DoclingDocument(name="test")
doc.add_text(
label=DocItemLabel.TITLE,
text="Website Page Title",
content_layer=ContentLayer.FURNITURE,
)
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"
def test_body_title(self, tmp_path):
"""TITLE on BODY layer (h1, PDF title) is extracted."""
doc = DoclingDocument(name="test")
doc.add_text(
label=DocItemLabel.TITLE,
text="Document Heading",
content_layer=ContentLayer.BODY,
)
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"
def test_section_header_fallback(self, tmp_path):
"""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"
def test_no_title_or_headers(self, tmp_path):
"""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
def test_furniture_title_preferred_over_body_title(self, tmp_path):
"""FURNITURE TITLE takes priority over BODY TITLE."""
doc = DoclingDocument(name="test")
doc.add_text(
label=DocItemLabel.TITLE,
text="Body H1 Title",
content_layer=ContentLayer.BODY,
)
doc.add_text(
label=DocItemLabel.TITLE,
text="HTML Page Title",
content_layer=ContentLayer.FURNITURE,
)
client = self._make_client(tmp_path)
result = client._extract_structural_title(doc)
assert result == "HTML Page Title"
def test_whitespace_stripped(self, tmp_path):
"""Whitespace is stripped from extracted titles."""
doc = DoclingDocument(name="test")
doc.add_text(
label=DocItemLabel.TITLE,
text=" Padded Title ",
content_layer=ContentLayer.BODY,
)
client = self._make_client(tmp_path)
result = client._extract_structural_title(doc)
assert result == "Padded Title"
def test_empty_title_text_skipped(self, tmp_path):
"""Empty or whitespace-only TITLE text is skipped."""
doc = DoclingDocument(name="test")
doc.add_text(
label=DocItemLabel.TITLE,
text=" ",
content_layer=ContentLayer.BODY,
)
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"
# =========================================================================
# _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):
"""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")
assert result is None
@pytest.mark.asyncio
async def test_structural_title_extracted(self, tmp_path):
"""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")
assert result == "Auto Extracted Title"
@pytest.mark.asyncio
async def test_llm_failure_returns_none(self, tmp_path, 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):
raise RuntimeError("LLM is down")
monkeypatch.setattr(HaikuRAG, "_generate_title_with_llm", exploding_llm)
result = await client._resolve_title(doc, "some content")
assert result is None
# =========================================================================
# Integration: create_document with auto_title
# =========================================================================
class TestCreateDocumentAutoTitle:
@pytest.mark.asyncio
async def test_auto_title_from_structural(self, temp_db_path):
"""create_document with auto_title=True extracts title from docling."""
config = AppConfig(processing=ProcessingConfig(auto_title=True))
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await client.create_document(
"# My Document\n\nSome content here.", uri="test://auto-title"
)
assert doc.title == "My Document"
@pytest.mark.asyncio
async def test_auto_title_disabled(self, temp_db_path):
"""create_document with auto_title=False leaves title as None."""
config = AppConfig(processing=ProcessingConfig(auto_title=False))
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await client.create_document(
"# My Document\n\nSome content here.", uri="test://no-auto-title"
)
assert doc.title is None
@pytest.mark.asyncio
async def test_explicit_title_not_overridden(self, temp_db_path):
"""Explicit title is never overridden by auto-generation."""
config = AppConfig(processing=ProcessingConfig(auto_title=True))
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await client.create_document(
"# Auto Title\n\nSome content here.",
uri="test://explicit-title",
title="My Explicit Title",
)
assert doc.title == "My Explicit Title"
# =========================================================================
# Integration: import_document with auto_title
# =========================================================================
class TestImportDocumentAutoTitle:
@pytest.mark.asyncio
async def test_auto_title_from_structural(self, temp_db_path):
"""import_document with auto_title=True extracts title from docling."""
config = AppConfig(processing=ProcessingConfig(auto_title=True))
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
docling_doc = await client.convert("# Imported Doc\n\nContent here.")
chunks = await client.chunk(docling_doc)
doc = await client.import_document(
docling_doc, chunks, uri="test://import-auto-title"
)
assert doc.title == "Imported Doc"
@pytest.mark.asyncio
async def test_explicit_title_preserved(self, temp_db_path):
"""import_document explicit title is not overridden."""
config = AppConfig(processing=ProcessingConfig(auto_title=True))
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
docling_doc = await client.convert("# Auto Title\n\nContent here.")
chunks = await client.chunk(docling_doc)
doc = await client.import_document(
docling_doc,
chunks,
uri="test://import-explicit",
title="Keep This Title",
)
assert doc.title == "Keep This Title"
# =========================================================================
# generate_title() public method
# =========================================================================
class TestGenerateTitle:
@pytest.mark.asyncio
async def test_structural_title(self, temp_db_path):
"""generate_title extracts structural title from document."""
config = AppConfig(processing=ProcessingConfig(auto_title=True))
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await client.create_document(
"# Great Heading\n\nSome content.", uri="test://gen-title"
)
title = await client.generate_title(doc)
assert title == "Great Heading"
@pytest.mark.asyncio
async def test_no_structural_title_no_llm(self, temp_db_path):
"""generate_title raises when no structural title and LLM unavailable."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
"Just plain text without headings.",
uri="test://gen-no-title",
format="plain",
)
with pytest.raises(RuntimeError):
await client.generate_title(doc)
@pytest.mark.asyncio
async def test_bypasses_auto_title_config(self, temp_db_path):
"""generate_title works even when auto_title is False."""
config = AppConfig(processing=ProcessingConfig(auto_title=False))
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await client.create_document(
"# Heading Present\n\nBody text.", uri="test://gen-bypass"
)
title = await client.generate_title(doc)
assert title == "Heading Present"
# =========================================================================
# rebuild --title-only
# =========================================================================
class TestRebuildTitleOnly:
@pytest.mark.asyncio
async def test_generates_titles_for_untitled_docs(self, temp_db_path):
"""TITLE_ONLY mode generates titles for documents without one."""
from haiku.rag.client import RebuildMode
config = AppConfig(processing=ProcessingConfig(auto_title=True))
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc1 = await client.create_document(
"# Doc With Heading\n\nContent.",
uri="test://titled",
)
assert doc1.title == "Doc With Heading"
# Simulate an untitled doc
doc1.title = None
await client.document_repository.update(doc1)
doc2 = await client.create_document(
"# Another Heading\n\nMore content.",
uri="test://also-titled",
)
assert doc2.title == "Another Heading"
processed_ids = []
async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY):
processed_ids.append(doc_id)
assert len(processed_ids) == 1
assert doc1.id in processed_ids
updated_doc1 = await client.get_document_by_id(doc1.id)
assert updated_doc1 is not None
assert updated_doc1.title == "Doc With Heading"
@pytest.mark.asyncio
async def test_skips_already_titled_docs(self, temp_db_path):
"""TITLE_ONLY mode skips documents that already have titles."""
from haiku.rag.client import RebuildMode
config = AppConfig(processing=ProcessingConfig(auto_title=True))
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document(
"# Has Title\n\nContent.",
uri="test://has-title",
)
processed_ids = []
async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY):
processed_ids.append(doc_id)
assert len(processed_ids) == 0
@pytest.mark.asyncio
async def test_continues_on_per_document_failure(self, temp_db_path, monkeypatch):
"""TITLE_ONLY mode continues when generate_title fails for a document."""
from haiku.rag.client import RebuildMode
config = AppConfig(processing=ProcessingConfig(auto_title=True))
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc1 = await client.create_document(
"# First Heading\n\nContent.",
uri="test://first",
)
doc2 = await client.create_document(
"# Second Heading\n\nContent.",
uri="test://second",
)
# Clear both titles
doc1.title = None
await client.document_repository.update(doc1)
doc2.title = None
await client.document_repository.update(doc2)
# Make generate_title fail for the first doc, succeed for the second
original = HaikuRAG.generate_title
call_count = 0
async def flaky_generate(self, document):
nonlocal call_count
call_count += 1
if call_count == 1:
raise RuntimeError("LLM failed")
return await original(self, document)
monkeypatch.setattr(HaikuRAG, "generate_title", flaky_generate)
processed_ids = []
async for doc_id in client.rebuild_database(mode=RebuildMode.TITLE_ONLY):
processed_ids.append(doc_id)
# Only the second doc should have been processed
assert len(processed_ids) == 1