wehub/omnigent-ai--omnigent
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
omnigent-ai--omnigent/omnigent/inner/openai_agents_sdk_executor.py
T
wehub-resource-sync 3de48288cb chore: import upstream snapshot with attribution
2026-07-13 13:12:00 +08:00

1877 lines
83 KiB
Python
Raw Permalink Blame History

"""OpenAIAgentsSDKExecutor: run agents using the OpenAI Agents SDK.
This executor uses the OpenAI Agents SDK as the agent runtime while keeping
Omnigent as the system of record for tools, policies, and session history.
Omnigent tools are exposed as SDK FunctionTools whose callbacks route back
through the Session's tool registry.
Conversation state is delegated to the SDK's session layer. Each Omnigent
Session gets its own SDK SQLiteSession keyed by the Omnigent session id.
"""
from __future__ import annotations
import asyncio
import base64
import contextlib
import json
import logging
import os
import subprocess
import time
from collections.abc import AsyncIterator, Awaitable, Callable, Generator
from dataclasses import dataclass
from types import ModuleType
from typing import Any, Literal, Protocol, TypeAlias, cast
import httpx
from omnigent.llms._usage_observer import notify_from_dict as _notify_usage_from_dict
from omnigent.llms.errors import is_context_length_exceeded as _is_context_length_exceeded
from omnigent.reasoning_effort import OPENAI_AGENTS_EFFORTS, validate_effort
from omnigent.spec.types import RetryPolicy
from .executor import (
Executor,
ExecutorConfig,
ExecutorError,
ExecutorEvent,
Message,
ReasoningChunk,
TextChunk,
ToolCallComplete,
ToolCallRequest,
ToolSpec,
TurnComplete,
classify_tool_result,
split_transient_tail,
)
from .open_responses_sdk import (
_OPENAI_KEY_PLACEHOLDER,
_convert_messages_to_responses,
_databricks_openai_base_url,
)
logger = logging.getLogger(__name__)
_OPENAI_AGENTS_DEFAULT_MODEL = "gpt-5.3-codex"
_DATABRICKS_OPENAI_AGENTS_DEFAULT_MODEL = "databricks-gpt-5-5"
# Total run attempts per turn (1 initial + retries). The Databricks
# gateway occasionally returns a completed-but-empty turn (status
# completed, no text / no tool calls / no output items); a single
# retry recovers it. Kept small and hardcoded: more retries would
# risk wedging long e2e turns under the shard timeout, and the
# transient is rare.
_EMPTY_TURN_MAX_ATTEMPTS = 2
# SDK ``RunItem.type`` values that are bookkeeping, not user-visible
# output. A turn whose only new items are these (and which produced no
# text and no tool activity) is treated as empty and retried. Excluding
# by name (rather than allow-listing output types) is the conservative
# direction: an unknown future item type counts as output and is NOT
# retried.
_NON_OUTPUT_ITEM_TYPES: frozenset[str] = frozenset({"reasoning_item", "compaction_item"})
# Replay items persisted to the SDK Session — heterogeneous Responses-API
# input items (function_call / function_call_output / message / etc.).
# The SDK declares this as ``TResponseInputItem``, which is a TypedDict
# union; we only care about two fields (``id``, ``call_id``) and pass the
# rest through opaquely.
ReplayItem: TypeAlias = dict[str, Any] # type: ignore[explicit-any]
# Raw pydantic tool-call / tool-output items from the Responses API. Their
# concrete union is huge (``ToolCallItemTypes`` / ``ToolCallOutputTypes``)
# and the fields we touch depend on the runtime member, so we keep it
# opaque at the boundary and narrow with ``isinstance(..., dict)`` + a
# Protocol cast at each read site.
RawToolItem: TypeAlias = Any # type: ignore[explicit-any]
# The openai ``AsyncOpenAI`` client. Kept as ``Any`` so ``openai`` stays
# an optional import at type-check time — the executor only constructs
# one when instantiated.
AsyncOpenAIClient: TypeAlias = Any # type: ignore[explicit-any]
# Tool executor callable wired in by ``omnigent.Session``. The result
# is JSON-ish (dict[str, Any]) but the static type leaks ``Any`` through
# the kwargs dict. Constrained here to match omnigent.executor.
ToolExecutor: TypeAlias = Callable[ # type: ignore[explicit-any]
[str, dict[str, Any]], Awaitable[dict[str, Any]]
]
# SDK objects we treat as opaque: agent/tool instances and tool results
# are structurally duck-typed by callers and the Protocols above.
SDKAgent: TypeAlias = Any # type: ignore[explicit-any]
SDKTool: TypeAlias = Any # type: ignore[explicit-any]
ToolResult: TypeAlias = Any # type: ignore[explicit-any]
# Parsed tool arguments / tool-call-output dict. JSON-shaped bag handed
# to the Omnigent tool executor.
ToolArgs: TypeAlias = dict[str, Any] # type: ignore[explicit-any]
def _normalize_responses_items_for_chat(
items: list[dict[str, Any]],
) -> list[dict[str, Any]]:
"""Apply :func:`_normalize_content_blocks_for_chat` to every message item.
Walks the full Responses-API item list produced by
:func:`~omnigent.inner.open_responses_sdk._convert_messages_to_responses`
and normalises ``input_file`` blocks in any ``message`` item's ``content``
list. Non-message items (function_call, function_call_output, …) pass
through unchanged.
:param items: Responses-API input item dicts.
:returns: New list with normalised ``input_file`` blocks in message
content. Items without ``input_file`` blocks are returned as-is.
"""
result: list[dict[str, Any]] = []
for item in items:
if item.get("type") == "message":
raw_content = item.get("content")
if isinstance(raw_content, list):
normalized_content = _normalize_content_blocks_for_chat(raw_content)
if normalized_content is not raw_content:
item = {**item, "content": normalized_content}
result.append(item)
return result
def _normalize_content_blocks_for_chat(
content: list[dict[str, Any]],
) -> list[dict[str, Any]]:
"""Normalize content blocks before handing them to the openai-agents Runner.
Converts ``input_file`` blocks whose ``file_data`` is a ``data:`` URI into
``input_text`` blocks by decoding the base64 payload. Also strips
Omnigent-only metadata (for example ``filename``) from known Responses
content block types before the OpenAI Agents SDK serializes them to the
provider API.
Background: :mod:`omnigent.runtime.content_resolver` resolves uploaded
file IDs to inline ``data:<mime>;base64,<data>`` URIs stored in
``file_data`` / ``image_url`` while preserving UI/history metadata such as
``filename``. That metadata is useful in Omnigent, but the OpenAI API
rejects unknown content-part fields (for example
``input[0].content[0].filename``). The openai-agents SDK's
``chatcmpl_converter`` also passes raw ``file_data`` strings to the Chat
Completions ``file`` content type, which expects **plain base64** — not a
data URI. Additionally, the Databricks GPT endpoint may not support the
``file`` content block type at all, so textual files are converted to
``input_text``. Images are preserved as ``input_image`` with a
conventional ``data:image/<type>;base64,...`` URL so vision-capable
Chat Completions providers can see the pixels.
:param content: List of Responses-API content-block dicts, possibly
containing ``input_file`` blocks with ``file_data`` data URIs or
metadata-bearing ``input_image`` blocks.
:returns: New list with each ``input_file`` block replaced by an
``input_text`` block carrying the decoded text content, and known
content blocks stripped to provider-supported keys. Binary
(non-``text/*``) file blocks, undecodable, or empty file blocks are
silently dropped. Unknown block types pass through unchanged.
"""
result: list[dict[str, Any]] = []
changed = False
for block in content:
block_type = block.get("type")
if block_type == "input_file":
changed = True
file_data: str = block.get("file_data", "")
if file_data.startswith("data:"):
try:
meta, b64 = file_data.split(",", 1)
mime = meta.split(";")[0].replace("data:", "")
if not mime.startswith("text/"):
continue # binary files (PDF, etc.) can't be inlined as text
text = base64.b64decode(b64).decode("utf-8", errors="replace")
except Exception: # noqa: BLE001
text = ""
else:
text = file_data
if text:
result.append({"type": "input_text", "text": text})
# else: drop empty / undecodable blocks
elif block_type in {"input_text", "output_text", "text"}:
sanitized = _copy_known_keys(block, ("type", "text"))
changed = changed or sanitized != block
result.append(sanitized)
elif block_type == "input_image":
# Preserve inline images; strip Omnigent-only metadata.
sanitized = _copy_known_keys(block, ("type", "image_url", "file_id", "detail"))
changed = changed or sanitized != block
result.append(sanitized)
else:
result.append(block)
# Return original list when nothing changed so callers can use
# identity check (``is``) to avoid unnecessary copies.
return result if changed else content
def _copy_known_keys(block: dict[str, Any], keys: tuple[str, ...]) -> dict[str, Any]:
"""Return *block* with only keys accepted by provider content schemas.
Omnigent content blocks can carry UI/history metadata such as
``filename``. Known OpenAI/Responses block types are strict at the API
boundary, so those metadata fields must be removed from the payload handed
to the OpenAI Agents SDK while leaving persisted conversation state intact.
"""
return {key: block[key] for key in keys if key in block}
class _SDKSession(Protocol):
"""Structural view of ``agents.memory.Session`` — the four async
methods + two attributes the SDK's Session protocol requires.
Using a local Protocol lets ``_SanitizingSession`` satisfy the SDK
contract without inheriting from it, and lets tests plug in fake
sessions without importing the SDK.
"""
session_id: str
session_settings: Any # type: ignore[explicit-any] # SDK's SessionSettings | None; opaque here
async def get_items(self, limit: int | None = None) -> list[ReplayItem]: ...
async def add_items(self, items: list[ReplayItem]) -> None: ...
async def pop_item(self) -> ReplayItem | None: ...
async def clear_session(self) -> None: ...
class _AgentsSDK(Protocol):
"""Structural view of the ``agents`` module.
Tests swap in a ``SimpleNamespace`` with the same attribute names, so
we can't rely on ``ModuleType`` with real-type attributes. Attributes
are typed as ``type`` factories that return the actual SDK objects
the caller touches — kept deliberately loose (``Any`` factories in a
few places) so the duck-typed test doubles also satisfy the
Protocol structurally.
"""
# Factories invoked with **kwargs; the results are used via Protocols
# (_RunResult, _RunState, _StreamEvent) or via attribute access that
# works on both the real dataclasses and the test fakes.
Agent: Any # type: ignore[explicit-any] # agents.Agent[Any] factory
Runner: Any # type: ignore[explicit-any] # agents.Runner class with run_streamed classmethod
RunConfig: Any # type: ignore[explicit-any] # agents.RunConfig dataclass factory
ModelSettings: Any # type: ignore[explicit-any] # agents.ModelSettings dataclass factory
OpenAIProvider: Any # type: ignore[explicit-any] # agents.OpenAIProvider class
FunctionTool: Any # type: ignore[explicit-any] # agents.FunctionTool dataclass factory
SQLiteSession: Any # type: ignore[explicit-any] # agents.SQLiteSession class
ItemHelpers: Any # type: ignore[explicit-any] # agents.ItemHelpers utility class
MaxTurnsExceeded: type[BaseException]
class _RunResult(Protocol):
"""Structural view of ``agents.result.RunResultStreaming``.
Mirrors the subset this executor actually consumes.
"""
final_output: Any # type: ignore[explicit-any] # SDK declares this as Any; may be str or arbitrary pydantic output
new_items: list[Any] # type: ignore[explicit-any] # list[RunItem] union the SDK narrows internally
raw_responses: list[Any] # type: ignore[explicit-any] # list[ModelResponse]; each has .usage.input_tokens
def cancel(self, mode: str = ...) -> None: ...
def stream_events(self) -> AsyncIterator[_StreamEvent]: ...
def to_state(self) -> _RunState: ...
class _RunState(Protocol):
"""Structural view of ``agents.run_state.RunState`` — only the three
private counters this executor tweaks when resuming a stepwise turn.
"""
_current_turn: int
_max_turns: int
_current_turn_persisted_item_count: int
class _StreamEvent(Protocol):
"""Minimal shape shared by the SDK's stream events: each one has a
``type`` discriminator. We narrow on it before reading ``data`` /
``item`` via ``cast``.
"""
type: str
class _RawResponseEvent(Protocol):
type: str
data: _RawResponseData
class _RawResponseData(Protocol):
type: str
delta: str
class _RunItemEvent(Protocol):
type: str
item: _RunItem
class _RunItem(Protocol):
type: str
raw_item: Any # type: ignore[explicit-any] # ToolCallItemTypes | ToolCallOutputTypes union, narrowed on .type
output: Any # type: ignore[explicit-any] # ToolCallOutputItem.output is declared Any by the SDK
class _ToolCallRawItem(Protocol):
"""Duck-typed view of non-dict tool-call raw items (pydantic models
like ``ResponseFunctionToolCall``). Attribute access returns the
corresponding field — ``None`` on items where it isn't set.
"""
name: str | None
arguments: str | None
call_id: str | None
id: str | None
class _ToolCallOutputRawItem(Protocol):
call_id: str | None
id: str | None
class _CallModelData(Protocol):
model_data: _ModelInputData
class _ModelInputData(Protocol):
input: list[ReplayItem]
def _patch_openai_sse_keepalive_tolerance() -> None:
"""Tolerate empty SSE keepalive frames from OpenAI-compatible proxies.
Some OpenAI-compatible proxies and relays (for example cli-proxy-api) keep
long streaming responses alive by periodically emitting an *empty* ``data:``
frame rather than an SSE comment. The stock openai client then calls
``json.loads("")`` on that frame and raises ``Expecting value: line 1
column 1 (char 0)``, which aborts the whole turn with
``Error streaming response: ...``.
SSE *comment* keepalives (``: ping``) are already ignored by the decoder;
this shim extends the same tolerance to empty-``data`` keepalives by
dropping any decoded event whose payload is blank, so a single keepalive no
longer kills a long-running streamed response. Best-effort and idempotent:
if the openai internals change shape it silently no-ops.
"""
try:
from openai import _streaming as _oai_streaming
except Exception: # noqa: BLE001 - openai internals are best-effort.
return
decoder = getattr(_oai_streaming, "SSEDecoder", None)
if decoder is None or getattr(decoder, "_omnigent_keepalive_patch", False):
return
_orig_decode = decoder.decode
def decode(self: Any, line: str) -> Any: # type: ignore[explicit-any]
sse = _orig_decode(self, line)
if sse is not None and not (getattr(sse, "data", None) or "").strip():
# Empty-data keepalive frame: drop it so the client never calls
# json.loads("") on a blank payload mid-stream.
return None
return sse
decoder.decode = decode # type: ignore[method-assign]
decoder._omnigent_keepalive_patch = True # type: ignore[attr-defined]
def _ensure_agents_sdk() -> ModuleType:
try:
import agents
_patch_openai_sse_keepalive_tolerance()
return agents
except ImportError as exc:
raise ImportError(
"OpenAIAgentsSDKExecutor requires the 'openai-agents' package. "
"Install it with: pip install openai-agents"
) from exc
def _get_openai_async_client(
profile: str | None = None,
api_key: str | None = None,
retry_policy: RetryPolicy | None = None,
base_url_override: str | None = None,
host_override: str | None = None,
databricks_auth_command: str | None = None,
model: str | None = None,
) -> AsyncOpenAIClient:
"""Construct an AsyncOpenAI client for direct or Databricks-hosted use.
For Databricks-backed profiles, the returned client uses an httpx
``Auth`` callback that calls ``Config.authenticate()`` on every
HTTP request, so OAuth tokens are refreshed transparently and
long-running sessions survive the 1-hour access-token lifetime.
:param profile: Optional ``~/.databrickscfg`` profile name for the
Databricks path, e.g. ``"dev"``.
:param api_key: Direct OpenAI-compatible API key, e.g.
``"sk-proj-..."``. When set, constructs
``AsyncOpenAI(api_key=…)`` directly, bypassing all profile and
env-var lookups. Intended for specs that declare
``executor.auth: {type: api_key, api_key: …}``.
:param retry_policy: Optional retry policy. When provided, its
``policy.openai.kwargs()`` (max_retries, timeout) are spread
into the ``AsyncOpenAI(...)`` constructor so L0 retry uses
the spec's budget. ``None`` falls back to a default policy
with the project-wide budget.
:param base_url_override: When set, use this as the client base URL.
In gateway host mode this is required and is populated from
``HARNESS_OPENAI_AGENTS_GATEWAY_BASE_URL``.
:param host_override: Databricks workspace host, e.g.
``"https://example.databricks.com"``. When set, skips profile host
lookup and requires ucode-provided gateway URL and auth command
values.
:param databricks_auth_command: Shell command from ucode state that
prints a bearer token, e.g.
``"databricks auth token --host https://example.databricks.com ..."``.
:param model: The model name that will be used. When set to a
non-Databricks model (i.e. does not start with ``"databricks-"``),
both the explicit ``profile`` block and the ambient Databricks
credential fallback are skipped — the caller must supply
``OPENAI_API_KEY`` (and optionally ``OPENAI_BASE_URL``) instead.
:raises DatabricksAuthError: When an explicit ``profile`` is given,
authentication fails, and no ``OPENAI_BASE_URL`` / ``OPENAI_API_KEY``
env-var fallbacks are available.
:raises OSError: If ucode host state is present but missing the
corresponding base URL or auth command.
"""
try:
from openai import AsyncOpenAI
except ImportError as exc:
raise ImportError(
"The 'openai' package is required for OpenAIAgentsSDKExecutor. "
"Install it with: pip install openai"
) from exc
policy = retry_policy if retry_policy is not None else RetryPolicy()
retry_kwargs = policy.openai.kwargs()
if host_override:
if base_url_override is None:
raise OSError(
"OpenAIAgentsSDKExecutor with a gateway workspace host requires "
"HARNESS_OPENAI_AGENTS_GATEWAY_BASE_URL."
)
if databricks_auth_command is None:
raise OSError(
"OpenAIAgentsSDKExecutor with a gateway workspace host requires "
"HARNESS_OPENAI_AGENTS_GATEWAY_AUTH_COMMAND."
)
host = host_override.rstrip("/")
return AsyncOpenAI(
base_url=base_url_override,
api_key=_OPENAI_KEY_PLACEHOLDER,
http_client=httpx.AsyncClient(auth=_ShellCommandBearerAuth(databricks_auth_command)),
**retry_kwargs,
)
# Explicit spec-level api_key (executor.auth: {type: api_key, …}).
# Checked before profile and env-var lookups so the spec is self-contained.
# base_url_override is populated from HARNESS_OPENAI_AGENTS_GATEWAY_BASE_URL
# when the spec also declares executor.auth.base_url.
#
# Fall back to the ambient OPENAI_BASE_URL when no override reached us.
# The api_key is frequently a gateway credential (e.g. a Databricks AI
# Gateway PAT detected from OPENAI_API_KEY), and the companion base_url
# can be dropped anywhere on the daemon → runner → harness propagation
# chain (the spec auth bake omits it when OPENAI_BASE_URL is absent at
# materialization time; a reused local daemon may predate the env var).
# Without this fallback, a missing base_url silently routes the gateway
# token to api.openai.com and every request 401s; honoring the ambient
# OPENAI_BASE_URL the runner inherits keeps the gateway target present on
# every turn even when the override is lost. base_url=None still defaults
# to api.openai.com for a genuine OpenAI key with no gateway configured.
if api_key and api_key.strip():
return AsyncOpenAI(
api_key=api_key,
base_url=base_url_override or os.environ.get("OPENAI_BASE_URL") or None,
**retry_kwargs,
)
is_databricks_model = model is None or model.startswith("databricks-")
# Databricks profile auth only applies when the model is Databricks-hosted.
# An explicit spec/provider profile wins over ambient env vars.
if is_databricks_model and profile:
from .databricks_executor import DatabricksAuthError, _resolve_databricks_auth
try:
auth, host = _resolve_databricks_auth(profile)
return AsyncOpenAI(
base_url=base_url_override or _databricks_openai_base_url(host),
api_key=_OPENAI_KEY_PLACEHOLDER,
http_client=httpx.AsyncClient(auth=auth),
**retry_kwargs,
)
except DatabricksAuthError as exc:
# Fall through to env-var credentials when available so that
# CI environments (OIDC tokens injected via OPENAI_BASE_URL /
# OPENAI_API_KEY) still work even if the named profile is not
# present in ~/.databrickscfg. Log a warning so the fallback
# is visible rather than silent. Re-raise only when there are
# no env-var fallbacks to fall through to.
if os.environ.get("OPENAI_BASE_URL") or os.environ.get("OPENAI_API_KEY"):
logger.warning(
"Databricks profile %r authentication failed (%s); "
"falling back to OPENAI_BASE_URL/OPENAI_API_KEY.",
profile,
exc,
)
else:
raise
except ImportError:
logger.warning(
"databricks-sdk is not installed; cannot resolve Databricks "
"profile %r. Falling back to OPENAI_BASE_URL/OPENAI_API_KEY.",
profile,
)
if os.environ.get("OPENAI_BASE_URL"):
return AsyncOpenAI(
base_url=os.environ["OPENAI_BASE_URL"],
api_key=os.environ.get("OPENAI_API_KEY", _OPENAI_KEY_PLACEHOLDER),
**retry_kwargs,
)
api_key = os.environ.get("OPENAI_API_KEY")
if api_key:
return AsyncOpenAI(api_key=api_key, **retry_kwargs)
# Non-Databricks model with no OpenAI credentials — fail loudly rather
# than silently routing to the Databricks AI Gateway, which will 404.
if not is_databricks_model:
raise ValueError(
f"Model {model!r} is not a Databricks-hosted model but no OpenAI "
"credentials were found. Set OPENAI_API_KEY (and optionally "
"OPENAI_BASE_URL) to use this model, or use a 'databricks-' "
"prefixed model name to route through Databricks."
)
# No profile, no env — final fallback via ambient Databricks credentials.
try:
from .databricks_executor import _resolve_databricks_auth
auth, host = _resolve_databricks_auth(profile)
except ImportError as exc:
raise ImportError(
"The 'databricks-sdk' package is required for Databricks "
"authentication but is not installed, and no OPENAI_API_KEY or "
"OPENAI_BASE_URL environment variables are set. Either install "
"the package (`pip install 'omnigent[databricks]'`) or set "
"OPENAI_API_KEY/OPENAI_BASE_URL for non-Databricks OpenAI access."
) from exc
return AsyncOpenAI(
base_url=base_url_override or _databricks_openai_base_url(host),
api_key=_OPENAI_KEY_PLACEHOLDER,
http_client=httpx.AsyncClient(auth=auth),
**retry_kwargs,
)
class _ShellCommandBearerAuth(httpx.Auth):
"""httpx Auth that refreshes a bearer token through a shell command.
:param command: Shell command that prints the access token to stdout,
e.g. ``"databricks auth token --host https://example.databricks.com ..."``.
"""
def __init__(self, command: str) -> None:
"""Store the command for per-request token refresh.
:param command: Shell command that prints the access token.
"""
self._command = command
def auth_flow(
self,
request: httpx.Request,
) -> Generator[httpx.Request, httpx.Response, None]:
"""Inject a fresh ``Authorization: Bearer`` header.
:param request: The outgoing httpx request.
:yields: The request with the auth header set.
:raises RuntimeError: When the command fails or prints no token.
"""
result = subprocess.run(
["sh", "-c", self._command],
check=False,
capture_output=True,
text=True,
)
token = result.stdout.strip()
if result.returncode != 0 or not token:
raise RuntimeError("Databricks auth command failed to return a bearer token.")
request.headers["Authorization"] = f"Bearer {token}"
yield request
# Sentinel pushed onto the per-turn event queue when
# :meth:`OpenAIAgentsSDKExecutor._consume_stream_into_queue`
# exits (normally OR via cancellation). The consumer loop in
# :meth:`OpenAIAgentsSDKExecutor.run_turn` breaks on it and
# proceeds to await the producer task, which surfaces any
# non-cancel exception. A unique module-level object avoids
# false positives from real ``StreamEvent`` instances.
_STREAM_DONE = object()
@dataclass
class _AgentsSessionState:
"""
Per-session state for the OpenAI Agents SDK executor.
:param sdk_session: The SDK's persistent session (SQLite-
backed). Holds the conversational history items.
:param started: ``True`` once the SDK has accepted the
first turn for this session.
:param agent: Cached ``agents.Agent[Any]`` instance for
reuse across turns; rebuilt when *agent_signature*
changes.
:param agent_signature: ``(model, system_prompt, tools,
parallel_tool_calls)``-derived key. A signature change
forces an agent rebuild.
:param resume_state: Optional SDK ``_RunState`` carried
across stepwise-internal turns.
:param active_result: The in-flight ``RunResultStreaming``
for the current turn, or ``None`` between turns.
:param interrupt_requested: Set by :meth:`interrupt_session`
to signal the in-flight turn to stop emitting events.
Cleared at the top of the next ``run_turn``.
:param stream_consumer_task: The asyncio.Task running
``_consume_stream_into_queue`` for the in-flight turn.
:meth:`interrupt_session` cancels this task, which lands
``CancelledError`` inside the SDK's
``await self._event_queue.get()`` (agents/result.py:721)
— the SDK's own ``except CancelledError`` then calls
``self.cancel()`` to close the underlying httpx stream.
Without this task indirection, the only cancel signal
is ``result.cancel(mode="immediate")``, which sets a
flag but does NOT close the network connection — events
keep streaming for 15+ seconds. ``None`` between turns.
:param history_cursor: Index into the persisted-message
suffix marking how much we've already replayed.
:param run_item_count_before: ``len(sdk_session.items)``
snapshot taken before the in-flight turn so a cancelled
turn can be rolled back to this point.
:param rollback_to_item_count: When non-``None``, the next
``run_turn`` rewinds the SDK session to this length
before starting. Set by :meth:`interrupt_session`.
"""
sdk_session: Any # type: ignore[explicit-any] # _SanitizingSession or OpenAIResponsesCompactionSession
started: bool = False
# agents.Agent[Any] instance; cached for reuse across turns.
agent: SDKAgent = None
agent_signature: tuple[str, str, str, str, str] | None = None
resume_state: _RunState | None = None
active_result: _RunResult | None = None
interrupt_requested: bool = False
stream_consumer_task: asyncio.Task[None] | None = None
history_cursor: int = 0
run_item_count_before: int | None = None
rollback_to_item_count: int | None = None
# ``_sanitize_replay_item`` walks replay values recursively. At the
# boundary these are either SDK TypedDicts, pydantic models that have
# already been ``model_dump``-ed to dicts, or primitive JSON values — so
# ``Any`` is the right annotation. The recursion preserves shape, so we
# return the same static type we received.
_ReplayValue: TypeAlias = Any # type: ignore[explicit-any]
def _sanitize_replay_item(value: _ReplayValue) -> _ReplayValue:
"""Strip provider-only fields Databricks rejects during replay."""
if isinstance(value, list):
return [_sanitize_replay_item(item) for item in value]
if not isinstance(value, dict):
return value
return {
key: _sanitize_replay_item(item)
for key, item in value.items()
if not (key == "id" and isinstance(item, str) and len(item) > 64)
}
class _SanitizingSession:
"""Session wrapper that normalizes replay items before persistence."""
def __init__(self, underlying: _SDKSession) -> None:
self._underlying = underlying
self.session_id = underlying.session_id
# ``session_settings`` is declared on the SDK Session protocol but
# some fake sessions (tests) / older SDK revs omit it. Fall back
# to ``None`` rather than propagating the AttributeError.
self.session_settings = (
underlying.session_settings if hasattr(underlying, "session_settings") else None
)
async def get_items(self, limit: int | None = None) -> list[ReplayItem]:
return await self._underlying.get_items(limit)
async def add_items(self, items: list[ReplayItem]) -> None:
await self._underlying.add_items([_sanitize_replay_item(item) for item in items])
async def pop_item(self) -> ReplayItem | None:
item = await self._underlying.pop_item()
if item is None:
return None
return cast(ReplayItem, _sanitize_replay_item(item))
async def clear_session(self) -> None:
await self._underlying.clear_session()
@dataclass(frozen=True)
class RawToolItemParts:
"""Normalized view of a tool-call raw item from the Agents SDK.
The SDK hands tool calls back in one of two shapes — a dict (e.g. the
Responses-API input-item TypedDict) or a pydantic model such as
``ResponseFunctionToolCall`` — and this dataclass is the single
extracted form we pass downstream as a ``ToolCallRequest``.
:param name: The tool name. Empty string when the raw item didn't
carry a name.
:param args: Parsed JSON arguments. Falls back to
``{"raw": <args_text>}`` when the argument string isn't valid JSON.
:param call_id: The SDK-assigned call id (``call_id`` or ``id``), or
``None`` when the raw item didn't carry one.
"""
name: str
args: ToolArgs
call_id: str | None
def _tool_args_from_raw_item(raw_item: RawToolItem) -> RawToolItemParts:
"""Extract ``name`` / ``args`` / ``call_id`` from an SDK tool raw item.
:param raw_item: Either a dict (Responses-API TypedDict input item) or
a pydantic ``_ToolCallRawItem`` produced by the Agents SDK.
"""
name: str | None = None
args_text: str | None = None
call_id: str | None = None
if isinstance(raw_item, dict):
raw_dict: ToolArgs = raw_item
name_val = raw_dict.get("name")
name = name_val if isinstance(name_val, str) else None
args_val = raw_dict.get("arguments")
args_text = args_val if isinstance(args_val, str) else None
cid = raw_dict.get("call_id") or raw_dict.get("id")
call_id = cid if isinstance(cid, str) else None
else:
# Pydantic BaseModel-style raw item (e.g. ResponseFunctionToolCall).
# Only ``ResponseFunctionToolCall`` in the ToolCallItemTypes union
# carries all three fields; other members don't emit ToolCallRequest
# events in practice (this executor only wires function tools).
proto = cast(_ToolCallRawItem, raw_item)
name = proto.name
args_text = proto.arguments
call_id = proto.call_id or proto.id
args: ToolArgs
try:
args = cast(ToolArgs, json.loads(args_text)) if args_text else {}
except (TypeError, json.JSONDecodeError):
args = {"raw": args_text}
# ``RawToolItemParts.name`` is typed ``str`` because every
# downstream consumer emits it as the tool-call name; widening to
# ``str | None`` here would cascade into ``ToolCallRequest`` and
# force ``None``-handling on every callsite for a case that
# indicates a malformed SDK payload. Keep the coercion localized
# with an explicit ternary at this dataclass boundary.
resolved_name: str = name if name is not None else ""
return RawToolItemParts(name=resolved_name, args=args, call_id=call_id)
def _build_reasoning_model_settings(effort: str | None) -> dict[str, object]:
"""Build ModelSettings kwargs for reasoning effort."""
if effort is None:
return {}
try:
from openai.types.shared import Reasoning
except ImportError:
return {"reasoning": {"effort": effort}}
openai_effort = cast(
Literal["none", "minimal", "low", "medium", "high", "xhigh"],
effort,
)
return {"reasoning": Reasoning(effort=openai_effort)}
def _is_databricks_openai_client(client: AsyncOpenAIClient) -> bool:
"""Return whether *client* targets a Databricks AI Gateway base URL."""
return "/ai-gateway/" in str(getattr(client, "base_url", ""))
class _ReasoningBlockFilterStream:
"""Async stream wrapper that converts list-type ``delta.content`` to ``None``.
Kimi K2 (and similar reasoning models served via Databricks) sends
thinking/reasoning blocks as list-type content in streaming deltas:
``choices[0].delta.content = [{'type': 'reasoning', ...}]``.
The openai-agents SDK's ``ChatCmplStreamHandler`` creates a
``ResponseTextDeltaEvent(delta=content)`` which Pydantic rejects for
list input — ``delta`` must be ``str``.
This wrapper replaces list-type content with ``None`` (the SDK treats
``None`` as "no text delta") so the reasoning tokens are silently
discarded at the stream boundary and only the final answer text
reaches the SDK.
:param stream: The underlying ``AsyncStream[ChatCompletionChunk]``
from ``AsyncOpenAI.chat.completions.create(stream=True)``.
"""
def __init__(self, stream: Any) -> None: # type: ignore[explicit-any]
"""
:param stream: Underlying async stream of chat completion chunks.
"""
self._stream = stream
def __aiter__(self) -> _ReasoningBlockFilterStream:
return self
async def __anext__(self) -> Any: # type: ignore[explicit-any]
"""
:returns: The next chunk, with list-type ``delta.content`` replaced
by ``None``.
:raises StopAsyncIteration: When the underlying stream is exhausted.
"""
chunk = await self._stream.__anext__()
for choice in chunk.choices:
if isinstance(getattr(choice.delta, "content", None), list):
# Bypass Pydantic's __setattr__ validation — the field type
# in ChoiceDelta is ``str | None`` and the list value from
# reasoning models fails validation if assigned normally.
object.__setattr__(choice.delta, "content", None)
return chunk
async def __aenter__(self) -> _ReasoningBlockFilterStream:
if hasattr(self._stream, "__aenter__"):
await self._stream.__aenter__()
return self
async def __aexit__(self, *args: object) -> None:
if hasattr(self._stream, "__aexit__"):
await self._stream.__aexit__(*args)
def __getattr__(self, name: str) -> Any: # type: ignore[explicit-any]
return getattr(self._stream, name)
class _ReasoningBlockFilterCompletions:
"""Wraps ``AsyncCompletions`` to inject :class:`_ReasoningBlockFilterStream`.
Sits between the openai-agents SDK's ``OpenAIProvider`` and the real
``AsyncOpenAI.chat.completions`` so reasoning-model list content never
reaches ``ChatCmplStreamHandler``.
:param completions: The real ``AsyncCompletions`` object.
"""
def __init__(self, completions: Any) -> None: # type: ignore[explicit-any]
self._completions = completions
async def create(self, **kwargs: Any) -> Any: # type: ignore[explicit-any]
"""
Proxy ``create()``; wrap the result in
:class:`_ReasoningBlockFilterStream` when streaming is enabled.
:param kwargs: Forwarded verbatim to the underlying ``create()``.
:returns: A :class:`_ReasoningBlockFilterStream` when ``stream=True``,
otherwise the raw ``ChatCompletion`` response.
"""
result = await self._completions.create(**kwargs)
if kwargs.get("stream") and hasattr(result, "__anext__"):
return _ReasoningBlockFilterStream(result)
return result
def __getattr__(self, name: str) -> Any: # type: ignore[explicit-any]
return getattr(self._completions, name)
class _ReasoningBlockFilterChat:
"""Wraps ``AsyncChat`` to expose a :class:`_ReasoningBlockFilterCompletions`.
:param chat: The real ``AsyncOpenAI.chat`` object.
"""
def __init__(self, chat: Any) -> None: # type: ignore[explicit-any]
self._chat = chat
self.completions: _ReasoningBlockFilterCompletions = _ReasoningBlockFilterCompletions(
chat.completions
)
def __getattr__(self, name: str) -> Any: # type: ignore[explicit-any]
return getattr(self._chat, name)
def _wrap_client_for_reasoning_models(client: AsyncOpenAIClient) -> AsyncOpenAIClient:
"""Wrap *client* so reasoning-model list content is filtered from streams.
Replaces ``client.chat`` with a :class:`_ReasoningBlockFilterChat`
proxy that intercepts ``chat.completions.create(stream=True)`` and
wraps the result in :class:`_ReasoningBlockFilterStream`. All other
attribute accesses fall through to the real client.
:param client: The ``AsyncOpenAI`` (or compatible) client to wrap.
:returns: The same *client* object with ``chat`` replaced by the filter
proxy. The object is modified in-place and returned for chaining.
"""
# Patch ``chat`` directly on the client instance so the proxy intercepts
# every ``client.chat.completions.create()`` call the SDK makes.
# ``object.__setattr__`` bypasses both the OpenAI SDK's own ``__setattr__``
# and any Pydantic frozen-instance guard on the client class.
object.__setattr__(client, "chat", _ReasoningBlockFilterChat(client.chat))
return client
def _count_output_items(new_items: list[Any]) -> int: # type: ignore[explicit-any]
"""Count run items that represent user-visible output.
Excludes bookkeeping items (reasoning, compaction) per
:data:`_NON_OUTPUT_ITEM_TYPES` so a reasoning-only ghost turn
counts as zero output items.
:param new_items: ``RunResult.new_items`` for the completed run.
:returns: Number of items whose ``.type`` is not in
:data:`_NON_OUTPUT_ITEM_TYPES`.
"""
return sum(
1 for item in new_items if getattr(item, "type", None) not in _NON_OUTPUT_ITEM_TYPES
)
def _sum_output_tokens(raw_responses: list[Any] | None) -> int: # type: ignore[explicit-any]
"""Sum ``output_tokens`` across a run's raw model responses.
:param raw_responses: ``RunResult.raw_responses``; each element has a
``.usage.output_tokens``. ``None`` or empty yields ``0``.
:returns: Total output tokens reported across all sub-turn responses.
"""
if not raw_responses:
return 0
return sum(getattr(r.usage, "output_tokens", 0) or 0 for r in raw_responses)
def _is_empty_turn(
final_text: str,
saw_tool_activity: bool,
new_items: list[Any], # type: ignore[explicit-any]
) -> bool:
"""Whether a completed run produced literally nothing worth surfacing.
A turn is empty when it has no final text, no tool activity, and no
output-bearing items. A turn that called tools but produced no text
is NOT empty (tool activity is legitimate output).
Token usage is deliberately NOT part of this predicate: it must not
suppress a retry (gateway token accounting can be unreliable, and a
token quirk should not stop us from re-running an empty turn). The
fail-loud decision after the retry loop applies a *separate*,
narrower gate (``output_tokens == 0``) on top of this one — see the
call site for why an empty turn that still billed tokens is left to
complete silently rather than raised as an error.
:param final_text: The assembled assistant text for the turn.
:param saw_tool_activity: ``True`` if any tool call was streamed.
:param new_items: ``RunResult.new_items`` for the completed run.
:returns: ``True`` if the turn is empty and should be retried.
"""
return not final_text and not saw_tool_activity and _count_output_items(new_items) == 0
class OpenAIAgentsSDKExecutor(Executor):
"""Execute turns using the OpenAI Agents SDK."""
def __init__(
self,
*,
client: AsyncOpenAIClient = None,
profile: str | None = None,
api_key: str | None = None,
use_responses: bool = True,
model: str | None = None,
retry_policy: RetryPolicy | None = None,
base_url_override: str | None = None,
gateway_host: str | None = None,
gateway_auth_command: str | None = None,
) -> None:
"""Create an OpenAIAgentsSDKExecutor.
:param client: A preconfigured ``openai.AsyncOpenAI`` client. When
``None`` the executor calls :func:`_get_openai_async_client`.
:param profile: Optional ``~/.databrickscfg`` profile name for the
Databricks fallback path, e.g. ``"<your-profile>"``.
:param api_key: Direct OpenAI-compatible API key, e.g.
``"sk-proj-..."``. When set, constructs the client with this key
directly, bypassing profile resolution and env-var lookups. Set
from ``HARNESS_OPENAI_AGENTS_API_KEY`` when the agent spec
declares ``executor.auth: {type: api_key, api_key: …}``.
:param use_responses: When ``True`` the executor talks to the
OpenAI ``/responses`` endpoint; when ``False`` it falls back to
``/chat/completions``.
:param model: Optional constructor-level model default. Acts
as the spec-level default — applied when a per-turn
``ExecutorConfig.model`` is not set. A non-None
``cfg.model`` (set by AP's
the harness HTTP client when the request carries
a ``model_override``) wins over this value, so the REPL's
``/model`` slash command can shadow the spec model
per-request. The harness wrap reads
``HARNESS_OPENAI_AGENTS_MODEL`` from the spawn env and
threads the value here so the harness can pin a model
independently of the per-call ``request.model`` (which AP
uses to identify the AGENT, not the underlying LLM).
``None`` falls back to ``cfg.model`` then the executor's
built-in default.
:param base_url_override: Override the OpenAI-compatible base URL
instead of deriving it from the Databricks profile host. Set
from ``HARNESS_OPENAI_AGENTS_GATEWAY_BASE_URL`` (written by
the Omnigent workflow layer). Required whenever ``gateway_host`` is set.
:param gateway_host: Gateway workspace host origin, e.g.
``"https://example.databricks.com"``. Set from
``HARNESS_OPENAI_AGENTS_GATEWAY_HOST`` (written by the AP
workflow layer). When set, skips profile host lookup and requires
the gateway base URL and auth command values.
:param gateway_auth_command: Shell command that prints a bearer token,
e.g.
``"databricks auth token --host https://example.databricks.com ..."``
or ``"printf %s sk-..."``. Set from
``HARNESS_OPENAI_AGENTS_GATEWAY_AUTH_COMMAND``.
"""
self._retry_policy = retry_policy if retry_policy is not None else RetryPolicy()
raw_client = (
client
if client is not None
else _get_openai_async_client(
profile=profile,
api_key=api_key,
retry_policy=self._retry_policy,
base_url_override=base_url_override,
host_override=gateway_host,
databricks_auth_command=gateway_auth_command,
model=model,
)
)
# Wrap the chat.completions path to strip list-type delta.content
# (reasoning blocks emitted by models like Kimi K2). The SDK's
# ChatCmplStreamHandler validates delta as str; list input raises
# ValidationError and the turn silently produces no output.
# The wrapper is a no-op for models that always return str content.
# Only needed for the chat-completions path; the Responses API path
# has its own event handling that doesn't go through ChatCmplStreamHandler.
self._client = (
_wrap_client_for_reasoning_models(raw_client) if not use_responses else raw_client
)
self._profile = profile
self._use_responses = use_responses
self._model_override = model
self._databricks = _is_databricks_openai_client(self._client)
self._tool_executor: ToolExecutor | None = None
self._session_states: dict[str, _AgentsSessionState] = {}
def supports_streaming(self) -> bool:
return True
def supports_tool_calling(self) -> bool:
return True
def handles_tools_internally(self) -> bool:
return True
def supports_stepwise_internal_turns(self) -> bool:
return True
def max_context_tokens(self) -> int | None:
return None
def _session_key(self, messages: list[Message]) -> str:
if messages:
if messages[-1].get("session_id"):
return str(messages[-1]["session_id"])
metadata = messages[-1].get("metadata", {})
if isinstance(metadata, dict) and metadata.get("session_id"):
return str(metadata["session_id"])
return "default"
def _get_or_create_session_state(
self, agents_sdk: _AgentsSDK, session_key: str
) -> _AgentsSessionState:
state = self._session_states.get(session_key)
if state is not None:
return state
underlying = _SanitizingSession(agents_sdk.SQLiteSession(session_key))
sdk_session: Any = underlying # type: ignore[explicit-any]
# Wrap with compaction session when the client targets a real
# HTTP endpoint. Skip for bare object() clients in unit tests
# that lack base_url.
_base = str(getattr(self._client, "base_url", ""))
if _base.startswith("http"):
try:
from agents.memory import OpenAIResponsesCompactionSession
class _SafeCompactionSession(OpenAIResponsesCompactionSession):
"""Compaction session that treats compaction failures as non-fatal.
The responses.compact API may not be available on all
endpoints (mock servers, some proxies). A 404 or other
failure should not kill the turn — the session continues
without compaction.
"""
async def run_compaction(self, args=None): # type: ignore[override]
try:
await super().run_compaction(args)
except Exception: # noqa: BLE001
logger.debug(
"Compaction call failed (endpoint may not support "
"responses.compact), continuing without compaction",
exc_info=True,
)
sdk_session = _SafeCompactionSession(
session_id=session_key,
underlying_session=underlying, # type: ignore[arg-type]
client=self._client,
)
except (ImportError, AttributeError, ValueError) as exc:
logger.debug(
"Compaction session setup failed, falling back to plain session: %s",
exc,
)
state = _AgentsSessionState(sdk_session=sdk_session)
self._session_states[session_key] = state
return state
async def close_session(self, session_key: str) -> None:
self._session_states.pop(session_key, None)
async def interrupt_session(self, session_key: str) -> bool:
"""
Halt the in-flight turn for *session_key* and roll the
SDK session back to its pre-turn state.
Cancellation strategy: cancel the per-turn
``stream_consumer_task`` task (set up in :meth:`run_turn`).
That task is parked on
``await self._event_queue.get()`` deep inside the SDK's
``stream_events`` (agents/result.py:721). Cancelling our
task lands ``CancelledError`` exactly there — and the
SDK's own ``except CancelledError`` block (result.py:722-
725) catches it, calls ``self.cancel()`` on the
``RunResultStreaming``, and re-raises. ``self.cancel()``
in turn cancels the SDK's run-loop task, which closes the
underlying httpx response stream when its async-with
unwinds — finally stopping the network bytes that were
producing 1,500+ post-cancel deltas in the original bug.
Calling ``active.cancel(mode="immediate")`` on the result
directly does *not* close the network stream — it only
sets a flag (``_cancel_mode``) and pushes a sentinel; if
the consumer is currently awaiting ``queue.get()``,
``_drain_event_queue`` is skipped and pre-buffered events
keep flowing. The task-cancel path above is the SDK's
documented contract for "really" stopping a stream.
:param session_key: Key for the per-session state, e.g.
``"conv_abc123"``. ``"default"`` for un-keyed turns.
:returns: ``True`` if a turn was in flight and the
cancel was issued; ``False`` if no active result
exists for this session.
"""
state = self._session_states.get(session_key)
if state is None or state.active_result is None:
return False
state.interrupt_requested = True
# Cancel the consumer task so CancelledError lands inside
# the SDK's stream_events()'s queue.get() — see the class
# docstring above for why this is the *only* path that
# actually halts the underlying httpx stream.
consumer = state.stream_consumer_task
if consumer is not None and not consumer.done():
consumer.cancel()
state.resume_state = None
state.rollback_to_item_count = state.run_item_count_before
return True
async def _rewind_sdk_session(
self,
state: _AgentsSessionState,
target_item_count: int,
) -> None:
items = await state.sdk_session.get_items()
while len(items) > target_item_count:
await state.sdk_session.pop_item()
items = await state.sdk_session.get_items()
async def _prepare_sdk_session_for_turn(
self,
state: _AgentsSessionState,
messages: list[Message],
) -> None:
if state.rollback_to_item_count is not None:
await self._rewind_sdk_session(state, state.rollback_to_item_count)
state.rollback_to_item_count = None
if state.started and len(split_transient_tail(messages).persisted) < state.history_cursor:
await self._rewind_sdk_session(state, 0)
state.started = False
state.resume_state = None
state.history_cursor = 0
def _build_input_for_turn(
self,
state: _AgentsSessionState,
messages: list[Message],
*,
allow_resume_without_new_input: bool,
) -> str | list[Message] | _RunState:
if state.resume_state is not None and allow_resume_without_new_input:
resume_state = state.resume_state
state.resume_state = None
resume_state._current_turn = 0
resume_state._max_turns = 1
resume_state._current_turn_persisted_item_count = 0
return resume_state
if state.resume_state is not None:
state.resume_state = None
if not state.started:
return _normalize_responses_items_for_chat(_convert_messages_to_responses(messages))
split = split_transient_tail(messages)
delta_persisted = split.persisted[state.history_cursor :]
delta_messages = list(delta_persisted) + list(split.transient)
if not delta_messages:
return ""
if len(delta_messages) == 1 and delta_messages[0].get("role") == "user":
content = delta_messages[0].get("content")
if content is None:
return ""
if isinstance(content, str):
return content
if isinstance(content, list):
# Normalize input_file blocks before handing to the SDK.
# The openai-agents SDK's chatcmpl_converter passes
# ``file_data`` directly to the Chat Completions ``file``
# content type, which expects plain base64 — not the data:
# URI that content_resolver produces. The Databricks GPT
# endpoint may not support ``file`` content blocks at all.
# Converting to ``input_text`` is the universally compatible
# path: the model sees the file content as plain text.
normalized = _normalize_content_blocks_for_chat(content)
return [{"type": "message", "role": "user", "content": normalized}]
return json.dumps(content)
return _normalize_responses_items_for_chat(_convert_messages_to_responses(delta_messages))
def _build_tools(
self,
agents_sdk: _AgentsSDK,
tools: list[ToolSpec],
) -> list[SDKTool]:
sdk_tools: list[SDKTool] = []
for tool in tools:
raw_name = tool.get("name")
# The OpenAI Agents SDK's ``FunctionTool`` requires a
# non-empty ``name``; drop malformed specs rather than
# registering an unnamed tool the SDK will reject.
if not isinstance(raw_name, str) or not raw_name:
continue
tool_name: str = raw_name
raw_desc = tool.get("description")
tool_desc: str = raw_desc if isinstance(raw_desc, str) else ""
params = tool.get("parameters", {"type": "object", "properties": {}})
async def on_invoke_tool(
_ctx: ToolResult,
input_json: str,
*,
_tool_name: str = tool_name,
) -> ToolResult:
if self._tool_executor is None:
return {"error": f"No tool executor for '{_tool_name}'"}
try:
args = json.loads(input_json) if input_json else {}
except (TypeError, json.JSONDecodeError):
args = {}
result = self._tool_executor(_tool_name, args)
if hasattr(result, "__await__"):
result = await result # type: ignore[assignment]
return result
sdk_tools.append(
agents_sdk.FunctionTool(
name=tool_name,
description=tool_desc,
params_json_schema=params,
on_invoke_tool=on_invoke_tool,
strict_json_schema=False,
)
)
return sdk_tools
@staticmethod
def _tool_signature(tools: list[ToolSpec]) -> str:
normalized = [
{
# Signature is used only as a cache key; missing fields
# are normalised to ``None`` so tool specs that differ
# only by absent keys still produce distinct signatures
# from ones that explicitly set empty strings.
"name": tool.get("name"),
"description": tool.get("description"),
"parameters": tool.get("parameters", {"type": "object", "properties": {}}),
}
for tool in tools
]
return json.dumps(normalized, sort_keys=True, separators=(",", ":"))
def _get_or_create_agent(
self,
agents_sdk: _AgentsSDK,
state: _AgentsSessionState,
*,
model: str,
system_prompt: str,
tools: list[ToolSpec],
parallel_tool_calls: bool | None,
reasoning_effort: str | None,
max_tokens: int | None,
) -> SDKAgent:
signature = (
model,
system_prompt,
self._tool_signature(tools),
json.dumps(parallel_tool_calls),
reasoning_effort or "",
json.dumps(max_tokens),
)
if state.agent is not None and state.agent_signature == signature:
return state.agent
state.agent = agents_sdk.Agent(
name="Omnigent",
instructions=system_prompt or None,
model=model,
model_settings=agents_sdk.ModelSettings(
parallel_tool_calls=parallel_tool_calls,
max_tokens=max_tokens,
**_build_reasoning_model_settings(reasoning_effort),
),
tools=self._build_tools(agents_sdk, tools),
tool_use_behavior="run_llm_again",
)
state.agent_signature = signature
return state.agent
async def _consume_stream_into_queue(
self,
result: _RunResult,
queue: asyncio.Queue[object],
) -> None:
"""
Pump the SDK's ``stream_events`` into a local queue.
Runs as a separate :class:`asyncio.Task` so
:meth:`interrupt_session` can cancel us while
:meth:`run_turn` is blocked on the queue. Cancellation
lands inside ``await self._event_queue.get()`` deep in
the SDK's ``stream_events`` (agents/result.py:721) — and
the SDK's own ``except CancelledError`` (result.py:722-
725) catches it, calls ``self.cancel()`` on the
``RunResultStreaming`` (which closes the underlying
httpx response stream by cancelling the run-loop task),
and re-raises.
Without this task indirection, the only way to signal
the SDK to stop is the public
``result.cancel(mode="immediate")`` API, which sets a
flag and pushes a sentinel but does NOT close the
network connection — events keep streaming for 15+
seconds.
:param result: The SDK's ``RunResultStreaming``
instance for the in-flight turn.
:param queue: Per-turn unbounded :class:`asyncio.Queue`
owned by :meth:`run_turn`. Receives every
``StreamEvent`` the SDK emits, then a single
:data:`_STREAM_DONE` sentinel in this method's
``finally`` (so the consumer sees a terminal even
on cancellation / exception).
:raises asyncio.CancelledError: When :meth:`interrupt_session`
cancels this task. Caller (run_turn) awaits the
task after the sentinel and expects this on the
interrupt path.
"""
try:
async for event in result.stream_events():
queue.put_nowait(event)
finally:
# Always push the sentinel — even on CancelledError /
# exception — so the consumer's ``await queue.get()``
# unblocks and the run_turn finally can clean up.
queue.put_nowait(_STREAM_DONE)
async def run_turn(
self,
messages: list[Message],
tools: list[ToolSpec],
system_prompt: str,
config: ExecutorConfig | None = None,
) -> AsyncIterator[ExecutorEvent]:
cfg = config or ExecutorConfig()
# cfg.model (per-request /model override; agent name no longer
# leaks here) wins over the spec default
# (HARNESS_OPENAI_AGENTS_MODEL → self._model_override).
model = (
cfg.model
or self._model_override
or (
_DATABRICKS_OPENAI_AGENTS_DEFAULT_MODEL
if self._databricks
else _OPENAI_AGENTS_DEFAULT_MODEL
)
)
agents_sdk = cast(_AgentsSDK, _ensure_agents_sdk())
session_key = self._session_key(messages)
try:
reasoning_effort = validate_effort(
cfg.extra.get("reasoning_effort"), "OpenAI Agents SDK", OPENAI_AGENTS_EFFORTS
)
except ValueError as exc:
yield ExecutorError(message=str(exc), retryable=False)
return
state = self._get_or_create_session_state(agents_sdk, session_key)
state.interrupt_requested = False
await self._prepare_sdk_session_for_turn(state, messages)
stepwise_internal_turns = bool(cfg.extra.get("stepwise_internal_turns"))
input_value = self._build_input_for_turn(
state,
messages,
allow_resume_without_new_input=(
stepwise_internal_turns and not bool(cfg.extra.get("new_user_messages_flushed"))
),
)
parallel_tool_calls = cfg.extra.get("parallel_tool_calls", None)
max_tokens_raw = cfg.extra.get("max_tokens")
max_tokens = int(max_tokens_raw) if max_tokens_raw is not None else None
provider = agents_sdk.OpenAIProvider(
openai_client=self._client,
use_responses=self._use_responses,
)
agent = self._get_or_create_agent(
agents_sdk,
state,
model=model,
system_prompt=system_prompt,
tools=tools,
parallel_tool_calls=parallel_tool_calls,
reasoning_effort=reasoning_effort,
max_tokens=max_tokens,
)
current_item_count = len(await state.sdk_session.get_items())
run_config = agents_sdk.RunConfig(
model=model,
model_provider=provider,
tracing_disabled=True,
reasoning_item_id_policy="omit",
call_model_input_filter=self._filter_model_input,
)
max_turns = 1 if stepwise_internal_turns else int(cfg.extra.get("max_turns", 1000))
# ── LLM_REQUEST policy evaluation ────────────────────────
# If the executor adapter installed a ``_policy_evaluator``
# callback, call it with the request data so the Omnigent server
# can evaluate LLM_REQUEST policies before the LLM call.
# This mirrors the identical block in ``OpenResponsesExecutor``
# and ``ClaudeSDKExecutor``.
_policy_eval = getattr(self, "_policy_evaluator", None)
if _policy_eval is not None:
_last_user_msg = ""
for _msg in reversed(messages):
if _msg.get("role") == "user":
_content = _msg.get("content")
if isinstance(_content, str):
_last_user_msg = _content[:500]
elif isinstance(_content, list):
_parts = [
b.get("text", "")
for b in _content
if isinstance(b, dict) and b.get("type") in ("input_text", "text")
]
_last_user_msg = " ".join(_parts)[:500]
break
_req_data: dict[str, Any] = {
"model": model,
"messages_count": len(messages),
"tools_count": len(tools),
"system_prompt_preview": (system_prompt[:200] if system_prompt else ""),
"last_user_message": _last_user_msg,
}
_req_verdict = await _policy_eval("PHASE_LLM_REQUEST", _req_data)
if _req_verdict.action == "POLICY_ACTION_DENY":
_deny_reason = _req_verdict.reason or "no reason given"
yield ExecutorError(message=f"LLM call denied by policy: {_deny_reason}")
return
# Run the turn, retrying once if the gateway returns a
# completed-but-empty turn (see ``_EMPTY_TURN_MAX_ATTEMPTS``).
# ``_is_empty_turn`` guarantees an empty attempt streamed nothing
# (no text, no tool calls), so retrying never double-emits — we
# stream each attempt's events live and only defer the terminal
# TurnComplete until the emptiness decision after the loop.
result: _RunResult | None = None
response_text = ""
saw_tool_activity = False
final_text = ""
for attempt in range(_EMPTY_TURN_MAX_ATTEMPTS):
response_text = ""
pending_tools: dict[str, tuple[str, float]] = {}
saw_tool_activity = False
result = None
# Per-turn unbounded queue bridging the stream-pump task
# (producer) with this generator (consumer). Unbounded so
# the producer's ``put_nowait`` never blocks — keeping the
# producer parked on ``stream_events()``'s ``queue.get()``
# is what makes :meth:`interrupt_session`'s ``task.cancel()``
# land at the SDK boundary instead of in our queue plumbing.
event_queue: asyncio.Queue[object] = asyncio.Queue()
try:
# Reuse the count already fetched above; avoids a second
# ``get_items()`` round-trip.
state.run_item_count_before = current_item_count
result = cast(
_RunResult,
agents_sdk.Runner.run_streamed(
agent,
input=input_value,
session=state.sdk_session,
max_turns=max_turns,
run_config=run_config,
),
)
state.active_result = result
state.stream_consumer_task = asyncio.create_task(
self._consume_stream_into_queue(result, event_queue),
name=f"openai-agents-stream-pump:{session_key}",
)
while True:
item = await event_queue.get()
if item is _STREAM_DONE:
break
# Drop everything queued after interrupt, but keep
# draining so the sentinel can land and the finally
# block can run cleanly.
if state.interrupt_requested:
continue
event = cast(_StreamEvent, item)
if event.type == "raw_response_event":
raw_event = cast(_RawResponseEvent, event)
data = raw_event.data
if data.type == "response.output_text.delta":
text = data.delta
if text:
response_text += text
yield TextChunk(text=text)
elif data.type in (
"response.reasoning_summary_text.delta",
"response.reasoning_text.delta",
):
reasoning_delta = data.delta
if reasoning_delta:
yield ReasoningChunk(
delta=reasoning_delta, event_type="reasoning_text"
)
elif event.type == "run_item_stream_event":
item_event = cast(_RunItemEvent, event)
sdk_item = item_event.item
if sdk_item.type == "tool_call_item":
parts = _tool_args_from_raw_item(sdk_item.raw_item)
pending_tools[parts.call_id or parts.name] = (
parts.name,
time.monotonic(),
)
saw_tool_activity = True
yield ToolCallRequest(
name=parts.name,
args=parts.args,
metadata={"call_id": parts.call_id} if parts.call_id else {},
)
elif sdk_item.type == "tool_call_output_item":
raw_item = sdk_item.raw_item
call_id = None
if isinstance(raw_item, dict):
raw_dict: ToolArgs = raw_item
maybe_cid = raw_dict.get("call_id") or raw_dict.get("id")
call_id = maybe_cid if isinstance(maybe_cid, str) else None
else:
out_proto = cast(_ToolCallOutputRawItem, raw_item)
call_id = out_proto.call_id or out_proto.id
name = ""
started: float | None = None
# ``call_id`` is ``None`` when the SDK didn't
# surface one; in that case we can't look up
# a pending tool and emit with defaults.
if call_id is not None and call_id in pending_tools:
name, started = pending_tools.pop(call_id)
duration_ms = ((time.monotonic() - started) * 1000) if started else 0.0
classification = classify_tool_result(sdk_item.output)
yield ToolCallComplete(
name=name,
status=classification.status,
result=sdk_item.output,
error=classification.error,
duration_ms=duration_ms,
metadata={"call_id": call_id} if call_id else {},
)
# Producer signalled completion via _STREAM_DONE.
# Surface any non-cancel exception it raised so the
# outer try/except below can translate it to the
# MaxTurnsExceeded / generic-error paths exactly as
# before — these used to fire from inside the inline
# ``async for`` loop; now they fire from the producer.
pump = state.stream_consumer_task
assert pump is not None and pump.done()
# On the interrupt path, ``pump.result()`` raises
# CancelledError — state.interrupt_requested is True,
# the consumer already drained without emitting, and
# the rollback was set up by interrupt_session.
# Nothing more to do here. Non-cancel exceptions from
# the producer (MaxTurnsExceeded, generic SDK errors)
# propagate to the outer try/except blocks below.
with contextlib.suppress(asyncio.CancelledError):
pump.result()
state.started = True
state.resume_state = None
except agents_sdk.MaxTurnsExceeded:
if state.interrupt_requested:
return
if stepwise_internal_turns and saw_tool_activity and result is not None:
state.started = True
state.resume_state = result.to_state()
state.history_cursor = len(split_transient_tail(messages).persisted)
yield TurnComplete(continue_turn=True)
return
logger.error("OpenAIAgentsSDKExecutor: max turns exceeded")
yield ExecutorError(message=f"OpenAI Agents SDK exceeded max turns ({max_turns})")
return
except Exception as exc: # re-raises context overflow, classifies the rest
if state.interrupt_requested:
return
if _is_context_length_exceeded(exc):
# Re-raise so the ExecutorAdapter's error classifier
# maps it to ``context_length_exceeded`` and the
# runner surfaces the overflow to the user.
raise
from .databricks_executor import DatabricksAuthError
if isinstance(exc, DatabricksAuthError) or (
exc.__cause__ is not None and isinstance(exc.__cause__, DatabricksAuthError)
):
# When exc IS the DatabricksAuthError, use str(exc) — it
# carries the actionable "Run: databricks auth login -p X"
# message. Its __cause__ is the raw SDK exception (e.g.
# ValueError("token expired")), which is NOT actionable.
# When exc.__cause__ IS the DatabricksAuthError (exc is a
# wrapper), str(exc.__cause__) is correct.
auth_msg = (
str(exc) if isinstance(exc, DatabricksAuthError) else str(exc.__cause__)
)
logger.error("OpenAIAgentsSDKExecutor: auth failed: %s", auth_msg)
yield ExecutorError(message=auth_msg)
else:
logger.error("OpenAIAgentsSDKExecutor: run failed: %s", exc)
yield ExecutorError(message=f"OpenAI Agents SDK error: {exc}")
return
finally:
# If the outer generator was aclose'd before the
# drain completed (e.g. ExecutorAdapter returned
# without going through interrupt_session), the
# pump task is still alive — cancel it here so it
# doesn't outlive run_turn.
pump_task = state.stream_consumer_task
if pump_task is not None and not pump_task.done():
pump_task.cancel()
with contextlib.suppress(asyncio.CancelledError, Exception):
await pump_task
state.stream_consumer_task = None
state.active_result = None
state.run_item_count_before = None
if state.interrupt_requested:
return
assert result is not None # mypy: try block either assigns or raises/returns
final_output = result.final_output
if isinstance(final_output, str):
final_text = final_output
else:
final_text = agents_sdk.ItemHelpers.text_message_outputs(result.new_items)
if not final_text:
final_text = response_text
if not _is_empty_turn(final_text, saw_tool_activity, result.new_items):
break # got real output: text, tool activity, or output items
# Empty turn. Retry from a clean state if attempts remain;
# otherwise fall through to the fail-loud gate below.
if attempt + 1 < _EMPTY_TURN_MAX_ATTEMPTS:
logger.warning(
"OpenAIAgentsSDKExecutor: empty completion (attempt %d/%d), retrying",
attempt + 1,
_EMPTY_TURN_MAX_ATTEMPTS,
)
# Roll the SDK session back to its pre-turn item count so
# the retry re-runs from an identical state — the empty
# attempt may have appended a stray empty assistant item.
await self._rewind_sdk_session(state, current_item_count)
# ``final_text`` / ``result`` are from the surfaced (last) attempt.
# If still empty AND the gateway reported zero output tokens, the
# turn is a gateway hiccup, not a deliberate empty answer (which
# still bills output tokens): surface a loud retryable error so
# the workflow's retry policy can reissue, rather than a silent
# empty turn.
#
# The zero-token condition deliberately narrows fail-loud to the
# signature we actually observed (status=completed,
# output=[], 0 output tokens). An empty turn that DID bill tokens
# falls through to the silent ``TurnComplete("")`` below — exactly
# today's behavior, so this change never makes a billed-but-empty
# turn worse. If a billed-but-empty gateway failure shows up later,
# widen the gate then, with that evidence in hand.
assert result is not None
if _is_empty_turn(final_text, saw_tool_activity, result.new_items):
if _sum_output_tokens(getattr(result, "raw_responses", None)) == 0:
logger.error(
"OpenAIAgentsSDKExecutor: empty completion after %d attempts",
_EMPTY_TURN_MAX_ATTEMPTS,
)
yield ExecutorError(
message=(
"openai-agents returned an empty completion after "
f"{_EMPTY_TURN_MAX_ATTEMPTS} attempts"
),
retryable=True,
)
return
if not response_text and final_text:
yield TextChunk(text=final_text)
state.history_cursor = len(split_transient_tail(messages).persisted) + 1
# Aggregate usage across all raw responses for this turn. The SDK
# accumulates per-request usage on RunResult.raw_responses.
#
# ``input_tokens``, ``output_tokens``, and ``total_tokens`` are
# summed across all sub-turns for billing accuracy.
#
# ``context_tokens`` is set to the *last* sub-turn's total only,
# ``context_tokens`` is the last sub-turn's total (input +
# output) — the best proxy for how full the context window
# will be on the NEXT request (prior conversation + the
# response just generated). For multi-tool turns, summing
# input_tokens across sub-turns double-counts (each sub-turn
# re-sends the full history); the last sub-turn's total is
# stable. Always set so the REPL and compaction don't fall
# back to total_tokens (which sums across ALL sub-turns).
turn_usage: dict[str, Any] | None = None
raw_responses = getattr(result, "raw_responses", None)
if raw_responses:
in_tok = sum(getattr(r.usage, "input_tokens", 0) or 0 for r in raw_responses)
out_tok = sum(getattr(r.usage, "output_tokens", 0) or 0 for r in raw_responses)
total_tok = sum(getattr(r.usage, "total_tokens", 0) or 0 for r in raw_responses)
# OpenAI's ``input_tokens`` (aka ``prompt_tokens``) is the
# TOTAL input count *including* cached tokens, whereas
# ``compute_llm_cost`` expects Anthropic semantics where
# ``input_tokens`` is the non-cached portion and
# ``cache_read_input_tokens`` is additive. Extract cached
# tokens from ``prompt_tokens_details.cached_tokens`` and
# subtract so downstream billing uses the cheaper cache rate.
cached_tok = 0
for r in raw_responses:
details = getattr(r.usage, "prompt_tokens_details", None)
if details is not None:
cached = getattr(details, "cached_tokens", None)
if cached is None and isinstance(details, dict):
cached = details.get("cached_tokens")
cached_tok += cached or 0
last_r = raw_responses[-1]
last_in = getattr(last_r.usage, "input_tokens", 0) or 0
last_out = getattr(last_r.usage, "output_tokens", 0) or 0
last_total = getattr(last_r.usage, "total_tokens", 0) or 0
context_tok = last_total if last_total else last_in + last_out
if in_tok or out_tok:
turn_usage = {
"input_tokens": in_tok - cached_tok, # non-cached portion
"output_tokens": out_tok,
"total_tokens": total_tok if total_tok else in_tok + out_tok,
"context_tokens": context_tok,
# Harness-reported model for cost pricing when the spec pins no model.
"model": model,
}
if cached_tok:
turn_usage["cache_read_input_tokens"] = cached_tok
_notify_usage_from_dict(model=model, usage=turn_usage)
# Emit CompactionComplete if the SDK compacted this turn.
if result is not None:
for item in result.new_items:
if getattr(item, "type", None) == "compaction_item":
from omnigent.inner.executor import CompactionComplete
_compaction_tokens = 0
if turn_usage is not None:
_compaction_tokens = turn_usage.get("context_tokens", 0) or 0
_compacted: list[dict[str, Any]] | None = None
try:
_compacted = await state.sdk_session.get_items()
except Exception: # noqa: BLE001
logger.warning(
"Failed to read compacted session items",
exc_info=True,
)
yield CompactionComplete(
summary=(
"[OpenAI Responses API compaction"
" — context was automatically compacted]"
),
token_count=_compaction_tokens,
model=model,
compacted_messages=_compacted,
)
break
yield TurnComplete(response=final_text, usage=turn_usage)
@staticmethod
def _filter_model_input(data: _CallModelData) -> _ModelInputData:
# Sanitize so identity fields (long ``id`` strings) are stripped
# and content blocks are normalized for the chat-completions path.
sanitized = [_sanitize_replay_item(item) for item in data.model_data.input]
data.model_data.input = _normalize_responses_items_for_chat(sanitized)
return data.model_data
Reference in New Issue View Git Blame Copy Permalink