From b06b0248ffb69c1121755369bf7551b06b1bf9cc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 23 Nov 2024 14:11:54 +0530 Subject: [PATCH] fix: fix pass through testing --- .../vertex_passthrough_logging_handler.py | 5 ++++ .../test_unit_test_anthropic_pass_through.py | 27 ++----------------- .../test_unit_test_streaming.py | 1 + 3 files changed, 8 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 75a0d04ec..2773979ad 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -153,6 +153,11 @@ class VertexPassthroughLoggingHandler: verbose_proxy_logger.error( "Unable to build complete streaming response for Vertex passthrough endpoint, not logging..." ) + return { + "result": None, + "kwargs": kwargs, + } + kwargs = VertexPassthroughLoggingHandler._create_vertex_response_logging_payload_for_generate_content( litellm_model_response=complete_streaming_response, model=model, diff --git a/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py b/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py index afb77f718..ecd289005 100644 --- a/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py +++ b/tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py @@ -73,7 +73,7 @@ async def test_anthropic_passthrough_handler( start_time = datetime.now() end_time = datetime.now() - await AnthropicPassthroughLoggingHandler.anthropic_passthrough_handler( + result = AnthropicPassthroughLoggingHandler.anthropic_passthrough_handler( httpx_response=mock_httpx_response, response_body=mock_response, logging_obj=mock_logging_obj, @@ -84,30 +84,7 @@ async def test_anthropic_passthrough_handler( cache_hit=False, ) - # Assert that async_success_handler was called - assert mock_logging_obj.async_success_handler.called - - call_args = mock_logging_obj.async_success_handler.call_args - call_kwargs = call_args.kwargs - print("call_kwargs", call_kwargs) - - # Assert required fields are present in call_kwargs - assert "result" in call_kwargs - assert "start_time" in call_kwargs - assert "end_time" in call_kwargs - assert "cache_hit" in call_kwargs - assert "response_cost" in call_kwargs - assert "model" in call_kwargs - assert "standard_logging_object" in call_kwargs - - # Assert specific values and types - assert isinstance(call_kwargs["result"], litellm.ModelResponse) - assert isinstance(call_kwargs["start_time"], datetime) - assert isinstance(call_kwargs["end_time"], datetime) - assert isinstance(call_kwargs["cache_hit"], bool) - assert isinstance(call_kwargs["response_cost"], float) - assert call_kwargs["model"] == "claude-3-opus-20240229" - assert isinstance(call_kwargs["standard_logging_object"], dict) + assert isinstance(result["result"], litellm.ModelResponse) def test_create_anthropic_response_logging_payload(mock_logging_obj): diff --git a/tests/pass_through_unit_tests/test_unit_test_streaming.py b/tests/pass_through_unit_tests/test_unit_test_streaming.py index bbbc465fc..61b71b56d 100644 --- a/tests/pass_through_unit_tests/test_unit_test_streaming.py +++ b/tests/pass_through_unit_tests/test_unit_test_streaming.py @@ -64,6 +64,7 @@ async def test_chunk_processor_yields_raw_bytes(endpoint_type, url_route): litellm_logging_obj = MagicMock() start_time = datetime.now() passthrough_success_handler_obj = MagicMock() + litellm_logging_obj.async_success_handler = AsyncMock() # Capture yielded chunks and perform detailed assertions received_chunks = []