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

haiku/rag/a2a/__init__.py

@@ -57,12 +57,12 @@ def create_a2a_app(
     """
     base_storage = InMemoryStorage()
     storage = LRUMemoryStorage(
-        storage=base_storage, max_contexts=Config.A2A_MAX_CONTEXTS
+        storage=base_storage, max_contexts=Config.a2a.max_contexts
     )
     broker = InMemoryBroker()
 
     # Create the agent with native search tool
-    model = get_model(Config.QA_PROVIDER, Config.QA_MODEL)
+    model = get_model(Config.qa.provider, Config.qa.model)
     agent = Agent(
         model=model,
         deps_type=AgentDependencies,
@@ -120,7 +120,7 @@ def create_a2a_app(
     # Create FastA2A app with custom worker lifecycle
     @asynccontextmanager
     async def lifespan(app):
-        logger.info(f"Started A2A server (max contexts: {Config.A2A_MAX_CONTEXTS})")
+        logger.info(f"Started A2A server (max contexts: {Config.a2a.max_contexts})")
         async with app.task_manager:
             async with worker.run():
                 yield

haiku/rag/a2a/client.py

@@ -7,9 +7,18 @@ from rich.console import Console
 from rich.markdown import Markdown
 from rich.prompt import Prompt
 
+try:
+    from fasta2a.client import A2AClient as FastA2AClient
+    from fasta2a.schema import Message, TextPart
+except ImportError as e:
+    raise ImportError(
+        "A2A support requires the 'a2a' extra. "
+        "Install with: uv pip install 'haiku.rag[a2a]'"
+    ) from e
+
 
 class A2AClient:
-    """Simple A2A protocol client."""
+    """Interactive A2A protocol client."""
 
     def __init__(self, base_url: str = "http://localhost:8000"):
         """Initialize A2A client.
@@ -18,11 +27,12 @@ class A2AClient:
             base_url: Base URL of the A2A server
         """
         self.base_url = base_url.rstrip("/")
-        self.client = httpx.AsyncClient(timeout=60.0)
+        http_client = httpx.AsyncClient(timeout=60.0)
+        self._client = FastA2AClient(base_url=base_url, http_client=http_client)
 
     async def close(self):
         """Close the HTTP client."""
-        await self.client.aclose()
+        await self._client.http_client.aclose()
 
     async def get_agent_card(self) -> dict[str, Any]:
         """Fetch the agent card from the A2A server.
@@ -30,7 +40,9 @@ class A2AClient:
         Returns:
             Agent card dictionary with agent capabilities and metadata
         """
-        response = await self.client.get(f"{self.base_url}/.well-known/agent-card.json")
+        response = await self._client.http_client.get(
+            f"{self.base_url}/.well-known/agent-card.json"
+        )
         response.raise_for_status()
         return response.json()
 
@@ -53,46 +65,38 @@ class A2AClient:
         if context_id is None:
             context_id = str(uuid.uuid4())
 
-        message_id = str(uuid.uuid4())
-
-        payload: dict[str, Any] = {
-            "jsonrpc": "2.0",
-            "method": "message/send",
-            "params": {
-                "contextId": context_id,
-                "message": {
-                    "kind": "message",
-                    "role": "user",
-                    "messageId": message_id,
-                    "parts": [{"kind": "text", "text": text}],
-                },
-            },
-            "id": 1,
-        }
+        message = Message(
+            kind="message",
+            role="user",
+            message_id=str(uuid.uuid4()),
+            parts=[TextPart(kind="text", text=text)],
+        )
 
+        metadata: dict[str, Any] = {"contextId": context_id}
         if skill_id:
-            payload["params"]["skillId"] = skill_id
+            metadata["skillId"] = skill_id
 
-        response = await self.client.post(
-            self.base_url,
-            json=payload,
-            headers={"Content-Type": "application/json"},
-        )
-        response.raise_for_status()
-        initial_response = response.json()
+        response = await self._client.send_message(message, metadata=metadata)
 
-        # Extract task ID from response
-        result = initial_response.get("result", {})
-        task_id = result.get("id")
+        if "error" in response:
+            return {"error": response["error"]}
 
-        if not task_id:
-            return initial_response
+        result = response.get("result")
+        if not result:
+            return {"result": result}
 
-        # Poll for task completion
-        return await self.wait_for_task(task_id)
+        # Result can be either Task or Message - check if it's a Task with an id
+        if result.get("kind") == "task":
+            task_id = result.get("id")
+            if task_id:
+                # Poll for task completion
+                return await self.wait_for_task(task_id)
+
+        # Return the message directly
+        return {"result": result}
 
     async def wait_for_task(
-        self, task_id: str, max_wait: int = 60, poll_interval: float = 0.5
+        self, task_id: str, max_wait: int = 120, poll_interval: float = 0.5
     ) -> dict[str, Any]:
         """Poll for task completion.
 
@@ -109,27 +113,19 @@ class A2AClient:
         start_time = time.time()
 
         while time.time() - start_time < max_wait:
-            payload = {
-                "jsonrpc": "2.0",
-                "method": "tasks/get",
-                "params": {"id": task_id},
-                "id": 2,
-            }
-
-            response = await self.client.post(
-                self.base_url,
-                json=payload,
-                headers={"Content-Type": "application/json"},
-            )
-            response.raise_for_status()
-            task = response.json()
-
-            result = task.get("result", {})
-            status = result.get("status", {})
-            state = status.get("state")
+            task_response = await self._client.get_task(task_id)
+
+            if "error" in task_response:
+                return {"error": task_response["error"]}
+
+            task = task_response.get("result")
+            if not task:
+                raise Exception("No task in response")
+
+            state = task.get("status", {}).get("state")
 
             if state == "completed":
-                return task
+                return {"result": task}
             elif state == "failed":
                 raise Exception(f"Task failed: {task}")
 
@@ -191,6 +187,7 @@ def print_response(response: dict[str, Any], console: Console):
 
     # Print artifacts summary with details
     if artifacts:
+        console.rule("[dim]Artifacts generated[/dim]")
         summary_lines = []
 
         for artifact in artifacts:

haiku/rag/app.py

@@ -160,13 +160,20 @@ class HaikuRAGApp:
         self, source: str, title: str | None = None, metadata: dict | None = None
     ):
         async with HaikuRAG(db_path=self.db_path) as self.client:
