asyncio.wait_for触发TimeoutError致异步生成器耗尽的解决方法

问题根源 #

原代码中每次触发超时后，都会重新调用response_async_generator.__anext__()，但此前的__anext__()协程可能在后台已经完成，导致异步生成器的内部状态被推进（消耗了一个元素）。当再次调用__anext__()时，生成器可能已经没有剩余元素，从而触发asyncio.StopAsyncIteration错误。

解决方案 #

核心是复用同一个__anext__()任务，直到它完成或确定生成器耗尽，避免重复调用__anext__()导致生成器状态异常。

修正后的代码 #

import openai
import asyncio
import os

os.environ['OPENAI_API_KEY'] = "XXXXXXXXXXXXXXXX"
openai.api_key = os.getenv('OPENAI_API_KEY')

def streaming_chat(history):
    return openai.ChatCompletion.acreate(
        model="gpt-3.5-turbo",
        messages=history,
        stream=True
    )

history = [
    {'role': 'user', 'content': "Can you give me three tips about cooking"}
]

# 无需await，streaming_chat返回的是异步生成器
response_async_generator = streaming_chat(history=history)

async def main():
    # 创建第一个获取下一个元素的任务
    next_task = asyncio.create_task(response_async_generator.__anext__())
    
    while True:
        try:
            # 等待任务完成，超时1秒
            next_element = await asyncio.wait_for(next_task, timeout=1.0)
            print(next_element)
            # 元素获取成功，创建新任务取下一个元素
            next_task = asyncio.create_task(response_async_generator.__anext__())
        except asyncio.TimeoutError:
            print("Waiting...")
            # 超时后继续等待同一个任务，不创建新任务
            continue
        except asyncio.StopAsyncIteration:
            print("生成器已耗尽，结束循环")
            break
        except Exception as e:
            print(f"发生错误: {e}")
            break

asyncio.run(main())

关键改动说明 #

1. 复用任务对象：首次创建__anext__()任务后，超时情况下不会重新创建任务，而是继续等待同一个任务完成，避免生成器被意外推进。
2. 修复原代码错误：原代码中await streaming_chat是错误的，因为streaming_chat返回的是异步生成器而非协程，直接赋值即可。
3. 异常处理完善：新增通用异常捕获，处理API调用可能出现的其他错误。

内容的提问来源于stack exchange，提问作者Guillermo Álvarez