Stamina doesn't seem to catch openai.APITimeoutError

[paw-lu]

First, thanks a bunch for creating and sharing this wrapper! It's really useful and I love the api shared here. Really appreciate it!

Using stamina with pydantic_ai and open ai. Despite having on=openai.APITimeoutError, it seems like the error is raised from the inner function before stamina retries it n times.

I'm guessing i'm misusing asyncio here.

import asyncio
import datetime
import random
import string

import openai
import pydantic
import pydantic_ai
import stamina
import structlog
from pydantic import Field
from pydantic_ai.models.openai import OpenAIModel

log = structlog.get_logger()

log.info("Hey i'm starting")


class Output(pydantic.BaseModel):
    unscrambled: str = Field(description="This is the scrambled text")


def make_agent():
    openai_model = OpenAIModel("gpt-4o-mini")
    system_prompt = "You'll get bunch of text, scramble it"
    return pydantic_ai.Agent(
        model=openai_model,
        system_prompt=system_prompt,
        retries=1,
        result_type=Output,
    )


def generate_random_string(length):
    # Define the possible characters: letters (both lowercase and uppercase) and digits
    characters = string.ascii_letters + string.digits
    # Use random.choices to pick 'length' number of characters from the characters string
    random_string = "".join(random.choices(characters, k=length))
    return random_string


async def get():
    agent = make_agent()
    stuff = (agent.run(generate_random_string(2000)) for _ in range(2))

    for thing in asyncio.as_completed(stuff):
        async for wrap in stamina.retry_context(
            on=openai.APITimeoutError,
            attempts=3,
            wait_initial=1,
            timeout=datetime.timedelta(hours=3),
        ):
            with wrap:
                await thing


def main():
    asyncio.run(get())


if __name__ == "__main__":
    main()

Logs and traceback

 uv run agent_me.py
