diff --git a/litellm/llms/snowflake/embedding/transformation.py b/litellm/llms/snowflake/embedding/transformation.py index 45bec1283a3..85dd85d38de 100644 --- a/litellm/llms/snowflake/embedding/transformation.py +++ b/litellm/llms/snowflake/embedding/transformation.py @@ -55,9 +55,11 @@ class SnowflakeEmbeddingConfig(SnowflakeBaseConfig, BaseEmbeddingConfig): litellm_params: dict, ) -> EmbeddingResponse: response_json: Final = raw_response.json() - # convert embeddings to 1d array - for item in response_json["data"]: - item["embedding"] = item["embedding"][0] + # convert embeddings to 1d array if nested + for item in response_json.get("data") or (): + embedding_val = item.get("embedding") + if isinstance(embedding_val, list) and len(embedding_val) > 0 and isinstance(embedding_val[0], list): + item["embedding"] = embedding_val[0] returned_response: Final = EmbeddingResponse.model_validate(_JSON_OBJECT.validate_python(response_json)) returned_response.model = "snowflake/" + (returned_response.model or "") diff --git a/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py b/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py index 8ddbb7c1a6b..61355fe66e0 100644 --- a/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py +++ b/tests/unit/llms/snowflake/test_snowflake_native_endpoints.py @@ -11,7 +11,7 @@ Run: """ import json -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import httpx import pytest @@ -20,8 +20,8 @@ from litellm.llms.snowflake.chat.transformation import ( SnowflakeConfig, is_snowflake_claude_model, ) -from litellm.types.utils import ModelResponse - +from litellm.llms.snowflake.embedding.transformation import SnowflakeEmbeddingConfig +from litellm.types.utils import EmbeddingResponse, ModelResponse # ─── Fixtures ────────────────────────────────────────────────────────────── @@ -778,3 +778,84 @@ class TestAnthropicMultiTurnToolMessages: tool_result = body["messages"][2]["content"][0] assert tool_result["type"] == "tool_result" assert json.loads(tool_result["content"]) == {"result_key": "result_value"} + + +# ─── SnowflakeEmbeddingConfig ────────────────────────────────────────────── + + +class TestSnowflakeEmbeddingConfig: + def setup_method(self): + self.cfg = SnowflakeEmbeddingConfig() + + def test_embedding_url(self): + url = self.cfg.get_complete_url( + api_base=API_BASE, + api_key=PAT_TOKEN, + model="snowflake/snowflake-arctic-embed-m", + optional_params={}, + litellm_params={}, + ) + assert url == f"{API_BASE}/api/v2/cortex/inference:embed" + + def test_embedding_request_transformation(self): + data = self.cfg.transform_embedding_request( + model="snowflake-arctic-embed-m", + input=["Hello world"], + optional_params={}, + headers={}, + ) + assert data == {"text": ["Hello world"], "model": "snowflake-arctic-embed-m"} + + def test_embedding_response_with_1d_vectors(self): + """Standard 1D embedding vectors should remain 1D and not be indexed to a float scalar.""" + raw_response = httpx.Response( + 200, + json={ + "data": [ + {"embedding": [0.1, 0.2, 0.3], "index": 0, "object": "embedding"}, + {"embedding": [0.4, 0.5, 0.6], "index": 1, "object": "embedding"}, + ], + "model": "snowflake-arctic-embed-m", + "usage": {"prompt_tokens": 5, "total_tokens": 5}, + }, + ) + result = self.cfg.transform_embedding_response( + model="snowflake/snowflake-arctic-embed-m", + raw_response=raw_response, + model_response=EmbeddingResponse(), + logging_obj=_mock_logging(), + api_key=PAT_TOKEN, + request_data={}, + optional_params={}, + litellm_params={}, + ) + assert isinstance(result, EmbeddingResponse) + assert result.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert result.data[1]["embedding"] == [0.4, 0.5, 0.6] + assert result.model == "snowflake/snowflake-arctic-embed-m" + + def test_embedding_response_with_nested_2d_vectors(self): + """Nested 2D embedding vectors ([[0.1, 0.2, 0.3]]) should be unwrapped to 1D.""" + raw_response = httpx.Response( + 200, + json={ + "data": [ + {"embedding": [[0.1, 0.2, 0.3]], "index": 0, "object": "embedding"}, + ], + "model": "snowflake-arctic-embed-m", + "usage": {"prompt_tokens": 5, "total_tokens": 5}, + }, + ) + result = self.cfg.transform_embedding_response( + model="snowflake/snowflake-arctic-embed-m", + raw_response=raw_response, + model_response=EmbeddingResponse(), + logging_obj=_mock_logging(), + api_key=PAT_TOKEN, + request_data={}, + optional_params={}, + litellm_params={}, + ) + assert isinstance(result, EmbeddingResponse) + assert result.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert result.model == "snowflake/snowflake-arctic-embed-m"