axio 0.1.0__py3-none-any.whl

axio/__init__.py

@@ -0,0 +1 @@
+

axio/agent.py

@@ -0,0 +1,239 @@
+"""Agent: the core agentic loop orchestrating transport, tools, and context."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass, field
+from typing import Any
+
+from axio.blocks import TextBlock, ToolResultBlock, ToolUseBlock
+from axio.context import ContextStore
+from axio.events import (
+    Error,
+    IterationEnd,
+    SessionEndEvent,
+    StreamEvent,
+    TextDelta,
+    ToolInputDelta,
+    ToolResult,
+    ToolUseStart,
+)
+from axio.messages import Message
+from axio.stream import AgentStream
+from axio.tool import Tool
+from axio.transport import CompletionTransport
+from axio.types import StopReason, Usage
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(slots=True)
+class Agent:
+    system: str
+    tools: list[Tool]
+    transport: CompletionTransport
+    max_iterations: int = field(default=50)
+
+    def run_stream(self, user_message: str, context: ContextStore) -> AgentStream:
+        return AgentStream(self._run_loop(user_message, context))
+
+    async def run(self, user_message: str, context: ContextStore) -> str:
+        return await self.run_stream(user_message, context).get_final_text()
+
+    async def dispatch_tools(self, blocks: list[ToolUseBlock], iteration: int) -> list[ToolResultBlock]:
+        tool_names = [b.name for b in blocks]
+        logger.info("Dispatching %d tool(s): %s", len(blocks), tool_names)
+
+        async def _run_one(block: ToolUseBlock) -> ToolResultBlock:
+            tool = self._find_tool(block.name)
+            if tool is None:
+                logger.warning("Unknown tool requested: %s", block.name)
+                return ToolResultBlock(tool_use_id=block.id, content=f"Unknown tool: {block.name}", is_error=True)
+            logger.debug("Tool %s (id=%s) args=%s", block.name, block.id, json.dumps(block.input)[:200])
+            try:
+                result = await tool(**block.input)
+                content = result if isinstance(result, str) else str(result)
+            except Exception as exc:
+                logger.error("Tool %s raised %s: %s", block.name, type(exc).__name__, exc, exc_info=True)
+                return ToolResultBlock(tool_use_id=block.id, content=str(exc), is_error=True)
+            return ToolResultBlock(tool_use_id=block.id, content=content)
+
+        results = list(await asyncio.gather(*[_run_one(b) for b in blocks]))
+        error_count = sum(1 for r in results if r.is_error)
+        logger.info("Tools complete: %d total, %d errors", len(results), error_count)
+        return results
+
+    def _find_tool(self, name: str) -> Tool | None:
+        for tool in self.tools:
+            if tool.name == name:
+                return tool
+        return None
+
+    async def _append(self, context: ContextStore, message: Message) -> None:
+        await context.append(message)
+
+    @staticmethod
+    def _accumulate_text(content: list[TextBlock | ToolUseBlock], delta: str) -> None:
+        """Append text delta — merge into last TextBlock or start a new one."""
+        if content and isinstance(content[-1], TextBlock):
+            content[-1] = TextBlock(text=content[-1].text + delta)
+        else:
+            content.append(TextBlock(text=delta))
+
+    @staticmethod
+    def _finalize_pending_tools(
+        pending: dict[str, dict[str, Any]],
+        usage: Usage,
+    ) -> tuple[list[ToolUseBlock], set[str]]:
+        """Convert streamed tool-call fragments into ToolUseBlocks.
+
+        Returns (blocks, malformed_ids).
+        """
+        blocks: list[ToolUseBlock] = []
+        malformed: set[str] = set()
+        for tid, info in pending.items():
+            raw = "".join(info["json_parts"])
+            if not raw:
+                logger.warning(
+                    "Tool %s (id=%s) received empty arguments (output may be truncated, output_tokens=%d)",
+                    info["name"],
+                    tid,
+                    usage.output_tokens,
+                )
+                inp: dict[str, Any] = {}
+            else:
+                try:
+                    inp = json.loads(raw)
+                except json.JSONDecodeError as exc:
+                    logger.warning(
+                        "Tool %s (id=%s) has malformed JSON arguments: %s\nRaw: %s",
+                        info["name"],
+                        tid,
+                        exc,
+                        raw,
+                    )
+                    malformed.add(tid)
+                    inp = {}
+            blocks.append(ToolUseBlock(id=tid, name=info["name"], input=inp))
+        return blocks, malformed
+
+    async def _run_loop(self, user_message: str, context: ContextStore) -> AsyncGenerator[StreamEvent, None]:
+        total_usage = Usage(0, 0)
+        session_end_emitted = False
+        await self._append(context, Message(role="user", content=[TextBlock(text=user_message)]))
+
+        try:
+            for iteration in range(1, self.max_iterations + 1):
+                history = await context.get_history()
+                logger.info("Iteration %d, history length=%d", iteration, len(history))
+                active_tools = self.tools
+
+                content: list[TextBlock | ToolUseBlock] = []
+                pending: dict[str, dict[str, Any]] = {}
+                stop_reason = StopReason.end_turn
+                malformed: set[str] = set()
+
+                try:
+                    async for event in self.transport.stream(history, active_tools, self.system):
+                        yield event
+                        match event:
+                            case TextDelta(delta=delta):
+                                self._accumulate_text(content, delta)
+                            case ToolUseStart(tool_use_id=tid, name=name):
+                                pending[tid] = {"name": name, "json_parts": []}
+                            case ToolInputDelta(tool_use_id=tid, partial_json=pj):
+                                if tid in pending:
+                                    pending[tid]["json_parts"].append(pj)
+                            case IterationEnd(usage=usage, stop_reason=sr):
+                                blocks, malformed = self._finalize_pending_tools(pending, usage)
+                                content.extend(blocks)
+                                pending.clear()
+                                total_usage = total_usage + usage
+                                await context.add_context_tokens(usage.input_tokens, usage.output_tokens)
+                                stop_reason = sr
+                except Exception as exc:
+                    logger.error("Transport error: %s", exc, exc_info=True)
+                    yield Error(exception=exc)
+                    yield SessionEndEvent(stop_reason=StopReason.error, total_usage=total_usage)
+                    session_end_emitted = True
+                    return
+
+                tool_blocks = [b for b in content if isinstance(b, ToolUseBlock)]
+
+                if tool_blocks:
+                    if stop_reason != StopReason.tool_use:
+                        logger.warning(
+                            "Dispatching %d tool(s) despite stop_reason=%s",
+                            len(tool_blocks),
+                            stop_reason,
+                        )
+
+                    # Dispatch tools BEFORE appending to context — cancellation
+                    # between here and the two appends below cannot leave orphan
+                    # ToolUseBlocks in the persistent context store.
+                    valid = [b for b in tool_blocks if b.id not in malformed]
+                    error_results = [
+                        ToolResultBlock(
+                            tool_use_id=b.id,
+                            content=(
+                                f"Malformed JSON arguments for tool {b.name}."
+                                f" Raw input could not be parsed. Please retry the tool call"
+                                f" with valid JSON arguments."
+                            ),
+                            is_error=True,
+                        )
+                        for b in tool_blocks
+                        if b.id in malformed
+                    ]
+                    dispatched = await self.dispatch_tools(valid, iteration) if valid else []
+                    results = dispatched + error_results
+
+                    # Append both messages atomically (assistant + tool results)
+                    await self._append(context, Message(role="assistant", content=list(content)))
+                    await self._append(context, Message(role="user", content=list(results)))
+
+                    # Yield ToolResult events
+                    by_id = {b.id: b for b in tool_blocks}
+                    for r in results:
+                        block = by_id.get(r.tool_use_id)
+                        result_content = (
+                            r.content
+                            if isinstance(r.content, str)
+                            else "\n".join(b.text for b in r.content if isinstance(b, TextBlock))
+                        )
+                        yield ToolResult(
+                            tool_use_id=r.tool_use_id,
+                            name=block.name if block else "",
+                            is_error=r.is_error,
+                            content=result_content,
+                            input=block.input if block else {},
+                        )
+                    continue
+
+                await self._append(context, Message(role="assistant", content=list(content)))
+
+                match stop_reason:
+                    case StopReason.end_turn:
+                        logger.debug("End turn: total_usage=%s", total_usage)
+                        yield SessionEndEvent(stop_reason=StopReason.end_turn, total_usage=total_usage)
+                        session_end_emitted = True
+                        return
+                    case StopReason.max_tokens | StopReason.error:
+                        yield Error(exception=RuntimeError(f"Transport stopped with: {stop_reason}"))
+                        yield SessionEndEvent(stop_reason=StopReason.error, total_usage=total_usage)
+                        session_end_emitted = True
+                        return
+
+            logger.warning("Max iterations (%d) reached", self.max_iterations)
+            yield SessionEndEvent(stop_reason=StopReason.error, total_usage=total_usage)
+            session_end_emitted = True
+
+        except GeneratorExit:
+            return
+        except BaseException:
+            if not session_end_emitted:
+                yield SessionEndEvent(stop_reason=StopReason.error, total_usage=total_usage)
+            raise

axio/blocks.py

@@ -0,0 +1,98 @@
+"""Content blocks: TextBlock, ImageBlock, ToolUseBlock, ToolResultBlock."""
+
+from __future__ import annotations
+
+import base64
+from dataclasses import dataclass
+from functools import singledispatch
+from typing import Any, Literal
+
+from axio.types import ToolCallID, ToolName
+
+
+class ContentBlock:
+    """Base class for all content blocks."""
+
+    __slots__ = ()
+
+
+@dataclass(frozen=True, slots=True)
+class TextBlock(ContentBlock):
+    text: str
+
+
+@dataclass(frozen=True, slots=True)
+class ImageBlock(ContentBlock):
+    media_type: Literal["image/jpeg", "image/png", "image/gif", "image/webp"]
+    data: bytes
+
+
+@dataclass(frozen=True, slots=True)
+class ToolUseBlock(ContentBlock):
+    id: ToolCallID
+    name: ToolName
+    input: dict[str, Any]
+
+
+@dataclass(frozen=True, slots=True)
+class ToolResultBlock(ContentBlock):
+    tool_use_id: ToolCallID
+    content: str | list[TextBlock | ImageBlock]
+    is_error: bool = False
+
+
+@singledispatch
+def to_dict(block: ContentBlock) -> dict[str, Any]:
+    """Serialize a ContentBlock to a plain dict."""
+    msg = f"Unknown block type: {type(block).__name__}"
+    raise TypeError(msg)
+
+
+@to_dict.register(TextBlock)
+def _text_to_dict(block: TextBlock) -> dict[str, Any]:
+    return {"type": "text", "text": block.text}
+
+
+@to_dict.register(ImageBlock)
+def _image_to_dict(block: ImageBlock) -> dict[str, Any]:
+    return {"type": "image", "media_type": block.media_type, "data": base64.b64encode(block.data).decode()}
+
+
+@to_dict.register(ToolUseBlock)
+def _tool_use_to_dict(block: ToolUseBlock) -> dict[str, Any]:
+    return {"type": "tool_use", "id": block.id, "name": block.name, "input": block.input}
+
+
+@to_dict.register(ToolResultBlock)
+def _tool_result_to_dict(block: ToolResultBlock) -> dict[str, Any]:
+    if isinstance(block.content, str):
+        serialized_content: str | list[dict[str, Any]] = block.content
+    else:
+        serialized_content = [to_dict(b) for b in block.content]
+    return {
+        "type": "tool_result",
+        "tool_use_id": block.tool_use_id,
+        "content": serialized_content,
+        "is_error": block.is_error,
+    }
+
+
+def from_dict(data: dict[str, Any]) -> ContentBlock:
+    """Deserialize a plain dict to a ContentBlock."""
+    match data["type"]:
+        case "text":
+            return TextBlock(text=data["text"])
+        case "image":
+            return ImageBlock(media_type=data["media_type"], data=base64.b64decode(data["data"]))
+        case "tool_use":
+            return ToolUseBlock(id=data["id"], name=data["name"], input=data["input"])
+        case "tool_result":
+            raw = data["content"]
+            if isinstance(raw, str):
+                content: str | list[TextBlock | ImageBlock] = raw
+            else:
+                content = [from_dict(b) for b in raw]  # type: ignore[misc]
+            return ToolResultBlock(tool_use_id=data["tool_use_id"], content=content, is_error=data["is_error"])
+        case _:
+            msg = f"Unknown block type: {data['type']}"
+            raise ValueError(msg)

axio/context.py

@@ -0,0 +1,197 @@
+"""ContextStore: protocol for conversation history storage."""
+
+from __future__ import annotations
+
+import copy
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Self
+from uuid import uuid4
+
+from axio.blocks import TextBlock, ToolResultBlock
+from axio.messages import Message
+from axio.transport import CompletionTransport
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class SessionInfo:
+    session_id: str
+    message_count: int
+    preview: str
+    created_at: str
+    input_tokens: int = 0
+    output_tokens: int = 0
+
+
+class ContextStore(ABC):
+    @property
+    @abstractmethod
+    def session_id(self) -> str: ...
+
+    @abstractmethod
+    async def append(self, message: Message) -> None: ...
+
+    @abstractmethod
+    async def get_history(self) -> list[Message]: ...
+
+    @abstractmethod
+    async def clear(self) -> None: ...
+
+    @abstractmethod
+    async def fork(self) -> ContextStore: ...
+
+    @abstractmethod
+    async def set_context_tokens(self, input_tokens: int, output_tokens: int) -> None: ...
+
+    @abstractmethod
+    async def get_context_tokens(self) -> tuple[int, int]: ...
+
+    @abstractmethod
+    async def close(self) -> None: ...
+
+    async def list_sessions(self) -> list[SessionInfo]:
+        """List available sessions. Default: returns a single entry for the current session."""
+        history = await self.get_history()
+        in_tok, out_tok = await self.get_context_tokens()
+        preview = "(empty)"
+        for msg in history:
+            if msg.role == "user":
+                for block in msg.content:
+                    if isinstance(block, TextBlock):
+                        text = block.text
+                        preview = text[:80] + ("..." if len(text) > 80 else "")
+                        break
+                break
+        return [
+            SessionInfo(
+                session_id=self.session_id,
+                message_count=len(history),
+                preview=preview,
+                created_at="",
+                input_tokens=in_tok,
+                output_tokens=out_tok,
+            ),
+        ]
+
+    async def add_context_tokens(self, input_tokens: int, output_tokens: int) -> None:
+        cur_in, cur_out = await self.get_context_tokens()
+        await self.set_context_tokens(cur_in + input_tokens, cur_out + output_tokens)
+
+    @classmethod
+    async def from_history(cls, history: list[Message]) -> Self:
+        """Create a new ContextStore pre-populated with *history*."""
+        store = cls()
+        for message in history:
+            await store.append(message)
+        return store
+
+    @classmethod
+    async def from_context(cls, context: ContextStore) -> Self:
+        return await cls.from_history(await context.get_history())
+
+
+class MemoryContextStore(ContextStore):
+    """Simple in-memory context store. fork() returns a deep copy."""
+
+    def __init__(self, history: list[Message] | None = None) -> None:
+        self._session_id = uuid4().hex
+        self._history: list[Message] = list(history or [])
+        self._input_tokens: int = 0
+        self._output_tokens: int = 0
+
+    @property
+    def session_id(self) -> str:
+        return self._session_id
+
+    async def append(self, message: Message) -> None:
+        self._history.append(message)
+
+    async def get_history(self) -> list[Message]:
+        return list(self._history)
+
+    async def clear(self) -> None:
+        self._history.clear()
+        self._input_tokens = 0
+        self._output_tokens = 0
+
+    async def fork(self) -> MemoryContextStore:
+        store = MemoryContextStore(copy.deepcopy(self._history))
+        store._input_tokens = self._input_tokens
+        store._output_tokens = self._output_tokens
+        return store
+
+    async def set_context_tokens(self, input_tokens: int, output_tokens: int) -> None:
+        self._input_tokens = input_tokens
+        self._output_tokens = output_tokens
+
+    async def get_context_tokens(self) -> tuple[int, int]:
+        return self._input_tokens, self._output_tokens
+
+    async def close(self) -> None:
+        pass
+
+
+_DEFAULT_COMPACTION_PROMPT = (
+    "You are a conversation summarizer. You will see a conversation between"
+    " a user and an AI assistant, including tool calls and their results."
+    " Produce a concise summary preserving: user goals, decisions made,"
+    " key facts, tool outcomes, and state changes. Write as narrative prose,"
+    " not as a transcript."
+)
+
+
+async def compact_context(
+    context: ContextStore,
+    transport: CompletionTransport,
+    *,
+    max_messages: int = 20,
+    keep_recent: int = 6,
+    system_prompt: str | None = None,
+) -> list[Message] | None:
+    """Summarize old messages from *context*, keeping recent ones verbatim.
+
+    Returns a compacted message list ready to populate a fresh store,
+    or ``None`` when no compaction is needed.
+    """
+    history = await context.get_history()
+    if len(history) <= max_messages:
+        return None
+
+    split = _find_safe_boundary(history, keep_recent)
+    if split <= 0:
+        return None
+
+    old, recent = history[:split], history[split:]
+
+    # Deferred import to avoid circular dependency (context ↔ agent)
+    from axio.agent import Agent
+
+    summary_ctx = MemoryContextStore(old)
+    agent = Agent(
+        system=system_prompt or _DEFAULT_COMPACTION_PROMPT,
+        tools=[],
+        transport=transport,
+        max_iterations=1,
+    )
+    try:
+        summary = await agent.run("Summarize the conversation above.", summary_ctx)
+    except Exception:
+        logger.warning("Context compaction failed, keeping original history", exc_info=True)
+        return None
+
+    return [
+        Message(role="user", content=[TextBlock(text=summary)]),
+        Message(role="assistant", content=[TextBlock(text="Understood, context restored.")]),
+        *recent,
+    ]
+
+
+def _find_safe_boundary(history: list[Message], keep_recent: int) -> int:
+    """Return a split index that never separates a tool_use from its tool_result."""
+    split = len(history) - keep_recent
+    while split > 0 and any(isinstance(b, ToolResultBlock) for b in history[split].content):
+        split -= 1
+    return split

axio/events.py

@@ -0,0 +1,66 @@
+"""Stream events: all variants emitted by AgentStream."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from axio.types import StopReason, ToolCallID, ToolName, Usage
+
+
+@dataclass(frozen=True, slots=True)
+class ReasoningDelta:
+    index: int
+    delta: str
+
+
+@dataclass(frozen=True, slots=True)
+class TextDelta:
+    index: int
+    delta: str
+
+
+@dataclass(frozen=True, slots=True)
+class ToolUseStart:
+    index: int
+    tool_use_id: ToolCallID
+    name: ToolName
+
+
+@dataclass(frozen=True, slots=True)
+class ToolInputDelta:
+    index: int
+    tool_use_id: ToolCallID
+    partial_json: str
+
+
+@dataclass(frozen=True, slots=True)
+class ToolResult:
+    tool_use_id: ToolCallID
+    name: ToolName
+    is_error: bool
+    content: str = ""
+    input: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True, slots=True)
+class IterationEnd:
+    iteration: int
+    stop_reason: StopReason
+    usage: Usage
+
+
+@dataclass(frozen=True, slots=True)
+class Error:
+    exception: BaseException
+
+
+@dataclass(frozen=True, slots=True)
+class SessionEndEvent:
+    stop_reason: StopReason
+    total_usage: Usage
+
+
+type StreamEvent = (
+    ReasoningDelta | TextDelta | ToolUseStart | ToolInputDelta | ToolResult | IterationEnd | Error | SessionEndEvent
+)

