Clients

Client for Interactions with LLMs (Language Models) This module provides a client interface for interacting with various LLMs. It includes functionality for sending requests, receiving responses, and handling errors.

AbstractClient

AbstractClient(preset: Optional[str] = None, tools: Optional[List[Union[str, AbstractTool]]] = None, use_tools: bool = False, debug: bool = True, tool_manager: Optional[ToolManager] = None, **kwargs)

Bases: EventEmitterMixin, ABC

Abstract base Class for LLM models.

Source code in packages/ai-parrot/src/parrot/clients/base.py

def __init__(
    self,
    preset: Optional[str] = None,
    tools: Optional[List[Union[str, AbstractTool]]] = None,
    use_tools: bool = False,
    debug: bool = True,
    tool_manager: Optional[ToolManager] = None,
    **kwargs,
):
    self.__name__ = self.__class__.__name__
    self.model: str = kwargs.get("model", None)
    # Per-loop client cache: keyed by id(asyncio.get_running_loop()).
    # Each entry holds the SDK client instance and a weakref to the loop.
    self._clients_by_loop: dict[int, _LoopClientEntry] = {}
    # Per-loop locks: asyncio.Lock() is loop-bound, so one per loop.
    self._locks_by_loop: dict[int, asyncio.Lock] = {}
    self.session: Optional[aiohttp.ClientSession] = None
    self.use_session: bool = kwargs.get("use_session", self.use_session)
    if preset:
        preset_config = LLM_PRESETS.get(preset, LLM_PRESETS["default"])
        # define temp, top_k, top_p, max_tokens from selected preset:
        self.temperature = preset_config.get("temperature", 0.4)
        self.top_k = preset_config.get("top_k", 30)
        self.top_p = preset_config.get("top_p", 0.2)
        self.max_tokens = preset_config.get("max_tokens")
    else:
        # define default values from preset default:
        self.temperature = kwargs.get("temperature", 0)
        self.top_k = kwargs.get("top_k", 30)
        self.top_p = kwargs.get("top_p", 0.2)
        # ``None`` means "not configured" — the per-client
        # _default_max_tokens / _invoke_max_tokens defaults take over. Do
        # NOT reintroduce a literal here: a framework-wide default assigned
        # eagerly is indistinguishable from a deliberate caller choice, and
        # would shadow every per-client default downstream.
        self.max_tokens = kwargs.get("max_tokens")
    # Whether max_tokens was EXPLICITLY configured by the caller (directly
    # or via a preset) rather than left unset above. invoke() only honours
    # an explicit value, so that a client's _invoke_max_tokens default is
    # never shadowed by ask()'s generic one.
    self._max_tokens_configured: bool = "max_tokens" in kwargs or preset is not None
    # Per-instance override for invoke()'s output-token budget specifically.
    # ``None`` means "fall back to an explicit self.max_tokens, then the
    # class default" (see _resolve_max_tokens).
    self.invoke_max_tokens: Optional[int] = kwargs.get("invoke_max_tokens", None)
    # FEAT-550 §2.1: constructor-level question budget defaults. Immutable
    # record, never a mutable counter on the shared client instance.
    _tb = kwargs.get("token_budget", None)
    self._budget_defaults_value: BudgetDefaults = BudgetDefaults(
        token_budget=_tb,
        budget_mode=kwargs.get("budget_mode", "estimated"),
        final_answer_reserve=kwargs.get("final_answer_reserve", 0.15),
        registry=kwargs.get("budget_registry", None),
    )
    self._budget_defaults_active: bool = _tb is not None
    self._budget_registry = self._budget_defaults_value.registry
    if _tb is not None:
        # Validate eagerly by constructing TokenBudgetPolicy(...) so a bad constructor value fails at
        # construction, not at first call — bounded by spec §2.1 "Reject booleans, strings, negative…".
        from parrot.models.token_budget import TokenBudgetPolicy

        TokenBudgetPolicy(
            token_budget=_tb,
            budget_mode=self._budget_defaults_value.budget_mode,
            final_answer_reserve=self._budget_defaults_value.final_answer_reserve,
        )
    self.base_headers.update(kwargs.get("headers", {}))
    self.api_key = kwargs.get("api_key", None)
    self.version = kwargs.get("version", self.version)
    self._config = config
    self.logger: logging.Logger = logging.getLogger(self.__name__)
    self._json: Any = JSONContent()
    self.client_type: str = kwargs.get("client_type", self.client_type)
    self._debug: bool = debug
    self._program: str = kwargs.get("program", "parrot")  # Default program slug
    # Use provided tool_manager or create a new one
    # This allows Agent to pass its tool_manager as a reference
    if tool_manager is not None:
        self._tool_manager = tool_manager
    else:
        self._tool_manager = ToolManager(logger=self.logger, debug=self._debug)
    self.tools: Dict[str, Union[ToolDefinition, AbstractTool]] = {}
    self.enable_tools: bool = use_tools
    # FEAT-438 G5: only create an instance attribute when the caller
    # explicitly passed fallback_model= (including explicit None) —
    # an unconditional assignment here would shadow a subclass's
    # class-level _fallback_model default on every instance (the
    # class attribute declared above, :~272, is the fallback when no
    # instance attribute is created).
    if "fallback_model" in kwargs:
        self._fallback_model = kwargs.pop("fallback_model")
    # Initialize tools if provided
    if use_tools and tools:
        self._tool_manager.default_tools(tools)
        self.enable_tools = True
    # FEAT-176: initialise per-instance lifecycle event registry.
    # forward_to_global=False so client events stay isolated unless the
    # owning bot wires them up via a parent trace.
    self._init_events(forward_to_global=False)

tool_manager property writable

tool_manager: ToolManager

Get the tool manager.

client property writable

client: Optional[Any]

Return the SDK client bound to the current event loop, or None.

The cache key is id(asyncio.get_running_loop()). Returns None when called outside a running loop (e.g. in __init__ or sync code).

RETURNS	DESCRIPTION
Optional[Any]	The loop-local SDK client instance, or None if no client has
Optional[Any]	been created yet for the current loop.

default_model property

default_model: str

Return the default model for the client.

has_server_credentials classmethod

has_server_credentials() -> bool

Whether the server itself holds a usable credential for this provider (any credential_env resolves).

Source code in packages/ai-parrot/src/parrot/clients/base.py

@classmethod
def has_server_credentials(cls) -> bool:
    """Whether the server itself holds a usable credential for this provider (any ``credential_env`` resolves)."""
    return any(config.get(name) for name in (cls.credential_env or ()))

get_client abstractmethod async

get_client() -> Any

Return the client instance.

Source code in packages/ai-parrot/src/parrot/clients/base.py

@abstractmethod
async def get_client(self) -> Any:
    """Return the client instance."""
    raise NotImplementedError

complete async

complete(prompt: str, *, model: Optional[str] = None, system_prompt: Optional[str] = None, max_tokens: Optional[int] = None, temperature: Optional[float] = None) -> str

Send a prompt, return the model's textual reply as a plain string.

Thin convenience wrapper around ask() for single-shot, tool-less text generation. Use cases:

• Interop with components that expect async complete(prompt) -> str (e.g. parrot_tools.scraping.PlanGenerator).
• Quick prompts where you don't need the full AIMessage.

Handles two things ask() does NOT:

1. Auto-enters the async context manager if the client isn't already initialized, so callers don't need async with client: before calling complete(). If the client was already entered (e.g. inside a surrounding async with), we reuse it and don't tear it down.
2. Extracts plain text from the response — ask() returns an AIMessage pydantic model (provider-specific clients) or a MessageResponse TypedDict with a content-block list. Both shapes collapse to a single string here.

PARAMETER	DESCRIPTION
prompt	User prompt. TYPE: str
model	Override the instance default model. TYPE: Optional[str] DEFAULT: None
system_prompt	Optional system prompt. TYPE: Optional[str] DEFAULT: None
max_tokens	Override default max tokens. TYPE: Optional[int] DEFAULT: None
temperature	Override default sampling temperature. TYPE: Optional[float] DEFAULT: None

RETURNS	DESCRIPTION
str	The model's textual response.

RAISES	DESCRIPTION
RuntimeError	If the response has no extractable text.

Source code in packages/ai-parrot/src/parrot/clients/base.py

