From 2f01930a5e4bd8a517c78c0c5e2b7846159101b0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E3=83=84=E3=83=B3=E3=83=87=E3=83=AC?= Date: Fri, 20 Mar 2026 15:48:07 +0200 Subject: [PATCH] feat: add timeout to Agent.generate_response() Closes #1 --- rooms/agent.py | 4 ++++ rooms/config.py | 1 + 2 files changed, 5 insertions(+) diff --git a/rooms/agent.py b/rooms/agent.py index 06b2139..bb801ac 100644 --- a/rooms/agent.py +++ b/rooms/agent.py @@ -80,11 +80,15 @@ def generate_response(self, context_messages: List[Dict[str, str]], override_par if self.config.max_tokens: litellm_params["max_tokens"] = self.config.max_tokens + litellm_params["timeout"] = self.config.timeout litellm_params.update(params) response = litellm.completion(**litellm_params) return response.choices[0].message.content.strip() + except litellm.Timeout as e: + logger.error(f"Timeout logic executed for agent '{self.name}' on model '{self.model}': {e}") + return f"[Timeout Error: The model '{self.model}' took too long to respond ({self.config.timeout}s)]" except Exception as e: logger.error(f"Error getting response from agent '{self.name}' on model '{self.model}': {e}") return f"[Error: Could not generate response. Details: {str(e)}]" diff --git a/rooms/config.py b/rooms/config.py index 0067ca3..b9250d4 100644 --- a/rooms/config.py +++ b/rooms/config.py @@ -19,6 +19,7 @@ class AgentConfig(BaseModel): model: str = Field(default="ollama/llama3", description="LiteLLM compatible model name or placeholder for custom") temperature: float = Field(default=0.7, description="Generation temperature") max_tokens: Optional[int] = Field(default=None, description="Max generated tokens") + timeout: int = Field(default=30, description="Timeout in seconds for model generation") color: str = Field(default="blue", description="CLI output color for this agent") custom_function_path: Optional[str] = Field(default=None, description="Path to .py file if model_type is custom_function") custom_function_name: Optional[str] = Field(default=None, description="Name of the python function to call")