diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bed3d218df..1fc8bc377e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5081,7 +5081,7 @@ async def embeddings( # noqa: PLR0915 ### CALL HOOKS ### - modify incoming data / reject request before calling the model data = await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_dict, data=data, call_type="embeddings" + user_api_key_dict=user_api_key_dict, data=data, call_type="aembedding" ) tasks = [] @@ -5089,7 +5089,7 @@ async def embeddings( # noqa: PLR0915 proxy_logging_obj.during_call_hook( data=data, user_api_key_dict=user_api_key_dict, - call_type="embeddings", + call_type="aembedding", ) ) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 1f4bf806c1..c41ab391df 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -545,7 +545,27 @@ def test_embedding(mock_aembedding, client_no_auth): "input": ["good morning from litellm"], } - response = client_no_auth.post("/v1/embeddings", json=test_data) + pre_call_return_value = { + **test_data, + "metadata": {"source": "unit-test"}, + "proxy_server_request": {"path": "/v1/embeddings"}, + "secret_fields": [], + } + + with patch.object( + litellm.proxy.proxy_server.proxy_logging_obj, + "pre_call_hook", + new=AsyncMock(return_value=pre_call_return_value), + ) as mock_pre_call_hook, patch.object( + litellm.proxy.proxy_server.proxy_logging_obj, + "during_call_hook", + new=AsyncMock(return_value=None), + ) as mock_during_hook, patch.object( + litellm.proxy.proxy_server.proxy_logging_obj, + "post_call_success_hook", + new=AsyncMock(return_value=None), + ): + response = client_no_auth.post("/v1/embeddings", json=test_data) mock_aembedding.assert_called_once_with( model="azure/text-embedding-ada-002", @@ -559,6 +579,16 @@ def test_embedding(mock_aembedding, client_no_auth): result = response.json() print(len(result["data"][0]["embedding"])) assert len(result["data"][0]["embedding"]) > 10 # this usually has len==1536 so + + pre_call_kwargs = mock_pre_call_hook.await_args_list[0].kwargs + assert ( + pre_call_kwargs.get("call_type") == "aembedding" + ), f"expected pre_call_hook to receive call_type='aembedding', got {pre_call_kwargs.get('call_type')}" + + during_call_kwargs = mock_during_hook.await_args_list[0].kwargs + assert ( + during_call_kwargs.get("call_type") == "aembedding" + ), f"expected during_call_hook to receive call_type='aembedding', got {during_call_kwargs.get('call_type')}" except Exception as e: pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}")