async def complete(
    self,
    prompt: str,
    *,
    model: Optional[str] = None,
    system_prompt: Optional[str] = None,
    max_tokens: Optional[int] = None,
    temperature: Optional[float] = None,
) -> str:
    """Send a prompt, return the model's textual reply as a plain string.

    Thin convenience wrapper around ``ask()`` for single-shot,
    tool-less text generation. Use cases:

    - Interop with components that expect ``async complete(prompt) -> str``
      (e.g. ``parrot_tools.scraping.PlanGenerator``).
    - Quick prompts where you don't need the full ``AIMessage``.

    Handles two things ``ask()`` does NOT:

    1. **Auto-enters the async context manager** if the client isn't
       already initialized, so callers don't need ``async with client:``
       before calling ``complete()``. If the client was already entered
       (e.g. inside a surrounding ``async with``), we reuse it and
       don't tear it down.
    2. **Extracts plain text** from the response — ``ask()`` returns
       an ``AIMessage`` pydantic model (provider-specific clients) or
       a ``MessageResponse`` TypedDict with a content-block list.
       Both shapes collapse to a single string here.

    Args:
        prompt: User prompt.
        model: Override the instance default model.
        system_prompt: Optional system prompt.
        max_tokens: Override default max tokens.
        temperature: Override default sampling temperature.

    Returns:
        The model's textual response.

    Raises:
        RuntimeError: If the response has no extractable text.
    """
    need_enter = self.client is None
    if need_enter:
        await self.__aenter__()
    try:
        resolved_model = (
            model or self.model or getattr(self, "default_model", None) or getattr(self, "_default_model", None)
        )
        kwargs: Dict[str, Any] = {"prompt": prompt, "model": resolved_model}
        if system_prompt is not None:
            kwargs["system_prompt"] = system_prompt
        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens
        if temperature is not None:
            kwargs["temperature"] = temperature
        response = await self.ask(**kwargs)
    finally:
        if need_enter:
            await self.__aexit__(None, None, None)

    text = self._extract_text(response)
    if not text:
        raise RuntimeError(f"LLM returned no extractable text " f"(response type: {type(response).__name__})")
    return text

close async

close() -> None

Close all per-loop SDK clients.

Delegates to close_all(). Both close and close_all exist so callers can be explicit about intent (spec §8 Q1).

Source code in packages/ai-parrot/src/parrot/clients/base.py

async def close(self) -> None:
    """Close all per-loop SDK clients.

    Delegates to ``close_all()``.  Both ``close`` and ``close_all`` exist
    so callers can be explicit about intent (spec §8 Q1).
    """
    await self.close_all()

close_all async

close_all() -> None

Tear down every per-loop SDK client entry.

Safely handles dead / foreign loops: entries whose loop has been garbage-collected or belongs to a different running context are dropped without awaiting their close() coroutine.

After this call, _clients_by_loop and _locks_by_loop are empty.

Source code in packages/ai-parrot/src/parrot/clients/base.py

async def close_all(self) -> None:
    """Tear down every per-loop SDK client entry.

    Safely handles dead / foreign loops: entries whose loop has been
    garbage-collected or belongs to a different running context are dropped
    without awaiting their ``close()`` coroutine.

    After this call, ``_clients_by_loop`` and ``_locks_by_loop`` are empty.
    """
    current = self._get_current_loop()
    current_id = id(current) if current is not None else None
    for loop_id, entry in list(self._clients_by_loop.items()):
        await self._safe_close_entry(entry, is_current_loop=(loop_id == current_id))
    self._clients_by_loop.clear()
    self._locks_by_loop.clear()

set_program

set_program(program_slug: str) -> None

Set the program slug for the client.

Source code in packages/ai-parrot/src/parrot/clients/base.py

def set_program(self, program_slug: str) -> None:
    """Set the program slug for the client."""
    self._program = program_slug

set_tools

set_tools(tools: List[Union[str, AbstractTool]]) -> None

Set complete list of tools, replacing existing.

Source code in packages/ai-parrot/src/parrot/clients/base.py

def set_tools(self, tools: List[Union[str, AbstractTool]]) -> None:
    """Set complete list of tools, replacing existing."""
    self.tool_manager.clear_tools()
    self.tools.clear()
    self.register_tools(tools)

get_tool

get_tool(name: str) -> Optional[AbstractTool]

Get a tool by name from ToolManager or legacy tools.

Source code in packages/ai-parrot/src/parrot/clients/base.py

def get_tool(self, name: str) -> Optional[AbstractTool]:
    """Get a tool by name from ToolManager or legacy tools."""
    # Try ToolManager first
    if tool := self.tool_manager.get_tool(name):
        return tool

    # Fall back to legacy tools
    legacy_tool = self.tools.get(name)
    return legacy_tool if isinstance(legacy_tool, AbstractTool) else None

register_tool

register_tool(tool: Union[ToolDefinition, AbstractTool] = None, name: str = None, description: str = None, input_schema: Dict[str, Any] = None, function: Callable = None) -> None

Register a Python function as a tool for LLM to call.

Source code in packages/ai-parrot/src/parrot/clients/base.py

def register_tool(
    self,
    tool: Union[ToolDefinition, AbstractTool] = None,
    name: str = None,
    description: str = None,
    input_schema: Dict[str, Any] = None,
    function: Callable = None,
) -> None:
    """Register a Python function as a tool for LLM to call."""
    self.tool_manager.register_tool(
        tool=tool, name=name, description=description, input_schema=input_schema, function=function
    )

register_tools

register_tools(tools: List[Union[ToolDefinition, AbstractTool]]) -> None

Register multiple tools at once.

Source code in packages/ai-parrot/src/parrot/clients/base.py

def register_tools(self, tools: List[Union[ToolDefinition, AbstractTool]]) -> None:
    """Register multiple tools at once."""
    self.tool_manager.register_tools(tools)
    self.enable_tools = True

register_python_tool

register_python_tool(report_dir: Optional[Path] = None) -> PythonREPLTool

Register Python REPL tool with a ClaudeAPIClient.

PARAMETER	DESCRIPTION
client	The ClaudeAPIClient instance
report_dir	Directory for saving reports TYPE: Optional[Path] DEFAULT: None

RETURNS	DESCRIPTION
PythonREPLTool	The PythonREPLTool instance

Source code in packages/ai-parrot/src/parrot/clients/base.py

def register_python_tool(
    self,
    report_dir: Optional[Path] = None,
) -> PythonREPLTool:
    """Register Python REPL tool with a ClaudeAPIClient.

    Args:
        client: The ClaudeAPIClient instance
        report_dir: Directory for saving reports

    Returns:
        The PythonREPLTool instance
    """
    if "python_repl" in self.tools:
        return self.tools["python_repl"]

    from ..tools.pythonrepl import PythonREPLTool

    tool = PythonREPLTool(
        report_dir=report_dir,
        debug=self._debug,
    )
    self.tool_manager.add_tool(tool)
    return tool

list_tools

list_tools() -> List[str]

Get a list of all registered tool names.

Source code in packages/ai-parrot/src/parrot/clients/base.py

def list_tools(self) -> List[str]:
    """Get a list of all registered tool names."""
    tool_names = self.tool_manager.list_tools()
    legacy_names = list(self.tools.keys())
    return tool_names + [name for name in legacy_names if name not in tool_names]

remove_tool

remove_tool(name: str) -> bool

Remove a tool by name.

PARAMETER	DESCRIPTION
name	Tool name to remove TYPE: str

RETURNS	DESCRIPTION
bool	True if tool was removed, False if not found

Source code in packages/ai-parrot/src/parrot/clients/base.py

def remove_tool(self, name: str) -> bool:
    """
    Remove a tool by name.

    Args:
        name: Tool name to remove

    Returns:
        True if tool was removed, False if not found
    """
    self.tool_manager.remove_tool(name)

clear_tools

clear_tools() -> None

Clear all registered tools.

Source code in packages/ai-parrot/src/parrot/clients/base.py

def clear_tools(self) -> None:
    """Clear all registered tools."""
    self.tool_manager.clear_tools()
    self.tools.clear()
    self.logger.info("Cleared all tools")

ask abstractmethod async

ask(prompt: str, model: str, max_tokens: Optional[int] = None, temperature: float = 0.7, files: Optional[List[Union[str, Path]]] = None, system_prompt: Optional[str] = None, history: Optional[Sequence[HistoryMessage]] = None, structured_output: Union[type, StructuredOutputConfig, None] = None, tools: Optional[List[Dict[str, Any]]] = None, use_tools: Optional[bool] = None, deep_research: bool = False, background: bool = False, lazy_loading: bool = False) -> MessageResponse

Send a prompt to the model and return the response.

PARAMETER	DESCRIPTION
prompt	The input prompt for the model TYPE: str
model	The model to use TYPE: str
max_tokens	Maximum number of tokens in the response. None (the default) resolves via :meth:_resolve_max_tokens — the per-instance max_tokens, then the client's :attr:_default_max_tokens. TYPE: Optional[int] DEFAULT: None
temperature	Sampling temperature for response generation TYPE: float DEFAULT: 0.7
files	Optional files to include in the request TYPE: Optional[List[Union[str, Path]]] DEFAULT: None
system_prompt	Optional system prompt to guide the model TYPE: Optional[str] DEFAULT: None
history	Already-rendered conversation history from the owning bot (FEAT-524). The client only formats it for its provider — it never loads or persists history itself. TYPE: Optional[Sequence[HistoryMessage]] DEFAULT: None
structured_output	Optional structured output configuration TYPE: Union[type, StructuredOutputConfig, None] DEFAULT: None
tools	Optional tools to register for this call TYPE: Optional[List[Dict[str, Any]]] DEFAULT: None
use_tools	Whether to use tools TYPE: Optional[bool] DEFAULT: None
deep_research	If True, use deep research mode (provider-specific) TYPE: bool DEFAULT: False
background	If True, execute research in background (async mode) TYPE: bool DEFAULT: False
lazy_loading	If True, enabled dynamic tool searching TYPE: bool DEFAULT: False

