diff --git a/docs/agents.md b/docs/agents.md index 4d2d82c7..59e314cc 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -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: diff --git a/docs/cli.md b/docs/cli.md index 29304493..79ae0fbb 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -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 diff --git a/docs/index.md b/docs/index.md index a788105c..514fe5bf 100644 --- a/docs/index.md +++ b/docs/index.md @@ -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 diff --git a/docs/python.md b/docs/python.md index 94209792..78649480 100644 --- a/docs/python.md +++ b/docs/python.md @@ -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)). diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index c41d6267..18a0d1f6 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -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) diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 4fc5b144..66aa733c 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -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") diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 9a6b89ec..9cbb6e8d 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -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" diff --git a/src/haiku/rag/mcp.py b/src/haiku/rag/mcp.py index 16a4d4a4..654cf8e8 100644 --- a/src/haiku/rag/mcp.py +++ b/src/haiku/rag/mcp.py @@ -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), diff --git a/src/haiku/rag/qa/agent.py b/src/haiku/rag/qa/agent.py index 3493b0ce..b3ad2248 100644 --- a/src/haiku/rag/qa/agent.py +++ b/src/haiku/rag/qa/agent.py @@ -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 ] diff --git a/src/haiku/rag/qa/prompts.py b/src/haiku/rag/qa/prompts.py index 19896d2c..b25b36af 100644 --- a/src/haiku/rag/qa/prompts.py +++ b/src/haiku/rag/qa/prompts.py @@ -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] diff --git a/src/haiku/rag/research/models.py b/src/haiku/rag/research/models.py index 3f789842..30cb97d1 100644 --- a/src/haiku/rag/research/models.py +++ b/src/haiku/rag/research/models.py @@ -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, ) diff --git a/src/haiku/rag/research/nodes/search.py b/src/haiku/rag/research/nodes/search.py index 0943338a..ae863d37 100644 --- a/src/haiku/rag/research/nodes/search.py +++ b/src/haiku/rag/research/nodes/search.py @@ -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 ] diff --git a/src/haiku/rag/research/prompts.py b/src/haiku/rag/research/prompts.py index 6d9aad48..29e10424 100644 --- a/src/haiku/rag/research/prompts.py +++ b/src/haiku/rag/research/prompts.py @@ -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. diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index 4d736834..517b139f 100644 --- a/src/haiku/rag/store/engine.py +++ b/src/haiku/rag/store/engine.py @@ -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: "") diff --git a/src/haiku/rag/store/models/chunk.py b/src/haiku/rag/store/models/chunk.py index 383f178e..42ea1368 100644 --- a/src/haiku/rag/store/models/chunk.py +++ b/src/haiku/rag/store/models/chunk.py @@ -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 diff --git a/src/haiku/rag/store/models/document.py b/src/haiku/rag/store/models/document.py index 39a46271..e23d5d3c 100644 --- a/src/haiku/rag/store/models/document.py +++ b/src/haiku/rag/store/models/document.py @@ -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) diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 567a5a34..fe3bd540 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -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), ) diff --git a/src/haiku/rag/store/repositories/document.py b/src/haiku/rag/store/repositories/document.py index bf8160e8..ffc66a1b 100644 --- a/src/haiku/rag/store/repositories/document.py +++ b/src/haiku/rag/store/repositories/document.py @@ -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, }, diff --git a/src/haiku/rag/store/upgrades/__init__.py b/src/haiku/rag/store/upgrades/__init__.py index 5888cdba..cf9b8e19 100644 --- a/src/haiku/rag/store/upgrades/__init__.py +++ b/src/haiku/rag/store/upgrades/__init__.py @@ -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) diff --git a/src/haiku/rag/store/upgrades/v0_10_2.py b/src/haiku/rag/store/upgrades/v0_10_2.py new file mode 100644 index 00000000..bb9b5113 --- /dev/null +++ b/src/haiku/rag/store/upgrades/v0_10_2.py @@ -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", +) diff --git a/tests/test_app.py b/tests/test_app.py index 5e7f80d1..ce0db165 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -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]" diff --git a/tests/test_cli.py b/tests/test_cli.py index d4e5a455..7c2b4a80 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -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() diff --git a/tests/test_client.py b/tests/test_client.py index 60038067..44ec66c9 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -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 diff --git a/tests/test_search.py b/tests/test_search.py index f97086e9..683649c0 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -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."""