RomanNum3ral
/
cai
mirror of https://github.com/aliasrobotics/cai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: enforce LiteLLM completion timeout

This commit is contained in:
avir4er 2026-07-03 11:57:22 +08:00
parent cdb6bce77f
commit 3cfbf06313
6 changed files with 193 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
src/cai/continuation.py
View File

@ -11,6 +11,7 @@ from typing import List, Dict, Any, Optional
from rich.console import Console
from cai.config import get_config
from cai.sdk.agents.models.chatcompletions.litellm_adapter import acompletion_with_timeout
logger = logging.getLogger(__name__)
@ -108,9 +109,7 @@ Generate a specific, actionable continuation prompt that:
IMPORTANT: Respond with ONLY the continuation prompt. No explanations, no "Here's a prompt:", just the direct instruction."""
try:
# Use litellm directly, which is how the rest of the codebase handles API calls
import litellm
# Use LiteLLM through CAI's timeout wrapper.
# Enable debug logging for litellm if in debug mode
if logger.isEnabledFor(logging.DEBUG):
logger.debug(f"Generating continuation advice with model: {model_name}")
@ -138,7 +137,7 @@ IMPORTANT: Respond with ONLY the continuation prompt. No explanations, no "Here'
# Make the API call
logger.debug(f"Making API call with kwargs: {kwargs.get('model')}, provider: {kwargs.get('custom_llm_provider', 'default')}")
response = await litellm.acompletion(**kwargs)
response = await acompletion_with_timeout(kwargs, stream=False, model_name=model_name)
# Extract content safely
continuation_prompt = None

5
src/cai/ctr/digest.py
View File

@ -19,6 +19,8 @@ from pathlib import Path
from typing import Dict, List, Tuple, Optional
from rich.console import Console
from cai.sdk.agents.models.chatcompletions.litellm_adapter import acompletion_with_timeout
console = Console()
# Global cache for digest results (per-run caching)
@ -519,7 +521,6 @@ OUTPUT REQUIREMENTS:
- Maximum 350 words"""
# Use LiteLLM for model compatibility (handles alias1, OpenRouter, etc.)
import litellm
model = os.getenv("CAI_CTR_DIGEST_MODEL", "alias1")
@ -549,7 +550,7 @@ OUTPUT REQUIREMENTS:
kwargs["api_key"] = os.getenv("ALIAS_API_KEY", "sk-alias-1234567890")
try:
response = await litellm.acompletion(**kwargs)
response = await acompletion_with_timeout(kwargs, stream=False, model_name=model)
# Extract content (handle reasoning models like alias1/o1 that use reasoning_content)
message = response.choices[0].message

112
src/cai/sdk/agents/models/chatcompletions/litellm_adapter.py
View File

