haiku.rag 0.9.3__py3-none-any.whl → 0.10.1__py3-none-any.whl

haiku/rag/app.py

@@ -9,7 +9,13 @@ from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
 from haiku.rag.mcp import create_mcp_server
 from haiku.rag.monitor import FileWatcher
-from haiku.rag.research.orchestrator import ResearchOrchestrator
+from haiku.rag.research.dependencies import ResearchContext
+from haiku.rag.research.graph import (
+    PlanNode,
+    ResearchDeps,
+    ResearchState,
+    build_research_graph,
+)
 from haiku.rag.store.models.chunk import Chunk
 from haiku.rag.store.models.document import Document
 
@@ -33,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]"
@@ -80,28 +86,53 @@ class HaikuRAGApp:
                 self.console.print(f"[red]Error: {e}[/red]")
 
     async def research(
-        self, question: str, max_iterations: int = 3, verbose: bool = False
+        self,
+        question: str,
+        max_iterations: int = 3,
+        confidence_threshold: float = 0.8,
+        max_concurrency: int = 1,
+        verbose: bool = False,
     ):
-        """Run multi-agent research on a question."""
+        """Run research via the pydantic-graph pipeline (default)."""
         async with HaikuRAG(db_path=self.db_path) as client:
             try:
-                # Create orchestrator with default config or fallback to QA
-                orchestrator = ResearchOrchestrator()
-
                 if verbose:
-                    self.console.print(
-                        f"[bold cyan]Starting research with {orchestrator.provider}:{orchestrator.model}[/bold cyan]"
-                    )
+                    self.console.print("[bold cyan]Starting research[/bold cyan]")
                     self.console.print(f"[bold blue]Question:[/bold blue] {question}")
                     self.console.print()
 
-                # Conduct research
-                report = await orchestrator.conduct_research(
+                graph = build_research_graph()
+                state = ResearchState(
                     question=question,
-                    client=client,
+                    context=ResearchContext(original_question=question),
                     max_iterations=max_iterations,
-                    verbose=verbose,
+                    confidence_threshold=confidence_threshold,
+                    max_concurrency=max_concurrency,
+                )
+                deps = ResearchDeps(
+                    client=client, console=self.console if verbose else None
+                )
+
+                start = PlanNode(
+                    provider=Config.RESEARCH_PROVIDER or Config.QA_PROVIDER,
+                    model=Config.RESEARCH_MODEL or Config.QA_MODEL,
                 )
+                # Prefer graph.run; fall back to iter if unavailable
+                report = None
+                try:
+                    result = await graph.run(start, state=state, deps=deps)
+                    report = result.output
+                except Exception:
+                    from pydantic_graph import End
+
+                    async with graph.iter(start, state=state, deps=deps) as run:
+                        node = run.next_node
+                        while not isinstance(node, End):
+                            node = await run.next(node)
+                        if run.result:
+                            report = run.result.output
+                if report is None:
+                    raise RuntimeError("Graph did not produce a report")
 
                 # Display the report
                 self.console.print("[bold green]Research Report[/bold green]")
@@ -114,6 +145,12 @@ class HaikuRAGApp:
                 self.console.print(report.executive_summary)
                 self.console.print()
 
+                # Confidence (from last evaluation)
+                if state.last_eval:
+                    conf = state.last_eval.confidence_score  # type: ignore[attr-defined]
+                    self.console.print(f"[bold cyan]Confidence:[/bold cyan] {conf:.1%}")
+                    self.console.print()
+
                 # Main Findings
                 if report.main_findings:
                     self.console.print("[bold cyan]Main Findings:[/bold cyan]")
@@ -215,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}"
@@ -235,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)
@@ -252,8 +300,6 @@ class HaikuRAGApp:
             try:
                 if transport == "stdio":
                     await server.run_stdio_async()
-                elif transport == "sse":
-                    await server.run_sse_async()
                 else:
                     await server.run_http_async(transport="streamable-http")
             except KeyboardInterrupt:

haiku/rag/cli.py

@@ -3,28 +3,16 @@ import warnings
 from importlib.metadata import version
 from pathlib import Path
 
-import logfire
 import typer
-from rich.console import Console
 
-from haiku.rag.app import HaikuRAGApp
 from haiku.rag.config import Config
 from haiku.rag.logging import configure_cli_logging
-from haiku.rag.migration import migrate_sqlite_to_lancedb
 from haiku.rag.utils import is_up_to_date
 
-logfire.configure(send_to_logfire="if-token-present")
-logfire.instrument_pydantic_ai()
-
-if not Config.ENV == "development":
-    warnings.filterwarnings("ignore")
-
 cli = typer.Typer(
     context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True
 )
 
-console = Console()
-
 
 def complete_document_ids(ctx: typer.Context, incomplete: str):
     """Autocomplete document IDs from the selected DB."""
@@ -89,16 +77,16 @@ async def check_version():
     """Check if haiku.rag is up to date and show warning if not."""
     up_to_date, current_version, latest_version = await is_up_to_date()
     if not up_to_date:
-        console.print(
-            f"[yellow]Warning: haiku.rag is outdated. Current: {current_version}, Latest: {latest_version}[/yellow]"
+        typer.echo(
+            f"Warning: haiku.rag is outdated. Current: {current_version}, Latest: {latest_version}",
         )
-        console.print("[yellow]Please update.[/yellow]")
+        typer.echo("Please update.")
 
 
 def version_callback(value: bool):
     if value:
         v = version("haiku.rag")
-        console.print(f"haiku.rag version {v}")
+        typer.echo(f"haiku.rag version {v}")
         raise typer.Exit()
 
 
@@ -113,10 +101,26 @@ def main(
     ),
 ):
     """haiku.rag CLI - Vector database RAG system"""
-    # Ensure only haiku.rag logs are emitted in CLI context
-    configure_cli_logging()
+    # Configure logging minimally for CLI context
+    if Config.ENV == "development":
+        # Lazy import logfire only in development
+        try:
+            import logfire  # type: ignore
+
+            logfire.configure(send_to_logfire="if-token-present")
+            logfire.instrument_pydantic_ai()
+        except Exception:
+            pass
+    else:
+        configure_cli_logging()
+        warnings.filterwarnings("ignore")
+
     # Run version check before any command
-    asyncio.run(check_version())
+    try:
+        asyncio.run(check_version())
+    except Exception:
+        # Do not block CLI on version check issues
+        pass
 
 
 @cli.command("list", help="List all stored documents")
@@ -127,6 +131,8 @@ def list_documents(
         help="Path to the LanceDB database file",
     ),
 ):
+    from haiku.rag.app import HaikuRAGApp
+
     app = HaikuRAGApp(db_path=db)
     asyncio.run(app.list_documents())
 
@@ -142,6 +148,8 @@ def add_document_text(
         help="Path to the LanceDB database file",
     ),
 ):
+    from haiku.rag.app import HaikuRAGApp
+
     app = HaikuRAGApp(db_path=db)
     asyncio.run(app.add_document_from_text(text=text))
 
@@ -152,14 +160,21 @@ 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",
         help="Path to the LanceDB database file",
     ),
 ):
+    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")
@@ -174,6 +189,8 @@ def get_document(
         help="Path to the LanceDB database file",
     ),
 ):
+    from haiku.rag.app import HaikuRAGApp
+
     app = HaikuRAGApp(db_path=db)
     asyncio.run(app.get_document(doc_id=doc_id))
 
@@ -190,6 +207,8 @@ def delete_document(
         help="Path to the LanceDB database file",
     ),
 ):
+    from haiku.rag.app import HaikuRAGApp
+
     app = HaikuRAGApp(db_path=db)
     asyncio.run(app.delete_document(doc_id=doc_id))
 
@@ -215,6 +234,8 @@ def search(
         help="Path to the LanceDB database file",
     ),
 ):
+    from haiku.rag.app import HaikuRAGApp
+
     app = HaikuRAGApp(db_path=db)
     asyncio.run(app.search(query=query, limit=limit))
 
@@ -235,6 +256,8 @@ def ask(
         help="Include citations in the response",
     ),
 ):
+    from haiku.rag.app import HaikuRAGApp
+
     app = HaikuRAGApp(db_path=db)
     asyncio.run(app.ask(question=question, cite=cite))
 
@@ -250,6 +273,16 @@ def research(
         "-n",
         help="Maximum search/analyze iterations",
     ),
+    confidence_threshold: float = typer.Option(
+        0.8,
+        "--confidence-threshold",
+        help="Minimum confidence (0-1) to stop",
+    ),
+    max_concurrency: int = typer.Option(
+        1,
+        "--max-concurrency",
+        help="Max concurrent searches per iteration (planned)",
+    ),
     db: Path = typer.Option(
         Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         "--db",
@@ -261,11 +294,15 @@ def research(
         help="Show verbose progress output",
     ),
 ):
