axon-protocol 0.1.0__py3-none-any.whl

axon/__init__.py

@@ -0,0 +1,51 @@
+from axon.client import AxonClient, AxonSyncClient
+from axon.messages import MessagesClient, SyncMessagesClient
+from axon.types import (
+    MemoryResult,
+    MemorySearchResponse,
+    StoredMemory,
+    LockInfo,
+    LockStatus,
+    ReasoningStep,
+    StepsLogger,
+    ReceiptInfo,
+    ReceiptVerifyResult,
+)
+from axon.exceptions import (
+    AxonError,
+    AuthError,
+    LockConflictError,
+    NotFoundError,
+    AxonPermissionError,
+    RateLimitError,
+    ServerError,
+    AxonConnectionError,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "AxonClient",
+    "AxonSyncClient",
+    "MessagesClient",
+    "SyncMessagesClient",
+    # Types
+    "MemoryResult",
+    "MemorySearchResponse",
+    "StoredMemory",
+    "LockInfo",
+    "LockStatus",
+    "ReasoningStep",
+    "StepsLogger",
+    "ReceiptInfo",
+    "ReceiptVerifyResult",
+    # Exceptions
+    "AxonError",
+    "AuthError",
+    "LockConflictError",
+    "NotFoundError",
+    "AxonPermissionError",
+    "RateLimitError",
+    "ServerError",
+    "AxonConnectionError",
+]

axon/_base.py

@@ -0,0 +1,142 @@
+import httpx
+import time
+import asyncio
+from axon.exceptions import (
+    AuthError,
+    NotFoundError,
+    AxonPermissionError,
+    LockConflictError,
+    RateLimitError,
+    ServerError,
+    AxonConnectionError,
+)
+
+
+class _BaseClient:
+    def __init__(self, http: httpx.AsyncClient, base_url: str):
+        self._http = http
+        self._base_url = base_url.rstrip("/")
+
+    async def _request(self, method: str, path: str, **kwargs) -> dict:
+        url = f"{self._base_url}{path}"
+        retries = 3
+        backoff = 0.5
+
+        for attempt in range(retries):
+            try:
+                response = await self._http.request(method, url, **kwargs)
+                if response.status_code >= 500 and attempt < retries - 1:
+                    await asyncio.sleep(backoff * (2 ** attempt))
+                    continue
+                break
+            except httpx.ConnectError:
+                if attempt == retries - 1:
+                    raise AxonConnectionError(
+                        f"Cannot connect to Axon server at {self._base_url}. "
+                        f"Make sure the server is running."
+                    )
+                await asyncio.sleep(backoff * (2 ** attempt))
+            except httpx.TimeoutException:
+                if attempt == retries - 1:
+                    raise AxonConnectionError(
+                        f"Request to {url} timed out. Server may be overloaded."
+                    )
+                await asyncio.sleep(backoff * (2 ** attempt))
+
+        # Success
+        if response.status_code == 200:
+            return response.json()
+
+        # Parse error detail from response body
+        try:
+            detail = response.json().get("detail", response.text)
+        except Exception:
+            detail = response.text
+
+        # Map status codes to specific exceptions
+        if response.status_code == 401:
+            raise AuthError(f"Authentication failed: {detail}", 401)
+        elif response.status_code == 403:
+            raise AxonPermissionError(f"Permission denied: {detail}", 403)
+        elif response.status_code == 404:
+            raise NotFoundError(f"Not found: {detail}", 404)
+        elif response.status_code == 409:
+            raise LockConflictError("", detail=detail)
+        elif response.status_code == 429:
+            retry_after = int(response.headers.get("Retry-After", 60))
+            raise RateLimitError(retry_after)
+        elif response.status_code >= 500:
+            raise ServerError(
+                f"Server error ({response.status_code}): {detail}",
+                response.status_code,
+            )
+        else:
+            raise ServerError(
+                f"Unexpected status {response.status_code}: {detail}",
+                response.status_code,
+            )
+
+
+class _BaseSyncClient:
+    def __init__(self, http: httpx.Client, base_url: str):
+        self._http = http
+        self._base_url = base_url.rstrip("/")
+
+    def _request(self, method: str, path: str, **kwargs) -> dict:
+        url = f"{self._base_url}{path}"
+        retries = 3
+        backoff = 0.5
+
+        for attempt in range(retries):
+            try:
+                response = self._http.request(method, url, **kwargs)
+                if response.status_code >= 500 and attempt < retries - 1:
+                    time.sleep(backoff * (2 ** attempt))
+                    continue
+                break
+            except httpx.ConnectError:
+                if attempt == retries - 1:
+                    raise AxonConnectionError(
+                        f"Cannot connect to Axon server at {self._base_url}. "
+                        f"Make sure the server is running."
+                    )
+                time.sleep(backoff * (2 ** attempt))
+            except httpx.TimeoutException:
+                if attempt == retries - 1:
+                    raise AxonConnectionError(
+                        f"Request to {url} timed out. Server may be overloaded."
+                    )
+                time.sleep(backoff * (2 ** attempt))
+
+        # Success
+        if response.status_code == 200:
+            return response.json()
+
+        # Parse error detail from response body
+        try:
+            detail = response.json().get("detail", response.text)
+        except Exception:
+            detail = response.text
+
+        # Map status codes to specific exceptions
+        if response.status_code == 401:
+            raise AuthError(f"Authentication failed: {detail}", 401)
+        elif response.status_code == 403:
+            raise AxonPermissionError(f"Permission denied: {detail}", 403)
+        elif response.status_code == 404:
+            raise NotFoundError(f"Not found: {detail}", 404)
+        elif response.status_code == 409:
+            raise LockConflictError("", detail=detail)
+        elif response.status_code == 429:
+            retry_after = int(response.headers.get("Retry-After", 60))
+            raise RateLimitError(retry_after)
+        elif response.status_code >= 500:
+            raise ServerError(
+                f"Server error ({response.status_code}): {detail}",
+                response.status_code,
+            )
+        else:
+            raise ServerError(
+                f"Unexpected status {response.status_code}: {detail}",
+                response.status_code,
+            )

