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

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:

haiku/rag/cli.py

@@ -128,10 +128,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,

haiku/rag/client.py

@@ -106,8 +106,8 @@ class HaikuRAG:
 
     async def create_document_from_source(
         self, source: str | Path, title: str | None = None, metadata: dict | None = None
-    ) -> Document:
-        """Create or update a document from a file path or URL.
+    ) -> Document | list[Document]:
+        """Create or update document(s) from a file path, directory, or URL.
 
         Checks if a document with the same URI already exists:
         - If MD5 is unchanged, returns existing document
@@ -115,11 +115,13 @@ class HaikuRAG:
         - If no document exists, creates a new one
 
         Args:
-            source: File path (as string or Path) or URL to parse
+            source: File path, directory (as string or Path), or URL to parse
+            title: Optional title (only used for single files, not directories)
             metadata: Optional metadata dictionary
 
         Returns:
-            Document instance (created, updated, or existing)
+            Document instance (created, updated, or existing) for single files/URLs
+            List of Document instances for directories
 
         Raises:
             ValueError: If the file/URL cannot be parsed or doesn't exist
@@ -142,6 +144,45 @@ class HaikuRAG:
         else:
             # Handle as regular file path
             source_path = Path(source) if isinstance(source, str) else source
+
+        # Handle directories
+        if source_path.is_dir():
+            documents = []
+            supported_extensions = set(FileReader.extensions)
+            for file_path in source_path.rglob("*"):
+                if (
+                    file_path.is_file()
+                    and file_path.suffix.lower() in supported_extensions
+                ):
+                    doc = await self._create_document_from_file(
+                        file_path, title=None, metadata=metadata
+                    )
+                    documents.append(doc)
+            return documents
+
+        # Handle single file
+        return await self._create_document_from_file(
+            source_path, title=title, metadata=metadata
+        )
+
+    async def _create_document_from_file(
+        self, source_path: Path, title: str | None = None, metadata: dict | None = None
+    ) -> Document:
+        """Create or update a document from a single file path.
+
+        Args:
+            source_path: Path to the file
+            title: Optional title
+            metadata: Optional metadata dictionary
+
+        Returns:
+            Document instance (created, updated, or existing)
+
+        Raises:
+            ValueError: If the file cannot be parsed or doesn't exist
+        """
+        metadata = metadata or {}
+
         if source_path.suffix.lower() not in FileReader.extensions:
             raise ValueError(f"Unsupported file extension: {source_path.suffix}")
 
@@ -592,6 +633,8 @@ class HaikuRAG:
                         new_doc = await self.create_document_from_source(
                             source=doc.uri, metadata=doc.metadata or {}
                         )
+                        # URIs always point to single files/URLs, never directories
+                        assert isinstance(new_doc, Document)
                         assert new_doc.id is not None, (
                             "New document ID should not be None"
                         )

haiku/rag/embeddings/base.py

@@ -1,3 +1,5 @@
+from typing import overload
+
 from haiku.rag.config import Config
 
 
@@ -9,6 +11,12 @@ class EmbedderBase:
         self._model = model
         self._vector_dim = vector_dim
 