axio/exceptions.py

@@ -0,0 +1,21 @@
+"""Exception hierarchy for axio."""
+
+
+class AxioError(Exception):
+    """Base exception for all axio errors."""
+
+
+class ToolError(AxioError):
+    """Base for tool-related errors."""
+
+
+class GuardError(ToolError):
+    """Guard denied or crashed during permission check."""
+
+
+class HandlerError(ToolError):
+    """Handler raised during execution."""
+
+
+class StreamError(AxioError):
+    """Error during stream collection."""

axio/messages.py

@@ -0,0 +1,21 @@
+"""Message: the fundamental unit of conversation history."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+from axio.blocks import ContentBlock, from_dict, to_dict
+
+
+@dataclass(slots=True)
+class Message:
+    role: Literal["user", "assistant"]
+    content: list[ContentBlock] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {"role": self.role, "content": [to_dict(b) for b in self.content]}
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Message:
+        return cls(role=data["role"], content=[from_dict(b) for b in data["content"]])

axio/models.py

@@ -0,0 +1,102 @@
+"""Transport-agnostic model types: Capability, ModelSpec, ModelRegistry."""
+
+from __future__ import annotations
+
+from collections.abc import ItemsView, Iterable, Iterator, KeysView, MutableMapping, ValuesView
+from dataclasses import dataclass
+from enum import StrEnum
+
+
+@dataclass(frozen=True, slots=True)
+class TransportMeta:
+    """Metadata a transport plugin declares about itself."""
+
+    label: str
+    api_key_env: str
+    role_defaults: dict[str, str]
+
+
+class Capability(StrEnum):
+    text = "text"
+    vision = "vision"
+    reasoning = "reasoning"
+    tool_use = "tool_use"
+    json_mode = "json_mode"
+    structured_outputs = "structured_outputs"
+    embedding = "embedding"
+
+
+@dataclass(frozen=True, slots=True)
+class ModelSpec:
+    id: str
+    capabilities: frozenset[Capability] = frozenset()
+    max_output_tokens: int = 8192
+    context_window: int = 128000
+    input_cost: float = 0.0
+    output_cost: float = 0.0
+
+
+class ModelRegistry(MutableMapping[str, ModelSpec]):
+    __slots__ = ("_models",)
+
+    def __init__(self, models: Iterable[ModelSpec] | None = None) -> None:
+        self._models: dict[str, ModelSpec] = {m.id: m for m in (models or [])}
+
+    def __setitem__(self, key: str, value: ModelSpec, /) -> None:
+        if not isinstance(value, ModelSpec):
+            raise ValueError("ModelRegistry values must be ModelSpec instances")
+        self._models[key] = value
+
+    def __delitem__(self, key: str, /) -> None:
+        del self._models[key]
+
+    def __getitem__(self, key: str, /) -> ModelSpec:
+        return self._models[key]
+
+    def __len__(self) -> int:
+        return len(self._models)
+
+    def __iter__(self) -> Iterator[ModelSpec]:  # type: ignore[override]
+        return iter(self._models.values())
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, ModelRegistry):
+            return self._models == other._models
+        if isinstance(other, dict):
+            return self._models == other
+        return NotImplemented
+
+    def __repr__(self) -> str:
+        return f"ModelRegistry({self._models!r})"
+
+    def clear(self) -> None:
+        self._models.clear()
+
+    def keys(self) -> KeysView[str]:
+        return self._models.keys()
+
+    def values(self) -> ValuesView[ModelSpec]:
+        return self._models.values()
+
+    def items(self) -> ItemsView[str, ModelSpec]:
+        return self._models.items()
+
+    def by_prefix(self, prefix: str) -> ModelRegistry:
+        return ModelRegistry(v for k, v in self._models.items() if k.startswith(prefix))
+
+    def by_capability(self, *caps: Capability) -> ModelRegistry:
+        required = frozenset(caps)
+        return ModelRegistry(v for v in self._models.values() if required <= v.capabilities)
+
+    def search(self, *q: str) -> ModelRegistry:
+        """search by parts of id"""
+        return ModelRegistry(v for k, v in self._models.items() if all(part in k for part in q))
+
+    def by_cost(self, *, output: bool = False, desc: bool = False) -> ModelRegistry:
+        """Return registry ordered by cost (input by default, output if *output=True*)."""
+        attr = "output_cost" if output else "input_cost"
+        items = sorted(self._models.values(), key=lambda v: getattr(v, attr), reverse=desc)
+        return ModelRegistry(items)
+
+    def ids(self) -> list[str]:
+        return list(self._models)

axio/permission.py

@@ -0,0 +1,50 @@
+"""Permission system: guards that gate tool execution."""
+
+from __future__ import annotations
+
+import asyncio
+from abc import ABC, abstractmethod
+from typing import Any
+
+from axio.exceptions import GuardError
+
+
+class PermissionGuard(ABC):
+    """Gate for tool calls. Return handler to allow, raise to deny.
+
+    Tool calls guards via ``await guard(instance)``.
+    """
+
+    async def __call__(self, handler: Any) -> Any:
+        return await self.check(handler)
+
+    @abstractmethod
+    async def check(self, handler: Any) -> Any: ...
+
+
+class ConcurrentGuard(PermissionGuard, ABC):
+    """Guard with concurrency control.
+
+    Subclass and override ``check()``.  ``__call__`` acquires the semaphore
+    then delegates to ``check()``.  Set ``concurrency`` to control parallelism
+    (default 1 — one check at a time).
+    """
+
+    concurrency: int = 1
+
+    def __init__(self) -> None:
+        self._semaphore = asyncio.Semaphore(self.concurrency)
+
+    async def __call__(self, handler: Any) -> Any:
+        async with self._semaphore:
+            return await self.check(handler)
+
+
+class AllowAllGuard(PermissionGuard):
+    async def check(self, handler: Any) -> Any:
+        return handler
+
+
+class DenyAllGuard(PermissionGuard):
+    async def check(self, handler: Any) -> Any:
+        raise GuardError("denied")

axio/selector.py

@@ -0,0 +1,121 @@
+"""Tool selector: choose relevant tools per query via embedding similarity."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import AsyncIterator
+from typing import Protocol, runtime_checkable
+
+from axio.blocks import TextBlock, ToolResultBlock
+from axio.events import StreamEvent
+from axio.messages import Message
+from axio.tool import Tool
+from axio.transport import CompletionTransport, EmbeddingTransport
+
+
+@runtime_checkable
+class ToolSelector(Protocol):
+    async def select(self, messages: list[Message], tools: list[Tool]) -> list[Tool]: ...
+
+
+def _cosine_similarity(a: list[float], b: list[float]) -> float:
+    dot = sum(x * y for x, y in zip(a, b))
+    norm_a = math.sqrt(sum(x * x for x in a))
+    norm_b = math.sqrt(sum(x * x for x in b))
+    if norm_a == 0.0 or norm_b == 0.0:
+        return 0.0
+    return dot / (norm_a * norm_b)
+
+
+def _extract_last_user_text(messages: list[Message]) -> str | None:
+    """Return joined text from the last user message, or None if it's a tool-result iteration."""
+    for msg in reversed(messages):
+        if msg.role != "user":
+            continue
+        texts = [b.text for b in msg.content if isinstance(b, TextBlock)]
+        if texts:
+            return " ".join(texts)
+        # User message with only ToolResultBlocks → tool-result iteration
+        if any(isinstance(b, ToolResultBlock) for b in msg.content):
+            return None
+        return None
+    return None
+
+
+class EmbeddingToolSelector:
+    def __init__(
+        self,
+        transport: EmbeddingTransport,
+        *,
+        top_k: int = 5,
+        pinned: frozenset[str] = frozenset(),
+    ) -> None:
+        self._transport = transport
+        self._top_k = top_k
+        self._pinned = pinned
+        self._tool_embeddings: dict[str, list[float]] = {}
+        self._tool_descriptions: dict[str, str] = {}
+
+    async def _ensure_embeddings(self, tools: list[Tool]) -> None:
+        to_embed: list[tuple[str, str]] = []
+        for tool in tools:
+            cached_desc = self._tool_descriptions.get(tool.name)
+            if cached_desc is None or cached_desc != tool.description:
+                to_embed.append((tool.name, tool.description))
+
+        if not to_embed:
+            return
+
+        texts = [desc for _, desc in to_embed]
+        vectors = await self._transport.embed(texts)
+        for (name, desc), vec in zip(to_embed, vectors):
+            self._tool_embeddings[name] = vec
+            self._tool_descriptions[name] = desc
+
+    async def select(self, messages: list[Message], tools: list[Tool]) -> list[Tool]:
+        if len(tools) <= self._top_k:
+            return tools
+
+        query = _extract_last_user_text(messages)
+        if query is None:
+            return tools
+
+        await self._ensure_embeddings(tools)
+
+        query_vec = (await self._transport.embed([query]))[0]
+
+        pinned_tools: list[Tool] = []
+        scorable: list[tuple[float, Tool]] = []
+
+        for tool in tools:
+            if tool.name in self._pinned:
+                pinned_tools.append(tool)
+            elif tool.name in self._tool_embeddings:
+                score = _cosine_similarity(query_vec, self._tool_embeddings[tool.name])
+                scorable.append((score, tool))
+
+        scorable.sort(key=lambda x: x[0], reverse=True)
+        remaining_slots = max(0, self._top_k - len(pinned_tools))
+        selected = pinned_tools + [t for _, t in scorable[:remaining_slots]]
+        return selected
+
+
+class ToolFilteringTransport:
+    """CompletionTransport decorator that filters tools via a ToolSelector."""
+
+    def __init__(self, transport: CompletionTransport, selector: ToolSelector) -> None:
+        self._transport = transport
+        self._selector = selector
+
+    def stream(self, messages: list[Message], tools: list[Tool], system: str) -> AsyncIterator[StreamEvent]:
+        return self._filtered_stream(messages, tools, system)
+
+    async def _filtered_stream(
+        self,
+        messages: list[Message],
+        tools: list[Tool],
+        system: str,
+    ) -> AsyncIterator[StreamEvent]:
+        selected = await self._selector.select(messages, tools)
+        async for event in self._transport.stream(messages, selected, system):
+            yield event

