phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

Fixed azure ad token not being processed properly in embedding models

Browse source
This commit is contained in:
Dmitry Supranovich 2024-02-22 23:23:06 +00:00 • committed by Dmitry Supranovich
parent 0d18f3c0ca
commit 57ebb9582e
4 changed files with 535 additions and 383 deletions
Download patch file Download diff file Expand all files Collapse all files

27
litellm/tests/test_embedding.py
View file

@ -191,6 +191,33 @@ def test_openai_azure_embedding():
pytest.fail(f"Error occurred: {e}")
def test_openai_azure_embedding_optional_arg(mocker):
mocked_create_embeddings = mocker.patch.object(
openai.resources.embeddings.Embeddings,
"create",
return_value=openai.types.create_embedding_response.CreateEmbeddingResponse(
data=[],
model="azure/test",
object="list",
usage=openai.types.create_embedding_response.Usage(
prompt_tokens=1, total_tokens=2
),
),
)
_ = litellm.embedding(
model="azure/test",
input=["test"],
api_version="test",
api_base="test",
azure_ad_token="test",
)
assert mocked_create_embeddings.called_once_with(
model="test", input=["test"], timeout=600
)
assert "azure_ad_token" not in mocked_create_embeddings.call_args.kwargs
# test_openai_azure_embedding()
# test_openai_embedding()