Merge pull request #1551 from mito-ds/http_client_timeout

mito-ai: update http request timeout

Author: aarondr77

mito-ai/mito_ai/utils/open_ai_utils.py

@@ -6,7 +6,9 @@
 import json
 from typing import Any, Dict, List, Optional, Type, Final, Union
 from datetime import datetime, timedelta
+import os
 
+from mito_ai.utils.utils import is_running_test
 from pydantic import BaseModel
 from tornado.httpclient import AsyncHTTPClient
 from mito_ai.models import MessageType
@@ -107,8 +109,20 @@ async def get_ai_completion_from_mito_server(
     headers = {
         "Content-Type": "application/json",
     }
-
-    http_client = AsyncHTTPClient(defaults=dict(user_agent="Mito-AI client"))
+    
+    http_client = None
+    if is_running_test():
+        # If we are running in a test environment, setting the request_timeout fails for some reason.
+        http_client = AsyncHTTPClient(defaults=dict(user_agent="Mito-AI client"))
+    else:
+        
+        # The HTTP client timesout after 20 seconds by default. We update this to match the timeout
+        # we give to OpenAI. The OpenAI timeouts are denoted in seconds, wherease the HTTP client
+        # expects milliseconds. We also give the HTTP client a 10 second buffer to account for
+        # the time it takes to send the request, etc.
+        http_client_timeout = timeout * 1000 * max_retries + 10000
+        http_client = AsyncHTTPClient(defaults=dict(user_agent="Mito-AI client"), request_timeout=http_client_timeout)
+    
     try:
         res = await http_client.fetch(
             # Important: DO NOT CHANGE MITO_AI_URL. If you want to use the dev endpoint,