Source code in packages/ai-parrot/src/parrot/clients/base.py

@abstractmethod
async def ask(
    self,
    prompt: str,
    model: str,
    max_tokens: Optional[int] = None,
    temperature: float = 0.7,
    files: Optional[List[Union[str, Path]]] = None,
    system_prompt: Optional[str] = None,
    history: Optional[Sequence[HistoryMessage]] = None,
    structured_output: Union[type, StructuredOutputConfig, None] = None,
    tools: Optional[List[Dict[str, Any]]] = None,
    use_tools: Optional[bool] = None,
    deep_research: bool = False,
    background: bool = False,
    lazy_loading: bool = False,
) -> MessageResponse:
    """Send a prompt to the model and return the response.

    Args:
        prompt: The input prompt for the model
        model: The model to use
        max_tokens: Maximum number of tokens in the response. ``None`` (the
            default) resolves via :meth:`_resolve_max_tokens` — the
            per-instance ``max_tokens``, then the client's
            :attr:`_default_max_tokens`.
        temperature: Sampling temperature for response generation
        files: Optional files to include in the request
        system_prompt: Optional system prompt to guide the model
        history: Already-rendered conversation history from the owning bot
            (FEAT-524). The client only formats it for its provider — it
            never loads or persists history itself.
        structured_output: Optional structured output configuration
        tools: Optional tools to register for this call
        use_tools: Whether to use tools
        deep_research: If True, use deep research mode (provider-specific)
        background: If True, execute research in background (async mode)
        lazy_loading: If True, enabled dynamic tool searching
    """
    raise NotImplementedError("Subclasses must implement this method.")

ask_stream abstractmethod async

ask_stream(prompt: str, model: str = None, max_tokens: Optional[int] = None, temperature: float = 0.7, files: Optional[List[Union[str, Path]]] = None, system_prompt: Optional[str] = None, history: Optional[Sequence[HistoryMessage]] = None, tools: Optional[List[Dict[str, Any]]] = None, deep_research: bool = False, agent_config: Optional[Dict[str, Any]] = None, lazy_loading: bool = False) -> AsyncIterator[Union[str, AIMessage]]

Stream the model's response.

Yields successive string chunks of the model response followed by a single final :class:~parrot.models.responses.AIMessage carrying full response metadata (token usage, stop reason, model, provider, turn_id, etc.).

Implementors MUST yield at least one str chunk before the final AIMessage. Consumers can detect the end-of-stream sentinel via isinstance(chunk, AIMessage).

Source code in packages/ai-parrot/src/parrot/clients/base.py

@abstractmethod
async def ask_stream(
    self,
    prompt: str,
    model: str = None,
    max_tokens: Optional[int] = None,
    temperature: float = 0.7,
    files: Optional[List[Union[str, Path]]] = None,
    system_prompt: Optional[str] = None,
    history: Optional[Sequence[HistoryMessage]] = None,
    tools: Optional[List[Dict[str, Any]]] = None,
    deep_research: bool = False,
    agent_config: Optional[Dict[str, Any]] = None,
    lazy_loading: bool = False,
) -> AsyncIterator[Union[str, AIMessage]]:
    """Stream the model's response.

    Yields successive string chunks of the model response followed by a
    single final :class:`~parrot.models.responses.AIMessage` carrying full
    response metadata (token usage, stop reason, model, provider, turn_id,
    etc.).

    Implementors MUST yield at least one ``str`` chunk before the final
    ``AIMessage``.  Consumers can detect the end-of-stream sentinel via
    ``isinstance(chunk, AIMessage)``.
    """
    raise NotImplementedError("Subclasses must implement this method.")

resume abstractmethod async

resume(session_id: str, user_input: str, state: Dict[str, Any]) -> MessageResponse

Resume a suspended model execution.

PARAMETER	DESCRIPTION
session_id	The session ID TYPE: str
user_input	The user's input to inject as tool result TYPE: str
state	The suspended state containing messages and tool_call_id TYPE: Dict[str, Any]

RETURNS	DESCRIPTION
MessageResponse	The response from the LLM TYPE: MessageResponse

Source code in packages/ai-parrot/src/parrot/clients/base.py

@abstractmethod
async def resume(self, session_id: str, user_input: str, state: Dict[str, Any]) -> MessageResponse:
    """Resume a suspended model execution.

    Args:
        session_id: The session ID
        user_input: The user's input to inject as tool result
        state: The suspended state containing messages and tool_call_id

    Returns:
        MessageResponse: The response from the LLM
    """
    raise NotImplementedError("Subclasses must implement this method.")

batch_ask async

batch_ask(requests: List[Any]) -> List[Any]

Process multiple requests in batch.

Source code in packages/ai-parrot/src/parrot/clients/base.py

async def batch_ask(self, requests: List[Any]) -> List[Any]:
    """Process multiple requests in batch."""
    raise NotImplementedError("Subclasses must implement batch processing.")

invoke abstractmethod async

invoke(prompt: str, *, output_type: Optional[type] = None, structured_output: Optional[StructuredOutputConfig] = None, model: Optional[str] = None, system_prompt: Optional[str] = None, max_tokens: Optional[int] = None, temperature: float = 0.0, use_tools: bool = False, tools: Optional[list] = None) -> InvokeResult

Lightweight stateless invocation — no retry, no history, no prompt builder.

Each concrete client implements this method using provider-native structured output. Use this instead of ask() when you need fast, stateless structured extraction without conversation history overhead.

PARAMETER	DESCRIPTION
prompt	The user prompt to send. TYPE: str
output_type	A Pydantic model or dataclass class to parse the response into. Mutually exclusive with structured_output (structured_output wins). TYPE: Optional[type] DEFAULT: None
structured_output	Full :class:StructuredOutputConfig, including optional custom_parser. Takes precedence over output_type. TYPE: Optional[StructuredOutputConfig] DEFAULT: None
model	Override the model for this call. Falls back to an explicitly selected self.model, then _lightweight_model, then _default_model (see :meth:_resolve_invoke_model). TYPE: Optional[str] DEFAULT: None
system_prompt	Override the system prompt. Falls back to BASIC_SYSTEM_PROMPT rendered with instance attributes. TYPE: Optional[str] DEFAULT: None
max_tokens	Maximum completion tokens. None (the default) resolves via :meth:_resolve_invoke_max_tokens — per-instance invoke_max_tokens, then max_tokens, then the client's :attr:_invoke_max_tokens class default. Pass an explicit value only to cap a single call. TYPE: Optional[int] DEFAULT: None
temperature	Sampling temperature (default 0.0 for deterministic output). TYPE: float DEFAULT: 0.0
use_tools	If True, inject registered tools into the request. TYPE: bool DEFAULT: False
tools	Additional tool definitions to pass directly to the provider. TYPE: Optional[list] DEFAULT: None

RETURNS	DESCRIPTION
InvokeResult	class:InvokeResult with the parsed output, model, usage,
InvokeResult	and optional raw_response.

RAISES	DESCRIPTION
	class:InvokeError: If the provider call fails for any reason.

Source code in packages/ai-parrot/src/parrot/clients/base.py

@abstractmethod
async def invoke(
    self,
    prompt: str,
    *,
    output_type: Optional[type] = None,
    structured_output: Optional[StructuredOutputConfig] = None,
    model: Optional[str] = None,
    system_prompt: Optional[str] = None,
    max_tokens: Optional[int] = None,
    temperature: float = 0.0,
    use_tools: bool = False,
    tools: Optional[list] = None,
) -> InvokeResult:
    """Lightweight stateless invocation — no retry, no history, no prompt builder.

    Each concrete client implements this method using provider-native structured
    output. Use this instead of ``ask()`` when you need fast, stateless structured
    extraction without conversation history overhead.

    Args:
        prompt: The user prompt to send.
        output_type: A Pydantic model or dataclass class to parse the response into.
            Mutually exclusive with ``structured_output`` (``structured_output`` wins).
        structured_output: Full :class:`StructuredOutputConfig`, including optional
            ``custom_parser``. Takes precedence over ``output_type``.
        model: Override the model for this call. Falls back to an
            explicitly selected ``self.model``, then ``_lightweight_model``,
            then ``_default_model`` (see :meth:`_resolve_invoke_model`).
        system_prompt: Override the system prompt. Falls back to
            ``BASIC_SYSTEM_PROMPT`` rendered with instance attributes.
        max_tokens: Maximum completion tokens. ``None`` (the default) resolves
            via :meth:`_resolve_invoke_max_tokens` — per-instance
            ``invoke_max_tokens``, then ``max_tokens``, then the client's
            :attr:`_invoke_max_tokens` class default. Pass an explicit value
            only to cap a single call.
        temperature: Sampling temperature (default 0.0 for deterministic output).
        use_tools: If ``True``, inject registered tools into the request.
        tools: Additional tool definitions to pass directly to the provider.

    Returns:
        :class:`InvokeResult` with the parsed ``output``, ``model``, ``usage``,
        and optional ``raw_response``.

    Raises:
        :class:`InvokeError`: If the provider call fails for any reason.
    """
    raise NotImplementedError("Subclasses must implement invoke().")

