phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

support timeouts on http handler

Browse source
This commit is contained in:
Ishaan Jaff 2024-07-30 22:49:00 -07:00
parent 62d4d5e746
commit da494da12c
1 changed files with 39 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

47
litellm/llms/custom_httpx/http_handler.py
View file

@ -80,12 +80,18 @@ class AsyncHTTPHandler:
json: Optional[dict] = None, json: Optional[dict] = None,
params: Optional[dict] = None, params: Optional[dict] = None,
headers: Optional[dict] = None, headers: Optional[dict] = None,
timeout: Optional[Union[float, httpx.Timeout]] = None,
stream: bool = False, stream: bool = False,
): ):
try: try:
req = self.client.build_request( if timeout is not None:
"POST", url, data=data, json=json, params=params, headers=headers # type: ignore req = self.client.build_request(
) "POST", url, data=data, json=json, params=params, headers=headers, timeout=timeout # type: ignore
)
else:
req = self.client.build_request(
"POST", url, data=data, json=json, params=params, headers=headers # type: ignore
)
response = await self.client.send(req, stream=stream) response = await self.client.send(req, stream=stream)
response.raise_for_status() response.raise_for_status()
return response return response
@ -104,6 +110,14 @@ class AsyncHTTPHandler:
) )
finally: finally:
await new_client.aclose() await new_client.aclose()
except httpx.ConnectTimeout:
if data is None:
data = {}
raise litellm.Timeout(
message=f"Connection timed out after {timeout} seconds.",
model=data.get("model"),
llm_provider="litellm-httpx-handler",
)
except httpx.HTTPStatusError as e: except httpx.HTTPStatusError as e:
setattr(e, "status_code", e.response.status_code) setattr(e, "status_code", e.response.status_code)
if stream is True: if stream is True:
@ -192,13 +206,30 @@ class HTTPHandler:
params: Optional[dict] = None, params: Optional[dict] = None,
headers: Optional[dict] = None, headers: Optional[dict] = None,
stream: bool = False, stream: bool = False,
timeout: Optional[Union[float, httpx.Timeout]] = None,
): ):
try:
req = self.client.build_request( if timeout is not None:
"POST", url, data=data, json=json, params=params, headers=headers # type: ignore req = self.client.build_request(
) "POST", url, data=data, json=json, params=params, headers=headers, timeout=timeout # type: ignore
response = self.client.send(req, stream=stream) )
return response else:
req = self.client.build_request(
"POST", url, data=data, json=json, params=params, headers=headers # type: ignore
)
response = self.client.send(req, stream=stream)
return response
except httpx.ConnectTimeout:
if data is None:
data = {}
raise litellm.Timeout(
message=f"Connection timed out after {timeout} seconds.",
model=data.get("model"),
llm_provider="litellm-httpx-handler",
)
except Exception as e:
raise e
def __del__(self) -> None: def __del__(self) -> None:
try: try: