fix: bound LiteLLM model wait time

This commit is contained in:
avir4er 2026-07-03 11:32:18 +08:00
parent c78b489edb
commit cdb6bce77f
2 changed files with 99 additions and 0 deletions

View File

@ -8,6 +8,7 @@ Extracted from openai_chatcompletions.py [F] to reduce monolith size.
from __future__ import annotations
import os
import time
from typing import TYPE_CHECKING, Any, Literal, cast
@ -28,6 +29,41 @@ if TYPE_CHECKING:
from ...model_settings import ModelSettings
_DEFAULT_MODEL_TIMEOUT = 180.0
def _configured_model_timeout() -> float | None:
"""Return CAI's LiteLLM request timeout in seconds.
``CAI_MODEL_TIMEOUT`` is the public name. ``CAI_LLM_TIMEOUT`` is accepted
as a compatibility alias for local configs/scripts. Values <= 0 disable the
injected timeout and defer entirely to LiteLLM/provider defaults.
"""
raw = os.getenv("CAI_MODEL_TIMEOUT")
if raw is None:
raw = os.getenv("CAI_LLM_TIMEOUT")
if raw is None or raw == "":
return _DEFAULT_MODEL_TIMEOUT
try:
timeout = float(raw)
except (TypeError, ValueError):
return _DEFAULT_MODEL_TIMEOUT
if timeout <= 0:
return None
return timeout
def _apply_litellm_timeouts(kwargs: dict, *, stream: bool) -> dict:
"""Add bounded model request timeouts unless the caller already set them."""
timeout = _configured_model_timeout()
if timeout is None:
return kwargs
kwargs.setdefault("timeout", timeout)
if stream:
kwargs.setdefault("stream_timeout", timeout)
return kwargs
def _build_response_obj(
model: str,
model_settings: "ModelSettings",
@ -66,6 +102,8 @@ async def fetch_response_litellm_openai(
too long, truncate all tool_call ids in the messages to 40 characters
and retry once silently.
"""
kwargs = _apply_litellm_timeouts(kwargs, stream=stream)
try:
if stream:
stream_obj = await litellm.acompletion(**kwargs)
@ -148,6 +186,8 @@ async def fetch_response_litellm_ollama(
api_base = get_ollama_api_base()
ollama_kwargs = _apply_litellm_timeouts(ollama_kwargs, stream=stream)
if stream:
response = _build_response_obj(model_name, model_settings, tool_choice, parallel_tool_calls)
stream_obj = await litellm.acompletion(

View File

@ -84,3 +84,62 @@ async def test_litellm_streaming_tool_call_id_retry_opens_one_retry_stream(monke
assert len(calls) == 2
assert kwargs["messages"][0]["tool_call_id"] == long_id[:40]
assert kwargs["messages"][1]["tool_calls"][0]["id"] == long_id[:40]
@pytest.mark.asyncio
async def test_litellm_streaming_applies_default_model_timeout(monkeypatch):
monkeypatch.delenv("CAI_MODEL_TIMEOUT", raising=False)
monkeypatch.delenv("CAI_LLM_TIMEOUT", raising=False)
calls = []
sentinel_stream = object()
async def fake_acompletion(**kwargs):
calls.append(kwargs.copy())
return sentinel_stream
monkeypatch.setattr(
"cai.sdk.agents.models.chatcompletions.litellm_adapter.litellm.acompletion",
fake_acompletion,
)
_response, stream = await fetch_response_litellm_openai(
kwargs={"model": "deepseek/deepseek-v4-pro", "messages": [], "stream": True},
model_name="deepseek/deepseek-v4-pro",
model_settings=ModelSettings(),
tool_choice=NOT_GIVEN,
stream=True,
parallel_tool_calls=False,
)
assert stream is sentinel_stream
assert calls[0]["timeout"] == 180.0
assert calls[0]["stream_timeout"] == 180.0
@pytest.mark.asyncio
async def test_litellm_model_timeout_uses_env_override(monkeypatch):
monkeypatch.setenv("CAI_MODEL_TIMEOUT", "45")
calls = []
sentinel_response = object()
async def fake_acompletion(**kwargs):
calls.append(kwargs.copy())
return sentinel_response
monkeypatch.setattr(
"cai.sdk.agents.models.chatcompletions.litellm_adapter.litellm.acompletion",
fake_acompletion,
)
response = await fetch_response_litellm_openai(
kwargs={"model": "deepseek/deepseek-v4-pro", "messages": [], "stream": False},
model_name="deepseek/deepseek-v4-pro",
model_settings=ModelSettings(),
tool_choice=NOT_GIVEN,
stream=False,
parallel_tool_calls=False,
)
assert response is sentinel_response
assert calls[0]["timeout"] == 45.0
assert "stream_timeout" not in calls[0]