axon/client.py

@@ -0,0 +1,132 @@
+import httpx
+from axon.memory import MemoryClient, SyncMemoryClient
+from axon.coordination import CoordinationClient, SyncCoordinationClient
+from axon.receipts import ReceiptsClient, SyncReceiptsClient
+from axon.events import EventsClient, SyncEventsClient
+from axon.messages import MessagesClient, SyncMessagesClient
+
+
+class AxonClient:
+    """
+    Async entry point for the Axon Protocol Python SDK.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        project_id: str,
+        agent_token: str = None,
+        agent_id: str = None,
+        base_url: str = "http://localhost:8000",
+        timeout: float = 30.0,
+    ):
+        self.api_key = api_key
+        self.project_id = project_id
+        self.agent_token = agent_token
+        self.agent_id = agent_id
+        self.base_url = base_url.rstrip("/")
+
+        headers = {}
+        if api_key:
+            headers["X-API-Key"] = api_key
+        if agent_token:
+            headers["Authorization"] = f"Bearer {agent_token}"
+
+        self._http = httpx.AsyncClient(
+            headers=headers,
+            timeout=httpx.Timeout(timeout),
+            limits=httpx.Limits(
+                max_connections=20,
+                max_keepalive_connections=10,
+                keepalive_expiry=30,
+            ),
+        )
+
+        self.memory = MemoryClient(self._http, self.base_url)
+        self.lock = CoordinationClient(self._http, self.base_url)
+        self.receipts = ReceiptsClient(self._http, self.base_url)
+        self.events = EventsClient(self.base_url, api_key, project_id)
+        self.messages = MessagesClient(self._http, self.base_url)
+
+    async def ping(self) -> bool:
+        """
+        Check if the Axon server is reachable (async).
+        """
+        try:
+            response = await self._http.get(f"{self.base_url}/v1/health")
+            return response.status_code == 200
+        except Exception:
+            return False
+
+    async def close(self):
+        """Close the underlying HTTP connection pool."""
+        await self._http.aclose()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
+
+class AxonSyncClient:
+    """
+    Sync entry point for the Axon Protocol Python SDK.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        project_id: str,
+        agent_token: str = None,
+        agent_id: str = None,
+        base_url: str = "http://localhost:8000",
+        timeout: float = 30.0,
+    ):
+        self.api_key = api_key
+        self.project_id = project_id
+        self.agent_token = agent_token
+        self.agent_id = agent_id
+        self.base_url = base_url.rstrip("/")
+
+        headers = {}
+        if api_key:
+            headers["X-API-Key"] = api_key
+        if agent_token:
+            headers["Authorization"] = f"Bearer {agent_token}"
+
+        self._http = httpx.Client(
+            headers=headers,
+            timeout=httpx.Timeout(timeout),
+            limits=httpx.Limits(
+                max_connections=20,
+                max_keepalive_connections=10,
+                keepalive_expiry=30,
+            ),
+        )
+
+        self.memory = SyncMemoryClient(self._http, self.base_url)
+        self.lock = SyncCoordinationClient(self._http, self.base_url)
+        self.receipts = SyncReceiptsClient(self._http, self.base_url)
+        self.events = SyncEventsClient(self.base_url, api_key, project_id)
+        self.messages = SyncMessagesClient(self._http, self.base_url)
+
+    def ping(self) -> bool:
+        """
+        Check if the Axon server is reachable (sync).
+        """
+        try:
+            response = self._http.get(f"{self.base_url}/v1/health")
+            return response.status_code == 200
+        except Exception:
+            return False
+
+    def close(self):
+        """Close the underlying HTTP connection pool."""
+        self._http.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()