+    @overload
+    async def embed(self, text: str) -> list[float]: ...
+
+    @overload
+    async def embed(self, text: list[str]) -> list[list[float]]: ...
+
     async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
         raise NotImplementedError(
             "Embedder is an abstract class. Please implement the embed method in a subclass."

haiku/rag/embeddings/ollama.py

@@ -1,3 +1,5 @@
+from typing import overload
+
 from openai import AsyncOpenAI
 
 from haiku.rag.config import Config
@@ -5,6 +7,12 @@ from haiku.rag.embeddings.base import EmbedderBase
 
 
 class Embedder(EmbedderBase):
+    @overload
+    async def embed(self, text: str) -> list[float]: ...
+
+    @overload
+    async def embed(self, text: list[str]) -> list[list[float]]: ...
+
     async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
         client = AsyncOpenAI(base_url=f"{Config.OLLAMA_BASE_URL}/v1", api_key="dummy")
         if not text:

haiku/rag/embeddings/openai.py

@@ -1,9 +1,17 @@
+from typing import overload
+
 from openai import AsyncOpenAI
 
 from haiku.rag.embeddings.base import EmbedderBase
 
 
 class Embedder(EmbedderBase):
+    @overload
+    async def embed(self, text: str) -> list[float]: ...
+
+    @overload
+    async def embed(self, text: list[str]) -> list[list[float]]: ...
+
     async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
         client = AsyncOpenAI()
         if not text:

haiku/rag/embeddings/vllm.py

@@ -1,3 +1,5 @@
+from typing import overload
+
 from openai import AsyncOpenAI
 
 from haiku.rag.config import Config
@@ -5,6 +7,12 @@ from haiku.rag.embeddings.base import EmbedderBase
 
 
 class Embedder(EmbedderBase):
+    @overload
+    async def embed(self, text: str) -> list[float]: ...
+
+    @overload
+    async def embed(self, text: list[str]) -> list[list[float]]: ...
+
     async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
         client = AsyncOpenAI(
             base_url=f"{Config.VLLM_EMBEDDINGS_BASE_URL}/v1", api_key="dummy"

haiku/rag/embeddings/voyageai.py

@@ -1,9 +1,17 @@
 try:
+    from typing import overload
+
     from voyageai.client import Client  # type: ignore
 
     from haiku.rag.embeddings.base import EmbedderBase
 
     class Embedder(EmbedderBase):
+        @overload
+        async def embed(self, text: str) -> list[float]: ...
+
+        @overload
+        async def embed(self, text: list[str]) -> list[list[float]]: ...
+
         async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
             client = Client()
             if not text:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.12.0
+Version: 0.12.1
 Summary: Agentic Retrieval Augmented Generation (RAG) with LanceDB
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT
@@ -24,7 +24,7 @@ Requires-Dist: httpx>=0.28.1
 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: pydantic>=2.12.2
 Requires-Dist: python-dotenv>=1.1.1
 Requires-Dist: rich>=14.2.0
 Requires-Dist: tiktoken>=0.12.0
@@ -211,4 +211,6 @@ Full documentation at: https://ggozad.github.io/haiku.rag/
 - [CLI](https://ggozad.github.io/haiku.rag/cli/) - Command reference
 - [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs
 - [Agents](https://ggozad.github.io/haiku.rag/agents/) - QA agent and multi-agent research
+- [MCP Server](https://ggozad.github.io/haiku.rag/mcp/) - Model Context Protocol integration
+- [A2A Agent](https://ggozad.github.io/haiku.rag/a2a/) - Agent-to-Agent protocol support
 - [Benchmarks](https://ggozad.github.io/haiku.rag/benchmarks/) - Performance Benchmarks

{haiku_rag-0.12.0.dist-info → haiku_rag-0.12.1.dist-info}/RECORD

@@ -1,8 +1,8 @@
 haiku/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-haiku/rag/app.py,sha256=lOoWMqQLNokMXcGh_6tYAeq_qaIV1dPOqDlQuetHLl0,21217
+haiku/rag/app.py,sha256=nrfg3iGWP_HJBfwUFPv60_1Do8YK3WQYWZhq22r001s,21548
 haiku/rag/chunker.py,sha256=PVe6ysv8UlacUd4Zb3_8RFWIaWDXnzBAy2VDJ4TaUsE,1555
-haiku/rag/cli.py,sha256=7pf8lcLlu2S7KW-M8JqOPG7Xxy1aVbv6e8CPZbCL21w,13169
-haiku/rag/client.py,sha256=tk3BWa9u2mtUEMEOEcgF2ebUqdZRv4jSnu4S6Zho-Xg,24210
+haiku/rag/cli.py,sha256=ghmfvCmoitgySZsF6t5UQjsm3_rul0KUh0L774BzeuI,13196
+haiku/rag/client.py,sha256=GVXHq9weIaFdcZvO9a4YO1WnrroJJUXVVriDGdMxpH4,25855
 haiku/rag/config.py,sha256=FBsMMijl5PxIfPGifk_AJVRjL4omb03jfoZm0P_VqxI,2743
 haiku/rag/logging.py,sha256=dm65AwADpcQsH5OAPtRA-4hsw0w5DK-sGOvzYkj6jzw,1720
 haiku/rag/mcp.py,sha256=DZk-IJgVjAesu-vvqVd5BYnfDWKWNR6TQugKgdoFrvg,8976
@@ -11,7 +11,7 @@ haiku/rag/monitor.py,sha256=VP3bqY0mEodOP60eN4RMldgrL1ti5gMjuDuQ-_vBvFc,2759
 haiku/rag/reader.py,sha256=aW8LG0X31kVWS7kU2tKVpe8RqP3Ne_oIidd_X3UDLH0,3307
 haiku/rag/utils.py,sha256=dBzhKaOHI9KRiJqHErcXUnqtnXY2AgOK8PCLA3rhO0A,6115
 haiku/rag/a2a/__init__.py,sha256=4SlJBr9GUVZ0879o5VI6-qpcBKpieP2hW4hmNbm8NGg,5933
-haiku/rag/a2a/client.py,sha256=HIHF6HkIWU9UDmsl-hjJwMBRyNdHnSm0Z4K2d688CuQ,8644
+haiku/rag/a2a/client.py,sha256=awuiHXgVHn1uzaEXE98RIqqKHj1JjszOvn9WI3Jtth8,8760
 haiku/rag/a2a/context.py,sha256=SofkFUZcGonoJcgZh-RGqHTh0UWT4J7Zl4Mz6WDkMl4,2053
 haiku/rag/a2a/models.py,sha256=XhGYj2g3rgVM4JoCDXlll0YjaysqdalybJrBqFXSwl4,689
 haiku/rag/a2a/prompts.py,sha256=yCla8x0hbOhKrkuaqVrF1upn-YjQM3-2NsE2TSnet0M,3030
@@ -19,11 +19,11 @@ haiku/rag/a2a/skills.py,sha256=dwyD2Bn493eL3Vf4uQzmyxj_9IUSb66kQ-085FBAuCs,2701
 haiku/rag/a2a/storage.py,sha256=c8vmGCiZ3nuV9wUuTnwpoRD2HVVvK2JPySQOc5PVMvg,2759
 haiku/rag/a2a/worker.py,sha256=S9hiA1ncpJPdtN0eEmMjsvr5LQ4wMVN5R8CjYkTeohU,12367
 haiku/rag/embeddings/__init__.py,sha256=44IfDITGIFTflGT6UEmiYOwpWFVbYv5smLY59D0YeCs,1419
-haiku/rag/embeddings/base.py,sha256=BnSviKrlzjv3L0sZJs_T-pxfawd-bcTak-rsX-D2f3A,497
-haiku/rag/embeddings/ollama.py,sha256=c1BeKTgpymniZw1sm4iAIdK5vA0MYoRzHLcd2_pFA44,638
-haiku/rag/embeddings/openai.py,sha256=bwoUVlzu9UtbDpN7CtG6OPt0d5tfJNeje4lR81Btpl0,546
-haiku/rag/embeddings/vllm.py,sha256=7ocp9D9bD1R5rqRIC4-Vih9VlKQNuD429k8-9wu234E,669
-haiku/rag/embeddings/voyageai.py,sha256=I4kVdT2KPtwcbjxD22GWJmgcIQIEEHpkOY2_QbFh7mQ,712
+haiku/rag/embeddings/base.py,sha256=Aw4kjfVn2can0R17pdiAgpPRyk5BpdBgMXuor5mstDY,682
+haiku/rag/embeddings/ollama.py,sha256=KXq-eJ58co5rwYchIO3kpvIv0OBwMJkwMXq1xDsETz0,823
+haiku/rag/embeddings/openai.py,sha256=BfmPni567DH8KqwLCPiOmr3q-dpzpOJkvFFoUuTR5as,731
+haiku/rag/embeddings/vllm.py,sha256=wgul0nMWTn6Q1aKA4DJe03EktsRoBxEgtB7gfpWVOyQ,854
+haiku/rag/embeddings/voyageai.py,sha256=6vEuk6q510AJv-K2lL93P2dVrziAjELTOe_w_Zp5YT4,917
 haiku/rag/graph/__init__.py,sha256=BHfMchuUO_UhHKpjjGHjd6xPxNkrIwJzHn4YJiLqG1g,62
 haiku/rag/graph/base.py,sha256=DepZqLF9E64YCCkjmbqmgyp28oNp69WfJCXp614xzh0,819
 haiku/rag/graph/common.py,sha256=xTejucXei3x9tqbal3ZS_64lZAC6Bw3-QfXPniZcZEw,986
@@ -69,8 +69,8 @@ haiku/rag/store/repositories/settings.py,sha256=ObrDrzxHn-yA1WcbgIoJoVmAbVvQHAFv
 haiku/rag/store/upgrades/__init__.py,sha256=RQ8A6rEXBASLb5PD9vdDnEas_m_GgRzzdVu4B88Snqc,1975
 haiku/rag/store/upgrades/v0_10_1.py,sha256=qNGnxj6hoHaHJ1rKTiALfw0c9NQOi0KAK-VZCD_073A,1959
 haiku/rag/store/upgrades/v0_9_3.py,sha256=NrjNilQSgDtFWRbL3ZUtzQzJ8tf9u0dDRJtnDFwwbdw,3322
-haiku_rag-0.12.0.dist-info/METADATA,sha256=RNOibOVVUHbG7NMf_isKEKmsCnrcyCDXKN-34qZu_fU,7295
-haiku_rag-0.12.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haiku_rag-0.12.0.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
-haiku_rag-0.12.0.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
-haiku_rag-0.12.0.dist-info/RECORD,,
+haiku_rag-0.12.1.dist-info/METADATA,sha256=POFHzbGYiVj7UkX_1VSA8zUByIiQEG1dPePWO55T7nU,7477
+haiku_rag-0.12.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+haiku_rag-0.12.1.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
+haiku_rag-0.12.1.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
+haiku_rag-0.12.1.dist-info/RECORD,,