StreamingRetryConfig

StreamingRetryConfig(max_retries: int = 3, base_delay: float = 1.0, max_delay: float = 60.0, backoff_factor: float = 2.0, jitter: bool = True, auto_retry_on_max_tokens: bool = True, token_increase_factor: float = 1.5, retry_on_rate_limit: bool = True, retry_on_server_error: bool = True)

Configuration for streaming retry behavior.

Source code in packages/ai-parrot/src/parrot/clients/base.py

def __init__(
    self,
    max_retries: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    backoff_factor: float = 2.0,
    jitter: bool = True,
    auto_retry_on_max_tokens: bool = True,
    token_increase_factor: float = 1.5,
    retry_on_rate_limit: bool = True,
    retry_on_server_error: bool = True,
):
    self.max_retries = max_retries
    self.base_delay = base_delay
    self.max_delay = max_delay
    self.backoff_factor = backoff_factor
    self.jitter = jitter
    self.auto_retry_on_max_tokens = auto_retry_on_max_tokens
    self.token_increase_factor = token_increase_factor
    self.retry_on_rate_limit = retry_on_rate_limit
    self.retry_on_server_error = retry_on_server_error

OpenAIBaseClient

OpenAIBaseClient(api_key: str | None = None, base_url: str | None = None, **kwargs)

Bases: AbstractClient

OpenAI-compatible wire protocol; carries NO OpenAI-provider defaults.

Subclasses that speak the OpenAI chat-completions wire protocol under a provider-specific label (Bedrock Mantle, OpenRouter, Moonshot, Nvidia, LocalLLM/vLLM, and — in Phase 2 — Groq/Zai via their native SDKs) should inherit from this class instead of :class:~parrot.clients.openai.OpenAIClient directly, so they never inherit OpenAI-the-provider defaults (OpenAI-only model ids, Responses-API routing, Sora, etc.).

Initialize the OpenAI-compatible wire client.

PARAMETER	DESCRIPTION
api_key	Bearer token for the target endpoint. Providers supply their own environment-variable default in their own __init__ — this base does not read any env var. TYPE: str | None DEFAULT: None
base_url	Base URL of the OpenAI-compatible endpoint. Providers supply their own default in their own __init__. TYPE: str | None DEFAULT: None
**kwargs	Forwarded to :class:~parrot.clients.base.AbstractClient. May include model (normalized via :meth:_normalize_model before being forwarded) and timeout (SDK request timeout, defaulting to this class's :attr:_default_timeout, which subclasses raise for slower endpoints). DEFAULT: {}

Source code in packages/ai-parrot/src/parrot/clients/openai_base.py

def __init__(
    self,
    api_key: str | None = None,
    base_url: str | None = None,
    **kwargs,
):
    """Initialize the OpenAI-compatible wire client.

    Args:
        api_key: Bearer token for the target endpoint. Providers supply
            their own environment-variable default in their own
            ``__init__`` — this base does not read any env var.
        base_url: Base URL of the OpenAI-compatible endpoint. Providers
            supply their own default in their own ``__init__``.
        **kwargs: Forwarded to :class:`~parrot.clients.base.AbstractClient`.
            May include ``model`` (normalized via :meth:`_normalize_model`
            before being forwarded) and ``timeout`` (SDK request timeout,
            defaulting to this class's :attr:`_default_timeout`, which
            subclasses raise for slower endpoints).
    """
    self.api_key = api_key
    self.base_url = base_url
    self.base_headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {self.api_key}",
    }
    self._timeout = kwargs.pop("timeout", self._default_timeout)
    if "model" in kwargs:
        kwargs["model"] = self._normalize_model(kwargs["model"])
    super().__init__(**kwargs)

get_client async

get_client() -> Any

Build the default OpenAI-SDK-shaped async client.

Lazily imports openai.AsyncOpenAI so the SDK is only required when an OpenAI-compatible client is actually instantiated. Subclasses that wrap a native SDK (Groq, Zai) override this hook.

RETURNS	DESCRIPTION
Any	An AsyncOpenAI instance configured with this client's
Any	api_key/base_url/timeout.

RAISES	DESCRIPTION
ImportError	If the openai package is not installed.

Source code in packages/ai-parrot/src/parrot/clients/openai_base.py

async def get_client(self) -> Any:
    """Build the default OpenAI-SDK-shaped async client.

    Lazily imports ``openai.AsyncOpenAI`` so the SDK is only required
    when an OpenAI-compatible client is actually instantiated. Subclasses
    that wrap a native SDK (Groq, Zai) override this hook.

    Returns:
        An ``AsyncOpenAI`` instance configured with this client's
        ``api_key``/``base_url``/timeout.

    Raises:
        ImportError: If the ``openai`` package is not installed.
    """
    try:
        from openai import AsyncOpenAI
    except ImportError as exc:
        raise ImportError(
            "OpenAIBaseClient requires the 'openai' SDK. " "Install with: pip install ai-parrot[openai]"
        ) from exc
    return AsyncOpenAI(
        api_key=self.api_key,
        base_url=self.base_url,
        timeout=self._timeout,
    )

ask async

ask(prompt: str, model: Any | None = None, max_tokens: int | None = None, temperature: float | None = None, files: list[str | Path] | None = None, system_prompt: str | None = None, history: Sequence[HistoryMessage] | None = None, structured_output: type | StructuredOutputConfig | None = None, tools: list[dict[str, Any]] | None = None, use_tools: bool | None = None, lazy_loading: bool = False) -> AIMessage

Ask the OpenAI-compatible endpoint a question with optional conversation memory.

Generic chat-completions implementation shared by every OpenAIBaseClient subclass that speaks the plain OpenAI wire protocol. Subclasses needing Responses-API routing, deep-research dispatch, or other OpenAI-provider-only behavior override this method (currently only :class:~parrot.clients.openai.OpenAIClient, which reuses :meth:_chat_completion/:meth:_run_tool_call_loop from this base while keeping its own richer ask()).

PARAMETER	DESCRIPTION
prompt	The prompt to send to the model. TYPE: str
model	The model to use, or None to use the configured one. TYPE: Any | None DEFAULT: None
max_tokens	Maximum tokens for the response. TYPE: int | None DEFAULT: None
temperature	Sampling temperature. TYPE: float | None DEFAULT: None
files	Files to upload before the call. TYPE: list[str | Path] | None DEFAULT: None
system_prompt	System prompt to prepend. TYPE: str | None DEFAULT: None
structured_output	Structured output definition (Pydantic model, dataclass, or explicit StructuredOutputConfig). TYPE: type | StructuredOutputConfig | None DEFAULT: None
user_id	User ID for conversation memory.
session_id	Session ID for conversation memory.
tools	Tools to register for this call. TYPE: list[dict[str, Any]] | None DEFAULT: None
use_tools	Whether to use tools; defaults to self.enable_tools. TYPE: bool | None DEFAULT: None
lazy_loading	If True, enable dynamic tool searching. TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
AIMessage	The response from the model.

RAISES	DESCRIPTION
NotImplementedError	If the resolved model requires Responses-API routing (:meth:_is_responses_model) and this class does not override ask() to handle it.

Source code in packages/ai-parrot/src/parrot/clients/openai_base.py

