knowledge2 0.4.0__py3-none-any.whl

sdk/resources/a2a.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+
+import uuid
+from typing import Any
+
+from sdk._preview import preview_resource
+from sdk._request_options import RequestOptions
+from sdk._validation import require_str
+from sdk.resources._mixin_base import RequesterMixin
+from sdk.types import (
+    A2AAgentCardResponse,
+    A2AJsonRpcResponse,
+)
+
+
+@preview_resource
+class A2AMixin(RequesterMixin):
+    """SDK mixin for the K2 A2A (Agent-to-Agent) protocol adapter."""
+
+    def a2a_get_agent_card(
+        self,
+        corpus_id: str,
+        request_options: RequestOptions | None = None,
+    ) -> A2AAgentCardResponse:
+        """Fetch the A2A Agent Card for a corpus."""
+        corpus_id = require_str(corpus_id, "corpus_id")
+        data = self._request(
+            "GET",
+            f"/a2a/v1/agents/{corpus_id}/.well-known/agent.json",
+            request_options=request_options,
+        )
+        return self._maybe_validate(data, "A2AAgentCardResponse")
+
+    def a2a_send_message(
+        self,
+        corpus_id: str,
+        *,
+        operation: str,
+        text: str | None = None,
+        data: dict[str, Any] | None = None,
+        message_id: str | None = None,
+        context_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        configuration: dict[str, Any] | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> A2AJsonRpcResponse:
+        """Send an A2A message/send request to a corpus agent.
+
+        Args:
+            corpus_id: Target corpus ID.
+            operation: K2 operation (retrieve, answer, ingest.urls, ingest.text_batch).
+            text: Query text (required for retrieve/answer).
+            data: Structured data (search config, URLs, documents, etc.).
+            message_id: Optional message ID (auto-generated if omitted).
+            context_id: Optional context/session ID.
+            metadata: Additional metadata merged with operation.
+            configuration: Optional A2A configuration passed through to the agent.
+        """
+        corpus_id = require_str(corpus_id, "corpus_id")
+        parts: list[dict[str, Any]] = []
+        if text is not None:
+            parts.append({"kind": "text", "text": text})
+        if data is not None:
+            parts.append({"kind": "data", "data": data})
+
+        msg_meta: dict[str, Any] = dict(metadata or {})
+        # Keep operation authoritative even if metadata contains an `operation` key.
+        msg_meta["operation"] = operation
+
+        message: dict[str, Any] = {
+            "messageId": message_id or str(uuid.uuid4()),
+            "role": "user",
+            "parts": parts,
+            "metadata": msg_meta,
+        }
+        if context_id is not None:
+            message["contextId"] = context_id
+
+        params: dict[str, Any] = {"message": message}
+        if configuration is not None:
+            params["configuration"] = configuration
+
+        body = {
+            "jsonrpc": "2.0",
+            "id": str(uuid.uuid4()),
+            "method": "message/send",
+            "params": params,
+        }
+        resp = self._request(
+            "POST", f"/a2a/v1/agents/{corpus_id}", json=body, request_options=request_options
+        )
+        return self._maybe_validate(resp, "A2AJsonRpcResponse")
+
+    def a2a_retrieve(
+        self,
+        corpus_id: str,
+        query: str,
+        *,
+        top_k: int = 10,
+        filters: dict[str, Any] | None = None,
+        hybrid: dict[str, Any] | None = None,
+        rerank: dict[str, Any] | None = None,
+        return_config: dict[str, Any] | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> A2AJsonRpcResponse:
+        """Convenience wrapper for the A2A 'retrieve' operation."""
+        corpus_id = require_str(corpus_id, "corpus_id")
+        data: dict[str, Any] = {"top_k": top_k}
+        if filters is not None:
+            data["filters"] = filters
+        if hybrid is not None:
+            data["hybrid"] = hybrid
+        if rerank is not None:
+            data["rerank"] = rerank
+        if return_config is not None:
+            data["return"] = return_config
+        return self.a2a_send_message(
+            corpus_id, operation="retrieve", text=query, data=data, request_options=request_options
+        )
+
+    def a2a_answer(
+        self,
+        corpus_id: str,
+        query: str,
+        *,
+        top_k: int = 10,
+        filters: dict[str, Any] | None = None,
+        hybrid: dict[str, Any] | None = None,
+        rerank: dict[str, Any] | None = None,
+        generation: dict[str, Any] | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> A2AJsonRpcResponse:
+        """Convenience wrapper for the A2A 'answer' operation."""
+        corpus_id = require_str(corpus_id, "corpus_id")
+        data: dict[str, Any] = {"top_k": top_k}
+        if filters is not None:
+            data["filters"] = filters
+        if hybrid is not None:
+            data["hybrid"] = hybrid
+        if rerank is not None:
+            data["rerank"] = rerank
+        if generation is not None:
+            data["generation"] = generation
+        return self.a2a_send_message(
+            corpus_id, operation="answer", text=query, data=data, request_options=request_options
+        )
+
+    def a2a_ingest_urls(
+        self,
+        corpus_id: str,
+        urls: list[str | dict[str, Any]],
+        *,
+        auto_index: bool | None = None,
+        message_id: str | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> A2AJsonRpcResponse:
+        """Convenience wrapper for the A2A 'ingest.urls' operation."""
+        corpus_id = require_str(corpus_id, "corpus_id")
+        data: dict[str, Any] = {"urls": urls}
+        if auto_index is not None:
+            data["auto_index"] = auto_index
+        return self.a2a_send_message(
+            corpus_id,
+            operation="ingest.urls",
+            data=data,
+            message_id=message_id,
+            request_options=request_options,
+        )
+
+    def a2a_ingest_text_batch(
+        self,
+        corpus_id: str,
+        documents: list[dict[str, Any]],
+        *,
+        auto_index: bool | None = None,
+        message_id: str | None = None,
+        request_options: RequestOptions | None = None,
+    ) -> A2AJsonRpcResponse:
+        """Convenience wrapper for the A2A 'ingest.text_batch' operation."""
+        corpus_id = require_str(corpus_id, "corpus_id")
+        data: dict[str, Any] = {"documents": documents}
+        if auto_index is not None:
+            data["auto_index"] = auto_index
+        return self.a2a_send_message(
+            corpus_id,
+            operation="ingest.text_batch",
+            data=data,
+            message_id=message_id,
+            request_options=request_options,
+        )
+
+    def a2a_get_task(
+        self,
+        corpus_id: str,
+        task_id: str,
+        request_options: RequestOptions | None = None,
+    ) -> A2AJsonRpcResponse:
+        """Get the status of an async A2A task (maps to tasks/get)."""
+        corpus_id = require_str(corpus_id, "corpus_id")
+        task_id = require_str(task_id, "task_id")
+        body = {
+            "jsonrpc": "2.0",
+            "id": str(uuid.uuid4()),
+            "method": "tasks/get",
+            "params": {"id": task_id},
+        }
+        resp = self._request(
+            "POST", f"/a2a/v1/agents/{corpus_id}", json=body, request_options=request_options
+        )
+        return self._maybe_validate(resp, "A2AJsonRpcResponse")
+
+    def a2a_cancel_task(
+        self,
+        corpus_id: str,
+        task_id: str,
+        request_options: RequestOptions | None = None,
+    ) -> A2AJsonRpcResponse:
+        """Cancel an async A2A task (maps to tasks/cancel)."""
+        corpus_id = require_str(corpus_id, "corpus_id")
+        task_id = require_str(task_id, "task_id")
+        body = {
+            "jsonrpc": "2.0",
+            "id": str(uuid.uuid4()),
+            "method": "tasks/cancel",
+            "params": {"id": task_id},
+        }
+        resp = self._request(
+            "POST", f"/a2a/v1/agents/{corpus_id}", json=body, request_options=request_options
+        )
+        return self._maybe_validate(resp, "A2AJsonRpcResponse")