fix: properly connected client timeout option

balrog/client.py

@@ -143,9 +143,9 @@ def _initialize_client(self):
         """Initialize the OpenAI client if not already initialized."""
         if not self._initialized:
             if self.client_name.lower() == "vllm":
-                self.client = OpenAI(api_key="EMPTY", base_url=self.base_url)
+                self.client = OpenAI(api_key="EMPTY", base_url=self.base_url, timeout=self.timeout)
             elif self.client_name.lower() == "openai":
-                self.client = OpenAI()
+                self.client = OpenAI(timeout=self.timeout)
             self._initialized = True
 
     def convert_messages(self, messages):
@@ -275,6 +275,7 @@ def get_completion(self, converted_messages, max_retries=5, delay=5):
                 response = self.model.generate_content(
                     converted_messages,
                     generation_config=self.generation_config,
+                    request_options={"timeout": self.timeout},
                 )
                 return response
             except Exception as e:
@@ -374,7 +375,7 @@ def __init__(self, client_config):
     def _initialize_client(self):
         """Initialize the Claude client if not already initialized."""
         if not self._initialized:
-            self.client = Anthropic()
+            self.client = Anthropic(timeout=self.timeout)
             self._initialized = True
 
     def convert_messages(self, messages):

balrog/config/config.yaml

@@ -31,7 +31,7 @@ client:
   generate_kwargs:
     temperature: 0.0            # Sampling temperature; 0.0 makes the output deterministic
     max_tokens: 4096            # Max tokens to generate in the response
-  timeout: 60                   # Timeout for API requests in seconds
+  timeout: 300                  # Timeout for API requests in seconds
   max_retries: 5                # Max number of retries for failed API calls
   delay: 2                      # Exponential backoff factor between retries in seconds
   alternate_roles: False        # Whether the client requires alternating between the agent and the environment