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

test(test_embedding.py): add test for passing extra headers via embedding

Browse source
This commit is contained in:
Krrish Dholakia 2024-11-23 00:32:40 +05:30
parent d788c3c37f
commit 1a3fb18a64
1 changed files with 17 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

17
tests/local_testing/test_embedding.py
View file

@ -1080,3 +1080,20 @@ def test_cohere_img_embeddings(input, input_type):
assert response.usage.prompt_tokens_details.image_tokens > 0
else:
assert response.usage.prompt_tokens_details.text_tokens > 0
def test_embedding_with_extra_headers():
input = ["hello world"]
from litellm.llms.custom_httpx.http_handler import HTTPHandler
client = HTTPHandler()
with patch.object(client, "post") as mock_post:
embedding(
model="cohere/embed-english-v3.0",
input=input,
extra_headers={"my-test-param": "hello-world"},
client=client,
)
assert "my-test-param" in mock_post.call_args.kwargs["headers"]