zonix 0.2.1__py3-none-any.whl

zonix/__init__.py

@@ -0,0 +1,44 @@
+from .events import (
+    ApprovalRequired,
+    ErrorEvent,
+    Event,
+    Finish,
+    ReasoningDelta,
+    TextDelta,
+    TextEnd,
+    TextStart,
+    ToolInputAvailable,
+    ToolInputDelta,
+    ToolInputStart,
+    ToolOutputAvailable,
+)
+from .spec import Agent, agent, router, team, workflow
+from .types import Message, PendingApproval, Route, RunResult, RunState, Span, ToolCall, Usage
+
+__all__ = [
+    "Agent",
+    "ApprovalRequired",
+    "ErrorEvent",
+    "Event",
+    "Finish",
+    "Message",
+    "PendingApproval",
+    "ReasoningDelta",
+    "Route",
+    "RunResult",
+    "RunState",
+    "Span",
+    "TextDelta",
+    "TextEnd",
+    "TextStart",
+    "ToolCall",
+    "ToolInputAvailable",
+    "ToolInputDelta",
+    "ToolInputStart",
+    "ToolOutputAvailable",
+    "Usage",
+    "agent",
+    "router",
+    "team",
+    "workflow",
+]

zonix/engine.py

@@ -0,0 +1,335 @@
+from __future__ import annotations
+
+import asyncio
+import inspect
+import json
+from typing import Any
+
+from pydantic import BaseModel, TypeAdapter
+
+from .events import (
+    ApprovalRequired,
+    TextDelta,
+    TextEnd,
+    TextStart,
+    ToolInputAvailable,
+    ToolInputStart,
+    ToolOutputAvailable,
+)
+from .exceptions import (
+    MaxToolRoundsExceeded,
+    OutputValidationError,
+    RunPaused,
+    ToolApprovalRejected,
+)
+from .hitl import approval_key, pending_from_call
+from .models import ModelRequest, ModelResponse
+from .serialization import to_jsonable
+from .tools import ToolContext, ToolDefinition
+from .types import Message, RunState, ToolCall
+
+
+class RunEngine:
+    def __init__(self, agent: Any) -> None:
+        self.agent = agent
+
+    async def invoke(self, task: Any, st: RunState) -> Any:
+        attempts = self.agent.retry_attempts + 1
+        last_error: Exception | None = None
+        for attempt in range(attempts):
+            try:
+                if self.agent.timeout_seconds is None:
+                    return await self._invoke_once(task, st)
+                return await asyncio.wait_for(
+                    self._invoke_once(task, st),
+                    timeout=self.agent.timeout_seconds,
+                )
+            except self.agent.retry_on as exc:
+                last_error = exc
+                st.trace.record({"attempt": attempt + 1, "error": str(exc)})
+                if attempt + 1 >= attempts:
+                    break
+        if self.agent.fallback_node is not None:
+            return await self.agent.fallback_node.invoke(task, st.scoped(self.agent.fallback_node.name))
+        if last_error is not None:
+            raise last_error
+        raise RuntimeError("Agent invocation failed without an exception.")
+
+    async def _invoke_once(self, task: Any, st: RunState) -> Any:
+        messages = await self._build_messages(task, st)
+        tools = {tool.name: tool for tool in self.agent.tools}
+        tool_rounds = 0
+        repair_rounds = 0
+
+        while tool_rounds <= self.agent.max_tool_rounds:
+            request = ModelRequest(
+                messages=messages,
+                tools=[tool.model_tool_schema() for tool in tools.values()],
+                output_schema=self._output_schema(),
+                output_name=self._output_name(),
+                metadata={"agent": self.agent.name, "role": self.agent.role},
+            )
+            response = await self._call_model(request, st)
+            st.usage += response.usage
+
+            if response.tool_calls:
+                messages.append(
+                    Message(
+                        role="assistant",
+                        content=response.text,
+                        data={
+                            "tool_calls": [
+                                call.model_dump(mode="json") for call in response.tool_calls
+                            ]
+                        },
+                    )
+                )
+                tool_rounds += 1
+                for call in response.tool_calls:
+                    await self._run_tool(call, tools, messages, st)
+                continue
+
+            if response.text:
+                messages.append(Message(role="assistant", content=response.text))
+
+            try:
+                output = self._validate_output(response)
+            except OutputValidationError as exc:
+                if repair_rounds >= self.agent.output_repair_attempts:
+                    raise
+                repair_rounds += 1
+                messages.append(
+                    Message(
+                        role="user",
+                        content=(
+                            "The previous response did not validate. "
+                            "Return a corrected JSON value only, with no Markdown or prose. "
+                            f"Validation error: {exc}"
+                        ),
+                        data={"kind": "output_repair", "attempt": repair_rounds},
+                    )
+                )
+                continue
+            st.messages = messages
+            st.scratch[self.agent.name] = output
+            await self._remember(task, output, st)
+            return output
+
+        raise MaxToolRoundsExceeded(
+            f"Agent {self.agent.name!r} exceeded {self.agent.max_tool_rounds} tool rounds."
+        )
+
+    async def _build_messages(self, task: Any, st: RunState) -> list[Message]:
+        messages: list[Message] = []
+        instructions = await self._instructions(task, st)
+        if instructions:
+            messages.append(Message(role="system", content=instructions))
+        if st.session is not None:
+            history = await st.session.recall(task, st.ctx, memory=self.agent.memory)
+            messages.extend(history)
+        messages.append(Message(role="user", content=str(task)))
+        if st.extra:
+            messages.append(Message(role="user", content=st.extra, data={"kind": "extra"}))
+        return messages
+
+    async def _instructions(self, task: Any, st: RunState) -> str:
+        parts: list[str] = []
+        if self.agent.role:
+            parts.append(f"Role: {self.agent.role}")
+        for prompt in self.agent.prompts:
+            if isinstance(prompt, str):
+                parts.append(prompt)
+                continue
+            value = self._call_prompt(prompt, task, st)
+            if inspect.isawaitable(value):
+                value = await value
+            if value:
+                parts.append(str(value))
+        schema = self._output_schema()
+        if schema is not None:
+            parts.append(
+                "Return exactly one valid JSON value that validates against this output schema. "
+                "Do not include Markdown, prose, comments, or trailing text. "
+                "Use double quotes and valid JSON escaping.\n"
+                + json.dumps(schema, ensure_ascii=False)
+            )
+        return "\n\n".join(parts)
+
+    def _call_prompt(self, prompt: Any, task: Any, st: RunState) -> Any:
+        signature = inspect.signature(prompt)
+        params = list(signature.parameters)
+        if len(params) == 0:
+            return prompt()
+        if len(params) == 1:
+            return prompt(st.ctx)
+        return prompt(st.ctx, task)
+
+    async def _call_model(self, request: ModelRequest, st: RunState) -> ModelResponse:
+        if st.bus._emit is not None:
+            return await self.agent.model.stream_complete(request, st.bus.publish, st.path)
+        response = await self.agent.model.complete(request)
+        await self._emit_response(response, st)
+        return response
+
+    async def _emit_response(self, response: ModelResponse, st: RunState) -> None:
+        if response.text:
+            await st.bus.publish(TextStart(st.path, "text_0"))
+            await st.bus.publish(TextDelta(st.path, "text_0", response.text))
+            await st.bus.publish(TextEnd(st.path, "text_0"))
+        for call in response.tool_calls:
+            await st.bus.publish(ToolInputStart(st.path, call.call_id, call.tool))
+            await st.bus.publish(ToolInputAvailable(st.path, call.call_id, call.tool, call.input))
+
+    async def _run_tool(
+        self,
+        call: ToolCall,
+        tools: dict[str, ToolDefinition],
+        messages: list[Message],
+        st: RunState,
+    ) -> None:
+        if call.tool not in tools:
+            raise KeyError(f"Agent {self.agent.name!r} has no tool named {call.tool!r}.")
+
+        tool = tools[call.tool]
+        parsed_input = tool.parse_input(call.input)
+        normalized_call = ToolCall(call_id=call.call_id, tool=call.tool, input=parsed_input)
+        approved_input = await self._approval_for(tool, normalized_call, st)
+        if approved_input is not None and isinstance(approved_input, dict):
+            parsed_input = tool.parse_input(approved_input)
+
+        ctx = ToolContext(deps=st.ctx, usage=st.usage, state=st, agent=self.agent)
+        output = await tool.invoke(ctx, parsed_input)
+        st.usage.tool_calls += 1
+        await st.bus.publish(ToolOutputAvailable(st.path, call.call_id, to_jsonable(output)))
+        messages.append(
+            Message(
+                role="tool",
+                name=tool.name,
+                tool_call_id=call.call_id,
+                content=json.dumps(to_jsonable(output), ensure_ascii=False),
+            )
+        )
+
+    async def _approval_for(
+        self,
+        tool: ToolDefinition,
+        call: ToolCall,
+        st: RunState,
+    ) -> dict[str, Any] | bool | None:
+        if not tool.approval:
+            return None
+        key = approval_key(call.tool, call.input)
+        decision = st.approvals.get(call.call_id, st.approvals.get(key))
+        if decision is False:
+            raise ToolApprovalRejected(f"Tool call {call.call_id} was rejected.")
+        if decision is True or isinstance(decision, dict):
+            return decision
+
+        pending = pending_from_call(call)
+        await st.bus.publish(ApprovalRequired(st.path, call.call_id, call.tool, call.input))
+        snapshot = {
+            "run_id": st.run_id,
+            "path": st.path,
+            "pending": pending,
+            "messages": messages_dump(st.messages),
+            "scratch": to_jsonable(st.scratch),
+            "trace": to_jsonable(st.trace),
+            "usage": to_jsonable(st.usage),
+        }
+        raise RunPaused(pending=pending, snapshot=snapshot)
+
+    def _output_schema(self) -> dict[str, Any] | None:
+        output_type = self.agent.output_type
+        if output_type is None or output_type is Any:
+            return None
+        return TypeAdapter(output_type).json_schema()
+
+    def _output_name(self) -> str | None:
+        output_type = self.agent.output_type
+        if output_type is None:
+            return None
+        return getattr(output_type, "__name__", repr(output_type))
+
+    def _validate_output(self, response: ModelResponse) -> Any:
+        output_type = self.agent.output_type
+        raw = response.output if response.output is not None else response.text
+        if output_type is None or output_type is Any:
+            return raw
+        try:
+            if isinstance(raw, output_type):
+                return raw
+        except TypeError:
+            pass
+        data = raw
+        if isinstance(raw, str) and output_type is not str:
+            try:
+                data = _load_json_text(raw)
+            except json.JSONDecodeError as exc:
+                raise OutputValidationError(
+                    f"Agent {self.agent.name!r} expected {self._output_name()} JSON, got text."
+                ) from exc
+        try:
+            if isinstance(output_type, type) and issubclass(output_type, BaseModel):
+                return output_type.model_validate(data)
+        except TypeError:
+            pass
+        try:
+            return TypeAdapter(output_type).validate_python(data)
+        except Exception as exc:
+            raise OutputValidationError(
+                f"Agent {self.agent.name!r} could not validate model output as {self._output_name()}."
+            ) from exc
+
+    async def _remember(self, task: Any, output: Any, st: RunState) -> None:
+        if st.session is None:
+            return
+        await st.session.remember(Message(role="user", content=str(task)))
+        await st.session.remember(
+            Message(role="assistant", content=json.dumps(to_jsonable(output), ensure_ascii=False))
+        )
+
+
+def messages_dump(messages: list[Message]) -> list[dict[str, Any]]:
+    return [message.model_dump(mode="json") for message in messages]
+
+
+def _load_json_text(text: str) -> Any:
+    stripped = text.strip()
+    try:
+        return json.loads(stripped)
+    except json.JSONDecodeError:
+        candidate = _extract_json_candidate(stripped)
+        if candidate is None:
+            raise
+        return json.loads(candidate)
+
+
+def _extract_json_candidate(text: str) -> str | None:
+    starts = [index for index in (text.find("{"), text.find("[")) if index >= 0]
+    if not starts:
+        return None
+    start = min(starts)
+    opener = text[start]
+    closer = "}" if opener == "{" else "]"
+    depth = 0
+    in_string = False
+    escaped = False
+
+    for index, char in enumerate(text[start:], start=start):
+        if in_string:
+            if escaped:
+                escaped = False
+            elif char == "\\":
+                escaped = True
+            elif char == '"':
+                in_string = False
+            continue
+        if char == '"':
+            in_string = True
+        elif char == opener:
+            depth += 1
+        elif char == closer:
+            depth -= 1
+            if depth == 0:
+                return text[start : index + 1]
+    return None

zonix/events.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from .serialization import to_jsonable
+from .types import Usage
+
+
+@dataclass(frozen=True)
+class Event:
+    path: tuple[str, ...]
+
+    @property
+    def type(self) -> str:
+        name = type(self).__name__
+        chars: list[str] = []
+        for i, char in enumerate(name):
+            if char.isupper() and i:
+                chars.append("-")
+            chars.append(char.lower())
+        return "".join(chars)
+
+    def dump(self) -> dict[str, Any]:
+        data = to_jsonable(self)
+        if isinstance(data, dict):
+            data["type"] = self.type
+            return data
+        return {"type": self.type, "value": data}
+
+
+@dataclass(frozen=True)
+class NodeStart(Event):
+    name: str
+
+
+@dataclass(frozen=True)
+class NodeEnd(Event):
+    name: str
+    status: str = "ok"
+
+
+@dataclass(frozen=True)
+class TextStart(Event):
+    id: str
+
+
+@dataclass(frozen=True)
+class TextDelta(Event):
+    id: str
+    delta: str
+
+
+@dataclass(frozen=True)
+class TextEnd(Event):
+    id: str
+
+
+@dataclass(frozen=True)
+class ReasoningDelta(Event):
+    id: str
+    delta: str
+
+
+@dataclass(frozen=True)
+class ToolInputStart(Event):
+    call_id: str
+    tool: str
+
+
+@dataclass(frozen=True)
+class ToolInputDelta(Event):
+    call_id: str
+    delta: str
+
+
+@dataclass(frozen=True)
+class ToolInputAvailable(Event):
+    call_id: str
+    tool: str
+    input: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class ToolOutputAvailable(Event):
+    call_id: str
+    output: Any
+
+
+@dataclass(frozen=True)
+class ApprovalRequired(Event):
+    call_id: str
+    tool: str
+    input: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class ErrorEvent(Event):
+    message: str
+    error_type: str = "error"
+
+
+@dataclass(frozen=True)
+class Finish(Event):
+    output: Any
+    usage: Usage

zonix/exceptions.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from .types import PendingApproval
+
+
+class ZonixError(Exception):
+    """Base error for Zonix."""
+
+
+class ModelError(ZonixError):
+    """Raised when a model adapter cannot complete a request."""
+
+
+class ToolError(ZonixError):
+    """Raised when a tool fails."""
+
+
+class OutputValidationError(ZonixError):
+    """Raised when a model response cannot be validated as the requested output."""
+
+
+class MaxToolRoundsExceeded(ZonixError):
+    """Raised when an agent keeps asking for tools beyond its configured limit."""
+
+
+class MaxStepsExceeded(ZonixError):
+    """Raised when a team router exceeds its step budget."""
+
+
+class ToolApprovalRejected(ZonixError):
+    """Raised when a paused tool call is resumed with approval=False."""
+
+
+@dataclass
+class RunPaused(ZonixError):
+    pending: PendingApproval
+    snapshot: dict[str, Any]
+
+    def __str__(self) -> str:
+        return f"Run paused for approval: {self.pending.tool}"

zonix/hitl.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from .serialization import to_jsonable
+from .types import PendingApproval, ToolCall
+
+
+def approval_key(tool: str, input: dict[str, Any]) -> str:
+    payload = json.dumps(to_jsonable(input), ensure_ascii=False, sort_keys=True)
+    digest = hashlib.sha256(payload.encode("utf-8")).hexdigest()
+    return f"{tool}:{digest}"
+
+
+def pending_from_call(call: ToolCall) -> PendingApproval:
+    return PendingApproval(
+        call_id=call.call_id,
+        tool=call.tool,
+        input=call.input,
+        approval_key=approval_key(call.tool, call.input),
+    )
+
+
+@dataclass
+class CheckpointStore:
+    root: Path
+
+    def __post_init__(self) -> None:
+        self.root.mkdir(parents=True, exist_ok=True)
+
+    def path_for(self, run_id: str) -> Path:
+        return self.root / f"{run_id}.json"
+
+    def save(self, run_id: str, snapshot: dict[str, Any]) -> Path:
+        path = self.path_for(run_id)
+        path.write_text(
+            json.dumps(to_jsonable(snapshot), ensure_ascii=False, indent=2),
+            encoding="utf-8",
+        )
+        return path
+
+    def load(self, run_id: str) -> dict[str, Any]:
+        return json.loads(self.path_for(run_id).read_text(encoding="utf-8"))

