mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-11 17:05:43 +00:00
Andrey
4dbbec233c
* added oauth mcp to docs * added azure ai/grok-4 model family * Revert "added oauth mcp to docs" This reverts commit 950b7cef44f14b2db1429f6fbd32548a7c95d325. * fix: arize ui integration * need to remove a file This reverts commit d6c877b73ac763464f204b77135f3786342373b7. * snowflake support PAT, account_id and embeddings * format * test embeddings * format * complete test * fix: add arize from ui * updated clarifai functions to openai compatible (#15615) * fix: npm build errors * update tests * SnowflakeBaseConfig moved to utils * rename pat_key => api_ke * key_type=PAT => 'pat/key' * fix if api_key is None * doc update * doc update --------- Co-authored-by: mubashir1osmani <mubashir.osmani777@gmail.com> Co-authored-by: Mubashir Osmani <ilikewafflesomcuh@gmail.com> Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com> Co-authored-by: mogith-pn <143642606+mogith-pn@users.noreply.github.com>
97 lines
2.7 KiB
Python
97 lines
2.7 KiB
Python
import os
|
|
import json
|
|
import copy
|
|
|
|
from unittest.mock import patch
|
|
|
|
import litellm
|
|
|
|
model_name = "snowflake-arctic-embed"
|
|
|
|
embed_response = {
|
|
"object": "list",
|
|
"data": [
|
|
{
|
|
"object": "embedding",
|
|
"embedding": [[0.1, 0.2, 0.3]],
|
|
"index": 0,
|
|
}
|
|
],
|
|
"model": model_name,
|
|
"usage": {"total_tokens": 4},
|
|
}
|
|
|
|
|
|
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
|
|
def test_snowflake_jwt_account_id(mock_post):
|
|
mock_post().json.return_value = copy.deepcopy(embed_response)
|
|
|
|
response = litellm.embedding(
|
|
f"snowflake/{model_name}",
|
|
input=["document"],
|
|
api_key="00000",
|
|
account_id="AAAA-BBBB",
|
|
)
|
|
assert len(response.data) == 1
|
|
assert response.data[0]["embedding"] == [0.1, 0.2, 0.3]
|
|
|
|
# check request
|
|
post_kwargs = mock_post.call_args_list[-1][1]
|
|
body = json.loads(post_kwargs["data"])
|
|
assert body["model"] == model_name
|
|
assert body["text"][0] == "document"
|
|
|
|
# JWT key was used
|
|
assert "00000" in post_kwargs["headers"]["Authorization"]
|
|
# account id was used
|
|
assert "AAAA-BBBB" in post_kwargs["url"]
|
|
# is embedding
|
|
assert post_kwargs["url"].endswith("cortex/inference:embed")
|
|
|
|
|
|
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
|
|
def test_snowflake_pat_key_account_id(mock_post):
|
|
mock_post().json.return_value = copy.deepcopy(embed_response)
|
|
|
|
response = litellm.embedding(
|
|
f"snowflake/{model_name}",
|
|
input=["document"],
|
|
api_key="pat/xxxxx",
|
|
account_id="AAAA-BBBB",
|
|
)
|
|
assert len(response.data) == 1
|
|
assert response.data[0]["embedding"] == [0.1, 0.2, 0.3]
|
|
|
|
# PAT key was used
|
|
post_kwargs = mock_post.call_args_list[-1][1]
|
|
assert "xxxxx" in post_kwargs["headers"]["Authorization"]
|
|
assert (
|
|
post_kwargs["headers"]["X-Snowflake-Authorization-Token-Type"]
|
|
== "PROGRAMMATIC_ACCESS_TOKEN"
|
|
)
|
|
|
|
# account id was used
|
|
assert "AAAA-BBBB" in post_kwargs["url"]
|
|
|
|
|
|
@patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post")
|
|
def test_snowflake_env(mock_post):
|
|
mock_post().json.return_value = copy.deepcopy(embed_response)
|
|
|
|
os.environ["SNOWFLAKE_ACCOUNT_ID"] = "AAAA-BBBB"
|
|
os.environ["SNOWFLAKE_JWT"] = "00000"
|
|
|
|
response = litellm.embedding(f"snowflake/{model_name}", input=["document"])
|
|
|
|
assert len(response.data) == 1
|
|
assert response.data[0]["embedding"] == [0.1, 0.2, 0.3]
|
|
|
|
# JWT key was used
|
|
post_kwargs = mock_post.call_args_list[-1][1]
|
|
assert "00000" in post_kwargs["headers"]["Authorization"]
|
|
# account id was used
|
|
assert "AAAA-BBBB" in post_kwargs["url"]
|
|
|
|
os.environ.pop("SNOWFLAKE_ACCOUNT_ID", None)
|
|
os.environ.pop("SNOWFLAKE_JWT", None)
|