axio/stream.py

@@ -0,0 +1,57 @@
+"""AgentStream: async iterator wrapper over the agent event generator."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator
+
+from axio.events import Error, SessionEndEvent, StreamEvent, TextDelta
+from axio.exceptions import StreamError
+
+
+class AgentStream:
+    def __init__(self, generator: AsyncGenerator[StreamEvent, None]) -> None:
+        self._generator = generator
+        self._closed = False
+
+    def __aiter__(self) -> AgentStream:
+        return self
+
+    async def __anext__(self) -> StreamEvent:
+        if self._closed:
+            raise StopAsyncIteration
+        try:
+            return await self._generator.__anext__()
+        except StopAsyncIteration:
+            self._closed = True
+            raise
+
+    async def aclose(self) -> None:
+        if not self._closed:
+            self._closed = True
+            await self._generator.aclose()
+
+    async def get_final_text(self) -> str:
+        parts: list[str] = []
+        try:
+            async for event in self:
+                if isinstance(event, Error):
+                    raise StreamError(str(event.exception)) from event.exception
+                if isinstance(event, TextDelta):
+                    parts.append(event.delta)
+        finally:
+            await self.aclose()
+        return "".join(parts)
+
+    async def get_session_end(self) -> SessionEndEvent:
+        result: SessionEndEvent | None = None
+        try:
+            async for event in self:
+                if isinstance(event, Error):
+                    raise StreamError(str(event.exception)) from event.exception
+                if isinstance(event, SessionEndEvent):
+                    result = event
+        finally:
+            await self.aclose()
+        if result is None:
+            raise StreamError("Stream ended without SessionEndEvent")
+        return result

axio/testing.py

@@ -0,0 +1,87 @@
+"""Shared test helpers: StubTransport, fixtures, response builders."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import AsyncIterator
+from typing import Any
+
+from axio.context import MemoryContextStore
+from axio.events import IterationEnd, StreamEvent, TextDelta, ToolInputDelta, ToolUseStart
+from axio.messages import Message
+from axio.tool import Tool, ToolHandler
+from axio.types import StopReason, Usage
+
+
+class MsgInput(ToolHandler):
+    msg: str
+
+    async def __call__(self) -> str:
+        return self.model_dump_json()
+
+
+class StubTransport:
+    """A CompletionTransport that yields pre-configured event sequences.
+
+    Each call to stream() pops the next sequence from the list.
+    """
+
+    def __init__(self, responses: list[list[StreamEvent]] | None = None) -> None:
+        self._responses: list[list[StreamEvent]] = list(responses or [])
+        self._call_count = 0
+
+    async def _generate(self, events: list[StreamEvent]) -> AsyncIterator[StreamEvent]:
+        for event in events:
+            yield event
+
+    def stream(self, messages: list[Message], tools: list[Tool], system: str) -> AsyncIterator[StreamEvent]:
+        idx = min(self._call_count, len(self._responses) - 1)
+        events = self._responses[idx]
+        self._call_count += 1
+        return self._generate(events)
+
+
+def make_tool_use_response(
+    tool_name: str = "echo",
+    tool_id: str = "call_1",
+    tool_input: dict[str, Any] | None = None,
+    iteration: int = 1,
+    usage: Usage | None = None,
+) -> list[StreamEvent]:
+    """Build a standard tool_use response event sequence."""
+    inp = tool_input or {"msg": "hi"}
+    u = usage or Usage(10, 5)
+    return [
+        ToolUseStart(0, tool_id, tool_name),
+        ToolInputDelta(0, tool_id, json.dumps(inp)),
+        IterationEnd(iteration, StopReason.tool_use, u),
+    ]
+
+
+def make_text_response(text: str = "Done", iteration: int = 2, usage: Usage | None = None) -> list[StreamEvent]:
+    """Build a standard end_turn text response event sequence."""
+    u = usage or Usage(10, 5)
+    return [
+        TextDelta(0, text),
+        IterationEnd(iteration, StopReason.end_turn, u),
+    ]
+
+
+def make_stub_transport() -> StubTransport:
+    return StubTransport(
+        [
+            [
+                TextDelta(0, "Hello"),
+                TextDelta(0, " world"),
+                IterationEnd(1, StopReason.end_turn, Usage(10, 5)),
+            ]
+        ]
+    )
+
+
+def make_ephemeral_context() -> MemoryContextStore:
+    return MemoryContextStore()
+
+
+def make_echo_tool() -> Tool:
+    return Tool(name="echo", description="Returns input as JSON", handler=MsgInput)