axon/coordination.py

@@ -0,0 +1,161 @@
+from contextlib import asynccontextmanager, contextmanager
+from axon._base import _BaseClient, _BaseSyncClient
+from axon.types import LockInfo, LockStatus
+
+
+class CoordinationClient(_BaseClient):
+
+    async def acquire(
+        self,
+        resource_id: str,
+        timeout: int = 300,
+        metadata: dict = None,
+    ) -> LockInfo:
+        """
+        Acquire an exclusive lock on a resource (async).
+        """
+        body = {"resource_id": resource_id, "timeout": timeout}
+        if metadata is not None:
+            body["metadata"] = metadata
+
+        data = await self._request("POST", "/v1/lock/acquire", json=body)
+        return LockInfo(
+            lock_id=data["lock_id"],
+            resource_id=data["resource_id"],
+            expires_at=data["expires_at"],
+        )
+
+    async def release(self, resource_id: str) -> bool:
+        """
+        Release a lock that you hold (async).
+        """
+        data = await self._request(
+            "POST",
+            "/v1/lock/release",
+            params={"resource_id": resource_id},
+        )
+        return data.get("released", False)
+
+    async def status(self, resource_id: str) -> LockStatus:
+        """
+        Check lock status of a resource (async).
+        """
+        data = await self._request("GET", f"/v1/lock/status/{resource_id}")
+        return LockStatus(
+            locked=data["locked"],
+            resource_id=data["resource_id"],
+            holder_agent_id=data.get("holder_agent_id"),
+            locked_at=data.get("locked_at"),
+            expires_at=data.get("expires_at"),
+        )
+
+    async def list_active(self) -> list[dict]:
+        """
+        List all active locks in the project (async).
+        """
+        data = await self._request("GET", "/v1/lock/list")
+        return data.get("locks", [])
+
+    @asynccontextmanager
+    async def hold(
+        self,
+        resource_id: str,
+        timeout: int = 300,
+        metadata: dict = None,
+    ):
+        """
+        Context manager that acquires a lock and automatically releases it on exit (async).
+        """
+        lock_info = await self.acquire(resource_id, timeout=timeout, metadata=metadata)
+        try:
+            yield lock_info
+        finally:
+            try:
+                await self.release(resource_id)
+            except Exception:
+                pass
+
+    def __call__(self, resource_id: str, timeout: int = 300, metadata: dict = None):
+        """
+        Redirect axon.lock("resource_id") to hold(resource_id)
+        """
+        return self.hold(resource_id, timeout=timeout, metadata=metadata)
+
+
+class SyncCoordinationClient(_BaseSyncClient):
+
+    def acquire(
+        self,
+        resource_id: str,
+        timeout: int = 300,
+        metadata: dict = None,
+    ) -> LockInfo:
+        """
+        Acquire an exclusive lock on a resource (sync).
+        """
+        body = {"resource_id": resource_id, "timeout": timeout}
+        if metadata is not None:
+            body["metadata"] = metadata
+
+        data = self._request("POST", "/v1/lock/acquire", json=body)
+        return LockInfo(
+            lock_id=data["lock_id"],
+            resource_id=data["resource_id"],
+            expires_at=data["expires_at"],
+        )
+
+    def release(self, resource_id: str) -> bool:
+        """
+        Release a lock that you hold (sync).
+        """
+        data = self._request(
+            "POST",
+            "/v1/lock/release",
+            params={"resource_id": resource_id},
+        )
+        return data.get("released", False)
+
+    def status(self, resource_id: str) -> LockStatus:
+        """
+        Check lock status of a resource (sync).
+        """
+        data = self._request("GET", f"/v1/lock/status/{resource_id}")
+        return LockStatus(
+            locked=data["locked"],
+            resource_id=data["resource_id"],
+            holder_agent_id=data.get("holder_agent_id"),
+            locked_at=data.get("locked_at"),
+            expires_at=data.get("expires_at"),
+        )
+
+    def list_active(self) -> list[dict]:
+        """
+        List all active locks in the project (sync).
+        """
+        data = self._request("GET", "/v1/lock/list")
+        return data.get("locks", [])
+
+    @contextmanager
+    def hold(
+        self,
+        resource_id: str,
+        timeout: int = 300,
+        metadata: dict = None,
+    ):
+        """
+        Context manager that acquires a lock and automatically releases it on exit (sync).
+        """
+        lock_info = self.acquire(resource_id, timeout=timeout, metadata=metadata)
+        try:
+            yield lock_info
+        finally:
+            try:
+                self.release(resource_id)
+            except Exception:
+                pass
+
+    def __call__(self, resource_id: str, timeout: int = 300, metadata: dict = None):
+        """
+        Redirect axon.lock("resource_id") to hold(resource_id)
+        """
+        return self.hold(resource_id, timeout=timeout, metadata=metadata)