async def ask(
    self,
    prompt: str,
    model: Any | None = None,
    max_tokens: int | None = None,
    temperature: float | None = None,
    files: list[str | Path] | None = None,
    system_prompt: str | None = None,
    history: Sequence[HistoryMessage] | None = None,
    structured_output: type | StructuredOutputConfig | None = None,
    tools: list[dict[str, Any]] | None = None,
    use_tools: bool | None = None,
    lazy_loading: bool = False,
) -> AIMessage:
    """Ask the OpenAI-compatible endpoint a question with optional conversation memory.

    Generic chat-completions implementation shared by every
    ``OpenAIBaseClient`` subclass that speaks the plain OpenAI wire
    protocol. Subclasses needing Responses-API routing, deep-research
    dispatch, or other OpenAI-provider-only behavior override this
    method (currently only :class:`~parrot.clients.openai.OpenAIClient`,
    which reuses :meth:`_chat_completion`/:meth:`_run_tool_call_loop`
    from this base while keeping its own richer ``ask()``).

    Args:
        prompt: The prompt to send to the model.
        model: The model to use, or ``None`` to use the configured one.
        max_tokens: Maximum tokens for the response.
        temperature: Sampling temperature.
        files: Files to upload before the call.
        system_prompt: System prompt to prepend.
        structured_output: Structured output definition (Pydantic model,
            dataclass, or explicit ``StructuredOutputConfig``).
        user_id: User ID for conversation memory.
        session_id: Session ID for conversation memory.
        tools: Tools to register for this call.
        use_tools: Whether to use tools; defaults to ``self.enable_tools``.
        lazy_loading: If ``True``, enable dynamic tool searching.

    Returns:
        The response from the model.

    Raises:
        NotImplementedError: If the resolved model requires Responses-API
            routing (:meth:`_is_responses_model`) and this class does not
            override ``ask()`` to handle it.
    """
    turn_id = str(uuid.uuid4())
    original_prompt = prompt
    _use_tools = use_tools if use_tools is not None else self.enable_tools

    model_str = self._resolve_model(model)

    if self._is_responses_model(model_str):
        raise NotImplementedError(
            f"{type(self).__name__} does not implement Responses-API routing "
            "for this model; override ask() to handle it."
        )

    messages = self._build_messages(prompt, files, history)
    _lc_tc = self._emit_before_call(
        client_name=self.client_name,
        model=model_str,
        temperature=temperature if temperature is not None else self.temperature,
        system_prompt=system_prompt,
        has_tools=bool(_use_tools),
        parent_trace=None,
    )
    _lc_t0 = time.perf_counter()

    if files:
        for file in files:
            if isinstance(file, str):
                file = Path(file)
            if isinstance(file, Path):
                await self._upload_file(file)

    if lazy_loading and system_prompt:
        system_prompt += (
            "\n\nYou have access to a library of tools. Use the 'search_tools' function to find relevant tools."
        )
    elif lazy_loading and not system_prompt:
        system_prompt = (
            "You have access to a library of tools. Use the 'search_tools' function to find relevant tools."
        )

    if system_prompt:
        if isinstance(system_prompt, list):
            system_prompt = "\n\n".join(s.text for s in system_prompt)
        messages.insert(0, {"role": "system", "content": system_prompt})

    # FEAT-524: do NOT append the current turn here. _build_messages()
    # already placed it last, after the rendered history; appending again
    # sent the prompt twice and produced two consecutive `user` messages,
    # which strict-alternation providers reject.

    output_config = self._get_structured_config(structured_output)

    if tools and isinstance(tools, list):
        for tool in tools:
            self.register_tool(tool)

    active_tool_names = set()
    prepared_tools = None

    if _use_tools:
        if lazy_loading:
            prepared_tools = self._prepare_lazy_tools()
            if prepared_tools:
                active_tool_names.add("search_tools")
        else:
            prepared_tools = self._prepare_tools()

    args: dict[str, Any] = {}
    if prepared_tools:
        args["tools"] = prepared_tools
        args["tool_choice"] = "auto"

    _resolved_max_tokens = self._resolve_max_tokens(max_tokens)
    if _resolved_max_tokens is not None:
        args["max_tokens"] = _resolved_max_tokens
    if temperature:
        args["temperature"] = temperature

    if output_config:
        args["response_format"] = self._build_response_format_from(output_config)

    _used_fallback = False
    _original_model = model_str
    # FEAT-550: one call_id per ask() invocation; the pre-loop call is
    # round 1 (spec §2.2/§2.4 — every physical attempt reserves anew).
    _budget_call_id = str(uuid.uuid4())
    _budget_forced = False
    _BUDGET_CALL_CTX.set({"call_id": _budget_call_id, "round_number": 1, "phase": "work"})

    _round_t0 = time.perf_counter()
    try:
        response = await self._chat_completion(model=model_str, messages=messages, use_tools=_use_tools, **args)
    except BudgetExhausted:
        # Ordinary work denied before the tool loop even starts: the
        # owner attempts ONE tools-disabled final attempt (spec §2.3).
        response = await self._finalize_budgeted_chat(
            messages,
            model_str=model_str,
            args=args,
            all_tool_calls=[],
            pending_tool_calls=[],
            partial_text="",
            stream=False,
        )
        _budget_forced = True
        _BUDGET_CALL_CTX.set({**(_BUDGET_CALL_CTX.get() or {}), "forced": True})
    except Exception as e:
        if self._should_use_fallback(model_str, e):
            self.logger.warning(
                "Model %s capacity error: %s. Retrying once with fallback: %s",
                model_str,
                e,
                self._fallback_model,
            )
            model_str = self._fallback_model
            _used_fallback = True
            try:
                response = await self._chat_completion(
                    model=model_str, messages=messages, use_tools=_use_tools, **args
                )
            except Exception as fallback_exc:
                await self._emit_failed_call_safe(
                    _lc_tc,
                    self.client_name,
                    model_str,
                    _lc_t0,
                    fallback_exc,
                )
                raise
        else:
            await self._emit_failed_call_safe(
                _lc_tc,
                self.client_name,
                model_str,
                _lc_t0,
                e,
            )
            raise
    _round_duration_ms = (time.perf_counter() - _round_t0) * 1000
    result = response.choices[0].message

    def _on_round(round_number, usage, raw_usage, tool_names, duration_ms):
        self._emit_round_event(
            _lc_tc,
            client_name=self.client_name,
            model=model_str,
            round_number=round_number,
            usage=usage,
            raw_usage=raw_usage,
            tool_calls=tool_names,
            duration_ms=duration_ms,
        )

    if _budget_forced:
        # The pre-loop call above was already the ONE tools-disabled
        # final attempt (spec §2.3) — never re-enter the ordinary tool
        # loop: `args` still carries the original `tools`/`tool_choice`,
        # so looping here could re-enable tools and execute one after
        # the terminal attempt (code-reviewer finding, FEAT-550
        # wrap-up). Mirrors `_run_tool_call_loop`'s own no-tool-calls
        # trivial return shape exactly.
        all_tool_calls: list[ToolCall] = []
        round_number = 1
        accumulated_usage, _ = self._extract_completion_usage(response)
        _bq_scope_pre = current_budget_scope()
        if _bq_scope_pre is not None and getattr(result, "tool_calls", None):
            await _bq_scope_pre.ledger.set_answer_complete(False)
    else:
        result, response, all_tool_calls, accumulated_usage, round_number = await self._run_tool_call_loop(
            result=result,
            response=response,
            messages=messages,
            model_str=model_str,
            use_tools=_use_tools,
            args=args,
            session_id=current_session_id.get(),
            lazy_loading=lazy_loading,
            active_tool_names=active_tool_names,
            track_usage=True,
            initial_duration_ms=_round_duration_ms,
            on_round=_on_round,
        )
    # Either the pre-loop call or a round inside the loop may have forced
    # finalization (spec §2.3) — the loop signals it via _BUDGET_CALL_CTX
    # rather than the shared 5-tuple return (other callers depend on it).
    _budget_forced = _budget_forced or bool((_BUDGET_CALL_CTX.get() or {}).get("forced"))

    messages.append({"role": "assistant", "content": result.content})

    response_text = result.content if isinstance(result.content, str) else self._json.dumps(result.content)
    final_output = None
    _bq_ask_scope = current_budget_scope()
    if output_config and not _budget_forced:
        try:
            # Known-truncated output must not reach a custom parser either.
            self._raise_if_truncated(self._extract_finish_reason(response), model=model_str)
            if output_config.custom_parser:
                final_output = output_config.custom_parser(response_text)
            else:
                final_output = await self._parse_structured_output(
                    response_text,
                    output_config,
                    finish_reason=self._extract_finish_reason(response),
                    model=model_str,
                )
        except InvokeError:
            raise
        except Exception:  # noqa: BLE001 pylint: disable=broad-except
            final_output = response_text
    elif _budget_forced and output_config:
        # FEAT-550 §2.3: a budget-forced final never reaches a custom
        # parser as if it were validated — raw text, marked incomplete.
        final_output = response_text
        if _bq_ask_scope is not None:
            await _bq_ask_scope.ledger.set_answer_complete(False)

    # FEAT-524: no memory write — AbstractBot.save_conversation_turn is the single writer.

    structured_payload = None
    if final_output is not None and not (isinstance(final_output, str) and final_output == response_text):
        structured_payload = final_output

    ai_message = AIMessageFactory.from_openai(
        response=response,
        input_text=original_prompt,
        model=model_str,
        user_id=current_user_id.get(),
        session_id=current_session_id.get(),
        turn_id=turn_id,
        structured_output=structured_payload,
    )

    if accumulated_usage is not None:
        if round_number > 1:
            accumulated_usage.extra_usage["rounds"] = round_number
        ai_message.usage = accumulated_usage

    ai_message.tool_calls = all_tool_calls
    if _used_fallback:
        ai_message.metadata["used_fallback_model"] = True
        ai_message.metadata["original_model"] = _original_model
        ai_message.metadata["fallback_model"] = self._fallback_model

    # FEAT-550 §2.3: every budgeted answer carries the ledger's report;
    # forced termination surfaces stop_reason="budget_exhausted" while the
    # provider's own value stays available via finish_reason below.
    if _bq_ask_scope is not None:
        if _budget_forced:
            ai_message.stop_reason = "budget_exhausted"
        ai_message.metadata["token_budget"] = (await _bq_ask_scope.ledger.report()).model_dump()

    _lc_usage = getattr(ai_message, "usage", None)
    await self._emit_after_call(
        _lc_tc,
        client_name=self.client_name,
        model=model_str,
        duration_ms=(time.perf_counter() - _lc_t0) * 1000,
        input_tokens=getattr(_lc_usage, "prompt_tokens", None) if _lc_usage else None,
        output_tokens=getattr(_lc_usage, "completion_tokens", None) if _lc_usage else None,
        finish_reason=getattr(ai_message, "stop_reason", None),
    )
    return ai_message

resume async

resume(session_id: str, user_input: str, state: dict[str, Any]) -> AIMessage

Resume a suspended model execution.

PARAMETER	DESCRIPTION
session_id	The session ID. TYPE: str
user_input	The user's input to inject as a tool result. TYPE: str
state	The suspended state containing messages and tool_call_id. TYPE: dict[str, Any]

RETURNS	DESCRIPTION
AIMessage	The response from the model.

Source code in packages/ai-parrot/src/parrot/clients/openai_base.py

async def resume(self, session_id: str, user_input: str, state: dict[str, Any]) -> AIMessage:
    """Resume a suspended model execution.

    Args:
        session_id: The session ID.
        user_input: The user's input to inject as a tool result.
        state: The suspended state containing messages and tool_call_id.

    Returns:
        The response from the model.
    """
    await self._ensure_client()

    # FEAT-550 §2.6: deep-copy (never mutate the caller's stored suspended
    # state). Reattachment of a suspended budget scope from the
    # TOKEN_BUDGET_STATE_KEY envelope already happened in the client-level
    # entry adapter (AbstractClient.__init_subclass__ wraps resume() and
    # calls registry.resume(state, ...) before this body runs — see
    # budget_scope.py::_enter_scope) — current_budget_scope() below
    # already reflects it; no manual reattachment needed here.
    messages = copy.deepcopy(state["messages"])
    tool_call_id = state["tool_call_id"]
    model_str = state.get("agent_name", self.model or self.default_model)
    _scope = current_budget_scope()

    messages.append({"role": "tool", "tool_call_id": tool_call_id, "name": "handoff_tool", "content": user_input})

    turn_id = str(uuid.uuid4())

    response = await self._chat_completion(model=model_str, messages=messages, use_tools=True)
    result = response.choices[0].message

    _track_usage = _scope is not None
    result, response, all_tool_calls, accumulated_usage, _round_number = await self._run_tool_call_loop(
        result=result,
        response=response,
        messages=messages,
        model_str=model_str,
        use_tools=True,
        args={},
        session_id=session_id,
        record_malformed_tool_calls=False,
        default_tool_name="unknown",
        track_usage=_track_usage,
    )

    ai_message = AIMessageFactory.from_openai(
        response=response,
        input_text="[Resumed Conversation]",
        model=model_str,
        user_id="unknown",
        session_id=session_id,
        turn_id=turn_id,
    )
    ai_message.tool_calls = all_tool_calls

    # FEAT-550: Attach accumulated usage and budget report
    if accumulated_usage is not None and _scope is not None:
        ai_message.usage = accumulated_usage
        ai_message.metadata["token_budget"] = (await _scope.ledger.report()).model_dump()

    return ai_message

batch_ask async

batch_ask(requests: list[dict[str, Any]]) -> list[AIMessage]

Process multiple ask() requests sequentially.

No native batch API exists for the OpenAI wire protocol; requests are processed one at a time via :meth:ask.

PARAMETER	DESCRIPTION
requests	A list of kwargs dicts, each forwarded to :meth:ask. TYPE: list[dict[str, Any]]

RETURNS	DESCRIPTION
list[AIMessage]	The list of :class:AIMessage responses, in request order.

Source code in packages/ai-parrot/src/parrot/clients/openai_base.py

async def batch_ask(self, requests: list[dict[str, Any]]) -> list[AIMessage]:
    """Process multiple ``ask()`` requests sequentially.

    No native batch API exists for the OpenAI wire protocol; requests
    are processed one at a time via :meth:`ask`.

    Args:
        requests: A list of kwargs dicts, each forwarded to :meth:`ask`.

    Returns:
        The list of :class:`AIMessage` responses, in request order.
    """
    results = []
    for request in requests:
        result = await self.ask(**request)
        results.append(result)
    return results

ask_stream async

ask_stream(prompt: str, model: Any | None = None, max_tokens: int | None = None, temperature: float | None = None, files: list[str | Path] | None = None, system_prompt: str | None = None, history: Sequence[HistoryMessage] | None = None, tools: list[dict[str, Any]] | None = None, use_tools: bool = True, structured_output: type | StructuredOutputConfig | None = None, lazy_loading: bool = False, **kwargs) -> AsyncIterator[str | AIMessage]

Stream a response with tool-use support and conversation memory.

Generic chat-completions streaming implementation shared by every OpenAIBaseClient subclass. Subclasses needing Responses-API routing or other OpenAI-provider-only streaming behavior override this method (currently only :class:~parrot.clients.openai.OpenAIClient). Routes through :meth:_chat_completion with stream=True (FEAT-438 G3 — the single completion funnel) instead of calling the SDK directly, so a subclass's funnel override applies to streaming too.

When the model requests tool calls during streaming, the tool-call chunks are accumulated, the tools executed between rounds, and a new streaming request is issued with the tool results — mirroring the _run_tool_call_loop() used by :meth:ask.

Yields successive string chunks followed by a final :class:~parrot.models.responses.AIMessage.

PARAMETER	DESCRIPTION
use_tools	Whether to include tool definitions and run the streaming tool-call loop. Defaults to True. TYPE: bool DEFAULT: True

RAISES	DESCRIPTION
NotImplementedError	If the resolved model requires Responses-API routing and this class does not override ask_stream().

Source code in packages/ai-parrot/src/parrot/clients/openai_base.py