axio/tool.py

@@ -0,0 +1,74 @@
+"""Tool: frozen dataclass binding a ToolHandler to a name, guard, and concurrency."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, field
+from typing import Any
+
+from pydantic import BaseModel
+
+from axio.exceptions import GuardError, HandlerError
+from axio.permission import PermissionGuard
+from axio.types import ToolName
+
+type JSONSchema = dict[str, Any]
+
+
+class ToolHandler(BaseModel):
+    """Base for tool handlers.
+
+    Subclass fields define the input JSON-schema.
+    Override ``async def __call__`` to implement execution logic.
+    Pydantic provides ``__repr__`` automatically — override for custom display.
+    """
+
+    async def __call__(self) -> str:
+        raise NotImplementedError
+
+
+@dataclass(frozen=True, slots=True)
+class Tool:
+    name: ToolName
+    description: str
+    handler: type[ToolHandler]
+    guards: tuple[PermissionGuard, ...] = ()
+    concurrency: int | None = None
+
+    _semaphore: asyncio.Semaphore | None = field(init=False, default=None, repr=False, compare=False)
+
+    def __post_init__(self) -> None:
+        if self.concurrency is not None:
+            object.__setattr__(self, "_semaphore", asyncio.Semaphore(self.concurrency))
+
+    @asynccontextmanager
+    async def _acquire(self) -> AsyncGenerator[None, None]:
+        if self._semaphore is None:
+            yield
+            return
+
+        async with self._semaphore:
+            yield
+
+    @property
+    def input_schema(self) -> JSONSchema:
+        return self.handler.model_json_schema()
+
+    async def __call__(self, **kwargs: Any) -> Any:
+        async with self._acquire():
+            instance = self.handler.model_validate(kwargs)
+            for guard in self.guards:
+                try:
+                    instance = await guard(instance)
+                except GuardError:
+                    raise
+                except Exception as exc:
+                    raise GuardError(str(exc)) from exc
+            try:
+                return await instance()
+            except HandlerError:
+                raise
+            except Exception as exc:
+                raise HandlerError(str(exc)) from exc

axio/transport.py

@@ -0,0 +1,35 @@
+"""Transport protocols: completion, image gen, TTS, STT."""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from typing import Protocol, runtime_checkable
+
+from axio.events import StreamEvent
+from axio.messages import Message
+from axio.tool import Tool
+
+
+@runtime_checkable
+class CompletionTransport(Protocol):
+    def stream(self, messages: list[Message], tools: list[Tool], system: str) -> AsyncIterator[StreamEvent]: ...
+
+
+@runtime_checkable
+class ImageGenTransport(Protocol):
+    async def generate(self, prompt: str, *, size: tuple[int, int] | None = None, n: int = 1) -> list[bytes]: ...
+
+
+@runtime_checkable
+class TTSTransport(Protocol):
+    def synthesize(self, text: str, *, voice: str | None = None) -> AsyncIterator[bytes]: ...
+
+
+@runtime_checkable
+class STTTransport(Protocol):
+    async def transcribe(self, audio: bytes, media_type: str = "audio/wav") -> str: ...
+
+
+@runtime_checkable
+class EmbeddingTransport(Protocol):
+    async def embed(self, texts: list[str]) -> list[list[float]]: ...

axio/types.py

@@ -0,0 +1,28 @@
+"""Primitive types: ToolName, ToolCallID, StopReason, Usage."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import StrEnum
+
+type ToolName = str
+type ToolCallID = str
+
+
+class StopReason(StrEnum):
+    end_turn = "end_turn"
+    tool_use = "tool_use"
+    max_tokens = "max_tokens"
+    error = "error"
+
+
+@dataclass(frozen=True, slots=True)
+class Usage:
+    input_tokens: int
+    output_tokens: int
+
+    def __add__(self, other: Usage) -> Usage:
+        return Usage(
+            input_tokens=self.input_tokens + other.input_tokens,
+            output_tokens=self.output_tokens + other.output_tokens,
+        )

