From cec4535cd8a0bcbb87e746e3a36732d937b5776a Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 16 May 2025 10:33:44 -0700 Subject: [PATCH] fix: fix unit tests --- .../llms/base_llm/realtime/transformation.py | 3 +++ litellm/llms/gemini/realtime/transformation.py | 4 ++++ .../test_gemini_realtime_transformation.py | 18 ++++++++++++++++-- 3 files changed, 23 insertions(+), 2 deletions(-) diff --git a/litellm/llms/base_llm/realtime/transformation.py b/litellm/llms/base_llm/realtime/transformation.py index 759cda4744..db98b7e56a 100644 --- a/litellm/llms/base_llm/realtime/transformation.py +++ b/litellm/llms/base_llm/realtime/transformation.py @@ -72,4 +72,7 @@ class BaseRealtimeConfig(ABC): logging_obj: LiteLLMLoggingObj, realtime_response_transform_input: RealtimeResponseTransformInput, ) -> RealtimeResponseTypedDict: # message sent to setup the realtime session + """ + Keep this state less - leave the state management (e.g. tracking current_output_item_id, current_response_id, current_conversation_id, current_delta_chunks) to the caller. + """ pass diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 4d29a68c2f..abddc766af 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -464,6 +464,9 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): logging_obj: LiteLLMLoggingObj, realtime_response_transform_input: RealtimeResponseTransformInput, ) -> RealtimeResponseTypedDict: + """ + Keep this state less - leave the state management (e.g. tracking current_output_item_id, current_response_id, current_conversation_id, current_delta_chunks) to the caller. + """ try: json_message = json.loads(message) except json.JSONDecodeError: @@ -474,6 +477,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): raise ValueError(f"Invalid JSON message: {message_str}") logging_session_id = logging_obj.litellm_trace_id + current_output_item_id = realtime_response_transform_input[ "current_output_item_id" ] diff --git a/tests/litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py index 4a7ef4d85d..18dd35aba2 100644 --- a/tests/litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -33,7 +33,14 @@ def test_gemini_realtime_transformation_session_created(): session_created_message_str, "gemini-1.5-flash", logging_obj, - session_configuration_request_str, + realtime_response_transform_input={ + "session_configuration_request": session_configuration_request_str, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + }, ) assert transformed_message["response"]["type"] == "session.created" @@ -66,7 +73,14 @@ def test_gemini_realtime_transformation_content_delta(): session_created_message_str, "gemini-1.5-flash", logging_obj, - session_configuration_request_str, + realtime_response_transform_input={ + "session_configuration_request": session_configuration_request_str, + "current_output_item_id": None, + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + }, ) transformed_message = returned_object["response"] assert isinstance(transformed_message, list)