scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 66

Fix websocker close

Browse source
This commit is contained in:
Sameer Kankute 2026-03-12 09:37:58 +05:30
parent 401f107210
commit 2738715ec0
5 changed files with 34 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
.circleci/config.yml
View file

@ -4781,12 +4781,12 @@ workflows:
only:
- main
- /litellm_.*/
- publish_proxy_extras:
filters:
branches:
only:
- main
- /litellm_release_day_.*/
# - publish_proxy_extras:
# filters:
# branches:
# only:
# - main
# - /litellm_release_day_.*/
- publish_to_pypi:
requires:
- mypy_linting
@ -4840,5 +4840,5 @@ workflows:
- proxy_build_from_pip_tests
- proxy_pass_through_endpoint_tests
- check_code_and_doc_quality
- publish_proxy_extras
# - publish_proxy_extras
- guardrails_testing

26
tests/image_gen_tests/test_image_generation.py
View file

@ -272,19 +272,19 @@ class TestRunwaymlImageGeneration(BaseImageGenTest):
return {"model": "runwayml/gen4_image"}
class TestAzureOpenAIDalle3(BaseImageGenTest):
def get_base_image_generation_call_args(self) -> dict:
return {
"model": "azure/dall-e-3",
"api_version": "2024-02-01",
"api_base": os.getenv("AZURE_API_BASE"),
"api_key": os.getenv("AZURE_API_KEY"),
"metadata": {
"model_info": {
"base_model": "azure/dall-e-3",
}
},
}
# class TestAzureOpenAIDalle3(BaseImageGenTest):
# def get_base_image_generation_call_args(self) -> dict:
# return {
# "model": "azure/dall-e-3",
# "api_version": "2024-02-01",
# "api_base": os.getenv("AZURE_API_BASE"),
# "api_key": os.getenv("AZURE_API_KEY"),
# "metadata": {
# "model_info": {
# "base_model": "azure/dall-e-3",
# }
# },
# }
@pytest.mark.skip(reason="model EOL")

17
tests/llm_translation/realtime/base_realtime_tests.py
View file

@ -12,7 +12,7 @@ from abc import ABC, abstractmethod
from typing import Optional
import pytest
import websockets
from websockets import ConnectionClosed
sys.path.insert(0, os.path.abspath("../../.."))
@ -33,9 +33,8 @@ class RealTimeWebSocketClient:
self.connection_successful = False
self.close_code = None
self.close_reason = None
# Required by realtime_streaming.py - import exceptions module
from websockets import exceptions as websockets_exceptions
self.exceptions = websockets_exceptions
# Required by realtime_streaming.py - use ConnectionClosed for websockets 15+ compatibility
self.exceptions = type("exceptions", (), {"ConnectionClosed": ConnectionClosed})()
async def accept(self):
"""Accept the WebSocket connection"""
@ -112,7 +111,7 @@ class RealTimeWebSocketClient:
print(f"TEST COMPLETE - Closing connection")
print(f"Total messages received from API: {len(self.messages_received)}")
print(f"{'='*80}\n")
raise websockets.exceptions.ConnectionClosed(None, None)
raise ConnectionClosed(None, None)
def queue_client_message(self, message: str):
"""Queue a message to be sent from 'client' to backend"""
@ -190,7 +189,7 @@ class BaseRealtimeTest(ABC):
api_key=os.environ.get(self.get_api_key_env_var()),
timeout=60
)
except websockets.exceptions.ConnectionClosed:
except ConnectionClosed:
pass
except Exception as e:
print(f"\nException: {type(e).__name__}: {e}\n")
@ -249,7 +248,7 @@ class BaseRealtimeTest(ABC):
query_params=query_params,
timeout=60
)
except websockets.exceptions.ConnectionClosed:
except ConnectionClosed:
pass
except Exception as e:
caught_exception = e
@ -364,7 +363,7 @@ class BaseRealtimeTest(ABC):
print(f"CLOSING CONNECTION")
print(f"Total messages received: {len(self.messages_received)}")
print(f"{'='*80}\n")
raise websockets.exceptions.ConnectionClosed(None, None)
raise ConnectionClosed(None, None)
websocket_client = InteractiveWebSocketClient()
caught_exception = None
@ -382,7 +381,7 @@ class BaseRealtimeTest(ABC):
api_key=os.environ.get(self.get_api_key_env_var()),
timeout=60
)
except websockets.exceptions.ConnectionClosed:
except ConnectionClosed:
pass
except Exception as e:
print(f"\nException: {type(e).__name__}: {e}\n")

9
tests/llm_translation/realtime/test_openai_realtime.py
View file

@ -3,6 +3,7 @@ import sys
from unittest.mock import AsyncMock, MagicMock
import pytest
from websockets import ConnectionClosed
sys.path.insert(
0, os.path.abspath("../..")
@ -86,7 +87,7 @@ async def test_openai_realtime_direct_call_no_intent():
if not self.connection_successful:
await asyncio.sleep(3.0)
raise websockets.exceptions.ConnectionClosed(None, None)
raise ConnectionClosed(None, None)
async def close(self, code=1000, reason=""):
self.close_code = code
@ -106,7 +107,7 @@ async def test_openai_realtime_direct_call_no_intent():
api_key=os.environ.get("OPENAI_API_KEY"),
timeout=60
)
except websockets.exceptions.ConnectionClosed:
except ConnectionClosed:
pass
except Exception as e:
caught_exception = e
@ -220,7 +221,7 @@ async def test_openai_realtime_direct_call_with_intent():
if not self.connection_successful:
await asyncio.sleep(3.0)
raise websockets.exceptions.ConnectionClosed(None, None)
raise ConnectionClosed(None, None)
async def close(self, code=1000, reason=""):
self.close_code = code
@ -246,7 +247,7 @@ async def test_openai_realtime_direct_call_with_intent():
query_params=query_params,
timeout=60
)
except websockets.exceptions.ConnectionClosed:
except ConnectionClosed:
pass
except Exception as e:
caught_exception = e

2
tests/llm_translation/realtime/test_realtime_guardrails_openai.py
View file

@ -220,7 +220,7 @@ async def test_voice_transcript_blocked_by_guardrail():
Simulate a backend-side voice transcription event containing the blocked phrase.
Guardrail must block it - no response.create sent to OpenAI.
"""
from websockets.exceptions import ConnectionClosed
from websockets import ConnectionClosed
guardrail = _make_guardrail(GuardrailEventHooks.realtime_input_transcription)
litellm.callbacks = [guardrail]