2025-04-10 14:07:23 [info     ] Hey i'm starting
2025-04-10 14:07:25 [warning  ] stamina.retry_scheduled        args=() callable='<context block>' caused_by="APITimeoutError('Request timed out.')" kwargs={} retry_num=1 wait_for=1.33 waited_so_far=0.0
Traceback (most recent call last):
  File "/Users/pawlu/Documents/scratch/rate_limie_me/agent_me.py", line 60, in <module>
    main()
  File "/Users/pawlu/Documents/scratch/rate_limie_me/agent_me.py", line 56, in main
    asyncio.run(get())
  File "/Users/pawlu/.local/share/uv/python/cpython-3.12.9-macos-aarch64-none/lib/python3.12/asyncio/runners.py", line 195, in run
    return runner.run(main)
           ^^^^^^^^^^^^^^^^
  File "/Users/pawlu/.local/share/uv/python/cpython-3.12.9-macos-aarch64-none/lib/python3.12/asyncio/runners.py", line 118, in run
    return self._loop.run_until_complete(task)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/.local/share/uv/python/cpython-3.12.9-macos-aarch64-none/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/agent_me.py", line 46, in get
    async for wrap in stamina.retry_context(
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/stamina/_core.py", line 549, in __anext__
    await self._t_a_retrying.__anext__(),
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/tenacity/asyncio/__init__.py", line 166, in __anext__
    do = await self.iter(retry_state=self._retry_state)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/tenacity/asyncio/__init__.py", line 153, in iter
    result = await action(retry_state)
             ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/tenacity/_utils.py", line 99, in inner
    return call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/tenacity/__init__.py", line 400, in <lambda>
    self._add_action_func(lambda rs: rs.outcome.result())
                                     ^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/.local/share/uv/python/cpython-3.12.9-macos-aarch64-none/lib/python3.12/concurrent/futures/_base.py", line 449, in result
    return self.__get_result()
           ^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/.local/share/uv/python/cpython-3.12.9-macos-aarch64-none/lib/python3.12/concurrent/futures/_base.py", line 401, in __get_result
    raise self._exception
  File "/Users/pawlu/Documents/scratch/rate_limie_me/agent_me.py", line 52, in get
    await thing
RuntimeError: cannot reuse already awaited coroutine
Task exception was never retrieved
future: <Task finished name='Task-3' coro=<Agent.run() done, defined at /Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/pydantic_ai/agent.py:272> exception=APITimeoutError('Request timed out.')>
Traceback (most recent call last):
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/openai/_base_client.py", line 1500, in _request
    raise httpx.TimeoutException("Timeout")
httpx.TimeoutException: Timeout

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/openai/_base_client.py", line 1500, in _request
    raise httpx.TimeoutException("Timeout")
httpx.TimeoutException: Timeout

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/openai/_base_client.py", line 1500, in _request
    raise httpx.TimeoutException("Timeout")
httpx.TimeoutException: Timeout

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/pydantic_ai/agent.py", line 329, in run
    async for _ in agent_run:
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/pydantic_ai/agent.py", line 1454, in __anext__
    next_node = await self._graph_run.__anext__()
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/pydantic_graph/graph.py", line 790, in __anext__
    return await self.next(self._next_node)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/pydantic_graph/graph.py", line 763, in next
    self._next_node = await node.run(ctx)
                      ^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/pydantic_ai/_agent_graph.py", line 264, in run
    return await self._make_request(ctx)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/pydantic_ai/_agent_graph.py", line 317, in _make_request
    model_response, request_usage = await ctx.deps.model.request(
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/pydantic_ai/models/openai.py", line 193, in request
    response = await self._completions_create(
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/pydantic_ai/models/openai.py", line 266, in _completions_create
    return await self.client.chat.completions.create(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/openai/resources/chat/completions/completions.py", line 2000, in create
    return await self._post(
           ^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/openai/_base_client.py", line 1768, in post
    return await self.request(cast_to, opts, stream=stream, stream_cls=stream_cls)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/openai/_base_client.py", line 1461, in request
    return await self._request(
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/openai/_base_client.py", line 1510, in _request
    return await self._retry_request(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/openai/_base_client.py", line 1595, in _retry_request
    return await self._request(
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/openai/_base_client.py", line 1510, in _request
    return await self._retry_request(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/openai/_base_client.py", line 1595, in _retry_request
    return await self._request(
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/pawlu/Documents/scratch/rate_limie_me/.venv/lib/python3.12/site-packages/openai/_base_client.py", line 1520, in _request
    raise APITimeoutError(request=request) from err
openai.APITimeoutError: Request timed out.

[paw-lu]

ah maybe it's because the job is defined outside of the context manager? this seems to work as expected

async def get():
    agent = make_agent()
-  stuff = (agent.run(generate_random_string(2000)) for _ in range(2))

-  for thing in asyncio.as_completed(stuff):
    async for wrap in stamina.retry_context(
        on=openai.APITimeoutError,
        attempts=3,
        wait_initial=1,
        timeout=datetime.timedelta(hours=3),
    ):
        with wrap:
-          await thing
+         await agent.run(generate_random_string(10000))

[hynek]

yeah I think as_completed works a bit unintuitively – have you considered using task groups? See also https://hynek.me/articles/waiting-in-asyncio/

[paw-lu]

have you considered using task groups?

i actually haven't 😅 . we were on an old python and only recently upgraded. (As usual) fantastic article. I'll experiment with it. Thanks a bunch for writing it up and sharing.

[paw-lu]

dumb q. the reason for using as_completed here is because for each task, i want to give it a shot, then retry if it hits some exceptions, catch and continue if it hits others, etc. seems like TaskGroup—to me—seems more like a more controllable alternative to gather? where you can specify all your tasks then you run them all at once? I don't see where you could add a try/except

[hynek]

Right, it depends on the exact semantics you need.

But you could probably wrap the tasks using a RetryingCaller (in your generator). Your problem is simply that the exceptions are not coming from the await it seems. Also your stuff variable is a generator — not sure that has an influence on how as_completed works.