axon/events.py

@@ -0,0 +1,116 @@
+import json
+import time
+import asyncio
+from typing import AsyncIterator, Iterator
+import websockets
+from websockets.sync.client import connect as ws_sync_connect
+from axon.exceptions import AxonConnectionError
+
+
+class EventsClient:
+    """
+    Connects to the Axon server via WebSocket and streams real-time events (async).
+    """
+
+    def __init__(self, base_url: str, api_key: str, project_id: str):
+        ws_base = base_url.replace("https://", "wss://").replace("http://", "ws://")
+        self._ws_url = ws_base.rstrip("/")
+        self._api_key = api_key
+        self._project_id = project_id
+
+    async def listen(self, reconnect: bool = True) -> AsyncIterator[dict]:
+        url = f"{self._ws_url}/v1/events/{self._project_id}"
+
+        while True:
+            try:
+                async with websockets.connect(
+                    url,
+                    additional_headers={"X-API-Key": self._api_key},
+                    ping_interval=30,
+                    ping_timeout=10,
+                ) as ws:
+                    async for raw_message in ws:
+                        try:
+                            event = json.loads(raw_message)
+
+                            # Skip heartbeat pings
+                            if event.get("type") == "ping":
+                                continue
+
+                            yield event
+
+                        except json.JSONDecodeError:
+                            continue
+
+            except websockets.exceptions.ConnectionClosed:
+                if not reconnect:
+                    return
+                await asyncio.sleep(2)
+
+            except Exception as e:
+                if not reconnect:
+                    raise AxonConnectionError(
+                        f"WebSocket connection failed: {str(e)}"
+                    )
+                await asyncio.sleep(2)
+
+    async def listen_once(self, timeout: float = 10.0) -> dict | None:
+        try:
+            async with asyncio.timeout(timeout):
+                async for event in self.listen(reconnect=False):
+                    return event
+        except (asyncio.TimeoutError, StopAsyncIteration):
+            return None
+
+
+class SyncEventsClient:
+    """
+    Connects to the Axon server via WebSocket and streams real-time events (sync).
+    """
+
+    def __init__(self, base_url: str, api_key: str, project_id: str):
+        ws_base = base_url.replace("https://", "wss://").replace("http://", "ws://")
+        self._ws_url = ws_base.rstrip("/")
+        self._api_key = api_key
+        self._project_id = project_id
+
+    def listen(self, reconnect: bool = True) -> Iterator[dict]:
+        url = f"{self._ws_url}/v1/events/{self._project_id}"
+
+        while True:
+            try:
+                with ws_sync_connect(
+                    url,
+                    additional_headers={"X-API-Key": self._api_key},
+                ) as ws:
+                    for raw_message in ws:
+                        try:
+                            event = json.loads(raw_message)
+
+                            # Skip heartbeat pings
+                            if event.get("type") == "ping":
+                                continue
+
+                            yield event
+
+                        except json.JSONDecodeError:
+                            continue
+
+            except Exception as e:
+                if not reconnect:
+                    raise AxonConnectionError(
+                        f"WebSocket connection failed: {str(e)}"
+                    )
+                time.sleep(2)
+
+    def listen_once(self, timeout: float = 10.0) -> dict | None:
+        iterator = self.listen(reconnect=False)
+        start_time = time.time()
+        while time.time() - start_time < timeout:
+            try:
+                return next(iterator)
+            except StopIteration:
+                return None
+            except Exception:
+                raise
+        return None

