From 4890aa56a64332ccf5550c551caaf0d445177bc2 Mon Sep 17 00:00:00 2001
From: Hemang <hemang@invariantlabs.ai>
Date: Fri, 14 Feb 2025 17:23:30 +0100
Subject: [PATCH] Add timeout to httpx.AsyncClient

---
 proxy/routes/anthropic.py | 4 ++--
 proxy/routes/open_ai.py   | 4 ++--
 proxy/utils/constants.py  | 2 ++
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/proxy/routes/anthropic.py b/proxy/routes/anthropic.py
index f2f0a2b..b407c6b 100644
--- a/proxy/routes/anthropic.py
+++ b/proxy/routes/anthropic.py
@@ -5,7 +5,7 @@ from typing import Any
 
 import httpx
 from fastapi import APIRouter, Depends, Header, HTTPException, Request
-from utils.constants import IGNORED_HEADERS
+from utils.constants import CLIENT_TIMEOUT, IGNORED_HEADERS
 from utils.explorer import push_trace
 
 proxy = APIRouter()
@@ -50,7 +50,7 @@ async def anthropic_proxy(
     request_body_json = json.loads(request_body)
 
     anthropic_url = f"https://api.anthropic.com/{endpoint}"
-    client = httpx.AsyncClient()
+    client = httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT))
 
     anthropic_request = client.build_request(
         "POST", anthropic_url, headers=headers, data=request_body
diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py
index 12c6b56..8bce55b 100644
--- a/proxy/routes/open_ai.py
+++ b/proxy/routes/open_ai.py
@@ -6,7 +6,7 @@ from typing import Any
 import httpx
 from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response
 from starlette.responses import StreamingResponse
-from utils.constants import IGNORED_HEADERS
+from utils.constants import CLIENT_TIMEOUT, IGNORED_HEADERS
 from utils.explorer import push_trace
 
 ALLOWED_OPEN_AI_ENDPOINTS = {"chat/completions"}
@@ -54,7 +54,7 @@ async def openai_proxy(
     is_streaming = request_body_json.get("stream", False)
     invariant_authorization = request.headers.get("invariant-authorization")
 
-    client = httpx.AsyncClient()
+    client = httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT))
     open_ai_request = client.build_request(
         "POST",
         f"https://api.openai.com/v1/{endpoint}",
diff --git a/proxy/utils/constants.py b/proxy/utils/constants.py
index f690498..74dd72e 100644
--- a/proxy/utils/constants.py
+++ b/proxy/utils/constants.py
@@ -11,3 +11,5 @@ IGNORED_HEADERS = [
     "x-forwarded-server",
     "x-real-ip",
 ]
+
+CLIENT_TIMEOUT = 60.0
\ No newline at end of file