set timeouts on llm clients

This commit is contained in:
AI Christianson 2025-02-10 11:41:27 -05:00
parent 26b0f1bdd5
commit 149e8e2251
2 changed files with 25 additions and 3 deletions

View File

@ -10,8 +10,8 @@ from langchain_openai import ChatOpenAI
class ChatDeepseekReasoner(ChatOpenAI):
"""ChatDeepseekReasoner with custom overrides for R1/reasoner models."""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def __init__(self, *args, timeout: int = 180, max_retries: int = 5, **kwargs):
super().__init__(*args, timeout=timeout, max_retries=max_retries, **kwargs)
def invocation_params(
self, options: Optional[Dict[str, Any]] = None, **kwargs: Any

View File

@ -20,6 +20,10 @@ known_temp_providers = {
"deepseek",
}
# Constants for API request configuration
LLM_REQUEST_TIMEOUT = 180
LLM_MAX_RETRIES = 5
logger = get_logger(__name__)
@ -51,6 +55,8 @@ def create_deepseek_client(
if is_expert
else (temperature if temperature is not None else 1),
model=model_name,
timeout=LLM_REQUEST_TIMEOUT,
max_retries=LLM_MAX_RETRIES,
)
return ChatOpenAI(
@ -58,6 +64,8 @@ def create_deepseek_client(
base_url=base_url,
temperature=0 if is_expert else (temperature if temperature is not None else 1),
model=model_name,
timeout=LLM_REQUEST_TIMEOUT,
max_retries=LLM_MAX_RETRIES,
)
@ -76,12 +84,16 @@ def create_openrouter_client(
if is_expert
else (temperature if temperature is not None else 1),
model=model_name,
timeout=LLM_REQUEST_TIMEOUT,
max_retries=LLM_MAX_RETRIES,
)
return ChatOpenAI(
api_key=api_key,
base_url="https://openrouter.ai/api/v1",
model=model_name,
timeout=LLM_REQUEST_TIMEOUT,
max_retries=LLM_MAX_RETRIES,
**({"temperature": temperature} if temperature is not None else {}),
)
@ -188,11 +200,17 @@ def create_llm_client(
}
if is_expert:
openai_kwargs["reasoning_effort"] = "high"
return ChatOpenAI(**openai_kwargs)
return ChatOpenAI(**{
**openai_kwargs,
"timeout": LLM_REQUEST_TIMEOUT,
"max_retries": LLM_MAX_RETRIES,
})
elif provider == "anthropic":
return ChatAnthropic(
api_key=config["api_key"],
model_name=model_name,
timeout=LLM_REQUEST_TIMEOUT,
max_retries=LLM_MAX_RETRIES,
**temp_kwargs,
)
elif provider == "openai-compatible":
@ -200,12 +218,16 @@ def create_llm_client(
api_key=config["api_key"],
base_url=config["base_url"],
model=model_name,
timeout=LLM_REQUEST_TIMEOUT,
max_retries=LLM_MAX_RETRIES,
**temp_kwargs,
)
elif provider == "gemini":
return ChatGoogleGenerativeAI(
api_key=config["api_key"],
model=model_name,
timeout=LLM_REQUEST_TIMEOUT,
max_retries=LLM_MAX_RETRIES,
**temp_kwargs,
)
else: