haiku.rag 0.11.3__py3-none-any.whl → 0.12.0__py3-none-any.whl

haiku/rag/cli.py

@@ -16,65 +16,6 @@ cli = typer.Typer(
 )
 
 
-def complete_document_ids(ctx: typer.Context, incomplete: str):
-    """Autocomplete document IDs from the selected DB."""
-    db_path = ctx.params.get("db") or (Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb")
-
-    try:
-        from haiku.rag.client import HaikuRAG
-
-        async def _list_ids():
-            async with HaikuRAG(db_path) as client:
-                docs = await client.list_documents()
-                return [d.id for d in docs if d.id]
-
-        ids = asyncio.run(_list_ids())
-    except Exception:
-        return []
-
-    return [i for i in ids if i and i.startswith(incomplete)]
-
-
-def complete_local_paths(ctx: typer.Context, incomplete: str) -> list[str]:
-    """Autocomplete local filesystem paths.
-
-    Provides directory/file suggestions based on the current incomplete input.
-    Does not validate or restrict to specific extensions to keep it flexible
-    (URLs are still allowed to be typed manually).
-    """
-    try:
-        text = incomplete or ""
-
-        # Expand user home
-        from os.path import expanduser
-
-        expanded = expanduser(text)
-        p = Path(expanded)
-
-        # Choose directory to list and prefix to filter
-        if text == "" or text.endswith(("/", "\\")):
-            directory = p
-            prefix = ""
-        else:
-            directory = p.parent
-            prefix = p.name
-
-        if not directory.exists():
-            return []
-
-        suggestions: list[str] = []
-        for entry in directory.iterdir():
-            name = entry.name
-            if not prefix or name.startswith(prefix):
-                suggestion = str(directory / name)
-                if entry.is_dir():
-                    suggestion += "/"
-                suggestions.append(suggestion)
-        return suggestions
-    except Exception:
-        return []
-
-
 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()
@@ -191,7 +132,6 @@ def add_document_text(
 def add_document_src(
     source: str = typer.Argument(
         help="The file path or URL of the document to add",
-        autocompletion=complete_local_paths,
     ),
     title: str | None = typer.Option(
         None,
@@ -225,7 +165,6 @@ def add_document_src(
 def get_document(
     doc_id: str = typer.Argument(
         help="The ID of the document to get",
-        autocompletion=complete_document_ids,
     ),
     db: Path = typer.Option(
         Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
@@ -243,7 +182,6 @@ def get_document(
 def delete_document(
     doc_id: str = typer.Argument(
         help="The ID of the document to delete",
-        autocompletion=complete_document_ids,
     ),
     db: Path = typer.Option(
         Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
@@ -428,7 +366,8 @@ def download_models_cmd():
 
 
 @cli.command(
-    "serve", help="Start the haiku.rag MCP server (by default in streamable HTTP mode)"
+    "serve",
+    help="Start haiku.rag server. Use --monitor, --mcp, and/or --a2a to enable services.",
 )
 def serve(
     db: Path = typer.Option(
@@ -436,22 +375,71 @@ def serve(
         "--db",
         help="Path to the LanceDB database file",
     ),
+    monitor: bool = typer.Option(
+        False,
+        "--monitor",
+        help="Enable file monitoring",
+    ),
+    mcp: bool = typer.Option(
+        False,
+        "--mcp",
+        help="Enable MCP server",
+    ),
     stdio: bool = typer.Option(
         False,
         "--stdio",
-        help="Run MCP server on stdio Transport",
+        help="Run MCP server on stdio Transport (requires --mcp)",
+    ),
+    mcp_port: int = typer.Option(
+        8001,
+        "--mcp-port",
+        help="Port to bind MCP server to (ignored with --stdio)",
+    ),
+    a2a: bool = typer.Option(
+        False,
+        "--a2a",
+        help="Enable A2A (Agent-to-Agent) server",
+    ),
+    a2a_host: str = typer.Option(
+        "127.0.0.1",
+        "--a2a-host",
+        help="Host to bind A2A server to",
+    ),
+    a2a_port: int = typer.Option(
+        8000,
+        "--a2a-port",
+        help="Port to bind A2A server to",
     ),
 ) -> None:
-    """Start the MCP server."""
+    """Start the server with selected services."""
+    # Require at least one service flag
+    if not (monitor or mcp or a2a):
+        typer.echo(
+            "Error: At least one service flag (--monitor, --mcp, or --a2a) must be specified"
+        )
+        raise typer.Exit(1)
+
+    if stdio and not mcp:
+        typer.echo("Error: --stdio requires --mcp")
+        raise typer.Exit(1)
+
     from haiku.rag.app import HaikuRAGApp
 
     app = HaikuRAGApp(db_path=db)
 
-    transport = None
-    if stdio:
-        transport = "stdio"
+    transport = "stdio" if stdio else None
 
-    asyncio.run(app.serve(transport=transport))
+    asyncio.run(
+        app.serve(
+            enable_monitor=monitor,
+            enable_mcp=mcp,
+            mcp_transport=transport,
+            mcp_port=mcp_port,
+            enable_a2a=a2a,
+            a2a_host=a2a_host,
+            a2a_port=a2a_port,
+        )
+    )
 
 
 @cli.command("migrate", help="Migrate an SQLite database to LanceDB")
@@ -472,5 +460,27 @@ def migrate(
         raise typer.Exit(1)
 
 
+@cli.command(
+    "a2aclient", help="Run interactive client to chat with haiku.rag's A2A server"
+)
+def a2aclient(
+    url: str = typer.Option(
+        "http://localhost:8000",
+        "--url",
+        help="Base URL of the A2A server",
+    ),
+):
+    try:
+        from haiku.rag.a2a.client import run_interactive_client
+    except ImportError:
+        typer.echo(
+            "Error: A2A support requires the 'a2a' extra. "
+            "Install with: uv pip install 'haiku.rag[a2a]'"
+        )
+        raise typer.Exit(1)
+
+    asyncio.run(run_interactive_client(url=url))
+
+
 if __name__ == "__main__":
     cli()

haiku/rag/client.py

@@ -46,6 +46,9 @@ class HaikuRAG:
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):  # noqa: ARG002
         """Async context manager exit."""
+        # Wait for any pending vacuum to complete before closing
+        async with self.store._vacuum_lock:
+            pass
         self.close()
         return False
 
@@ -522,19 +525,22 @@ class HaikuRAG:
         merged.append(current)
         return merged
 
-    async def ask(self, question: str, cite: bool = False) -> str:
+    async def ask(
+        self, question: str, cite: bool = False, system_prompt: str | None = None
+    ) -> str:
         """Ask a question using the configured QA agent.
 
         Args:
             question: The question to ask.
             cite: Whether to include citations in the response.
+            system_prompt: Optional custom system prompt for the QA agent.
 
         Returns:
             The generated answer as a string.
         """
         from haiku.rag.qa import get_qa_agent
 
-        qa_agent = get_qa_agent(self, use_citations=cite)
+        qa_agent = get_qa_agent(self, use_citations=cite, system_prompt=system_prompt)
         return await qa_agent.answer(question)
 
     async def rebuild_database(self) -> AsyncGenerator[str, None]:
@@ -617,13 +623,13 @@ class HaikuRAG:
 
         # Final maintenance: centralized vacuum to curb disk usage
         try:
-            self.store.vacuum()
+            await self.store.vacuum()
         except Exception:
             pass
 
     async def vacuum(self) -> None:
         """Optimize and clean up old versions across all tables."""
-        self.store.vacuum()
+        await self.store.vacuum()
 
     def close(self):
         """Close the underlying store connection."""

haiku/rag/config.py

@@ -57,6 +57,15 @@ class AppConfig(BaseModel):
     # and error out when the database does not already exist.
     DISABLE_DB_AUTOCREATE: bool = False
 
+    # Vacuum retention threshold in seconds. Only versions older than this
+    # threshold will be removed during vacuum operations. Default is 60 seconds
+    # to allow concurrent connections to safely use recent versions.
+    VACUUM_RETENTION_SECONDS: int = 60
+
+    # Maximum number of A2A contexts to keep in memory. When exceeded, least
+    # recently used contexts will be evicted. Default is 1000.
+    A2A_MAX_CONTEXTS: int = 1000
+
     @field_validator("MONITOR_DIRECTORIES", mode="before")
     @classmethod
     def parse_monitor_directories(cls, v):

haiku/rag/mcp.py

@@ -5,6 +5,8 @@ from fastmcp import FastMCP
 from pydantic import BaseModel
 
 from haiku.rag.client import HaikuRAG
+from haiku.rag.config import Config
+from haiku.rag.research.models import ResearchReport
 
 
 class SearchResult(BaseModel):
@@ -153,4 +155,101 @@ def create_mcp_server(db_path: Path) -> FastMCP:
         except Exception:
             return False
 
+    @mcp.tool()
+    async def ask_question(
+        question: str,
+        cite: bool = False,
+        deep: bool = False,
+    ) -> str:
+        """Ask a question using the QA agent.
+
+        Args:
+            question: The question to ask.
+            cite: Whether to include citations in the response.
+            deep: Use deep multi-agent QA for complex questions that require decomposition.
+
+        Returns:
+            The answer as a string.
+        """
+        try:
+            async with HaikuRAG(db_path) as rag:
+                if deep:
+                    from haiku.rag.config import Config
+                    from haiku.rag.qa.deep.dependencies import DeepQAContext
+                    from haiku.rag.qa.deep.graph import build_deep_qa_graph
+                    from haiku.rag.qa.deep.nodes import DeepQAPlanNode
+                    from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
+
+                    graph = build_deep_qa_graph()
+                    context = DeepQAContext(
+                        original_question=question, use_citations=cite
+                    )
+                    state = DeepQAState(context=context)
+                    deps = DeepQADeps(client=rag)
+
+                    start_node = DeepQAPlanNode(
+                        provider=Config.QA_PROVIDER,
+                        model=Config.QA_MODEL,
+                    )
+
+                    result = await graph.run(
+                        start_node=start_node, state=state, deps=deps
+                    )
+                    answer = result.output.answer
+                else:
+                    answer = await rag.ask(question, cite=cite)
+                return answer
+        except Exception as e:
+            return f"Error answering question: {e!s}"
+
+    @mcp.tool()
+    async def research_question(
+        question: str,
+        max_iterations: int = 3,
+        confidence_threshold: float = 0.8,
+        max_concurrency: int = 1,
+    ) -> ResearchReport | None:
+        """Run multi-agent research to investigate a complex question.
+
+        The research process uses multiple agents to plan, search, evaluate, and synthesize
+        information iteratively until confidence threshold is met or max iterations reached.
+
+        Args:
+            question: The research question to investigate.
+            max_iterations: Maximum search/analyze iterations (default: 3).
+            confidence_threshold: Minimum confidence score (0-1) to stop early (default: 0.8).
+            max_concurrency: Maximum concurrent searches per iteration (default: 1).
+
+        Returns:
+            A research report with findings, or None if an error occurred.
+        """
+        try:
+            from haiku.rag.graph.nodes.plan import PlanNode
+            from haiku.rag.research.dependencies import ResearchContext
+            from haiku.rag.research.graph import build_research_graph
+            from haiku.rag.research.state import ResearchDeps, ResearchState
+
+            async with HaikuRAG(db_path) as rag:
+                graph = build_research_graph()
+                state = ResearchState(
+                    context=ResearchContext(original_question=question),
+                    max_iterations=max_iterations,
+                    confidence_threshold=confidence_threshold,
+                    max_concurrency=max_concurrency,
+                )
+                deps = ResearchDeps(client=rag)
+
+                result = await graph.run(
+                    PlanNode(
+                        provider=Config.RESEARCH_PROVIDER or Config.QA_PROVIDER,
+                        model=Config.RESEARCH_MODEL or Config.QA_MODEL,
+                    ),
+                    state=state,
+                    deps=deps,
+                )
+
+                return result.output
+        except Exception:
+            return None
+
     return mcp

haiku/rag/migration.py

@@ -27,7 +27,7 @@ class SQLiteToLanceDBMigrator:
         self.lancedb_path = lancedb_path
         self.console = Console()
 
-    def migrate(self) -> bool:
+    async def migrate(self) -> bool:
         """Perform the migration."""
         try:
             self.console.print(
@@ -94,7 +94,7 @@ class SQLiteToLanceDBMigrator:
             # Optimize and cleanup using centralized vacuum
             self.console.print("[cyan]Optimizing LanceDB...[/cyan]")
             try:
-                lance_store.vacuum()
+                await lance_store.vacuum()
                 self.console.print("[green]✅ Optimization completed[/green]")
             except Exception as e:
                 self.console.print(
@@ -313,4 +313,4 @@ async def migrate_sqlite_to_lancedb(
         lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
 
     migrator = SQLiteToLanceDBMigrator(sqlite_path, lancedb_path)
-    return migrator.migrate()
+    return await migrator.migrate()

haiku/rag/qa/__init__.py

@@ -3,7 +3,11 @@ from haiku.rag.config import Config
 from haiku.rag.qa.agent import QuestionAnswerAgent
 
 
-def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswerAgent:
+def get_qa_agent(
+    client: HaikuRAG,
+    use_citations: bool = False,
+    system_prompt: str | None = None,
+) -> QuestionAnswerAgent:
     provider = Config.QA_PROVIDER
     model_name = Config.QA_MODEL
 
@@ -12,4 +16,5 @@ def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswe
         provider=provider,
         model=model_name,
         use_citations=use_citations,
+        system_prompt=system_prompt,
     )

haiku/rag/qa/agent.py

@@ -30,18 +30,21 @@ class QuestionAnswerAgent:
         model: str,
         use_citations: bool = False,
         q: float = 0.0,
+        system_prompt: str | None = None,
     ):
         self._client = client
 
-        system_prompt = (
-            QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
-        )
+        if system_prompt is None:
+            system_prompt = (
+                QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
+            )
         model_obj = self._get_model(provider, model)
 
         self._agent = Agent(
             model=model_obj,
             deps_type=Dependencies,
             system_prompt=system_prompt,
+            retries=3,
         )
 
         @self._agent.tool
@@ -51,9 +54,6 @@ class QuestionAnswerAgent:
             limit: int = 3,
         ) -> list[SearchResult]:
             """Search the knowledge base for relevant documents."""
-
-            # Remove quotes from queries as this requires positional indexing in lancedb
-            query = query.replace('"', "")
             search_results = await ctx.deps.client.search(query, limit=limit)
             expanded_results = await ctx.deps.client.expand_context(search_results)
 

haiku/rag/store/engine.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import logging
 from datetime import timedelta
@@ -51,6 +52,7 @@ class Store:
     def __init__(self, db_path: Path, skip_validation: bool = False):
         self.db_path: Path = db_path
         self.embedder = get_embedder()
+        self._vacuum_lock = asyncio.Lock()
 
         # Create the ChunkRecord model with the correct vector dimension
         self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
@@ -78,14 +80,40 @@ class Store:
         if not skip_validation:
             self._validate_configuration()
 
-    def vacuum(self) -> None:
-        """Optimize and clean up old versions across all tables to reduce disk usage."""
+    async def vacuum(self, retention_seconds: int | None = None) -> None:
+        """Optimize and clean up old versions across all tables to reduce disk usage.
+
+        Args:
+            retention_seconds: Retention threshold in seconds. Only versions older
+                              than this will be removed. If None, uses Config.VACUUM_RETENTION_SECONDS.
+
+        Note:
+            If vacuum is already running, this method returns immediately without blocking.
+            Use asyncio.create_task(store.vacuum()) for non-blocking background execution.
+        """
         if self._has_cloud_config() and str(Config.LANCEDB_URI).startswith("db://"):
             return
 
-        # Perform maintenance per table using optimize() with cleanup_older_than 0
-        for table in [self.documents_table, self.chunks_table, self.settings_table]:
-            table.optimize(cleanup_older_than=timedelta(0))
+        # Skip if already running (non-blocking)
+        if self._vacuum_lock.locked():
+            return
+
+        async with self._vacuum_lock:
+            try:
+                # Evaluate config at runtime to allow dynamic changes
+                if retention_seconds is None:
+                    retention_seconds = Config.VACUUM_RETENTION_SECONDS
+                # Perform maintenance per table using optimize() with configurable retention
+                retention = timedelta(seconds=retention_seconds)
+                for table in [
+                    self.documents_table,
+                    self.chunks_table,
+                    self.settings_table,
+                ]:
+                    table.optimize(cleanup_older_than=retention)
+            except (RuntimeError, OSError) as e:
+                # Handle resource errors gracefully
+                logger.debug(f"Vacuum skipped due to resource constraints: {e}")
 
     def _connect_to_lancedb(self, db_path: Path):
         """Establish connection to LanceDB (local, cloud, or object storage)."""

haiku/rag/store/repositories/chunk.py

@@ -1,4 +1,3 @@
-import asyncio
 import inspect
 import json
 import logging
@@ -23,7 +22,6 @@ class ChunkRepository:
     def __init__(self, store: Store) -> None:
         self.store = store
         self.embedder = get_embedder()
-        self._optimize_lock = asyncio.Lock()
 
     def _ensure_fts_index(self) -> None:
         """Ensure FTS index exists on the content column."""
@@ -35,21 +33,6 @@ class ChunkRepository:
             # Log the error but don't fail - FTS might already exist
             logger.debug(f"FTS index creation skipped: {e}")
 
-    async def _optimize(self) -> None:
-        """Optimize the chunks table to refresh indexes."""
-        # Skip optimization for LanceDB Cloud as it handles this automatically
-        if Config.LANCEDB_URI and Config.LANCEDB_URI.startswith("db://"):
-            return
-
-        async with self._optimize_lock:
-            try:
-                self.store.chunks_table.optimize()
-            except (RuntimeError, OSError) as e:
-                # Handle "too many open files" and other resource errors gracefully
-                logger.debug(
-                    f"Table optimization skipped due to resource constraints: {e}"
-                )
-
     async def create(self, entity: Chunk) -> Chunk:
         """Create a chunk in the database."""
         assert entity.document_id, "Chunk must have a document_id to be created"
@@ -77,11 +60,6 @@ class ChunkRepository:
         self.store.chunks_table.add([chunk_record])
 
         entity.id = chunk_id
-
-        # Try to optimize if not currently locked (non-blocking)
-        if not self._optimize_lock.locked():
-            asyncio.create_task(self._optimize())
-
         return entity
 
     async def get_by_id(self, entity_id: str) -> Chunk | None:
@@ -125,10 +103,6 @@ class ChunkRepository:
                 "vector": embedding,
             },
         )
-        # Try to optimize if not currently locked (non-blocking)
-        if not self._optimize_lock.locked():
-            asyncio.create_task(self._optimize())
-
         return entity
 
     async def delete(self, entity_id: str) -> bool:
@@ -227,8 +201,6 @@ class ChunkRepository:
         if chunk_records:
             self.store.chunks_table.add(chunk_records)
 
-        # Force optimization once at the end for bulk operations
-        await self._optimize()
         return created_chunks
 
     async def delete_all(self) -> None:

haiku/rag/store/repositories/document.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 from datetime import datetime
 from typing import TYPE_CHECKING
@@ -200,6 +201,9 @@ class DocumentRepository:
                     chunk.order = order
                     await self.chunk_repository.create(chunk)
 
+            # Vacuum old versions in background (non-blocking)
+            asyncio.create_task(self.store.vacuum())
+
             return created_doc
         except Exception:
             # Roll back to the captured versions and re-raise
@@ -230,6 +234,9 @@ class DocumentRepository:
                 updated_doc.id, docling_document
             )
 
+            # Vacuum old versions in background (non-blocking)
+            asyncio.create_task(self.store.vacuum())
+
             return updated_doc
         except Exception:
             # Roll back to the captured versions and re-raise

{haiku_rag-0.11.3.dist-info → haiku_rag-0.12.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.11.3
+Version: 0.12.0
 Summary: Agentic Retrieval Augmented Generation (RAG) with LanceDB
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT
@@ -18,18 +18,20 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Typing :: Typed
 Requires-Python: >=3.12
-Requires-Dist: docling>=2.52.0
-Requires-Dist: fastmcp>=2.12.3
+Requires-Dist: docling>=2.56.1
+Requires-Dist: fastmcp>=2.12.4
 Requires-Dist: httpx>=0.28.1
-Requires-Dist: lancedb>=0.25.0
-Requires-Dist: pydantic-ai>=1.0.8
-Requires-Dist: pydantic-graph>=1.0.8
-Requires-Dist: pydantic>=2.11.9
+Requires-Dist: lancedb>=0.25.2
+Requires-Dist: pydantic-ai>=1.0.18
+Requires-Dist: pydantic-graph>=1.0.18
+Requires-Dist: pydantic>=2.12.1
 Requires-Dist: python-dotenv>=1.1.1
-Requires-Dist: rich>=14.1.0
-Requires-Dist: tiktoken>=0.11.0
-Requires-Dist: typer>=0.16.1
+Requires-Dist: rich>=14.2.0
+Requires-Dist: tiktoken>=0.12.0
+Requires-Dist: typer>=0.19.2
 Requires-Dist: watchfiles>=1.1.0
+Provides-Extra: a2a
+Requires-Dist: fasta2a>=0.1.0; extra == 'a2a'
 Provides-Extra: mxbai
 Requires-Dist: mxbai-rerank>=0.1.6; extra == 'mxbai'
 Provides-Extra: voyageai
@@ -56,6 +58,7 @@ Retrieval-Augmented Generation (RAG) library built on LanceDB.
 - **File monitoring**: Auto-index files when run as server
 - **40+ file formats**: PDF, DOCX, HTML, Markdown, code files, URLs
 - **MCP server**: Expose as tools for AI assistants
+- **A2A agent**: Conversational agent with context and multi-turn dialogue
 - **CLI & Python API**: Use from command line or Python
 
 ## Quick Start
@@ -181,6 +184,24 @@ haiku-rag serve --stdio
 
 Provides tools for document management and search directly in your AI assistant.
 
+## A2A Agent
+
+Run as a conversational agent with the Agent-to-Agent protocol:
+
+```bash
+# Start the A2A server
+haiku-rag serve --a2a
+
+# Connect with the interactive client (in another terminal)
+haiku-rag a2aclient
+```
+
+The A2A agent provides:
+- Multi-turn dialogue with context
+- Intelligent multi-search for complex questions
+- Source citations with titles and URIs
+- Full document retrieval on request
+
 ## Documentation
 
 Full documentation at: https://ggozad.github.io/haiku.rag/