class Agent:
"""
Main orchestrator for LLM interactions with automatic tool execution.
Coordinates the provider, session management, and tool execution to
provide a complete conversational interface.
The agent auto-initializes on first `run()` call, verifying the model
and setting up the session database. You can also call `initialize()`
manually if you want to catch initialization errors early.
Example:
provider = Provider(...)
session = SessionManager(Path("sessions.db"))
agent = Agent(
provider=provider,
session_manager=session,
tools=[my_tool_func],
system_prompt="You are a helpful assistant.",
)
# Just use it - auto-initializes on first run()
async for event in agent.run("Hello!"):
if isinstance(event, TextChunkEvent):
print(event.chunk, end="")
elif isinstance(event, ToolCallEvent):
print(f"Calling tool: {event.name}")
elif isinstance(event, DoneEvent):
print(f"\\nSession: {event.session_id}")
# Don't forget to close when done
await agent.close()
# Or pre-initialize to catch errors early:
await agent.initialize()
async for event in agent.run("Hello!"):
...
Batch Mode:
When batch=True, the agent uses batch processing API (50% cost discount).
The interface remains the same - just add batch=True:
agent = Agent(
provider=provider,
session_manager=session,
batch=True, # Enable batch mode
)
async for event in agent.run("Hello!"):
# Events are yielded after batch completes
...
"""
# Set by the Harness, or an embedding application, to allow bounded workspace
# file reads (for example outbound channel attachments). None disables them.
workspace: Path | None = None
_execution_bridge: "_DeliveryExecution | None" = None
def __init__(
self,
provider: Provider,
session_manager: SessionManager,
tools: list[Callable[..., Any]] | None = None,
system_prompt: str | None = None,
max_tool_rounds: int = 1000,
streaming: bool = False,
log_file: Path | str | None = None,
fail_on_invalid_tool: bool = False,
batch: bool = False,
batch_poll_interval: float = 30.0,
stt_service: "STTService | None" = None,
unsupported_audio: UnsupportedAudioBehavior = UnsupportedAudioBehavior.RAISE,
audio: "AudioDuplex | None" = None,
retry_config: RetryConfig | None = None,
plugins: Iterable[AgentPlugin] = (),
compaction_strategy: CompactionStrategy | _DefaultCompactionStrategy | None = _DEFAULT_COMPACTION_STRATEGY,
tool_executor: ToolExecutor | None = None,
save_tool_outputs: bool = True,
skill_discoverer: SkillDiscoverer | None = None,
skill_token_limit: int = DEFAULT_SKILL_TOKEN_LIMIT,
delegation_agent: "Agent | None" = None,
):
"""
Initialize the agent.
Args:
provider: The LLM provider to use
session_manager: Session manager for conversation history.
tools: Optional list of tool functions to register
system_prompt: Optional system prompt to prepend to conversations
max_tool_rounds: Maximum number of tool execution rounds
streaming: Whether to stream responses (default: False)
log_file: Optional path to log file for HTTP/SSE traffic debugging.
Creates parent directories if they don't exist.
fail_on_invalid_tool: If True, emit an error and stop when the LLM
tries to call a tool that doesn't exist (tool
hallucination). If False (default), the error
is passed back to the LLM so it can recover.
batch: If True, use batch processing API (50% cost discount).
Results are delayed but significantly cheaper.
batch_poll_interval: Seconds between batch status checks (default: 30)
stt_service: Optional Speech-to-Text service for audio transcription.
If provided, audio content in messages will be transcribed
to text before being sent to the LLM.
unsupported_audio: How to handle audio sent to models that don't support
inline audio input. Options:
- RAISE (default): Raise UnsupportedAudioError
- DROP: Replace audio with a text warning
- TRANSCRIBE: Auto-transcribe using stt_service
(requires stt_service to be set)
audio: Optional :class:`~nagents.audio.AudioDuplex` pairing an
:class:`~nagents.audio.AudioInput` and
:class:`~nagents.audio.AudioOutput`. When set, calling
:meth:`run` with no message runs a Realtime speech-to-speech
conversation. Either end may be ``None`` to disable it.
retry_config: Optional retry configuration for rate limit (HTTP 429)
and server error (5xx) handling. If provided, overrides
the provider's retry config. Retries are enabled by
default (3 retries with exponential backoff). Set
RetryConfig(max_retries=0) to disable.
compaction_strategy: Strategy for automatic and forced context replacement.
Omitted selects a fresh ModelCompaction; None disables it.
plugins: Trusted, ordered hooks for the text run lifecycle.
tool_executor: Optional authorization/execution boundary. The consumer
must ensure its registry matches this agent's tools;
assigning an executor after construction is supported.
save_tool_outputs: Enable the reserved _save_to filesystem convention.
Disable for executor-controlled filesystem access.
skill_discoverer: Optional live text-only skill source. Installs skill(name)
and refreshes the catalog at message, model and tool boundaries.
skill_token_limit: Per-load and aggregate explicit-activation content budget,
estimated as ceil(UTF-8 bytes / 4), default 10,000.
delegation_agent: Client-owned GPT-Live backend. Runs normal text/tool
inference independently of audio. Caller owns its lifetime.
"""
self.provider = provider
self.session = session_manager
self.system_prompt = system_prompt
self.max_tool_rounds = max_tool_rounds
self.streaming = streaming
self.fail_on_invalid_tool = fail_on_invalid_tool
self.batch = batch
self.batch_poll_interval = batch_poll_interval
self.stt_service = stt_service
self.unsupported_audio = unsupported_audio
self.audio = audio
self.delegation_agent = delegation_agent
self._live_updates: _LiveUpdates | None = None
self._last_live_updates: _LiveUpdates | None = None
self.plugins: list[AgentPlugin] = list(plugins)
self._default_compaction_strategy: ModelCompaction | None = None
if isinstance(compaction_strategy, _DefaultCompactionStrategy):
self._default_compaction_strategy = ModelCompaction()
self.compaction_strategy: CompactionStrategy | None = self._default_compaction_strategy
else:
self.compaction_strategy = compaction_strategy
self.save_tool_outputs = save_tool_outputs
validate_skill_token_limit(skill_token_limit)
self.skill_token_limit = skill_token_limit
self._skill_discoverer = skill_discoverer
self._skills: Mapping[str, Skill] = MappingProxyType({})
self._skill_lock = asyncio.Lock()
# Forward retry config to provider if provided
if retry_config is not None:
self.provider.retry_config = retry_config
# Validate that TRANSCRIBE mode has an STT service
if self.unsupported_audio == UnsupportedAudioBehavior.TRANSCRIBE and not self.stt_service:
raise ValueError(
"unsupported_audio='transcribe' requires an stt_service to be provided. "
"Use OpenAISTTService or GeminiSTTService."
)
self.tool_registry = ToolRegistry()
self._default_tool_executor = ToolExecutor(self.tool_registry)
self.tool_executor = tool_executor if tool_executor is not None else self._default_tool_executor
if self.batch:
self._check_extension_mode("batch")
# Track actual prompt_tokens from API responses per session
# Maps session_id -> last known prompt_tokens
self._session_tokens: dict[str, int] = {}
# Main-generation usage observed by this agent, not persisted history or compaction billing.
self._session_usage: dict[str, TokenUsage] = {}
# Last single generation's real usage per session, for context-stat comparison.
self._last_usage: dict[str, TokenUsage] = {}
# Flag to request compaction during run (set by trigger_compaction())
self._compaction_requested: bool = False
self._session_compaction_requests: set[str] = set()
self._initialized = False
self._http_logger: FileHTTPLogger | None = None
self._batch_client: BatchClient | None = None
self._channels: dict[str, Channel] = {}
self._channel_listening = False
self._channel_listener_task: asyncio.Task[None] | None = None
self._channel_execution_task: asyncio.Task[None] | None = None
self._active_runs = 0
self._cleanup_tasks: set[asyncio.Task[None]] = set()
self._closing = False
# Retain the completed task for idempotent close until resources are reused.
self._close_task: asyncio.Task[None] | None = None
self._close_provider = provider
# Set up HTTP logging if log_file is provided
if log_file:
log_path = Path(log_file) if isinstance(log_file, str) else log_file
self._http_logger = FileHTTPLogger(log_path)
self.provider.set_http_logger(self._http_logger)
logger.info(f"HTTP logging enabled: {log_path}")
if tools:
for tool in tools:
self.tool_registry.register(tool)
self._ensure_skill_tool()
# Initialize batch client if batch mode enabled
if self.batch:
self._init_batch_client()
def _init_batch_client(self) -> None:
"""Initialize the batch client for batch mode."""
if self.provider.bearer_token_provider is not None:
raise ValueError("Batch mode does not support bearer_token_provider")
valid_types = (
ProviderType.OPENAI_COMPATIBLE,
ProviderType.AZURE_OPENAI_COMPATIBLE,
ProviderType.AZURE_OPENAI_COMPATIBLE_V1,
ProviderType.ANTHROPIC,
)
if self.provider.provider_type not in valid_types:
raise ValueError(
f"Batch mode only supported for OPENAI_COMPATIBLE, AZURE_OPENAI, and ANTHROPIC providers. "
f"Got: {self.provider.provider_type}"
)
self._batch_client = BatchClient(
provider_type=self.provider.provider_type,
api_key=self.provider.api_key,
model=self.provider.model,
base_url=self.provider.base_url,
logger=self._http_logger,
)
logger.info("Batch mode enabled - requests will be processed with 50% cost discount")
def _check_extension_mode(self, mode: str) -> None:
if (
self.plugins
or (
self.compaction_strategy is not None
and self.compaction_strategy is not self._default_compaction_strategy
)
or self.tool_executor is not self._default_tool_executor
or not self.save_tool_outputs
or self.skill_discoverer is not None
):
raise ValueError(f"Agent extensions are not supported in {mode} mode; use text run()")
@property
def skill_discoverer(self) -> SkillDiscoverer | None:
return self._skill_discoverer
@skill_discoverer.setter
def skill_discoverer(self, discoverer: SkillDiscoverer | None) -> None:
"""Install or replace a source; the old catalog is immediately invalidated."""
if discoverer is not None and self.batch:
raise ValueError("Agent skills are not supported in batch mode; use text run()")
self._skill_discoverer = discoverer
self._skills = MappingProxyType({})
self._ensure_skill_tool()
@property
def skills(self) -> Mapping[str, Skill]:
"""Read-only snapshot of the most recently completed catalog replacement."""
return self._skills
def _ensure_skill_tool(self) -> None:
existing = self.tool_registry.get("skill")
if self.skill_discoverer is not None and existing is None:
self.tool_registry.register(self.load_skill, name="skill")
elif self.skill_discoverer is None and existing is not None and existing.func == self.load_skill:
self.tool_registry.unregister("skill")
async def _refresh_skills_locked(self) -> Mapping[str, Skill]:
while (source := self.skill_discoverer) is not None:
catalog = skill_catalog(await source.discover())
if source is self.skill_discoverer:
self._skills = MappingProxyType(catalog)
self._ensure_skill_tool()
return self._skills
# A source was replaced while awaiting discovery. Never publish it.
self._skills = MappingProxyType({})
return self._skills
async def refresh_skills(self) -> Mapping[str, Skill]:
"""Atomically replace the entire live catalog; discovery errors propagate.
Sources should skip malformed/disappearing entries with diagnostics.
Protocol failures leave the last successful snapshot visible to callers,
but abort the current run rather than send a stale catalog to a model.
"""
if self.skill_discoverer is None:
return self._skills
async with self._skill_lock:
return await self._refresh_skills_locked()
async def load_skill(self, name: str) -> SkillLoadResult:
"""Load current skill text as task context, within the configured token budget.
Args:
name: Exact name from the available skills catalog.
"""
async with self._skill_lock:
await self._refresh_skills_locked()
source = self.skill_discoverer
skill = self.skills.get(name)
if source is None or skill is None:
raise ValueError(f"Unknown skill {name!r}; refresh the available skills catalog")
content = await source.load(skill)
if source is not self.skill_discoverer:
raise ValueError("Skill source changed while loading; retry the load")
return budget_skill_content(skill, content, self.skill_token_limit)
async def _execute_text_tool(
self, call: ToolCall, context: RunContext, plugins: tuple[AgentPlugin, ...]
) -> ToolResultEvent:
identity = (call.id, call.name)
await self.refresh_skills()
try:
result = await self.tool_executor.execute(call)
if not isinstance(result, ToolResultEvent) or (result.id, result.name) != identity:
raise ValueError("Tool executor must preserve tool call id/name")
for plugin in plugins:
result = await plugin.after_tool(context, result)
if not isinstance(result, ToolResultEvent) or (result.id, result.name) != identity:
raise ValueError(f"{type(plugin).__name__}.after_tool must preserve tool call id/name")
except Exception as error:
# Keep the execution failure primary if discovery also fails. Do not
# start more discovery on cancellation/GeneratorExit during teardown.
try:
await self.refresh_skills()
except Exception as refresh_error:
error.add_note(f"After-tool skill refresh failed: {refresh_error}")
raise
await self.refresh_skills()
return result
async def _activate_skills(
self,
names: tuple[str, ...],
context: RunContext,
plugins: tuple[AgentPlugin, ...],
task_context: list[Message],
) -> AsyncGenerator[Event, None]:
remaining = self.skill_token_limit
# Bound metadata/event overhead even for empty skill bodies.
for name in names[:32]:
if remaining <= 0:
break
skill = self.skills.get(name)
if skill is None:
continue
call = ToolCall(f"skill-ref-{uuid.uuid4().hex[:12]}", "skill", {"name": name})
identity = (call.id, call.name)
for plugin in plugins:
call = await plugin.before_tool(context, call)
if not isinstance(call, ToolCall) or (call.id, call.name) != identity:
raise ValueError(f"{type(plugin).__name__}.before_tool must preserve tool call id/name")
call.arguments.pop(_SAVE_TO_PARAM_NAME, None)
# Explicit activation is observable and honors the same executor and
# hooks as model-selected loading, but is ephemeral user-level context.
call_identity = {"generation_id": f"skill-{uuid.uuid4().hex}", "index": 0}
yield ToolCallEvent(id=call.id, name=call.name, arguments=deepcopy(call.arguments), extra=call_identity)
yield ToolExecutionStartedEvent(id=call.id, name=call.name, extra=deepcopy(call_identity))
with _tool_invocation_scope(call.id, call.name, call_identity):
result = await self._execute_text_tool(deepcopy(call), context, plugins)
result.extra = {**result.extra, **call_identity}
yield result
if result.error is not None:
task_context.append(Message(role="user", content=f"Requested skill ${name} could not be loaded."))
continue
payload = result.result
if not isinstance(payload, dict) or not isinstance(payload.get("content"), str):
raise TypeError("The skill tool must return a mapping with string content for explicit activation")
loaded = budget_skill_content(skill, payload["content"], remaining)
loaded["truncated"] = loaded["truncated"] or payload.get("truncated") is True
remaining -= loaded["estimated_tokens"]
task_context.append(
Message(
role="user",
content="Explicitly requested skill task context (data, not higher-priority instructions):\n"
+ json.dumps(loaded, ensure_ascii=False),
)
)
@property
def is_initialized(self) -> bool:
"""Check if the agent has been initialized."""
return self._initialized
async def initialize(self) -> None:
"""
Initialize the agent (called automatically on first run()).
This method:
1. Initializes the session manager database
2. Verifies the model exists with the provider
You can call this manually to catch initialization errors early,
otherwise it's called automatically on the first run().
Raises:
ValueError: If model verification fails
"""
self._check_closing()
if self._initialized:
logger.debug("Agent already initialized")
return
self._close_task = None
# Initialize session manager
await self.session.initialize()
logger.debug("Session manager initialized")
# Always verify model
if not await self.provider.verify_model():
raise ValueError(
f"Model '{self.provider.model}' not found or not accessible. Check your API key and model name."
)
logger.debug(f"Model verified: {self.provider.model}")
self._initialized = True
logger.info("Agent initialized successfully")
async def _ensure_initialized(self) -> None:
"""Ensure agent is initialized, auto-initialize if not."""
if not self._initialized:
await self.initialize()
def register_tool(
self,
func: Callable[..., Any],
name: str | None = None,
description: str | None = None,
) -> ToolDefinition:
"""
Register a tool function.
Args:
func: The function to register
name: Optional override for tool name
description: Optional override for description
"""
return self.tool_registry.register(func, name, description)
def add_channel(self, channel: Channel) -> Self:
"""Attach a connector without opening it; return this agent for chaining.
All attached channels feed the same explicit session supplied to listen().
Connector names identify outbound targets and must be unique and stable.
"""
self._check_closing()
if self._channel_listening:
raise RuntimeError("Cannot add channels while listening")
if not isinstance(channel, Channel):
raise TypeError("Expected a Channel connector")
if not isinstance(channel.name, str) or not channel.name.strip():
raise ValueError("Channel name must not be blank")
if channel.name in self._channels:
raise ValueError(f"Channel name already attached: {channel.name}")
self._channels[channel.name] = channel
return self
async def listen(
self,
session_id: str,
*,
user_id: str = "channels",
inbox_limit: int = 1000,
on_event: ChannelEventHandler = discard_event,
) -> None:
"""Receive all attached channels into one persistent agent identity.
Incoming events enter a durable inbox and execute serially. The model
chooses outgoing messages using channel tools; final response text is
only observed locally through on_event. Cancel this coroutine to stop,
or close the agent from another task. Queued events survive restart;
interrupted executions are not automatically replayed.
The listener owns this Agent until shutdown. Independent run(), compact(),
and session mutations are rejected during that time. This is text mode;
batch execution is unsupported. Use one process per session database.
"""
from .channels.runtime import ChannelRuntime
self._check_closing()
if self._channel_listening or self._active_runs:
raise RuntimeError("Agent is already executing or listening")
if self.batch:
raise ValueError("Channel listeners require text execution, not batch mode")
if not self._channels:
raise ValueError("Attach at least one channel before listening")
if not isinstance(session_id, str) or not session_id.strip():
raise ValueError("An explicit nonblank session_id is required for the agent identity")
if any(name != channel.name for name, channel in self._channels.items()):
raise ValueError("An attached channel changed its name")
runtime = ChannelRuntime(
self, tuple(self._channels.values()), session_id, user_id, inbox_limit, workspace=self.workspace
)
self._close_task = None
self._channel_listening = True
self._channel_listener_task = asyncio.current_task()
try:
await runtime.listen(on_event=on_event)
finally:
self._channel_listening = False
self._channel_listener_task = None
self._channel_execution_task = None
def _check_closing(self, *, allow_cleanup: bool = False) -> None:
if self._closing and not (allow_cleanup and asyncio.current_task() in self._cleanup_tasks):
raise RuntimeError("Agent is closing; wait for resource cleanup before starting new work")
def _check_channel_access(self, *, allow_cleanup: bool = False) -> None:
self._check_closing(allow_cleanup=allow_cleanup)
task = asyncio.current_task()
if self._channel_listening and task is not self._channel_execution_task and task not in self._cleanup_tasks:
raise RuntimeError("The channel listener owns this Agent; submit input through its channels")
async def _finish_cleanup(self, operation: Awaitable[None]) -> None:
"""Track each cleanup task, including nested trusted run cleanup."""
async def cleanup() -> None:
task = asyncio.current_task()
assert task is not None
self._cleanup_tasks.add(task)
try:
await operation
finally:
self._cleanup_tasks.discard(task)
await _join_cleanup(asyncio.create_task(cleanup()))
@asynccontextmanager
async def _closing_events(self, events: AsyncIterator[Event]) -> AsyncIterator[AsyncIterator[Event]]:
"""Close provider iterators when supported, without narrowing their API."""
try:
yield events
finally:
if isinstance(events, _AsyncClosable):
await self._finish_cleanup(events.aclose())
def realtime_session(
self,
model: str | None = None,
voice: str | None = None,
config: "RealtimeConfig | None" = None,
audio: "AudioDuplex | None" = None,
base_url: str = "wss://api.openai.com/v1/realtime",
heartbeat: float | None = None,
log_file: Path | str | None = None,
) -> "RealtimeSession":
"""
Create a Realtime speech-to-speech session for this agent.
Reuses this agent's provider credentials, system prompt, and registered
tools. The model and its configuration come from the provider's
:attr:`~nagents.Provider.realtime_config` (or ``config``), with the
model falling back to ``model``, then the provider's model, then the
default Realtime model.
Requires an OpenAI-compatible provider.
Example:
async with agent.realtime_session() as session:
async for event in session:
...
Args:
model: Optional Realtime model override.
voice: Optional output voice override.
config: Optional :class:`~nagents.realtime.RealtimeConfig` override
(defaults to ``provider.realtime_config``). Empty instructions
inherit the agent's system prompt; the supplied config is not mutated.
audio: Optional :class:`~nagents.audio.AudioDuplex` (defaults to
the agent's ``audio``).
base_url: WebSocket endpoint (defaults to OpenAI).
heartbeat: Optional WebSocket heartbeat interval in seconds.
log_file: Optional path to a JSON-lines log of all raw events.
Returns:
A connectable :class:`~nagents.realtime.RealtimeSession`.
Raises:
ValueError: If the provider is not OpenAI-compatible.
"""
self._check_extension_mode("voice")
from .realtime import DEFAULT_REALTIME_MODEL
from .realtime import RealtimeConfig
from .realtime import RealtimeSession
if self.provider.bearer_token_provider is not None:
raise ValueError("Realtime does not support bearer_token_provider; use GPT-Live with live_config")
if self.provider.provider_type != ProviderType.OPENAI_COMPATIBLE:
raise ValueError(
"Realtime speech-to-speech sessions require an OPENAI_COMPATIBLE provider. "
f"Got: {self.provider.provider_type}"
)
cfg = deepcopy(config or self.provider.realtime_config or RealtimeConfig())
cfg.instructions = cfg.instructions or self.system_prompt or ""
if model is not None:
cfg.model = model
if voice is not None:
cfg.voice = voice
if cfg.model is None:
cfg.model = self.provider.model or DEFAULT_REALTIME_MODEL
duplex = audio or self.audio
tool_funcs = [td.func for td in self.tool_registry.get_all() if td.func is not None]
return RealtimeSession(
api_key=self.provider.api_key,
tools=tool_funcs,
config=cfg,
audio_in=duplex.input if duplex else None,
audio_out=duplex.output if duplex else None,
base_url=base_url,
heartbeat=heartbeat,
log_file=log_file,
)
@property
def live(self) -> "_LiveUpdates":
"""Native Live controls and latest session status (available after startup)."""
controls = self._live_updates or self._last_live_updates
if controls is None:
raise RuntimeError("This agent has not started a Live connection")
return controls
def live_configuration(self, *, media: bool = False) -> dict[str, object]:
"""Session setup for native Live HTTP provisioning; media omits PCM format."""
from .live.runtime import build_session
return build_session(self, media=media)
async def add_comment(self, content: str) -> str:
"""Queue a short spoken update for GPT-Live; returns its event ID.
Usable as a backend tool or from application code. During a delegation,
the original delegation ID is included automatically. Otherwise the
update is session-wide. Queuing does not mean playback has completed.
"""
from .live import append_update
return await append_update(self, "commentary", content)
async def add_thinking(self, content: str) -> str:
"""Queue factual context or quiet progress for GPT-Live; return its event ID."""
from .live import append_update
return await append_update(self, "thinking", content)
async def add_instructions(self, content: str) -> str:
"""Queue application instructions for GPT-Live; return its event ID.
This can redirect speech, but does not cancel backend work.
"""
from .live import append_update
return await append_update(self, "instructions", content)
async def _run_voice(
self,
*,
auto_commit: bool | None = None,
log_file: Path | str | None = None,
) -> AsyncGenerator[Event, None]:
"""
Run a full-duplex speech-to-speech conversation through this agent.
Uses the agent's ``audio`` :class:`~nagents.audio.AudioDuplex` (set
at construction) and the provider's Realtime model/config. Audio flows
in through the duplex input and out through its output, while
transcripts and tool events are yielded — the same event stream as
:meth:`run`.
"""
duplex = self.audio
if duplex is None and self.provider.live_config is not None and self.provider.live_config.attach_to:
from .audio import AudioDuplex
duplex = AudioDuplex()
if duplex is None or (
duplex.input is None
and duplex.output is None
and not (self.provider.live_config and self.provider.live_config.attach_to)
):
raise ValueError(
"Voice mode requires an AudioDuplex with at least an input or output. "
"Set audio=AudioDuplex(input=..., output=...) on the Agent."
)
if self.provider.live_config is not None:
from .live import run_live
if auto_commit:
raise ValueError("GPT-Live uses continuous audio, not Realtime input commits")
async with aclosing(run_live(self, duplex, log_file=log_file)) as events:
async for event in events:
yield event
return
session = self.realtime_session(audio=duplex, log_file=log_file)
try:
await session.connect()
async for event in session.run_duplex(duplex.input, duplex.output, auto_commit=auto_commit):
yield event
finally:
await session.close()
async def _process_multimodal_content(
self,
content: str | list[ContentPart] | None,
) -> str | list[ContentPart] | None:
"""
Process multimodal content, handling audio based on provider capabilities
and the ``unsupported_audio`` configuration.
Processing order for each AudioContent part:
1. If provider supports the audio format natively → pass through.
2. If provider supports audio but not *this* format → transcode to WAV.
3. If provider does NOT support audio at all → apply ``unsupported_audio`` policy:
- RAISE: raise UnsupportedAudioError
- DROP: replace with a TextContent warning
- TRANSCRIBE: use stt_service to transcribe to text
When ``stt_service`` is set independently (without ``unsupported_audio=TRANSCRIBE``),
audio is transcribed only when the provider doesn't support the format and
transcoding also isn't an option.
Args:
content: Message content (str, list of ContentParts, or None)
Returns:
Processed content with audio handled appropriately
Raises:
UnsupportedAudioError: If unsupported_audio is RAISE and audio is sent
to a model that doesn't support it
"""
if content is None or isinstance(content, str):
return content
capabilities = self.provider.supported_media_formats
processed_parts: list[ContentPart] = []
for part in content:
if not isinstance(part, AudioContent):
processed_parts.append(part)
continue
# Case 1: Provider supports this exact audio format → pass through
if capabilities.supports_audio(part.format):
processed_parts.append(part)
continue
# Case 2: Provider supports audio, but not this format → transcode to WAV
if capabilities.audio_formats:
logger.info(
f"Audio format '{part.format}' not supported by provider "
f"(supported: {sorted(capabilities.audio_formats)}). "
"Transcoding to WAV via ffmpeg..."
)
try:
wav_data = await transcode_audio_to_wav(part.base64_data, part.format)
processed_parts.append(AudioContent(base64_data=wav_data, format="wav"))
logger.info(f"Transcoded audio from '{part.format}' to 'wav' successfully")
except RuntimeError as e:
logger.error(f"Audio transcoding failed: {e}")
processed_parts.append(TextContent(text=f"[Audio transcoding failed: {e}]"))
continue
# Case 3: Provider does NOT support audio at all → apply policy
if self.unsupported_audio == UnsupportedAudioBehavior.RAISE:
raise UnsupportedAudioError(self.provider.model, part.format)
elif self.unsupported_audio == UnsupportedAudioBehavior.DROP:
logger.info(
f"Model '{self.provider.model}' does not support inline audio input. "
"Audio replaced with placeholder text (unsupported_audio='drop')."
)
processed_parts.append(
TextContent(
text="system-message: You are receiving this because the user sent an audio message, "
"but your model does not support direct audio input. "
"Audio-capable models include the Gemini family or GPT voice models (gpt-audio, gpt-4o-audio-preview). "
"Alternatively, STT can be enabled to receive transcriptions instead. "
"Audio omitted."
)
)
elif self.unsupported_audio == UnsupportedAudioBehavior.TRANSCRIBE:
# stt_service is guaranteed present (validated in __init__)
assert self.stt_service is not None
logger.info("Auto-transcribing audio (unsupported_audio='transcribe')...")
try:
result = await self.stt_service.transcribe_base64(
part.base64_data,
part.format,
)
logger.info(f"Transcription: {result.text[:100]}...")
if result.text:
processed_parts.append(TextContent(text=result.text))
else:
processed_parts.append(TextContent(text="[Audio transcription returned empty]"))
except Exception as e:
logger.error(f"Failed to transcribe audio: {e}")
processed_parts.append(TextContent(text=f"[Audio transcription failed: {e}]"))
return processed_parts
def _filter_unsupported_audio(self, messages: list[Message]) -> list[Message]:
"""
Filter unsupported AudioContent from message history.
When a model doesn't support inline audio (empty audio_formats),
replaces AudioContent parts in *history* messages with TextContent warnings.
This prevents 400 errors from old audio in session history being sent
to non-audio models.
Only applies when ``unsupported_audio`` is DROP or TRANSCRIBE.
When RAISE, history audio should not exist (since the initial send would
have raised), but we still filter defensively.
"""
capabilities = self.provider.supported_media_formats
if capabilities.audio_formats:
# Model supports audio — no filtering needed
return messages
filtered: list[Message] = []
for msg in messages:
if not isinstance(msg.content, list):
filtered.append(msg)
continue
has_audio = any(isinstance(p, AudioContent) for p in msg.content)
if not has_audio:
filtered.append(msg)
continue
# Replace AudioContent with text warning
new_parts: list[ContentPart] = []
for part in msg.content:
if isinstance(part, AudioContent):
new_parts.append(
TextContent(
text="system-message: You are receiving this because the user sent an audio message, "
"but your model does not support direct audio input. "
"Audio-capable models include the Gemini family or GPT voice models (gpt-audio, gpt-4o-audio-preview). "
"Alternatively, STT can be enabled to receive transcriptions instead. "
"Audio omitted."
)
)
else:
new_parts.append(part)
filtered.append(
Message(
role=msg.role,
content=new_parts,
tool_calls=msg.tool_calls,
tool_call_id=msg.tool_call_id,
name=msg.name,
)
)
return filtered
async def _maybe_compact(self, messages: list[Message], session_id: str) -> AsyncIterator[Event]:
if self.compaction_strategy is not None:
async for event in self._compact_with_strategy(session_id):
yield event
async def _do_compact(self, messages: list[Message], session_id: str) -> AsyncIterator[Event]:
if self.compaction_strategy is None:
raise ValueError("No compaction strategy configured")
async for event in self._compact_with_strategy(session_id, force=True):
yield event
async def _compact_with_strategy(self, session_id: str, force: bool = False) -> AsyncIterator[Event]:
strategy = self.compaction_strategy
assert strategy is not None
history = await self.session.get_history(session_id)
original_count = len(history)
original_tokens = estimate_messages_tokens(history)
request = CompactionRequest(
messages=deepcopy(history),
session_id=session_id,
provider=self.provider,
estimated_tokens=original_tokens,
force=force,
actual_input_tokens=self._session_tokens.get(session_id, 0),
)
if not force and not await strategy.should_compact(request):
return
yield CompactionStartedEvent(
message_count=original_count,
estimated_tokens=original_tokens,
session_id=session_id,
)
result = await strategy.compact(request)
if not isinstance(result, CompactionResult):
raise TypeError("CompactionStrategy.compact must return CompactionResult")
if not result.messages:
raise ValueError("Compaction strategy returned empty history")
await self.session.replace_context(session_id, result.messages)
new_tokens = estimate_messages_tokens(result.messages)
self._session_tokens[session_id] = new_tokens
yield CompactionDoneEvent(
original_message_count=original_count,
original_token_count=original_tokens,
new_message_count=len(result.messages),
summary_tokens=new_tokens,
summary_text=result.summary,
compactor_used=type(strategy).__name__,
session_id=session_id,
)
def trigger_compaction(self, session_id: str = "") -> None:
"""Request compaction before the next generation cycle.
Call this to trigger compaction before the next model call.
The compaction happens after all pending tool calls complete,
before the next generation round.
Useful when a tool knows it has generated a lot of context.
Pass the active session ID to scope the request to that conversation.
Scoped requests are discarded when its run ends or is cancelled.
Example:
# During run(), call this before the next generation:
agent.trigger_compaction()
# Compaction will happen before next model call
"""
if session_id:
self._session_compaction_requests.add(session_id)
else:
self._compaction_requested = True
async def compact(self, session_id: str) -> "CompactionDoneEvent":
"""Manually trigger compaction for a session.
Call this outside of run() to compact a session on demand.
Args:
session_id: Session to compact
Returns:
CompactionDoneEvent with summary and stats
Raises:
ValueError: If no compaction strategy is configured or session not found
Example:
# Compact a session before continuing
result = await agent.compact("session-123")
print(f"Compacted: {result.original_message_count} -> {result.new_message_count}")
print(f"Summary: {result.summary_text[:100]}...")
"""
self._check_channel_access()
# Manual compaction can use the provider without calling initialize().
self._close_task = None
self._active_runs += 1
try:
return await self._compact_session(session_id)
finally:
self._active_runs -= 1
async def _compact_session(self, session_id: str) -> "CompactionDoneEvent":
if not await self.session.session_exists(session_id):
raise ValueError(f"Session '{session_id}' not found")
# Get current messages
messages: list[Message] = []
if self.system_prompt:
messages.append(Message(role="system", content=self.system_prompt))
history = await self.session.get_history(session_id)
messages.extend(history)
# Run compaction and collect result
result: CompactionDoneEvent | None = None
async for event in self._do_compact(messages, session_id):
if isinstance(event, CompactionDoneEvent):
result = event
if result is None:
raise ValueError("Compaction did not complete")
return result
async def context_stats(self, session_id: str) -> ContextStats:
"""Estimate the token breakdown of the request this session would send.
Read-only: it assembles the same system prompt, tool schemas, skill
catalog and history that the next text request uses, without calling the
provider or changing session state. Values are character/byte heuristics;
see :func:`nagents.estimate_context_stats` for what is included.
Explicit ``$skill`` activation text is request-local and never persisted,
so it is not part of an idle estimate.
Args:
session_id: Session to account for.
Raises:
ValueError: If the session does not exist.
"""
self._check_channel_access()
if not await self.session.session_exists(session_id):
raise ValueError(f"Session '{session_id}' not found")
history = await self.session.get_history(session_id)
messages = self._filter_unsupported_audio(list(history))
tools = self.tool_registry.get_all() if self.tool_registry.has_tools() else []
if tools and self.save_tool_outputs:
tools = _inject_save_to(tools)
skill_manifest = render_skill_manifest(self.skills.values()) if self.skill_discoverer is not None else ""
observed = self._last_usage.get(session_id)
return estimate_context_stats(
system_prompt=self.system_prompt or "",
tools=tools,
skill_manifest=skill_manifest,
messages=messages,
context_window=get_model_context_limit(self.provider),
provider=self.provider.provider_type.value,
model=self.provider.model,
observed=observed,
)
async def run(
self,
user_message: str | list[ContentPart] | Message | None = None,
session_id: str | None = None,
user_id: str = "default",
config: GenerationConfig | None = None,
*,
auto_commit: bool | None = None,
log_file: Path | str | None = None,
) -> AsyncGenerator[Event, None]:
"""
Run a complete interaction with automatic tool execution.
Auto-initializes the agent on first call (verifies model, sets up database).
Two modes:
- **Text**: pass ``user_message`` (str, list of content parts, or a
:class:`Message`). Yields text chunks, tool calls/results, etc.
- **Voice**: pass ``user_message=None`` and configure ``audio=`` on the
agent. Runs a Realtime speech-to-speech conversation, yielding the
same event stream (transcripts, tool calls/results, done, errors).
Args:
user_message: The user's message, or ``None`` for voice mode. Can be:
- str: Simple text message
- list[ContentPart]: Multimodal content (text, images, audio)
- Message: Full message object
If audio content is present and stt_service is configured,
audio will be transcribed to text.
session_id: Optional session identifier. If provided, verifies it exists
or creates a new one. If None, creates a new session.
user_id: User identifier (default: "default")
config: Optional generation configuration (text mode)
auto_commit: Voice mode: whether to commit input and request a response
when the input is exhausted (push-to-talk). Defaults to the
inverse of the Realtime config's turn detection.
log_file: Voice mode: optional JSON-lines log of all raw events.
Yields:
Events as they occur
Raises:
ValueError: If model verification fails during auto-initialization
UnsupportedAudioError: If audio content is sent to a model that doesn't
support it and unsupported_audio is RAISE (default)
Text runs repair unresolved persisted calls without replaying them.
Cancellation records unknown outcomes, not exactly-once execution:
an interrupted synchronous tool may still finish in its worker thread.
Consumers must serialize runs/compaction for the same session and close
the iterator explicitly when stopping consumption early.
"""
self._check_channel_access()
self._active_runs += 1
try:
with self._execution_bridge.turn(session_id) if self._execution_bridge else nullcontext():
async with aclosing(
self._run_interaction(
user_message, session_id, user_id, config, auto_commit=auto_commit, log_file=log_file
)
) as events:
async for event in events:
yield event
finally:
self._active_runs -= 1
async def _run_interaction(
self,
user_message: str | list[ContentPart] | Message | None,
session_id: str | None,
user_id: str,
config: GenerationConfig | None,
*,
auto_commit: bool | None,
log_file: Path | str | None,
) -> AsyncGenerator[Event, None]:
if user_message is None:
if self.provider.live_config is None:
self._check_extension_mode("voice")
elif (
self.plugins
or (
self.compaction_strategy is not None
and self.compaction_strategy is not self._default_compaction_strategy
)
or self.skill_discoverer
):
raise ValueError("Configure text lifecycle extensions on the Live backend agent")
async with aclosing(self._run_voice(auto_commit=auto_commit, log_file=log_file)) as events:
async for event in events:
yield event
return
if self.batch:
self._check_extension_mode("batch")
await self._ensure_initialized()
# Normalize user_message to a Message object
message = Message(role="user", content=user_message) if isinstance(user_message, str | list) else user_message
message = deepcopy(message)
if self.batch:
message.content = await self._process_multimodal_content(message.content)
async for event in self._run_batch(message, session_id, user_id, config):
yield event
return
if session_id is None:
session_id = f"session-{uuid.uuid4().hex[:12]}"
await self.session.get_or_create_session(session_id, user_id)
context = RunContext(self, session_id, user_id)
plugins = tuple(self.plugins)
text_run: AsyncGenerator[Event, None] | None = None
try:
await self._repair_tool_calls(session_id)
await self.refresh_skills()
# Capture only original incoming user text, never tool output, prior
# history, transcribed audio or plugin-injected instructions.
text = ""
if message.role == "user":
if isinstance(message.content, str):
text = message.content
elif isinstance(message.content, list):
text = "\n".join(part.text for part in message.content if isinstance(part, TextContent))
references = explicit_skill_names(text) if self.skill_discoverer is not None else ()
for plugin in plugins:
message = await plugin.before_run(context, message)
if not isinstance(message, Message):
raise TypeError(f"{type(plugin).__name__}.before_run must return Message")
text_run = self._run_text(message, context, config, plugins, references)
async for event in text_run:
for plugin in plugins:
await plugin.on_event(context, deepcopy(event))
yield event
finally:
async def cleanup() -> None:
errors: list[BaseException] = []
try:
if text_run is not None:
await text_run.aclose()
except BaseException as error:
errors.append(error)
try:
# Re-read durable state, including commits that raced cancellation.
await self._repair_tool_calls(session_id)
except BaseException as error:
errors.append(error)
for plugin in plugins:
try:
await plugin.after_run(context)
except BaseException as error:
error.add_note(f"{type(plugin).__name__}.after_run failed")
errors.append(error)
self._session_compaction_requests.discard(session_id)
if errors:
raise BaseExceptionGroup("Agent run cleanup failed", errors)
await self._finish_cleanup(cleanup())
async def _repair_tool_calls(self, session_id: str) -> None:
"""Close interrupted tool blocks without assuming whether side effects ran."""
pending: dict[str, ToolCall] = {}
for message in await self.session.get_history(session_id):
for call in message.tool_calls:
pending[call.id] = call
if message.role == "tool" and message.tool_call_id in pending:
del pending[message.tool_call_id]
for call in pending.values():
await self.session.add_message(
session_id,
Message(
role="tool",
tool_call_id=call.id,
name=call.name,
content="Error: Tool outcome unknown (run interrupted or cancelled); call was not replayed.",
),
)
async def _run_text(
self,
message: Message,
context: RunContext,
config: GenerationConfig | None,
plugins: tuple[AgentPlugin, ...],
skill_references: tuple[str, ...] = (),
) -> AsyncGenerator[Event, None]:
session_id = context.session_id
skill_context: list[Message] = []
async with aclosing(self._activate_skills(skill_references, context, plugins, skill_context)) as activations:
async for event in activations:
yield event
# Process audio content through STT if present
processed_content = await self._process_multimodal_content(message.content)
# Create processed message
processed_message = Message(
role=message.role,
content=processed_content,
tool_calls=message.tool_calls,
tool_call_id=message.tool_call_id,
name=message.name,
)
# Add user message to history
await self.session.add_message(session_id, processed_message)
# Set session ID for HTTP logging
self.provider.set_session_id(session_id)
session_usage = self._session_usage.setdefault(session_id, TokenUsage())
last_usage = Usage()
for round_num in range(self.max_tool_rounds):
context.round_number = round_num
logger.debug(f"Generation round {round_num + 1}/{self.max_tool_rounds}")
# Get current history
history = await self.session.get_history(session_id)
# Prepend system prompt if set
messages: list[Message] = []
if self.system_prompt:
messages.append(Message(role="system", content=self.system_prompt))
messages.extend(history)
# Filter unsupported audio from all messages (including history)
messages = self._filter_unsupported_audio(messages)
# Get stored prompt_tokens from previous API calls (or estimate if first call)
stored_prompt_tokens = self._session_tokens.get(session_id, 0)
# Estimate tokens for new user message (not yet in history)
new_message_tokens = estimate_messages_tokens([processed_message])
# Total context: stored prompt_tokens + estimated new tokens
# Use stored tokens from last API response (actual) + estimate for new message
if stored_prompt_tokens > 0:
# We have actual tokens from previous response
context_tokens = stored_prompt_tokens + new_message_tokens
tokens_source = "actual"
else:
# First call or after compaction - estimate everything
context_tokens = estimate_messages_tokens(messages)
tokens_source = "estimated"
# Get the model context limit for display.
display_limit = get_model_context_limit(self.provider)
usage_percent = (context_tokens / display_limit * 100) if display_limit > 0 else 0
logger.info(
f"Context usage: {len(messages)} messages, {context_tokens} tokens ({tokens_source}) "
f"({usage_percent:.1f}% of {display_limit})"
)
# Apply context compaction if configured
compaction = (
self._do_compact(messages, session_id)
if self._compaction_requested or session_id in self._session_compaction_requests
else self._maybe_compact(messages, session_id)
)
self._compaction_requested = False
self._session_compaction_requests.discard(session_id)
async for event in compaction:
yield event
# Update messages if compaction was done
if isinstance(event, CompactionDoneEvent):
# Reload compacted history from session
messages = []
if self.system_prompt:
messages.append(Message(role="system", content=self.system_prompt))
history = await self.session.get_history(session_id)
messages.extend(history)
messages = self._filter_unsupported_audio(messages)
# Reset token tracking after compaction (context was summarized)
# The next API call will give us new actual prompt_tokens
self._session_tokens[session_id] = event.summary_tokens
observe("context_after_compaction", messages=messages, session_id=session_id)
# Generate response
pending_tool_calls: list[ToolCall] = []
# Exact streamed event identities for this generation only. A retry
# discards these with its uncommitted calls, even if IDs are reused.
pending_tool_events: list[ToolCallEvent] = []
full_text = ""
has_error = False
fatal_error = False
finish_reason = FinishReason.UNKNOWN
last_usage = Usage()
# Copy schemas, not callable owners (which can hold locks or live clients).
await self.refresh_skills()
# Re-resolve after refresh/compaction so mid-run tool installation
# is visible in the very next model request.
tools = self.tool_registry.get_all() if self.tool_registry.has_tools() else None
if tools and self.save_tool_outputs:
tools = _inject_save_to(tools)
if self.skill_discoverer is not None:
# Ephemeral context is rebuilt after compaction and never stored
# in session history, so removed/renamed catalog entries vanish.
# This is tool-selection metadata, not a fabricated user turn.
# Descriptors are escaped/framed as data; loaded bodies remain
# exclusively in user-level skill_context or tool results.
messages.insert(
1 if self.system_prompt else 0,
Message(role="system", content=render_skill_manifest(self.skills.values())),
)
messages.extend(deepcopy(skill_context))
request = ModelRequest(
deepcopy(messages),
[replace(tool, parameters=deepcopy(tool.parameters)) for tool in tools or []],
deepcopy(config),
)
model_call_id = uuid.uuid4().hex
observation_scope.set(
{
**observation_scope.get(),
"model_call_id": model_call_id,
"round": round_num,
"session_id": context.session_id,
}
)
for plugin in plugins:
request = await plugin.before_model(context, request)
if not isinstance(request, ModelRequest):
raise TypeError(f"{type(plugin).__name__}.before_model must return ModelRequest")
observe(
"model_context",
messages=request.messages,
tools=[
{"name": tool.name, "description": tool.description, "parameters": tool.parameters}
for tool in request.tools
],
config=request.config,
)
async with self._closing_events(
self.provider.generate(
messages=request.messages,
tools=request.tools or None,
config=request.config,
stream=self.streaming,
)
) as events:
async for event in events:
# Providers repeat cumulative snapshots within a generation, not additive deltas.
if event.usage.has_usage():
session_usage.prompt_tokens += event.usage.prompt_tokens - last_usage.prompt_tokens
session_usage.completion_tokens += event.usage.completion_tokens - last_usage.completion_tokens
session_usage.total_tokens += event.usage.total_tokens - last_usage.total_tokens
last_usage = replace(event.usage)
# Context size remains the latest prompt, not accumulated billing usage.
self._session_tokens[session_id] = last_usage.prompt_tokens
self._last_usage[session_id] = TokenUsage(
prompt_tokens=last_usage.prompt_tokens,
completion_tokens=last_usage.completion_tokens,
total_tokens=last_usage.total_tokens,
)
event.usage = replace(last_usage, session=replace(session_usage))
if isinstance(event, TextChunkEvent):
full_text += event.chunk
elif isinstance(event, TextDoneEvent):
full_text = event.text
finish_reason = event.finish_reason
elif isinstance(event, ToolCallEvent):
generation, index = event.extra.get("generation_id"), event.extra.get("index")
if not (
isinstance(generation, str) and generation and type(index) is int and 0 <= index < 1024
):
# Providers without argument previews still need a
# fresh, exact identity for direct approval events.
event.extra = {
**event.extra,
"generation_id": f"agent-{model_call_id}",
"index": len(pending_tool_calls),
}
call = ToolCall(
id=event.id,
name=event.name,
arguments=deepcopy(event.arguments),
metadata=deepcopy(event.metadata),
)
if not call.id or not call.name or any(tc.id == call.id for tc in pending_tool_calls):
raise ValueError("Tool calls require unique nonempty IDs and nonempty names")
for plugin in plugins:
call = await plugin.before_tool(context, call)
if not isinstance(call, ToolCall) or (call.id, call.name) != (event.id, event.name):
raise ValueError(f"{type(plugin).__name__}.before_tool must preserve tool call id/name")
if not self.save_tool_outputs:
call.arguments.pop(_SAVE_TO_PARAM_NAME, None)
pending_tool_calls.append(deepcopy(call))
event.arguments = deepcopy(call.arguments)
event.metadata = deepcopy(call.metadata)
pending_tool_events.append(event)
yield event
if isinstance(event, ErrorEvent):
has_error = True
if not event.recoverable:
if self._execution_bridge:
await self._execution_bridge.abandon(tuple(pending_tool_events))
fatal_error = True
break
for plugin in plugins:
await plugin.after_model(context)
if fatal_error:
yield DoneEvent(
final_text="",
session_id=session_id,
finish_reason=FinishReason.UNKNOWN,
usage=replace(last_usage, session=replace(session_usage)),
)
return
# If we hit an error, don't continue
if has_error:
if self._execution_bridge:
await self._execution_bridge.abandon(tuple(pending_tool_events))
continue
# If we got tool calls, execute them
if pending_tool_calls:
# Add assistant message with tool calls to history
assistant_block = Message(role="assistant", content=full_text or None, tool_calls=pending_tool_calls)
bridge = self._execution_bridge
with bridge.reserve(assistant_block, tuple(pending_tool_events)) if bridge else nullcontext():
await self.session.add_message(session_id, assistant_block)
# Execute each tool
for call_position, tool_call in enumerate(pending_tool_calls):
logger.debug(f"Executing tool: {tool_call.name}")
# _save_to convention: remove from args before tool execution
execution_call = deepcopy(tool_call)
yield ToolExecutionStartedEvent(
id=execution_call.id,
name=execution_call.name,
extra=deepcopy(pending_tool_events[call_position].extra),
usage=replace(last_usage, session=replace(session_usage)),
)
observe("tool_started", call=execution_call)
save_path = _extract_save_path(execution_call) if self.save_tool_outputs else None
# Approval callbacks may publish before queued provider
# events. Give them this exact invocation's display identity.
with (
_tool_invocation_scope(
execution_call.id, execution_call.name, pending_tool_events[call_position].extra
),
bridge.executing(execution_call, call_position) if bridge else nullcontext(),
):
result_event = await self._execute_text_tool(execution_call, context, plugins)
result_event.extra = {
**result_event.extra,
"generation_id": pending_tool_events[call_position].extra["generation_id"],
"index": pending_tool_events[call_position].extra["index"],
}
observe("tool_finished", result=result_event)
# Attach last known usage info to tool result events
result_event.usage = replace(last_usage, session=replace(session_usage))
# Add tool result to history
if self.save_tool_outputs and save_path and result_event.error is None:
result_content = _save_and_return(result_event, save_path, session_id)
else:
result_content = (
_serialize_tool_result(result_event.result)
if result_event.error is None
else f"Error: {result_event.error}"
)
await self.session.add_message(
session_id,
Message(
role="tool",
content=result_content,
tool_call_id=tool_call.id,
name=tool_call.name,
),
)
yield result_event
# Persist the error before either UI interruption or fail-fast exit.
if result_event.error and self.fail_on_invalid_tool and "does not exist" in result_event.error:
logger.error(f"Tool hallucination detected: {tool_call.name}")
raise ToolHallucinationError(
tool_name=tool_call.name,
available_tools=self.tool_registry.names(),
message=result_event.error,
)
for plugin in plugins:
await plugin.after_tools(context)
# Continue to next round (model will see tool results)
continue
# No tool calls - we're done
if full_text:
await self.session.add_message(session_id, Message(role="assistant", content=full_text))
yield DoneEvent(
final_text=full_text,
session_id=session_id,
finish_reason=finish_reason,
usage=replace(last_usage, session=replace(session_usage)),
)
return
# Hit max rounds
logger.warning(f"Max tool rounds ({self.max_tool_rounds}) exceeded")
yield ErrorEvent(
message=f"Max tool rounds ({self.max_tool_rounds}) exceeded",
code="MAX_TOOL_ROUNDS",
recoverable=False,
extra={"max_tool_rounds": self.max_tool_rounds}
if type(self.max_tool_rounds) is int and self.max_tool_rounds > 0
else {},
usage=replace(last_usage, session=replace(session_usage)),
)
yield DoneEvent(
final_text="",
session_id=session_id,
finish_reason=FinishReason.UNKNOWN,
usage=replace(last_usage, session=replace(session_usage)),
)
async def _run_batch(
self,
user_message: Message,
session_id: str | None,
user_id: str | None,
config: GenerationConfig | None,
) -> AsyncIterator[Event]:
"""Run with batch processing for cost-effective large-scale operations.
This method submits the request as a batch job, waits for completion,
and yields events from the results. Supports multi-turn tool execution
by submitting subsequent batches with tool results.
"""
self._check_extension_mode("batch")
if not self._batch_client:
yield ErrorEvent(message="Batch client not initialized")
return
# Handle session ID
if session_id is None:
session_id = f"batch-session-{uuid.uuid4().hex[:12]}"
await self.session.get_or_create_session(session_id, user_id or "default")
logger.info(f"Created new batch session: {session_id}")
else:
if not await self.session.session_exists(session_id):
await self.session.get_or_create_session(session_id, user_id or "default")
# Add user message to history
await self.session.add_message(session_id, user_message)
# Set session ID for HTTP logging in batch client
if self._batch_client:
self._batch_client.set_session_id(session_id)
# Track accumulated usage and final text
total_usage = Usage()
full_text = ""
finish_reason = FinishReason.UNKNOWN
# Loop until no more tool calls
iteration = 0
while True:
iteration += 1
logger.info(f"Batch iteration {iteration}")
# Get current history for the batch request
history = await self.session.get_history(session_id)
# Build messages for batch request
messages: list[Message] = []
if self.system_prompt:
messages.append(Message(role="system", content=self.system_prompt))
messages.extend(history)
# Create batch request
batch_request = BatchRequest(
custom_id=f"{session_id}-{uuid.uuid4().hex[:8]}",
messages=messages,
config=config,
)
# Submit batch
tools = self.tool_registry.get_all() if self.tool_registry.has_tools() else None
if tools:
tools = _inject_save_to(tools)
logger.info(f"Submitting batch request: {batch_request.custom_id}")
try:
job = await self._batch_client.create_batch(
requests=[batch_request],
config=BatchConfig(),
tools=tools,
generation_config=config,
)
logger.info(f"Batch job created: {job.id}, status: {job.status}")
except HTTPError as e:
error_msg = f"Failed to create batch: {e}"
if e.body:
error_msg += f"\nResponse: {e.body}"
logger.error(error_msg)
yield ErrorEvent(message=error_msg)
return
except Exception as e:
error_msg = f"Failed to create batch: {e}"
logger.error(error_msg)
yield ErrorEvent(message=error_msg)
return
# Wait for completion
logger.info("Waiting for batch completion...")
job = await self._batch_client.wait_for_completion(
job.id,
poll_interval=self.batch_poll_interval,
)
logger.info(f"Batch completed: {job.status}")
# Check status
if job.status != BatchStatus.COMPLETED:
yield ErrorEvent(
message=f"Batch job failed with status: {job.status}",
code=str(job.status.value),
)
return
# Get results
iteration_text = ""
iteration_usage = Usage()
iteration_finish_reason = FinishReason.UNKNOWN
tool_calls_to_execute: list[ToolCall] = []
async for result in self._batch_client.get_results(job):
if result.result_type == "succeeded":
iteration_finish_reason = result.finish_reason
# Parse detailed usage
if result.usage:
iteration_usage = Usage(
prompt_tokens=result.usage.get("prompt_tokens", 0),
completion_tokens=result.usage.get("completion_tokens", 0),
total_tokens=result.usage.get("total_tokens", 0),
cached_tokens=result.usage.get("cached_tokens", 0),
audio_tokens=result.usage.get("audio_tokens", 0),
reasoning_tokens=result.usage.get("reasoning_tokens", 0),
)
# Accumulate usage
total_usage.prompt_tokens += iteration_usage.prompt_tokens
total_usage.completion_tokens += iteration_usage.completion_tokens
total_usage.total_tokens += iteration_usage.total_tokens
total_usage.cached_tokens += iteration_usage.cached_tokens
total_usage.audio_tokens += iteration_usage.audio_tokens
total_usage.reasoning_tokens += iteration_usage.reasoning_tokens
# Handle tool calls
if result.tool_calls:
for tc in result.tool_calls:
import json as json_mod
from .types import ToolArguments
args_str = tc.get("function", {}).get("arguments", "{}")
try:
args: ToolArguments = (
json_mod.loads(args_str) if isinstance(args_str, str) else args_str
)
except json_mod.JSONDecodeError:
args = {"raw": args_str}
tool_call = ToolCall(
id=tc.get("id", ""),
name=tc.get("function", {}).get("name", ""),
arguments=args,
)
tool_calls_to_execute.append(tool_call)
yield ToolCallEvent(
id=tool_call.id,
name=tool_call.name,
arguments=tool_call.arguments,
usage=iteration_usage,
finish_reason=FinishReason.TOOL_CALLS,
metadata=tool_call.metadata,
)
# Handle text content
if result.content:
iteration_text = result.content
yield TextDoneEvent(
text=result.content,
usage=iteration_usage,
finish_reason=iteration_finish_reason,
)
elif result.result_type == "errored":
yield ErrorEvent(
message=result.error_message or "Unknown batch error",
code=result.error_type,
)
return
# Check if we have tool calls to execute
if tool_calls_to_execute:
# Execute tools
for tool_call in tool_calls_to_execute:
# _save_to convention: remove from args before tool execution
save_path = _extract_save_path(tool_call)
yield ToolExecutionStartedEvent(id=tool_call.id, name=tool_call.name)
with _tool_invocation_scope(tool_call.id, tool_call.name):
result_event = await self.tool_executor.execute(tool_call)
yield result_event
# Add tool result to history
if save_path and result_event.result is not None:
result_content = _save_and_return(result_event, save_path, session_id)
else:
result_content = (
_serialize_tool_result(result_event.result)
if result_event.result is not None
else (result_event.error or "Error")
)
await self.session.add_message(
session_id,
Message(
role="tool",
content=result_content,
tool_call_id=tool_call.id,
name=tool_call.name,
),
)
# Continue the loop - will submit new batch with tool results
logger.info(
f"Batch iteration {iteration} completed with {len(tool_calls_to_execute)} tool calls, continuing..."
)
continue
# No tool calls - we're done
full_text = iteration_text
finish_reason = iteration_finish_reason
break
# Add final assistant message to history
if full_text:
await self.session.add_message(session_id, Message(role="assistant", content=full_text))
yield DoneEvent(
final_text=full_text,
session_id=session_id,
usage=total_usage,
finish_reason=finish_reason,
)
async def run_simple(
self,
messages: list[Message],
config: GenerationConfig | None = None,
) -> AsyncGenerator[Event, None]:
"""
Run a simple generation without session management.
Auto-initializes the agent on first call.
Useful for one-off generations or when managing history externally.
Does NOT automatically execute tools.
Args:
messages: List of messages to send
config: Optional generation configuration
Yields:
Events as they occur
"""
self._check_channel_access()
self._active_runs += 1
try:
async with aclosing(self._run_simple(messages, config)) as events:
async for event in events:
yield event
finally:
self._active_runs -= 1
async def _run_simple(
self, messages: list[Message], config: GenerationConfig | None
) -> AsyncGenerator[Event, None]:
self._check_extension_mode("simple")
await self._ensure_initialized()
tools = self.tool_registry.get_all() if self.tool_registry.has_tools() else None
if tools:
tools = _inject_save_to(tools)
# For simple runs, session_usage equals usage (single generation)
session_usage = TokenUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
# Track last known usage to use for events without usage data
last_usage = Usage()
async with self._closing_events(
self.provider.generate(messages=messages, tools=tools, config=config, stream=self.streaming)
) as events:
async for event in events:
# Track usage from events that have actual token counts
if event.usage.has_usage():
# For simple runs, session_usage = current usage
session_usage = TokenUsage(
prompt_tokens=event.usage.prompt_tokens,
completion_tokens=event.usage.completion_tokens,
total_tokens=event.usage.total_tokens,
)
last_usage = Usage(
prompt_tokens=event.usage.prompt_tokens,
completion_tokens=event.usage.completion_tokens,
total_tokens=event.usage.total_tokens,
)
else:
# Use last known usage for events without usage data
event.usage = Usage(
prompt_tokens=last_usage.prompt_tokens,
completion_tokens=last_usage.completion_tokens,
total_tokens=last_usage.total_tokens,
)
# Always attach session usage to the event
event.usage.session = session_usage
yield event
async def clear_session(self, session_id: str) -> None:
"""
Clear a session's conversation history.
Args:
session_id: Session identifier to clear
"""
self._check_channel_access(allow_cleanup=True)
self._active_runs += 1
try:
await self._ensure_initialized()
await self.session.clear_session(session_id)
finally:
self._active_runs -= 1
async def close(self) -> None:
"""Join one resource teardown, even if callers cancel or close concurrently.
New work is rejected until teardown finishes. Repeated closes reuse its
result until provider replacement, initialization, listening, or manual
compaction reopens the agent. Owned run cleanup may clear history, but
cannot start a new model run during shutdown.
"""
listener = self._channel_listener_task
task = asyncio.current_task()
if task is self._close_task:
raise RuntimeError("Cannot close the Agent from its own resource cleanup")
if listener is not None and (task in (listener, self._channel_execution_task) or task in self._cleanup_tasks):
raise RuntimeError("Stop a listener by cancelling it or closing the Agent from another task")
if self._close_task is not None and self._close_task.done() and self.provider is not self._close_provider:
# Harness login/logout can replace a closed provider without a run
# or Agent.initialize(). Its resources still need their own teardown.
self._close_task = None
if self._close_task is None:
self._closing = True
self._close_provider = self.provider
self._close_task = asyncio.create_task(self._close_resources(self._close_provider))
await _join_cleanup(self._close_task)
async def _close_resources(self, provider: Provider) -> None:
errors: list[BaseException] = []
try:
listener = self._channel_listener_task
try:
if listener is not None:
await self._stop_channel_listener(listener)
except BaseException as error:
errors.append(error)
try:
await provider.close()
except BaseException as error:
errors.append(error)
if self._batch_client:
try:
await self._batch_client.close()
except BaseException as error:
errors.append(error)
if len(errors) == 1:
raise errors[0]
if errors:
raise BaseExceptionGroup("Agent resource cleanup failed", errors)
finally:
self._initialized = False
self._closing = False
async def _stop_channel_listener(self, listener: asyncio.Task[None]) -> None:
if not listener.done():
listener.cancel()
cancelled = False
while not listener.done():
try:
await asyncio.shield(listener)
except asyncio.CancelledError:
task = asyncio.current_task()
cancelled = cancelled or (task is not None and task.cancelling() > 0)
if not listener.cancelled():
listener.result()
if cancelled:
raise asyncio.CancelledError