From 7730520fb03f97f06203ff6b12e8284a93e065ac Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 26 Apr 2024 14:56:58 -0700
Subject: [PATCH] fix(router.py): allow passing httpx.timeout to timeout param
 in router

Closes https://github.com/BerriAI/litellm/issues/3162
---
 litellm/tests/test_router.py | 31 +++++++++++++++++++++++++++++++
 litellm/types/router.py      |  7 +++++--
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/litellm/tests/test_router.py b/litellm/tests/test_router.py
index 7beb1d67c7b..26843b50b94 100644
--- a/litellm/tests/test_router.py
+++ b/litellm/tests/test_router.py
@@ -14,10 +14,41 @@ from litellm.router import Deployment, LiteLLM_Params, ModelInfo
 from concurrent.futures import ThreadPoolExecutor
 from collections import defaultdict
 from dotenv import load_dotenv
+import os, httpx
 
 load_dotenv()
 
 
+@pytest.mark.parametrize(
+    "timeout", [10, 1.0, httpx.Timeout(timeout=300.0, connect=20.0)]
+)
+def test_router_timeout_init(timeout):
+    """
+    Allow user to pass httpx.Timeout
+
+    related issue - https://github.com/BerriAI/litellm/issues/3162
+    """
+
+    router = Router(
+        model_list=[
+            {
+                "model_name": "test-model",
+                "litellm_params": {
+                    "model": "azure/chatgpt-v-2",
+                    "api_key": os.getenv("AZURE_API_KEY"),
+                    "api_base": os.getenv("AZURE_API_BASE"),
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "timeout": timeout,
+                },
+            }
+        ]
+    )
+
+    router.completion(
+        model="test-model", messages=[{"role": "user", "content": "Hey!"}]
+    )
+
+
 def test_exception_raising():
     # this tests if the router raises an exception when invalid params are set
     # in this test both deployments have bad keys - Keep this test. It validates if the router raises the most recent exception
diff --git a/litellm/types/router.py b/litellm/types/router.py
index c5ec47091a1..7fa15ac365e 100644
--- a/litellm/types/router.py
+++ b/litellm/types/router.py
@@ -1,5 +1,5 @@
 from typing import List, Optional, Union, Dict, Tuple, Literal
-
+import httpx
 from pydantic import BaseModel, validator
 from .completion import CompletionRequest
 from .embedding import EmbeddingRequest
@@ -104,7 +104,9 @@ class LiteLLM_Params(BaseModel):
     api_key: Optional[str] = None
     api_base: Optional[str] = None
     api_version: Optional[str] = None
-    timeout: Optional[Union[float, str]] = None  # if str, pass in as os.environ/
+    timeout: Optional[Union[float, str, httpx.Timeout]] = (
+        None  # if str, pass in as os.environ/
+    )
     stream_timeout: Optional[Union[float, str]] = (
         None  # timeout when making stream=True calls, if str, pass in as os.environ/
     )
@@ -154,6 +156,7 @@ class LiteLLM_Params(BaseModel):
 
     class Config:
         extra = "allow"
+        arbitrary_types_allowed = True
 
     def __contains__(self, key):
         # Define custom behavior for the 'in' operator