chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,20 @@
|
||||
# omnigent-client
|
||||
|
||||
Python client SDK for the [omnigent](https://github.com/omnigent-ai/omnigent)
|
||||
server API.
|
||||
|
||||
`omnigent-client` is a typed client for driving omnigent sessions over the
|
||||
server's HTTP + SSE API — creating sessions, sending turns, and streaming
|
||||
responses. It shares the `StreamEvent` / `SessionStreamEventType` types that the
|
||||
server emits, so streamed envelopes are validated against a single source of
|
||||
truth.
|
||||
|
||||
It is released in lockstep with the core `omnigent` package at a matching
|
||||
version:
|
||||
|
||||
```bash
|
||||
pip install omnigent-client
|
||||
```
|
||||
|
||||
See the [omnigent repository](https://github.com/omnigent-ai/omnigent) for full
|
||||
documentation.
|
||||
@@ -0,0 +1,129 @@
|
||||
"""omnigent client SDK — Python client for the omnigent server API.
|
||||
|
||||
Headless HTTP/SSE client for invoking agents, tracking conversation
|
||||
state, and consuming the response stream as either raw events or
|
||||
semantic blocks. No UI or terminal dependencies — frontends layer
|
||||
on top of this.
|
||||
|
||||
Usage::
|
||||
|
||||
from omnigent_client import OmnigentClient
|
||||
|
||||
async with OmnigentClient(base_url="http://localhost:8080") as client:
|
||||
session = client.session(model="archer")
|
||||
async for event in session.send("hello"):
|
||||
...
|
||||
|
||||
Or consume semantic blocks via :class:`BlockStream`::
|
||||
|
||||
from omnigent_client import BlockStream, pipe, skip_intermediate_ends
|
||||
|
||||
stream = BlockStream()
|
||||
async for block in pipe(
|
||||
stream.stream(session, "hello"),
|
||||
skip_intermediate_ends(),
|
||||
):
|
||||
...
|
||||
"""
|
||||
|
||||
from ._blocks import (
|
||||
AnyBlock,
|
||||
BlockContext,
|
||||
CompactionBlock,
|
||||
ErrorBlock,
|
||||
FileBlock,
|
||||
NativeToolBlock,
|
||||
ReasoningBlock,
|
||||
ReasoningChunk,
|
||||
ReasoningStartBlock,
|
||||
ResponseEndBlock,
|
||||
ResponseStartBlock,
|
||||
RetryBlock,
|
||||
StreamBlock,
|
||||
TextChunk,
|
||||
TextDone,
|
||||
ToolExecution,
|
||||
ToolGroup,
|
||||
ToolResultBlock,
|
||||
)
|
||||
from ._child_status import (
|
||||
TERMINAL_TASK_STATUSES,
|
||||
child_session_busy,
|
||||
child_summary_busy,
|
||||
)
|
||||
from ._client import OmnigentClient
|
||||
from ._errors import OmnigentError, ToolCallDenied
|
||||
from ._events import MCP_ELICITATION_METHOD, ElicitationRequest
|
||||
from ._query import QueryResult, QueryStream
|
||||
from ._server import LocalServer
|
||||
from ._session import Session
|
||||
from ._sessions import SessionsNamespace
|
||||
from ._sessions_chat import SessionsChat, SessionToolCallInfo, ToolCallable
|
||||
from ._stream import BlockStream, format_tool_args_brief
|
||||
from ._tool_handler import (
|
||||
ElicitationRequestCtx,
|
||||
StreamHooks,
|
||||
ToolCallInfo,
|
||||
ToolHandler,
|
||||
)
|
||||
from ._transforms import (
|
||||
merge_text_across_iterations,
|
||||
only_agent,
|
||||
pipe,
|
||||
skip_blocks,
|
||||
skip_intermediate_ends,
|
||||
)
|
||||
from ._types import File
|
||||
from .tools import ToolMetadata, ToolState, tool
|
||||
|
||||
__all__ = [
|
||||
"MCP_ELICITATION_METHOD",
|
||||
"TERMINAL_TASK_STATUSES",
|
||||
"AnyBlock",
|
||||
"BlockContext",
|
||||
"BlockStream",
|
||||
"CompactionBlock",
|
||||
"ElicitationRequest",
|
||||
"ElicitationRequestCtx",
|
||||
"ErrorBlock",
|
||||
"File",
|
||||
"FileBlock",
|
||||
"LocalServer",
|
||||
"NativeToolBlock",
|
||||
"OmnigentClient",
|
||||
"OmnigentError",
|
||||
"QueryResult",
|
||||
"QueryStream",
|
||||
"ReasoningBlock",
|
||||
"ReasoningChunk",
|
||||
"ReasoningStartBlock",
|
||||
"ResponseEndBlock",
|
||||
"ResponseStartBlock",
|
||||
"RetryBlock",
|
||||
"Session",
|
||||
"SessionToolCallInfo",
|
||||
"SessionsChat",
|
||||
"SessionsNamespace",
|
||||
"StreamBlock",
|
||||
"StreamHooks",
|
||||
"TextChunk",
|
||||
"TextDone",
|
||||
"ToolCallDenied",
|
||||
"ToolCallInfo",
|
||||
"ToolCallable",
|
||||
"ToolExecution",
|
||||
"ToolGroup",
|
||||
"ToolHandler",
|
||||
"ToolMetadata",
|
||||
"ToolResultBlock",
|
||||
"ToolState",
|
||||
"child_session_busy",
|
||||
"child_summary_busy",
|
||||
"format_tool_args_brief",
|
||||
"merge_text_across_iterations",
|
||||
"only_agent",
|
||||
"pipe",
|
||||
"skip_blocks",
|
||||
"skip_intermediate_ends",
|
||||
"tool",
|
||||
]
|
||||
@@ -0,0 +1,282 @@
|
||||
"""Stream block types with context.
|
||||
|
||||
Every block carries a ``BlockContext`` describing which agent produced
|
||||
it, at what depth, and in which turn. The simple case ignores context.
|
||||
Multi-agent frontends route by ``block.ctx.agent``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from ._types import Response
|
||||
|
||||
|
||||
@dataclass
|
||||
class BlockContext:
|
||||
"""Metadata attached to every stream block.
|
||||
|
||||
:param agent: Name of the agent that produced this block, e.g.
|
||||
``"coder.researcher"``. ``None`` for the root agent.
|
||||
:param depth: Nesting depth of the agent in the sub-agent tree.
|
||||
``0`` for the root agent.
|
||||
:param turn: Turn number within the current response.
|
||||
:param timestamp: Monotonic timestamp when the block was created.
|
||||
"""
|
||||
|
||||
agent: str | None = None
|
||||
depth: int = 0
|
||||
turn: int = 0
|
||||
timestamp: float = field(default_factory=time.monotonic)
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class StreamBlock:
|
||||
"""Base for all stream blocks."""
|
||||
|
||||
ctx: BlockContext = field(default_factory=BlockContext)
|
||||
|
||||
|
||||
# ── Response lifecycle ───────────────────────────────────
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ResponseStartBlock(StreamBlock):
|
||||
"""The response has started.
|
||||
|
||||
:param model: Agent model name, e.g. ``"coder"``.
|
||||
:param response_id: Server-assigned response ID.
|
||||
"""
|
||||
|
||||
model: str
|
||||
response_id: str
|
||||
|
||||
|
||||
# ── Tool calls ───────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ToolExecution:
|
||||
"""A single tool call paired with its result.
|
||||
|
||||
:param name: Tool name, e.g. ``"Read"``.
|
||||
:param arguments: Parsed arguments dict.
|
||||
:param args_summary: One-line summary of the arguments, e.g. ``"test.py"``.
|
||||
:param call_id: Server-assigned call ID.
|
||||
:param agent_name: Name of the agent that invoked the tool.
|
||||
:param executed_by: ``"server"`` or ``"client"``.
|
||||
:param output: Tool output text, or ``None`` if not yet available.
|
||||
"""
|
||||
|
||||
name: str
|
||||
arguments: dict[str, object] = field(default_factory=dict)
|
||||
args_summary: str
|
||||
call_id: str
|
||||
agent_name: str
|
||||
executed_by: str = "server"
|
||||
output: str | None = None
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ToolGroup(StreamBlock):
|
||||
"""A batch of tool calls from one iteration.
|
||||
|
||||
:param executions: The tool calls in this group.
|
||||
:param iteration: The iteration number within the response.
|
||||
"""
|
||||
|
||||
executions: list[ToolExecution] = field(default_factory=list)
|
||||
iteration: int = 0
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ToolResultBlock(StreamBlock):
|
||||
"""A tool result, emitted after the tool executes.
|
||||
|
||||
:param name: Tool name, e.g. ``"Read"``.
|
||||
:param call_id: Server-assigned call ID.
|
||||
:param agent_name: Name of the agent that invoked the tool.
|
||||
:param output: Tool output text.
|
||||
:param arguments: Parsed arguments from the matching tool call,
|
||||
retained so result-only renderers can use call metadata.
|
||||
:param args_summary: One-line summary of the matching tool call
|
||||
arguments.
|
||||
"""
|
||||
|
||||
name: str
|
||||
call_id: str
|
||||
agent_name: str
|
||||
output: str
|
||||
arguments: dict[str, object] = field(default_factory=dict)
|
||||
args_summary: str = ""
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class NativeToolBlock(StreamBlock):
|
||||
"""A provider-native tool output (web_search, mcp, etc.).
|
||||
|
||||
:param tool_type: Provider tool type, e.g. ``"web_search_call"``.
|
||||
:param label: Human-readable label for display, e.g. ``"search"``.
|
||||
:param data: Raw provider data dict.
|
||||
"""
|
||||
|
||||
tool_type: str
|
||||
label: str
|
||||
data: dict[str, object] = field(default_factory=dict)
|
||||
|
||||
|
||||
# ── Text ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class TextChunk(StreamBlock):
|
||||
"""A flushed chunk of streamed text.
|
||||
|
||||
:param text: The text content of this chunk.
|
||||
"""
|
||||
|
||||
text: str
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class TextDone(StreamBlock):
|
||||
"""Complete text from a text-streaming section.
|
||||
|
||||
:param full_text: The complete accumulated text.
|
||||
:param has_code_blocks: Whether the text contains fenced code blocks.
|
||||
"""
|
||||
|
||||
full_text: str
|
||||
has_code_blocks: bool = False
|
||||
|
||||
|
||||
# ── Reasoning ────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ReasoningStartBlock(StreamBlock):
|
||||
"""Reasoning has started — show a thinking indicator."""
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ReasoningChunk(StreamBlock):
|
||||
"""An incremental reasoning chunk — analog of :class:`TextChunk`.
|
||||
|
||||
Emitted while reasoning is still in progress so renderers can
|
||||
show live progress (e.g. Codex's command/reasoning stream during
|
||||
the long tool-call window). The eventual :class:`ReasoningBlock`
|
||||
is suppressed when any chunks were emitted, to avoid the
|
||||
formatter re-rendering the same text as a summary panel.
|
||||
|
||||
:param text: The incremental reasoning text, e.g. one flushed
|
||||
line of summary tokens.
|
||||
"""
|
||||
|
||||
text: str
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ReasoningBlock(StreamBlock):
|
||||
"""A completed reasoning/thinking block.
|
||||
|
||||
Emitted only when no :class:`ReasoningChunk` was streamed for
|
||||
this reasoning section. Carries the full accumulated reasoning
|
||||
so non-streaming renderers (logs, web UIs that prefer cards)
|
||||
still get a single summary block.
|
||||
|
||||
:param reasoning_text: The raw reasoning text.
|
||||
:param summary_text: A summary of the reasoning.
|
||||
"""
|
||||
|
||||
reasoning_text: str
|
||||
summary_text: str
|
||||
|
||||
|
||||
# ── Status ───────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ErrorBlock(StreamBlock):
|
||||
"""An error during the response.
|
||||
|
||||
:param message: The free-form error message from the server's
|
||||
:class:`ErrorInfo`. May be empty when the server emitted
|
||||
``response.error`` without populating it — renderers should
|
||||
fall back to :attr:`code` in that case so the user sees at
|
||||
least the error classification instead of a blank panel.
|
||||
:param source: Where the error originated, e.g. ``"llm"``.
|
||||
:param code: The machine-readable error code from the server's
|
||||
:class:`ErrorInfo` payload (e.g. ``"llm_auth_failed"``,
|
||||
``"executor_error"``). Empty when the server omitted it.
|
||||
Carried separately from ``message`` so renderers can show
|
||||
a useful label even when the free-form message is blank.
|
||||
"""
|
||||
|
||||
message: str
|
||||
source: str
|
||||
code: str = ""
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class RetryBlock(StreamBlock):
|
||||
"""The server is retrying.
|
||||
|
||||
:param source: What is being retried, e.g. ``"tool"``.
|
||||
:param attempt: Current attempt number.
|
||||
:param max_attempts: Maximum retry attempts.
|
||||
:param delay_seconds: Delay before the next attempt.
|
||||
"""
|
||||
|
||||
source: str
|
||||
attempt: int
|
||||
max_attempts: int
|
||||
delay_seconds: float
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class CompactionBlock(StreamBlock):
|
||||
"""Conversation is being compacted."""
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class FileBlock(StreamBlock):
|
||||
"""A file artifact produced by the agent.
|
||||
|
||||
:param file_id: Server-assigned file ID.
|
||||
:param filename: Original filename, or ``None`` if unknown.
|
||||
"""
|
||||
|
||||
file_id: str
|
||||
filename: str | None = None
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class ResponseEndBlock(StreamBlock):
|
||||
"""The response reached a terminal state.
|
||||
|
||||
:param status: Terminal status, e.g. ``"completed"`` or ``"failed"``.
|
||||
:param response: The full response object, or ``None``.
|
||||
"""
|
||||
|
||||
status: str
|
||||
response: Response | None = None
|
||||
|
||||
|
||||
# Union of all block types (for type hints).
|
||||
AnyBlock = (
|
||||
ResponseStartBlock
|
||||
| ToolGroup
|
||||
| ToolResultBlock
|
||||
| NativeToolBlock
|
||||
| TextChunk
|
||||
| TextDone
|
||||
| ReasoningStartBlock
|
||||
| ReasoningChunk
|
||||
| ReasoningBlock
|
||||
| ErrorBlock
|
||||
| RetryBlock
|
||||
| CompactionBlock
|
||||
| FileBlock
|
||||
| ResponseEndBlock
|
||||
)
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Canonical sub-agent "busy" predicate — the single source of truth shared
|
||||
by the CLI REPL and SDK drivers.
|
||||
|
||||
A parent session's own ``status`` is per-session: a delegating agent fans out
|
||||
sub-agents and returns to its own prompt, so the parent reads ``idle`` while
|
||||
its children are still working. To answer *"is anything in this subtree still
|
||||
working?"* both the CLI (the ``state: N agents running`` badge / ``↓`` menu in
|
||||
:mod:`omnigent_ui_sdk.terminal`) and SDK rollups
|
||||
(:meth:`SessionsNamespace.subtree_busy` / :meth:`SessionsChat.tree_busy`) need
|
||||
one agreed definition of a *single* child being busy. That definition lives
|
||||
here so the two can't drift.
|
||||
|
||||
The predicate mirrors the web ``SubagentsPanel`` ``childStatus`` semantics
|
||||
(``web/src/shell/SubagentsPanel.tsx``): awaiting input outranks everything
|
||||
(a sub-agent parked on an elicitation is still mid-turn), then ``launching``,
|
||||
then the live ``busy`` flag, then any non-terminal ``current_task_status``.
|
||||
|
||||
This is deliberately stateless — no linger/debounce. The terminal host wraps it
|
||||
with a UI-only linger window to stop the badge flickering between a child's
|
||||
turns; an SDK eval driver wants the un-debounced ground truth instead.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
# Task statuses that mean the child's latest run has settled. Shared with the
|
||||
# terminal host so the CLI's terminal set and this predicate stay in lockstep.
|
||||
TERMINAL_TASK_STATUSES: tuple[str, ...] = ("completed", "failed", "cancelled")
|
||||
|
||||
|
||||
def child_session_busy(
|
||||
*,
|
||||
busy: bool,
|
||||
current_task_status: str | None,
|
||||
pending_elicitations_count: int = 0,
|
||||
) -> bool:
|
||||
"""Return whether a single sub-agent counts as still working.
|
||||
|
||||
Mirrors the web ``SubagentsPanel`` / CLI ``N agents running`` semantics:
|
||||
|
||||
* ``pending_elicitations_count > 0`` — parked on an elicitation it must
|
||||
answer first; still mid-turn, so it counts as busy (outranks ``busy``).
|
||||
* ``current_task_status == "launching"`` — spawned, not yet running.
|
||||
* ``busy`` — the live "session loop is running" flag from the server.
|
||||
* any non-terminal ``current_task_status`` (e.g. ``in_progress`` /
|
||||
``queued``) — a busy-flag gap between turns still counts.
|
||||
|
||||
A warm-idle child (loop idle, no in-progress task, no pending input) and a
|
||||
terminal one (``completed`` / ``failed`` / ``cancelled``) return ``False``.
|
||||
|
||||
:param busy: ``ChildSessionSummary.busy`` — server status in
|
||||
``("running", "waiting")``.
|
||||
:param current_task_status: ``ChildSessionSummary.current_task_status`` —
|
||||
``launching`` / ``in_progress`` / ``completed`` / ``failed`` /
|
||||
``cancelled`` / ``None``.
|
||||
:param pending_elicitations_count: number of approvals/prompts blocking the
|
||||
child.
|
||||
:returns: ``True`` while the child is still working.
|
||||
"""
|
||||
if pending_elicitations_count > 0 or busy or current_task_status == "launching":
|
||||
return True
|
||||
return current_task_status is not None and current_task_status not in TERMINAL_TASK_STATUSES
|
||||
|
||||
|
||||
def child_summary_busy(summary: Mapping[str, Any]) -> bool:
|
||||
""":func:`child_session_busy` applied to a raw ``ChildSessionSummary`` dict
|
||||
as returned by :meth:`SessionsNamespace.child_sessions`.
|
||||
|
||||
Tolerates missing / ``None`` fields (e.g. a poll row with no task yet).
|
||||
"""
|
||||
return child_session_busy(
|
||||
busy=bool(summary.get("busy")),
|
||||
current_task_status=summary.get("current_task_status"),
|
||||
pending_elicitations_count=int(summary.get("pending_elicitations_count") or 0),
|
||||
)
|
||||
@@ -0,0 +1,348 @@
|
||||
"""OmnigentClient — the top-level client tying all namespaces together."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any, Literal, overload
|
||||
|
||||
import httpx
|
||||
|
||||
from omnigent.runner.identity import OMNIGENT_INTERNAL_WS_ORIGIN
|
||||
|
||||
from ._files import FilesNamespace
|
||||
from ._query import QueryResult, QueryStream
|
||||
from ._responses import ResponsesNamespace
|
||||
from ._session import Session
|
||||
from ._sessions import SessionsNamespace
|
||||
from ._sessions_chat import SessionsChat, ToolCallable
|
||||
from ._tool_handler import StreamHooks, ToolHandler
|
||||
|
||||
|
||||
class OmnigentClient:
|
||||
"""Typed Python client for the omnigent server API.
|
||||
|
||||
One-shot::
|
||||
|
||||
async with OmnigentClient(base_url="http://localhost:8080") as client:
|
||||
result = await client.query(model="archer", input="hello")
|
||||
print(result.text) # the assistant's reply
|
||||
print(result.files) # any files the agent produced
|
||||
|
||||
Streaming::
|
||||
|
||||
stream = await client.query(model="archer", input="hi", stream=True)
|
||||
async for chunk in stream:
|
||||
print(chunk, end="", flush=True)
|
||||
print(stream.files) # populated after the stream ends
|
||||
|
||||
Multi-turn conversation::
|
||||
|
||||
session = client.session(model="archer")
|
||||
await session.query("hello")
|
||||
await session.query("what did I just say?")
|
||||
|
||||
For access to raw events or semantic blocks (tool-call display,
|
||||
reasoning, lifecycle), drop to :attr:`responses` or
|
||||
:class:`BlockStream`.
|
||||
|
||||
:param base_url: Server base URL, e.g. ``"http://localhost:8080"``.
|
||||
:param headers: Extra headers sent on every request (e.g. auth).
|
||||
:param auth: Optional ``httpx.Auth`` for per-request
|
||||
authentication. When set, the auth handler runs on every
|
||||
request, allowing transparent token refresh for OAuth
|
||||
flows. ``None`` (default) relies on static ``headers``.
|
||||
:param timeout: Default timeout for HTTP requests in seconds.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
*,
|
||||
headers: dict[str, str] | None = None,
|
||||
auth: httpx.Auth | None = None,
|
||||
# Public SDK API surface. Removing would be a breaking
|
||||
# change for downstream consumers; leaving it
|
||||
# accepted-but-ignored preserves compatibility while the
|
||||
# SSE client internally uses a fixed 600s read timeout
|
||||
# (tool calls can legitimately hold the stream open for
|
||||
# minutes).
|
||||
timeout: float = 30.0,
|
||||
) -> None:
|
||||
self._base_url = base_url.rstrip("/")
|
||||
# Long read timeout for SSE streams (tool execution can
|
||||
# pause the stream for minutes).
|
||||
sse_timeout = httpx.Timeout(
|
||||
connect=30.0,
|
||||
read=600.0,
|
||||
write=30.0,
|
||||
pool=30.0,
|
||||
)
|
||||
# Announce this as a first-party non-browser client via the sentinel
|
||||
# Origin. The server's require_trusted_origin CSRF guard on the
|
||||
# multipart routes (POST /v1/sessions bundle create, file upload)
|
||||
# requires a trusted Origin; the SDK sends none of its own, so the
|
||||
# sentinel is what lets it through. Caller-supplied headers win on
|
||||
# conflict (so an explicit Origin override is still honored).
|
||||
default_headers = {"Origin": OMNIGENT_INTERNAL_WS_ORIGIN}
|
||||
if headers:
|
||||
default_headers.update(headers)
|
||||
self._http = httpx.AsyncClient(
|
||||
headers=default_headers,
|
||||
auth=auth,
|
||||
timeout=sse_timeout,
|
||||
)
|
||||
|
||||
self.sessions = SessionsNamespace(self._http, self._base_url)
|
||||
self.files = FilesNamespace(self._http, self._base_url)
|
||||
self.responses = ResponsesNamespace(self._http, self._base_url)
|
||||
|
||||
def session(
|
||||
self,
|
||||
model: str,
|
||||
*,
|
||||
tool_handler: ToolHandler | None = None,
|
||||
hooks: StreamHooks | None = None,
|
||||
) -> Session:
|
||||
"""Create a conversation session.
|
||||
|
||||
A session tracks ``previous_response_id`` automatically.
|
||||
``send()`` auto-steers if a response is in progress, or
|
||||
starts a new turn if the response is terminal.
|
||||
|
||||
:param model: Agent name.
|
||||
:param tool_handler: Optional client-side tool execution config.
|
||||
:param hooks: Optional lifecycle hooks.
|
||||
:returns: A new :class:`Session`.
|
||||
"""
|
||||
return Session(
|
||||
client=self,
|
||||
model=model,
|
||||
tool_handler=tool_handler,
|
||||
hooks=hooks,
|
||||
)
|
||||
|
||||
@overload
|
||||
async def query(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
input: str | list[dict[str, object]],
|
||||
tools: list[Callable[..., Any]] | None = ...,
|
||||
tool_handler: ToolHandler | None = ...,
|
||||
files: list[str] | None = ...,
|
||||
reasoning: dict[str, str] | None = ...,
|
||||
model_override: str | None = ...,
|
||||
stream: Literal[False] = ...,
|
||||
) -> QueryResult: ...
|
||||
|
||||
@overload
|
||||
async def query(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
input: str | list[dict[str, object]],
|
||||
tools: list[Callable[..., Any]] | None = ...,
|
||||
tool_handler: ToolHandler | None = ...,
|
||||
files: list[str] | None = ...,
|
||||
reasoning: dict[str, str] | None = ...,
|
||||
model_override: str | None = ...,
|
||||
stream: Literal[True],
|
||||
) -> QueryStream: ...
|
||||
|
||||
async def query(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
input: str | list[dict[str, object]],
|
||||
tools: list[Callable[..., Any]] | None = None,
|
||||
tool_handler: ToolHandler | None = None,
|
||||
files: list[str] | None = None,
|
||||
reasoning: dict[str, str] | None = None,
|
||||
model_override: str | None = None,
|
||||
stream: bool = False,
|
||||
) -> QueryResult | QueryStream:
|
||||
"""One-shot invocation: send a prompt, get text (plus any files) back.
|
||||
|
||||
Non-streaming (default) returns a :class:`QueryResult`::
|
||||
|
||||
result = await client.query(model="archer", input="hi")
|
||||
print(result.text)
|
||||
for f in result.files:
|
||||
await client.files.for_session("<session-id>").download(
|
||||
f.id, f"./out/{f.filename}"
|
||||
)
|
||||
|
||||
Streaming returns a :class:`QueryStream`::
|
||||
|
||||
stream = await client.query(model="archer", input="hi", stream=True)
|
||||
async for chunk in stream:
|
||||
print(chunk, end="", flush=True)
|
||||
# After iteration, stream.files holds the produced files.
|
||||
|
||||
With client-side tools, pass ``@tool``-decorated functions::
|
||||
|
||||
from omnigent_client import tool
|
||||
|
||||
@tool
|
||||
def get_time() -> str:
|
||||
'''Return the current time.'''
|
||||
return datetime.now().isoformat()
|
||||
|
||||
result = await client.query(
|
||||
model="archer", input="what time?", tools=[get_time],
|
||||
)
|
||||
|
||||
Creates a single-turn session internally. For multi-turn
|
||||
conversations, call :meth:`session` and use its ``query()``.
|
||||
|
||||
:param model: Agent name, e.g. ``"archer"``.
|
||||
:param input: User text or a list of content-block dicts.
|
||||
:param tools: List of ``@tool``-decorated Python functions
|
||||
the agent may call. Mutually exclusive with ``tool_handler``.
|
||||
:param tool_handler: Low-level escape hatch — a pre-built
|
||||
:class:`ToolHandler` with custom schemas/dispatch. Most
|
||||
callers should use ``tools=`` instead.
|
||||
:param files: Optional list of local file paths to attach.
|
||||
:param reasoning: Optional Responses API reasoning config, e.g. {"effort": "high"}.
|
||||
:param model_override: Optional per-request LLM model override
|
||||
(e.g. ``"openai/gpt-5.4-mini"``). Shadows the spec model
|
||||
for this one-shot call; mirrors
|
||||
:meth:`Session.set_model_override`.
|
||||
:param stream: If True, return a :class:`QueryStream`. If
|
||||
False (default), return a :class:`QueryResult`.
|
||||
:returns: :class:`QueryResult` (``stream=False``) or
|
||||
:class:`QueryStream` (``stream=True``).
|
||||
:raises ValueError: If both ``tools`` and ``tool_handler``
|
||||
are provided.
|
||||
"""
|
||||
handler = _resolve_tool_handler(tools=tools, tool_handler=tool_handler)
|
||||
session = self.session(model=model, tool_handler=handler)
|
||||
effort = reasoning.get("effort") if reasoning is not None else None
|
||||
if effort is not None:
|
||||
session.set_reasoning_effort(effort)
|
||||
if model_override is not None:
|
||||
session.set_model_override(model_override)
|
||||
if stream:
|
||||
return await session.query(input, files=files, stream=True)
|
||||
return await session.query(input, files=files)
|
||||
|
||||
async def sessions_chat(
|
||||
self,
|
||||
bundle: bytes,
|
||||
*,
|
||||
filename: str = "agent.tar.gz",
|
||||
tool_callables: dict[str, ToolCallable] | None = None,
|
||||
hooks: StreamHooks | None = None,
|
||||
) -> SessionsChat:
|
||||
"""Create a sessions-API-native chat helper bound to a new session.
|
||||
|
||||
Counterpart to :meth:`session` but built on ``/v1/sessions``
|
||||
rather than ``/v1/responses``. Use this for new code; the
|
||||
legacy :meth:`session` is preserved for in-flight migrations.
|
||||
|
||||
:param bundle: Gzipped agent tarball bytes uploaded through
|
||||
multipart ``POST /v1/sessions``.
|
||||
:param filename: Filename for the multipart upload, e.g.
|
||||
``"agent.tar.gz"``.
|
||||
:param tool_callables: Optional mapping from tool name to
|
||||
an executable callable (sync or async) for client-side
|
||||
tool execution. Validated against the agent's
|
||||
spec-declared tools at stream-start time (the first
|
||||
``send()`` / ``query()`` / ``stream()`` call), not at
|
||||
construction. See :class:`SessionsChat` for the
|
||||
validation rules.
|
||||
:param hooks: Optional lifecycle hooks fired from sessions
|
||||
stream events.
|
||||
:returns: A :class:`SessionsChat` ready for use.
|
||||
:raises OmnigentError: If session creation fails.
|
||||
"""
|
||||
return await SessionsChat.create(
|
||||
namespace=self.sessions,
|
||||
bundle=bundle,
|
||||
filename=filename,
|
||||
files_namespace=self.files,
|
||||
tool_callables=tool_callables,
|
||||
agent_tools_getter=self._fetch_agent_tools,
|
||||
hooks=hooks,
|
||||
)
|
||||
|
||||
async def _fetch_agent_tools(
|
||||
self, agent_id: str, session_id: str | None = None
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Fetch the spec-declared tool entries for an agent.
|
||||
|
||||
Used as the ``agent_tools_getter`` injection for
|
||||
:class:`SessionsChat`. Reads the tool list off the
|
||||
:class:`Agent` returned by ``GET /api/agents/{agent_id}``.
|
||||
|
||||
The server's :class:`AgentObject`
|
||||
carries a ``tools`` list where each entry has a ``name``
|
||||
and a ``runtime`` discriminator
|
||||
(``"server"`` or ``"client"``). When that field is not yet
|
||||
present, the SDK's :class:`Agent` dataclass simply lacks
|
||||
the field and this returns ``[]`` — which means
|
||||
validation succeeds for any caller that doesn't pass
|
||||
``tool_callables``, and fails loud (with a clear "extra
|
||||
callable" message) for any caller that does. That is the
|
||||
correct degraded behavior: in F1's absence we cannot
|
||||
verify the spec, but we will never silently accept a
|
||||
broken setup.
|
||||
|
||||
:param agent_id: The agent's durable identifier, e.g.
|
||||
``"ag_abc123"``.
|
||||
:returns: List of tool-entry dicts with at least ``name``
|
||||
and (post-F1) ``runtime`` keys. Empty if the agent
|
||||
declares no tools or the server response shape
|
||||
predates F1.
|
||||
:raises OmnigentError: If the agents endpoint returns
|
||||
a non-2xx (e.g. 404).
|
||||
"""
|
||||
if session_id is None:
|
||||
return [] # No session context — cannot resolve agent tools
|
||||
path = f"{self._base_url}/v1/sessions/{session_id}/agent"
|
||||
resp = await self._http.get(path)
|
||||
if resp.status_code != 200:
|
||||
return []
|
||||
agent_data = resp.json()
|
||||
tools = agent_data.get("tools")
|
||||
if isinstance(tools, list):
|
||||
return [t for t in tools if isinstance(t, dict)]
|
||||
return []
|
||||
|
||||
async def close(self) -> None:
|
||||
"""Close the underlying HTTP client."""
|
||||
await self._http.aclose()
|
||||
|
||||
async def __aenter__(self) -> OmnigentClient:
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc: object) -> None:
|
||||
await self.close()
|
||||
|
||||
|
||||
def _resolve_tool_handler(
|
||||
*,
|
||||
tools: list[Callable[..., Any]] | None,
|
||||
tool_handler: ToolHandler | None,
|
||||
) -> ToolHandler | None:
|
||||
"""Pick one of ``tools=`` or ``tool_handler=``; reject both.
|
||||
|
||||
:param tools: High-level list of ``@tool``-decorated functions.
|
||||
:param tool_handler: Low-level pre-built handler.
|
||||
:returns: The handler to use, or ``None`` if neither was given.
|
||||
:raises ValueError: If both were provided.
|
||||
"""
|
||||
if tools is not None and tool_handler is not None:
|
||||
raise ValueError(
|
||||
"Pass either `tools=[...]` or `tool_handler=...`, not both. "
|
||||
"`tools=` is the high-level API (auto-builds a handler from "
|
||||
"@tool-decorated functions); `tool_handler=` is the low-level "
|
||||
"escape hatch."
|
||||
)
|
||||
if tools is not None:
|
||||
# Local import keeps the dep inside the tools subpackage.
|
||||
from .tools import build_tool_handler
|
||||
|
||||
return build_tool_handler(tools)
|
||||
return tool_handler
|
||||
@@ -0,0 +1,180 @@
|
||||
"""Typed exceptions for the omnigent client."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
|
||||
_BODY_PREVIEW_CHARS = 200
|
||||
|
||||
|
||||
class OmnigentError(Exception):
|
||||
"""Base exception for all omnigent client errors."""
|
||||
|
||||
def __init__(self, message: str, status_code: int | None = None, code: str | None = None):
|
||||
super().__init__(message)
|
||||
self.status_code = status_code
|
||||
self.code = code
|
||||
|
||||
|
||||
class AgentNotFoundError(OmnigentError):
|
||||
"""Agent name or ID not found (HTTP 404)."""
|
||||
|
||||
|
||||
class ResponseNotFoundError(OmnigentError):
|
||||
"""Response ID not found (HTTP 404)."""
|
||||
|
||||
|
||||
class FileNotFoundError(OmnigentError):
|
||||
"""File ID not found (HTTP 404)."""
|
||||
|
||||
|
||||
class ConversationNotFoundError(OmnigentError):
|
||||
"""Conversation ID not found (HTTP 404)."""
|
||||
|
||||
|
||||
class InvalidInputError(OmnigentError):
|
||||
"""Bad request — invalid input, missing fields, etc. (HTTP 400)."""
|
||||
|
||||
|
||||
class ConflictError(OmnigentError):
|
||||
"""Resource conflict — duplicate name, stale state, etc. (HTTP 409)."""
|
||||
|
||||
|
||||
class BundleInvalidError(OmnigentError):
|
||||
"""Agent bundle is invalid — corrupt tarball, bad config, etc. (HTTP 400)."""
|
||||
|
||||
|
||||
class ServerError(OmnigentError):
|
||||
"""Internal server error (HTTP 5xx)."""
|
||||
|
||||
|
||||
class ToolCallDenied(Exception):
|
||||
"""Raised by ``on_tool_call_start`` hook to deny a client-side tool call.
|
||||
|
||||
The exception message is sent back to the agent as the tool's output,
|
||||
so the agent knows the call was denied and can adapt.
|
||||
"""
|
||||
|
||||
|
||||
def raise_for_status(status_code: int, body: dict[str, object] | str) -> None:
|
||||
"""Raise a typed exception based on HTTP status code and error body.
|
||||
|
||||
Uses the server's error ``code`` field for classification when
|
||||
available, falling back to status code only. Never relies on
|
||||
substring matching in error messages.
|
||||
"""
|
||||
if status_code < 400:
|
||||
return
|
||||
|
||||
if isinstance(body, dict):
|
||||
error = body.get("error", {})
|
||||
if isinstance(error, dict):
|
||||
code = str(error.get("code", ""))
|
||||
message = str(error.get("message", str(body)))
|
||||
else:
|
||||
code = ""
|
||||
message = str(body)
|
||||
else:
|
||||
code = ""
|
||||
message = str(body)
|
||||
|
||||
# Use the server's error code for precise classification.
|
||||
_CODE_MAP: dict[str, type[OmnigentError]] = {
|
||||
"not_found": OmnigentError,
|
||||
"invalid_input": InvalidInputError,
|
||||
"conflict": ConflictError,
|
||||
"server_error": ServerError,
|
||||
}
|
||||
|
||||
if status_code == 404:
|
||||
# 404 is always "not found" — the specific type (agent, response,
|
||||
# file, conversation) is determined by the endpoint the caller
|
||||
# hit, not the error message. Callers can catch the base
|
||||
# OmnigentError or the specific subclass at the call site.
|
||||
raise OmnigentError(message, status_code, code)
|
||||
|
||||
if status_code == 409:
|
||||
raise ConflictError(message, status_code, code)
|
||||
|
||||
if status_code == 400:
|
||||
raise InvalidInputError(message, status_code, code)
|
||||
|
||||
if status_code >= 500:
|
||||
raise ServerError(message, status_code, code)
|
||||
|
||||
raise OmnigentError(message, status_code, code)
|
||||
|
||||
|
||||
def response_body(resp: httpx.Response) -> dict[str, object] | str:
|
||||
"""
|
||||
Decode an HTTP response body for status handling.
|
||||
|
||||
Server errors and proxy/auth failures are not guaranteed to
|
||||
return JSON. This helper preserves structured JSON error
|
||||
envelopes when available and falls back to raw response text
|
||||
otherwise.
|
||||
|
||||
:param resp: HTTP response returned by httpx.
|
||||
:returns: Parsed JSON object, or raw text when the body is not
|
||||
a JSON object.
|
||||
"""
|
||||
try:
|
||||
data = resp.json()
|
||||
except ValueError:
|
||||
return resp.text
|
||||
if isinstance(data, dict):
|
||||
return data
|
||||
return resp.text
|
||||
|
||||
|
||||
def require_json_object(resp: httpx.Response, endpoint: str) -> dict[str, object]:
|
||||
"""
|
||||
Return a successful endpoint body as a JSON object.
|
||||
|
||||
Used after :func:`raise_for_status` when an endpoint's success
|
||||
contract is a JSON object. A non-JSON response at this point is a
|
||||
protocol/base-URL/auth problem, so raise :class:`OmnigentError`
|
||||
with status, content type, and a short body preview instead of
|
||||
leaking ``JSONDecodeError``.
|
||||
|
||||
:param resp: HTTP response returned by httpx.
|
||||
:param endpoint: Human-readable endpoint label, e.g.
|
||||
``"GET /v1/conversations"``.
|
||||
:returns: Parsed JSON object body.
|
||||
:raises OmnigentError: If the body is not a JSON object.
|
||||
"""
|
||||
try:
|
||||
data = resp.json()
|
||||
except ValueError as exc:
|
||||
content_type = resp.headers.get("content-type")
|
||||
content_type_detail = (
|
||||
f"content-type={content_type}"
|
||||
if content_type is not None
|
||||
else "no content-type header"
|
||||
)
|
||||
preview = _response_text_preview(resp.text)
|
||||
raise OmnigentError(
|
||||
f"{endpoint} returned non-JSON response "
|
||||
f"(status={resp.status_code}, {content_type_detail}): {preview}",
|
||||
resp.status_code,
|
||||
) from exc
|
||||
if not isinstance(data, dict):
|
||||
raise OmnigentError(
|
||||
f"{endpoint} returned JSON {type(data).__name__}, expected object",
|
||||
resp.status_code,
|
||||
)
|
||||
return data
|
||||
|
||||
|
||||
def _response_text_preview(text: str) -> str:
|
||||
"""
|
||||
Return a compact, single-line response body preview.
|
||||
|
||||
:param text: Raw response text from HTTPX.
|
||||
:returns: Whitespace-normalized body preview, capped at
|
||||
``_BODY_PREVIEW_CHARS`` characters, or ``"<empty body>"``.
|
||||
"""
|
||||
preview = " ".join(text.split())[:_BODY_PREVIEW_CHARS]
|
||||
if preview:
|
||||
return preview
|
||||
return "<empty body>"
|
||||
@@ -0,0 +1,341 @@
|
||||
"""Typed event dataclasses for SSE stream events.
|
||||
|
||||
The client parses raw SSE frames into these types. Consumers
|
||||
iterate over them via ``async for event in stream``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from ._types import ErrorInfo, Response
|
||||
|
||||
# ── Native tool type constants ───────────────────────────
|
||||
|
||||
NATIVE_TOOL_TYPES: frozenset[str] = frozenset(
|
||||
{
|
||||
"web_search_call",
|
||||
"file_search_call",
|
||||
"code_interpreter_call",
|
||||
"computer_call",
|
||||
"image_generation_call",
|
||||
"mcp_call",
|
||||
"mcp_list_tools",
|
||||
}
|
||||
)
|
||||
|
||||
# JSON-RPC method name MCP uses for elicitation requests. The
|
||||
# server's ``response.elicitation_request`` SSE event carries this
|
||||
# verbatim under ``method`` so MCP-aware consumers can route on
|
||||
# the same name they already recognize.
|
||||
MCP_ELICITATION_METHOD = "elicitation/create"
|
||||
|
||||
|
||||
# ── Response lifecycle events ────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResponseCreated:
|
||||
"""``response.created`` — always first (sequence 0)."""
|
||||
|
||||
response: Response
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResponseQueued:
|
||||
"""``response.queued`` — only when ``background=True``."""
|
||||
|
||||
response: Response
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResponseInProgress:
|
||||
"""``response.in_progress`` — execution started."""
|
||||
|
||||
response: Response
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResponseCompleted:
|
||||
"""``response.completed`` — agent finished successfully."""
|
||||
|
||||
response: Response
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResponseFailed:
|
||||
"""``response.failed`` — unrecoverable error."""
|
||||
|
||||
response: Response
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResponseIncomplete:
|
||||
"""``response.incomplete`` — stopped early."""
|
||||
|
||||
response: Response
|
||||
reason: str # "max_iterations", "execution_timeout", etc.
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResponseCancelled:
|
||||
"""``response.cancelled`` — cancelled via POST /cancel."""
|
||||
|
||||
response: Response
|
||||
|
||||
|
||||
# ── Text streaming ───────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextDelta:
|
||||
"""``response.output_text.delta`` — incremental text token."""
|
||||
|
||||
delta: str
|
||||
|
||||
|
||||
# ── Reasoning ────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReasoningStarted:
|
||||
"""``response.reasoning.started`` — reasoning block opened."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReasoningDelta:
|
||||
"""``response.reasoning_text.delta`` — reasoning token."""
|
||||
|
||||
delta: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReasoningSummaryDelta:
|
||||
"""``response.reasoning_summary_text.delta`` — summary token."""
|
||||
|
||||
delta: str
|
||||
|
||||
|
||||
# ── Parsed output items ─────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
"""A tool call from ``output_item.done`` (type ``function_call``)."""
|
||||
|
||||
name: str
|
||||
arguments: dict[str, object]
|
||||
call_id: str
|
||||
status: str # "completed", "action_required", "incomplete"
|
||||
agent_name: str # "coder" or "coder.researcher"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolResult:
|
||||
"""A tool result from ``output_item.done`` (type ``function_call_output``).
|
||||
|
||||
``arguments`` is optional because the OpenAI-compatible
|
||||
``function_call_output`` item only requires ``call_id`` and ``output``.
|
||||
Omnigent producers may include it as a convenience copy of the
|
||||
originating ``function_call.arguments`` so result-only renderers can use
|
||||
call metadata (for example, rendering ``sys_os_edit`` as a diff).
|
||||
"""
|
||||
|
||||
call_id: str
|
||||
output: str
|
||||
arguments: dict[str, object] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ElicitationRequest:
|
||||
"""
|
||||
A server-initiated elicitation, MCP shape.
|
||||
|
||||
Parsed out of the ``response.elicitation_request`` SSE event;
|
||||
the event's ``params`` block matches MCP's
|
||||
``ElicitRequestFormParams`` field-for-field, plus extras
|
||||
(``phase``, ``policy_name``, ``content_preview``) under MCP's
|
||||
``extra="allow"`` config that today's only producer (the
|
||||
policy ASK flow) populates for renderer use.
|
||||
|
||||
Consumers respond via
|
||||
``POST /v1/sessions/{session_id}/events`` with
|
||||
``type == "approval"`` and MCP-shape ``ElicitResult`` fields
|
||||
(``action`` + optional ``content``) in ``data``. The client
|
||||
handles the POST automatically when an
|
||||
``on_elicitation_request`` hook is registered on
|
||||
:class:`StreamHooks`.
|
||||
|
||||
:param elicitation_id: Server-assigned id. Used in the
|
||||
approval event payload, e.g. ``"elicit_abc123"``.
|
||||
:param message: Human-readable prompt the consumer should
|
||||
render to the user. For the policy ASK producer this is
|
||||
the combined reason string from deciding ASKing policies.
|
||||
:param requested_schema: A restricted subset of JSON Schema
|
||||
defining the structure of the expected response. Empty
|
||||
``{}`` for binary approve/reject elicitations (the verdict
|
||||
is in the consumer's ``action``).
|
||||
:param mode: MCP elicitation mode — ``"form"`` (inline) or
|
||||
``"url"`` (standalone approval page). e.g. ``"url"``.
|
||||
:param phase: Producer-supplied extra (policy ASK only):
|
||||
which enforcement point produced the ASK — one of
|
||||
``"input"``, ``"tool_call"``, ``"tool_result"``,
|
||||
``"output"``. Empty string when the elicitation came from
|
||||
a producer that doesn't populate this extra.
|
||||
:param policy_name: Producer-supplied extra (policy ASK
|
||||
only): name of the deciding (first-in-YAML-order) ASKing
|
||||
policy, e.g. ``"approve_web_search"``. Empty string when
|
||||
not applicable.
|
||||
:param content_preview: Producer-supplied extra (policy ASK
|
||||
only): truncated snapshot of the gated content. Safe to
|
||||
display verbatim. Empty string when not applicable.
|
||||
:param target_session_id: Session whose resolve endpoint owns
|
||||
this elicitation, e.g. ``"conv_child123"``. Set when a
|
||||
sub-agent's prompt was mirrored into an ancestor stream so the
|
||||
consumer can route the verdict back to the child that parked on
|
||||
it. ``None`` means resolve against the session the event arrived
|
||||
on.
|
||||
"""
|
||||
|
||||
elicitation_id: str
|
||||
message: str
|
||||
requested_schema: dict[str, object]
|
||||
mode: str
|
||||
phase: str
|
||||
policy_name: str
|
||||
content_preview: str
|
||||
url: str | None = None
|
||||
target_session_id: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class NativeToolCall:
|
||||
"""A provider-native tool output (web_search, mcp, etc.)."""
|
||||
|
||||
tool_type: str # e.g. "web_search_call"
|
||||
data: dict[str, object]
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageDone:
|
||||
"""The final assistant message from ``output_item.done`` (type ``message``)."""
|
||||
|
||||
content: list[dict[str, object]] = field(default_factory=list)
|
||||
|
||||
|
||||
# ── File output ──────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class OutputFileDone:
|
||||
"""``response.output_file.done`` — file artifact produced."""
|
||||
|
||||
file_id: str
|
||||
filename: str | None = None
|
||||
content_type: str | None = None
|
||||
|
||||
|
||||
# ── Error and retry ──────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetryEvent:
|
||||
"""``response.retry`` — a retryable failure, will retry."""
|
||||
|
||||
source: str # "llm" or "tool"
|
||||
tool_name: str | None
|
||||
attempt: int
|
||||
max_attempts: int
|
||||
delay_seconds: float
|
||||
error: ErrorInfo
|
||||
|
||||
|
||||
@dataclass
|
||||
class ErrorEvent:
|
||||
"""``response.error`` — an error during execution."""
|
||||
|
||||
source: str # "llm" or "tool"
|
||||
tool_name: str | None
|
||||
error: ErrorInfo
|
||||
|
||||
|
||||
# ── Compaction ───────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompactionInProgress:
|
||||
"""``response.compaction.in_progress`` — server started compacting."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompactionCompleted:
|
||||
"""``response.compaction.completed`` — compaction finished successfully."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompactionFailed:
|
||||
"""``response.compaction.failed`` — compaction failed; history unchanged."""
|
||||
|
||||
|
||||
# ── Async client-tool cancel (Phase 5) ───────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClientTaskCancel:
|
||||
"""
|
||||
``response.client_task.cancel`` — server-to-client cancel
|
||||
notification for async client-tool dispatches.
|
||||
|
||||
Emitted when an async client tool task
|
||||
(``kind="client_tool"``) was cancelled mid-flight, either
|
||||
via direct ``sys_cancel_task`` or via parent-cancel
|
||||
propagation. The SDK should cancel the matching local
|
||||
``asyncio.Task`` running the tool body. The body's
|
||||
eventual PATCH back is a no-op (G3 first-write-wins), so
|
||||
the cancel is best-effort from the SDK's perspective —
|
||||
the server has already decided the task is cancelled.
|
||||
|
||||
:param task_id: The server-issued client-tool task id from
|
||||
the original ``function_call_output`` handle, e.g.
|
||||
``"resp_async_xyz"``.
|
||||
:param call_id: The synthesized ``function_call.call_id``
|
||||
the SDK saw on the action_required event for this
|
||||
dispatch, e.g. ``"call_async_b2c4..."``. Populated by
|
||||
the server from ``pending_tool_calls`` so the SDK can
|
||||
look up the local ``asyncio.Task`` it spawned. ``None``
|
||||
on legacy emissions where the lookup wasn't done; the
|
||||
SDK falls back to no-op when both fields can't be
|
||||
matched to a tracked task.
|
||||
"""
|
||||
|
||||
task_id: str
|
||||
call_id: str | None = None
|
||||
|
||||
|
||||
# ── Union type for all events ────────────────────────────
|
||||
|
||||
StreamEvent = (
|
||||
ResponseCreated
|
||||
| ResponseQueued
|
||||
| ResponseInProgress
|
||||
| ResponseCompleted
|
||||
| ResponseFailed
|
||||
| ResponseIncomplete
|
||||
| ResponseCancelled
|
||||
| TextDelta
|
||||
| ReasoningStarted
|
||||
| ReasoningDelta
|
||||
| ReasoningSummaryDelta
|
||||
| ToolCall
|
||||
| ToolResult
|
||||
| NativeToolCall
|
||||
| MessageDone
|
||||
| OutputFileDone
|
||||
| RetryEvent
|
||||
| ErrorEvent
|
||||
| CompactionInProgress
|
||||
| CompactionCompleted
|
||||
| CompactionFailed
|
||||
| ClientTaskCancel
|
||||
| ElicitationRequest
|
||||
)
|
||||
@@ -0,0 +1,240 @@
|
||||
"""Session-scoped files namespace — upload, list, get, delete."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import mimetypes
|
||||
import pathlib
|
||||
from urllib.parse import quote
|
||||
|
||||
import httpx
|
||||
|
||||
from ._errors import raise_for_status, require_json_object, response_body
|
||||
from ._types import File, PaginatedList
|
||||
|
||||
|
||||
class FilesNamespace:
|
||||
"""
|
||||
Factory for session-scoped file namespaces.
|
||||
|
||||
:param http: Shared async HTTP client.
|
||||
:param base_url: Server base URL, e.g. ``"http://localhost:8080"``.
|
||||
"""
|
||||
|
||||
def __init__(self, http: httpx.AsyncClient, base_url: str) -> None:
|
||||
"""
|
||||
Initialize the unbound files namespace.
|
||||
|
||||
:param http: Shared async HTTP client.
|
||||
:param base_url: Server base URL, e.g. ``"http://localhost:8080"``.
|
||||
:returns: None.
|
||||
"""
|
||||
self._http = http
|
||||
self._base = base_url
|
||||
|
||||
def for_session(self, session_id: str) -> SessionFilesNamespace:
|
||||
"""
|
||||
Return a file namespace bound to one session.
|
||||
|
||||
``/v1/files`` has been removed; callers must scope uploads and
|
||||
downloads to the session/conversation that owns the file.
|
||||
|
||||
:param session_id: Session/conversation id, e.g. ``"conv_abc123"``.
|
||||
:returns: A :class:`SessionFilesNamespace` bound to that session.
|
||||
"""
|
||||
return SessionFilesNamespace(self._http, self._base, session_id)
|
||||
|
||||
async def upload(self, path: str) -> File:
|
||||
"""Legacy global upload removed; use ``for_session(id).upload``."""
|
||||
raise RuntimeError(
|
||||
"/v1/files was removed; use client.files.for_session(session_id).upload(path)"
|
||||
)
|
||||
|
||||
async def list(self, *args: object, **kwargs: object) -> list[File]:
|
||||
"""Legacy global list removed; use ``for_session(id).list``."""
|
||||
raise RuntimeError(
|
||||
"/v1/files was removed; use client.files.for_session(session_id).list()"
|
||||
)
|
||||
|
||||
async def get(self, file_id: str) -> File:
|
||||
"""Legacy global get removed; use ``for_session(id).get``."""
|
||||
raise RuntimeError(
|
||||
"/v1/files was removed; use client.files.for_session(session_id).get(file_id)"
|
||||
)
|
||||
|
||||
async def get_content(self, file_id: str) -> bytes:
|
||||
"""Legacy global content download removed; use ``for_session(id).get_content``."""
|
||||
raise RuntimeError(
|
||||
"/v1/files was removed; use client.files.for_session(session_id).get_content(file_id)"
|
||||
)
|
||||
|
||||
async def download(self, file_id: str, to_path: str | pathlib.Path) -> pathlib.Path:
|
||||
"""Legacy global download removed; use ``for_session(id).download``."""
|
||||
raise RuntimeError(
|
||||
"/v1/files was removed; use "
|
||||
"client.files.for_session(session_id).download(file_id, path)"
|
||||
)
|
||||
|
||||
async def delete(self, file_id: str) -> None:
|
||||
"""Legacy global delete removed; use ``for_session(id).delete``."""
|
||||
raise RuntimeError(
|
||||
"/v1/files was removed; use client.files.for_session(session_id).delete(file_id)"
|
||||
)
|
||||
|
||||
|
||||
class SessionFilesNamespace:
|
||||
"""
|
||||
File operations scoped to a single session.
|
||||
|
||||
:param http: Shared async HTTP client.
|
||||
:param base_url: Server base URL, e.g. ``"http://localhost:8080"``.
|
||||
:param session_id: Session/conversation id, e.g. ``"conv_abc123"``.
|
||||
"""
|
||||
|
||||
def __init__(self, http: httpx.AsyncClient, base_url: str, session_id: str) -> None:
|
||||
"""
|
||||
Initialize the session-scoped files namespace.
|
||||
|
||||
:param http: Shared async HTTP client.
|
||||
:param base_url: Server base URL, e.g. ``"http://localhost:8080"``.
|
||||
:param session_id: Session/conversation id, e.g. ``"conv_abc123"``.
|
||||
:returns: None.
|
||||
"""
|
||||
self._http = http
|
||||
self._base = base_url
|
||||
self._session_id = session_id
|
||||
|
||||
@property
|
||||
def session_id(self) -> str:
|
||||
"""
|
||||
The session/conversation id this namespace is bound to.
|
||||
|
||||
:returns: Session id, e.g. ``"conv_abc123"``.
|
||||
"""
|
||||
return self._session_id
|
||||
|
||||
@property
|
||||
def _path(self) -> str:
|
||||
"""
|
||||
Session-scoped files collection URL.
|
||||
|
||||
:returns: Fully qualified endpoint URL.
|
||||
"""
|
||||
return f"{self._base}/v1/sessions/{quote(self._session_id, safe='')}/resources/files"
|
||||
|
||||
@staticmethod
|
||||
def _resource_to_file(data: dict[str, object]) -> File:
|
||||
"""
|
||||
Convert a session file resource response to SDK ``File``.
|
||||
|
||||
:param data: JSON object returned by a session file endpoint.
|
||||
:returns: The corresponding :class:`File`.
|
||||
"""
|
||||
metadata = data.get("metadata")
|
||||
metadata_dict = metadata if isinstance(metadata, dict) else {}
|
||||
return File.from_dict(
|
||||
{
|
||||
"id": data.get("id", ""),
|
||||
"filename": metadata_dict.get("filename", data.get("name", "")),
|
||||
"bytes": metadata_dict.get("bytes", 0),
|
||||
"created_at": metadata_dict.get("created_at", 0),
|
||||
}
|
||||
)
|
||||
|
||||
async def upload(self, path: str) -> File:
|
||||
"""
|
||||
Upload a local file into this session's file namespace.
|
||||
|
||||
:param path: Path to the local file, e.g. ``"./report.pdf"``.
|
||||
:returns: The uploaded file metadata.
|
||||
"""
|
||||
p = pathlib.Path(path)
|
||||
content_type = mimetypes.guess_type(str(p))[0]
|
||||
with open(p, "rb") as f:
|
||||
resp = await self._http.post(
|
||||
self._path,
|
||||
files={"file": (p.name, f, content_type)},
|
||||
timeout=30.0,
|
||||
)
|
||||
raise_for_status(resp.status_code, response_body(resp))
|
||||
return self._resource_to_file(
|
||||
require_json_object(resp, "POST /v1/sessions/{session_id}/resources/files")
|
||||
)
|
||||
|
||||
async def list(
|
||||
self,
|
||||
*,
|
||||
limit: int = 20,
|
||||
after: str | None = None,
|
||||
order: str = "desc",
|
||||
) -> list[File]:
|
||||
"""
|
||||
List files owned by this session.
|
||||
|
||||
:param limit: Maximum number of files to return.
|
||||
:param after: Cursor for forward pagination,
|
||||
e.g. ``"file_abc123"``.
|
||||
:param order: Sort order, e.g. ``"desc"``.
|
||||
:returns: File metadata entries.
|
||||
"""
|
||||
params: dict[str, object] = {"limit": limit, "order": order}
|
||||
if after is not None:
|
||||
params["after"] = after
|
||||
resp = await self._http.get(self._path, params=params)
|
||||
raise_for_status(resp.status_code, response_body(resp))
|
||||
page = PaginatedList.from_dict(
|
||||
require_json_object(resp, "GET /v1/sessions/{session_id}/resources/files")
|
||||
)
|
||||
return [self._resource_to_file(d) for d in page.data]
|
||||
|
||||
async def get(self, file_id: str) -> File:
|
||||
"""
|
||||
Get session file metadata by ID.
|
||||
|
||||
:param file_id: Server-issued file id, e.g. ``"file_abc123"``.
|
||||
:returns: File metadata.
|
||||
"""
|
||||
resp = await self._http.get(f"{self._path}/{quote(file_id, safe='')}")
|
||||
raise_for_status(resp.status_code, response_body(resp))
|
||||
return self._resource_to_file(
|
||||
require_json_object(resp, "GET /v1/sessions/{session_id}/resources/files/{file_id}")
|
||||
)
|
||||
|
||||
async def get_content(self, file_id: str) -> bytes:
|
||||
"""
|
||||
Download session file content.
|
||||
|
||||
:param file_id: Server-issued file id, e.g. ``"file_abc123"``.
|
||||
:returns: Raw file bytes.
|
||||
"""
|
||||
resp = await self._http.get(
|
||||
f"{self._path}/{quote(file_id, safe='')}/content",
|
||||
timeout=30.0,
|
||||
)
|
||||
if resp.status_code >= 400:
|
||||
raise_for_status(resp.status_code, response_body(resp))
|
||||
return resp.content
|
||||
|
||||
async def download(self, file_id: str, to_path: str | pathlib.Path) -> pathlib.Path:
|
||||
"""
|
||||
Download file content and write it to disk.
|
||||
|
||||
:param file_id: Server-issued file id, e.g. ``"file_abc123"``.
|
||||
:param to_path: Local output path, e.g. ``"./out/report.pdf"``.
|
||||
:returns: The path that was written.
|
||||
"""
|
||||
content = await self.get_content(file_id)
|
||||
path = pathlib.Path(to_path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_bytes(content)
|
||||
return path
|
||||
|
||||
async def delete(self, file_id: str) -> None:
|
||||
"""
|
||||
Delete a session-scoped file.
|
||||
|
||||
:param file_id: Server-issued file id, e.g. ``"file_abc123"``.
|
||||
:returns: None.
|
||||
"""
|
||||
resp = await self._http.delete(f"{self._path}/{quote(file_id, safe='')}")
|
||||
if resp.status_code >= 400:
|
||||
raise_for_status(resp.status_code, response_body(resp))
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Return types for the ``query()`` convenience API.
|
||||
|
||||
:class:`QueryResult` is returned from ``query(stream=False)`` (the
|
||||
default). It carries the final assistant text plus any file artifacts
|
||||
the agent produced on this turn.
|
||||
|
||||
:class:`QueryStream` is returned from ``query(stream=True)``. It is
|
||||
an async-iterable of text chunks. After the iteration completes, its
|
||||
``files`` attribute holds the same list of produced files.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from ._types import File
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class QueryResult:
|
||||
"""Non-streaming result of a ``query()`` call.
|
||||
|
||||
:param text: The assistant's final text, joined across tool-loop
|
||||
iterations. Empty string if the agent produced no text.
|
||||
:param files: Files the agent produced on this turn (via the
|
||||
``upload_file`` builtin or equivalent). Empty list if none.
|
||||
Each entry is a :class:`File` with ``id``, ``filename``,
|
||||
``bytes``, and ``created_at``; use
|
||||
:meth:`FilesNamespace.download` or
|
||||
:meth:`FilesNamespace.get_content` to retrieve the bytes.
|
||||
"""
|
||||
|
||||
text: str
|
||||
files: list[File] = field(default_factory=list)
|
||||
|
||||
|
||||
class QueryStream:
|
||||
"""Async-iterable of text chunks from ``query(stream=True)``.
|
||||
|
||||
Iterating yields ``str`` chunks in order as the agent emits them.
|
||||
After the iterator is exhausted, :attr:`files` holds the file
|
||||
artifacts produced this turn. Files discovered during streaming
|
||||
are appended as they arrive, so an in-progress iteration may
|
||||
expose a partial list.
|
||||
|
||||
Single-use: iterating a second time raises ``RuntimeError``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chunks: AsyncIterator[str],
|
||||
files: list[File],
|
||||
) -> None:
|
||||
"""Wrap a text-chunk iterator with a reference to a shared file list.
|
||||
|
||||
:param chunks: Async iterator that yields assistant text
|
||||
chunks. Typically produced by the ``Session`` internals.
|
||||
:param files: List that will be populated with produced
|
||||
:class:`File` entries as the iterator is consumed.
|
||||
Passed by reference so the caller and the backing
|
||||
generator share the same list.
|
||||
"""
|
||||
self._chunks = chunks
|
||||
self._files = files
|
||||
self._consumed = False
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[str]:
|
||||
"""Return the underlying async iterator.
|
||||
|
||||
:raises RuntimeError: If the stream has already been iterated.
|
||||
:returns: The wrapped async iterator of text chunks.
|
||||
"""
|
||||
if self._consumed:
|
||||
raise RuntimeError(
|
||||
"QueryStream has already been iterated. Each stream is "
|
||||
"single-use — call client.query(..., stream=True) again "
|
||||
"to get a fresh one."
|
||||
)
|
||||
self._consumed = True
|
||||
return self._chunks
|
||||
|
||||
@property
|
||||
def files(self) -> list[File]:
|
||||
"""Files the agent has produced so far this turn.
|
||||
|
||||
Empty until the first file arrives. Fully populated once the
|
||||
iterator is exhausted.
|
||||
|
||||
:returns: A shallow copy of the internal list.
|
||||
"""
|
||||
return list(self._files)
|
||||
@@ -0,0 +1,977 @@
|
||||
"""Responses namespace — streaming, tool loop, polling, cancel, steer.
|
||||
|
||||
.. deprecated::
|
||||
``ResponsesNamespace`` (``client.responses.*``) targets the
|
||||
deprecated ``/v1/responses`` surface. New code should use
|
||||
``SessionsNamespace`` (``client.sessions.*``) instead.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import warnings
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from ._errors import ToolCallDenied, raise_for_status, require_json_object, response_body
|
||||
from ._events import (
|
||||
ClientTaskCancel,
|
||||
CompactionInProgress,
|
||||
ElicitationRequest,
|
||||
ErrorEvent,
|
||||
MessageDone,
|
||||
NativeToolCall,
|
||||
OutputFileDone,
|
||||
ReasoningDelta,
|
||||
ReasoningStarted,
|
||||
ReasoningSummaryDelta,
|
||||
ResponseCancelled,
|
||||
ResponseCompleted,
|
||||
ResponseCreated,
|
||||
ResponseFailed,
|
||||
ResponseIncomplete,
|
||||
RetryEvent,
|
||||
StreamEvent,
|
||||
TextDelta,
|
||||
ToolCall,
|
||||
ToolResult,
|
||||
)
|
||||
from ._sse import parse_sse_stream
|
||||
from ._tool_handler import (
|
||||
CompactionStartCtx,
|
||||
ElicitationRequestCtx,
|
||||
FileOutputCtx,
|
||||
MessageEndCtx,
|
||||
MessageStartCtx,
|
||||
NativeToolCallCtx,
|
||||
ReasoningEndCtx,
|
||||
ResponseEndCtx,
|
||||
ResponseStartCtx,
|
||||
RetryCtx,
|
||||
ServerErrorCtx,
|
||||
StreamHooks,
|
||||
ToolCallEndCtx,
|
||||
ToolCallInfo,
|
||||
ToolCallStartCtx,
|
||||
ToolHandler,
|
||||
ToolResultInfo,
|
||||
ToolResultsReadyCtx,
|
||||
)
|
||||
from ._tool_handler import (
|
||||
ReasoningStartCtx as ReasoningStartHookCtx,
|
||||
)
|
||||
from ._types import Response
|
||||
|
||||
_log = logging.getLogger("omnigent_client.responses")
|
||||
|
||||
# Terminal statuses — the response won't change further.
|
||||
_TERMINAL_STATUSES = frozenset({"completed", "failed", "incomplete", "cancelled"})
|
||||
|
||||
# Async dispatch of client tools flows through the
|
||||
# action_required handler (``_execute_and_patch``) — same path
|
||||
# the sub-agent tunnel uses. The SDK tracks the spawned local
|
||||
# tasks by call_id so :class:`ClientTaskCancel` events can find
|
||||
# and cancel them; everything else is fire-and-forget and
|
||||
# self-cleans on done.
|
||||
|
||||
|
||||
async def _call_hook(hook: Any, ctx: Any) -> Any:
|
||||
"""Call a hook (sync or async) and return its result."""
|
||||
if hook is None:
|
||||
return None
|
||||
result = hook(ctx)
|
||||
if inspect.isawaitable(result):
|
||||
return await result
|
||||
return result
|
||||
|
||||
|
||||
class ResponsesNamespace:
|
||||
"""Methods for ``/v1/responses`` endpoints.
|
||||
|
||||
.. deprecated::
|
||||
This namespace targets the deprecated ``/v1/responses``
|
||||
surface. Use :class:`SessionsNamespace`
|
||||
(``client.sessions.*``) instead — see
|
||||
``designs/SESSION_REARCHITECTURE.md``.
|
||||
"""
|
||||
|
||||
def __init__(self, http: httpx.AsyncClient, base_url: str) -> None:
|
||||
self._http = http
|
||||
self._base = base_url
|
||||
self._warned = False
|
||||
|
||||
def _deprecation_warn(self, method: str) -> None:
|
||||
"""Emit a one-shot ``DeprecationWarning`` per instance.
|
||||
|
||||
:param method: Name of the deprecated method, e.g.
|
||||
``"create"``.
|
||||
"""
|
||||
if not self._warned:
|
||||
warnings.warn(
|
||||
f"ResponsesNamespace.{method}() is deprecated; use client.sessions.* instead",
|
||||
DeprecationWarning,
|
||||
stacklevel=3,
|
||||
)
|
||||
self._warned = True
|
||||
|
||||
async def create(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
input: str | list[dict[str, object]],
|
||||
background: bool = False,
|
||||
instructions: str | None = None,
|
||||
previous_response_id: str | None = None,
|
||||
tools: list[dict[str, object]] | None = None,
|
||||
reasoning: dict[str, str] | None = None,
|
||||
model_override: str | None = None,
|
||||
) -> Response:
|
||||
"""Create a response (blocking, non-streaming).
|
||||
|
||||
:param model: Agent name.
|
||||
:param input: User text or content block list.
|
||||
:param background: If True, returns immediately (poll via get()).
|
||||
:param instructions: Per-request system instructions.
|
||||
:param previous_response_id: Prior response for multi-turn.
|
||||
:param tools: Client-specified tool schemas.
|
||||
:param reasoning: Reasoning config, e.g. ``{"effort": "high"}``.
|
||||
:param model_override: Optional per-request LLM model override,
|
||||
e.g. ``"openai/gpt-5.4-mini"``. Shadows the spec model.
|
||||
:returns: The response object.
|
||||
"""
|
||||
self._deprecation_warn("create")
|
||||
body = _build_body(
|
||||
model=model,
|
||||
input=input,
|
||||
stream=False,
|
||||
background=background,
|
||||
instructions=instructions,
|
||||
previous_response_id=previous_response_id,
|
||||
tools=tools,
|
||||
reasoning=reasoning,
|
||||
model_override=model_override,
|
||||
)
|
||||
resp = await self._http.post(f"{self._base}/v1/responses", json=body)
|
||||
raise_for_status(resp.status_code, response_body(resp))
|
||||
return Response.from_dict(require_json_object(resp, "POST /v1/responses"))
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
input: str | list[dict[str, object]],
|
||||
background: bool = False,
|
||||
instructions: str | None = None,
|
||||
previous_response_id: str | None = None,
|
||||
tool_handler: ToolHandler | None = None,
|
||||
hooks: StreamHooks | None = None,
|
||||
reasoning: dict[str, str] | None = None,
|
||||
model_override: str | None = None,
|
||||
) -> AsyncIterator[StreamEvent]:
|
||||
"""Stream a response, optionally running the tool loop.
|
||||
|
||||
If ``tool_handler`` is provided, the client runs the full tool
|
||||
execution loop: stream -> detect tool calls -> execute via handler
|
||||
-> POST results -> stream again. The consumer sees one continuous
|
||||
sequence of events.
|
||||
|
||||
If ``tool_handler`` is None, yields raw events for a single server
|
||||
response.
|
||||
|
||||
:param model: Agent name.
|
||||
:param input: User text or content block list.
|
||||
:param tool_handler: Optional client-side tool execution config.
|
||||
:param hooks: Optional lifecycle hooks.
|
||||
"""
|
||||
self._deprecation_warn("stream")
|
||||
if hooks is None:
|
||||
hooks = StreamHooks()
|
||||
|
||||
current_input: str | list[dict[str, object]] = input
|
||||
current_prev_id = previous_response_id
|
||||
iteration = 0
|
||||
# Track local asyncio.Tasks spawned by the action_required
|
||||
# handler so :class:`ClientTaskCancel` SSE events can find
|
||||
# and cancel them. Keyed by ``call_id`` because that's the
|
||||
# field the cancel event carries (looked up from the
|
||||
# server-side pending_tool_calls row when the cancel fires).
|
||||
# Persists across the outer ``while True`` so a task
|
||||
# dispatched in one iteration can be cancelled by an event
|
||||
# arriving in a later iteration's stream. Each entry
|
||||
# self-cleans on done via ``add_done_callback``.
|
||||
pending_local_tasks: dict[str, asyncio.Task[None]] = {}
|
||||
|
||||
while True:
|
||||
tools = tool_handler.schemas if tool_handler is not None else None
|
||||
pending_client_calls: list[ToolCall] = []
|
||||
completed_call_ids: set[str] = set()
|
||||
current_response_id: str | None = None
|
||||
current_session_id: str | None = None
|
||||
in_reasoning = False
|
||||
reasoning_text = ""
|
||||
summary_text = ""
|
||||
message_started = False
|
||||
|
||||
body = _build_body(
|
||||
model=model,
|
||||
input=current_input,
|
||||
stream=True,
|
||||
background=background,
|
||||
instructions=instructions,
|
||||
previous_response_id=current_prev_id,
|
||||
tools=tools,
|
||||
reasoning=reasoning,
|
||||
model_override=model_override,
|
||||
)
|
||||
|
||||
async with self._http.stream("POST", f"{self._base}/v1/responses", json=body) as resp:
|
||||
if resp.status_code >= 400:
|
||||
await resp.aread()
|
||||
raise_for_status(resp.status_code, response_body(resp))
|
||||
|
||||
async for event in parse_sse_stream(resp.aiter_bytes()):
|
||||
# ── Fire hooks and collect state ──────────
|
||||
|
||||
if isinstance(event, ResponseCreated):
|
||||
current_response_id = event.response.id
|
||||
if event.response.conversation is not None:
|
||||
current_session_id = event.response.conversation.id
|
||||
await _call_hook(
|
||||
hooks.on_response_start, ResponseStartCtx(response=event.response)
|
||||
)
|
||||
|
||||
elif isinstance(event, ReasoningStarted):
|
||||
in_reasoning = True
|
||||
reasoning_text = ""
|
||||
summary_text = ""
|
||||
await _call_hook(hooks.on_reasoning_start, ReasoningStartHookCtx())
|
||||
|
||||
elif isinstance(event, ReasoningDelta):
|
||||
reasoning_text += event.delta
|
||||
|
||||
elif isinstance(event, ReasoningSummaryDelta):
|
||||
summary_text += event.delta
|
||||
|
||||
elif isinstance(event, TextDelta):
|
||||
if in_reasoning:
|
||||
in_reasoning = False
|
||||
await _call_hook(
|
||||
hooks.on_reasoning_end,
|
||||
ReasoningEndCtx(
|
||||
reasoning_text=reasoning_text, summary_text=summary_text
|
||||
),
|
||||
)
|
||||
if not message_started:
|
||||
message_started = True
|
||||
await _call_hook(
|
||||
hooks.on_message_start,
|
||||
MessageStartCtx(response_id=current_response_id or ""),
|
||||
)
|
||||
|
||||
elif isinstance(event, CompactionInProgress):
|
||||
await _call_hook(hooks.on_compaction_start, CompactionStartCtx())
|
||||
|
||||
elif isinstance(event, ToolCall):
|
||||
is_client_side = (
|
||||
tool_handler is not None and event.call_id not in completed_call_ids
|
||||
)
|
||||
executed_by = "client" if is_client_side else "server"
|
||||
await _call_hook(
|
||||
hooks.on_tool_call_start,
|
||||
ToolCallStartCtx(
|
||||
name=event.name,
|
||||
arguments=event.arguments,
|
||||
call_id=event.call_id,
|
||||
agent_name=event.agent_name,
|
||||
executed_by=executed_by,
|
||||
),
|
||||
)
|
||||
# Action_required tool calls (sub-agent
|
||||
# tunnel + sys_call_async-dispatched
|
||||
# client tools) execute locally and PATCH
|
||||
# back while the stream continues.
|
||||
# Tracked by call_id so a follow-up
|
||||
# ``ClientTaskCancel`` event can cancel
|
||||
# the local task.
|
||||
if event.status == "action_required" and tool_handler is not None:
|
||||
completed_call_ids.add(event.call_id)
|
||||
local_task = asyncio.create_task(
|
||||
_execute_and_patch(
|
||||
self._http,
|
||||
self._base,
|
||||
tool_handler,
|
||||
hooks,
|
||||
event,
|
||||
current_response_id or "",
|
||||
iteration,
|
||||
)
|
||||
)
|
||||
pending_local_tasks[event.call_id] = local_task
|
||||
# Auto-clean once the task settles —
|
||||
# PATCH back, exception, or cancel all
|
||||
# land here. Captured in a closure so
|
||||
# the dict reference stays live.
|
||||
_local_task_call_id = event.call_id
|
||||
|
||||
def _drop_done(
|
||||
_t: asyncio.Task[None],
|
||||
_key: str = _local_task_call_id,
|
||||
) -> None:
|
||||
pending_local_tasks.pop(_key, None)
|
||||
|
||||
local_task.add_done_callback(_drop_done)
|
||||
elif is_client_side:
|
||||
# Synchronous client-side tool calls —
|
||||
# the LLM called the tool directly.
|
||||
# The SDK doesn't run them inline; it
|
||||
# accumulates them and the outer
|
||||
# ``stream()`` loop runs them after
|
||||
# the response terminates and PATCHes
|
||||
# ``tool_results`` to feed the next
|
||||
# iteration. Authors who want
|
||||
# background dispatch use
|
||||
# ``sys_call_async`` instead, which
|
||||
# produces an action_required event
|
||||
# handled by the branch above.
|
||||
pending_client_calls.append(event)
|
||||
|
||||
elif isinstance(event, ElicitationRequest):
|
||||
# MCP-shape elicitation — the server is
|
||||
# waiting on a verdict. Route to the
|
||||
# elicitation hook (not the tool_handler —
|
||||
# this is not a tool call). If no hook is
|
||||
# registered, fail-closed decline so the
|
||||
# parked workflow doesn't stall forever.
|
||||
# Fire-and-forget elicitation response —
|
||||
# see the peer tunneled-call site above for
|
||||
# why we intentionally drop the task
|
||||
# reference.
|
||||
asyncio.ensure_future( # noqa: RUF006 — see comment above
|
||||
_handle_elicitation_request(
|
||||
self._http,
|
||||
self._base,
|
||||
hooks,
|
||||
event,
|
||||
current_response_id or "",
|
||||
current_session_id or "",
|
||||
)
|
||||
)
|
||||
|
||||
elif isinstance(event, ToolResult):
|
||||
completed_call_ids.add(event.call_id)
|
||||
await _call_hook(
|
||||
hooks.on_tool_call_end,
|
||||
ToolCallEndCtx(
|
||||
name="",
|
||||
call_id=event.call_id,
|
||||
agent_name="",
|
||||
output=event.output,
|
||||
),
|
||||
)
|
||||
|
||||
elif isinstance(event, ClientTaskCancel):
|
||||
# Server is telling us to stop the local
|
||||
# body. Look up the local task by
|
||||
# call_id (the server includes the
|
||||
# synthetic call_id from the
|
||||
# pending_tool_call row in the SSE
|
||||
# payload) and cancel — the body's
|
||||
# ``except CancelledError`` runs to
|
||||
# cleanup. ``ensure_future``-style fires
|
||||
# without a tracked task remain
|
||||
# un-cancellable; that case is benign
|
||||
# because the server has already decided
|
||||
# the task is cancelled and the late
|
||||
# PATCH back will be a G3 no-op.
|
||||
if event.call_id is not None:
|
||||
tracked_task = pending_local_tasks.get(event.call_id)
|
||||
if tracked_task is not None and not tracked_task.done():
|
||||
tracked_task.cancel()
|
||||
|
||||
elif isinstance(event, NativeToolCall):
|
||||
await _call_hook(
|
||||
hooks.on_native_tool_call,
|
||||
NativeToolCallCtx(tool_type=event.tool_type, data=event.data),
|
||||
)
|
||||
|
||||
elif isinstance(event, MessageDone):
|
||||
if in_reasoning:
|
||||
in_reasoning = False
|
||||
await _call_hook(
|
||||
hooks.on_reasoning_end,
|
||||
ReasoningEndCtx(
|
||||
reasoning_text=reasoning_text, summary_text=summary_text
|
||||
),
|
||||
)
|
||||
await _call_hook(
|
||||
hooks.on_message_end,
|
||||
MessageEndCtx(content=event.content),
|
||||
)
|
||||
message_started = False
|
||||
|
||||
elif isinstance(event, OutputFileDone):
|
||||
await _call_hook(
|
||||
hooks.on_file_output,
|
||||
FileOutputCtx(
|
||||
file_id=event.file_id,
|
||||
filename=event.filename,
|
||||
content_type=event.content_type,
|
||||
),
|
||||
)
|
||||
|
||||
elif isinstance(event, RetryEvent):
|
||||
await _call_hook(
|
||||
hooks.on_retry,
|
||||
RetryCtx(
|
||||
source=event.source,
|
||||
tool_name=event.tool_name,
|
||||
attempt=event.attempt,
|
||||
max_attempts=event.max_attempts,
|
||||
delay_seconds=event.delay_seconds,
|
||||
error=event.error,
|
||||
),
|
||||
)
|
||||
|
||||
elif isinstance(event, ErrorEvent):
|
||||
await _call_hook(
|
||||
hooks.on_server_error,
|
||||
ServerErrorCtx(
|
||||
source=event.source,
|
||||
tool_name=event.tool_name,
|
||||
error=event.error,
|
||||
),
|
||||
)
|
||||
|
||||
elif isinstance(
|
||||
event,
|
||||
ResponseCompleted
|
||||
| ResponseFailed
|
||||
| ResponseIncomplete
|
||||
| ResponseCancelled,
|
||||
):
|
||||
if in_reasoning:
|
||||
in_reasoning = False
|
||||
await _call_hook(
|
||||
hooks.on_reasoning_end,
|
||||
ReasoningEndCtx(
|
||||
reasoning_text=reasoning_text, summary_text=summary_text
|
||||
),
|
||||
)
|
||||
status = event.response.status
|
||||
await _call_hook(
|
||||
hooks.on_response_end,
|
||||
ResponseEndCtx(response=event.response, status=status),
|
||||
)
|
||||
|
||||
# ── Yield event to consumer ──────────────
|
||||
yield event
|
||||
|
||||
# ── Post-stream: tool loop ───────────────────
|
||||
if current_response_id is not None:
|
||||
current_prev_id = current_response_id
|
||||
|
||||
# Filter out calls that already have server-side results.
|
||||
pending_client_calls = [
|
||||
tc for tc in pending_client_calls if tc.call_id not in completed_call_ids
|
||||
]
|
||||
|
||||
if not pending_client_calls or tool_handler is None:
|
||||
break
|
||||
|
||||
# Execute client-side tools and build results.
|
||||
results: list[dict[str, object]] = []
|
||||
result_infos: list[ToolResultInfo] = []
|
||||
|
||||
for tc in pending_client_calls:
|
||||
call_info = ToolCallInfo(
|
||||
name=tc.name,
|
||||
arguments=tc.arguments,
|
||||
call_id=tc.call_id,
|
||||
agent_name=tc.agent_name,
|
||||
response_id=current_response_id or "",
|
||||
iteration=iteration,
|
||||
)
|
||||
try:
|
||||
# Sync ``execute`` routes through
|
||||
# ``asyncio.to_thread`` so a blocking body
|
||||
# doesn't serialize the sync-tool loop
|
||||
# (each call would otherwise freeze the
|
||||
# event loop for its duration).
|
||||
output = await _call_execute_off_loop(tool_handler, call_info)
|
||||
except ToolCallDenied as exc:
|
||||
output = str(exc)
|
||||
|
||||
await _call_hook(
|
||||
hooks.on_tool_call_end,
|
||||
ToolCallEndCtx(
|
||||
name=tc.name,
|
||||
call_id=tc.call_id,
|
||||
agent_name=tc.agent_name,
|
||||
output=output,
|
||||
),
|
||||
)
|
||||
yield ToolResult(call_id=tc.call_id, output=output)
|
||||
|
||||
results.append(
|
||||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": tc.call_id,
|
||||
"output": output,
|
||||
}
|
||||
)
|
||||
result_infos.append(
|
||||
ToolResultInfo(
|
||||
call_id=tc.call_id,
|
||||
name=tc.name,
|
||||
output=output,
|
||||
agent_name=tc.agent_name,
|
||||
)
|
||||
)
|
||||
|
||||
await _call_hook(
|
||||
hooks.on_tool_results_ready,
|
||||
ToolResultsReadyCtx(results=result_infos, iteration=iteration),
|
||||
)
|
||||
|
||||
current_input = results
|
||||
iteration += 1
|
||||
|
||||
async def get(self, response_id: str) -> Response:
|
||||
"""Get a response by ID (poll for status).
|
||||
|
||||
:param response_id: The response/task ID.
|
||||
:returns: Current response state.
|
||||
"""
|
||||
self._deprecation_warn("get")
|
||||
resp = await self._http.get(f"{self._base}/v1/responses/{response_id}")
|
||||
raise_for_status(resp.status_code, response_body(resp))
|
||||
return Response.from_dict(require_json_object(resp, "GET /v1/responses/{response_id}"))
|
||||
|
||||
async def poll(
|
||||
self,
|
||||
response_id: str,
|
||||
*,
|
||||
interval: float = 0.5,
|
||||
tool_handler: ToolHandler | None = None,
|
||||
) -> Response:
|
||||
"""Poll a background response until it reaches a terminal status.
|
||||
|
||||
If ``tool_handler`` is provided, tunneled tool calls
|
||||
(``status: "action_required"``) are executed and PATCHed back.
|
||||
|
||||
:param response_id: The response/task ID.
|
||||
:param interval: Seconds between polls.
|
||||
:param tool_handler: Optional client-side tool handler.
|
||||
:returns: The terminal response.
|
||||
"""
|
||||
while True:
|
||||
response = await self.get(response_id)
|
||||
if response.status in _TERMINAL_STATUSES:
|
||||
return response
|
||||
# Check for tunneled tool calls needing client execution.
|
||||
if tool_handler is not None:
|
||||
await self._handle_polling_tool_calls(response_id, response, tool_handler)
|
||||
await asyncio.sleep(interval)
|
||||
|
||||
async def _handle_polling_tool_calls(
|
||||
self,
|
||||
response_id: str,
|
||||
response: Response,
|
||||
tool_handler: ToolHandler,
|
||||
) -> None:
|
||||
"""Execute action_required tool calls found during polling."""
|
||||
action_required = [
|
||||
item
|
||||
for item in response.output
|
||||
if isinstance(item, dict)
|
||||
and item.get("type") == "function_call"
|
||||
and item.get("status") == "action_required"
|
||||
]
|
||||
if not action_required:
|
||||
return
|
||||
|
||||
tool_results = []
|
||||
for fc in action_required:
|
||||
name = str(fc.get("name", ""))
|
||||
call_id = str(fc.get("call_id", ""))
|
||||
args_str = str(fc.get("arguments", "{}"))
|
||||
try:
|
||||
arguments = json.loads(args_str)
|
||||
except json.JSONDecodeError:
|
||||
arguments = {}
|
||||
|
||||
call_info = ToolCallInfo(
|
||||
name=name,
|
||||
arguments=arguments,
|
||||
call_id=call_id,
|
||||
agent_name=str(fc.get("model", "")),
|
||||
response_id=response_id,
|
||||
iteration=0,
|
||||
)
|
||||
output = tool_handler.execute(call_info)
|
||||
if inspect.isawaitable(output):
|
||||
output = await output
|
||||
tool_results.append({"call_id": call_id, "output": output})
|
||||
|
||||
if tool_results:
|
||||
await self._patch_tool_results(response_id, tool_results)
|
||||
|
||||
async def _patch_tool_results(
|
||||
self,
|
||||
response_id: str,
|
||||
tool_results: list[dict[str, str]],
|
||||
) -> None:
|
||||
"""PATCH tool results back to the server."""
|
||||
resp = await self._http.patch(
|
||||
f"{self._base}/v1/responses/{response_id}",
|
||||
json={"tool_results": tool_results},
|
||||
timeout=60.0,
|
||||
)
|
||||
if resp.status_code not in (200, 404, 409):
|
||||
_log.warning(
|
||||
"PATCH tool results failed (%d): %s",
|
||||
resp.status_code,
|
||||
resp.text[:200],
|
||||
)
|
||||
|
||||
async def steer(
|
||||
self,
|
||||
response_id: str,
|
||||
input: str,
|
||||
*,
|
||||
model: str,
|
||||
reasoning: dict[str, str] | None = None,
|
||||
model_override: str | None = None,
|
||||
) -> Response:
|
||||
"""Send a steering message to an in-progress response.
|
||||
|
||||
:param response_id: The in-progress response ID.
|
||||
:param input: Steering text.
|
||||
:param model: Agent name.
|
||||
:param reasoning: Reasoning config if this steer races into a new response.
|
||||
:param model_override: LLM model override if this steer races
|
||||
into a new response (server ignores it otherwise).
|
||||
:returns: The response (same ID if delivered, new ID if agent finished).
|
||||
"""
|
||||
self._deprecation_warn("steer")
|
||||
body: dict[str, object] = {
|
||||
"model": model,
|
||||
"input": input,
|
||||
"previous_response_id": response_id,
|
||||
"stream": False,
|
||||
"background": True,
|
||||
}
|
||||
if reasoning is not None:
|
||||
body["reasoning"] = reasoning
|
||||
if model_override is not None:
|
||||
body["model_override"] = model_override
|
||||
resp = await self._http.post(
|
||||
f"{self._base}/v1/responses",
|
||||
json=body,
|
||||
timeout=120.0,
|
||||
)
|
||||
raise_for_status(resp.status_code, response_body(resp))
|
||||
return Response.from_dict(require_json_object(resp, "POST /v1/responses"))
|
||||
|
||||
async def cancel(self, response_id: str) -> Response:
|
||||
"""Cancel an in-progress response.
|
||||
|
||||
:param response_id: The response ID to cancel.
|
||||
:returns: The cancelled response.
|
||||
"""
|
||||
self._deprecation_warn("cancel")
|
||||
resp = await self._http.post(
|
||||
f"{self._base}/v1/responses/{response_id}/cancel",
|
||||
timeout=10.0,
|
||||
)
|
||||
raise_for_status(resp.status_code, response_body(resp))
|
||||
return Response.from_dict(
|
||||
require_json_object(resp, "POST /v1/responses/{response_id}/cancel")
|
||||
)
|
||||
|
||||
async def delete(self, response_id: str) -> None:
|
||||
"""Delete a response.
|
||||
|
||||
:param response_id: The response ID to delete.
|
||||
"""
|
||||
self._deprecation_warn("delete")
|
||||
resp = await self._http.delete(
|
||||
f"{self._base}/v1/responses/{response_id}",
|
||||
)
|
||||
if resp.status_code >= 400:
|
||||
raise_for_status(resp.status_code, response_body(resp))
|
||||
|
||||
|
||||
# ── Helpers ──────────────────────────────────────────────
|
||||
|
||||
|
||||
def _build_body(
|
||||
*,
|
||||
model: str,
|
||||
input: str | list[dict[str, object]],
|
||||
stream: bool,
|
||||
background: bool,
|
||||
instructions: str | None,
|
||||
previous_response_id: str | None,
|
||||
tools: list[dict[str, object]] | None,
|
||||
reasoning: dict[str, str] | None,
|
||||
model_override: str | None = None,
|
||||
) -> dict[str, object]:
|
||||
"""Build the request body for POST /v1/responses."""
|
||||
body: dict[str, object] = {
|
||||
"model": model,
|
||||
"input": input,
|
||||
"stream": stream,
|
||||
}
|
||||
if background:
|
||||
body["background"] = True
|
||||
if instructions is not None:
|
||||
body["instructions"] = instructions
|
||||
if previous_response_id is not None:
|
||||
body["previous_response_id"] = previous_response_id
|
||||
if tools is not None:
|
||||
body["tools"] = tools
|
||||
if reasoning is not None:
|
||||
body["reasoning"] = reasoning
|
||||
if model_override is not None:
|
||||
body["model_override"] = model_override
|
||||
return body
|
||||
|
||||
|
||||
async def _execute_and_patch(
|
||||
http: httpx.AsyncClient,
|
||||
base_url: str,
|
||||
tool_handler: ToolHandler,
|
||||
hooks: StreamHooks,
|
||||
tool_call: ToolCall,
|
||||
root_response_id: str,
|
||||
iteration: int,
|
||||
) -> None:
|
||||
"""Execute a tunneled sub-agent tool call and PATCH the result back.
|
||||
|
||||
Runs in the background via ``asyncio.ensure_future`` while the
|
||||
SSE stream continues.
|
||||
"""
|
||||
call_info = ToolCallInfo(
|
||||
name=tool_call.name,
|
||||
arguments=tool_call.arguments,
|
||||
call_id=tool_call.call_id,
|
||||
agent_name=tool_call.agent_name,
|
||||
response_id=root_response_id,
|
||||
iteration=iteration,
|
||||
)
|
||||
try:
|
||||
# Sync ``execute`` goes through asyncio.to_thread so a
|
||||
# blocking body doesn't stall the SSE stream or the
|
||||
# TUI's render loop. See ``_call_execute_off_loop``.
|
||||
output = await _call_execute_off_loop(tool_handler, call_info)
|
||||
except ToolCallDenied as exc:
|
||||
output = str(exc)
|
||||
except Exception:
|
||||
_log.exception("Error executing tunneled tool call %s", tool_call.name)
|
||||
output = f"Error executing tool: {tool_call.name}"
|
||||
|
||||
await _call_hook(
|
||||
hooks.on_tool_call_end,
|
||||
ToolCallEndCtx(
|
||||
name=tool_call.name,
|
||||
call_id=tool_call.call_id,
|
||||
agent_name=tool_call.agent_name,
|
||||
output=output,
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
resp = await http.patch(
|
||||
f"{base_url}/v1/responses/{root_response_id}",
|
||||
json={"tool_results": [{"call_id": tool_call.call_id, "output": output}]},
|
||||
timeout=60.0,
|
||||
)
|
||||
if resp.status_code not in (200, 404, 409):
|
||||
_log.warning(
|
||||
"PATCH failed for call_id %s: %s",
|
||||
tool_call.call_id,
|
||||
resp.text[:200],
|
||||
)
|
||||
except Exception:
|
||||
_log.exception("Error PATCHing tool result for call_id %s", tool_call.call_id)
|
||||
|
||||
|
||||
async def _call_execute_off_loop(
|
||||
tool_handler: ToolHandler,
|
||||
call_info: ToolCallInfo,
|
||||
) -> Any:
|
||||
"""
|
||||
Invoke ``tool_handler.execute(call_info)`` without blocking
|
||||
the event loop.
|
||||
|
||||
``execute`` may be ``async def`` (returns a coroutine) or
|
||||
sync ``def`` (returns a plain value). If async, we await
|
||||
the coroutine on the loop; async bodies are already
|
||||
loop-cooperative. If sync, we run it via
|
||||
:func:`asyncio.to_thread` so blocking calls inside —
|
||||
``time.sleep``, file I/O, subprocess, ``requests`` — don't
|
||||
stall the loop.
|
||||
|
||||
Sync-vs-async is detected via
|
||||
:func:`inspect.iscoroutinefunction` so we only invoke
|
||||
``execute`` once. Without the thread bounce, a single
|
||||
sync body with ``time.sleep(5)`` would serialize every
|
||||
concurrent tool call AND freeze any render loop (like the
|
||||
``omnigent chat`` TUI) sharing the event loop.
|
||||
|
||||
:param tool_handler: Handler whose ``execute`` will run.
|
||||
:param call_info: Call context passed through to the handler.
|
||||
:returns: Whatever ``execute`` returns (typically a string).
|
||||
"""
|
||||
if inspect.iscoroutinefunction(tool_handler.execute):
|
||||
return await tool_handler.execute(call_info)
|
||||
return await asyncio.to_thread(tool_handler.execute, call_info)
|
||||
|
||||
|
||||
async def _invoke_elicitation_hook(
|
||||
hooks: StreamHooks,
|
||||
ctx: ElicitationRequestCtx,
|
||||
) -> bool:
|
||||
"""
|
||||
Run the consumer's ``on_elicitation_request`` hook fail-closed.
|
||||
|
||||
Returns True only when the hook explicitly accepts. A missing
|
||||
hook, an exception from the hook, or a falsy return all map to
|
||||
False (decline) — the parked workflow must not stall on an
|
||||
elicitation no consumer can answer.
|
||||
|
||||
:param hooks: Stream hooks — the ``on_elicitation_request``
|
||||
callable to invoke.
|
||||
:param ctx: Context carrying the parsed elicitation, passed to
|
||||
the hook verbatim.
|
||||
:returns: ``True`` iff the hook explicitly accepted; ``False``
|
||||
on missing-hook, exception, or any falsy return.
|
||||
"""
|
||||
if hooks.on_elicitation_request is None:
|
||||
_log.info(
|
||||
"No on_elicitation_request hook registered; declining elicitation_id=%s policy=%s",
|
||||
ctx.elicitation_id,
|
||||
ctx.policy_name,
|
||||
)
|
||||
return False
|
||||
try:
|
||||
raw = hooks.on_elicitation_request(ctx)
|
||||
if inspect.isawaitable(raw):
|
||||
raw = await raw
|
||||
return bool(raw)
|
||||
except Exception:
|
||||
_log.exception(
|
||||
"on_elicitation_request hook raised for elicitation_id %s; declining fail-closed",
|
||||
ctx.elicitation_id,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
async def _post_elicitation_result(
|
||||
http: httpx.AsyncClient,
|
||||
base_url: str,
|
||||
session_id: str,
|
||||
elicitation_id: str,
|
||||
accepted: bool,
|
||||
) -> None:
|
||||
"""
|
||||
POST a verdict to one elicitation's dedicated resolve URL.
|
||||
|
||||
URL-based elicitation: the verdict goes to
|
||||
``POST /v1/sessions/{id}/elicitations/{eid}/resolve`` with the
|
||||
MCP ``ElicitResult`` body (``{"action": "accept"}`` or
|
||||
``{"action": "decline"}``) — the elicitation id rides in the URL
|
||||
path. ``content`` is omitted for binary accept/decline (future
|
||||
richer hooks would populate it). This matches the REPL and web
|
||||
clients; the server routes all of them through the same
|
||||
``_resolve_elicitation``.
|
||||
|
||||
Status codes the server returns (202/404/409) are all benign —
|
||||
the verdict is either applied or moot. Anything else gets a
|
||||
warning log; the call still returns normally because the
|
||||
background task is fire-and-forget.
|
||||
|
||||
:param http: Shared HTTPX client.
|
||||
:param base_url: Server base URL.
|
||||
:param session_id: Session/conversation id that emitted the
|
||||
elicitation. The resolve URL is scoped to this session so
|
||||
the server can apply normal session authorization.
|
||||
:param elicitation_id: Server-assigned elicitation id.
|
||||
:param accepted: Hook verdict — ``True`` → ``"accept"``,
|
||||
``False`` → ``"decline"``.
|
||||
"""
|
||||
if not session_id:
|
||||
_log.warning(
|
||||
"Cannot resolve elicitation_id %s: missing session id",
|
||||
elicitation_id,
|
||||
)
|
||||
return
|
||||
body = {"action": "accept" if accepted else "decline"}
|
||||
try:
|
||||
resp = await http.post(
|
||||
f"{base_url}/v1/sessions/{session_id}/elicitations/{elicitation_id}/resolve",
|
||||
json=body,
|
||||
timeout=60.0,
|
||||
)
|
||||
# 202 = delivered, 404 = session/parked workflow already gone
|
||||
# (cancel race), 409 = already completed. All three are
|
||||
# benign: the verdict is either applied or moot.
|
||||
if resp.status_code not in (202, 404, 409):
|
||||
_log.warning(
|
||||
"Resolve elicitation_id %s failed: %s",
|
||||
elicitation_id,
|
||||
resp.text[:200],
|
||||
)
|
||||
except Exception:
|
||||
_log.exception(
|
||||
"Error POSTing approval event for elicitation_id %s",
|
||||
elicitation_id,
|
||||
)
|
||||
|
||||
|
||||
async def _handle_elicitation_request(
|
||||
http: httpx.AsyncClient,
|
||||
base_url: str,
|
||||
hooks: StreamHooks,
|
||||
event: ElicitationRequest,
|
||||
response_id: str,
|
||||
session_id: str,
|
||||
) -> None:
|
||||
"""
|
||||
Route an elicitation through the hook and POST the verdict.
|
||||
|
||||
Runs on a background task so the SSE stream continues to drain
|
||||
while the user is deciding. Two-step: invoke the hook
|
||||
(:func:`_invoke_elicitation_hook`), POST the approval event
|
||||
(:func:`_post_elicitation_result`).
|
||||
|
||||
:param http: Shared HTTPX client.
|
||||
:param base_url: Server base URL.
|
||||
:param hooks: Stream hooks — the ``on_elicitation_request``
|
||||
hook gets the decision.
|
||||
:param event: The parsed elicitation request.
|
||||
:param response_id: Root response id (the parked workflow);
|
||||
carried into the ctx for the hook's audit/logging only.
|
||||
:param session_id: Session/conversation id that emitted the
|
||||
elicitation. The verdict is posted back as a session
|
||||
``approval`` event.
|
||||
"""
|
||||
ctx = ElicitationRequestCtx(
|
||||
elicitation_id=event.elicitation_id,
|
||||
message=event.message,
|
||||
requested_schema=event.requested_schema,
|
||||
mode=event.mode,
|
||||
phase=event.phase,
|
||||
policy_name=event.policy_name,
|
||||
content_preview=event.content_preview,
|
||||
response_id=response_id,
|
||||
url=event.url,
|
||||
)
|
||||
accepted = await _invoke_elicitation_hook(hooks, ctx)
|
||||
await _post_elicitation_result(http, base_url, session_id, event.elicitation_id, accepted)
|
||||
@@ -0,0 +1,170 @@
|
||||
"""LocalServer — start and manage a local omnigent server."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._client import OmnigentClient
|
||||
|
||||
|
||||
def _find_free_port() -> int:
|
||||
"""Find a free TCP port."""
|
||||
import socket
|
||||
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind(("127.0.0.1", 0))
|
||||
return s.getsockname()[1]
|
||||
|
||||
|
||||
class LocalServer:
|
||||
"""Context manager that starts a local omnigent server.
|
||||
|
||||
Usage::
|
||||
|
||||
async with LocalServer(agent_path="./my-agent/") as server:
|
||||
client = server.client
|
||||
async for event in client.responses.stream(...):
|
||||
...
|
||||
|
||||
:param agent_path: Path to the agent directory or tarball.
|
||||
The server pre-registers this agent at startup.
|
||||
:param host: Bind address.
|
||||
:param port: Port number (0 = auto-assign).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent_path: str,
|
||||
*,
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 0,
|
||||
) -> None:
|
||||
self._agent_path = agent_path
|
||||
self._host = host
|
||||
self._port = port if port != 0 else _find_free_port()
|
||||
self._proc: subprocess.Popen[bytes] | None = None
|
||||
self._tmpdir: str | None = None
|
||||
self._client: OmnigentClient | None = None
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
"""The server's base URL."""
|
||||
return f"http://{self._host}:{self._port}"
|
||||
|
||||
@property
|
||||
def client(self) -> OmnigentClient:
|
||||
"""A pre-configured client connected to this server."""
|
||||
if self._client is None:
|
||||
raise RuntimeError("Server not started — use 'async with LocalServer(...) as server:'")
|
||||
return self._client
|
||||
|
||||
async def __aenter__(self) -> LocalServer:
|
||||
self._start()
|
||||
# Import here to avoid circular import at module level.
|
||||
from ._client import OmnigentClient
|
||||
|
||||
self._client = OmnigentClient(base_url=self.base_url)
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc: object) -> None:
|
||||
self._stop()
|
||||
if self._client is not None:
|
||||
await self._client.close()
|
||||
self._client = None
|
||||
|
||||
def _start(self) -> None:
|
||||
"""Launch the server subprocess."""
|
||||
self._tmpdir = tempfile.mkdtemp(prefix="omnigent-client-")
|
||||
db_uri = f"sqlite:///{self._tmpdir}/chat.db"
|
||||
art_loc = f"{self._tmpdir}/artifacts"
|
||||
|
||||
# Resolve the omnigent project root so Alembic can find
|
||||
# its migrations directory. Walk up from the agent path or
|
||||
# from this file's location looking for omnigent/cli.py.
|
||||
project_root = self._find_project_root()
|
||||
|
||||
self._proc = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"omnigent.cli",
|
||||
"server",
|
||||
"--host",
|
||||
self._host,
|
||||
"--port",
|
||||
str(self._port),
|
||||
"--database-uri",
|
||||
db_uri,
|
||||
"--artifact-location",
|
||||
art_loc,
|
||||
"--agent",
|
||||
os.path.abspath(self._agent_path),
|
||||
],
|
||||
env={**os.environ},
|
||||
cwd=project_root,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
)
|
||||
self._wait_for_ready()
|
||||
|
||||
def _find_project_root(self) -> str:
|
||||
"""Find the omnigent project root directory.
|
||||
|
||||
Walks up from the agent path looking for a directory that
|
||||
contains ``omnigent/cli.py``.
|
||||
"""
|
||||
# Try from agent path first.
|
||||
candidate = pathlib.Path(self._agent_path).resolve()
|
||||
for parent in [candidate, *list(candidate.parents)]:
|
||||
if (parent / "omnigent" / "cli.py").exists():
|
||||
return str(parent)
|
||||
|
||||
# Try from this file's location (inside frontends/clients/python/).
|
||||
candidate = pathlib.Path(__file__).resolve()
|
||||
for parent in candidate.parents:
|
||||
if (parent / "omnigent" / "cli.py").exists():
|
||||
return str(parent)
|
||||
|
||||
# Fallback to cwd.
|
||||
return os.getcwd()
|
||||
|
||||
def _wait_for_ready(self, timeout: float = 15.0) -> None:
|
||||
"""Poll until the server responds."""
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
if self._proc is not None and self._proc.poll() is not None:
|
||||
out = (
|
||||
self._proc.stdout.read().decode(errors="replace") if self._proc.stdout else ""
|
||||
)
|
||||
raise RuntimeError(
|
||||
f"Server exited with code {self._proc.returncode}.\n{out[-3000:]}"
|
||||
)
|
||||
try:
|
||||
resp = httpx.get(f"{self.base_url}/health", timeout=2.0)
|
||||
if resp.status_code == 200:
|
||||
return
|
||||
except httpx.ConnectError:
|
||||
pass
|
||||
time.sleep(0.5)
|
||||
raise RuntimeError(f"Server did not start within {timeout}s")
|
||||
|
||||
def _stop(self) -> None:
|
||||
"""Stop the server subprocess."""
|
||||
if self._proc is None:
|
||||
return
|
||||
self._proc.send_signal(signal.SIGINT)
|
||||
try:
|
||||
self._proc.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
self._proc.kill()
|
||||
self._proc = None
|
||||
@@ -0,0 +1,473 @@
|
||||
"""Session helper — tracks conversation state for interactive use."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import mimetypes
|
||||
import pathlib
|
||||
from collections.abc import AsyncIterator, Callable, Iterator
|
||||
|
||||
# Import at the type level to avoid circular imports at runtime.
|
||||
from typing import TYPE_CHECKING, Any, Literal, overload
|
||||
|
||||
from ._events import (
|
||||
ResponseCancelled,
|
||||
ResponseCompleted,
|
||||
ResponseCreated,
|
||||
ResponseFailed,
|
||||
ResponseIncomplete,
|
||||
StreamEvent,
|
||||
)
|
||||
from ._query import QueryResult, QueryStream
|
||||
from ._tool_handler import StreamHooks, ToolHandler
|
||||
from ._types import File, Response
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._client import OmnigentClient
|
||||
|
||||
|
||||
_TERMINAL_STATUSES = frozenset({"completed", "failed", "incomplete", "cancelled"})
|
||||
_VALID_REASONING_EFFORTS = frozenset({"none", "minimal", "low", "medium", "high", "xhigh", "max"})
|
||||
|
||||
|
||||
class Session:
|
||||
"""Tracks conversation state for interactive use.
|
||||
|
||||
Holds the model name, last response ID, and whether the current
|
||||
response is still running. ``send()`` automatically steers if
|
||||
a response is in progress, or starts a new turn if the response
|
||||
is terminal.
|
||||
|
||||
:param client: The underlying :class:`OmnigentClient`.
|
||||
:param model: Agent name to use for requests.
|
||||
:param tool_handler: Optional client-side tool execution config.
|
||||
:param hooks: Optional lifecycle hooks.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: OmnigentClient,
|
||||
model: str,
|
||||
tool_handler: ToolHandler | None = None,
|
||||
hooks: StreamHooks | None = None,
|
||||
) -> None:
|
||||
self._client = client
|
||||
self._model = model
|
||||
self._tool_handler = tool_handler
|
||||
self._hooks = hooks
|
||||
self._previous_response_id: str | None = None
|
||||
self._current_response_id: str | None = None
|
||||
self._is_terminal: bool = True
|
||||
self._reasoning_effort: str | None = None
|
||||
# Session-local /model override. Forwarded as model_override
|
||||
# on POST /v1/responses. Mirrors _reasoning_effort's shape.
|
||||
self._model_override: str | None = None
|
||||
# Lazily populated from the first /v1/responses event so
|
||||
# produced-file lookups can use the session-scoped file endpoints.
|
||||
self._conversation_id: str | None = None
|
||||
|
||||
@property
|
||||
def model(self) -> str:
|
||||
"""The agent name for this session."""
|
||||
return self._model
|
||||
|
||||
@property
|
||||
def current_response_id(self) -> str | None:
|
||||
"""The most recent response ID, or None if no messages sent."""
|
||||
return self._current_response_id
|
||||
|
||||
@property
|
||||
def is_streaming(self) -> bool:
|
||||
"""True if a response is currently in progress."""
|
||||
return not self._is_terminal
|
||||
|
||||
@property
|
||||
def reasoning_effort(self) -> str | None:
|
||||
"""Current per-request reasoning-effort override, or None for default."""
|
||||
return self._reasoning_effort
|
||||
|
||||
def set_reasoning_effort(self, effort: str | None) -> None:
|
||||
"""Set or clear the per-request reasoning-effort override."""
|
||||
if effort is not None:
|
||||
normalized = effort.strip().lower()
|
||||
if normalized not in _VALID_REASONING_EFFORTS:
|
||||
valid = ", ".join(sorted(_VALID_REASONING_EFFORTS))
|
||||
raise ValueError(f"Invalid reasoning effort {effort!r}; expected one of: {valid}")
|
||||
effort = normalized
|
||||
self._reasoning_effort = effort
|
||||
|
||||
def _reasoning_request(self) -> dict[str, str] | None:
|
||||
"""Return the Responses API reasoning payload for this session."""
|
||||
if self._reasoning_effort is None:
|
||||
return None
|
||||
return {"effort": self._reasoning_effort}
|
||||
|
||||
@property
|
||||
def model_override(self) -> str | None:
|
||||
"""Current per-request LLM model override, or ``None`` for agent default."""
|
||||
return self._model_override
|
||||
|
||||
def set_model_override(self, model: str | None) -> None:
|
||||
"""
|
||||
Set or clear the per-request LLM model override.
|
||||
|
||||
:param model: A non-empty model identifier, e.g.
|
||||
``"databricks-claude-sonnet-4-6"``. Whitespace is trimmed;
|
||||
``None`` clears so subsequent requests fall back to the
|
||||
agent spec's model.
|
||||
:raises ValueError: When *model* is empty after trimming. No
|
||||
other validation — server / runtime surfaces unknown-model
|
||||
errors at request time.
|
||||
"""
|
||||
if model is not None:
|
||||
normalized = model.strip()
|
||||
if not normalized:
|
||||
raise ValueError("model override must be a non-empty string")
|
||||
model = normalized
|
||||
self._model_override = model
|
||||
|
||||
async def send(
|
||||
self,
|
||||
input: str | list[dict[str, object]],
|
||||
*,
|
||||
files: list[str] | None = None,
|
||||
instructions: str | None = None,
|
||||
) -> AsyncIterator[StreamEvent]:
|
||||
"""Send a message — auto-steers if a response is in progress.
|
||||
|
||||
Always returns an async iterator. The caller always does
|
||||
``async for event in session.send(text): ...`` regardless
|
||||
of whether it steered or started a new turn.
|
||||
|
||||
Three cases:
|
||||
|
||||
1. **Response in progress, steer delivered**: The server
|
||||
accepted the message into the running agent's inbox.
|
||||
Yields nothing — the existing stream (from the original
|
||||
``send()`` call) will surface the agent's reaction.
|
||||
|
||||
2. **Response in progress, agent already finished**: The
|
||||
server created a new response instead of steering.
|
||||
Yields the full event stream for that new response.
|
||||
|
||||
3. **No response in progress**: Starts a new turn. Yields
|
||||
the full event stream.
|
||||
|
||||
:param input: User text or content block list.
|
||||
:param files: Optional file paths to upload and attach.
|
||||
:param instructions: Per-request system instructions.
|
||||
:yields: Stream events.
|
||||
"""
|
||||
if files:
|
||||
input = await self._build_input_with_files(input, files)
|
||||
|
||||
# Auto-steer if response is in progress.
|
||||
if not self._is_terminal and self._current_response_id is not None:
|
||||
steer_resp = await self._client.responses.steer(
|
||||
self._current_response_id,
|
||||
input if isinstance(input, str) else str(input),
|
||||
model=self._model,
|
||||
reasoning=self._reasoning_request(),
|
||||
model_override=self._model_override,
|
||||
)
|
||||
if steer_resp.id == self._current_response_id:
|
||||
# Case 1: steering delivered. Nothing to yield.
|
||||
return
|
||||
yield
|
||||
|
||||
# Case 2: agent finished — server created a new response.
|
||||
# Stream it like a normal turn. The input was already
|
||||
# included in the steer POST, so the new response has it.
|
||||
async for event in self._stream_and_track(steer_resp.id, instructions):
|
||||
yield event
|
||||
return
|
||||
|
||||
# Case 3: no response in progress — new turn.
|
||||
async for event in self._stream_and_track(
|
||||
None,
|
||||
instructions,
|
||||
input=input,
|
||||
):
|
||||
yield event
|
||||
|
||||
async def _stream_and_track(
|
||||
self,
|
||||
previous_response_id: str | None,
|
||||
instructions: str | None,
|
||||
*,
|
||||
input: str | list[dict[str, object]] | None = None,
|
||||
) -> AsyncIterator[StreamEvent]:
|
||||
"""Stream a response and track session state.
|
||||
|
||||
If ``input`` is None, uses an empty string (for steer-fallback
|
||||
where the input was already sent via the steer POST). If
|
||||
``previous_response_id`` is None, uses the session's stored ID.
|
||||
"""
|
||||
self._is_terminal = False
|
||||
prev_id = (
|
||||
previous_response_id
|
||||
if previous_response_id is not None
|
||||
else self._previous_response_id
|
||||
)
|
||||
actual_input: str | list[dict[str, object]] = input if input is not None else ""
|
||||
|
||||
async for event in self._client.responses.stream(
|
||||
model=self._model,
|
||||
input=actual_input,
|
||||
previous_response_id=prev_id,
|
||||
tool_handler=self._tool_handler,
|
||||
hooks=self._hooks,
|
||||
instructions=instructions,
|
||||
reasoning=self._reasoning_request(),
|
||||
model_override=self._model_override,
|
||||
):
|
||||
if isinstance(event, ResponseCreated):
|
||||
self._current_response_id = event.response.id
|
||||
if event.response.conversation is not None:
|
||||
self._conversation_id = event.response.conversation.id
|
||||
|
||||
if isinstance(
|
||||
event, ResponseCompleted | ResponseFailed | ResponseIncomplete | ResponseCancelled
|
||||
):
|
||||
self._is_terminal = True
|
||||
self._previous_response_id = event.response.id
|
||||
self._current_response_id = event.response.id
|
||||
|
||||
yield event
|
||||
|
||||
@overload
|
||||
async def query(
|
||||
self,
|
||||
input: str | list[dict[str, object]],
|
||||
*,
|
||||
files: list[str] | None = ...,
|
||||
tools: list[Callable[..., Any]] | None = ...,
|
||||
stream: Literal[False] = ...,
|
||||
) -> QueryResult: ...
|
||||
|
||||
@overload
|
||||
async def query(
|
||||
self,
|
||||
input: str | list[dict[str, object]],
|
||||
*,
|
||||
files: list[str] | None = ...,
|
||||
tools: list[Callable[..., Any]] | None = ...,
|
||||
stream: Literal[True],
|
||||
) -> QueryStream: ...
|
||||
|
||||
async def query(
|
||||
self,
|
||||
input: str | list[dict[str, object]],
|
||||
*,
|
||||
files: list[str] | None = None,
|
||||
tools: list[Callable[..., Any]] | None = None,
|
||||
stream: bool = False,
|
||||
) -> QueryResult | QueryStream:
|
||||
"""Send a prompt and get text (plus any files) back.
|
||||
|
||||
Non-streaming (default) returns a :class:`QueryResult`::
|
||||
|
||||
result = await session.query("make me a chart")
|
||||
print(result.text)
|
||||
for f in result.files:
|
||||
await client.files.for_session("<session-id>").download(
|
||||
f.id, f"./out/{f.filename}"
|
||||
)
|
||||
|
||||
Streaming returns a :class:`QueryStream`::
|
||||
|
||||
stream = await session.query("hello", stream=True)
|
||||
async for chunk in stream:
|
||||
print(chunk, end="", flush=True)
|
||||
# After iteration, stream.files holds the produced files.
|
||||
|
||||
Client-side tools can be passed per-call via ``tools=``, or
|
||||
configured session-wide via ``client.session(tool_handler=...)``.
|
||||
If this turn's ``tools=`` is given, it OVERRIDES any session
|
||||
handler for this call only.
|
||||
|
||||
:param input: User text or a list of content-block dicts,
|
||||
e.g. ``"hello"`` or
|
||||
``[{"type": "input_text", "text": "hi"}]``.
|
||||
:param files: Optional list of local file paths to upload and
|
||||
attach to the turn, e.g. ``["./data.csv"]``.
|
||||
:param tools: Optional list of ``@tool``-decorated functions
|
||||
the agent may call on this turn. Overrides the session's
|
||||
configured ``tool_handler`` for this call only.
|
||||
:param stream: If True, return a :class:`QueryStream`. If
|
||||
False (default), return a :class:`QueryResult`.
|
||||
:returns: :class:`QueryResult` (``stream=False``) or
|
||||
:class:`QueryStream` (``stream=True``).
|
||||
:raises OmnigentError: If the response ends in an error.
|
||||
"""
|
||||
if stream:
|
||||
return self._stream_query(input, files=files, tools=tools)
|
||||
return await self._collect_query(input, files=files, tools=tools)
|
||||
|
||||
async def _collect_query(
|
||||
self,
|
||||
input: str | list[dict[str, object]],
|
||||
*,
|
||||
files: list[str] | None,
|
||||
tools: list[Callable[..., Any]] | None,
|
||||
) -> QueryResult:
|
||||
"""Run a turn; return final text + produced files as a QueryResult."""
|
||||
# Local imports avoid a circular dep at module load — _stream
|
||||
# imports _session under TYPE_CHECKING.
|
||||
from ._blocks import FileBlock, TextDone
|
||||
from ._stream import BlockStream
|
||||
from ._transforms import merge_text_across_iterations, pipe, skip_intermediate_ends
|
||||
|
||||
with _per_call_tool_override(self, tools):
|
||||
block_stream = BlockStream()
|
||||
final_text = ""
|
||||
produced: list[File] = []
|
||||
async for block in pipe(
|
||||
block_stream.stream(self, input, files=files),
|
||||
# Merges per-iteration text into one TextDone per response,
|
||||
# so we don't truncate to just the last iteration's text.
|
||||
merge_text_across_iterations(),
|
||||
skip_intermediate_ends(),
|
||||
):
|
||||
if isinstance(block, TextDone):
|
||||
final_text = block.full_text
|
||||
elif isinstance(block, FileBlock):
|
||||
produced.append(await self._get_file(block.file_id))
|
||||
return QueryResult(text=final_text, files=produced)
|
||||
|
||||
def _stream_query(
|
||||
self,
|
||||
input: str | list[dict[str, object]],
|
||||
*,
|
||||
files: list[str] | None,
|
||||
tools: list[Callable[..., Any]] | None,
|
||||
) -> QueryStream:
|
||||
"""Build a QueryStream that yields text chunks and collects files."""
|
||||
produced: list[File] = []
|
||||
chunks = self._stream_chunks(input, files=files, tools=tools, produced=produced)
|
||||
return QueryStream(chunks=chunks, files=produced)
|
||||
|
||||
async def _stream_chunks(
|
||||
self,
|
||||
input: str | list[dict[str, object]],
|
||||
*,
|
||||
files: list[str] | None,
|
||||
tools: list[Callable[..., Any]] | None,
|
||||
produced: list[File],
|
||||
) -> AsyncIterator[str]:
|
||||
"""Async generator backing QueryStream.
|
||||
|
||||
Yields text chunks as they arrive. Appends produced files to
|
||||
the shared ``produced`` list as :class:`FileBlock` events
|
||||
surface, so the owning :class:`QueryStream` sees them through
|
||||
its reference to the same list.
|
||||
"""
|
||||
from ._blocks import FileBlock, TextChunk
|
||||
from ._stream import BlockStream
|
||||
from ._transforms import pipe, skip_intermediate_ends
|
||||
|
||||
with _per_call_tool_override(self, tools):
|
||||
block_stream = BlockStream()
|
||||
async for block in pipe(
|
||||
block_stream.stream(self, input, files=files),
|
||||
skip_intermediate_ends(),
|
||||
):
|
||||
if isinstance(block, TextChunk):
|
||||
yield block.text
|
||||
elif isinstance(block, FileBlock):
|
||||
produced.append(await self._get_file(block.file_id))
|
||||
|
||||
async def _get_file(self, file_id: str) -> File:
|
||||
"""Fetch produced-file metadata from this session's file namespace."""
|
||||
if self._conversation_id is None:
|
||||
raise RuntimeError("cannot fetch produced file before conversation is known")
|
||||
return await self._client.files.for_session(self._conversation_id).get(file_id)
|
||||
|
||||
async def cancel(self) -> Response | None:
|
||||
"""Cancel the current in-progress response.
|
||||
|
||||
:returns: The cancelled response, or None if no response is active.
|
||||
"""
|
||||
if self._current_response_id is None:
|
||||
return None
|
||||
response = await self._client.responses.cancel(self._current_response_id)
|
||||
self._is_terminal = True
|
||||
self._previous_response_id = response.id
|
||||
return response
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset the session — start a new conversation."""
|
||||
self._previous_response_id = None
|
||||
self._current_response_id = None
|
||||
self._is_terminal = True
|
||||
|
||||
def resume_from_response(self, response_id: str) -> None:
|
||||
"""Resume conversation from a specific response ID."""
|
||||
self._previous_response_id = response_id
|
||||
self._current_response_id = response_id
|
||||
self._is_terminal = True
|
||||
|
||||
async def _build_input_with_files(
|
||||
self,
|
||||
text: str | list[dict[str, object]],
|
||||
file_paths: list[str],
|
||||
) -> list[dict[str, object]]:
|
||||
"""Upload files and build content blocks."""
|
||||
blocks: list[dict[str, object]] = []
|
||||
|
||||
# Add text block.
|
||||
if isinstance(text, str) and text:
|
||||
blocks.append({"type": "input_text", "text": text})
|
||||
elif isinstance(text, list):
|
||||
blocks.extend(text)
|
||||
|
||||
# Upload and add file blocks.
|
||||
for path in file_paths:
|
||||
if self._conversation_id is None:
|
||||
raise RuntimeError(
|
||||
"file attachments require an established conversation; "
|
||||
"use the sessions API for first-turn attachments"
|
||||
)
|
||||
uploaded = await self._client.files.for_session(self._conversation_id).upload(path)
|
||||
content_type = mimetypes.guess_type(path)[0]
|
||||
if content_type and content_type.startswith("image/"):
|
||||
blocks.append({"type": "input_image", "file_id": uploaded.id})
|
||||
else:
|
||||
blocks.append(
|
||||
{
|
||||
"type": "input_file",
|
||||
"file_id": uploaded.id,
|
||||
"filename": pathlib.Path(path).name,
|
||||
}
|
||||
)
|
||||
|
||||
return blocks
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _per_call_tool_override(
|
||||
session: Session,
|
||||
tools: list[Callable[..., Any]] | None,
|
||||
) -> Iterator[None]:
|
||||
"""Temporarily override ``session._tool_handler`` for one call.
|
||||
|
||||
If ``tools`` is ``None``, the session's configured handler is
|
||||
used unchanged. Otherwise a handler is built from the decorated
|
||||
functions and swapped in for the duration of the ``with`` block;
|
||||
the original is restored on exit, even on exception.
|
||||
|
||||
:param session: The session whose ``_tool_handler`` to override.
|
||||
:param tools: List of ``@tool``-decorated functions, or ``None``
|
||||
to leave the session's handler in place.
|
||||
"""
|
||||
if tools is None:
|
||||
yield
|
||||
return
|
||||
from .tools import build_tool_handler
|
||||
|
||||
previous = session._tool_handler
|
||||
session._tool_handler = build_tool_handler(tools)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
session._tool_handler = previous
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,366 @@
|
||||
"""SSE frame parser — converts raw byte chunks into typed events.
|
||||
|
||||
Handles the ``event:`` / ``data:`` / ``[DONE]`` framing from the
|
||||
server's ``text/event-stream`` responses.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
from omnigent.server import schemas as _srv_events
|
||||
|
||||
from ._events import (
|
||||
NATIVE_TOOL_TYPES,
|
||||
ClientTaskCancel,
|
||||
CompactionInProgress,
|
||||
ElicitationRequest,
|
||||
ErrorEvent,
|
||||
MessageDone,
|
||||
NativeToolCall,
|
||||
OutputFileDone,
|
||||
ReasoningDelta,
|
||||
ReasoningStarted,
|
||||
ReasoningSummaryDelta,
|
||||
ResponseCancelled,
|
||||
ResponseCompleted,
|
||||
ResponseCreated,
|
||||
ResponseFailed,
|
||||
ResponseIncomplete,
|
||||
ResponseInProgress,
|
||||
ResponseQueued,
|
||||
RetryEvent,
|
||||
StreamEvent,
|
||||
TextDelta,
|
||||
ToolCall,
|
||||
ToolResult,
|
||||
)
|
||||
from ._types import ErrorInfo, Response
|
||||
|
||||
_log = logging.getLogger("omnigent_client.sse")
|
||||
|
||||
|
||||
def _wire_type(cls: type) -> str:
|
||||
"""
|
||||
Extract the wire ``type`` literal from a server event class.
|
||||
|
||||
Each :mod:`omnigent.server.schemas` event class pins its
|
||||
``type`` field as ``Literal["..."]``; this helper unwraps that
|
||||
to a plain string so the SDK's ``str == str`` dispatch table
|
||||
stays a ``str`` comparison rather than introducing a class-side
|
||||
isinstance check.
|
||||
|
||||
:param cls: A subclass of the server's ``_SSEEventBase``,
|
||||
e.g. :class:`omnigent.server.schemas.OutputTextDeltaEvent`.
|
||||
:returns: The wire ``type`` literal, e.g.
|
||||
``"response.output_text.delta"``.
|
||||
"""
|
||||
return cls.model_fields["type"].annotation.__args__[0] # type: ignore[union-attr]
|
||||
|
||||
|
||||
# Wire type literals — pulled from the server's typed source of
|
||||
# truth so a rename there is a one-edit change here as well.
|
||||
_T_RESPONSE_CREATED = _wire_type(_srv_events.CreatedEvent)
|
||||
_T_RESPONSE_QUEUED = _wire_type(_srv_events.QueuedEvent)
|
||||
_T_RESPONSE_IN_PROGRESS = _wire_type(_srv_events.InProgressEvent)
|
||||
_T_RESPONSE_COMPLETED = _wire_type(_srv_events.CompletedEvent)
|
||||
_T_RESPONSE_FAILED = _wire_type(_srv_events.FailedEvent)
|
||||
_T_RESPONSE_INCOMPLETE = _wire_type(_srv_events.IncompleteEvent)
|
||||
_T_RESPONSE_CANCELLED = _wire_type(_srv_events.CancelledEvent)
|
||||
_T_RESPONSE_OUTPUT_TEXT_DELTA = _wire_type(_srv_events.OutputTextDeltaEvent)
|
||||
_T_RESPONSE_REASONING_STARTED = _wire_type(_srv_events.ReasoningStartedEvent)
|
||||
_T_RESPONSE_REASONING_TEXT_DELTA = _wire_type(_srv_events.ReasoningTextDeltaEvent)
|
||||
_T_RESPONSE_REASONING_SUMMARY_TEXT_DELTA = _wire_type(_srv_events.ReasoningSummaryTextDeltaEvent)
|
||||
_T_RESPONSE_OUTPUT_ITEM_DONE = _wire_type(_srv_events.OutputItemDoneEvent)
|
||||
_T_RESPONSE_OUTPUT_FILE_DONE = _wire_type(_srv_events.OutputFileDoneEvent)
|
||||
_T_RESPONSE_RETRY = _wire_type(_srv_events.RetryEvent)
|
||||
_T_RESPONSE_ERROR = _wire_type(_srv_events.ErrorEvent)
|
||||
_T_RESPONSE_COMPACTION_IN_PROGRESS = _wire_type(_srv_events.CompactionInProgressEvent)
|
||||
_T_RESPONSE_CLIENT_TASK_CANCEL = _wire_type(_srv_events.ClientTaskCancelEvent)
|
||||
_T_RESPONSE_ELICITATION_REQUEST = _wire_type(_srv_events.ElicitationRequestEvent)
|
||||
|
||||
|
||||
async def parse_sse_stream(
|
||||
byte_stream: AsyncIterator[bytes],
|
||||
) -> AsyncIterator[StreamEvent]:
|
||||
"""Parse an SSE byte stream into typed events.
|
||||
|
||||
:param byte_stream: Raw bytes from ``httpx.Response.aiter_bytes()``.
|
||||
:yields: Typed :class:`StreamEvent` instances.
|
||||
"""
|
||||
buf = ""
|
||||
current_event: str | None = None
|
||||
|
||||
async for chunk in byte_stream:
|
||||
buf += chunk.decode("utf-8", errors="replace")
|
||||
while "\n" in buf:
|
||||
line, buf = buf.split("\n", 1)
|
||||
line = line.rstrip("\r")
|
||||
|
||||
if line.startswith("event: "):
|
||||
current_event = line[7:]
|
||||
elif line.startswith("data: "):
|
||||
data_str = line[6:]
|
||||
if data_str.strip() == "[DONE]":
|
||||
# The server's terminal sentinel is a bare ``data: [DONE]``
|
||||
# with no preceding ``event:`` line, so detect it regardless
|
||||
# of ``current_event``.
|
||||
return
|
||||
# A ``data:`` line is only a parseable event when an ``event:``
|
||||
# line preceded it; a lone data line (other than [DONE]) is
|
||||
# ignored.
|
||||
if current_event is not None:
|
||||
try:
|
||||
data = json.loads(data_str)
|
||||
except json.JSONDecodeError:
|
||||
_log.warning("Failed to parse SSE data: %s", data_str[:200])
|
||||
current_event = None
|
||||
continue
|
||||
event = _parse_event(current_event, data)
|
||||
if event is not None:
|
||||
yield event
|
||||
current_event = None
|
||||
elif line == "":
|
||||
current_event = None
|
||||
|
||||
|
||||
def _normalize_event_type(event_type: str) -> str:
|
||||
"""Normalize event type to handle server enum rendering.
|
||||
|
||||
The server builds terminal events as ``f"response.{task.status}"``
|
||||
where ``task.status`` may be a Python enum (rendering as
|
||||
``response.TaskStatus.COMPLETED``) instead of the expected
|
||||
``response.completed``. Normalize by extracting and lowercasing
|
||||
the enum value.
|
||||
"""
|
||||
if ".TaskStatus." in event_type:
|
||||
# "response.TaskStatus.COMPLETED" → "response.completed"
|
||||
parts = event_type.split(".")
|
||||
status = parts[-1].lower()
|
||||
return f"response.{status}"
|
||||
return event_type
|
||||
|
||||
|
||||
def _parse_event(event_type: str, data: dict[str, Any]) -> StreamEvent | None:
|
||||
"""Convert a raw SSE event type + JSON data into a typed event.
|
||||
|
||||
:param event_type: Wire name of the SSE ``event:`` field, e.g.
|
||||
``"response.output_text.delta"``. Compared against the
|
||||
``_T_RESPONSE_*`` wire-type constants (sourced from
|
||||
:mod:`omnigent.server.schemas`) to dispatch.
|
||||
:param data: Decoded JSON payload from the SSE ``data:`` field.
|
||||
:returns: A typed :class:`StreamEvent` for known event names, or
|
||||
``None`` when the payload is missing required fields or the
|
||||
event type is unrecognized (forward-compatible skip).
|
||||
"""
|
||||
event_type = _normalize_event_type(event_type)
|
||||
|
||||
# Response lifecycle
|
||||
if event_type == _T_RESPONSE_CREATED:
|
||||
return ResponseCreated(response=_parse_response(data))
|
||||
if event_type == _T_RESPONSE_QUEUED:
|
||||
return ResponseQueued(response=_parse_response(data))
|
||||
if event_type == _T_RESPONSE_IN_PROGRESS:
|
||||
return ResponseInProgress(response=_parse_response(data))
|
||||
if event_type == _T_RESPONSE_COMPLETED:
|
||||
return ResponseCompleted(response=_parse_response(data))
|
||||
if event_type == _T_RESPONSE_FAILED:
|
||||
return ResponseFailed(response=_parse_response(data))
|
||||
if event_type == _T_RESPONSE_INCOMPLETE:
|
||||
resp = _parse_response(data)
|
||||
reason = ""
|
||||
if resp.incomplete_details is not None:
|
||||
reason = resp.incomplete_details.reason
|
||||
return ResponseIncomplete(response=resp, reason=reason)
|
||||
if event_type == _T_RESPONSE_CANCELLED:
|
||||
return ResponseCancelled(response=_parse_response(data))
|
||||
|
||||
# Text streaming
|
||||
if event_type == _T_RESPONSE_OUTPUT_TEXT_DELTA:
|
||||
delta = data.get("delta")
|
||||
if isinstance(delta, str):
|
||||
return TextDelta(delta=delta)
|
||||
return None
|
||||
|
||||
# Reasoning
|
||||
if event_type == _T_RESPONSE_REASONING_STARTED:
|
||||
return ReasoningStarted()
|
||||
if event_type == _T_RESPONSE_REASONING_TEXT_DELTA:
|
||||
delta = data.get("delta")
|
||||
if isinstance(delta, str):
|
||||
return ReasoningDelta(delta=delta)
|
||||
return None
|
||||
if event_type == _T_RESPONSE_REASONING_SUMMARY_TEXT_DELTA:
|
||||
delta = data.get("delta")
|
||||
if isinstance(delta, str):
|
||||
return ReasoningSummaryDelta(delta=delta)
|
||||
return None
|
||||
|
||||
# Output items
|
||||
if event_type == _T_RESPONSE_OUTPUT_ITEM_DONE:
|
||||
return _parse_output_item(data)
|
||||
|
||||
# File output
|
||||
if event_type == _T_RESPONSE_OUTPUT_FILE_DONE:
|
||||
return OutputFileDone(
|
||||
file_id=str(data.get("file_id", "")),
|
||||
filename=str(data["filename"]) if data.get("filename") is not None else None,
|
||||
content_type=str(data["content_type"])
|
||||
if data.get("content_type") is not None
|
||||
else None,
|
||||
)
|
||||
|
||||
# Retry
|
||||
if event_type == _T_RESPONSE_RETRY:
|
||||
return RetryEvent(
|
||||
source=str(data.get("source", "")),
|
||||
tool_name=str(data["tool_name"]) if data.get("tool_name") is not None else None,
|
||||
attempt=int(data.get("attempt", 0)),
|
||||
max_attempts=int(data.get("max_attempts", 0)),
|
||||
delay_seconds=float(data.get("delay_seconds", 0.0)),
|
||||
error=_parse_error_info(data.get("error", {})),
|
||||
)
|
||||
|
||||
# Error
|
||||
if event_type == _T_RESPONSE_ERROR:
|
||||
return ErrorEvent(
|
||||
source=str(data.get("source", "")),
|
||||
tool_name=str(data["tool_name"]) if data.get("tool_name") is not None else None,
|
||||
error=_parse_error_info(data.get("error", {})),
|
||||
)
|
||||
|
||||
# Compaction
|
||||
if event_type == _T_RESPONSE_COMPACTION_IN_PROGRESS:
|
||||
return CompactionInProgress()
|
||||
|
||||
# Async client-tool cancel notification
|
||||
if event_type == _T_RESPONSE_CLIENT_TASK_CANCEL:
|
||||
task_id = data.get("task_id")
|
||||
if isinstance(task_id, str) and task_id:
|
||||
raw_call_id = data.get("call_id")
|
||||
call_id = raw_call_id if isinstance(raw_call_id, str) and raw_call_id else None
|
||||
return ClientTaskCancel(task_id=task_id, call_id=call_id)
|
||||
_log.warning("response.client_task.cancel missing task_id: %r", data)
|
||||
return None
|
||||
|
||||
# MCP-shape elicitation request (POLICIES.md §7 + the
|
||||
# universal API additions in
|
||||
# designs/SERVER_HARNESS_CONTRACT.md). The ``params`` block
|
||||
# mirrors MCP's ``ElicitRequestFormParams`` field-for-field;
|
||||
# we surface it as a typed event so the stream consumer can
|
||||
# route it to the elicitation hook (not into a ToolHandler).
|
||||
if event_type == "response.elicitation_request":
|
||||
elicitation_id = data.get("elicitation_id")
|
||||
if not isinstance(elicitation_id, str) or not elicitation_id:
|
||||
_log.warning(
|
||||
"response.elicitation_request missing elicitation_id: %r",
|
||||
data,
|
||||
)
|
||||
return None
|
||||
params = data.get("params")
|
||||
if not isinstance(params, dict):
|
||||
_log.warning(
|
||||
"response.elicitation_request missing params dict: %r",
|
||||
data,
|
||||
)
|
||||
return None
|
||||
requested_schema = params.get("requestedSchema")
|
||||
return ElicitationRequest(
|
||||
elicitation_id=elicitation_id,
|
||||
message=str(params.get("message") or ""),
|
||||
# MCP spec restricts requestedSchema to a JSON-Schema
|
||||
# dict — we accept anything dict-shaped and pass it
|
||||
# through; non-dict inputs (defensive against
|
||||
# malformed producers) become an empty schema.
|
||||
requested_schema=requested_schema if isinstance(requested_schema, dict) else {},
|
||||
mode=str(params.get("mode") or "form"),
|
||||
phase=str(params.get("phase") or ""),
|
||||
policy_name=str(params.get("policy_name") or ""),
|
||||
content_preview=str(params.get("content_preview") or ""),
|
||||
url=str(params["url"]) if isinstance(params.get("url"), str) else None,
|
||||
)
|
||||
|
||||
# Unknown event — skip gracefully for forward-compatibility
|
||||
_log.debug("Skipping unknown SSE event type: %s", event_type)
|
||||
return None
|
||||
|
||||
|
||||
def _parse_output_item(data: dict[str, Any]) -> StreamEvent | None:
|
||||
"""Parse a ``response.output_item.done`` event into a typed event."""
|
||||
item = data.get("item")
|
||||
if not isinstance(item, dict):
|
||||
return None
|
||||
|
||||
item_type = item.get("type", "")
|
||||
|
||||
if item_type == "function_call":
|
||||
args_str = str(item.get("arguments", "{}"))
|
||||
try:
|
||||
arguments = json.loads(args_str)
|
||||
except json.JSONDecodeError:
|
||||
arguments = {}
|
||||
name = str(item.get("name", ""))
|
||||
call_id = str(item.get("call_id", ""))
|
||||
return ToolCall(
|
||||
name=name,
|
||||
arguments=arguments,
|
||||
call_id=call_id,
|
||||
status=str(item.get("status", "")),
|
||||
agent_name=str(item.get("model", "")),
|
||||
)
|
||||
|
||||
if item_type == "function_call_output":
|
||||
raw_arguments = item.get("arguments")
|
||||
arguments: dict[str, object] = {}
|
||||
if isinstance(raw_arguments, dict):
|
||||
arguments = raw_arguments
|
||||
elif isinstance(raw_arguments, str) and raw_arguments:
|
||||
try:
|
||||
decoded_arguments = json.loads(raw_arguments)
|
||||
except json.JSONDecodeError:
|
||||
decoded_arguments = {}
|
||||
if isinstance(decoded_arguments, dict):
|
||||
arguments = decoded_arguments
|
||||
return ToolResult(
|
||||
call_id=str(item.get("call_id", "")),
|
||||
output=str(item.get("output", "")),
|
||||
arguments=arguments,
|
||||
)
|
||||
|
||||
if item_type == "message":
|
||||
content = item.get("content", [])
|
||||
return MessageDone(
|
||||
content=content if isinstance(content, list) else [],
|
||||
)
|
||||
|
||||
if item_type in NATIVE_TOOL_TYPES:
|
||||
return NativeToolCall(
|
||||
tool_type=item_type,
|
||||
data=item,
|
||||
)
|
||||
|
||||
# Compaction items, reasoning items, etc. — skip
|
||||
_log.debug("Skipping output item type: %s", item_type)
|
||||
return None
|
||||
|
||||
|
||||
def _parse_response(data: dict[str, Any]) -> Response:
|
||||
"""Extract the Response object from an SSE event payload."""
|
||||
resp_data = data.get("response")
|
||||
if isinstance(resp_data, dict):
|
||||
return Response.from_dict(resp_data)
|
||||
# Some events put fields at the top level
|
||||
return Response.from_dict(data)
|
||||
|
||||
|
||||
def _parse_error_info(raw: Any) -> ErrorInfo:
|
||||
"""Parse an ErrorInfo from a nested dict."""
|
||||
if isinstance(raw, dict):
|
||||
return ErrorInfo(
|
||||
code=str(raw.get("code", "")),
|
||||
message=str(raw.get("message", "")),
|
||||
)
|
||||
return ErrorInfo(code="", message=str(raw))
|
||||
@@ -0,0 +1,615 @@
|
||||
"""BlockStream — event dispatch state machine that emits semantic blocks.
|
||||
|
||||
Consumes the raw event stream from ``session.send()`` and produces
|
||||
a stream of typed blocks with context. Each block carries a
|
||||
``BlockContext`` identifying which agent produced it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from ._blocks import (
|
||||
AnyBlock,
|
||||
BlockContext,
|
||||
CompactionBlock,
|
||||
ErrorBlock,
|
||||
FileBlock,
|
||||
NativeToolBlock,
|
||||
ReasoningBlock,
|
||||
ReasoningChunk,
|
||||
ReasoningStartBlock,
|
||||
ResponseEndBlock,
|
||||
ResponseStartBlock,
|
||||
RetryBlock,
|
||||
TextChunk,
|
||||
TextDone,
|
||||
ToolExecution,
|
||||
ToolGroup,
|
||||
ToolResultBlock,
|
||||
)
|
||||
from ._events import (
|
||||
CompactionInProgress,
|
||||
ErrorEvent,
|
||||
MessageDone,
|
||||
NativeToolCall,
|
||||
OutputFileDone,
|
||||
ReasoningDelta,
|
||||
ReasoningStarted,
|
||||
ReasoningSummaryDelta,
|
||||
ResponseCancelled,
|
||||
ResponseCompleted,
|
||||
ResponseCreated,
|
||||
ResponseFailed,
|
||||
ResponseIncomplete,
|
||||
ResponseInProgress,
|
||||
ResponseQueued,
|
||||
RetryEvent,
|
||||
TextDelta,
|
||||
ToolCall,
|
||||
ToolResult,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ._session import Session
|
||||
|
||||
|
||||
def format_tool_args_brief(name: str, arguments: dict[str, object]) -> str:
|
||||
"""
|
||||
Format tool arguments for inline display next to the ``⏵ <name>``
|
||||
line.
|
||||
|
||||
Public so frontends that re-render historical tool calls (e.g. the
|
||||
REPL's conversation-resume preview) can produce the same
|
||||
``args_summary`` string the live stream produced when the call
|
||||
first ran. Without a single source of truth, "this is what you
|
||||
saw originally" diverges from "this is what the renderer
|
||||
chooses now" the moment either side changes.
|
||||
|
||||
:param name: Tool name, e.g. ``"Read"``.
|
||||
:param arguments: Parsed arguments dict, e.g.
|
||||
``{"file_path": "/x/y.py"}``.
|
||||
:returns: One-line summary, e.g. ``"y.py"`` for ``Read`` or the
|
||||
tool's primary user-meaningful field. Falls back to a JSON
|
||||
encoding of the full dict (truncated to 80 chars) for
|
||||
unrecognized tools.
|
||||
"""
|
||||
if not arguments:
|
||||
return ""
|
||||
_KEYS = {
|
||||
"Read": "file_path",
|
||||
"Write": "file_path",
|
||||
"Edit": "file_path",
|
||||
"Bash": "command",
|
||||
"Glob": "pattern",
|
||||
"Grep": "pattern",
|
||||
"web_search": "query",
|
||||
}
|
||||
key = _KEYS.get(name)
|
||||
if key and key in arguments:
|
||||
s = str(arguments[key])
|
||||
if key == "file_path" and "/" in s:
|
||||
s = s.rsplit("/", 1)[-1]
|
||||
return s[:80] + "…" if len(s) > 80 else s
|
||||
try:
|
||||
s = json.dumps(arguments, ensure_ascii=False)
|
||||
except (TypeError, ValueError):
|
||||
s = str(arguments)
|
||||
return s[:80] + "…" if len(s) > 80 else s
|
||||
|
||||
|
||||
def _format_native_label(tool_type: str, data: dict[str, object]) -> str:
|
||||
"""Format a native tool label."""
|
||||
if tool_type == "web_search_call":
|
||||
action = data.get("action")
|
||||
if isinstance(action, dict):
|
||||
at = action.get("type", "")
|
||||
if at == "search":
|
||||
return f"web search: {str(action.get('query', ''))[:80]}"
|
||||
if at == "open_page":
|
||||
return f"web open: {str(action.get('url', ''))[:80]}"
|
||||
return "web search"
|
||||
if tool_type == "mcp_call":
|
||||
n = data.get("name", "")
|
||||
return f"mcp: {n}" if n else "mcp call"
|
||||
return tool_type.replace("_", " ")
|
||||
|
||||
|
||||
def _output_text_from_message_content(content: list[dict[str, object]]) -> str:
|
||||
"""
|
||||
Extract completed assistant text from a message item.
|
||||
|
||||
:param content: Message content blocks from ``MessageDone``.
|
||||
:returns: Concatenated ``output_text`` content.
|
||||
"""
|
||||
parts: list[str] = []
|
||||
for block in content:
|
||||
if block.get("type") != "output_text":
|
||||
continue
|
||||
text = block.get("text")
|
||||
if isinstance(text, str):
|
||||
parts.append(text)
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
class BlockStream:
|
||||
"""Consumes a session event stream and emits semantic stream blocks.
|
||||
|
||||
:param text_flush_threshold: Min chars to buffer before flushing
|
||||
on a word boundary. Default 30.
|
||||
"""
|
||||
|
||||
def __init__(self, text_flush_threshold: int = 30) -> None:
|
||||
self._flush_threshold = text_flush_threshold
|
||||
|
||||
async def stream(
|
||||
self,
|
||||
session: Session,
|
||||
input: str | list[dict[str, object]],
|
||||
*,
|
||||
files: list[str] | None = None,
|
||||
) -> AsyncIterator[AnyBlock]:
|
||||
"""Stream semantic blocks for one turn.
|
||||
|
||||
:param session: The :class:`Session` whose event stream to fold.
|
||||
The session's ``send()`` method is invoked internally.
|
||||
:param input: User text or a list of content-block dicts,
|
||||
e.g. ``"hello"`` or ``[{"type": "input_text", "text": "hi"}]``.
|
||||
:param files: Optional list of file paths to upload and attach
|
||||
to the turn, e.g. ``["./data.csv"]``.
|
||||
:returns: Async iterator of :class:`AnyBlock` values in order.
|
||||
"""
|
||||
in_reasoning = False
|
||||
reasoning_text = ""
|
||||
summary_text = ""
|
||||
# Per-section accumulator for reasoning chunks awaiting a
|
||||
# newline / threshold flush. Mirrors ``accumulated`` for text.
|
||||
reasoning_accumulated = ""
|
||||
# Set when this reasoning section streamed any
|
||||
# ``ReasoningChunk``. The final ``ReasoningBlock`` is then
|
||||
# suppressed so renderers don't show the same text twice
|
||||
# (once live, once as a summary panel).
|
||||
reasoning_chunks_emitted = False
|
||||
in_text = False
|
||||
accumulated = ""
|
||||
full_text = ""
|
||||
pending_tools: dict[str, ToolExecution] = {}
|
||||
# Long-lived metadata for tool calls rendered in this stream.
|
||||
# Unlike ``pending_tools``, this survives text deltas so a
|
||||
# delayed ``ToolResult`` can still render with the original
|
||||
# tool name/arguments instead of being dropped.
|
||||
tool_executions_by_call_id: dict[str, ToolExecution] = {}
|
||||
# Dedup set: every ``call_id`` we've yielded a
|
||||
# ``ToolGroup`` for so far. Survives ``pending_tools``
|
||||
# clears (which fire on each ``TextDelta`` to flush
|
||||
# ``ToolResultBlock``s for completed tools), so the
|
||||
# post-stream action_required event for an MCP tool
|
||||
# whose inline observed event already rendered doesn't
|
||||
# silently re-render once the intervening text deltas
|
||||
# cleared ``pending_tools``. Call_ids are SDK-assigned
|
||||
# ``toolu_*`` ids — globally unique within this turn,
|
||||
# so the set never grows pathologically.
|
||||
seen_call_ids: set[str] = set()
|
||||
# Dedup set: every ``call_id`` we've yielded a
|
||||
# ``ToolResultBlock`` for. Same survives-clear semantics
|
||||
# as ``seen_call_ids`` for the call-line side, but on the
|
||||
# result side. Necessary because every action_required
|
||||
# tool's result fires TWICE on the SSE stream: once
|
||||
# inline from ``_dispatch_action_required`` (the moment
|
||||
# the dispatch returns) and once from the
|
||||
# ``response.completed`` flush in
|
||||
# ``_translate_omnigent_event``. Without this, the result
|
||||
# panel renders twice.
|
||||
seen_result_call_ids: set[str] = set()
|
||||
agent: str | None = None
|
||||
turn = 0
|
||||
started = False
|
||||
|
||||
def _ctx() -> BlockContext:
|
||||
depth = agent.count(".") if agent else 0
|
||||
return BlockContext(agent=agent, depth=depth, turn=turn)
|
||||
|
||||
async for event in session.send(input, files=files):
|
||||
# ── Response lifecycle ───────────────────
|
||||
if isinstance(event, ResponseCreated):
|
||||
# Tool calls were already yielded immediately. Emit
|
||||
# result-only groups for tools that got output between
|
||||
# ResponseCompleted and this ResponseCreated.
|
||||
for ex in list(pending_tools.values()):
|
||||
if ex.output is not None:
|
||||
yield ToolResultBlock(
|
||||
name=ex.name,
|
||||
call_id=ex.call_id,
|
||||
agent_name=ex.agent_name,
|
||||
output=ex.output,
|
||||
arguments=ex.arguments,
|
||||
args_summary=ex.args_summary,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
pending_tools.clear()
|
||||
tool_executions_by_call_id.clear()
|
||||
agent = event.response.model
|
||||
if not started:
|
||||
started = True
|
||||
yield ResponseStartBlock(
|
||||
model=agent,
|
||||
response_id=event.response.id,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
else:
|
||||
turn += 1
|
||||
|
||||
elif isinstance(event, ResponseQueued | ResponseInProgress):
|
||||
pass
|
||||
|
||||
# ── Reasoning ────────────────────────────
|
||||
elif isinstance(event, ReasoningStarted):
|
||||
# Already reasoning = a new summarized-thinking block in the
|
||||
# same section. Flush this block's tail + a separator so
|
||||
# consecutive thought items don't run together.
|
||||
if in_reasoning:
|
||||
yield ReasoningChunk(text=reasoning_accumulated + "\n\n", ctx=_ctx())
|
||||
reasoning_accumulated = ""
|
||||
reasoning_chunks_emitted = True
|
||||
continue
|
||||
# Entering reasoning closes open text (symmetric with
|
||||
# TextDelta closing reasoning) — else interleaved
|
||||
# think→speak→think orphans + concatenates text.
|
||||
if in_text:
|
||||
if accumulated:
|
||||
yield TextChunk(text=accumulated, ctx=_ctx())
|
||||
accumulated = ""
|
||||
yield TextDone(
|
||||
full_text=full_text,
|
||||
has_code_blocks="```" in full_text,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
in_text = False
|
||||
full_text = ""
|
||||
in_reasoning = True
|
||||
reasoning_text = ""
|
||||
summary_text = ""
|
||||
reasoning_accumulated = ""
|
||||
reasoning_chunks_emitted = False
|
||||
yield ReasoningStartBlock(ctx=_ctx())
|
||||
|
||||
elif isinstance(event, ReasoningDelta | ReasoningSummaryDelta):
|
||||
# An out-of-order delta (no preceding ReasoningStarted)
|
||||
# marks an implicit start — Codex's bridged events
|
||||
# arrive this way. Emit the start block once so the
|
||||
# formatter has its "thinking…" anchor.
|
||||
if not in_reasoning:
|
||||
# Close any open text first — same boundary as ReasoningStarted.
|
||||
if in_text:
|
||||
if accumulated:
|
||||
yield TextChunk(text=accumulated, ctx=_ctx())
|
||||
accumulated = ""
|
||||
yield TextDone(
|
||||
full_text=full_text,
|
||||
has_code_blocks="```" in full_text,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
in_text = False
|
||||
full_text = ""
|
||||
in_reasoning = True
|
||||
reasoning_text = ""
|
||||
summary_text = ""
|
||||
reasoning_accumulated = ""
|
||||
reasoning_chunks_emitted = False
|
||||
yield ReasoningStartBlock(ctx=_ctx())
|
||||
if isinstance(event, ReasoningDelta):
|
||||
reasoning_text += event.delta
|
||||
else:
|
||||
summary_text += event.delta
|
||||
# Stream the delta as a ReasoningChunk so the TUI can
|
||||
# render mid-flight. Mirrors the TextDelta line/
|
||||
# threshold flush so chunks land on natural breaks.
|
||||
reasoning_accumulated += event.delta
|
||||
while "\n" in reasoning_accumulated:
|
||||
line, reasoning_accumulated = reasoning_accumulated.split("\n", 1)
|
||||
yield ReasoningChunk(text=line + "\n", ctx=_ctx())
|
||||
reasoning_chunks_emitted = True
|
||||
if len(reasoning_accumulated) >= self._flush_threshold:
|
||||
last_space = reasoning_accumulated.rfind(" ")
|
||||
if last_space > 0:
|
||||
yield ReasoningChunk(
|
||||
text=reasoning_accumulated[: last_space + 1],
|
||||
ctx=_ctx(),
|
||||
)
|
||||
reasoning_accumulated = reasoning_accumulated[last_space + 1 :]
|
||||
reasoning_chunks_emitted = True
|
||||
|
||||
# ── Text ─────────────────────────────────
|
||||
elif isinstance(event, TextDelta):
|
||||
if in_reasoning:
|
||||
in_reasoning = False
|
||||
if reasoning_accumulated:
|
||||
yield ReasoningChunk(
|
||||
text=reasoning_accumulated,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
reasoning_chunks_emitted = True
|
||||
reasoning_accumulated = ""
|
||||
if not reasoning_chunks_emitted:
|
||||
yield ReasoningBlock(
|
||||
reasoning_text=reasoning_text,
|
||||
summary_text=summary_text,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
# Emit results for tools that completed.
|
||||
for ex in list(pending_tools.values()):
|
||||
if ex.output is not None:
|
||||
yield ToolResultBlock(
|
||||
name=ex.name,
|
||||
call_id=ex.call_id,
|
||||
agent_name=ex.agent_name,
|
||||
output=ex.output,
|
||||
arguments=ex.arguments,
|
||||
args_summary=ex.args_summary,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
pending_tools.clear()
|
||||
|
||||
in_text = True
|
||||
accumulated += event.delta
|
||||
full_text += event.delta
|
||||
|
||||
while "\n" in accumulated:
|
||||
line, accumulated = accumulated.split("\n", 1)
|
||||
yield TextChunk(text=line + "\n", ctx=_ctx())
|
||||
|
||||
if len(accumulated) >= self._flush_threshold:
|
||||
last_space = accumulated.rfind(" ")
|
||||
if last_space > 0:
|
||||
yield TextChunk(text=accumulated[: last_space + 1], ctx=_ctx())
|
||||
accumulated = accumulated[last_space + 1 :]
|
||||
|
||||
# ── Tool calls ───────────────────────────
|
||||
elif isinstance(event, ToolCall):
|
||||
# Dedupe by call_id. Under the claude-sdk harness's
|
||||
# MCP path, a tool call surfaces as TWO ToolCall
|
||||
# events with correlated call_ids: an inline
|
||||
# observed event (status="completed") emitted as
|
||||
# the inner SDK parses the tool_use block, and a
|
||||
# post-stream action_required event emitted when
|
||||
# the SDK invokes the MCP-server handler. The
|
||||
# adapter (omnigent/runtime/harnesses/
|
||||
# _executor_adapter.py) threads the SDK's
|
||||
# tool_use_id through both so they share a
|
||||
# call_id; this block keeps the first occurrence
|
||||
# (the inline render) and drops the second so the
|
||||
# REPL doesn't render ``⏵ tool_name`` twice. See
|
||||
# designs/RUN_OMNIGENT_REPL_PARITY.md.
|
||||
#
|
||||
# Non-MCP paths emit exactly one ToolCall per
|
||||
# call_id, so the second-arrival branch never
|
||||
# fires for them.
|
||||
if event.call_id in seen_call_ids:
|
||||
# Already rendered the ⏵ line for this call_id
|
||||
# (e.g. via the inline observed event from the
|
||||
# harness's content_block_stop, or via the
|
||||
# ``ToolCallInProgress`` event the Omnigent server
|
||||
# emits at action_required arrival). Re-register
|
||||
# in pending_tools so the eventual ``ToolResult``
|
||||
# can pair by call_id — the prior pending entry
|
||||
# may have been cleared by the
|
||||
# ``pending_tools.clear()`` that fires on every
|
||||
# ``TextDelta`` between observed and
|
||||
# action_required (or between action_required
|
||||
# and the post-PATCH function_call_output).
|
||||
execution = tool_executions_by_call_id.get(event.call_id)
|
||||
if execution is None:
|
||||
execution = ToolExecution(
|
||||
name=event.name,
|
||||
arguments=event.arguments,
|
||||
args_summary=format_tool_args_brief(event.name, event.arguments),
|
||||
call_id=event.call_id,
|
||||
agent_name=event.agent_name,
|
||||
executed_by="server",
|
||||
)
|
||||
tool_executions_by_call_id[event.call_id] = execution
|
||||
pending_tools[event.call_id] = execution
|
||||
continue
|
||||
seen_call_ids.add(event.call_id)
|
||||
|
||||
if in_reasoning:
|
||||
in_reasoning = False
|
||||
if reasoning_accumulated:
|
||||
yield ReasoningChunk(
|
||||
text=reasoning_accumulated,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
reasoning_chunks_emitted = True
|
||||
reasoning_accumulated = ""
|
||||
if not reasoning_chunks_emitted:
|
||||
yield ReasoningBlock(
|
||||
reasoning_text=reasoning_text,
|
||||
summary_text=summary_text,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
if in_text:
|
||||
if accumulated:
|
||||
yield TextChunk(text=accumulated, ctx=_ctx())
|
||||
accumulated = ""
|
||||
yield TextDone(
|
||||
full_text=full_text,
|
||||
has_code_blocks="```" in full_text,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
in_text = False
|
||||
full_text = ""
|
||||
|
||||
execution = ToolExecution(
|
||||
name=event.name,
|
||||
arguments=event.arguments,
|
||||
args_summary=format_tool_args_brief(event.name, event.arguments),
|
||||
call_id=event.call_id,
|
||||
agent_name=event.agent_name,
|
||||
executed_by="server",
|
||||
)
|
||||
pending_tools[event.call_id] = execution
|
||||
tool_executions_by_call_id[event.call_id] = execution
|
||||
# Yield immediately so the user sees the tool call
|
||||
# before execution. output=None means the formatter
|
||||
# shows the call line but no result panel.
|
||||
yield ToolGroup(executions=[execution], ctx=_ctx())
|
||||
|
||||
elif isinstance(event, ToolResult):
|
||||
if event.call_id in seen_result_call_ids:
|
||||
# Already rendered the result panel. The late
|
||||
# ``response.completed`` flush emission is
|
||||
# redundant; drop it so a subsequent
|
||||
# ``ResponseCreated`` doesn't yield it again
|
||||
# via the pending_tools sweep.
|
||||
continue
|
||||
ex = pending_tools.get(event.call_id) or tool_executions_by_call_id.get(
|
||||
event.call_id
|
||||
)
|
||||
if ex is None:
|
||||
# Result arrived after a TextDelta cleared
|
||||
# ``pending_tools``. If the matching call line
|
||||
# never rendered, we have no name/agent for the
|
||||
# panel — drop the event rather than render
|
||||
# with placeholders.
|
||||
continue
|
||||
ex.output = event.output
|
||||
ex.executed_by = "client"
|
||||
if event.arguments:
|
||||
ex.arguments = event.arguments
|
||||
# Yield the result panel IMMEDIATELY so multi-
|
||||
# tool turns don't bunch result rendering at
|
||||
# end-of-turn. Each dispatch's
|
||||
# ``_live_publish`` of the function_call_output
|
||||
# arrives here as the tool finishes.
|
||||
seen_result_call_ids.add(event.call_id)
|
||||
yield ToolResultBlock(
|
||||
name=ex.name,
|
||||
call_id=ex.call_id,
|
||||
agent_name=ex.agent_name,
|
||||
output=event.output,
|
||||
arguments=ex.arguments,
|
||||
args_summary=ex.args_summary,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
# Drop from ``pending_tools`` so the next
|
||||
# ``TextDelta`` / ``ResponseCreated`` sweep
|
||||
# doesn't re-yield this result.
|
||||
pending_tools.pop(event.call_id, None)
|
||||
|
||||
# ── Native tools ─────────────────────────
|
||||
elif isinstance(event, NativeToolCall):
|
||||
yield NativeToolBlock(
|
||||
tool_type=event.tool_type,
|
||||
label=_format_native_label(event.tool_type, event.data),
|
||||
data=event.data,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
|
||||
# ── Message done ─────────────────────────
|
||||
elif isinstance(event, MessageDone):
|
||||
if in_reasoning:
|
||||
in_reasoning = False
|
||||
if reasoning_accumulated:
|
||||
yield ReasoningChunk(
|
||||
text=reasoning_accumulated,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
reasoning_chunks_emitted = True
|
||||
reasoning_accumulated = ""
|
||||
if not reasoning_chunks_emitted:
|
||||
yield ReasoningBlock(
|
||||
reasoning_text=reasoning_text,
|
||||
summary_text=summary_text,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
if in_text:
|
||||
if accumulated:
|
||||
yield TextChunk(text=accumulated, ctx=_ctx())
|
||||
accumulated = ""
|
||||
yield TextDone(
|
||||
full_text=full_text,
|
||||
has_code_blocks="```" in full_text,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
in_text = False
|
||||
full_text = ""
|
||||
else:
|
||||
text = _output_text_from_message_content(event.content)
|
||||
if text:
|
||||
# No prior deltas: synthesize a TextChunk so
|
||||
# TextChunk-only consumers (Session._stream_chunks)
|
||||
# still see the message.
|
||||
yield TextChunk(text=text, ctx=_ctx())
|
||||
yield TextDone(
|
||||
full_text=text,
|
||||
has_code_blocks="```" in text,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
|
||||
# ── Status events ────────────────────────
|
||||
elif isinstance(event, CompactionInProgress):
|
||||
yield CompactionBlock(ctx=_ctx())
|
||||
|
||||
elif isinstance(event, RetryEvent):
|
||||
yield RetryBlock(
|
||||
source=event.source,
|
||||
attempt=event.attempt,
|
||||
max_attempts=event.max_attempts,
|
||||
delay_seconds=event.delay_seconds,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
|
||||
elif isinstance(event, ErrorEvent):
|
||||
# Pass ``code`` through too — renderers need it as a
|
||||
# fallback label when ``message`` is empty (otherwise
|
||||
# the error panel shows just ``[llm]`` with no hint
|
||||
# as to what went wrong).
|
||||
yield ErrorBlock(
|
||||
message=event.error.message,
|
||||
source=event.source,
|
||||
code=event.error.code,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
|
||||
elif isinstance(event, OutputFileDone):
|
||||
yield FileBlock(file_id=event.file_id, filename=event.filename, ctx=_ctx())
|
||||
|
||||
# ── Terminal events ──────────────────────
|
||||
elif isinstance(
|
||||
event,
|
||||
ResponseCompleted | ResponseFailed | ResponseIncomplete | ResponseCancelled,
|
||||
):
|
||||
if in_reasoning:
|
||||
in_reasoning = False
|
||||
if reasoning_accumulated:
|
||||
yield ReasoningChunk(
|
||||
text=reasoning_accumulated,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
reasoning_chunks_emitted = True
|
||||
reasoning_accumulated = ""
|
||||
if not reasoning_chunks_emitted:
|
||||
yield ReasoningBlock(
|
||||
reasoning_text=reasoning_text,
|
||||
summary_text=summary_text,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
if in_text:
|
||||
if accumulated:
|
||||
yield TextChunk(text=accumulated, ctx=_ctx())
|
||||
accumulated = ""
|
||||
yield TextDone(
|
||||
full_text=full_text,
|
||||
has_code_blocks="```" in full_text,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
in_text = False
|
||||
full_text = ""
|
||||
|
||||
yield ResponseEndBlock(
|
||||
status=event.response.status,
|
||||
response=event.response,
|
||||
ctx=_ctx(),
|
||||
)
|
||||
@@ -0,0 +1,323 @@
|
||||
"""Tool handler, hooks, and context types for client-side tool execution."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from ._types import ErrorInfo, Response
|
||||
|
||||
# ── Tool handler ─────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCallInfo:
|
||||
"""Context passed to ``ToolHandler.execute``."""
|
||||
|
||||
name: str
|
||||
arguments: dict[str, object]
|
||||
call_id: str
|
||||
agent_name: str # "coder" or "coder.researcher"
|
||||
response_id: str
|
||||
iteration: int # Tool loop iteration (0-based)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolHandler:
|
||||
"""Client-side tool execution configuration.
|
||||
|
||||
Passed to ``stream()`` to enable automatic tool execution.
|
||||
The client runs the full loop: stream -> detect tool calls ->
|
||||
call ``execute()`` -> send results -> continue streaming.
|
||||
"""
|
||||
|
||||
schemas: list[dict[str, object]]
|
||||
execute: Callable[[ToolCallInfo], Awaitable[str] | str]
|
||||
|
||||
|
||||
# ── Hook context types ───────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCallStartCtx:
|
||||
"""Context for ``on_tool_call_start``."""
|
||||
|
||||
name: str
|
||||
arguments: dict[str, object]
|
||||
call_id: str
|
||||
agent_name: str
|
||||
executed_by: str # "client" or "server"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCallEndCtx:
|
||||
"""Context for ``on_tool_call_end``."""
|
||||
|
||||
name: str
|
||||
call_id: str
|
||||
agent_name: str
|
||||
output: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class NativeToolCallCtx:
|
||||
"""Context for ``on_native_tool_call``."""
|
||||
|
||||
tool_type: str
|
||||
data: dict[str, object]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolResultInfo:
|
||||
"""A single tool result within ``ToolResultsReadyCtx``."""
|
||||
|
||||
call_id: str
|
||||
name: str
|
||||
output: str
|
||||
agent_name: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolResultsReadyCtx:
|
||||
"""Context for ``on_tool_results_ready``."""
|
||||
|
||||
results: list[ToolResultInfo]
|
||||
iteration: int
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReasoningStartCtx:
|
||||
"""Context for ``on_reasoning_start``."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReasoningEndCtx:
|
||||
"""Context for ``on_reasoning_end``."""
|
||||
|
||||
reasoning_text: str
|
||||
summary_text: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompactionStartCtx:
|
||||
"""Context for ``on_compaction_start``."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class CompactionEndCtx:
|
||||
"""Context for ``on_compaction_end``."""
|
||||
|
||||
item: dict[str, object]
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageStartCtx:
|
||||
"""Context for ``on_message_start``."""
|
||||
|
||||
response_id: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageEndCtx:
|
||||
"""Context for ``on_message_end``."""
|
||||
|
||||
content: list[dict[str, object]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class FileOutputCtx:
|
||||
"""Context for ``on_file_output``."""
|
||||
|
||||
file_id: str
|
||||
filename: str | None
|
||||
content_type: str | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetryCtx:
|
||||
"""Context for ``on_retry``."""
|
||||
|
||||
source: str
|
||||
tool_name: str | None
|
||||
attempt: int
|
||||
max_attempts: int
|
||||
delay_seconds: float
|
||||
error: ErrorInfo
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServerErrorCtx:
|
||||
"""Context for ``on_server_error``."""
|
||||
|
||||
source: str
|
||||
tool_name: str | None
|
||||
error: ErrorInfo
|
||||
|
||||
|
||||
@dataclass
|
||||
class TransportErrorCtx:
|
||||
"""Context for ``on_transport_error``."""
|
||||
|
||||
error: Exception
|
||||
|
||||
|
||||
@dataclass
|
||||
class SubAgentInfo:
|
||||
"""Info about a single spawned sub-agent."""
|
||||
|
||||
response_id: str
|
||||
agent_name: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class SubAgentSpawnedCtx:
|
||||
"""Context for ``on_sub_agent_spawned``."""
|
||||
|
||||
parent_response_id: str
|
||||
sub_agents: list[SubAgentInfo] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SubAgentCompletedCtx:
|
||||
"""Context for ``on_sub_agent_completed``."""
|
||||
|
||||
response_id: str
|
||||
agent_name: str
|
||||
status: str
|
||||
output_summary: str | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ElicitationRequestCtx:
|
||||
"""
|
||||
Context for ``on_elicitation_request``.
|
||||
|
||||
Surfaced when the server emits a ``response.elicitation_request``
|
||||
SSE event (matches MCP's ``elicitation/create`` semantics —
|
||||
see ``designs/SERVER_HARNESS_CONTRACT.md`` §"Universal API
|
||||
additions" + POLICIES.md §7). The hook returns ``True`` to
|
||||
accept or ``False`` to decline; the client submits the
|
||||
verdict to the server automatically as an MCP-shape
|
||||
``ElicitResult`` POST.
|
||||
|
||||
For richer elicitations (with form fields), the hook may
|
||||
return an ``ElicitationResult``-shaped value instead of a
|
||||
bool — bool is the convenience path for binary
|
||||
approve/decline.
|
||||
|
||||
:param elicitation_id: Server-assigned id. Used in the URL
|
||||
path of the reply endpoint, e.g. ``"elicit_abc123"``.
|
||||
:param message: Human-readable prompt to render. For the
|
||||
policy ASK producer this is the combined reason string
|
||||
from deciding ASKing policies (``"; "``-joined).
|
||||
:param requested_schema: A restricted subset of JSON Schema
|
||||
defining the structure of an expected response. Empty
|
||||
``{}`` for binary approve/reject elicitations.
|
||||
:param mode: MCP elicitation mode — ``"form"`` (inline) or
|
||||
``"url"`` (standalone approval page).
|
||||
:param phase: Producer-supplied extra (policy ASK only) —
|
||||
which enforcement point produced the ASK. Empty string
|
||||
when not applicable.
|
||||
:param policy_name: Producer-supplied extra (policy ASK
|
||||
only) — name of the deciding ASKing policy, e.g.
|
||||
``"approve_web_search"``. Empty string when not
|
||||
applicable.
|
||||
:param content_preview: Producer-supplied extra (policy ASK
|
||||
only) — truncated (1024 chars) snapshot of the gated
|
||||
content. Safe to display verbatim. Empty string when
|
||||
not applicable.
|
||||
:param response_id: The in-progress response id — for
|
||||
logging / audit purposes only. The client handles the
|
||||
elicitation reply POST automatically.
|
||||
"""
|
||||
|
||||
elicitation_id: str
|
||||
message: str
|
||||
requested_schema: dict[str, object]
|
||||
mode: str
|
||||
phase: str
|
||||
policy_name: str
|
||||
content_preview: str
|
||||
response_id: str
|
||||
url: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResponseStartCtx:
|
||||
"""Context for ``on_response_start``."""
|
||||
|
||||
response: Response
|
||||
|
||||
|
||||
@dataclass
|
||||
class ResponseEndCtx:
|
||||
"""Context for ``on_response_end``."""
|
||||
|
||||
response: Response
|
||||
status: str # "completed", "failed", "incomplete", "cancelled"
|
||||
|
||||
|
||||
# ── Stream hooks ─────────────────────────────────────────
|
||||
|
||||
# Hook type alias: sync or async callable, or None.
|
||||
_Hook = Callable[..., Awaitable[None] | None] | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamHooks:
|
||||
"""Lifecycle hooks for stream events.
|
||||
|
||||
Every class of event has a start/end hook pair where
|
||||
applicable. Hooks can be sync or async. All are optional.
|
||||
"""
|
||||
|
||||
# Tool calls (server-side and client-side)
|
||||
on_tool_call_start: Callable[[ToolCallStartCtx], Awaitable[None] | None] | None = None
|
||||
on_tool_call_end: Callable[[ToolCallEndCtx], Awaitable[None] | None] | None = None
|
||||
|
||||
# Native tool calls (provider-executed)
|
||||
on_native_tool_call: Callable[[NativeToolCallCtx], Awaitable[None] | None] | None = None
|
||||
|
||||
# Tool loop iteration
|
||||
on_tool_results_ready: Callable[[ToolResultsReadyCtx], Awaitable[None] | None] | None = None
|
||||
|
||||
# Reasoning
|
||||
on_reasoning_start: Callable[[ReasoningStartCtx], Awaitable[None] | None] | None = None
|
||||
on_reasoning_end: Callable[[ReasoningEndCtx], Awaitable[None] | None] | None = None
|
||||
|
||||
# Compaction
|
||||
on_compaction_start: Callable[[CompactionStartCtx], Awaitable[None] | None] | None = None
|
||||
on_compaction_end: Callable[[CompactionEndCtx], Awaitable[None] | None] | None = None
|
||||
|
||||
# Message (assistant response)
|
||||
on_message_start: Callable[[MessageStartCtx], Awaitable[None] | None] | None = None
|
||||
on_message_end: Callable[[MessageEndCtx], Awaitable[None] | None] | None = None
|
||||
|
||||
# File output
|
||||
on_file_output: Callable[[FileOutputCtx], Awaitable[None] | None] | None = None
|
||||
|
||||
# Retry and error
|
||||
on_retry: Callable[[RetryCtx], Awaitable[None] | None] | None = None
|
||||
on_server_error: Callable[[ServerErrorCtx], Awaitable[None] | None] | None = None
|
||||
on_transport_error: Callable[[TransportErrorCtx], Awaitable[bool] | bool] | None = None
|
||||
|
||||
# Sub-agent lifecycle
|
||||
on_sub_agent_spawned: Callable[[SubAgentSpawnedCtx], Awaitable[None] | None] | None = None
|
||||
on_sub_agent_completed: Callable[[SubAgentCompletedCtx], Awaitable[None] | None] | None = None
|
||||
|
||||
# Response lifecycle
|
||||
on_response_start: Callable[[ResponseStartCtx], Awaitable[None] | None] | None = None
|
||||
on_response_end: Callable[[ResponseEndCtx], Awaitable[None] | None] | None = None
|
||||
|
||||
# MCP-shape elicitations. Called when the server emits a
|
||||
# ``response.elicitation_request`` SSE event (POLICIES.md §7
|
||||
# surfaces policy ASKs as one producer; future producers
|
||||
# include credential prompts, parameter clarifications, and
|
||||
# any other MCP elicitation use case). The callback returns
|
||||
# ``True`` to accept or ``False`` to decline; the client POSTs
|
||||
# a session ``approval`` event carrying the MCP-shape
|
||||
# ``ElicitResult``. When no callback is registered, the client
|
||||
# fail-closed declines — an unhandled
|
||||
# elicitation becomes DENY, preserving fail-loud-rather-than-
|
||||
# block semantics.
|
||||
on_elicitation_request: Callable[[ElicitationRequestCtx], Awaitable[bool] | bool] | None = None
|
||||
@@ -0,0 +1,138 @@
|
||||
"""Composable stream transforms for stream blocks.
|
||||
|
||||
Each transform is an async generator function that wraps a block
|
||||
stream. Compose with ``pipe()``:
|
||||
|
||||
stream = pipe(
|
||||
block_stream.stream(session, text),
|
||||
skip_blocks(ReasoningBlock),
|
||||
skip_intermediate_ends(),
|
||||
)
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator
|
||||
|
||||
from ._blocks import (
|
||||
AnyBlock,
|
||||
ResponseEndBlock,
|
||||
TextDone,
|
||||
)
|
||||
|
||||
|
||||
def pipe(stream: AsyncIterator[AnyBlock], *transforms: object) -> AsyncIterator[AnyBlock]:
|
||||
"""Compose transforms left-to-right.
|
||||
|
||||
``pipe(stream, t1, t2)`` is equivalent to ``t2(t1(stream))``.
|
||||
"""
|
||||
for t in transforms:
|
||||
stream = t(stream)
|
||||
return stream
|
||||
|
||||
|
||||
def skip_blocks(*types: type) -> object:
|
||||
"""Drop blocks of the given types.
|
||||
|
||||
Usage::
|
||||
|
||||
stream = skip_blocks(ReasoningBlock)(renderer.stream(...))
|
||||
"""
|
||||
|
||||
def _transform(stream: AsyncIterator[AnyBlock]) -> AsyncIterator[AnyBlock]:
|
||||
async def _inner() -> AsyncIterator[AnyBlock]:
|
||||
async for block in stream:
|
||||
if not isinstance(block, tuple(types)):
|
||||
yield block
|
||||
|
||||
return _inner()
|
||||
|
||||
return _transform
|
||||
|
||||
|
||||
def skip_intermediate_ends() -> object:
|
||||
"""Suppress ``ResponseEndBlock`` events from tool loop iterations.
|
||||
|
||||
Only yields the final ``ResponseEndBlock`` — the one not followed
|
||||
by another block from the same turn.
|
||||
"""
|
||||
|
||||
def _transform(stream: AsyncIterator[AnyBlock]) -> AsyncIterator[AnyBlock]:
|
||||
async def _inner() -> AsyncIterator[AnyBlock]:
|
||||
pending_end: ResponseEndBlock | None = None
|
||||
async for block in stream:
|
||||
if isinstance(block, ResponseEndBlock):
|
||||
# Buffer it — might not be the last one.
|
||||
pending_end = block
|
||||
else:
|
||||
# A non-end block arrived, so the buffered end
|
||||
# was intermediate. Drop it.
|
||||
pending_end = None
|
||||
yield block
|
||||
# Stream ended. The buffered end IS the final one.
|
||||
if pending_end is not None:
|
||||
yield pending_end
|
||||
|
||||
return _inner()
|
||||
|
||||
return _transform
|
||||
|
||||
|
||||
def merge_text_across_iterations() -> object:
|
||||
"""Merge ``TextDone`` blocks across tool loop iterations.
|
||||
|
||||
When the tool loop runs N iterations, each may produce a
|
||||
``TextDone``. This merges them into a single ``TextDone`` at
|
||||
the end. ``TextChunk`` blocks pass through unchanged.
|
||||
"""
|
||||
|
||||
def _transform(stream: AsyncIterator[AnyBlock]) -> AsyncIterator[AnyBlock]:
|
||||
async def _inner() -> AsyncIterator[AnyBlock]:
|
||||
accumulated = ""
|
||||
async for block in stream:
|
||||
if isinstance(block, TextDone):
|
||||
accumulated += block.full_text
|
||||
elif isinstance(block, ResponseEndBlock):
|
||||
if accumulated:
|
||||
yield TextDone(
|
||||
full_text=accumulated,
|
||||
has_code_blocks="```" in accumulated,
|
||||
ctx=block.ctx,
|
||||
)
|
||||
accumulated = ""
|
||||
yield block
|
||||
else:
|
||||
yield block
|
||||
# Edge case: stream ends without ResponseEndBlock.
|
||||
if accumulated:
|
||||
yield TextDone(
|
||||
full_text=accumulated,
|
||||
has_code_blocks="```" in accumulated,
|
||||
)
|
||||
|
||||
return _inner()
|
||||
|
||||
return _transform
|
||||
|
||||
|
||||
def only_agent(agent_name: str | None) -> object:
|
||||
"""Filter to blocks from a specific agent.
|
||||
|
||||
Pass ``None`` to include all agents (no filtering).
|
||||
|
||||
Usage::
|
||||
|
||||
stream = only_agent("coder.researcher")(renderer.stream(...))
|
||||
|
||||
:param agent_name: Agent name to filter by, or ``None`` for all.
|
||||
"""
|
||||
|
||||
def _transform(stream: AsyncIterator[AnyBlock]) -> AsyncIterator[AnyBlock]:
|
||||
async def _inner() -> AsyncIterator[AnyBlock]:
|
||||
async for block in stream:
|
||||
if agent_name is None or block.ctx.agent == agent_name:
|
||||
yield block
|
||||
|
||||
return _inner()
|
||||
|
||||
return _transform
|
||||
@@ -0,0 +1,235 @@
|
||||
"""Typed dataclasses for API response objects."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConversationRef:
|
||||
"""Reference to a conversation, as returned on response objects."""
|
||||
|
||||
id: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class Usage:
|
||||
"""
|
||||
Token usage statistics for a completed response.
|
||||
|
||||
:param input_tokens: Number of input (prompt) tokens consumed
|
||||
(sum across all LLM sub-calls for multi-call turns).
|
||||
:param output_tokens: Number of output (completion) tokens
|
||||
generated (sum across all LLM sub-calls).
|
||||
:param total_tokens: Billing total — ``input_tokens + output_tokens``
|
||||
summed across all sub-calls.
|
||||
:param context_tokens: Context-fill estimate for the next turn.
|
||||
Present only for executors that make multiple LLM sub-calls per
|
||||
turn (e.g. ``openai-agents``); ``None`` for single-call
|
||||
executors where ``total_tokens`` serves the same purpose.
|
||||
"""
|
||||
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
total_tokens: int
|
||||
context_tokens: int | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ErrorInfo:
|
||||
"""Structured error information from the server."""
|
||||
|
||||
code: str
|
||||
message: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class IncompleteDetails:
|
||||
"""Details about why a response stopped early."""
|
||||
|
||||
reason: str # "max_iterations", "execution_timeout", "context_overflow", etc.
|
||||
|
||||
|
||||
@dataclass
|
||||
class Response:
|
||||
"""A response object from the server.
|
||||
|
||||
Mirrors the JSON shape from ``POST /v1/responses`` and
|
||||
``GET /v1/responses/{id}``.
|
||||
"""
|
||||
|
||||
id: str
|
||||
status: str # "queued", "in_progress", "completed", "failed", "incomplete", "cancelled"
|
||||
model: str
|
||||
output: list[dict[str, object]] = field(default_factory=list)
|
||||
created_at: int = 0
|
||||
completed_at: int | None = None
|
||||
previous_response_id: str | None = None
|
||||
conversation: ConversationRef | None = None
|
||||
usage: Usage | None = None
|
||||
error: ErrorInfo | None = None
|
||||
incomplete_details: IncompleteDetails | None = None
|
||||
background: bool = False
|
||||
instructions: str | None = None
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, object]) -> Response:
|
||||
"""Parse a response object from a JSON dict."""
|
||||
conv_raw = data.get("conversation")
|
||||
conversation = (
|
||||
ConversationRef(id=str(conv_raw["id"])) if isinstance(conv_raw, dict) else None
|
||||
)
|
||||
usage_raw = data.get("usage")
|
||||
usage = (
|
||||
Usage(
|
||||
input_tokens=int(usage_raw.get("input_tokens", 0)),
|
||||
output_tokens=int(usage_raw.get("output_tokens", 0)),
|
||||
total_tokens=int(usage_raw.get("total_tokens", 0)),
|
||||
context_tokens=(
|
||||
int(usage_raw["context_tokens"])
|
||||
if usage_raw.get("context_tokens") is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
if isinstance(usage_raw, dict)
|
||||
else None
|
||||
)
|
||||
error_raw = data.get("error")
|
||||
error = (
|
||||
ErrorInfo(
|
||||
code=str(error_raw.get("code", "")),
|
||||
message=str(error_raw.get("message", "")),
|
||||
)
|
||||
if isinstance(error_raw, dict)
|
||||
else None
|
||||
)
|
||||
inc_raw = data.get("incomplete_details")
|
||||
incomplete = (
|
||||
IncompleteDetails(reason=str(inc_raw.get("reason", "")))
|
||||
if isinstance(inc_raw, dict)
|
||||
else None
|
||||
)
|
||||
return cls(
|
||||
id=str(data.get("id", "")),
|
||||
status=str(data.get("status", "")),
|
||||
model=str(data.get("model", "")),
|
||||
output=data.get("output", []) if isinstance(data.get("output"), list) else [],
|
||||
created_at=int(data.get("created_at", 0)),
|
||||
completed_at=int(data["completed_at"])
|
||||
if data.get("completed_at") is not None
|
||||
else None,
|
||||
previous_response_id=(
|
||||
str(data["previous_response_id"])
|
||||
if data.get("previous_response_id") is not None
|
||||
else None
|
||||
),
|
||||
conversation=conversation,
|
||||
usage=usage,
|
||||
error=error,
|
||||
incomplete_details=incomplete,
|
||||
background=bool(data.get("background", False)),
|
||||
instructions=(
|
||||
str(data["instructions"]) if data.get("instructions") is not None else None
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Agent:
|
||||
"""An agent registered on the server."""
|
||||
|
||||
id: str
|
||||
name: str
|
||||
description: str | None
|
||||
created_at: int
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, object]) -> Agent:
|
||||
"""Parse an agent from a JSON dict."""
|
||||
return cls(
|
||||
id=str(data.get("id", "")),
|
||||
name=str(data.get("name", "")),
|
||||
description=str(data["description"]) if data.get("description") is not None else None,
|
||||
created_at=int(data.get("created_at", 0)),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class File:
|
||||
"""A file uploaded to the server."""
|
||||
|
||||
id: str
|
||||
filename: str
|
||||
bytes: int
|
||||
created_at: int
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, object]) -> File:
|
||||
"""Parse a file from a JSON dict."""
|
||||
return cls(
|
||||
id=str(data.get("id", "")),
|
||||
filename=str(data.get("filename", "")),
|
||||
bytes=int(data.get("bytes", 0)),
|
||||
created_at=int(data.get("created_at", 0)),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Conversation:
|
||||
"""
|
||||
A conversation on the server.
|
||||
|
||||
:param id: Server-assigned conversation id, e.g. ``"conv_abc123"``.
|
||||
:param title: Optional user-assigned title.
|
||||
:param created_at: Unix epoch timestamp of creation.
|
||||
:param labels: Guardrails labels on this conversation (keys are
|
||||
label names, values are their current string values). Empty
|
||||
dict when the server's PolicyEngine hasn't written any
|
||||
labels yet. Surfaced so callers like the REPL's Ctrl+O debug
|
||||
overlay can render them at parity with the legacy Ctrl+G
|
||||
overview. Older servers that don't include ``labels`` in
|
||||
their ``ConversationObject`` response parse to ``{}``.
|
||||
"""
|
||||
|
||||
id: str
|
||||
title: str | None
|
||||
created_at: int
|
||||
labels: dict[str, str] = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, object]) -> Conversation:
|
||||
"""Parse a conversation from a JSON dict."""
|
||||
raw_labels = data.get("labels")
|
||||
labels: dict[str, str] = {}
|
||||
if isinstance(raw_labels, dict):
|
||||
# Narrow to ``dict[str, str]`` — the server's schema declares
|
||||
# that shape but the JSON boundary only gives us
|
||||
# ``dict[str, object]``. Drop non-string entries rather than
|
||||
# crash the overlay render on a malformed response.
|
||||
labels = {str(k): str(v) for k, v in raw_labels.items() if isinstance(v, str)}
|
||||
return cls(
|
||||
id=str(data.get("id", "")),
|
||||
title=str(data["title"]) if data.get("title") is not None else None,
|
||||
created_at=int(data.get("created_at", 0)),
|
||||
labels=labels,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PaginatedList:
|
||||
"""A paginated list response from the server."""
|
||||
|
||||
data: list[dict[str, object]]
|
||||
first_id: str | None = None
|
||||
last_id: str | None = None
|
||||
has_more: bool = False
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, object]) -> PaginatedList:
|
||||
"""Parse a paginated list from a JSON dict."""
|
||||
return cls(
|
||||
data=data.get("data", []) if isinstance(data.get("data"), list) else [],
|
||||
first_id=str(data["first_id"]) if data.get("first_id") is not None else None,
|
||||
last_id=str(data["last_id"]) if data.get("last_id") is not None else None,
|
||||
has_more=bool(data.get("has_more", False)),
|
||||
)
|
||||
@@ -0,0 +1,34 @@
|
||||
"""Tool-authoring primitives for omnigent.
|
||||
|
||||
Use the :func:`tool` decorator to mark a module-level Python function
|
||||
as a tool the agent can call. The decorator derives the LLM-facing
|
||||
JSON schema from the function's signature and Google-style docstring;
|
||||
the caller just writes Python::
|
||||
|
||||
from omnigent_client import tool
|
||||
|
||||
@tool
|
||||
def get_current_time() -> dict[str, str]:
|
||||
\"\"\"Return the current UTC time as ISO-8601.\"\"\"
|
||||
return {"now": datetime.now(timezone.utc).isoformat()}
|
||||
|
||||
Pass decorated functions as the ``tools=`` argument to
|
||||
:meth:`OmnigentClient.query` or :meth:`Session.query`.
|
||||
|
||||
Server-side runtime (``omnigent.tools.local``) also consumes this
|
||||
decorator to load ``@tool``-decorated functions bundled inside agent
|
||||
images, so the same decorator powers both authoring and runtime.
|
||||
"""
|
||||
|
||||
from ._decorator import TOOL_MARKER_ATTR, ToolMetadata, get_tool_metadata, tool
|
||||
from ._handler import build_tool_handler
|
||||
from ._state import ToolState
|
||||
|
||||
__all__ = [
|
||||
"TOOL_MARKER_ATTR",
|
||||
"ToolMetadata",
|
||||
"ToolState",
|
||||
"build_tool_handler",
|
||||
"get_tool_metadata",
|
||||
"tool",
|
||||
]
|
||||
@@ -0,0 +1,238 @@
|
||||
"""
|
||||
The ``@tool`` decorator and its metadata.
|
||||
|
||||
A ``@tool``-decorated module-level function is the authoring
|
||||
contract for custom Python tools in omnigent. The decorator:
|
||||
|
||||
1. Validates that the target is a module-level ``def`` or
|
||||
``async def`` (rejects class methods, lambdas, and nested
|
||||
functions; see :func:`_validate_decorator_target`).
|
||||
2. Derives the function-calling JSON schema from the signature
|
||||
and Google-style docstring (see :mod:`._schema`).
|
||||
3. Attaches metadata to the function via the
|
||||
:data:`TOOL_MARKER_ATTR` attribute so the framework can
|
||||
discover decorated functions by scanning a module's namespace.
|
||||
|
||||
The decorator is intentionally pure metadata: it returns the
|
||||
original function unwrapped. The framework's executor handles
|
||||
sync vs async execution (wrapping plain ``def`` bodies in
|
||||
``asyncio.to_thread`` so they don't block the event loop).
|
||||
|
||||
**Decorator stacking**: ``@tool`` should be the outermost
|
||||
decorator. Inner decorators (``@retry``, ``@cache``, etc.) must
|
||||
use ``functools.wraps`` so the signature and docstring survive
|
||||
to the schema-derivation pass; otherwise the schema will reflect
|
||||
the wrapper, not the underlying function.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, ParamSpec, TypeVar, overload
|
||||
|
||||
from ._schema import build_function_schema
|
||||
|
||||
# Marker attribute name. The framework's loader scans
|
||||
# ``module.__dict__`` for objects carrying this attribute to
|
||||
# enumerate the tools a Python file exports.
|
||||
TOOL_MARKER_ATTR = "_omnigent_tool_metadata"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ToolMetadata:
|
||||
"""
|
||||
Metadata attached to a ``@tool``-decorated function.
|
||||
|
||||
Read by the framework loader and by the subprocess runner.
|
||||
Read-only by convention — re-decoration or hand-editing is
|
||||
not supported.
|
||||
|
||||
:param name: The tool name as the LLM sees it. Derived from
|
||||
the function's ``__name__``, e.g. ``"word_count"``.
|
||||
:param description: Human-readable description, derived from
|
||||
the function's docstring's leading paragraph.
|
||||
:param json_schema: The function-calling JSON schema for the
|
||||
tool's parameters, in the OpenAI function-calling shape
|
||||
(an ``object`` schema with ``properties`` / ``required``).
|
||||
Already strict-mode-normalized if ``strict=True`` was
|
||||
passed to the decorator.
|
||||
:param strict: Whether the schema was normalized to strict
|
||||
mode. Stored so the executor side can stay consistent with
|
||||
the LLM's expectations.
|
||||
:param return_annotation: The function's declared return type,
|
||||
used by the executor to deserialize the return value via
|
||||
``pydantic.TypeAdapter``. ``None`` if the function has no
|
||||
return annotation (executor falls back to
|
||||
``json.dumps(value, default=str)``).
|
||||
"""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
json_schema: dict[str, Any]
|
||||
strict: bool
|
||||
return_annotation: type[Any] | None
|
||||
|
||||
|
||||
P = ParamSpec("P")
|
||||
R = TypeVar("R")
|
||||
|
||||
|
||||
@overload
|
||||
def tool(fn: Callable[P, R]) -> Callable[P, R]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def tool(
|
||||
*,
|
||||
strict: bool = ...,
|
||||
) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
|
||||
|
||||
|
||||
def tool(
|
||||
fn: Callable[P, R] | None = None,
|
||||
*,
|
||||
strict: bool = True,
|
||||
) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
|
||||
"""
|
||||
Mark a module-level function as an omnigent tool.
|
||||
|
||||
The decorator infers the LLM-facing schema from the function's
|
||||
type hints and Google-style docstring, then attaches the
|
||||
derived metadata via the :data:`TOOL_MARKER_ATTR` attribute.
|
||||
The framework's loader scans modules for this marker to
|
||||
register tools at agent-image load time.
|
||||
|
||||
Every ``@tool``-decorated function is sync from the framework's
|
||||
perspective. Authors who want a tool dispatched as background
|
||||
work do not annotate the tool — the LLM picks per call site
|
||||
via ``sys_call_async(tool=..., args=...)`` (see
|
||||
``omnigent/tools/builtins/async_inbox.py``). The author-time
|
||||
``@tool(synchronous=False)`` decoration was removed once
|
||||
``sys_call_async`` shipped — keeping both surfaces would double
|
||||
the dispatch paths without adding capability.
|
||||
|
||||
Usage::
|
||||
|
||||
@tool
|
||||
async def word_count(text: str) -> dict[str, int]:
|
||||
\"\"\"Count words, characters, and lines.\"\"\"
|
||||
...
|
||||
|
||||
Restrictions:
|
||||
- Target must be a module-level ``def`` or ``async def``.
|
||||
Class methods, lambdas, and nested functions are rejected
|
||||
with ``TypeError``.
|
||||
- ``@tool`` should be the outermost decorator; inner
|
||||
decorators must use ``functools.wraps`` for schema derivation
|
||||
to see the underlying signature and docstring.
|
||||
|
||||
:param fn: When used as bare ``@tool`` (no parens), the
|
||||
decorated function. ``None`` when used as
|
||||
``@tool(strict=False)``.
|
||||
:param strict: If ``True`` (default), the derived schema is
|
||||
normalized to strict mode (``additionalProperties: false``
|
||||
on objects, all properties required). Authors who hit a
|
||||
schema strict mode breaks can opt out with
|
||||
``@tool(strict=False)``.
|
||||
:returns: The original function with the tool metadata
|
||||
attached (when called with ``fn``), or a decorator
|
||||
function (when called with keyword args).
|
||||
"""
|
||||
|
||||
def wrap(target: Callable[P, R]) -> Callable[P, R]:
|
||||
_validate_decorator_target(target)
|
||||
schema_result = build_function_schema(target, strict=strict)
|
||||
metadata = ToolMetadata(
|
||||
name=target.__name__,
|
||||
description=schema_result.description,
|
||||
json_schema=schema_result.parameters_json_schema,
|
||||
strict=strict,
|
||||
return_annotation=schema_result.return_annotation,
|
||||
)
|
||||
# Attach metadata via setattr (the dynamic attribute name
|
||||
# is intentional — it's the framework's discovery contract,
|
||||
# see TOOL_MARKER_ATTR).
|
||||
setattr(target, TOOL_MARKER_ATTR, metadata)
|
||||
return target
|
||||
|
||||
if fn is None:
|
||||
return wrap
|
||||
return wrap(fn)
|
||||
|
||||
|
||||
def _validate_decorator_target(target: Any) -> None:
|
||||
"""
|
||||
Reject decorator application to anything other than a
|
||||
module-level ``def`` or ``async def``.
|
||||
|
||||
Class methods would include ``self`` / ``cls`` in the schema,
|
||||
which the LLM has no way to fill. Lambdas have no name and
|
||||
no docstring. Nested functions can close over enclosing scope
|
||||
that doesn't survive subprocess invocation.
|
||||
|
||||
:param target: The object the decorator was applied to.
|
||||
:raises TypeError: If ``target`` is not a module-level
|
||||
function. The message names the offending construct so
|
||||
agent authors get an actionable error.
|
||||
"""
|
||||
if not callable(target):
|
||||
raise TypeError(f"@tool can only be applied to functions, got {type(target).__name__}.")
|
||||
|
||||
if isinstance(target, (staticmethod, classmethod)):
|
||||
raise TypeError(
|
||||
"@tool cannot be applied to staticmethod or classmethod. "
|
||||
"Define the tool as a module-level function instead — "
|
||||
"the framework has no way to bind 'self' or 'cls' from "
|
||||
"an LLM-supplied argument set."
|
||||
)
|
||||
|
||||
if not inspect.isfunction(target):
|
||||
# Covers callables, methods of class instances, etc.
|
||||
raise TypeError(
|
||||
f"@tool requires a plain Python function, got "
|
||||
f"{type(target).__name__}. Define the tool as a "
|
||||
f"module-level def or async def."
|
||||
)
|
||||
|
||||
if target.__name__ == "<lambda>":
|
||||
raise TypeError(
|
||||
"@tool cannot be applied to a lambda. Lambdas have no "
|
||||
"name or docstring; the LLM has nothing to call. Define "
|
||||
"the tool with `def` or `async def` instead."
|
||||
)
|
||||
|
||||
# Nested-function detection: __qualname__ contains a dot path
|
||||
# (e.g. "outer.<locals>.inner") for any function defined inside
|
||||
# another function or class body.
|
||||
qualname = target.__qualname__
|
||||
if qualname != target.__name__:
|
||||
# Allow class-level methods if someone re-binds them at module
|
||||
# scope (rare); the staticmethod/classmethod check above is
|
||||
# the primary defense. Otherwise reject as nested.
|
||||
raise TypeError(
|
||||
f"@tool cannot be applied to nested functions or methods "
|
||||
f"({qualname!r}). Define the tool at module scope so the "
|
||||
f"framework's subprocess runner can re-import it cleanly. "
|
||||
f"State that needs to persist across invocations belongs "
|
||||
f"in module-level globals, not closure variables."
|
||||
)
|
||||
|
||||
|
||||
def get_tool_metadata(obj: Any) -> ToolMetadata | None:
|
||||
"""
|
||||
Return the :class:`ToolMetadata` for an object if it is a
|
||||
``@tool``-decorated function, else ``None``.
|
||||
|
||||
Used by the framework loader to filter a module's namespace
|
||||
down to the decorated functions it should register.
|
||||
|
||||
:param obj: Any value pulled from a module's namespace.
|
||||
:returns: The attached metadata, or ``None`` if the object
|
||||
was not produced by ``@tool``.
|
||||
"""
|
||||
metadata = getattr(obj, TOOL_MARKER_ATTR, None)
|
||||
if isinstance(metadata, ToolMetadata):
|
||||
return metadata
|
||||
return None
|
||||
@@ -0,0 +1,189 @@
|
||||
"""
|
||||
Google-style docstring parsing for ``@tool``-decorated functions.
|
||||
|
||||
Extracts the function description (everything before the first
|
||||
section header) and per-parameter descriptions (from the
|
||||
``Args:`` / ``Arguments:`` / ``Parameters:`` section).
|
||||
|
||||
Used by the schema-derivation logic to populate the function-calling
|
||||
JSON schema's ``description`` and ``properties[name].description``
|
||||
fields.
|
||||
|
||||
We don't depend on a third-party docstring library because the
|
||||
Google-style format is simple and our needs are narrow. NumPy and
|
||||
Sphinx styles are intentionally not supported — authors should
|
||||
use Google style or the explicit ``Annotated[T, Field(description=...)]``
|
||||
form for per-param descriptions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from dataclasses import dataclass
|
||||
|
||||
# Recognized section headers that terminate the description and
|
||||
# the args section. Case-sensitive (Google convention).
|
||||
_SECTION_HEADERS = (
|
||||
"Args:",
|
||||
"Arguments:",
|
||||
"Parameters:",
|
||||
"Returns:",
|
||||
"Return:",
|
||||
"Yields:",
|
||||
"Yield:",
|
||||
"Raises:",
|
||||
"Raise:",
|
||||
"Note:",
|
||||
"Notes:",
|
||||
"Example:",
|
||||
"Examples:",
|
||||
"See Also:",
|
||||
"Warning:",
|
||||
"Warnings:",
|
||||
"Attributes:",
|
||||
)
|
||||
|
||||
_ARGS_HEADERS = ("Args:", "Arguments:", "Parameters:")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ParsedDocstring:
|
||||
"""
|
||||
Result of parsing a Google-style docstring.
|
||||
|
||||
:param description: The function-level description, taken from the
|
||||
text preceding the first section header. Whitespace-trimmed.
|
||||
:param param_descriptions: Mapping from parameter name to its
|
||||
description, extracted from the ``Args:`` / ``Arguments:`` /
|
||||
``Parameters:`` section. Empty if no such section exists.
|
||||
"""
|
||||
|
||||
description: str
|
||||
param_descriptions: dict[str, str]
|
||||
|
||||
|
||||
def parse_google_docstring(doc: str) -> ParsedDocstring:
|
||||
"""
|
||||
Parse a Google-style docstring into description and per-param docs.
|
||||
|
||||
Recognizes ``Args:`` / ``Arguments:`` / ``Parameters:`` as the
|
||||
parameter-list section header. Within that section, lines like
|
||||
`` name: description`` (or `` name (type): description``)
|
||||
are parsed as parameter entries; subsequent more-indented lines
|
||||
are treated as continuations of the current parameter's
|
||||
description. The args section ends at the next recognized
|
||||
section header.
|
||||
|
||||
:param doc: The raw docstring text (typically from
|
||||
``fn.__doc__``). May be ``None``-equivalent (empty string).
|
||||
:returns: A :class:`ParsedDocstring` with the extracted description
|
||||
and parameter descriptions. Returns empty values rather than
|
||||
raising for malformed input.
|
||||
"""
|
||||
if not doc:
|
||||
return ParsedDocstring(description="", param_descriptions={})
|
||||
|
||||
cleaned = inspect.cleandoc(doc)
|
||||
if not cleaned:
|
||||
return ParsedDocstring(description="", param_descriptions={})
|
||||
|
||||
lines = cleaned.split("\n")
|
||||
|
||||
# Locate the first recognized section header. Everything before
|
||||
# it is the description; the args section (if any) is what
|
||||
# contains parameter docs.
|
||||
description_lines: list[str] = []
|
||||
args_section_lines: list[str] = []
|
||||
in_args_section = False
|
||||
in_other_section = False
|
||||
|
||||
for line in lines:
|
||||
stripped = line.strip()
|
||||
if stripped in _SECTION_HEADERS:
|
||||
in_args_section = stripped in _ARGS_HEADERS
|
||||
in_other_section = not in_args_section
|
||||
continue
|
||||
if in_args_section:
|
||||
args_section_lines.append(line)
|
||||
elif in_other_section:
|
||||
# Skip non-args sections (Returns:, Raises:, etc.)
|
||||
continue
|
||||
else:
|
||||
description_lines.append(line)
|
||||
|
||||
description = "\n".join(description_lines).strip()
|
||||
param_descriptions = _parse_args_lines(args_section_lines)
|
||||
|
||||
return ParsedDocstring(
|
||||
description=description,
|
||||
param_descriptions=param_descriptions,
|
||||
)
|
||||
|
||||
|
||||
def _parse_args_lines(lines: list[str]) -> dict[str, str]:
|
||||
"""
|
||||
Parse the body of an ``Args:`` section into per-param descriptions.
|
||||
|
||||
Param lines have the form `` name: description`` or
|
||||
`` name (type): description`` at the section's base indent.
|
||||
Lines indented further are treated as continuations of the
|
||||
current parameter's description.
|
||||
|
||||
:param lines: The lines following the ``Args:`` header (not
|
||||
including the header itself), up to the next section.
|
||||
:returns: Mapping from parameter name to its (whitespace-collapsed)
|
||||
description.
|
||||
"""
|
||||
# Determine the base indent — the indent of the first non-empty line.
|
||||
base_indent: int | None = None
|
||||
for line in lines:
|
||||
if line.strip():
|
||||
base_indent = len(line) - len(line.lstrip())
|
||||
break
|
||||
|
||||
if base_indent is None:
|
||||
return {}
|
||||
|
||||
param_descriptions: dict[str, str] = {}
|
||||
current_name: str | None = None
|
||||
current_parts: list[str] = []
|
||||
|
||||
for line in lines:
|
||||
if not line.strip():
|
||||
# Blank lines within args section are continuation separators;
|
||||
# they don't terminate a param.
|
||||
continue
|
||||
line_indent = len(line) - len(line.lstrip())
|
||||
|
||||
if line_indent == base_indent:
|
||||
# Start of a new param entry.
|
||||
if current_name is not None:
|
||||
param_descriptions[current_name] = " ".join(current_parts).strip()
|
||||
current_name = None
|
||||
current_parts = []
|
||||
|
||||
stripped = line.strip()
|
||||
if ":" not in stripped:
|
||||
# Malformed entry; skip it.
|
||||
continue
|
||||
|
||||
name_part, _, desc = stripped.partition(":")
|
||||
# Handle "name (type)" form by trimming the parenthetical.
|
||||
if "(" in name_part:
|
||||
name_only = name_part.split("(", 1)[0].strip()
|
||||
else:
|
||||
name_only = name_part.strip()
|
||||
|
||||
if name_only.isidentifier():
|
||||
current_name = name_only
|
||||
current_parts = [desc.strip()]
|
||||
# Else: not a valid param entry; ignore.
|
||||
else:
|
||||
# Continuation line for the current param.
|
||||
if current_name is not None:
|
||||
current_parts.append(line.strip())
|
||||
|
||||
if current_name is not None:
|
||||
param_descriptions[current_name] = " ".join(current_parts).strip()
|
||||
|
||||
return param_descriptions
|
||||
@@ -0,0 +1,114 @@
|
||||
"""Adapter that turns ``@tool``-decorated functions into a ToolHandler.
|
||||
|
||||
The stream-layer ``ToolHandler`` takes a list of OpenAI-shape JSON
|
||||
schemas and a single ``execute`` callable. Users who have written
|
||||
tools with the ``@tool`` decorator (Python functions with type hints
|
||||
and Google-style docstrings) shouldn't have to hand-roll that shape:
|
||||
:func:`build_tool_handler` reads each function's tool metadata and
|
||||
builds the handler for them.
|
||||
|
||||
Dispatch is by tool name. Calling an unknown tool raises — the SDK
|
||||
surfaces the error back to the agent as a tool error.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from .._tool_handler import ToolCallInfo, ToolHandler
|
||||
from ._decorator import TOOL_MARKER_ATTR, ToolMetadata
|
||||
|
||||
|
||||
def build_tool_handler(functions: list[Callable[..., Any]]) -> ToolHandler:
|
||||
"""Build a :class:`ToolHandler` from ``@tool``-decorated functions.
|
||||
|
||||
Each function must carry tool metadata attached by the
|
||||
:func:`~omnigent_client.tool` decorator (checked via
|
||||
:data:`TOOL_MARKER_ATTR`). The returned handler exposes the
|
||||
OpenAI-shape schemas the SDK sends to the server, and an
|
||||
``execute`` callable that dispatches incoming tool calls by
|
||||
name.
|
||||
|
||||
:param functions: List of ``@tool``-decorated Python functions,
|
||||
e.g. ``[get_current_time, search_docs]``. Each must be a
|
||||
module-level ``def`` or ``async def`` decorated with
|
||||
``@tool``.
|
||||
:returns: A :class:`ToolHandler` ready to pass as
|
||||
``session.tool_handler`` or via the ``tools=`` keyword on
|
||||
``OmnigentClient.query`` / ``Session.query``.
|
||||
:raises TypeError: If any function is missing the ``@tool``
|
||||
marker (i.e. wasn't decorated).
|
||||
:raises ValueError: If two functions share the same tool name
|
||||
— tool names must be unique per handler.
|
||||
"""
|
||||
if not functions:
|
||||
raise ValueError("build_tool_handler() requires at least one function")
|
||||
|
||||
schemas: list[dict[str, object]] = []
|
||||
funcs_by_name: dict[str, Callable[..., Any]] = {}
|
||||
|
||||
for fn in functions:
|
||||
meta: ToolMetadata | None = getattr(fn, TOOL_MARKER_ATTR, None)
|
||||
if meta is None:
|
||||
raise TypeError(
|
||||
f"{fn.__module__}.{fn.__qualname__} is not decorated with "
|
||||
f"@tool. Decorate it with `from omnigent_client import tool` "
|
||||
f"and apply @tool above the function definition."
|
||||
)
|
||||
if meta.name in funcs_by_name:
|
||||
raise ValueError(
|
||||
f"Duplicate tool name {meta.name!r}: "
|
||||
f"{funcs_by_name[meta.name].__qualname__} and "
|
||||
f"{fn.__qualname__} both export the same name."
|
||||
)
|
||||
funcs_by_name[meta.name] = fn
|
||||
schema: dict[str, object] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": meta.name,
|
||||
"description": meta.description,
|
||||
"parameters": meta.json_schema,
|
||||
},
|
||||
}
|
||||
schemas.append(schema)
|
||||
|
||||
async def execute(call: ToolCallInfo) -> str:
|
||||
"""Dispatch ``call`` to the matching ``@tool`` function.
|
||||
|
||||
Async functions (``async def``) are awaited on the
|
||||
event loop. Sync functions (``def``) are dispatched to
|
||||
a worker thread via ``asyncio.to_thread`` so blocking
|
||||
calls inside — ``time.sleep``, file I/O, subprocess,
|
||||
``requests`` — don't stall the event loop. Without the
|
||||
thread bounce, several concurrent ``@tool`` invocations
|
||||
(e.g. a parallel fan-out of async client tools) would
|
||||
serialize: each body would block every sibling AND any
|
||||
caller render loop sharing the loop.
|
||||
|
||||
The return value is JSON-serialized unless the function
|
||||
already returned a string (which is passed through).
|
||||
"""
|
||||
fn = funcs_by_name.get(call.name)
|
||||
if fn is None:
|
||||
# The SDK will surface this back to the agent as a tool
|
||||
# error — this typically means the LLM invented a tool
|
||||
# name that wasn't in the schemas we sent.
|
||||
raise KeyError(f"Unknown tool {call.name!r}. Registered: {sorted(funcs_by_name)}")
|
||||
if inspect.iscoroutinefunction(fn):
|
||||
result = await fn(**call.arguments)
|
||||
else:
|
||||
# Sync body — route to a worker thread so it
|
||||
# doesn't block the event loop (see the fan-out
|
||||
# serialization case above).
|
||||
result = await asyncio.to_thread(lambda: fn(**call.arguments))
|
||||
if isinstance(result, str):
|
||||
return result
|
||||
# Pydantic models and dataclasses commonly aren't JSON-ready
|
||||
# out of the box — ``default=str`` handles datetime/UUID/etc.
|
||||
return json.dumps(result, default=str)
|
||||
|
||||
return ToolHandler(schemas=schemas, execute=execute)
|
||||
@@ -0,0 +1,267 @@
|
||||
"""
|
||||
Schema derivation for ``@tool``-decorated functions.
|
||||
|
||||
Given a typed Python function, produce the function-calling JSON
|
||||
schema the LLM sees. The pipeline:
|
||||
|
||||
1. Inspect the signature for parameters, annotations, and defaults.
|
||||
2. Parse the Google-style docstring for description and per-param
|
||||
descriptions (see :mod:`omnigent.tools._docstring`).
|
||||
3. Build a Pydantic model from the parameters via ``create_model``;
|
||||
Pydantic does the heavy lifting for type → schema (primitives,
|
||||
Pydantic models, ``Optional``, ``Literal``, ``Annotated[..., Field]``,
|
||||
etc.).
|
||||
4. Apply strict-mode normalization (see
|
||||
:mod:`omnigent.tools._strict`) when ``strict=True``.
|
||||
|
||||
Permissive types (``Any``, ``object``, missing annotations) are
|
||||
allowed but produce an INFO-level warning so authors can find them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import logging
|
||||
import typing
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Annotated, Any, get_args, get_origin
|
||||
|
||||
from pydantic import Field, create_model
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from ._docstring import parse_google_docstring
|
||||
from ._state import ToolState
|
||||
from ._strict import ensure_strict_schema
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
# Reserved parameter name for framework-injected per-conversation
|
||||
# per-agent tool state. A ``@tool`` function that declares a
|
||||
# parameter with this name receives a live :class:`ToolState` at
|
||||
# call time; the parameter is stripped from the LLM-facing JSON
|
||||
# schema. Convention over configuration — every stateful tool uses
|
||||
# the same identifier.
|
||||
STATE_PARAM_NAME = "tool_state"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FunctionSchemaResult:
|
||||
"""
|
||||
Output of :func:`build_function_schema`.
|
||||
|
||||
:param description: Function-level description, derived from
|
||||
the docstring's leading paragraph(s).
|
||||
:param parameters_json_schema: JSON schema for the function's
|
||||
parameters, in the OpenAI function-calling shape (an
|
||||
``object`` schema with ``properties`` and ``required``).
|
||||
Already normalized to strict mode if requested.
|
||||
:param return_annotation: The function's return type annotation,
|
||||
or ``None`` if no return annotation was provided. Used by
|
||||
the executor to deserialize the tool's return value via
|
||||
``pydantic.TypeAdapter``.
|
||||
"""
|
||||
|
||||
description: str
|
||||
parameters_json_schema: dict[str, Any]
|
||||
return_annotation: type[Any] | None
|
||||
|
||||
|
||||
def build_function_schema(
|
||||
fn: Callable[..., Any],
|
||||
*,
|
||||
strict: bool = True,
|
||||
) -> FunctionSchemaResult:
|
||||
"""
|
||||
Build the function-calling schema for a Python function.
|
||||
|
||||
:param fn: The Python function to derive a schema for. Must be
|
||||
a module-level ``def`` or ``async def`` (the ``@tool``
|
||||
decorator enforces this elsewhere; we do not re-validate
|
||||
here).
|
||||
:param strict: If ``True``, apply strict-mode normalization to
|
||||
the resulting schema (see :mod:`._strict`).
|
||||
:returns: A :class:`FunctionSchemaResult` with description,
|
||||
JSON schema, and return-type annotation.
|
||||
"""
|
||||
sig = inspect.signature(fn)
|
||||
type_hints = typing.get_type_hints(fn, include_extras=True)
|
||||
parsed_doc = parse_google_docstring(fn.__doc__ or "")
|
||||
|
||||
fields: dict[str, tuple[Any, FieldInfo]] = {}
|
||||
for name, param in sig.parameters.items():
|
||||
ann = type_hints.get(name, Any)
|
||||
|
||||
# Framework-injected parameter — reserved by convention:
|
||||
# any parameter named exactly ``tool_state`` is filled by
|
||||
# the runtime with a :class:`ToolState`. Skipped from the
|
||||
# LLM-facing schema; the LLM has no way to supply it.
|
||||
# Enforce the convention: if someone types a param as
|
||||
# ToolState but names it something else, fail loud so they
|
||||
# know the right contract.
|
||||
if name == STATE_PARAM_NAME:
|
||||
if ann is not ToolState and ann is not Any:
|
||||
raise TypeError(
|
||||
f"@tool function {fn.__name__!r} declares parameter "
|
||||
f"'{STATE_PARAM_NAME}' with unexpected type "
|
||||
f"{ann!r}. It must be typed as ToolState "
|
||||
f"(or left unannotated); any other type is a bug."
|
||||
)
|
||||
continue
|
||||
if ann is ToolState:
|
||||
raise TypeError(
|
||||
f"@tool function {fn.__name__!r} types parameter "
|
||||
f"{name!r} as ToolState but the parameter must be named "
|
||||
f"{STATE_PARAM_NAME!r}. Rename it and the framework "
|
||||
f"will inject a live ToolState at call time."
|
||||
)
|
||||
|
||||
_warn_if_permissive(fn.__name__, name, ann)
|
||||
|
||||
doc_desc = parsed_doc.param_descriptions.get(name)
|
||||
default = param.default if param.default is not inspect.Parameter.empty else ...
|
||||
field_info = _build_field_info(ann, default, doc_desc)
|
||||
fields[name] = (ann, field_info)
|
||||
|
||||
if fields:
|
||||
# Pydantic uses the model name when generating $defs refs;
|
||||
# capitalize so it looks reasonable in the schema output.
|
||||
model_name = f"{_pascal_case(fn.__name__)}Args"
|
||||
# mypy can't statically narrow create_model's overload for our
|
||||
# dynamic field dict, but pydantic accepts (Type, FieldInfo)
|
||||
# tuples here — they're the documented field-definition shape.
|
||||
Model = create_model(model_name, **fields) # type: ignore[call-overload]
|
||||
params_schema: dict[str, Any] = Model.model_json_schema()
|
||||
else:
|
||||
# Zero-arg tool: the schema is an empty object.
|
||||
params_schema = {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
}
|
||||
|
||||
if strict:
|
||||
params_schema = ensure_strict_schema(params_schema)
|
||||
|
||||
return_annotation = type_hints.get("return")
|
||||
|
||||
return FunctionSchemaResult(
|
||||
description=parsed_doc.description,
|
||||
parameters_json_schema=params_schema,
|
||||
return_annotation=return_annotation,
|
||||
)
|
||||
|
||||
|
||||
def _build_field_info(
|
||||
annotation: Any,
|
||||
default: Any,
|
||||
doc_description: str | None,
|
||||
) -> FieldInfo:
|
||||
"""
|
||||
Construct a Pydantic ``FieldInfo`` for one parameter.
|
||||
|
||||
Handles three description sources, with this priority:
|
||||
1. An explicit ``Field(description=...)`` in
|
||||
``Annotated[T, Field(description=...)]``.
|
||||
2. A bare string in ``Annotated[T, "desc"]`` (a common shorthand
|
||||
supported by some agent SDKs).
|
||||
3. The docstring entry for this parameter (Google-style ``Args:``).
|
||||
|
||||
:param annotation: The parameter's type annotation, possibly
|
||||
wrapped in ``Annotated[...]``.
|
||||
:param default: The parameter's default value, or ``...`` if
|
||||
the parameter is required.
|
||||
:param doc_description: Description from the docstring's
|
||||
``Args:`` section, or ``None`` if absent.
|
||||
:returns: A ``FieldInfo`` ready to pass to
|
||||
``pydantic.create_model``. The default (if any) and
|
||||
description are baked in at construction time so
|
||||
``model_json_schema`` picks them up correctly.
|
||||
"""
|
||||
# Pull metadata out of Annotated[T, ...] for description discovery.
|
||||
annotated_str_desc: str | None = None
|
||||
annotated_field: FieldInfo | None = None
|
||||
if get_origin(annotation) is Annotated:
|
||||
for extra in get_args(annotation)[1:]:
|
||||
if isinstance(extra, FieldInfo) and annotated_field is None:
|
||||
annotated_field = extra
|
||||
elif isinstance(extra, str) and annotated_str_desc is None:
|
||||
annotated_str_desc = extra
|
||||
|
||||
# Determine the effective description with the priority above.
|
||||
description: str | None = None
|
||||
if annotated_field is not None and annotated_field.description is not None:
|
||||
description = annotated_field.description
|
||||
elif annotated_str_desc is not None:
|
||||
description = annotated_str_desc
|
||||
elif doc_description:
|
||||
description = doc_description
|
||||
|
||||
# Field(default=PydanticUndefined) is the marker for "required";
|
||||
# we map our `...` sentinel to it via PydanticUndefined import.
|
||||
# Easier: build the constructor kwargs and let Pydantic translate
|
||||
# default=... directly (it accepts ``...`` as "required" too).
|
||||
field_kwargs: dict[str, Any] = {}
|
||||
if description is not None:
|
||||
field_kwargs["description"] = description
|
||||
if default is not ...:
|
||||
field_kwargs["default"] = default
|
||||
|
||||
if annotated_field is not None:
|
||||
# Merge: preserve other metadata (gt/lt/min_length/etc.) from
|
||||
# the author's Field, but override description and default.
|
||||
# merge_field_infos's stub return is too loose for mypy; cast.
|
||||
merged: FieldInfo = FieldInfo.merge_field_infos(annotated_field, FieldInfo(**field_kwargs))
|
||||
return merged
|
||||
|
||||
# pydantic.Field stub returns Any (it's polymorphic by default).
|
||||
# We're constructing a fresh FieldInfo; cast to the documented type.
|
||||
field_obj: FieldInfo = Field(**field_kwargs)
|
||||
return field_obj
|
||||
|
||||
|
||||
def _warn_if_permissive(fn_name: str, param_name: str, annotation: Any) -> None:
|
||||
"""
|
||||
Log a warning if a parameter's type provides no validation constraint.
|
||||
|
||||
``Any``, ``object``, and missing annotations all produce a
|
||||
permissive schema (no ``type`` field) that the LLM can fill
|
||||
with arbitrary structure. Useful but easy to write by accident.
|
||||
|
||||
:param fn_name: The decorated function's ``__name__``, for the
|
||||
log message, e.g. ``"process_payload"``.
|
||||
:param param_name: The offending parameter's name, e.g.
|
||||
``"data"``.
|
||||
:param annotation: The annotation as resolved by
|
||||
``typing.get_type_hints``.
|
||||
"""
|
||||
# Strip Annotated[...] so we inspect the underlying type.
|
||||
underlying = annotation
|
||||
if get_origin(underlying) is Annotated:
|
||||
underlying = get_args(underlying)[0]
|
||||
|
||||
if underlying is Any or underlying is object:
|
||||
type_name = (
|
||||
"Any" if underlying is Any else getattr(underlying, "__name__", str(underlying))
|
||||
)
|
||||
_logger.info(
|
||||
"Tool '%s' parameter '%s' has no concrete type annotation "
|
||||
"(resolved to %s); LLM will get a permissive schema with "
|
||||
"no validation.",
|
||||
fn_name,
|
||||
param_name,
|
||||
type_name,
|
||||
)
|
||||
|
||||
|
||||
def _pascal_case(snake: str) -> str:
|
||||
"""
|
||||
Convert a snake_case identifier to PascalCase.
|
||||
|
||||
Used to give the dynamically-created Pydantic model a readable
|
||||
name in schema ``$defs`` references.
|
||||
|
||||
:param snake: A snake_case identifier, e.g. ``"word_count"``.
|
||||
:returns: PascalCase form, e.g. ``"WordCount"``.
|
||||
"""
|
||||
return "".join(part.capitalize() for part in snake.split("_") if part)
|
||||
@@ -0,0 +1,300 @@
|
||||
"""Per-agent ToolState for stateful ``@tool`` functions.
|
||||
|
||||
See ``designs/TOOL_STATE.md`` for the full design. The primitive is
|
||||
a simple key-value store, JSON-serialized, scoped to one
|
||||
(conversation, agent) pair via the storage directory provided by
|
||||
the framework. The ``@tool`` decorator hides ``ToolState``-typed
|
||||
parameters from the LLM-facing schema; the subprocess runner
|
||||
reconstructs a ``ToolState`` from the directory path and injects it
|
||||
when the tool function is called.
|
||||
|
||||
Tool authors see::
|
||||
|
||||
from omnigent_client import tool, ToolState
|
||||
|
||||
@tool
|
||||
def add_task(desc: str, state: ToolState) -> str:
|
||||
with state.transaction("queue") as q:
|
||||
q = q or []
|
||||
q.append({"desc": desc})
|
||||
return f"#{len(q) - 1}"
|
||||
|
||||
and nothing else.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager, suppress
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import fcntl as _fcntl
|
||||
except ImportError: # pragma: no cover - exercised on native Windows
|
||||
_fcntl = None # type: ignore[assignment]
|
||||
|
||||
if os.name == "nt":
|
||||
import msvcrt as _msvcrt
|
||||
else: # pragma: no cover - exercised on POSIX
|
||||
_msvcrt = None # type: ignore[assignment]
|
||||
|
||||
# Subdirectory segments reserved by the framework — JSON key files
|
||||
# live at ``{root}/{key}.json``. Keys must not contain path
|
||||
# separators; we sanitize eagerly rather than allowing the bug to
|
||||
# surface as directory traversal.
|
||||
_KEY_SUFFIX = ".json"
|
||||
_THREAD_LOCKS_GUARD = threading.Lock()
|
||||
_THREAD_LOCKS: dict[Path, threading.Lock] = {}
|
||||
|
||||
|
||||
class ToolState:
|
||||
"""Per-agent, per-conversation key-value state for ``@tool`` functions.
|
||||
|
||||
Values are JSON-serialized. The keyspace is shared across every
|
||||
tool invoked for the same registered agent within the same
|
||||
conversation. Use :meth:`transaction` for atomic read-modify-write;
|
||||
plain :meth:`get` and :meth:`set` do not serialize concurrent
|
||||
writers on the same key.
|
||||
|
||||
Instances are constructed by the framework. Tool authors receive
|
||||
a ``ToolState`` by declaring a parameter of this type on their
|
||||
``@tool``-decorated function; the decorator strips the parameter
|
||||
from the LLM-facing schema and the subprocess runner injects the
|
||||
live ``ToolState`` at call time.
|
||||
|
||||
:param root: The directory this namespace lives in, e.g.
|
||||
``{workspace}/.tool_state/{agent_id}``. The directory does
|
||||
not need to exist yet; it is created lazily on first write.
|
||||
"""
|
||||
|
||||
def __init__(self, root: Path) -> None:
|
||||
self._root = root
|
||||
|
||||
# ── Primary API ──────────────────────────────────────────
|
||||
|
||||
def get(self, key: str, *, default: Any = None) -> Any:
|
||||
"""Return the stored value at ``key``, or ``default`` if absent.
|
||||
|
||||
:param key: The state key, e.g. ``"queue"``.
|
||||
:param default: Value to return when the key has never been
|
||||
written. ``None`` by default.
|
||||
:returns: The deserialized JSON value, or ``default``.
|
||||
"""
|
||||
path = self._path_for(key)
|
||||
if not path.exists():
|
||||
return default
|
||||
with path.open("r") as f:
|
||||
# Shared lock: allow parallel reads, block concurrent writers
|
||||
# briefly so we see a complete JSON payload.
|
||||
_lock_file(f, exclusive=False)
|
||||
try:
|
||||
return json.loads(f.read() or "null")
|
||||
finally:
|
||||
_unlock_file(f)
|
||||
|
||||
def set(self, key: str, value: Any) -> None:
|
||||
"""Replace (or create) the value at ``key``. JSON-serialized.
|
||||
|
||||
Non-atomic relative to concurrent writers on the same key —
|
||||
use :meth:`transaction` for read-modify-write sequences.
|
||||
|
||||
:param key: The state key.
|
||||
:param value: Any JSON-serializable value.
|
||||
"""
|
||||
path = self._path_for(key)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
# Write through a temp file + rename so a reader never sees
|
||||
# a half-written JSON payload even without a lock.
|
||||
tmp = path.with_suffix(_KEY_SUFFIX + ".tmp")
|
||||
with tmp.open("w") as f:
|
||||
json.dump(value, f)
|
||||
tmp.replace(path)
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
"""Remove ``key``. No-op if absent.
|
||||
|
||||
:param key: The state key to remove.
|
||||
"""
|
||||
path = self._path_for(key)
|
||||
# Idempotent delete — tools commonly don't know whether the
|
||||
# key was ever set.
|
||||
with suppress(FileNotFoundError):
|
||||
path.unlink()
|
||||
|
||||
def keys(self) -> list[str]:
|
||||
"""List all keys currently stored in this namespace.
|
||||
|
||||
:returns: Sorted list of keys, e.g. ``["counter", "queue"]``.
|
||||
Empty list if nothing has been written yet.
|
||||
"""
|
||||
if not self._root.exists():
|
||||
return []
|
||||
return sorted(p.stem for p in self._root.iterdir() if p.suffix == _KEY_SUFFIX)
|
||||
|
||||
def __contains__(self, key: object) -> bool:
|
||||
"""Return whether ``key`` currently exists in this namespace.
|
||||
|
||||
:param key: Candidate key, e.g. ``"queue"``.
|
||||
:returns: ``True`` when ``key`` is a valid stored key, else ``False``.
|
||||
"""
|
||||
if not isinstance(key, str):
|
||||
return False
|
||||
return self._path_for(key).exists()
|
||||
|
||||
@contextmanager
|
||||
def transaction(self, key: str, *, default: Any = None) -> Iterator[Any]:
|
||||
"""Atomic read-modify-write for one key.
|
||||
|
||||
Typical usage — supply a ``default`` so first-time callers
|
||||
get a usable container without a ``None`` check::
|
||||
|
||||
with state.transaction("queue", default=[]) as queue:
|
||||
queue.append(item)
|
||||
# queue is written back on normal exit.
|
||||
|
||||
The yielded value is the current contents, or a fresh
|
||||
``default`` if the key was never set. Mutating the yielded
|
||||
object in place is the expected pattern — the same object
|
||||
is serialized back on exit. Rebinding the local name inside
|
||||
the ``with`` block does NOT propagate (Python closures), so
|
||||
for "replace the value" semantics use :meth:`set` explicitly.
|
||||
|
||||
On a normal exit the yielded object is JSON-serialized and
|
||||
written back. On exception no write happens — the prior
|
||||
value is preserved.
|
||||
|
||||
:param key: The state key to lock + read + write.
|
||||
:param default: Value yielded when the key has never been
|
||||
written. Defaults to ``None``. Pass ``[]`` or ``{}``
|
||||
(or any JSON-serializable value) to skip the absent-key
|
||||
branch in caller code.
|
||||
:yields: The current value at ``key``, or ``default`` if
|
||||
the key has no stored value yet. Mutate in place.
|
||||
"""
|
||||
path = self._path_for(key)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
thread_lock = _get_thread_lock(path)
|
||||
# ``a+`` creates the file if missing and positions at end;
|
||||
# we seek to 0 before read. Opening with ``r+`` would fail
|
||||
# when the file doesn't exist yet, which is a common first-
|
||||
# call case for a tool.
|
||||
with thread_lock, path.open("a+") as f:
|
||||
_lock_file(f, exclusive=True)
|
||||
try:
|
||||
value = _read_transaction_value(f, default)
|
||||
yield value
|
||||
_write_transaction_value(f, value)
|
||||
finally:
|
||||
_unlock_file(f)
|
||||
|
||||
# ── Internals ────────────────────────────────────────────
|
||||
|
||||
def _path_for(self, key: str) -> Path:
|
||||
"""Resolve ``key`` to the on-disk path, rejecting traversal.
|
||||
|
||||
:param key: Caller-supplied key.
|
||||
:returns: The ``{root}/{key}.json`` path.
|
||||
:raises ValueError: If ``key`` is empty, contains a path
|
||||
separator, or starts with a dot (no hidden or
|
||||
escaped paths).
|
||||
"""
|
||||
if not key:
|
||||
raise ValueError("ToolState key must be a non-empty string")
|
||||
if "/" in key or "\\" in key or key.startswith("."):
|
||||
# Rejects traversal and hidden-file sigils. Authors who
|
||||
# really need slashes can encode them (e.g. "a__b") —
|
||||
# we'd rather break loudly than accept quiet bugs.
|
||||
raise ValueError(
|
||||
f"ToolState key {key!r} contains an illegal character. "
|
||||
f"Keys must be plain names (no '/', '\\', or leading '.')."
|
||||
)
|
||||
return self._root / f"{key}{_KEY_SUFFIX}"
|
||||
|
||||
|
||||
def _read_transaction_value(f: Any, default: Any) -> Any:
|
||||
"""Seek to 0 and decode the JSON value under the open file handle.
|
||||
|
||||
Returns ``default`` when the file is empty (first-time use of
|
||||
the key). Factored out of :meth:`ToolState.transaction` so the
|
||||
context manager stays short.
|
||||
|
||||
:param f: Open file handle positioned anywhere; will be seek(0)ed.
|
||||
:param default: Value to return on empty/whitespace content.
|
||||
:returns: Decoded JSON value or ``default``.
|
||||
"""
|
||||
f.seek(0)
|
||||
raw = f.read()
|
||||
if raw.strip():
|
||||
return json.loads(raw)
|
||||
return default
|
||||
|
||||
|
||||
def _write_transaction_value(f: Any, value: Any) -> None:
|
||||
"""Truncate and re-serialize ``value`` as JSON under the file handle.
|
||||
|
||||
Caller must hold the exclusive flock before calling. Factored
|
||||
out of :meth:`ToolState.transaction` so the context manager
|
||||
stays short.
|
||||
|
||||
:param f: Open file handle (must support ``r+``-style truncate).
|
||||
:param value: Any JSON-serializable value to persist.
|
||||
"""
|
||||
f.seek(0)
|
||||
f.truncate()
|
||||
json.dump(value, f)
|
||||
# Flush before releasing the flock held by the caller. Python file objects
|
||||
# buffer writes; if we unlock before flushing, another process can acquire
|
||||
# the lock and read stale on-disk contents, losing the prior update.
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
|
||||
|
||||
def _lock_file(f: Any, *, exclusive: bool) -> None:
|
||||
"""Acquire an advisory file lock for ``f``.
|
||||
|
||||
POSIX uses ``fcntl.flock``. Native Windows has no ``fcntl`` module, so it
|
||||
locks one byte with ``msvcrt.locking``; that API is exclusive-only, which
|
||||
is conservative for reads but preserves cross-process serialization.
|
||||
"""
|
||||
if _fcntl is not None:
|
||||
_fcntl.flock(f, _fcntl.LOCK_EX if exclusive else _fcntl.LOCK_SH)
|
||||
return
|
||||
if _msvcrt is None: # pragma: no cover - defensive for unusual platforms
|
||||
return
|
||||
f.seek(0)
|
||||
_msvcrt.locking(f.fileno(), _msvcrt.LK_LOCK, 1)
|
||||
|
||||
|
||||
def _unlock_file(f: Any) -> None:
|
||||
"""Release a lock acquired by :func:`_lock_file`."""
|
||||
if _fcntl is not None:
|
||||
_fcntl.flock(f, _fcntl.LOCK_UN)
|
||||
return
|
||||
if _msvcrt is None: # pragma: no cover - defensive for unusual platforms
|
||||
return
|
||||
f.seek(0)
|
||||
_msvcrt.locking(f.fileno(), _msvcrt.LK_UNLCK, 1)
|
||||
|
||||
|
||||
def _get_thread_lock(path: Path) -> threading.Lock:
|
||||
"""Return the per-key in-process mutex for ``path``.
|
||||
|
||||
``flock`` serializes across processes, but threads in the same
|
||||
process can still interleave on separate file descriptors. This
|
||||
helper layers a per-path ``threading.Lock`` on top so
|
||||
``transaction()`` is atomic under both thread and process
|
||||
contention.
|
||||
|
||||
:param path: The key file path, e.g. ``Path("/tmp/state/queue.json")``.
|
||||
:returns: The shared mutex guarding that path within this process.
|
||||
"""
|
||||
with _THREAD_LOCKS_GUARD:
|
||||
lock = _THREAD_LOCKS.get(path)
|
||||
if lock is None:
|
||||
lock = threading.Lock()
|
||||
_THREAD_LOCKS[path] = lock
|
||||
return lock
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Strict JSON-schema normalization for ``@tool``-derived schemas.
|
||||
|
||||
Strict-mode schemas have two extra constraints beyond ordinary
|
||||
JSON Schema:
|
||||
|
||||
- Every object schema sets ``additionalProperties: false``.
|
||||
- Every property in an object is listed in ``required`` (Python
|
||||
defaults are still applied on the executor side after the LLM
|
||||
emits a value).
|
||||
|
||||
This is the form most major LLM providers accept for
|
||||
function-calling tool schemas without further coercion. Authors
|
||||
who hit a real schema that strict mode breaks can opt out via
|
||||
``@tool(strict=False)``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def ensure_strict_schema(schema: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Recursively normalize a JSON schema to strict-mode rules.
|
||||
|
||||
Returns a new dict (does not mutate the input). Recurses into
|
||||
``properties``, ``items``, ``$defs``, and union variants
|
||||
(``anyOf`` / ``oneOf`` / ``allOf``).
|
||||
|
||||
:param schema: A JSON schema dict (typically produced by
|
||||
:func:`pydantic.BaseModel.model_json_schema`).
|
||||
:returns: A new dict with strict-mode constraints applied.
|
||||
"""
|
||||
if not isinstance(schema, dict):
|
||||
return schema
|
||||
|
||||
out: dict[str, Any] = dict(schema)
|
||||
|
||||
# Recurse into nested $defs FIRST so the recursion sees the
|
||||
# normalized definitions when it walks references.
|
||||
if "$defs" in out:
|
||||
out["$defs"] = {k: ensure_strict_schema(v) for k, v in out["$defs"].items()}
|
||||
|
||||
# Recurse into union variants.
|
||||
for union_key in ("anyOf", "oneOf", "allOf"):
|
||||
if union_key in out and isinstance(out[union_key], list):
|
||||
out[union_key] = [ensure_strict_schema(variant) for variant in out[union_key]]
|
||||
|
||||
# Recurse into array items.
|
||||
if "items" in out:
|
||||
out["items"] = ensure_strict_schema(out["items"])
|
||||
|
||||
if out.get("type") == "object":
|
||||
properties = out.get("properties", {}) or {}
|
||||
out["additionalProperties"] = False
|
||||
out["properties"] = {k: ensure_strict_schema(v) for k, v in properties.items()}
|
||||
# Strict mode requires every property to be listed as required.
|
||||
# The default value (if any) still applies on the executor side
|
||||
# via Pydantic — strict mode only governs what the LLM emits.
|
||||
if properties:
|
||||
out["required"] = list(properties.keys())
|
||||
|
||||
return out
|
||||
@@ -0,0 +1,25 @@
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "omnigent-client"
|
||||
version = "0.6.0.dev0"
|
||||
description = "Python client SDK for the omnigent server API"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
dependencies = [
|
||||
# Sibling package — provides the SessionStreamEventType /
|
||||
# StreamEvent source of truth that the SDK validates SSE
|
||||
# envelopes against (see _sessions.py, _sse.py). Resolved via
|
||||
# the relative path below so the SDK can be installed
|
||||
# editable alongside the root ``omnigent`` package. Version-locked
|
||||
# (==) so a published SDK always pairs with the server release it
|
||||
# shipped with (release-omnigent.yml verifies the pin).
|
||||
"omnigent==0.6.0.dev0",
|
||||
"httpx>=0.27",
|
||||
"pydantic>=2.0,<3",
|
||||
]
|
||||
|
||||
[tool.uv.sources]
|
||||
omnigent = { path = "../..", editable = true }
|
||||
Reference in New Issue
Block a user