-            doc = await self.client.create_document_from_source(
+            result = await self.client.create_document_from_source(
                 source, title=title, metadata=metadata
             )
-            self._rich_print_document(doc, truncate=True)
-            self.console.print(
-                f"[bold green]Document {doc.id} added successfully.[/bold green]"
-            )
+            if isinstance(result, list):
+                for doc in result:
+                    self._rich_print_document(doc, truncate=True)
+                self.console.print(
+                    f"[bold green]{len(result)} documents added successfully.[/bold green]"
+                )
+            else:
+                self._rich_print_document(result, truncate=True)
+                self.console.print(
+                    f"[bold green]Document {result.id} added successfully.[/bold green]"
+                )
 
     async def get_document(self, doc_id: str):
         async with HaikuRAG(db_path=self.db_path) as self.client:
@@ -224,8 +231,8 @@ class HaikuRAGApp:
                     )
 
                     start_node = DeepQAPlanNode(
-                        provider=Config.QA_PROVIDER,
-                        model=Config.QA_MODEL,
+                        provider=Config.qa.provider,
+                        model=Config.qa.model,
                     )
 
                     result = await graph.run(
@@ -271,8 +278,8 @@ class HaikuRAGApp:
                 )
 
                 start = PlanNode(
-                    provider=Config.RESEARCH_PROVIDER or Config.QA_PROVIDER,
-                    model=Config.RESEARCH_MODEL or Config.QA_MODEL,
+                    provider=Config.research.provider or Config.qa.provider,
+                    model=Config.research.model or Config.qa.model,
                 )
                 report = None
                 async for event in stream_research_graph(graph, start, state, deps):
@@ -467,7 +474,9 @@ class HaikuRAGApp:
 
             # Start file monitor if enabled
             if enable_monitor:
-                monitor = FileWatcher(paths=Config.MONITOR_DIRECTORIES, client=client)
+                monitor = FileWatcher(
+                    paths=Config.storage.monitor_directories, client=client
+                )
                 monitor_task = asyncio.create_task(monitor.observe())
                 tasks.append(monitor_task)
 

haiku/rag/chunker.py

@@ -22,7 +22,7 @@ class Chunker:
 
     def __init__(
         self,
-        chunk_size: int = Config.CHUNK_SIZE,
+        chunk_size: int = Config.processing.chunk_size,
     ):
         self.chunk_size = chunk_size
         tokenizer = OpenAITokenizer(

haiku/rag/cli.py

@@ -42,10 +42,21 @@ def main(
         callback=version_callback,
         help="Show version and exit",
     ),
+    config: Path | None = typer.Option(
+        None,
+        "--config",
+        help="Path to YAML configuration file",
+    ),
 ):
     """haiku.rag CLI - Vector database RAG system"""
+    # Store config path in environment for config loader to use
+    if config:
+        import os
+
+        os.environ["HAIKU_RAG_CONFIG_PATH"] = str(config.absolute())
+
     # Configure logging minimally for CLI context
-    if Config.ENV == "development":
+    if Config.environment == "development":
         # Lazy import logfire only in development
         try:
             import logfire  # type: ignore
@@ -69,7 +80,7 @@ def main(
 @cli.command("list", help="List all stored documents")
 def list_documents(
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        Config.storage.data_dir / "haiku.rag.lancedb",
         "--db",
         help="Path to the LanceDB database file",
     ),
@@ -116,7 +127,7 @@ def add_document_text(
         metavar="KEY=VALUE",
     ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        Config.storage.data_dir / "haiku.rag.lancedb",
         "--db",
         help="Path to the LanceDB database file",
     ),
@@ -128,10 +139,10 @@ def add_document_text(
     asyncio.run(app.add_document_from_text(text=text, metadata=metadata or None))
 
 
-@cli.command("add-src", help="Add a document from a file path or URL")
+@cli.command("add-src", help="Add a document from a file path, directory, or URL")
 def add_document_src(
     source: str = typer.Argument(
-        help="The file path or URL of the document to add",
+        help="The file path, directory, or URL of the document(s) to add",
     ),
     title: str | None = typer.Option(
         None,
@@ -145,7 +156,7 @@ def add_document_src(
         metavar="KEY=VALUE",
     ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        Config.storage.data_dir / "haiku.rag.lancedb",
         "--db",
         help="Path to the LanceDB database file",
     ),
@@ -167,7 +178,7 @@ def get_document(
         help="The ID of the document to get",
     ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        Config.storage.data_dir / "haiku.rag.lancedb",
         "--db",
         help="Path to the LanceDB database file",
     ),
@@ -184,7 +195,7 @@ def delete_document(
         help="The ID of the document to delete",
     ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        Config.storage.data_dir / "haiku.rag.lancedb",
         "--db",
         help="Path to the LanceDB database file",
     ),
@@ -211,7 +222,7 @@ def search(
         help="Maximum number of results to return",
     ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        Config.storage.data_dir / "haiku.rag.lancedb",
         "--db",
         help="Path to the LanceDB database file",
     ),
@@ -228,7 +239,7 @@ def ask(
         help="The question to ask",
     ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        Config.storage.data_dir / "haiku.rag.lancedb",
         "--db",
         help="Path to the LanceDB database file",
     ),
@@ -276,7 +287,7 @@ def research(
         help="Max concurrent searches per iteration (planned)",
     ),
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        Config.storage.data_dir / "haiku.rag.lancedb",
         "--db",
         help="Path to the LanceDB database file",
     ),
@@ -308,13 +319,61 @@ def settings():
     app.show_settings()
 
 
+@cli.command("init-config", help="Generate a YAML configuration file")
+def init_config(
+    output: Path = typer.Argument(
+        Path("haiku.rag.yaml"),
+        help="Output path for the config file",
+    ),
+    from_env: bool = typer.Option(
+        False,
+        "--from-env",
+        help="Migrate settings from .env file",
+    ),
+):
+    """Generate a YAML configuration file with defaults or from .env."""
+    import yaml
+
+    from haiku.rag.config.loader import generate_default_config, load_config_from_env
+
+    if output.exists():
+        typer.echo(
+            f"Error: {output} already exists. Remove it first or choose a different path."
+        )
+        raise typer.Exit(1)
+
+    if from_env:
+        # Load from environment variables (including .env if present)
+        from dotenv import load_dotenv
+
+        load_dotenv()
+        config_data = load_config_from_env()
+        if not config_data:
+            typer.echo("Warning: No environment variables found to migrate.")
+            typer.echo("Generating default configuration instead.")
+            config_data = generate_default_config()
+    else:
+        config_data = generate_default_config()
+
+    # Write YAML with comments
+    with open(output, "w") as f:
+        f.write("# haiku.rag configuration file\n")
+        f.write(
+            "# See https://ggozad.github.io/haiku.rag/configuration/ for details\n\n"
+        )
+        yaml.dump(config_data, f, default_flow_style=False, sort_keys=False)
+
+    typer.echo(f"Configuration file created: {output}")
+    typer.echo("Edit the file to customize your settings.")
+
+
 @cli.command(
     "rebuild",
     help="Rebuild the database by deleting all chunks and re-indexing all documents",
 )
 def rebuild(
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        Config.storage.data_dir / "haiku.rag.lancedb",
         "--db",
         help="Path to the LanceDB database file",
     ),
@@ -328,7 +387,7 @@ def rebuild(
 @cli.command("vacuum", help="Optimize and clean up all tables to reduce disk usage")
 def vacuum(
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        Config.storage.data_dir / "haiku.rag.lancedb",
         "--db",
         help="Path to the LanceDB database file",
     ),
@@ -342,7 +401,7 @@ def vacuum(
 @cli.command("info", help="Show read-only database info (no upgrades or writes)")
 def info(
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        Config.storage.data_dir / "haiku.rag.lancedb",
         "--db",
         help="Path to the LanceDB database file",
     ),
@@ -371,7 +430,7 @@ def download_models_cmd():
 )
 def serve(
     db: Path = typer.Option(
-        Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
+        Config.storage.data_dir / "haiku.rag.lancedb",
         "--db",
         help="Path to the LanceDB database file",
     ),
@@ -442,24 +501,6 @@ def serve(
     )
 
 
-@cli.command("migrate", help="Migrate an SQLite database to LanceDB")
-def migrate(
-    sqlite_path: Path = typer.Argument(
-        help="Path to the SQLite database file to 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:
-        raise typer.Exit(1)
-
-
 @cli.command(
     "a2aclient", help="Run interactive client to chat with haiku.rag's A2A server"
 )