@ -8,6 +8,7 @@ Extracted from openai_chatcompletions.py [F] to reduce monolith size.
from __future__ import annotations
import asyncio
import os
import time
from typing import TYPE_CHECKING, Any, Literal, cast
@ -16,6 +17,7 @@ import litellm
from openai import NOT_GIVEN, NotGiven
from openai.types.responses import Response
from cai.errors import LLMTimeout
from cai.util import get_ollama_api_base
from ..fake_id import FAKE_RESPONSES_ID
@ -32,7 +34,7 @@ if TYPE_CHECKING:
_DEFAULT_MODEL_TIMEOUT = 180.0
def _configured_model_timeout() -> float | None:
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
@ -53,9 +55,9 @@ def _configured_model_timeout() -> float | 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()
def apply_litellm_timeouts(kwargs: dict, *, stream: bool = False) -> dict:
"""Add bounded LiteLLM request timeouts unless the caller already set them."""
timeout = configured_model_timeout()
if timeout is None:
return kwargs
kwargs.setdefault("timeout", timeout)
@ -64,6 +66,84 @@ def _apply_litellm_timeouts(kwargs: dict, *, stream: bool) -> dict:
return kwargs
def _timeout_from_kwargs(kwargs: dict) -> float | None:
"""Return the effective numeric timeout for CAI's outer asyncio guard."""
raw_timeout = kwargs.get("timeout", configured_model_timeout())
if raw_timeout is None:
return None
try:
timeout = float(raw_timeout)
except (TypeError, ValueError):
return configured_model_timeout()
if timeout <= 0:
return None
return timeout
def wrap_stream_with_idle_timeout(stream_obj: Any, *, model_name: str, timeout: float | None = None) -> Any:
"""Bound waits for each streamed chunk.
Some LiteLLM/provider combinations return the stream object quickly, then
stall while the caller awaits the next SSE chunk. ``timeout``/
``stream_timeout`` do not consistently protect that phase, so CAI wraps the
async iterator itself. Non-async-iterable test doubles are returned as-is.
"""
if timeout is None:
timeout = configured_model_timeout()
if timeout is None or not hasattr(stream_obj, "__aiter__"):
return stream_obj
async def _iter_with_timeout():
iterator = stream_obj.__aiter__()
while True:
try:
chunk = await asyncio.wait_for(iterator.__anext__(), timeout=timeout)
except StopAsyncIteration:
return
except asyncio.TimeoutError as exc:
raise LLMTimeout(
f"Timed out waiting for streamed model chunk after {timeout:g}s "
f"[{model_name}]"
) from exc
yield chunk
return _iter_with_timeout()
async def acompletion_with_timeout(
kwargs: dict,
*,
stream: bool = False,
model_name: str | None = None,
) -> Any:
"""Call LiteLLM with CAI request and stream-idle timeouts applied."""
kwargs = apply_litellm_timeouts(kwargs, stream=stream)
timeout = _timeout_from_kwargs(kwargs)
model_label = str(model_name or kwargs.get("model") or "unknown model")
completion_coro = litellm.acompletion(**kwargs)
try:
if timeout is None:
result = await completion_coro
else:
result = await asyncio.wait_for(completion_coro, timeout=timeout)
except asyncio.TimeoutError as exc:
raise LLMTimeout(
f"Timed out waiting for model response after {timeout:g}s [{model_label}]"
) from exc
if stream:
return wrap_stream_with_idle_timeout(result, model_name=model_label, timeout=timeout)
return result
# Backward-compatible private aliases for local/internal imports.
_configured_model_timeout = configured_model_timeout
_apply_litellm_timeouts = apply_litellm_timeouts
_wrap_stream_with_idle_timeout = wrap_stream_with_idle_timeout
_acompletion_with_timeout = acompletion_with_timeout
def _build_response_obj(
model: str,
model_settings: "ModelSettings",
@ -106,10 +186,10 @@ async def fetch_response_litellm_openai(
try:
if stream:
stream_obj = await litellm.acompletion(**kwargs)
stream_obj = await acompletion_with_timeout(kwargs, stream=True, model_name=model_name)
return _build_response_obj(model_name, model_settings, tool_choice, parallel_tool_calls), stream_obj
else:
return await litellm.acompletion(**kwargs)
return await acompletion_with_timeout(kwargs, stream=False, model_name=model_name)
except Exception as e:
error_msg = str(e)
if (
@ -139,10 +219,10 @@ async def fetch_response_litellm_openai(
kwargs["messages"] = messages
if stream:
stream_obj = await litellm.acompletion(**kwargs)
stream_obj = await acompletion_with_timeout(kwargs, stream=True, model_name=model_name)
return _build_response_obj(model_name, model_settings, tool_choice, parallel_tool_calls), stream_obj
else:
return await litellm.acompletion(**kwargs)
return await acompletion_with_timeout(kwargs, stream=False, model_name=model_name)
else:
raise
@ -188,15 +268,15 @@ async def fetch_response_litellm_ollama(
ollama_kwargs = _apply_litellm_timeouts(ollama_kwargs, stream=stream)
call_kwargs = {
**ollama_kwargs,
"api_base": api_base,
"custom_llm_provider": "openai",
}
if stream:
response = _build_response_obj(model_name, model_settings, tool_choice, parallel_tool_calls)
stream_obj = await litellm.acompletion(
**ollama_kwargs, api_base=api_base, custom_llm_provider="openai"
)
stream_obj = await acompletion_with_timeout(call_kwargs, stream=True, model_name=model_name)
return response, stream_obj
else:
return await litellm.acompletion(
**ollama_kwargs,
api_base=api_base,
custom_llm_provider="openai",
)
return await acompletion_with_timeout(call_kwargs, stream=False, model_name=model_name)

17
src/cai/sdk/agents/models/openai_chatcompletions.py
View File

@ -130,6 +130,7 @@ from .chatcompletions.httpx_client import (
verbose_http_retries,
)
from .chatcompletions.litellm_adapter import (
acompletion_with_timeout,
fetch_response_litellm_openai as _fetch_litellm_openai_impl,
fetch_response_litellm_ollama as _fetch_litellm_ollama_impl,
)
@ -4116,10 +4117,14 @@ class OpenAIChatCompletionsModel(Model):
tools=[],
parallel_tool_calls=parallel_tool_calls or False,
)
stream_obj = await litellm.acompletion(**retry_kwargs)
stream_obj = await acompletion_with_timeout(
retry_kwargs, stream=True, model_name=str(self.model)
)
return response, stream_obj
else:
ret = await litellm.acompletion(**retry_kwargs)
ret = await acompletion_with_timeout(
retry_kwargs, stream=False, model_name=str(self.model)
)
return ret
except Exception:
# If retry also fails, raise the original error
@ -4168,11 +4173,15 @@ class OpenAIChatCompletionsModel(Model):
tools=[],
parallel_tool_calls=parallel_tool_calls or False,
)
stream_obj = await litellm.acompletion(**qwen_params)
stream_obj = await acompletion_with_timeout(
qwen_params, stream=True, model_name=str(self.model)
)
return response, stream_obj
else:
# Non-streaming case
ret = await litellm.acompletion(**qwen_params)
ret = await acompletion_with_timeout(
qwen_params, stream=False, model_name=str(self.model)
)
return ret
except Exception as direct_e:
# All approaches failed, log and raise the original error

16
src/cai/tui/components/agent_creator_panel.py
View File

@ -15,7 +15,7 @@ import json
import os
from typing import Dict, List, Any, Optional
from cai.agents.agent_builder import AgentBuilder
import litellm
from cai.sdk.agents.models.chatcompletions.litellm_adapter import acompletion_with_timeout
class AgentCreationConfirmed(Message):
@ -514,14 +514,18 @@ IMPORTANT: The "tools" field in your response must be exactly: {json.dumps(selec
"""
# Use litellm to generate the configuration
response = await litellm.acompletion(
model=os.getenv("CAI_MODEL", "gpt-4"),
messages=[
model_name = os.getenv("CAI_MODEL", "gpt-4")
kwargs = {
"model": model_name,
"messages": [
{"role": "system", "content": "You are an AI agent configuration generator. Always respond with valid JSON only."},
{"role": "user", "content": meta_prompt}
],
temperature=0.7,
max_tokens=2000
"temperature": 0.7,
"max_tokens": 2000,
}
response = await acompletion_with_timeout(
kwargs, stream=False, model_name=model_name
)
# Parse the response

68
tests/sdk/test_litellm_adapter_streaming.py
View File

@ -1,6 +1,10 @@
import asyncio
import pytest
from openai import NOT_GIVEN
from cai.errors import LLMTimeout
from cai.sdk.agents.model_settings import ModelSettings
from cai.sdk.agents.models.chatcompletions.litellm_adapter import (
fetch_response_litellm_openai,
@ -143,3 +147,67 @@ async def test_litellm_model_timeout_uses_env_override(monkeypatch):
assert response is sentinel_response
assert calls[0]["timeout"] == 45.0
assert "stream_timeout" not in calls[0]
@pytest.mark.asyncio
async def test_litellm_nonstream_times_out_when_completion_call_stalls(monkeypatch):
monkeypatch.setenv("CAI_MODEL_TIMEOUT", "0.01")
async def fake_acompletion(**_kwargs):
await asyncio.sleep(60)
return object()
monkeypatch.setattr(
"cai.sdk.agents.models.chatcompletions.litellm_adapter.litellm.acompletion",
fake_acompletion,
)
with pytest.raises(LLMTimeout, match="Timed out waiting for model response"):
await asyncio.wait_for(
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,
),
timeout=0.2,
)
@pytest.mark.asyncio
async def test_litellm_streaming_times_out_when_next_chunk_stalls(monkeypatch):
monkeypatch.setenv("CAI_MODEL_TIMEOUT", "0.01")
class StalledStream:
def __aiter__(self):
return self
async def __anext__(self):
await asyncio.sleep(60)
return object()
async def fake_acompletion(**_kwargs):
return StalledStream()
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,
)
with pytest.raises(LLMTimeout, match="Timed out waiting for streamed model chunk"):
await stream.__anext__()