Skip to content

fix: streaming timeout retry + progress keepalive during long inference #1749

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gdeyoung wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: streaming timeout retry + progress keepalive during long inference #1749

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 63 additions & 19 deletions agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -894,15 +894,37 @@ async def call_chat_model(
)

# call model
response, reasoning = await call_data["model"].unified_call(
messages=call_data["messages"],
reasoning_callback=call_data["reasoning_callback"],
response_callback=call_data["response_callback"],
rate_limiter_callback=(
self.rate_limiter_callback if not call_data["background"] else None
),
explicit_caching=call_data["explicit_caching"],
)
# Keepalive heartbeat during long inference
_ka_stop = asyncio.Event()
async def _ka_heartbeat():
_ka_count = 0
while not _ka_stop.is_set():
try:
await asyncio.wait_for(_ka_stop.wait(), timeout=10.0)
except asyncio.TimeoutError:
_ka_count += 1
try:
self.context.log.set_progress(f"Generating... ({_ka_count * 10}s)")
except Exception:
pass
_ka_task = asyncio.create_task(_ka_heartbeat())
try:
response, reasoning = await call_data["model"].unified_call(
messages=call_data["messages"],
reasoning_callback=call_data["reasoning_callback"],
response_callback=call_data["response_callback"],
rate_limiter_callback=(
self.rate_limiter_callback if not call_data["background"] else None
),
explicit_caching=call_data["explicit_caching"],
)
finally:
_ka_stop.set()
_ka_task.cancel()
try:
await _ka_task
except asyncio.CancelledError:
pass

await extension.call_extensions_async(
"chat_model_call_after", self, call_data=call_data, response=response, reasoning=reasoning
Expand Down Expand Up @@ -970,16 +992,38 @@ async def call_chat_model_turn(
if call_data.get(key) is not None:
turn_kwargs[key] = call_data.get(key)

llm_result = await call_data["model"].unified_turn(
messages=call_data["messages"],
reasoning_callback=call_data["reasoning_callback"],
response_callback=call_data["response_callback"],
rate_limiter_callback=(
self.rate_limiter_callback if not call_data["background"] else None
),
explicit_caching=call_data["explicit_caching"],
**turn_kwargs,
)
# Keepalive heartbeat during long inference
_ka_stop_t = asyncio.Event()
async def _ka_heartbeat_t():
_ka_count = 0
while not _ka_stop_t.is_set():
try:
await asyncio.wait_for(_ka_stop_t.wait(), timeout=10.0)
except asyncio.TimeoutError:
_ka_count += 1
try:
self.context.log.set_progress(f"Generating... ({_ka_count * 10}s)")
except Exception:
pass
_ka_task_t = asyncio.create_task(_ka_heartbeat_t())
try:
llm_result = await call_data["model"].unified_turn(
messages=call_data["messages"],
reasoning_callback=call_data["reasoning_callback"],
response_callback=call_data["response_callback"],
rate_limiter_callback=(
self.rate_limiter_callback if not call_data["background"] else None
),
explicit_caching=call_data["explicit_caching"],
**turn_kwargs,
)
finally:
_ka_stop_t.set()
_ka_task_t.cancel()
try:
await _ka_task_t
except asyncio.CancelledError:
pass

downgraded = llm_result.capability.get("builtin_tool_downgrades")
if downgraded:
Expand Down
22 changes: 15 additions & 7 deletions models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
from litellm import embedding
import litellm
import openai
import httpx

from helpers import dotenv
from helpers import settings, images
Expand Down Expand Up @@ -319,6 +320,10 @@ def _is_transient_litellm_error(exc: Exception) -> bool:
getattr(openai, "InternalServerError", Exception),
# Some providers map overloads to ServiceUnavailable-like errors
getattr(openai, "APIStatusError", Exception),
# Streaming socket timeouts (aiohttp.SocketTimeoutError maps to these via LiteLLM transport)
httpx.ReadTimeout,
httpx.ConnectTimeout,
httpx.PoolTimeout,
)
return isinstance(exc, transient_types)

Expand Down Expand Up @@ -631,8 +636,11 @@ async def unified_call(
except Exception as e:
import asyncio

# Retry only if no chunks received and error is transient
if got_any_chunk or not _is_transient_litellm_error(e) or attempt >= max_retries:
# Retry logic: allow retry for transient errors even with partial chunks
is_transient = _is_transient_litellm_error(e)
if attempt >= max_retries:
raise
if not is_transient and got_any_chunk:
raise
attempt += 1
await asyncio.sleep(retry_delay_s)
Expand Down Expand Up @@ -773,11 +781,11 @@ async def unified_turn(
except Exception as e:
import asyncio

if (
got_any_chunk
or not _is_transient_litellm_error(e)
or attempt >= max_retries
):
# Retry logic: allow retry for transient errors even with partial chunks
is_transient = _is_transient_litellm_error(e)
if attempt >= max_retries:
raise
if not is_transient and got_any_chunk:
raise
attempt += 1
await asyncio.sleep(retry_delay_s)
Expand Down