+    from haiku.rag.app import HaikuRAGApp
+
     app = HaikuRAGApp(db_path=db)
     asyncio.run(
         app.research(
             question=question,
             max_iterations=max_iterations,
+            confidence_threshold=confidence_threshold,
+            max_concurrency=max_concurrency,
             verbose=verbose,
         )
     )
@@ -273,6 +310,8 @@ def research(
 
 @cli.command("settings", help="Display current configuration settings")
 def settings():
+    from haiku.rag.app import HaikuRAGApp
+
     app = HaikuRAGApp(db_path=Path())  # Don't need actual DB for settings
     app.show_settings()
 
@@ -288,6 +327,8 @@ def rebuild(
         help="Path to the LanceDB database file",
     ),
 ):
+    from haiku.rag.app import HaikuRAGApp
+
     app = HaikuRAGApp(db_path=db)
     asyncio.run(app.rebuild())
 
@@ -300,6 +341,8 @@ def vacuum(
         help="Path to the LanceDB database file",
     ),
 ):
+    from haiku.rag.app import HaikuRAGApp
+
     app = HaikuRAGApp(db_path=db)
     asyncio.run(app.vacuum())
 
@@ -318,24 +361,15 @@ def serve(
         "--stdio",
         help="Run MCP server on stdio Transport",
     ),
-    sse: bool = typer.Option(
-        False,
-        "--sse",
-        help="Run MCP server on SSE transport",
-    ),
 ) -> None:
     """Start the MCP server."""
-    if stdio and sse:
-        console.print("[red]Error: Cannot use both --stdio and --http options[/red]")
-        raise typer.Exit(1)
+    from haiku.rag.app import HaikuRAGApp
 
     app = HaikuRAGApp(db_path=db)
 
     transport = None
     if stdio:
         transport = "stdio"
-    elif sse:
-        transport = "sse"
 
     asyncio.run(app.serve(transport=transport))
 
@@ -349,6 +383,9 @@ def migrate(
     # Generate LanceDB path in same parent directory
     lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
 
+    # Lazy import to avoid heavy deps on simple invocations
+    from haiku.rag.migration import migrate_sqlite_to_lancedb
+
     success = asyncio.run(migrate_sqlite_to_lancedb(sqlite_path, lancedb_path))
 
     if not success:

haiku/rag/client.py

@@ -33,8 +33,6 @@ class HaikuRAG:
             db_path: Path to the database file.
             skip_validation: Whether to skip configuration validation on database load.
         """
-        if not db_path.parent.exists():
-            Path.mkdir(db_path.parent, parents=True)
         self.store = Store(db_path, skip_validation=skip_validation)
         self.document_repository = DocumentRepository(self.store)
         self.chunk_repository = ChunkRepository(self.store)
@@ -52,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:
@@ -60,6 +59,7 @@ class HaikuRAG:
         document = Document(
             content=content,
             uri=uri,
+            title=title,
             metadata=metadata or {},
         )
         return await self.document_repository._create_with_docling(
@@ -70,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:
@@ -90,6 +91,7 @@ class HaikuRAG:
         document = Document(
             content=content,
             uri=uri,
+            title=title,
             metadata=metadata or {},
         )
         return await self.document_repository._create_with_docling(
@@ -97,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.
 
@@ -118,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)
@@ -138,37 +145,51 @@ class HaikuRAG:
         uri = source_path.absolute().as_uri()
         md5_hash = hashlib.md5(source_path.read_bytes()).hexdigest()
 
+        # 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, 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
+            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)
 
-        # Get content type from file extension
-        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})
-
         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.
 
@@ -188,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
             )
@@ -228,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(
@@ -418,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,
                 )
 
@@ -524,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"

haiku/rag/config.py

@@ -53,6 +53,10 @@ class AppConfig(BaseModel):
     ANTHROPIC_API_KEY: str = ""
     COHERE_API_KEY: str = ""
 
+    # If true, refuse to auto-create a new LanceDB database or tables
+    # and error out when the database does not already exist.
+    DISABLE_DB_AUTOCREATE: bool = False
+
     @field_validator("MONITOR_DIRECTORIES", mode="before")
     @classmethod
     def parse_monitor_directories(cls, v):

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