feat(sdk): add axons to OO SDK

Author: sid-rl

src/runloop_api_client/resources/axons.py

@@ -227,10 +227,12 @@ def subscribe_sse(
         """
         if not id:
             raise ValueError(f"Expected a non-empty value for `id` but received {id!r}")
+        default_headers: Headers = {"Accept": "text/event-stream"}
+        merged_headers = default_headers if extra_headers is None else {**default_headers, **extra_headers}
         return self._get(
             path_template("/v1/axons/{id}/subscribe/sse", id=id),
             options=make_request_options(
-                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+                extra_headers=merged_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
             ),
             cast_to=AxonEventView,
             stream=True,
@@ -437,10 +439,12 @@ async def subscribe_sse(
         """
         if not id:
             raise ValueError(f"Expected a non-empty value for `id` but received {id!r}")
+        default_headers: Headers = {"Accept": "text/event-stream"}
+        merged_headers = default_headers if extra_headers is None else {**default_headers, **extra_headers}
         return await self._get(
             path_template("/v1/axons/{id}/subscribe/sse", id=id),
             options=make_request_options(
-                extra_headers=extra_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
+                extra_headers=merged_headers, extra_query=extra_query, extra_body=extra_body, timeout=timeout
             ),
             cast_to=AxonEventView,
             stream=True,

src/runloop_api_client/sdk/__init__.py

@@ -5,7 +5,9 @@
 
 from __future__ import annotations
 
+from .axon import Axon
 from .sync import (
+    AxonOps,
     AgentOps,
     DevboxOps,
     ScorerOps,
@@ -23,6 +25,7 @@
 from .agent import Agent
 from ._types import ScenarioPreview
 from .async_ import (
+    AsyncAxonOps,
     AsyncAgentOps,
     AsyncDevboxOps,
     AsyncScorerOps,
@@ -45,6 +48,7 @@
 from .benchmark import Benchmark
 from .blueprint import Blueprint
 from .execution import Execution
+from .async_axon import AsyncAxon
 from .mcp_config import McpConfig
 from .async_agent import AsyncAgent
 from .async_devbox import AsyncDevbox, AsyncNamedShell
@@ -78,6 +82,8 @@
     # Management interfaces
     "AgentOps",
     "AsyncAgentOps",
+    "AxonOps",
+    "AsyncAxonOps",
     "BenchmarkOps",
     "AsyncBenchmarkOps",
     "DevboxOps",
@@ -103,6 +109,8 @@
     # Resource classes
     "Agent",
     "AsyncAgent",
+    "Axon",
+    "AsyncAxon",
     "AsyncSecret",
     "Benchmark",
     "AsyncBenchmark",

src/runloop_api_client/sdk/_types.py

@@ -5,9 +5,11 @@
     InputContext,
     ScenarioView,
     AgentListParams,
+    AxonCreateParams,
     DevboxListParams,
     ObjectListParams,
     AgentCreateParams,
+    AxonPublishParams,
     DevboxCreateParams,
     ObjectCreateParams,
     ScenarioListParams,
@@ -186,6 +188,14 @@ class SDKAgentListParams(AgentListParams, BaseRequestOptions):
     pass
 
 
+class SDKAxonCreateParams(AxonCreateParams, LongRequestOptions):
+    pass
+
+
+class SDKAxonPublishParams(AxonPublishParams, LongRequestOptions):
+    pass
+
+
 class SDKScenarioListParams(ScenarioListParams, BaseRequestOptions):
     pass
 

src/runloop_api_client/sdk/async_.py

@@ -14,6 +14,7 @@
     BaseRequestOptions,
     LongRequestOptions,
     SDKAgentListParams,
+    SDKAxonCreateParams,
     SDKDevboxListParams,
     SDKObjectListParams,
     SDKScorerListParams,
@@ -38,6 +39,7 @@
 from .._types import Timeout, NotGiven, not_given
 from .._client import DEFAULT_MAX_RETRIES, AsyncRunloop
 from ._helpers import detect_content_type
+from .async_axon import AsyncAxon
 from .async_agent import AsyncAgent
 from .async_devbox import AsyncDevbox
 from .async_scorer import AsyncScorer
@@ -521,6 +523,33 @@ async def upload_from_bytes(
         return obj
 
 
+class AsyncAxonOps:
+    """[Beta] Create and manage axons (async). Access via ``runloop.axon``.
+
+    Example:
+        >>> runloop = AsyncRunloopSDK()
+        >>> axon = await runloop.axon.create()
+        >>> await axon.publish(event_type="test", origin="USER_EVENT", payload="{}", source="sdk")
+    """
+
+    def __init__(self, client: AsyncRunloop) -> None:
+        self._client = client
+
+    async def create(self, **params: Unpack[SDKAxonCreateParams]) -> AsyncAxon:
+        """[Beta] Create a new axon."""
+        response = await self._client.axons.create(**params)
+        return AsyncAxon(self._client, response.id)
+
+    def from_id(self, axon_id: str) -> AsyncAxon:
+        """Get an AsyncAxon instance for an existing axon ID."""
+        return AsyncAxon(self._client, axon_id)
+
+    async def list(self, **options: Unpack[BaseRequestOptions]) -> list[AsyncAxon]:
+        """[Beta] List all active axons."""
+        result = await self._client.axons.list(**options)
+        return [AsyncAxon(self._client, axon.id) for axon in result.axons]
+
+
 class AsyncScorerOps:
     """Create and manage custom scorers (async). Access via ``runloop.scorer``.
 
@@ -1205,6 +1234,8 @@ class AsyncRunloopSDK:
     :vartype api: AsyncRunloop
     :ivar agent: High-level async interface for agent management.
     :vartype agent: AsyncAgentOps
+    :ivar axon: [Beta] High-level async interface for axon management
+    :vartype axon: AsyncAxonOps
     :ivar benchmark: High-level async interface for benchmark management
     :vartype benchmark: AsyncBenchmarkOps
     :ivar devbox: High-level async interface for devbox management
@@ -1238,6 +1269,7 @@ class AsyncRunloopSDK:
 
     api: AsyncRunloop
     agent: AsyncAgentOps
+    axon: AsyncAxonOps
     benchmark: AsyncBenchmarkOps
     devbox: AsyncDevboxOps
     blueprint: AsyncBlueprintOps
@@ -1289,6 +1321,7 @@ def __init__(
         )
 
         self.agent = AsyncAgentOps(self.api)
+        self.axon = AsyncAxonOps(self.api)
         self.benchmark = AsyncBenchmarkOps(self.api)
         self.devbox = AsyncDevboxOps(self.api)
         self.blueprint = AsyncBlueprintOps(self.api)

src/runloop_api_client/sdk/async_axon.py

@@ -0,0 +1,53 @@
+"""Axon resource class for asynchronous operations."""
+
+from __future__ import annotations
+
+from typing_extensions import Unpack, override
+
+from ._types import (
+    BaseRequestOptions,
+    SDKAxonPublishParams,
+)
+from .._client import AsyncRunloop
+from .._streaming import AsyncStream
+from ..types.axon_view import AxonView
+from ..types.axon_event_view import AxonEventView
+from ..types.publish_result_view import PublishResultView
+
+
+class AsyncAxon:
+    """[Beta] Wrapper around asynchronous axon operations.
+
+    Axons are event communication channels that support publishing events
+    and subscribing to event streams via server-sent events (SSE).
+    Obtain instances via ``runloop.axon.create()`` or ``runloop.axon.from_id()``.
+
+    Example:
+        >>> runloop = AsyncRunloopSDK()
+        >>> axon = await runloop.axon.create()
+        >>> await axon.publish(event_type="task_done", origin="AGENT_EVENT", payload="{}", source="my-agent")
+    """
+
+    def __init__(self, client: AsyncRunloop, axon_id: str) -> None:
+        self._client = client
+        self._id = axon_id
+
+    @override
+    def __repr__(self) -> str:
+        return f"<AsyncAxon id={self._id!r}>"
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    async def get_info(self, **options: Unpack[BaseRequestOptions]) -> AxonView:
+        """[Beta] Retrieve the latest axon information."""
+        return await self._client.axons.retrieve(self._id, **options)
+
+    async def publish(self, **params: Unpack[SDKAxonPublishParams]) -> PublishResultView:
+        """[Beta] Publish an event to this axon."""
+        return await self._client.axons.publish(self._id, **params)
+
+    async def subscribe_sse(self, **options: Unpack[BaseRequestOptions]) -> AsyncStream[AxonEventView]:
+        """[Beta] Subscribe to this axon's event stream via SSE."""
+        return await self._client.axons.subscribe_sse(self._id, **options)

src/runloop_api_client/sdk/axon.py

@@ -0,0 +1,53 @@
+"""Axon resource class for synchronous operations."""
+
+from __future__ import annotations
+
+from typing_extensions import Unpack, override
+
+from ._types import (
+    BaseRequestOptions,
+    SDKAxonPublishParams,
+)
+from .._client import Runloop
+from .._streaming import Stream
+from ..types.axon_view import AxonView
+from ..types.axon_event_view import AxonEventView
+from ..types.publish_result_view import PublishResultView
+
+
+class Axon:
+    """[Beta] Wrapper around synchronous axon operations.
+
+    Axons are event communication channels that support publishing events
+    and subscribing to event streams via server-sent events (SSE).
+    Obtain instances via ``runloop.axon.create()`` or ``runloop.axon.from_id()``.
+
+    Example:
+        >>> runloop = RunloopSDK()
+        >>> axon = runloop.axon.create()
+        >>> axon.publish(event_type="task_done", origin="AGENT_EVENT", payload="{}", source="my-agent")
+    """
+
+    def __init__(self, client: Runloop, axon_id: str) -> None:
+        self._client = client
+        self._id = axon_id
+
+    @override
+    def __repr__(self) -> str:
+        return f"<Axon id={self._id!r}>"
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    def get_info(self, **options: Unpack[BaseRequestOptions]) -> AxonView:
+        """[Beta] Retrieve the latest axon information."""
+        return self._client.axons.retrieve(self._id, **options)
+
+    def publish(self, **params: Unpack[SDKAxonPublishParams]) -> PublishResultView:
+        """[Beta] Publish an event to this axon."""
+        return self._client.axons.publish(self._id, **params)
+
+    def subscribe_sse(self, **options: Unpack[BaseRequestOptions]) -> Stream[AxonEventView]:
+        """[Beta] Subscribe to this axon's event stream via SSE."""
+        return self._client.axons.subscribe_sse(self._id, **options)

src/runloop_api_client/sdk/sync.py

@@ -9,11 +9,13 @@
 
 import httpx
 
+from .axon import Axon
 from .agent import Agent
 from ._types import (
     BaseRequestOptions,
     LongRequestOptions,
     SDKAgentListParams,
+    SDKAxonCreateParams,
     SDKDevboxListParams,
     SDKObjectListParams,
     SDKScorerListParams,
@@ -516,6 +518,33 @@ def upload_from_bytes(
         return obj
 
 
+class AxonOps:
+    """[Beta] Create and manage axons. Access via ``runloop.axon``.
+
+    Example:
+        >>> runloop = RunloopSDK()
+        >>> axon = runloop.axon.create()
+        >>> axon.publish(event_type="test", origin="USER_EVENT", payload="{}", source="sdk")
+    """
+
+    def __init__(self, client: Runloop) -> None:
+        self._client = client
+
+    def create(self, **params: Unpack[SDKAxonCreateParams]) -> Axon:
+        """[Beta] Create a new axon."""
+        response = self._client.axons.create(**params)
+        return Axon(self._client, response.id)
+
+    def from_id(self, axon_id: str) -> Axon:
+        """Get an Axon instance for an existing axon ID."""
+        return Axon(self._client, axon_id)
+
+    def list(self, **options: Unpack[BaseRequestOptions]) -> list[Axon]:
+        """[Beta] List all active axons."""
+        result = self._client.axons.list(**options)
+        return [Axon(self._client, axon.id) for axon in result.axons]
+
+
 class ScorerOps:
     """Create and manage custom scorers. Access via ``runloop.scorer``.
 
@@ -1230,6 +1259,8 @@ class RunloopSDK:
     :vartype api: Runloop
     :ivar agent: High-level interface for agent management.
     :vartype agent: AgentOps
+    :ivar axon: [Beta] High-level interface for axon management
+    :vartype axon: AxonOps
     :ivar benchmark: High-level interface for benchmark management
     :vartype benchmark: BenchmarkOps
     :ivar devbox: High-level interface for devbox management
@@ -1263,6 +1294,7 @@ class RunloopSDK:
 
     api: Runloop
     agent: AgentOps
+    axon: AxonOps
     benchmark: BenchmarkOps
     devbox: DevboxOps
     blueprint: BlueprintOps
@@ -1314,6 +1346,7 @@ def __init__(
         )
 
         self.agent = AgentOps(self.api)
+        self.axon = AxonOps(self.api)
         self.benchmark = BenchmarkOps(self.api)
         self.devbox = DevboxOps(self.api)
         self.blueprint = BlueprintOps(self.api)

tests/sdk/conftest.py

@@ -23,6 +23,7 @@
     "object": "obj_123",
     "scorer": "sco_123",
     "agent": "agt_123",
+    "axon": "axn_123",
     "scenario": "scn_123",
     "scenario_run": "scr_123",
     "benchmark": "bmd_123",
@@ -115,6 +116,23 @@ class MockAgentView:
     source: Any = None
 
 
+@dataclass
+class MockAxonView:
+    """Mock AxonView for testing."""
+
+    id: str = TEST_IDS["axon"]
+    created_at_ms: int = 1234567890000
+    name: str = "test-axon"
+
+
+@dataclass
+class MockPublishResultView:
+    """Mock PublishResultView for testing."""
+
+    sequence: int = 1
+    timestamp_ms: int = 1234567890000
+
+
 @dataclass
 class MockScenarioView:
     """Mock ScenarioView for testing."""
@@ -307,6 +325,12 @@ def agent_view() -> MockAgentView:
     return MockAgentView()
 
 
+@pytest.fixture
+def axon_view() -> MockAxonView:
+    """Create a mock AxonView."""
+    return MockAxonView()
+
+
 @pytest.fixture
 def scenario_view() -> MockScenarioView:
     """Create a mock ScenarioView."""