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

add generate_title(), rebuild --title-only, and --title on add

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-26 08:40:07 +02:00
parent ae3469b282
commit 49f9843db3
No known key found for this signature in database
4 changed files with 167 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

41
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
@ -291,7 +292,6 @@ class HaikuRAG:
from haiku.rag.utils import get_model
# Truncate content to limit token usage
truncated = content[:2000]
model = get_model(self._config.processing.title_model, self._config)
@ -337,6 +337,25 @@ class HaikuRAG:
return await self._generate_title_with_llm(content)
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,
@ -1629,6 +1648,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.
@ -1639,7 +1659,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:
@ -1660,6 +1683,20 @@ 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
title = await self.generate_title(doc)
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]:

101
tests/test_title_generation.py
View file

@ -256,3 +256,104 @@ class TestImportDocumentAutoTitle:
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 returns None 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",
)
title = await client.generate_title(doc)
# LLM blocked by ALLOW_MODEL_REQUESTS=False, so falls back to None
assert title is None
@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