Merge pull request #2142 from vilmar-hillow/azure_embedding_ad_token

Fixed azure ad token not being processed properly in embedding models
This commit is contained in:
Krish Dholakia 2024-03-19 11:51:28 -07:00 committed by GitHub
commit c840fecdeb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 535 additions and 383 deletions

View file

@ -193,6 +193,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()