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

Merge pull request #69 from ggozad/feat/display-title

Browse source 
Add optional  "title" to Documents. Replaces "uri" where appropriate.
This commit is contained in:
Yiorgis Gozadinos 2025-09-22 13:53:20 +03:00 committed by GitHub
commit a1cd64de2f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
24 changed files with 310 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
docs/agents.md
View file

@ -13,7 +13,8 @@ The simple QA agent answers a single question using the knowledge base. It retri
Key points:
- Uses a single `search_documents` tool to fetch relevant chunks
- Can be run with or without inline citations in the prompt
- Can be run with or without inline citations in the prompt (citations prefer
document titles when present, otherwise URIs)
- Returns a plain string answer
Python usage:

4
docs/cli.md
View file

@ -33,6 +33,9 @@ From file or URL:
```bash
haiku-rag add-src /path/to/document.pdf
haiku-rag add-src https://example.com/article.html
# Optionally set a human‑readable title stored in the DB schema
haiku-rag add-src /mnt/data/doc1.pdf --title "Q3 Financial Report"
```
!!! note
@ -83,6 +86,7 @@ haiku-rag ask "Who is the author of haiku.rag?" --cite
```
The QA agent will search your documents for relevant information and provide a comprehensive answer. With `--cite`, responses include citations showing which documents were used.
When available, citations use the document title; otherwise they fall back to the URI.
## Research

2
docs/index.md
View file

@ -15,6 +15,7 @@
- **Extended file format support**: Parse 40+ file formats including PDF, DOCX, HTML, Markdown, code files and more. Or add a URL!
- **MCP server**: Exposes functionality as MCP tools
- **CLI commands**: Access all functionality from your terminal
- Add sources from text, files, or URLs, optionally with a human‑readable title
- **Python client**: Call `haiku.rag` from your own python applications
## Quick Start
@ -42,6 +43,7 @@ async with HaikuRAG("database.lancedb") as client:
Or use the CLI:
```bash
haiku-rag add "Your document content"
haiku-rag add-src /path/to/document.pdf --title "Q3 Financial Report"
haiku-rag search "query"
haiku-rag ask "Who is the author of haiku.rag?"
haiku-rag migrate old_database.sqlite # Migrate from SQLite

12
docs/python.md
View file

