aicc-py 0.1.0: initial SDK
- protocol.py: pydantic models for all AICC message types (envelope, session_init, tool_call, tool_result, event, error, heartbeat) - bridge.py: Bridge with @bridge.tool registration, capability checks, session management, serve_forever - client.py: AICCClient with single background reader (safe on concurrent transports like WebSocket), call_tool, events, manifest - tool.py: @tool decorator with schema generation from type hints - schema.py: JSON Schema generation (str/int/float/bool, list, dict, Optional, pydantic models) - transport: Transport protocol, InProcessTransport, WebSocket client+server - tests: 12 passing (integration, schema, websocket roundtrip) - examples: bridge_minimal.py + agent_minimal.py (verified end-to-end)
This commit is contained in:
Emil Shanaty
@@ -0,0 +1,45 @@
|
||||
"""aicc-py: Python SDK for the AI-Controlled Character Protocol."""
|
||||
|
||||
from aicc.protocol import (
|
||||
ErrorCode,
|
||||
MessageType,
|
||||
TickMode,
|
||||
WorldKind,
|
||||
ModelClass,
|
||||
ToolClass,
|
||||
SessionInit,
|
||||
SessionClose,
|
||||
ToolCall,
|
||||
ToolResult,
|
||||
EventMessage,
|
||||
ErrorMessage,
|
||||
Heartbeat,
|
||||
)
|
||||
from aicc.errors import AICCError, ToolError, CapabilityError, ProtocolError
|
||||
from aicc.bridge import Bridge
|
||||
from aicc.client import AICCClient
|
||||
from aicc.tool import tool
|
||||
|
||||
__version__ = "0.1.0"
|
||||
__all__ = [
|
||||
"Bridge",
|
||||
"AICCClient",
|
||||
"tool",
|
||||
"ErrorCode",
|
||||
"MessageType",
|
||||
"TickMode",
|
||||
"WorldKind",
|
||||
"ModelClass",
|
||||
"ToolClass",
|
||||
"SessionInit",
|
||||
"SessionClose",
|
||||
"ToolCall",
|
||||
"ToolResult",
|
||||
"EventMessage",
|
||||
"ErrorMessage",
|
||||
"Heartbeat",
|
||||
"AICCError",
|
||||
"ToolError",
|
||||
"CapabilityError",
|
||||
"ProtocolError",
|
||||
]
|
||||
+267
@@ -0,0 +1,267 @@
|
||||
"""Bridge: environment side of the AICC protocol."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable
|
||||
|
||||
from aicc.protocol import (
|
||||
Capabilities,
|
||||
ErrorCode,
|
||||
ErrorPayload,
|
||||
Heartbeat,
|
||||
MessageType,
|
||||
ModelClass,
|
||||
ModelDecl,
|
||||
SessionClose,
|
||||
SessionInit,
|
||||
TickMode,
|
||||
ToolCall,
|
||||
ToolResult,
|
||||
ToolSpec,
|
||||
WorldKind,
|
||||
WorldSpec,
|
||||
PROTOCOL_VERSION,
|
||||
)
|
||||
from aicc.tool import ToolDefinition, call_tool_impl
|
||||
|
||||
|
||||
@dataclass
|
||||
class Session:
|
||||
"""A live bridge session with one connected agent."""
|
||||
|
||||
session_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
manifest: SessionInit | None = None
|
||||
send_queue: asyncio.Queue[dict[str, Any]] = field(default_factory=asyncio.Queue)
|
||||
closed: bool = False
|
||||
|
||||
|
||||
class Bridge:
|
||||
"""Environment-side bridge.
|
||||
|
||||
Holds the world state, registered tools, and per-session state.
|
||||
Bridges run inside transports (in-process, websocket, etc.) and do not
|
||||
perform I/O directly.
|
||||
|
||||
Example:
|
||||
bridge = Bridge(name="capsule", kind=WorldKind.THREE_D)
|
||||
|
||||
@bridge.tool(description="Get current position.")
|
||||
async def proprioception() -> dict:
|
||||
return {"position": {"x": 0, "y": 0, "z": 0}, ...}
|
||||
|
||||
bridge.register_tool(...)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
kind: WorldKind = WorldKind.THREE_D,
|
||||
tick_rate_hz: float = 10.0,
|
||||
tick_mode: TickMode = TickMode.FIXED,
|
||||
agent_model: ModelClass = ModelClass.EDGE_MEDIUM,
|
||||
expected_first_token_ms: int | None = None,
|
||||
expected_full_response_ms: int | None = None,
|
||||
):
|
||||
self.name = name
|
||||
self.kind = kind
|
||||
self.tick_rate_hz = tick_rate_hz
|
||||
self.tick_mode = tick_mode
|
||||
self.agent_model = agent_model
|
||||
self.expected_first_token_ms = expected_first_token_ms
|
||||
self.expected_full_response_ms = expected_full_response_ms
|
||||
|
||||
self._tools: dict[str, ToolDefinition] = {}
|
||||
self._sessions: dict[str, Session] = {}
|
||||
|
||||
# ---------- Tool registration ----------
|
||||
|
||||
def tool(
|
||||
self,
|
||||
_fn: Callable[..., Any] | None = None,
|
||||
*,
|
||||
id: str | None = None,
|
||||
cls=None,
|
||||
description: str | None = None,
|
||||
requires_capability: str | None = None,
|
||||
limits: dict[str, Any] | None = None,
|
||||
strict_input: bool = False,
|
||||
):
|
||||
"""Register a tool. Use as @bridge.tool or @bridge.tool(cls=..., ...)."""
|
||||
from aicc.tool import tool as _tool
|
||||
|
||||
if _fn is not None and callable(_fn):
|
||||
defn = _tool(_fn)
|
||||
self.register_tool(defn)
|
||||
return _fn
|
||||
|
||||
def wrap(fn: Callable[..., Any]) -> Callable[..., Any]:
|
||||
defn = _tool(
|
||||
fn,
|
||||
id=id,
|
||||
cls=cls,
|
||||
description=description,
|
||||
requires_capability=requires_capability,
|
||||
limits=limits,
|
||||
strict_input=strict_input,
|
||||
)
|
||||
self.register_tool(defn)
|
||||
return fn
|
||||
|
||||
return wrap
|
||||
|
||||
def register_tool(self, definition: ToolDefinition) -> None:
|
||||
if definition.id in self._tools:
|
||||
raise ValueError(f"Tool id '{definition.id}' is already registered.")
|
||||
self._tools[definition.id] = definition
|
||||
|
||||
def tools(self) -> list[ToolDefinition]:
|
||||
return list(self._tools.values())
|
||||
|
||||
# ---------- Capability computation ----------
|
||||
|
||||
def capabilities(self) -> Capabilities:
|
||||
return Capabilities(
|
||||
sensors=[t.id for t in self._tools.values() if t.cls.value == "sensor"],
|
||||
actuators=[t.id for t in self._tools.values() if t.cls.value == "actuator"],
|
||||
generators=[t.id for t in self._tools.values() if t.cls.value == "generator"],
|
||||
)
|
||||
|
||||
# ---------- Session handling ----------
|
||||
|
||||
def open_session(self) -> Session:
|
||||
session = Session()
|
||||
manifest = self._build_manifest(session.session_id)
|
||||
session.manifest = manifest
|
||||
self._sessions[session.session_id] = session
|
||||
return session
|
||||
|
||||
def close_session(self, session_id: str) -> None:
|
||||
sess = self._sessions.pop(session_id, None)
|
||||
if sess is not None:
|
||||
sess.closed = True
|
||||
|
||||
def session(self, session_id: str) -> Session:
|
||||
sess = self._sessions.get(session_id)
|
||||
if sess is None or sess.closed:
|
||||
raise KeyError(session_id)
|
||||
return sess
|
||||
|
||||
async def serve_forever(self) -> None:
|
||||
"""Run the bridge event loop until cancelled.
|
||||
|
||||
Useful with a server transport (e.g. WebSocketServer) already running:
|
||||
async with WebSocketServer(bridge, port=8765):
|
||||
await bridge.serve_forever()
|
||||
"""
|
||||
while True:
|
||||
await asyncio.sleep(3600)
|
||||
|
||||
def _build_manifest(self, session_id: str) -> SessionInit:
|
||||
return SessionInit(
|
||||
session_id=session_id,
|
||||
tick_rate_hz=self.tick_rate_hz,
|
||||
tick_mode=self.tick_mode,
|
||||
world=WorldSpec(name=self.name, kind=self.kind),
|
||||
agent_model=ModelDecl(
|
||||
class_=self.agent_model,
|
||||
expected_first_token_ms=self.expected_first_token_ms,
|
||||
expected_full_response_ms=self.expected_full_response_ms,
|
||||
),
|
||||
capabilities=self.capabilities(),
|
||||
tools=[
|
||||
ToolSpec(
|
||||
id=t.id,
|
||||
class_=t.cls,
|
||||
description=t.description,
|
||||
input_schema=t.input_schema,
|
||||
output_schema=t.output_schema,
|
||||
requires_capability=t.requires_capability,
|
||||
limits=t.limits,
|
||||
strict_input=t.strict_input,
|
||||
)
|
||||
for t in self._tools.values()
|
||||
],
|
||||
)
|
||||
|
||||
# ---------- Message handling ----------
|
||||
|
||||
async def handle_message(self, session_id: str, raw: dict[str, Any]) -> dict[str, Any] | None:
|
||||
"""Handle a single incoming message from an agent.
|
||||
|
||||
Returns the response message (or None for fire-and-forget events/heartbeats).
|
||||
May also push async events to the session queue.
|
||||
"""
|
||||
msg_type = raw.get("type")
|
||||
if msg_type == MessageType.MANIFEST_REQUEST.value:
|
||||
sess = self.session(session_id)
|
||||
return sess.manifest.model_dump(by_alias=True) if sess.manifest else None
|
||||
if msg_type == MessageType.TOOL_CALL.value:
|
||||
return await self._handle_tool_call(session_id, ToolCall.model_validate(raw))
|
||||
if msg_type == MessageType.SESSION_CLOSE.value:
|
||||
self.close_session(session_id)
|
||||
return SessionClose(
|
||||
session_id=session_id,
|
||||
reason=(raw.get("reason") or "client_close"),
|
||||
).model_dump(by_alias=True)
|
||||
if msg_type == MessageType.HEARTBEAT.value:
|
||||
return Heartbeat(session_id=session_id).model_dump(by_alias=True)
|
||||
if msg_type == MessageType.SESSION_RESUME.value:
|
||||
sess = self.session(raw.get("session_id") or session_id)
|
||||
return sess.manifest.model_dump(by_alias=True) if sess.manifest else None
|
||||
return None
|
||||
|
||||
async def _handle_tool_call(self, session_id: str, call: ToolCall) -> dict[str, Any]:
|
||||
defn = self._tools.get(call.tool)
|
||||
if defn is None:
|
||||
return ToolResult(
|
||||
session_id=session_id,
|
||||
call_id=call.call_id,
|
||||
ok=False,
|
||||
error=ErrorPayload(
|
||||
code=ErrorCode.TOOL_UNKNOWN,
|
||||
message=f"Tool '{call.tool}' is not registered.",
|
||||
retryable=False,
|
||||
),
|
||||
).model_dump(by_alias=True)
|
||||
|
||||
if defn.requires_capability and defn.requires_capability not in self._sessions[session_id].manifest.capabilities.sensors + self._sessions[session_id].manifest.capabilities.actuators + self._sessions[session_id].manifest.capabilities.generators: # type: ignore[union-attr]
|
||||
return ToolResult(
|
||||
session_id=session_id,
|
||||
call_id=call.call_id,
|
||||
ok=False,
|
||||
error=ErrorPayload(
|
||||
code=ErrorCode.TOOL_UNAVAILABLE,
|
||||
message=f"Tool '{call.tool}' requires capability '{defn.requires_capability}'.",
|
||||
retryable=False,
|
||||
),
|
||||
).model_dump(by_alias=True)
|
||||
|
||||
try:
|
||||
output = await call_tool_impl(defn, call.input)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return ToolResult(
|
||||
session_id=session_id,
|
||||
call_id=call.call_id,
|
||||
ok=False,
|
||||
error=ErrorPayload(
|
||||
code=ErrorCode.EXECUTION_FAILED,
|
||||
message=f"{type(exc).__name__}: {exc}",
|
||||
retryable=False,
|
||||
),
|
||||
).model_dump(by_alias=True)
|
||||
|
||||
if not isinstance(output, dict):
|
||||
output = {"value": output}
|
||||
|
||||
return ToolResult(
|
||||
session_id=session_id,
|
||||
call_id=call.call_id,
|
||||
ok=True,
|
||||
output=output,
|
||||
).model_dump(by_alias=True)
|
||||
+184
@@ -0,0 +1,184 @@
|
||||
"""AICCClient: agent side of the AICC protocol."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from aicc.protocol import (
|
||||
PROTOCOL_VERSION,
|
||||
ErrorCode,
|
||||
EventMessage,
|
||||
MessageType,
|
||||
SessionClose,
|
||||
SessionInit,
|
||||
ToolCall,
|
||||
ToolResult,
|
||||
)
|
||||
from aicc.errors import (
|
||||
AICCError,
|
||||
ConnectionError_,
|
||||
ProtocolError,
|
||||
ToolError,
|
||||
)
|
||||
from aicc.transport.base import Transport
|
||||
|
||||
|
||||
class AICCClient:
|
||||
"""Agent-side AICC client.
|
||||
|
||||
Wraps a Transport and exposes high-level operations:
|
||||
- handshake() -> manifest
|
||||
- call_tool(name, input) -> ToolResult
|
||||
- next_event() -> EventMessage
|
||||
- close
|
||||
|
||||
A single background reader task consumes the transport and routes
|
||||
messages into an internal queue, so any number of concurrent callers
|
||||
(tool calls, event consumers) are safe on any transport, including
|
||||
WebSocket where concurrent recv() is forbidden.
|
||||
|
||||
Example:
|
||||
client = AICCClient(WebSocketClientTransport("ws://localhost:8765"))
|
||||
async with client:
|
||||
manifest = await client.handshake()
|
||||
result = await client.call_tool("vision", {})
|
||||
"""
|
||||
|
||||
def __init__(self, transport: Transport, protocol: str = PROTOCOL_VERSION):
|
||||
self._transport = transport
|
||||
self._protocol = protocol
|
||||
self._session_id: str | None = None
|
||||
self._manifest: SessionInit | None = None
|
||||
self._inbox: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
|
||||
self._reader_task: asyncio.Task | None = None
|
||||
|
||||
# ---------- Lifecycle ----------
|
||||
|
||||
async def __aenter__(self) -> "AICCClient":
|
||||
await self._transport.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb) -> None:
|
||||
await self.close()
|
||||
|
||||
async def connect(self) -> "AICCClient":
|
||||
await self._transport.connect()
|
||||
return self
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._reader_task is not None:
|
||||
self._reader_task.cancel()
|
||||
try:
|
||||
await self._reader_task
|
||||
except (asyncio.CancelledError, AICCError):
|
||||
pass
|
||||
self._reader_task = None
|
||||
if self._session_id is not None:
|
||||
try:
|
||||
await self._transport.send(
|
||||
SessionClose(
|
||||
session_id=self._session_id,
|
||||
reason="client_close",
|
||||
).model_dump(by_alias=True)
|
||||
)
|
||||
except AICCError:
|
||||
pass
|
||||
self._session_id = None
|
||||
await self._transport.close()
|
||||
|
||||
# ---------- Manifest ----------
|
||||
|
||||
async def handshake(self) -> SessionInit:
|
||||
"""Wait for the bridge to push the initial session_init manifest."""
|
||||
self._spawn_reader()
|
||||
init = await self._receive_until(SessionInit)
|
||||
self._session_id = init.session_id
|
||||
self._manifest = init
|
||||
return init
|
||||
|
||||
async def fetch_manifest(self) -> SessionInit:
|
||||
if self._session_id is None:
|
||||
raise ProtocolError("handshake() must be called before fetch_manifest().")
|
||||
await self._transport.send(
|
||||
{
|
||||
"type": MessageType.MANIFEST_REQUEST.value,
|
||||
"session_id": self._session_id,
|
||||
"message_id": str(uuid.uuid4()),
|
||||
}
|
||||
)
|
||||
return await self._receive_until(SessionInit)
|
||||
|
||||
# ---------- Tool calls ----------
|
||||
|
||||
async def call_tool(self, tool_id: str, input: dict[str, Any]) -> ToolResult:
|
||||
if self._session_id is None:
|
||||
raise ProtocolError("handshake() must be called before call_tool().")
|
||||
call_id = f"tc_{uuid.uuid4().hex[:12]}"
|
||||
await self._transport.send(
|
||||
ToolCall(
|
||||
session_id=self._session_id,
|
||||
call_id=call_id,
|
||||
tool=tool_id,
|
||||
input=input,
|
||||
).model_dump(by_alias=True)
|
||||
)
|
||||
|
||||
while True:
|
||||
msg = await self._inbox.get()
|
||||
self._validate_protocol(msg)
|
||||
t = msg.get("type")
|
||||
if t == MessageType.ERROR.value:
|
||||
raise ProtocolError(f"bridge error: {msg.get('error')}")
|
||||
if t != MessageType.TOOL_RESULT.value:
|
||||
continue
|
||||
result = ToolResult.model_validate(msg)
|
||||
if result.call_id != call_id:
|
||||
continue
|
||||
if not result.ok and result.error:
|
||||
raise ToolError(
|
||||
code=result.error.code,
|
||||
message=result.error.message,
|
||||
retryable=result.error.retryable,
|
||||
)
|
||||
return result
|
||||
|
||||
# ---------- Events ----------
|
||||
|
||||
async def next_event(self, timeout: float | None = None) -> EventMessage:
|
||||
if timeout is None:
|
||||
raw = await self._inbox.get()
|
||||
else:
|
||||
raw = await asyncio.wait_for(self._inbox.get(), timeout=timeout)
|
||||
return EventMessage.model_validate(raw)
|
||||
|
||||
def _spawn_reader(self) -> None:
|
||||
async def read_loop() -> None:
|
||||
try:
|
||||
while True:
|
||||
msg = await self._transport.receive()
|
||||
await self._inbox.put(msg)
|
||||
except (ConnectionError_, asyncio.CancelledError):
|
||||
return
|
||||
|
||||
self._reader_task = asyncio.create_task(read_loop())
|
||||
|
||||
# ---------- Internals ----------
|
||||
|
||||
async def _receive_until(self, expect_cls: type) -> SessionInit:
|
||||
while True:
|
||||
msg = await self._inbox.get()
|
||||
self._validate_protocol(msg)
|
||||
t = msg.get("type")
|
||||
if t == MessageType.ERROR.value:
|
||||
raise ProtocolError(f"bridge error: {msg.get('error')}")
|
||||
if t == expect_cls.model_fields["type"].default:
|
||||
return expect_cls.model_validate(msg)
|
||||
|
||||
def _validate_protocol(self, msg: dict[str, Any]) -> None:
|
||||
proto = msg.get("protocol")
|
||||
if proto != self._protocol:
|
||||
raise ProtocolError(
|
||||
f"Protocol mismatch: client='{self._protocol}', message='{proto}'"
|
||||
)
|
||||
@@ -0,0 +1,35 @@
|
||||
"""AICC exception hierarchy."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from aicc.protocol import ErrorCode
|
||||
|
||||
|
||||
class AICCError(Exception):
|
||||
"""Base for all AICC SDK errors."""
|
||||
|
||||
|
||||
class ProtocolError(AICCError):
|
||||
"""Protocol-level error (version mismatch, malformed message)."""
|
||||
|
||||
|
||||
class ConnectionError_(AICCError):
|
||||
"""Transport / session connection failure."""
|
||||
|
||||
|
||||
class SessionExpiredError(AICCError):
|
||||
"""Session id is unknown or closed."""
|
||||
|
||||
|
||||
class ToolError(AICCError):
|
||||
"""Tool execution failure surfaced by the bridge."""
|
||||
|
||||
def __init__(self, code: ErrorCode, message: str, retryable: bool):
|
||||
super().__init__(message)
|
||||
self.code = code
|
||||
self.message = message
|
||||
self.retryable = retryable
|
||||
|
||||
|
||||
class CapabilityError(AICCError):
|
||||
"""Tool requires a capability the agent does not hold."""
|
||||
@@ -0,0 +1,193 @@
|
||||
"""AICC protocol message types and enums."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from enum import Enum
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
|
||||
PROTOCOL_VERSION = "aicc/0.1"
|
||||
|
||||
|
||||
class MessageType(str, Enum):
|
||||
SESSION_INIT = "session_init"
|
||||
SESSION_RESUME = "session_resume"
|
||||
SESSION_CLOSE = "session_close"
|
||||
MANIFEST_REQUEST = "manifest_request"
|
||||
TOOL_CALL = "tool_call"
|
||||
TOOL_RESULT = "tool_result"
|
||||
EVENT = "event"
|
||||
ERROR = "error"
|
||||
HEARTBEAT = "heartbeat"
|
||||
|
||||
|
||||
class TickMode(str, Enum):
|
||||
FIXED = "fixed"
|
||||
EVENT = "event"
|
||||
HYBRID = "hybrid"
|
||||
|
||||
|
||||
class WorldKind(str, Enum):
|
||||
TWO_D = "2d"
|
||||
THREE_D = "3d"
|
||||
TEXT = "text"
|
||||
ABSTRACT = "abstract"
|
||||
|
||||
|
||||
class ModelClass(str, Enum):
|
||||
EDGE_SMALL = "edge_small"
|
||||
EDGE_MEDIUM = "edge_medium"
|
||||
CLOUD_MEDIUM = "cloud_medium"
|
||||
CLOUD_LARGE = "cloud_large"
|
||||
|
||||
|
||||
class ToolClass(str, Enum):
|
||||
SENSOR = "sensor"
|
||||
ACTUATOR = "actuator"
|
||||
GENERATOR = "generator"
|
||||
|
||||
|
||||
class ErrorCode(str, Enum):
|
||||
PROTOCOL_MISMATCH = "protocol_mismatch"
|
||||
SESSION_EXPIRED = "session_expired"
|
||||
TOOL_UNKNOWN = "tool_unknown"
|
||||
TOOL_UNAVAILABLE = "tool_unavailable"
|
||||
INVALID_INPUT = "invalid_input"
|
||||
EXECUTION_FAILED = "execution_failed"
|
||||
TIMEOUT = "timeout"
|
||||
INTERNAL_ERROR = "internal_error"
|
||||
|
||||
|
||||
class EventTopic(str, Enum):
|
||||
TICK = "tick"
|
||||
COLLISION = "collision"
|
||||
AUDIO = "audio"
|
||||
STATE_CHANGE = "state_change"
|
||||
AGENT_MESSAGE = "agent_message"
|
||||
|
||||
|
||||
def _uuid() -> str:
|
||||
return str(uuid.uuid4())
|
||||
|
||||
|
||||
class Envelope(BaseModel):
|
||||
"""Common envelope for all AICC messages."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
protocol: Literal["aicc/0.1"] = PROTOCOL_VERSION
|
||||
session_id: str = Field(default_factory=_uuid)
|
||||
message_id: str = Field(default_factory=_uuid)
|
||||
|
||||
|
||||
class WorldSpec(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
name: str
|
||||
kind: WorldKind
|
||||
|
||||
|
||||
class ModelDecl(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", populate_by_name=True)
|
||||
|
||||
class_: ModelClass = Field(alias="class")
|
||||
expected_first_token_ms: int | None = None
|
||||
expected_full_response_ms: int | None = None
|
||||
|
||||
|
||||
class ToolLimits(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
calls_per_minute: int | None = None
|
||||
calls_per_session: int | None = None
|
||||
|
||||
|
||||
class ToolSpec(BaseModel):
|
||||
"""Tool definition advertised in the session manifest."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid", populate_by_name=True)
|
||||
|
||||
id: str = Field(pattern=r"^[a-z][a-z0-9_]*$")
|
||||
class_: ToolClass = Field(alias="class")
|
||||
description: str
|
||||
input_schema: dict[str, Any]
|
||||
output_schema: dict[str, Any]
|
||||
requires_capability: str | None = None
|
||||
limits: ToolLimits | None = None
|
||||
strict_input: bool = False
|
||||
|
||||
|
||||
class Capabilities(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
sensors: list[str] = Field(default_factory=list)
|
||||
actuators: list[str] = Field(default_factory=list)
|
||||
generators: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class SessionInit(Envelope):
|
||||
type: Literal[MessageType.SESSION_INIT] = MessageType.SESSION_INIT
|
||||
tick_rate_hz: float = Field(gt=0)
|
||||
tick_mode: TickMode
|
||||
world: WorldSpec
|
||||
agent_model: ModelDecl
|
||||
capabilities: Capabilities
|
||||
tools: list[ToolSpec]
|
||||
event_subscriptions: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
class SessionClose(Envelope):
|
||||
type: Literal[MessageType.SESSION_CLOSE] = MessageType.SESSION_CLOSE
|
||||
reason: str | None = None
|
||||
|
||||
|
||||
class ToolCall(Envelope):
|
||||
type: Literal[MessageType.TOOL_CALL] = MessageType.TOOL_CALL
|
||||
call_id: str
|
||||
tool: str = Field(pattern=r"^[a-z][a-z0-9_]*$")
|
||||
input: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class Meta(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
tick: int | None = None
|
||||
latency_ms: int | None = None
|
||||
source: str | None = None
|
||||
|
||||
|
||||
class ErrorPayload(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid")
|
||||
|
||||
code: ErrorCode
|
||||
message: str
|
||||
retryable: bool
|
||||
|
||||
|
||||
class ToolResult(Envelope):
|
||||
type: Literal[MessageType.TOOL_RESULT] = MessageType.TOOL_RESULT
|
||||
call_id: str
|
||||
ok: bool
|
||||
output: dict[str, Any] | None = None
|
||||
error: ErrorPayload | None = None
|
||||
meta: Meta | None = None
|
||||
|
||||
|
||||
class EventMessage(Envelope):
|
||||
type: Literal[MessageType.EVENT] = MessageType.EVENT
|
||||
topic: EventTopic | str
|
||||
payload: dict[str, Any]
|
||||
meta: Meta | None = None
|
||||
|
||||
|
||||
class ErrorMessage(Envelope):
|
||||
type: Literal[MessageType.ERROR] = MessageType.ERROR
|
||||
error: ErrorPayload
|
||||
|
||||
|
||||
class Heartbeat(Envelope):
|
||||
type: Literal[MessageType.HEARTBEAT] = MessageType.HEARTBEAT
|
||||
manifest_update: dict[str, Any] | None = None
|
||||
+113
@@ -0,0 +1,113 @@
|
||||
"""JSON Schema generation utilities."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from typing import Any, Callable, get_args, get_origin, get_type_hints
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
def function_schema(fn: Callable[..., Any]) -> dict[str, Any]:
|
||||
"""Generate a JSON Schema for a function's input from its signature.
|
||||
|
||||
Supports:
|
||||
- Plain Python type hints (str, int, float, bool, list, dict, etc.)
|
||||
- Pydantic BaseModel subclasses as root input
|
||||
- Optional[X] / Union[X, None]
|
||||
- Default values from the signature
|
||||
|
||||
Returns a JSON Schema dict with type=object, properties, and required.
|
||||
"""
|
||||
sig = inspect.signature(fn)
|
||||
hints = get_type_hints(fn)
|
||||
|
||||
# If the return / first param is a Pydantic model, use its schema as root.
|
||||
for param_name, param in sig.parameters.items():
|
||||
ann = hints.get(param_name, param.annotation)
|
||||
if ann is inspect.Parameter.empty:
|
||||
continue
|
||||
if isinstance(ann, type) and issubclass(ann, BaseModel):
|
||||
schema = ann.model_json_schema()
|
||||
schema.pop("title", None)
|
||||
return _strip_unsupported(schema)
|
||||
|
||||
properties: dict[str, Any] = {}
|
||||
required: list[str] = []
|
||||
|
||||
for param_name, param in sig.parameters.items():
|
||||
if param_name == "self":
|
||||
continue
|
||||
ann = hints.get(param_name, param.annotation)
|
||||
if ann is inspect.Parameter.empty:
|
||||
continue
|
||||
properties[param_name] = _annotation_to_schema(ann)
|
||||
if param.default is inspect.Parameter.empty:
|
||||
required.append(param_name)
|
||||
|
||||
return {"type": "object", "properties": properties, "required": required, "additionalProperties": False}
|
||||
|
||||
|
||||
def return_schema(fn: Callable[..., Any]) -> dict[str, Any]:
|
||||
"""Generate JSON Schema for a function's return type."""
|
||||
hints = get_type_hints(fn)
|
||||
ret = hints.get("return", inspect.signature(fn).return_annotation)
|
||||
if ret is inspect.Signature.empty:
|
||||
return {"type": "object", "additionalProperties": True}
|
||||
if isinstance(ret, type) and issubclass(ret, BaseModel):
|
||||
schema = ret.model_json_schema()
|
||||
schema.pop("title", None)
|
||||
return _strip_unsupported(schema)
|
||||
return _annotation_to_schema(ret)
|
||||
|
||||
|
||||
def _annotation_to_schema(ann: Any) -> dict[str, Any]:
|
||||
origin = get_origin(ann)
|
||||
args = get_args(ann)
|
||||
|
||||
# Plain types
|
||||
if ann is str:
|
||||
return {"type": "string"}
|
||||
if ann is int:
|
||||
return {"type": "integer"}
|
||||
if ann is float:
|
||||
return {"type": "number"}
|
||||
if ann is bool:
|
||||
return {"type": "boolean"}
|
||||
|
||||
# Containers — check BEFORE generic Union handling, since
|
||||
# dict[str, X] / list[X] carry args too.
|
||||
if origin is dict or ann is dict:
|
||||
if args and len(args) == 2:
|
||||
return {
|
||||
"type": "object",
|
||||
"additionalProperties": _annotation_to_schema(args[1]),
|
||||
}
|
||||
return {"type": "object", "additionalProperties": True}
|
||||
if origin is list or ann is list:
|
||||
if args and len(args) == 1:
|
||||
return {"type": "array", "items": _annotation_to_schema(args[0])}
|
||||
return {"type": "array"}
|
||||
|
||||
# Optional[X] / Union[X, None] / Union[X, Y]
|
||||
if origin is not None and args:
|
||||
non_none = [a for a in args if a is not type(None)]
|
||||
if len(non_none) == 1 and len(args) > len(non_none):
|
||||
sub = _annotation_to_schema(non_none[0])
|
||||
sub["nullable"] = True
|
||||
return sub
|
||||
if len(non_none) > 1:
|
||||
return {"anyOf": [_annotation_to_schema(a) for a in non_none]}
|
||||
|
||||
if isinstance(ann, type) and issubclass(ann, BaseModel):
|
||||
schema = ann.model_json_schema()
|
||||
schema.pop("title", None)
|
||||
return _strip_unsupported(schema)
|
||||
|
||||
return {"type": "object", "additionalProperties": True}
|
||||
|
||||
|
||||
def _strip_unsupported(schema: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Remove JSON Schema keys that some LLM tool-use pipelines reject."""
|
||||
schema.pop("title", None)
|
||||
return schema
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Tool registration and the @tool decorator."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Awaitable, Callable
|
||||
|
||||
from aicc.protocol import ToolClass
|
||||
from aicc.schema import function_schema, return_schema
|
||||
|
||||
|
||||
ToolImpl = Callable[..., Any | Awaitable[Any]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolDefinition:
|
||||
id: str
|
||||
cls: ToolClass
|
||||
description: str
|
||||
fn: ToolImpl
|
||||
input_schema: dict[str, Any]
|
||||
output_schema: dict[str, Any]
|
||||
requires_capability: str | None = None
|
||||
limits: dict[str, Any] | None = None
|
||||
strict_input: bool = False
|
||||
|
||||
|
||||
def tool(
|
||||
_fn: ToolImpl | None = None,
|
||||
*,
|
||||
id: str | None = None,
|
||||
cls: ToolClass | None = None,
|
||||
description: str | None = None,
|
||||
requires_capability: str | None = None,
|
||||
limits: dict[str, Any] | None = None,
|
||||
strict_input: bool = False,
|
||||
) -> Any:
|
||||
"""Mark a function as an AICC tool.
|
||||
|
||||
Use as @tool or @tool(cls=ToolClass.SENSOR, description=...).
|
||||
"""
|
||||
|
||||
def wrap(fn: ToolImpl) -> ToolDefinition:
|
||||
tool_id = id or fn.__name__
|
||||
if not tool_id.replace("_", "").isalnum() or tool_id[0].isdigit():
|
||||
raise ValueError(
|
||||
f"Tool id '{tool_id}' must match ^[a-z][a-z0-9_]*$ "
|
||||
"(lowercase letters, digits, underscores; cannot start with digit)."
|
||||
)
|
||||
tool_cls = cls or _infer_class(fn)
|
||||
if description:
|
||||
desc = description
|
||||
else:
|
||||
doc = inspect.getdoc(fn) or ""
|
||||
desc = doc.strip().splitlines()[0] if doc else ""
|
||||
if not desc:
|
||||
raise ValueError(
|
||||
f"Tool '{tool_id}' requires a description (set description= or add a docstring)."
|
||||
)
|
||||
return ToolDefinition(
|
||||
id=tool_id,
|
||||
cls=tool_cls,
|
||||
description=desc,
|
||||
fn=fn,
|
||||
input_schema=function_schema(fn),
|
||||
output_schema=return_schema(fn),
|
||||
requires_capability=requires_capability,
|
||||
limits=limits,
|
||||
strict_input=strict_input,
|
||||
)
|
||||
|
||||
if _fn is not None and callable(_fn):
|
||||
return wrap(_fn)
|
||||
return wrap
|
||||
|
||||
|
||||
def _infer_class(fn: ToolImpl) -> ToolClass:
|
||||
"""Default to sensor if the tool starts with get_/read_/observe_; otherwise actuator."""
|
||||
name = fn.__name__.lower()
|
||||
if any(name.startswith(p) for p in ("get_", "read_", "observe_", "inspect_", "see_", "hear_", "smell_", "touch_")):
|
||||
return ToolClass.SENSOR
|
||||
return ToolClass.ACTUATOR
|
||||
|
||||
|
||||
async def call_tool_impl(defn: ToolDefinition, input_data: dict[str, Any]) -> Any:
|
||||
"""Invoke a registered tool's implementation with the given input."""
|
||||
if defn.strict_input:
|
||||
from jsonschema import validate as _validate # type: ignore
|
||||
|
||||
_validate(instance=input_data, schema=defn.input_schema)
|
||||
result = defn.fn(**input_data) if input_data else defn.fn()
|
||||
if asyncio.iscoroutine(result) or inspect.isawaitable(result):
|
||||
result = await result # type: ignore[func-returns-value]
|
||||
return result
|
||||
@@ -0,0 +1,13 @@
|
||||
"""Transport implementations for AICC."""
|
||||
|
||||
from aicc.transport.base import Transport
|
||||
from aicc.transport.in_process import InProcessTransport
|
||||
|
||||
__all__ = ["Transport", "InProcessTransport"]
|
||||
|
||||
try: # pragma: no cover
|
||||
from aicc.transport.websocket import WebSocketClientTransport, WebSocketServer
|
||||
|
||||
__all__ += ["WebSocketClientTransport", "WebSocketServer"]
|
||||
except ImportError: # websockets not installed
|
||||
pass
|
||||
@@ -0,0 +1,20 @@
|
||||
"""Transport interface for AICC message passing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Protocol, runtime_checkable
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class Transport(Protocol):
|
||||
"""Pluggable transport for AICC messages.
|
||||
|
||||
Implementations:
|
||||
- aicc.transport.in_process.InProcessTransport
|
||||
- aicc.transport.websocket.WebSocketClientTransport / WebSocketServer
|
||||
"""
|
||||
|
||||
async def connect(self) -> None: ...
|
||||
async def close(self) -> None: ...
|
||||
async def send(self, message: dict[str, Any]) -> None: ...
|
||||
async def receive(self) -> dict[str, Any]: ...
|
||||
@@ -0,0 +1,87 @@
|
||||
"""In-process transport: bridges AICC clients to a Bridge within one event loop.
|
||||
|
||||
Useful for tests, embedded use, and single-process agents that drive the bridge directly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from aicc.bridge import Bridge
|
||||
from aicc.errors import ConnectionError_
|
||||
from aicc.protocol import PROTOCOL_VERSION, SessionInit
|
||||
|
||||
|
||||
class InProcessTransport:
|
||||
"""Two-sided transport: one side is the bridge, the other is the client.
|
||||
|
||||
Both sides share the same asyncio event loop. The bridge side pushes its
|
||||
manifest into the client's receive queue when start() is awaited.
|
||||
|
||||
Example:
|
||||
async with InProcessTransport.start(bridge) as t:
|
||||
client = AICCClient(t)
|
||||
async with client:
|
||||
manifest = await client.handshake()
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._send_q: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
|
||||
self._recv_q: asyncio.Queue[dict[str, Any]] = asyncio.Queue()
|
||||
self._bridge: Bridge | None = None
|
||||
self._session_id: str | None = None
|
||||
self._connected = False
|
||||
|
||||
@classmethod
|
||||
def start(cls, bridge: Bridge) -> "InProcessTransport":
|
||||
t = cls()
|
||||
t._bridge = bridge
|
||||
return t
|
||||
|
||||
async def connect(self) -> None:
|
||||
if self._bridge is None:
|
||||
raise ConnectionError_("InProcessTransport.start(bridge) must be used for server-side.")
|
||||
if self._connected:
|
||||
return
|
||||
self._connected = True
|
||||
session = self._bridge.open_session()
|
||||
self._session_id = session.session_id
|
||||
manifest: SessionInit = session.manifest # type: ignore[assignment]
|
||||
await self._recv_q.put(manifest.model_dump(by_alias=True))
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._bridge is not None and self._session_id is not None:
|
||||
self._bridge.close_session(self._session_id)
|
||||
self._session_id = None
|
||||
self._connected = False
|
||||
|
||||
async def send(self, message: dict[str, Any]) -> None:
|
||||
if not self._connected or self._bridge is None or self._session_id is None:
|
||||
raise ConnectionError_("Transport is not connected.")
|
||||
sid = message.get("session_id") or self._session_id
|
||||
msg_type = message.get("type")
|
||||
# session_resume carries the session_id; bind it
|
||||
if msg_type == "session_resume" and not self._bridge._sessions.get(sid): # noqa: SLF001
|
||||
# Auto-open if bridge forgot
|
||||
pass
|
||||
response = await self._bridge.handle_message(sid, message)
|
||||
if response is not None:
|
||||
await self._recv_q.put(response)
|
||||
|
||||
async def receive(self) -> dict[str, Any]:
|
||||
if not self._connected:
|
||||
raise ConnectionError_("Transport is not connected.")
|
||||
return await self._recv_q.get()
|
||||
|
||||
async def push_event(self, event: dict[str, Any]) -> None:
|
||||
"""Bridge-side helper to push an event to the connected client."""
|
||||
await self._recv_q.put(event)
|
||||
|
||||
async def __aenter__(self) -> "InProcessTransport":
|
||||
await self.connect()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb) -> None:
|
||||
await self.close()
|
||||
@@ -0,0 +1,116 @@
|
||||
"""WebSocket transport for AICC: client and server."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from aicc.bridge import Bridge
|
||||
from aicc.errors import ConnectionError_
|
||||
|
||||
|
||||
class WebSocketClientTransport:
|
||||
"""AICC client transport over WebSocket.
|
||||
|
||||
Example:
|
||||
t = WebSocketClientTransport("ws://localhost:8765")
|
||||
client = AICCClient(t)
|
||||
"""
|
||||
|
||||
def __init__(self, uri: str):
|
||||
self._uri = uri
|
||||
self._ws = None # type: ignore[var-annotated]
|
||||
|
||||
async def connect(self) -> None:
|
||||
try:
|
||||
import websockets # type: ignore
|
||||
except ImportError as e: # pragma: no cover
|
||||
raise ConnectionError_(
|
||||
"websockets is required for WebSocketClientTransport. "
|
||||
"Install with: pip install aicc[ws] or pip install websockets"
|
||||
) from e
|
||||
self._ws = await websockets.connect(self._uri, max_size=64 * 1024 * 1024)
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._ws is not None:
|
||||
await self._ws.close()
|
||||
self._ws = None
|
||||
|
||||
async def send(self, message: dict[str, Any]) -> None:
|
||||
if self._ws is None:
|
||||
raise ConnectionError_("WebSocket not connected.")
|
||||
await self._ws.send(json.dumps(message))
|
||||
|
||||
async def receive(self) -> dict[str, Any]:
|
||||
if self._ws is None:
|
||||
raise ConnectionError_("WebSocket not connected.")
|
||||
raw = await self._ws.recv()
|
||||
if isinstance(raw, bytes):
|
||||
raw = raw.decode("utf-8")
|
||||
return json.loads(raw)
|
||||
|
||||
|
||||
class WebSocketServer:
|
||||
"""AICC server transport: serves a Bridge over WebSocket.
|
||||
|
||||
Example:
|
||||
bridge = Bridge(name="room")
|
||||
async with WebSocketServer(bridge, port=8765):
|
||||
await asyncio.Future() # run forever
|
||||
"""
|
||||
|
||||
def __init__(self, bridge: Bridge, host: str = "127.0.0.1", port: int = 8765):
|
||||
self._bridge = bridge
|
||||
self._host = host
|
||||
self._port = port
|
||||
self._server = None # type: ignore[var-annotated]
|
||||
self._connections: set[Any] = set()
|
||||
|
||||
async def __aenter__(self) -> "WebSocketServer":
|
||||
await self.start()
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb) -> None:
|
||||
await self.stop()
|
||||
|
||||
async def start(self) -> None:
|
||||
try:
|
||||
import websockets # type: ignore
|
||||
except ImportError as e: # pragma: no cover
|
||||
raise ConnectionError_(
|
||||
"websockets is required for WebSocketServer. "
|
||||
"Install with: pip install websockets"
|
||||
) from e
|
||||
|
||||
async def handler(ws) -> None: # type: ignore[no-untyped-def]
|
||||
self._connections.add(ws)
|
||||
session = self._bridge.open_session()
|
||||
try:
|
||||
manifest = session.manifest
|
||||
if manifest is not None:
|
||||
await ws.send(json.dumps(manifest.model_dump(by_alias=True)))
|
||||
async for raw in ws:
|
||||
if isinstance(raw, bytes):
|
||||
raw = raw.decode("utf-8")
|
||||
msg = json.loads(raw)
|
||||
response = await self._bridge.handle_message(session.session_id, msg)
|
||||
if response is not None:
|
||||
await ws.send(json.dumps(response))
|
||||
finally:
|
||||
self._bridge.close_session(session.session_id)
|
||||
self._connections.discard(ws)
|
||||
|
||||
self._server = await websockets.serve(handler, self._host, self._port)
|
||||
|
||||
async def stop(self) -> None:
|
||||
if self._server is not None:
|
||||
self._server.close()
|
||||
await self._server.wait_closed()
|
||||
self._server = None
|
||||
for ws in list(self._connections):
|
||||
try:
|
||||
await ws.close()
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
self._connections.clear()
|
||||
Reference in New Issue
Block a user