async def ask_stream(
    self,
    prompt: str,
    model: Any | None = None,
    max_tokens: int | None = None,
    temperature: float | None = None,
    files: list[str | Path] | None = None,
    system_prompt: str | None = None,
    history: Sequence[HistoryMessage] | None = None,
    tools: list[dict[str, Any]] | None = None,
    use_tools: bool = True,
    structured_output: type | StructuredOutputConfig | None = None,
    lazy_loading: bool = False,
    **kwargs,
) -> AsyncIterator[str | AIMessage]:
    """Stream a response with tool-use support and conversation memory.

    Generic chat-completions streaming implementation shared by every
    ``OpenAIBaseClient`` subclass. Subclasses needing Responses-API
    routing or other OpenAI-provider-only streaming behavior override
    this method (currently only :class:`~parrot.clients.openai.OpenAIClient`).
    Routes through :meth:`_chat_completion` with ``stream=True``
    (FEAT-438 G3 — the single completion funnel) instead of calling the
    SDK directly, so a subclass's funnel override applies to streaming
    too.

    When the model requests tool calls during streaming, the tool-call
    chunks are accumulated, the tools executed between rounds, and a new
    streaming request is issued with the tool results — mirroring the
    ``_run_tool_call_loop()`` used by :meth:`ask`.

    Yields successive string chunks followed by a final
    :class:`~parrot.models.responses.AIMessage`.

    Args:
        use_tools: Whether to include tool definitions and run the
            streaming tool-call loop.  Defaults to ``True``.

    Raises:
        NotImplementedError: If the resolved model requires Responses-API
            routing and this class does not override ``ask_stream()``.
    """
    turn_id = str(uuid.uuid4())
    model_str = self._resolve_model(model)

    if self._is_responses_model(model_str):
        raise NotImplementedError(
            f"{type(self).__name__} does not implement Responses-API "
            "streaming for this model; override ask_stream() to handle it."
        )

    messages = self._build_messages(prompt, files, history)
    _lc_tc = self._emit_before_call(
        client_name=self.client_name,
        model=model_str,
        temperature=temperature if temperature is not None else self.temperature,
        system_prompt=system_prompt,
        has_tools=bool(self.tools) and use_tools,
        parent_trace=None,
    )
    _lc_t0 = time.perf_counter()

    if files:
        for file in files:
            if isinstance(file, str):
                file = Path(file)
            if isinstance(file, Path):
                await self._upload_file(file)

    if lazy_loading and system_prompt:
        system_prompt += (
            "\n\nYou have access to a library of tools. Use the 'search_tools' function to find relevant tools."
        )
    elif lazy_loading and not system_prompt:
        system_prompt = (
            "You have access to a library of tools. Use the 'search_tools' function to find relevant tools."
        )

    if system_prompt:
        if isinstance(system_prompt, list):
            system_prompt = "\n\n".join(s.text for s in system_prompt)
        messages.insert(0, {"role": "system", "content": system_prompt})

    if tools and isinstance(tools, list):
        for tool in tools:
            self.register_tool(tool)

    tools_payload = None
    if use_tools and self.tools:
        if lazy_loading:
            tools_payload = self._prepare_lazy_tools()
        else:
            tools_payload = self._prepare_tools()

    args: dict[str, Any] = {"stream_options": {"include_usage": True}}
    if tools_payload:
        args["tools"] = tools_payload
        args["tool_choice"] = "auto"

    max_tokens_value = self._resolve_max_tokens(max_tokens)
    if max_tokens_value is not None:
        args["max_tokens"] = max_tokens_value

    temperature_value = temperature if temperature is not None else self.temperature
    if temperature_value is not None:
        args["temperature"] = temperature_value

    output_config = self._get_structured_config(structured_output)
    if output_config:
        response_format = self._build_response_format_from(output_config)
        if response_format:
            args["response_format"] = response_format

    # ── Streaming tool-call loop ──────────────────────────────────
    # Mirrors _run_tool_call_loop() used by ask(). Each round streams
    # text and accumulates tool-call chunks; when the model finishes
    # with tool calls, they are executed between rounds and the
    # conversation continues.
    all_tool_calls: list[ToolCall] = []
    assistant_content = ""
    usage_data = None
    _max_tool_rounds = 25  # safety cap
    # FEAT-550: one call_id per ask_stream() invocation; each round is a
    # fresh physical attempt admitted before dispatch (spec §2.2/§2.4).
    _budget_call_id_s = str(uuid.uuid4())
    _budget_forced = False

    for _round in range(_max_tool_rounds):
        # Accumulate tool-call chunks by index. OpenAI streams them
        # incrementally: first chunk carries id+name, subsequent ones
        # carry argument fragments.
        _tc_accum: dict[int, dict[str, str]] = {}
        _round_content = ""
        _finish_reason: str | None = None

        _BUDGET_CALL_CTX.set({"call_id": _budget_call_id_s, "round_number": _round + 1, "phase": "work"})
        try:
            # `.parse()` cannot reliably stream every SDK's tool-calling/plain
            # responses; only prefer it when structured output was requested —
            # mirrors the pre-FEAT-438 dispatch this funnel now formalizes.
            response_stream = await self._chat_completion(
                model=model_str, messages=messages, use_tools=not bool(output_config), stream=True, **args
            )
        except asyncio.CancelledError:
            # spec §2.3 table row 4: cancellation never starts finalization.
            raise
        except BudgetExhausted:
            # Ordinary work denied: the owner attempts ONE tools-disabled
            # final attempt, streaming its text in-band (spec §2.3/§2.4).
            partial_text = assistant_content
            final_stream = await self._finalize_budgeted_chat(
                messages,
                model_str=model_str,
                args=args,
                all_tool_calls=all_tool_calls,
                pending_tool_calls=[],
                partial_text=partial_text,
                stream=True,
            )
            async for chunk in final_stream:
                if chunk.choices:
                    choice = chunk.choices[0]
                    delta = choice.delta
                    if delta and delta.content:
                        text_chunk = delta.content
                        _round_content += text_chunk
                        assistant_content += text_chunk
                        yield text_chunk
                    _finish_reason = getattr(choice, "finish_reason", None) or _finish_reason
                if hasattr(chunk, "usage") and chunk.usage is not None:
                    usage_data = chunk.usage
            _budget_forced = True
            if _round_content:
                messages.append({"role": "assistant", "content": _round_content})
            _bq_scope_s = current_budget_scope()
            if _bq_scope_s is not None and _finish_reason in ("tool_calls", "length"):
                await _bq_scope_s.ledger.set_answer_complete(False)
            break
        except Exception as _stream_exc:
            await self._emit_failed_call_safe(
                _lc_tc,
                self.client_name,
                model_str,
                _lc_t0,
                _stream_exc,
            )
            raise

        async for chunk in response_stream:
            if chunk.choices:
                choice = chunk.choices[0]
                delta = choice.delta
                if delta and delta.content:
                    text_chunk = delta.content
                    _round_content += text_chunk
                    assistant_content += text_chunk
                    yield text_chunk
                # Accumulate streamed tool-call fragments.
                if delta and getattr(delta, "tool_calls", None):
                    for tc_delta in delta.tool_calls:
                        idx = tc_delta.index
                        if idx not in _tc_accum:
                            _tc_accum[idx] = {
                                "id": getattr(tc_delta, "id", None) or "",
                                "name": "",
                                "arguments": "",
                            }
                        if tc_delta.function:
                            if tc_delta.function.name:
                                _tc_accum[idx]["name"] = tc_delta.function.name
                            if tc_delta.function.arguments:
                                _tc_accum[idx]["arguments"] += tc_delta.function.arguments
                _finish_reason = getattr(choice, "finish_reason", None) or _finish_reason
            if hasattr(chunk, "usage") and chunk.usage is not None:
                usage_data = chunk.usage

        # ── Tool-use round: execute and loop ──────────────────────
        if _tc_accum and _finish_reason in ("tool_calls", "stop"):
            # Build assistant message with tool_calls for conversation.
            tc_serialized = []
            for idx in sorted(_tc_accum):
                tc_info = _tc_accum[idx]
                tc_serialized.append(
                    {
                        "id": tc_info["id"],
                        "type": "function",
                        "function": {
                            "name": tc_info["name"],
                            "arguments": tc_info["arguments"],
                        },
                    }
                )
            messages.append(
                {
                    "role": "assistant",
                    "content": _round_content or None,
                    "tool_calls": tc_serialized,
                }
            )

            for idx in sorted(_tc_accum):
                tc_info = _tc_accum[idx]
                try:
                    tool_args = json.loads(tc_info["arguments"]) if tc_info["arguments"] else {}
                except (json.JSONDecodeError, TypeError):
                    tool_args = {}
                tc = ToolCall(
                    id=tc_info["id"],
                    name=tc_info["name"],
                    arguments=tool_args,
                )
                try:
                    start_time = time.time()
                    tool_result = await self._execute_tool(tc_info["name"], tool_args)
                    tc.result = tool_result
                    tc.execution_time = time.time() - start_time
                    messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": tc_info["id"],
                            "name": tc_info["name"],
                            "content": str(tool_result),
                        }
                    )
                except Exception as e:
                    from parrot.core.exceptions import HumanInteractionInterrupt

                    if isinstance(e, HumanInteractionInterrupt):
                        e.session_id = current_session_id.get()
                        e.messages = messages
                        e.tool_call_id = tc_info["id"]
                        e.agent_name = model_str
                        _scope = current_budget_scope()
                        if _scope is not None:
                            e.state = {
                                **(getattr(e, "state", None) or {}),
                                TOKEN_BUDGET_STATE_KEY: await _scope.registry.suspend(_scope),
                            }
                        raise
                    elif isinstance(e, BudgetError):
                        # FEAT-550 §3 M3: budget control from a tool's own
                        # inner budgeted call must propagate to the owner —
                        # never converted into a tool result (code-reviewer
                        # finding, FEAT-550 wrap-up).
                        raise
                    tc.error = str(e)
                    messages.append(
                        {
                            "role": "tool",
                            "tool_call_id": tc_info["id"],
                            "name": tc_info["name"],
                            "content": str(e),
                        }
                    )
                all_tool_calls.append(tc)
            # Continue to next streaming round.
        else:
            # No tool calls — append assistant content and break.
            if _round_content:
                messages.append({"role": "assistant", "content": _round_content})
            break

    if usage_data is not None:
        usage = CompletionUsage.from_openai(usage_data)
    else:
        usage = CompletionUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0)

    ai_message = AIMessage(
        input=prompt,
        output=assistant_content,
        response=assistant_content,
        model=model_str,
        provider=self.client_type,
        usage=usage,
        user_id=current_user_id.get(),
        session_id=current_session_id.get(),
        turn_id=turn_id,
        tool_calls=all_tool_calls,
    )

    # FEAT-550 §2.3/§2.4: attach the ledger report to the single terminal
    # AIMessage; forced termination sets stop_reason.
    _bq_stream_scope = current_budget_scope()
    if _bq_stream_scope is not None:
        if _budget_forced:
            ai_message.stop_reason = "budget_exhausted"
        ai_message.metadata["token_budget"] = (await _bq_stream_scope.ledger.report()).model_dump()

    await self._emit_after_call(
        _lc_tc,
        client_name=self.client_name,
        model=model_str,
        duration_ms=(time.perf_counter() - _lc_t0) * 1000,
        input_tokens=getattr(usage, "prompt_tokens", None),
        output_tokens=getattr(usage, "completion_tokens", None),
        finish_reason=_finish_reason,
    )
    yield ai_message