axio-0.1.0.dist-info/METADATA

@@ -0,0 +1,8 @@
+Metadata-Version: 2.4
+Name: axio
+Version: 0.1.0
+Summary: Minimal, streaming-first, protocol-driven foundation for LLM-powered agents
+License: MIT
+License-File: LICENSE
+Requires-Python: >=3.12
+Requires-Dist: pydantic>=2

axio-0.1.0.dist-info/RECORD

@@ -0,0 +1,19 @@
+axio/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+axio/agent.py,sha256=HqWptDuRdz1Hns-rcm43iiN3nEq35zmUBihWCH0EcOA,10691
+axio/blocks.py,sha256=MUOTqZcTWTi2d1b7WEoomueaCulkzoJ2QxJEmNiM5qA,3011
+axio/context.py,sha256=CTAaqJVp9aT6BkSkZbOO56b6-Kwaw4VcprpX6Rivzzc,6202
+axio/events.py,sha256=HimirQY5g9mSa1TWpSq1cyxl5SnUMoas3hrWPfT4GXs,1299
+axio/exceptions.py,sha256=pvgKVcVfJ-HsMr2JatDIE7lzdw-mMQZZLszzSCLj3Lk,422
+axio/messages.py,sha256=0-GFC7BNPTyrJaVA584EjNN3c6LjSMaftH9CQR_Tr04,656
+axio/models.py,sha256=FA8h3MQqQS0-oz50QST7ZiYyjiPWxJ73rivEQ1MJbeo,3325
+axio/permission.py,sha256=_BPXZv8QRljwV5wzpPEptvsbDxDtvoLEn97FFyg7GMw,1313
+axio/selector.py,sha256=ehXIYrZ5rUM-UDO6gctbQG8T2vp7nx6ez043IRPk20g,4243
+axio/stream.py,sha256=xLu_D2NefG6NBg4q5xwVjG7jhS-wZzWlyFcpbqKb-zI,1853
+axio/testing.py,sha256=jjr1dmOu4FYKfFOttFbkJDsCGDJV8dbdDpurehTIkvE,2637
+axio/tool.py,sha256=NTXQ3kz1x8mmKAAOntoF-TPgDZ7icsOJ29yR_pIsP1U,2237
+axio/transport.py,sha256=fWK8bUofM_GCZtfO8Cw5AJHjYsGs1GitYoRCCkGYb-4,1036
+axio/types.py,sha256=n18_fDTbZUI1hp6rjvithcJ_5tumKT6lr5Cm3ua0p-I,642
+axio-0.1.0.dist-info/METADATA,sha256=EVl7pr2sFcPvWlzAqB-OQtNDIeBF0RMuhAbwPAUsPHQ,219
+axio-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+axio-0.1.0.dist-info/licenses/LICENSE,sha256=ddOkAXgIM2QzvUDPydeis3tJXfhHxt1wKnmGjbwPPxQ,1074
+axio-0.1.0.dist-info/RECORD,,

axio-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

axio-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Axio contributors
+
+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.