mirror of https://github.com/aliasrobotics/cai.git
fix: bound LiteLLM model wait time
This commit is contained in:
parent
c78b489edb
commit
cdb6bce77f
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Reference in New Issue