zonix/memory/__init__.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import inspect
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+
+from zonix.types import Message
+
+
+class Memory(Protocol):
+    async def apply(self, history: list[Message], current: Any, ctx: Any) -> list[Message]:
+        ...
+
+
+def _as_list(memory: Any) -> list[Memory]:
+    if memory is None:
+        return []
+    if isinstance(memory, MemoryStack):
+        return list(memory.items)
+    if isinstance(memory, list | tuple):
+        return list(memory)
+    return [memory]
+
+
+@dataclass
+class MemoryStack:
+    items: list[Memory] = field(default_factory=list)
+
+    def __add__(self, other: Any) -> MemoryStack:
+        return MemoryStack([*self.items, *_as_list(other)])
+
+    async def apply(self, history: list[Message], current: Any, ctx: Any) -> list[Message]:
+        messages = list(history)
+        for memory in self.items:
+            messages = await memory.apply(messages, current, ctx)
+        return messages
+
+
+@dataclass
+class Window:
+    size: int = 20
+
+    def __add__(self, other: Any) -> MemoryStack:
+        return MemoryStack([self]) + other
+
+    async def apply(self, history: list[Message], current: Any, ctx: Any) -> list[Message]:
+        return list(history[-self.size :])
+
+
+@dataclass
+class Summarize:
+    over: int
+    keep: int = 20
+
+    def __add__(self, other: Any) -> MemoryStack:
+        return MemoryStack([self]) + other
+
+    async def apply(self, history: list[Message], current: Any, ctx: Any) -> list[Message]:
+        total_chars = sum(len(message.content or "") for message in history)
+        if total_chars <= self.over:
+            return list(history)
+        older = history[: -self.keep]
+        recent = history[-self.keep :]
+        summary = "Earlier conversation summary: " + " ".join(
+            (message.content or "")[:200] for message in older
+        )
+        return [Message(role="system", content=summary)] + recent
+
+
+@dataclass
+class Vector:
+    store: Any
+    k: int = 6
+
+    def __add__(self, other: Any) -> MemoryStack:
+        return MemoryStack([self]) + other
+
+    async def apply(self, history: list[Message], current: Any, ctx: Any) -> list[Message]:
+        results = await self._search(current)
+        if not results:
+            return list(history)
+        body = "\n".join(str(item) for item in results)
+        return [*history, Message(role="system", content=f"Relevant long-term memory:\n{body}")]
+
+    async def _search(self, query: Any) -> list[Any]:
+        for name in ("asearch", "asimilarity_search"):
+            method = getattr(self.store, name, None)
+            if method is not None:
+                value = method(query, k=self.k)
+                if inspect.isawaitable(value):
+                    value = await value
+                return list(value)
+        for name in ("search", "similarity_search"):
+            method = getattr(self.store, name, None)
+            if method is not None:
+                return list(method(query, k=self.k))
+        return []
+
+
+@dataclass
+class Session:
+    id: str = field(default_factory=lambda: f"session_{uuid.uuid4().hex}")
+    history: list[Message] = field(default_factory=list)
+    memory: Any = None
+
+    async def recall(self, current: Any, ctx: Any, memory: Any = None) -> list[Message]:
+        strategies = MemoryStack([*_as_list(self.memory), *_as_list(memory)])
+        return await strategies.apply(list(self.history), current, ctx)
+
+    async def remember(self, message: Message) -> None:
+        self.history.append(message)
+
+    def dump(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "history": [message.model_dump(mode="json") for message in self.history],
+        }