@ -23,6 +23,7 @@ From text:
doc = await client.create_document(
content="Your document content here",
uri="doc://example",
title="My Example Document", # optional human‑readable title
metadata={"source": "manual", "topic": "example"}
)
```
@ -54,12 +55,16 @@ doc = await client.create_document(
From file:
```python
doc = await client.create_document_from_source("path/to/document.pdf")
doc = await client.create_document_from_source(
"path/to/document.pdf", title="Project Brief"
)
```
From URL:
```python
doc = await client.create_document_from_source("https://example.com/article.html")
doc = await client.create_document_from_source(
"https://example.com/article.html", title="Example Article"
)
```
### Retrieving Documents
@ -159,6 +164,7 @@ for chunk, relevance_score in results:
print(f"Content: {chunk.content}")
print(f"From document: {chunk.document_id}")
print(f"Document URI: {chunk.document_uri}")
print(f"Document Title: {chunk.document_title}") # when available
print(f"Document metadata: {chunk.document_meta}")
```
@ -201,7 +207,7 @@ answer = await client.ask("Who is the author of haiku.rag?", cite=True)
print(answer)
```
The QA agent will search your documents for relevant information and use the configured LLM to generate a comprehensive answer. With `cite=True`, responses include citations showing which documents were used as sources.
The QA agent will search your documents for relevant information and use the configured LLM to generate a comprehensive answer. With `cite=True`, responses include citations showing which documents were used as sources. Citations prefer the document title when present, otherwise they use the URI.
The QA provider and model can be configured via environment variables (see [Configuration](configuration.md)).

17
src/haiku/rag/app.py
View file

@ -39,9 +39,9 @@ class HaikuRAGApp:
f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]"
)
async def add_document_from_source(self, source: str):
async def add_document_from_source(self, source: str, title: str | None = None):
async with HaikuRAG(db_path=self.db_path) as self.client:
doc = await self.client.create_document_from_source(source)
doc = await self.client.create_document_from_source(source, title=title)
self._rich_print_document(doc, truncate=True)
self.console.print(
f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]"
@ -252,8 +252,16 @@ class HaikuRAGApp:
content = Markdown(content)
else:
content = Markdown(doc.content)
title_part = (
f" [repr.attrib_name]title[/repr.attrib_name]: {doc.title}"
if doc.title
else ""
)
self.console.print(
f"[repr.attrib_name]id[/repr.attrib_name]: {doc.id} [repr.attrib_name]uri[/repr.attrib_name]: {doc.uri} [repr.attrib_name]meta[/repr.attrib_name]: {doc.metadata}"
f"[repr.attrib_name]id[/repr.attrib_name]: {doc.id} "
f"[repr.attrib_name]uri[/repr.attrib_name]: {doc.uri}"
+ title_part
+ f" [repr.attrib_name]meta[/repr.attrib_name]: {doc.metadata}"
)
self.console.print(
f"[repr.attrib_name]created at[/repr.attrib_name]: {doc.created_at} [repr.attrib_name]updated at[/repr.attrib_name]: {doc.updated_at}"
@ -272,6 +280,9 @@ class HaikuRAGApp:
if chunk.document_uri:
self.console.print("[repr.attrib_name]document uri[/repr.attrib_name]:")
self.console.print(chunk.document_uri)
if chunk.document_title:
self.console.print("[repr.attrib_name]document title[/repr.attrib_name]:")
self.console.print(chunk.document_title)
if chunk.document_meta:
self.console.print("[repr.attrib_name]document meta[/repr.attrib_name]:")
self.console.print(chunk.document_meta)

7
src/haiku/rag/cli.py
View file

@ -160,6 +160,11 @@ def add_document_src(
help="The file path or URL of the document to add",
autocompletion=complete_local_paths,
),
title: str | None = typer.Option(
None,
"--title",
help="Optional human-readable title to store with the document",
),
db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db",
@ -169,7 +174,7 @@ def add_document_src(
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(app.add_document_from_source(source=source))
asyncio.run(app.add_document_from_source(source=source, title=title))
@cli.command("get", help="Get and display a document by its ID")

84
src/haiku/rag/client.py
View file

@ -50,6 +50,7 @@ class HaikuRAG:
self,
docling_document,
uri: str | None = None,
title: str | None = None,
metadata: dict | None = None,
chunks: list[Chunk] | None = None,
) -> Document:
@ -58,6 +59,7 @@ class HaikuRAG:
document = Document(
content=content,
uri=uri,
title=title,
metadata=metadata or {},
)
return await self.document_repository._create_with_docling(
@ -68,6 +70,7 @@ class HaikuRAG:
self,
content: str,
uri: str | None = None,
title: str | None = None,
metadata: dict | None = None,
chunks: list[Chunk] | None = None,
) -> Document:
@ -88,6 +91,7 @@ class HaikuRAG:
document = Document(
content=content,
uri=uri,
title=title,
metadata=metadata or {},
)
return await self.document_repository._create_with_docling(
@ -95,7 +99,7 @@ class HaikuRAG:
)
async def create_document_from_source(
self, source: str | Path, metadata: dict = {}
self, source: str | Path, title: str | None = None, metadata: dict | None = None
) -> Document:
"""Create or update a document from a file path or URL.
@ -116,11 +120,16 @@ class HaikuRAG:
httpx.RequestError: If URL request fails
"""
# Normalize metadata
metadata = metadata or {}
# Check if it's a URL
source_str = str(source)
parsed_url = urlparse(source_str)
if parsed_url.scheme in ("http", "https"):
return await self._create_or_update_document_from_url(source_str, metadata)
return await self._create_or_update_document_from_url(
source_str, title=title, metadata=metadata
)
elif parsed_url.scheme == "file":
# Handle file:// URI by converting to path
source_path = Path(parsed_url.path)
@ -136,37 +145,51 @@ class HaikuRAG:
uri = source_path.absolute().as_uri()
md5_hash = hashlib.md5(source_path.read_bytes()).hexdigest()
# Check if document already exists
existing_doc = await self.get_document_by_uri(uri)
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
# MD5 unchanged, return existing document
return existing_doc
docling_document = FileReader.parse_file(source_path)
# Get content type from file extension
# Get content type from file extension (do before early return)
content_type, _ = mimetypes.guess_type(str(source_path))
if not content_type:
content_type = "application/octet-stream"
# Merge metadata with contentType and md5
metadata.update({"contentType": content_type, "md5": md5_hash})
# Check if document already exists
existing_doc = await self.get_document_by_uri(uri)
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
# MD5 unchanged; update title/metadata if provided
updated = False
if title is not None and title != existing_doc.title:
existing_doc.title = title
updated = True
if metadata:
existing_doc.metadata = {**(existing_doc.metadata or {}), **metadata}
updated = True
if updated:
return await self.document_repository.update(existing_doc)
return existing_doc
# Parse file only when content changed or new document
docling_document = FileReader.parse_file(source_path)
if existing_doc:
# Update existing document
existing_doc.content = docling_document.export_to_markdown()
existing_doc.metadata = metadata
if title is not None:
existing_doc.title = title
return await self.document_repository._update_with_docling(
existing_doc, docling_document
)
else:
# Create new document using DoclingDocument
return await self._create_document_with_docling(
docling_document=docling_document, uri=uri, metadata=metadata
docling_document=docling_document,
uri=uri,
title=title,
metadata=metadata,
)
async def _create_or_update_document_from_url(
self, url: str, metadata: dict = {}
self, url: str, title: str | None = None, metadata: dict | None = None
) -> Document:
"""Create or update a document from a URL by downloading and parsing the content.
@ -186,20 +209,35 @@ class HaikuRAG:
ValueError: If the content cannot be parsed
httpx.RequestError: If URL request fails
"""
metadata = metadata or {}
async with httpx.AsyncClient() as client:
response = await client.get(url)
response.raise_for_status()
md5_hash = hashlib.md5(response.content).hexdigest()
# Get content type early (used for potential no-op update)
content_type = response.headers.get("content-type", "").lower()
# Check if document already exists
existing_doc = await self.get_document_by_uri(url)
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
# MD5 unchanged, return existing document
# MD5 unchanged; update title/metadata if provided
updated = False
if title is not None and title != existing_doc.title:
existing_doc.title = title
updated = True
metadata.update({"contentType": content_type, "md5": md5_hash})
if metadata:
existing_doc.metadata = {
**(existing_doc.metadata or {}),
**metadata,
}
updated = True
if updated:
return await self.document_repository.update(existing_doc)
return existing_doc
# Get content type to determine file extension
content_type = response.headers.get("content-type", "").lower()
file_extension = self._get_extension_from_content_type_or_url(
url, content_type
)
@ -226,12 +264,17 @@ class HaikuRAG:
if existing_doc:
existing_doc.content = docling_document.export_to_markdown()
existing_doc.metadata = metadata
if title is not None:
existing_doc.title = title
return await self.document_repository._update_with_docling(
existing_doc, docling_document
)
else:
return await self._create_document_with_docling(
docling_document=docling_document, uri=url, metadata=metadata
docling_document=docling_document,
uri=url,
title=title,
metadata=metadata,
)
def _get_extension_from_content_type_or_url(
@ -416,6 +459,7 @@ class HaikuRAG:
content="".join(combined_content_parts),
metadata=original_chunk.metadata,
document_uri=original_chunk.document_uri,
document_title=original_chunk.document_title,
document_meta=original_chunk.document_meta,
)
@ -522,7 +566,7 @@ class HaikuRAG:
# Try to re-create from source (this creates the document with chunks)
new_doc = await self.create_document_from_source(
doc.uri, doc.metadata or {}
source=doc.uri, metadata=doc.metadata or {}
)
assert new_doc.id is not None, "New document ID should not be None"

24
src/haiku/rag/mcp.py
View file

@ -17,6 +17,7 @@ class DocumentResult(BaseModel):
id: str | None
content: str
uri: str | None = None
title: str | None = None
metadata: dict[str, Any] = {}
created_at: str
updated_at: str
@ -28,13 +29,15 @@ def create_mcp_server(db_path: Path) -> FastMCP:
@mcp.tool()
async def add_document_from_file(
file_path: str, metadata: dict[str, Any] | None = None
file_path: str,
metadata: dict[str, Any] | None = None,
title: str | None = None,
) -> str | None:
"""Add a document to the RAG system from a file path."""
try:
async with HaikuRAG(db_path) as rag:
document = await rag.create_document_from_source(
Path(file_path), metadata or {}
Path(file_path), title=title, metadata=metadata or {}
)
return document.id
except Exception:
@ -42,24 +45,31 @@ def create_mcp_server(db_path: Path) -> FastMCP:
@mcp.tool()
async def add_document_from_url(
url: str, metadata: dict[str, Any] | None = None
url: str, metadata: dict[str, Any] | None = None, title: str | None = None
) -> str | None:
"""Add a document to the RAG system from a URL."""
try:
async with HaikuRAG(db_path) as rag:
document = await rag.create_document_from_source(url, metadata or {})
document = await rag.create_document_from_source(
url, title=title, metadata=metadata or {}
)
return document.id
except Exception:
return None
@mcp.tool()
async def add_document_from_text(
content: str, uri: str | None = None, metadata: dict[str, Any] | None = None
content: str,
uri: str | None = None,
metadata: dict[str, Any] | None = None,
title: str | None = None,
) -> str | None:
"""Add a document to the RAG system from text content."""
try:
async with HaikuRAG(db_path) as rag:
document = await rag.create_document(content, uri, metadata or {})
document = await rag.create_document(
content, uri, title=title, metadata=metadata or {}
)
return document.id
except Exception:
return None
@ -102,6 +112,7 @@ def create_mcp_server(db_path: Path) -> FastMCP:
id=document.id,
content=document.content,
uri=document.uri,
title=document.title,
metadata=document.metadata,
created_at=str(document.created_at),
updated_at=str(document.updated_at),
@ -123,6 +134,7 @@ def create_mcp_server(db_path: Path) -> FastMCP:
id=doc.id,
content=doc.content,
uri=doc.uri,
title=doc.title,
metadata=doc.metadata,
created_at=str(doc.created_at),
updated_at=str(doc.updated_at),

6
src/haiku/rag/qa/agent.py
View file

@ -12,7 +12,9 @@ from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIO
class SearchResult(BaseModel):
content: str = Field(description="The document text content")
score: float = Field(description="Relevance score (higher is more relevant)")
document_uri: str = Field(description="Source URI/path of the document")
document_uri: str = Field(
description="Source title (if available) or URI/path of the document"
)
class Dependencies(BaseModel):
@ -59,7 +61,7 @@ class QuestionAnswerAgent:
SearchResult(
content=chunk.content,
score=score,
document_uri=chunk.document_uri or "",
document_uri=(chunk.document_title or chunk.document_uri or ""),
)
for chunk, score in expanded_results
]

4
src/haiku/rag/qa/prompts.py
View file

@ -44,9 +44,9 @@ Guidelines:
Citation Format:
After your answer, include a "Citations:" section that lists:
- The document URI from each search result used
- The document title (if available) or URI from each search result used
- A brief excerpt (first 50-100 characters) of the content that supported your answer
- Format: "Citations:\n- [document_uri]: [content_excerpt]..."
- Format: "Citations:\n- [document title or URI]: [content_excerpt]..."
Example response format:
[Your answer here]

4
src/haiku/rag/research/models.py
View file

@ -19,8 +19,8 @@ class SearchAnswer(BaseModel):
)
sources: list[str] = Field(
description=(
"Document URIs corresponding to the snippets actually used in the"
" answer (one URI per snippet; omit if none)"
"Document titles (if available) or URIs corresponding to the"
" snippets actually used in the answer (one per snippet; omit if none)"
),
default_factory=list,
)

4
src/haiku/rag/research/nodes/search.py
View file

@ -59,7 +59,9 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
{
"text": chunk.content,
"score": score,
"document_uri": (chunk.document_uri or ""),
"document_uri": (
chunk.document_title or chunk.document_uri or ""
),
}
for chunk, score in expanded
]

7
src/haiku/rag/research/prompts.py
View file

@ -27,13 +27,14 @@ Tasks:
Tool usage:
- Always call search_and_answer before drafting any answer.
- The tool returns snippets with verbatim `text`, a relevance `score`, and the
originating `document_uri`.
originating document identifier (document title if available, otherwise URI).
- You may call the tool multiple times to refine or broaden context, but do not
exceed 3 total calls. Favor precision over volume.
- Use scores to prioritize evidence, but include only the minimal subset of
snippet texts (verbatim) in SearchAnswer.context (typically 1‑4).
- Set SearchAnswer.sources to the corresponding document_uris for the snippets
you used (one URI per snippet; same order as context). Context must be text‑only.
- Set SearchAnswer.sources to the corresponding document identifiers for the
snippets you used (title if available, otherwise URI; one per snippet; same
order as context). Context must be text‑only.
- If no relevant information is found, clearly say so and return an empty
context list and sources list.

1
src/haiku/rag/store/engine.py
View file

@ -19,6 +19,7 @@ class DocumentRecord(LanceModel):
id: str = Field(default_factory=lambda: str(uuid4()))
content: str
uri: str | None = None
title: str | None = None
metadata: str = Field(default="{}")
created_at: str = Field(default_factory=lambda: "")
updated_at: str = Field(default_factory=lambda: "")

1
src/haiku/rag/store/models/chunk.py
View file

@ -12,5 +12,6 @@ class Chunk(BaseModel):
metadata: dict = {}
order: int = 0
document_uri: str | None = None
document_title: str | None = None
document_meta: dict = {}
embedding: list[float] | None = None

1
src/haiku/rag/store/models/document.py
View file

@ -11,6 +11,7 @@ class Document(BaseModel):
id: str | None = None
content: str
uri: str | None = None
title: str | None = None
metadata: dict = {}
created_at: datetime = Field(default_factory=datetime.now)
updated_at: datetime = Field(default_factory=datetime.now)

4
src/haiku/rag/store/repositories/chunk.py
View file

@ -317,6 +317,7 @@ class ChunkRepository:
)
doc_uri = doc_results[0].uri if doc_results else None
doc_title = doc_results[0].title if doc_results else None
doc_meta = doc_results[0].metadata if doc_results else "{}"
chunks: list[Chunk] = []
@ -330,6 +331,7 @@ class ChunkRepository:
metadata=md,
order=rec.order,
document_uri=doc_uri,
document_title=doc_title,
document_meta=json.loads(doc_meta),
)
)
@ -398,6 +400,7 @@ class ChunkRepository:
# Get document info from pre-fetched map
doc = documents_map.get(chunk_record.document_id)
doc_uri = doc.uri if doc else None
doc_title = doc.title if doc else None
doc_meta = doc.metadata if doc else "{}"
md = json.loads(chunk_record.metadata)
@ -409,6 +412,7 @@ class ChunkRepository:
metadata=md,
order=chunk_record.order,
document_uri=doc_uri,
document_title=doc_title,
document_meta=json.loads(doc_meta),
)

3
src/haiku/rag/store/repositories/document.py
View file

@ -34,6 +34,7 @@ class DocumentRepository:
id=record.id,
content=record.content,
uri=record.uri,
title=record.title,
metadata=json.loads(record.metadata),
created_at=datetime.fromisoformat(record.created_at)
if record.created_at
@ -56,6 +57,7 @@ class DocumentRepository:
id=doc_id,
content=entity.content,
uri=entity.uri,
title=entity.title,
metadata=json.dumps(entity.metadata),
created_at=now,
updated_at=now,
@ -97,6 +99,7 @@ class DocumentRepository:
values={
"content": entity.content,
"uri": entity.uri,
"title": entity.title,
"metadata": json.dumps(entity.metadata),
"updated_at": now,
},

2
src/haiku/rag/store/upgrades/__init__.py
View file

@ -55,6 +55,8 @@ def run_pending_upgrades(store: Store, from_version: str, to_version: str) -> No
from .v0_9_3 import upgrade_fts_phrase as upgrade_0_9_3_fts # noqa: E402
from .v0_9_3 import upgrade_order as upgrade_0_9_3_order # noqa: E402
from .v0_10_2 import upgrade_add_title as upgrade_0_10_2_add_title # noqa: E402
upgrades.append(upgrade_0_9_3_order)
upgrades.append(upgrade_0_9_3_fts)
upgrades.append(upgrade_0_10_2_add_title)

64
src/haiku/rag/store/upgrades/v0_10_2.py Normal file
View file

@ -0,0 +1,64 @@
import json
from lancedb.pydantic import LanceModel
from pydantic import Field
from haiku.rag.store.engine import Store
from haiku.rag.store.upgrades import Upgrade
def _apply_add_document_title(store: Store) -> None:
"""Add a nullable 'title' column to the documents table."""
# Read existing rows using Arrow for schema-agnostic access
try:
docs_arrow = store.documents_table.search().to_arrow()
rows = docs_arrow.to_pylist()
except Exception:
rows = []
class DocumentRecordV2(LanceModel):
id: str
content: str
uri: str | None = None
title: str | None = None
metadata: str = Field(default="{}")
created_at: str = Field(default_factory=lambda: "")
updated_at: str = Field(default_factory=lambda: "")
# Drop and recreate documents table with the new schema
try:
store.db.drop_table("documents")
except Exception:
pass
store.documents_table = store.db.create_table("documents", schema=DocumentRecordV2)
# Reinsert previous rows with title=None
if rows:
backfilled = []
for row in rows:
backfilled.append(
DocumentRecordV2(
id=row.get("id"),
content=row.get("content", ""),
uri=row.get("uri"),
title=None,
metadata=(
row.get("metadata")
if isinstance(row.get("metadata"), str)
else json.dumps(row.get("metadata") or {})
),
created_at=row.get("created_at", ""),
updated_at=row.get("updated_at", ""),
)
)
store.documents_table.add(backfilled)
upgrade_add_title = Upgrade(
version="0.10.2",
apply=_apply_add_document_title,
description="Add nullable 'title' column to documents table",
)

4
tests/test_app.py
View file

@ -78,7 +78,9 @@ async def test_add_document_from_source(app: HaikuRAGApp, monkeypatch):
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
await app.add_document_from_source(file_path)
mock_client.create_document_from_source.assert_called_once_with(file_path)
mock_client.create_document_from_source.assert_called_once_with(
file_path, title=None
)
mock_rich_print.assert_called_once_with(mock_doc, truncate=True)
mock_print.assert_called_once_with(
"[b]Document with id [cyan]1[/cyan] added successfully.[/b]"

16
tests/test_cli.py
View file

@ -45,6 +45,22 @@ def test_add_document_src():
mock_app_instance.add_document_from_source.assert_called_once()
def test_add_document_src_with_title():
with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock()
mock_app_instance.add_document_from_source = AsyncMock()
mock_app.return_value = mock_app_instance
result = runner.invoke(cli, ["add-src", "test.txt", "--title", "Nice Name"])
assert result.exit_code == 0
mock_app_instance.add_document_from_source.assert_called_once()
# Verify title is forwarded (inspect call kwargs)
_, kwargs = mock_app_instance.add_document_from_source.call_args
assert kwargs.get("title") == "Nice Name"
assert kwargs.get("source") == "test.txt"
def test_get_document():
with patch("haiku.rag.app.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock()

57
tests/test_client.py
View file

@ -105,6 +105,42 @@ async def test_client_create_document_from_source(temp_db_path):
assert "md5" in doc2.metadata
@pytest.mark.asyncio
async def test_client_create_document_from_source_with_title(temp_db_path):
"""Test creating a document from a file source with a title."""
async with HaikuRAG(temp_db_path) as client:
with tempfile.TemporaryDirectory() as temp_dir:
test_content = "This is test content from a file."
temp_path = Path(temp_dir) / "test_title.txt"
temp_path.write_text(test_content)
doc = await client.create_document_from_source(
source=temp_path, title="My Doc"
)
assert doc.id is not None
assert doc.title == "My Doc"
@pytest.mark.asyncio
async def test_client_update_title_noop_behavior(temp_db_path):
"""When content is unchanged, updating title should update document without re-chunking."""
async with HaikuRAG(temp_db_path) as client:
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir) / "test_update_title.txt"
temp_path.write_text("Original content")
doc1 = await client.create_document_from_source(temp_path, title="Title A")
assert doc1.id is not None
# Re-add with same content but new title
doc2 = await client.create_document_from_source(temp_path, title="Title B")
assert doc2.id == doc1.id
# Fetch and verify title updated
got = await client.get_document_by_id(doc1.id)
assert got is not None
assert got.title == "Title B"
@pytest.mark.asyncio
async def test_client_create_document_from_source_unsupported(temp_db_path):
"""Test creating a document from an unsupported file type."""
@ -536,18 +572,21 @@ async def test_client_expand_context(temp_db_path):
# Mock Config to have CONTEXT_CHUNK_RADIUS = 2
with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 2):
async with HaikuRAG(temp_db_path) as client:
# Create chunks manually
# Create chunks manually with precomputed embeddings to avoid network
dim = client.chunk_repository.embedder._vector_dim
z = [0.0] * dim
manual_chunks = [
Chunk(content="Chunk 0 content", order=0),
Chunk(content="Chunk 1 content", order=1),
Chunk(content="Chunk 2 content", order=2),
Chunk(content="Chunk 3 content", order=3),
Chunk(content="Chunk 4 content", order=4),
Chunk(content="Chunk 0 content", order=0, embedding=z),
Chunk(content="Chunk 1 content", order=1, embedding=z),
Chunk(content="Chunk 2 content", order=2, embedding=z),
Chunk(content="Chunk 3 content", order=3, embedding=z),
Chunk(content="Chunk 4 content", order=4, embedding=z),
]
doc = await client.create_document(
content="Full document content",
uri="test_doc.txt",
title="test_doc_title",
chunks=manual_chunks,
)
@ -560,16 +599,18 @@ async def test_client_expand_context(temp_db_path):
middle_chunk = next(c for c in chunks if c.order == 2)
search_results = [(middle_chunk, 0.8)]
# Test expand_context with radius=2
# Test expand_context with radius=2 and document title preserved
expanded_results = await client.expand_context(search_results, radius=2)
assert len(expanded_results) == 1
expanded_chunk, score = expanded_results[0]
# Check that the expanded chunk has combined content
# Check that the expanded chunk has combined content and preserves title/uri
assert expanded_chunk.id == middle_chunk.id
assert score == 0.8
assert "Chunk 2 content" in expanded_chunk.content
assert expanded_chunk.document_title == "test_doc_title"
assert expanded_chunk.document_uri == "test_doc.txt"
# Should include all chunks (radius=2 from chunk 2 = chunks 0,1,2,3,4)
assert "Chunk 0 content" in expanded_chunk.content

32
tests/test_search.py
View file

@ -106,6 +106,38 @@ async def test_chunks_include_document_info(temp_db_path):
store.close()
@pytest.mark.asyncio
async def test_chunks_include_document_title(temp_db_path):
"""Test that search results include the parent document title when present."""
store = Store(temp_db_path)
doc_repo = DocumentRepository(store)
chunk_repo = ChunkRepository(store)
# Create a document with URI and title
document = Document(
content="This is a test document with a custom title to verify enrichment.",
uri="file:///tmp/title-test.md",
title="My Custom Title",
)
# Create the document with chunks
from haiku.rag.utils import text_to_docling_document
dl = text_to_docling_document(document.content, name="title-test.md")
await doc_repo._create_with_docling(document, dl)
# Perform a search that should find this document
results = await chunk_repo.search("custom title", limit=3, search_type="hybrid")
assert results, "Expected at least one search result"
for chunk, _ in results:
# All returned chunks for this doc should carry the document title
if chunk.document_uri == "file:///tmp/title-test.md":
assert chunk.document_title == "My Custom Title"
store.close()
@pytest.mark.asyncio
async def test_search_score_types(temp_db_path):
"""Test that different search types return appropriate score ranges."""