invoke async

invoke(prompt: str, *, output_type: type | None = None, structured_output: StructuredOutputConfig | None = None, model: str | None = None, system_prompt: str | None = None, max_tokens: int | None = None, temperature: float = 0.0, use_tools: bool = False, tools: list | None = None) -> InvokeResult

Lightweight stateless invocation routed through the completion funnel.

Generic implementation shared by every OpenAIBaseClient subclass. A single call is made through :meth:_chat_completion (FEAT-438 G3 — previously this bypassed the funnel entirely and called the SDK directly) — no conversation history, no prompt builder; use_tools=True is passed to the funnel so it always uses .create() (matching the pre-FEAT-438 behavior of never using .parse() here).

PARAMETER	DESCRIPTION
prompt	User prompt. TYPE: str
output_type	Pydantic model or dataclass to parse the response into. TYPE: type | None DEFAULT: None
structured_output	Full StructuredOutputConfig; takes precedence over output_type. TYPE: StructuredOutputConfig | None DEFAULT: None
model	Model override. Falls back to an explicitly selected self.model, then _lightweight_model, then _default_model. TYPE: str | None DEFAULT: None
system_prompt	System prompt override. TYPE: str | None DEFAULT: None
max_tokens	Maximum completion tokens. TYPE: int | None DEFAULT: None
temperature	Sampling temperature. TYPE: float DEFAULT: 0.0
use_tools	Whether to inject registered tools. TYPE: bool DEFAULT: False
tools	Additional tool definitions (unused — kept for interface parity with :class:~parrot.clients.openai.OpenAIClient). TYPE: list | None DEFAULT: None

RETURNS	DESCRIPTION
InvokeResult	class:InvokeResult with parsed output.

RAISES	DESCRIPTION
InvokeError	On provider errors, or if the client is not initialized.

Source code in packages/ai-parrot/src/parrot/clients/openai_base.py

async def invoke(
    self,
    prompt: str,
    *,
    output_type: type | None = None,
    structured_output: StructuredOutputConfig | None = None,
    model: str | None = None,
    system_prompt: str | None = None,
    max_tokens: int | None = None,
    temperature: float = 0.0,
    use_tools: bool = False,
    tools: list | None = None,
) -> InvokeResult:
    """Lightweight stateless invocation routed through the completion funnel.

    Generic implementation shared by every ``OpenAIBaseClient``
    subclass. A single call is made through :meth:`_chat_completion`
    (FEAT-438 G3 — previously this bypassed the funnel entirely and
    called the SDK directly) — no conversation history, no prompt
    builder; ``use_tools=True`` is passed to the funnel so it always
    uses ``.create()`` (matching the pre-FEAT-438 behavior of never
    using ``.parse()`` here).

    Args:
        prompt: User prompt.
        output_type: Pydantic model or dataclass to parse the response into.
        structured_output: Full ``StructuredOutputConfig``; takes
            precedence over ``output_type``.
        model: Model override. Falls back to an explicitly selected
            ``self.model``, then ``_lightweight_model``, then
            ``_default_model``.
        system_prompt: System prompt override.
        max_tokens: Maximum completion tokens.
        temperature: Sampling temperature.
        use_tools: Whether to inject registered tools.
        tools: Additional tool definitions (unused — kept for interface
            parity with :class:`~parrot.clients.openai.OpenAIClient`).

    Returns:
        :class:`InvokeResult` with parsed output.

    Raises:
        InvokeError: On provider errors, or if the client is not
            initialized.
    """
    try:
        resolved_prompt = self._resolve_invoke_system_prompt(system_prompt)
        config = self._build_invoke_structured_config(output_type, structured_output)
        resolved_model = self._resolve_invoke_model(model)
        max_tokens = self._resolve_max_tokens(max_tokens, resolved_model, for_invoke=True)

        messages = [
            {"role": "system", "content": resolved_prompt},
            {"role": "user", "content": prompt},
        ]

        kwargs: dict[str, Any] = {
            "max_tokens": max_tokens,
            "temperature": temperature,
        }

        if config:
            response_format = self._build_response_format_from(config)
            if response_format:
                kwargs["response_format"] = response_format

        if use_tools:
            tool_defs = self._prepare_tools()
            if tool_defs:
                kwargs["tools"] = tool_defs

        if not self.client:
            raise RuntimeError(f"{type(self).__name__} not initialised. Use async context manager.")

        _budget_forced = False
        # FEAT-550: a stable call_id across the sole dispatch and any
        # follow-on finalization attempt — without this, `reserve()`'s
        # phase="final" owner check mints a fresh random call_id on each
        # `_BUDGET_CALL_CTX.get()` (empty default) and unconditionally
        # rejects the finalize dispatch as a foreign owner.
        _budget_call_id = str(uuid.uuid4())
        _BUDGET_CALL_CTX.set({"call_id": _budget_call_id, "round_number": 1, "phase": "work"})
        _lc_tc = self._emit_before_call(
            client_name=self.client_name,
            model=resolved_model,
            temperature=temperature,
            system_prompt=resolved_prompt,
            has_tools="tools" in kwargs,
            parent_trace=None,
        )
        _lc_t0 = time.perf_counter()
        _lc_terminal = False
        try:
            try:
                response = await self._chat_completion(
                    model=resolved_model, messages=messages, use_tools=True, **kwargs
                )
            except BudgetExhausted:
                # invoke() is a single-shot call: the frame is the payload
                # without tools; owner finalization still applies (spec §2.3).
                try:
                    response = await self._finalize_budgeted_chat(
                        messages,
                        model_str=resolved_model,
                        args=kwargs,
                        all_tool_calls=[],
                        pending_tool_calls=[],
                        partial_text="",
                        stream=False,
                    )
                    _budget_forced = True
                except BudgetExhausted as bx2:
                    _scope = current_budget_scope()
                    partial = (bx2.report or {}).pop("partial_text", "") or ""
                    invoke_result = InvokeResult(
                        output=partial,
                        output_type=None,
                        model=resolved_model,
                        usage=CompletionUsage(),
                        raw_response=None,
                    )
                    if _scope is not None:
                        invoke_result.budget_report = (await _scope.ledger.report()).model_dump()
                    _lc_terminal = True
                    await self._emit_after_call(
                        _lc_tc,
                        client_name=self.client_name,
                        model=resolved_model,
                        duration_ms=(time.perf_counter() - _lc_t0) * 1000,
                        input_tokens=None,
                        output_tokens=None,
                        finish_reason="budget_exhausted",
                    )
                    return invoke_result
        except Exception as exc:  # noqa: BLE001 — re-raised untouched below
            if not _lc_terminal:
                _lc_terminal = True
                await self._emit_failed_call_safe(_lc_tc, self.client_name, resolved_model, _lc_t0, exc)
            raise

        _lc_usage = CompletionUsage.from_openai(response.usage)
        _lc_terminal = True
        await self._emit_after_call(
            _lc_tc,
            client_name=self.client_name,
            model=resolved_model,
            duration_ms=(time.perf_counter() - _lc_t0) * 1000,
            input_tokens=_lc_usage.prompt_tokens,
            output_tokens=_lc_usage.completion_tokens,
            finish_reason="budget_exhausted" if _budget_forced else self._extract_finish_reason(response),
        )

        raw_text = response.choices[0].message.content or ""

        output: Any = raw_text
        if config and not _budget_forced:
            # Known-truncated output must not reach a custom parser either.
            self._raise_if_truncated(self._extract_finish_reason(response), model=resolved_model)
            if config.custom_parser:
                output = config.custom_parser(raw_text)
            else:
                output = await self._parse_structured_output(
                    raw_text,
                    config,
                    finish_reason=self._extract_finish_reason(response),
                    model=resolved_model,
                )
        elif _budget_forced:
            _scope = current_budget_scope()
            if _scope is not None:
                await _scope.ledger.set_answer_complete(False)

        usage = _lc_usage
        invoke_result = self._build_invoke_result(
            output if not _budget_forced else raw_text,
            None if _budget_forced else output_type,
            resolved_model,
            usage,
            response,
        )
        _scope = current_budget_scope()
        if _scope is not None:
            invoke_result.budget_report = (await _scope.ledger.report()).model_dump()
        return invoke_result
    except BudgetError:
        raise
    except InvokeError:
        raise
    except Exception as exc:  # noqa: BLE001 — funnel errors are wrapped into InvokeError
        raise self._handle_invoke_error(exc) from exc