chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
Reference in New Issue
Block a user