axon/exceptions.py

@@ -0,0 +1,51 @@
+class AxonError(Exception):
+    """Base exception for all Axon errors"""
+    def __init__(self, message: str, status_code: int = None):
+        super().__init__(message)
+        self.message = message
+        self.status_code = status_code
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(message={self.message!r}, status_code={self.status_code})"
+
+
+class AuthError(AxonError):
+    """Raised when authentication fails. Wrong or missing API key or token."""
+    pass
+
+
+class LockConflictError(AxonError):
+    """Raised when you try to acquire a lock that another agent already holds."""
+    def __init__(self, resource_id: str = "", detail: str = None):
+        msg = detail or f"Resource '{resource_id}' is already locked by another agent"
+        super().__init__(msg, 409)
+        self.resource_id = resource_id
+
+
+class NotFoundError(AxonError):
+    """Raised when the requested resource does not exist."""
+    pass
+
+
+class AxonPermissionError(AxonError):
+    """Raised when the agent does not have permission to perform the action."""
+    pass
+
+
+class RateLimitError(AxonError):
+    """Raised when the rate limit is exceeded."""
+    def __init__(self, retry_after: int = 60):
+        super().__init__(
+            f"Rate limit exceeded. Retry after {retry_after} seconds.", 429
+        )
+        self.retry_after = retry_after
+
+
+class ServerError(AxonError):
+    """Raised when the Axon server returns a 5xx error."""
+    pass
+
+
+class AxonConnectionError(AxonError):
+    """Raised when the SDK cannot connect to the Axon server at all."""
+    pass

axon/integrations/__init__.py

@@ -0,0 +1,15 @@
+from axon.integrations.langchain import AxonMemoryTool, AxonLockTool, AxonReceiptCallbackHandler
+from axon.integrations.crewai import (
+    AxonMemoryTool as CrewAxonMemoryTool,
+    AxonLockTool as CrewAxonLockTool,
+    AxonReceiptCallbackHandler as CrewAxonReceiptCallbackHandler,
+)
+
+__all__ = [
+    "AxonMemoryTool",
+    "AxonLockTool",
+    "AxonReceiptCallbackHandler",
+    "CrewAxonMemoryTool",
+    "CrewAxonLockTool",
+    "CrewAxonReceiptCallbackHandler",
+]