From 0cfb4277b52243df93ff1106a0ed0907535382fb Mon Sep 17 00:00:00 2001 From: Emil Shanaty Date: Sat, 8 Aug 2026 03:10:35 +0300 Subject: [PATCH] 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) --- .gitignore | 19 +++ LICENSE | 21 +++ README.md | 112 +++++++++++++++ aicc/__init__.py | 45 ++++++ aicc/bridge.py | 267 +++++++++++++++++++++++++++++++++++ aicc/client.py | 184 ++++++++++++++++++++++++ aicc/errors.py | 35 +++++ aicc/protocol.py | 193 +++++++++++++++++++++++++ aicc/py.typed | 0 aicc/schema.py | 113 +++++++++++++++ aicc/tool.py | 96 +++++++++++++ aicc/transport/__init__.py | 13 ++ aicc/transport/base.py | 20 +++ aicc/transport/in_process.py | 87 ++++++++++++ aicc/transport/websocket.py | 116 +++++++++++++++ examples/agent_minimal.py | 30 ++++ examples/bridge_minimal.py | 45 ++++++ pyproject.toml | 48 +++++++ tests/__init__.py | 1 + tests/test_integration.py | 164 +++++++++++++++++++++ tests/test_schema.py | 57 ++++++++ tests/test_websocket.py | 32 +++++ 22 files changed, 1698 insertions(+) create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 README.md create mode 100644 aicc/__init__.py create mode 100644 aicc/bridge.py create mode 100644 aicc/client.py create mode 100644 aicc/errors.py create mode 100644 aicc/protocol.py create mode 100644 aicc/py.typed create mode 100644 aicc/schema.py create mode 100644 aicc/tool.py create mode 100644 aicc/transport/__init__.py create mode 100644 aicc/transport/base.py create mode 100644 aicc/transport/in_process.py create mode 100644 aicc/transport/websocket.py create mode 100644 examples/agent_minimal.py create mode 100644 examples/bridge_minimal.py create mode 100644 pyproject.toml create mode 100644 tests/__init__.py create mode 100644 tests/test_integration.py create mode 100644 tests/test_schema.py create mode 100644 tests/test_websocket.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b2a5103 --- /dev/null +++ b/.gitignore @@ -0,0 +1,19 @@ +# Build / editor artifacts +__pycache__/ +*.pyc +*.pyo +.venv/ +venv/ +.pytest_cache/ +.ruff_cache/ +*.egg-info/ +dist/ +build/ +.coverage +htmlcov/ + +# IDE +.idea/ +.vscode/ +*.swp +.DS_Store diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..9d3a21d --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Emil Shanaty + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..4d9d0cb --- /dev/null +++ b/README.md @@ -0,0 +1,112 @@ +# aicc-py + +Python SDK for the [AICC Protocol](https://github.com/emil28092005/AICC-Protocol) (AI-Controlled Character). + +`aicc-py` provides client and bridge primitives for connecting language-model agents to virtual environments via the AICC wire protocol. Engine-agnostic, transport-pluggable, async-first. + +## Status + +`0.1.0` — matches AICC protocol `aicc/0.1`. Alpha. + +## Install + +```bash +pip install aicc +``` + +Or from source: + +```bash +git clone https://github.com/emil28092005/aicc-py +cd aicc-py +pip install -e ".[dev]" +``` + +## Minimal example + +### Bridge side + +```python +from aicc import Bridge + +bridge = Bridge(name="capsule-room", kind="3d") + +@bridge.tool(description="Get the agent's current position and rotation.") +async def proprioception() -> dict: + return { + "position": {"x": 1.0, "y": 0.5, "z": 2.0}, + "rotation": {"yaw": 0.0, "pitch": 0.0}, + "velocity": {"x": 0.0, "y": 0.0, "z": 0.0}, + "health": 100, + } + +@bridge.tool(description="Move the agent forward by the given distance in meters.") +async def move(forward: float = 0.0) -> dict: + # your physics / path-planning code here + return {"moved": forward} + +from aicc.transport import WebSocketServer +async with WebSocketServer(bridge, port=8765): + await bridge.serve_forever() +``` + +### Client side + +```python +import asyncio +from aicc import AICCClient +from aicc.transport import WebSocketClientTransport + +async def main(): + async with AICCClient(WebSocketClientTransport("ws://localhost:8765")) as client: + manifest = await client.handshake() + result = await client.call_tool("proprioception", {}) + print(result.output) + +asyncio.run(main()) +``` + +### In-process (for tests and embedding) + +```python +from aicc import Bridge, AICCClient +from aicc.transport import InProcessTransport + +bridge = Bridge(name="test") +@bridge.tool(description="noop") +async def ping() -> dict: + return {"pong": True} + +async with InProcessTransport(bridge) as transport: + client = AICCClient(transport) + manifest = await client.handshake() + result = await client.call_tool("ping", {}) + assert result.output == {"pong": True} +``` + +## Architecture + +``` +aicc/ + protocol.py # message types, envelope, errors (pydantic models) + client.py # AICCClient — agent side + bridge.py # Bridge — environment side, tool registration + tool.py # @tool decorator and tool metadata + schema.py # JSON Schema generation utilities + transport/ + base.py # Transport interface + in_process.py # In-process transport (tests, embedded) + websocket.py # WebSocket transport (client + server) +``` + +## Development + +```bash +pip install -e ".[dev]" +pytest +ruff check . +``` + +## License + +MIT — see `LICENSE`. diff --git a/aicc/__init__.py b/aicc/__init__.py new file mode 100644 index 0000000..5d27479 --- /dev/null +++ b/aicc/__init__.py @@ -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", +] diff --git a/aicc/bridge.py b/aicc/bridge.py new file mode 100644 index 0000000..fd8fcf2 --- /dev/null +++ b/aicc/bridge.py @@ -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) diff --git a/aicc/client.py b/aicc/client.py new file mode 100644 index 0000000..dc98b3b --- /dev/null +++ b/aicc/client.py @@ -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}'" + ) diff --git a/aicc/errors.py b/aicc/errors.py new file mode 100644 index 0000000..fe43d3f --- /dev/null +++ b/aicc/errors.py @@ -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.""" diff --git a/aicc/protocol.py b/aicc/protocol.py new file mode 100644 index 0000000..3d5a763 --- /dev/null +++ b/aicc/protocol.py @@ -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 diff --git a/aicc/py.typed b/aicc/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/aicc/schema.py b/aicc/schema.py new file mode 100644 index 0000000..05ab8be --- /dev/null +++ b/aicc/schema.py @@ -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 diff --git a/aicc/tool.py b/aicc/tool.py new file mode 100644 index 0000000..add5b29 --- /dev/null +++ b/aicc/tool.py @@ -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 diff --git a/aicc/transport/__init__.py b/aicc/transport/__init__.py new file mode 100644 index 0000000..3bad3cc --- /dev/null +++ b/aicc/transport/__init__.py @@ -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 diff --git a/aicc/transport/base.py b/aicc/transport/base.py new file mode 100644 index 0000000..87c1a4d --- /dev/null +++ b/aicc/transport/base.py @@ -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]: ... diff --git a/aicc/transport/in_process.py b/aicc/transport/in_process.py new file mode 100644 index 0000000..d240e67 --- /dev/null +++ b/aicc/transport/in_process.py @@ -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() diff --git a/aicc/transport/websocket.py b/aicc/transport/websocket.py new file mode 100644 index 0000000..ae3ef9b --- /dev/null +++ b/aicc/transport/websocket.py @@ -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() diff --git a/examples/agent_minimal.py b/examples/agent_minimal.py new file mode 100644 index 0000000..f92b9a9 --- /dev/null +++ b/examples/agent_minimal.py @@ -0,0 +1,30 @@ +"""Minimal AICC agent connecting to the bridge example. + +Run bridge_minimal.py first, then: + python examples/agent_minimal.py +""" + +import asyncio + +from aicc import AICCClient +from aicc.transport import WebSocketClientTransport + + +async def main() -> None: + async with AICCClient(WebSocketClientTransport("ws://127.0.0.1:8765")) as client: + manifest = await client.handshake() + print(f"Connected to world: {manifest.world.name} ({manifest.world.kind})") + print(f"Tools: {[t.id for t in manifest.tools]}") + + res = await client.call_tool("proprioception", {}) + print(f"proprioception -> {res.output}") + + res = await client.call_tool("move", {"forward": 2.5}) + print(f"move -> {res.output}") + + res = await client.call_tool("turn", {"yaw": -45}) + print(f"turn -> {res.output}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/bridge_minimal.py b/examples/bridge_minimal.py new file mode 100644 index 0000000..f1d224c --- /dev/null +++ b/examples/bridge_minimal.py @@ -0,0 +1,45 @@ +"""Minimal AICC bridge with a capsule in a room. + +Run: + python examples/bridge_minimal.py + +Then connect from another terminal: + python examples/agent_minimal.py +""" + +import asyncio + +from aicc import Bridge +from aicc.transport import WebSocketServer + + +async def main() -> None: + bridge = Bridge(name="capsule-room", kind="3d", tick_rate_hz=10.0) + + @bridge.tool(description="Get the agent's current position and rotation.") + async def proprioception() -> dict: + return { + "position": {"x": 1.0, "y": 0.5, "z": 2.0}, + "rotation": {"yaw": 0.0, "pitch": 0.0}, + "velocity": {"x": 0.0, "y": 0.0, "z": 0.0}, + "health": 100, + } + + @bridge.tool(description="Move the agent forward by the given distance in meters.") + async def move(forward: float = 0.0) -> dict: + return {"moved": forward} + + @bridge.tool(description="Turn the agent's heading by the given yaw angle in degrees.") + async def turn(yaw: float = 0.0) -> dict: + return {"turned": yaw} + + print(f"AICC bridge '{bridge.name}' listening on ws://127.0.0.1:8765") + async with WebSocketServer(bridge, port=8765): + await bridge.serve_forever() + + +if __name__ == "__main__": + try: + asyncio.run(main()) + except KeyboardInterrupt: + pass diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..f09ad79 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,48 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "aicc" +version = "0.1.0" +description = "Python SDK for the AI-Controlled Character Protocol" +readme = "README.md" +requires-python = ">=3.11" +license = { text = "MIT" } +authors = [{ name = "Emil Shanaty", email = "emil28092005@gmail.com" }] +keywords = ["aicc", "protocol", "agent", "llm", "gamedev", "embodied-ai"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Software Development :: Libraries :: Python Modules", +] +dependencies = [ + "pydantic>=2.5", + "websockets>=12", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8", + "pytest-asyncio>=0.23", + "ruff>=0.3", +] + +[project.urls] +Protocol = "https://github.com/emil28092005/AICC-Protocol" +Issues = "https://github.com/emil28092005/aicc-py/issues" + +[tool.setuptools.packages.find] +where = ["."] +include = ["aicc*"] + +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] + +[tool.ruff] +line-length = 100 +target-version = "py311" diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..d8df7b8 --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Package marker.""" diff --git a/tests/test_integration.py b/tests/test_integration.py new file mode 100644 index 0000000..efa39e6 --- /dev/null +++ b/tests/test_integration.py @@ -0,0 +1,164 @@ +"""Integration test: end-to-end AICC session via in-process transport.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from aicc import AICCClient, Bridge, ToolClass +from aicc.protocol import ModelClass, TickMode, WorldKind +from aicc.transport.in_process import InProcessTransport + + +@pytest.mark.asyncio +async def test_handshake_and_call_tool(): + bridge = Bridge( + name="test-room", + kind=WorldKind.TEXT, + tick_rate_hz=10.0, + tick_mode=TickMode.FIXED, + agent_model=ModelClass.EDGE_MEDIUM, + ) + + @bridge.tool(description="Get the agent's current position and rotation.") + async def proprioception() -> dict: + return { + "position": {"x": 1.0, "y": 0.5, "z": 2.0}, + "rotation": {"yaw": 0.0, "pitch": 0.0}, + "velocity": {"x": 0.0, "y": 0.0, "z": 0.0}, + "health": 100, + } + + @bridge.tool(description="Move the agent forward by the given distance.") + async def move(forward: float = 0.0) -> dict: + return {"moved": forward} + + @bridge.tool(cls=ToolClass.GENERATOR, description="Place an object in the world.") + async def place_object(prefab: str) -> dict: + return {"spawned": prefab} + + t = InProcessTransport.start(bridge) + async with t: + client = AICCClient(t) + async with client: + manifest = await client.handshake() + assert manifest.world.name == "test-room" + assert manifest.tick_rate_hz == 10.0 + tool_ids = {t.id for t in manifest.tools} + assert {"proprioception", "move", "place_object"} <= tool_ids + + res = await client.call_tool("proprioception", {}) + assert res.ok + assert res.output["position"] == {"x": 1.0, "y": 0.5, "z": 2.0} + + res = await client.call_tool("move", {"forward": 2.5}) + assert res.ok + assert res.output == {"moved": 2.5} + + res = await client.call_tool("place_object", {"prefab": "crate_01"}) + assert res.ok + assert res.output == {"spawned": "crate_01"} + + +@pytest.mark.asyncio +async def test_unknown_tool_returns_tool_error(): + bridge = Bridge(name="x") + + @bridge.tool(description="noop") + async def ping() -> dict: + return {"pong": True} + + t = InProcessTransport.start(bridge) + async with t: + client = AICCClient(t) + async with client: + await client.handshake() + from aicc.errors import ToolError + from aicc.protocol import ErrorCode + + with pytest.raises(ToolError) as ei: + await client.call_tool("does_not_exist", {}) + assert ei.value.code == ErrorCode.TOOL_UNKNOWN + assert not ei.value.retryable + + +@pytest.mark.asyncio +async def test_protocol_mismatch_raises(): + from aicc.errors import ProtocolError + + bridge = Bridge(name="x") + + @bridge.tool(description="noop") + async def ping() -> dict: + return {"pong": True} + + t = InProcessTransport.start(bridge) + async with t: + client = AICCClient(t, protocol="aicc/0.2") + async with client: + with pytest.raises(ProtocolError): + await client.handshake() + + +@pytest.mark.asyncio +async def test_capability_required_tool_is_unavailable(): + from aicc.errors import ToolError + from aicc.protocol import ErrorCode + + bridge = Bridge(name="x") + + @bridge.tool( + description="place", + cls=ToolClass.GENERATOR, + requires_capability="can_modify_world", + ) + async def place_object() -> dict: + return {"spawned": True} + + t = InProcessTransport.start(bridge) + async with t: + client = AICCClient(t) + async with client: + await client.handshake() + with pytest.raises(ToolError) as ei: + await client.call_tool("place_object", {}) + assert ei.value.code == ErrorCode.TOOL_UNAVAILABLE + + +@pytest.mark.asyncio +async def test_execution_failure_is_surfaced(): + from aicc.errors import ToolError + from aicc.protocol import ErrorCode + + bridge = Bridge(name="x") + + @bridge.tool(description="always fails") + async def boom() -> dict: + raise RuntimeError("kaboom") + + t = InProcessTransport.start(bridge) + async with t: + client = AICCClient(t) + async with client: + await client.handshake() + with pytest.raises(ToolError) as ei: + await client.call_tool("boom", {}) + assert ei.value.code == ErrorCode.EXECUTION_FAILED + + +@pytest.mark.asyncio +async def test_schema_generation_basic_types(): + bridge = Bridge(name="x") + + @bridge.tool(description="basic types") + async def fn(a: int, b: str, c: float = 1.0) -> dict: + return {"a": a, "b": b, "c": c} + + schema = bridge.tools()[0].input_schema + assert schema["properties"]["a"]["type"] == "integer" + assert schema["properties"]["b"]["type"] == "string" + assert schema["properties"]["c"]["type"] == "number" + assert "a" in schema["required"] + assert "b" in schema["required"] + assert "c" not in schema["required"] diff --git a/tests/test_schema.py b/tests/test_schema.py new file mode 100644 index 0000000..2582144 --- /dev/null +++ b/tests/test_schema.py @@ -0,0 +1,57 @@ +"""Schema unit tests: function_schema and return_schema.""" + +from aicc.schema import function_schema, return_schema + + +def test_basic_types(): + def fn(a: int, b: str, c: float = 1.0) -> dict: ... + + s = function_schema(fn) + assert s["type"] == "object" + assert s["properties"]["a"]["type"] == "integer" + assert s["properties"]["b"]["type"] == "string" + assert s["properties"]["c"]["type"] == "number" + assert set(s["required"]) == {"a", "b"} + + +def test_list_and_dict(): + def fn(items: list[int], mapping: dict[str, float]) -> dict: ... + + s = function_schema(fn) + assert s["properties"]["items"]["type"] == "array" + assert s["properties"]["items"]["items"]["type"] == "integer" + assert s["properties"]["mapping"]["type"] == "object" + + +def test_optional(): + def fn(x: int | None = None) -> dict: ... + + s = function_schema(fn) + assert s["properties"]["x"]["nullable"] is True + assert "x" not in s["required"] + + +def test_pydantic_model_as_root_input(): + from pydantic import BaseModel + + class Input(BaseModel): + x: int + y: str + + def fn(data: Input) -> dict: ... + + s = function_schema(fn) + assert "properties" in s + assert "x" in s["properties"] + assert "y" in s["properties"] + + +def test_return_schema_basic(): + def fn() -> int: ... + assert return_schema(fn)["type"] == "integer" + + def g() -> list[str]: ... + assert return_schema(g)["type"] == "array" + + def h() -> dict: ... + assert return_schema(h)["type"] == "object" diff --git a/tests/test_websocket.py b/tests/test_websocket.py new file mode 100644 index 0000000..f41667b --- /dev/null +++ b/tests/test_websocket.py @@ -0,0 +1,32 @@ +"""End-to-end test over WebSocket: server + client in the same event loop.""" + +from __future__ import annotations + +import asyncio + +import pytest + +from aicc import AICCClient, Bridge +from aicc.transport import WebSocketClientTransport, WebSocketServer + + +@pytest.mark.asyncio +async def test_websocket_roundtrip(): + bridge = Bridge(name="ws-room") + + @bridge.tool(description="Echo a value.") + async def echo(value: str = "") -> dict: + return {"value": value} + + async with WebSocketServer(bridge, port=0) as server: + # port=0 -> OS-assigned; read it back from the server object + port = server._server.sockets[0].getsockname()[1] # type: ignore[union-attr] + uri = f"ws://127.0.0.1:{port}" + + async with AICCClient(WebSocketClientTransport(uri)) as client: + manifest = await client.handshake() + assert manifest.world.name == "ws-room" + + res = await client.call_tool("echo", {"value": "hi"}) + assert res.ok + assert res.output == {"value": "hi"}