fix: Use valid CallTypes enum value in embeddings endpoint (#16328)
* Fix embeddings endpoint call_type to use valid CallTypes enum value Fixed bug where the `/embeddings` endpoint was passing `call_type="embeddings"` to guardrail hooks, but "embeddings" is not a valid value in the CallTypes enum. Changed to use `call_type="aembedding"` (async embedding) which is the correct CallTypes enum value and matches the route_type used in the same function. Added unit tests to verify: - "embeddings" is not a valid CallTypes enum value - "aembedding" is the correct valid value - The fix prevents ValueError when guardrails are enabled Fixes #16240 * Inline embeddings call type regression check * Ensure embedding test preserves proxy metadata
This commit is contained in:
parent
29e8d857f7
commit
16325024df
@ -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",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@ -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)}")
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user