fix: fix unit tests

This commit is contained in:
Krrish Dholakia
2025-05-16 10:33:44 -07:00
parent 803dda05ed
commit cec4535cd8
3 changed files with 23 additions and 2 deletions
@@ -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
@@ -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"
]
@@ -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)