From 2c1592ae73c0b1eb3b0aecdaae636867671b8aae Mon Sep 17 00:00:00 2001 From: Raushankumar0720 Date: Thu, 1 Oct 2026 16:49:08 +0530 Subject: [PATCH 1/2] fix(snowflake): safely preserve 1D vector embeddings in response transformation --- .../snowflake/embedding/transformation.py | 8 +- .../test_snowflake_native_endpoints.py | 91 ++++++++++++++++++- 2 files changed, 93 insertions(+), 6 deletions(-) diff --git a/litellm/llms/snowflake/embedding/transformation.py b/litellm/llms/snowflake/embedding/transformation.py index ec482337e52..901959bc8e8 100644 --- a/litellm/llms/snowflake/embedding/transformation.py +++ b/litellm/llms/snowflake/embedding/transformation.py @@ -54,9 +54,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 344b8e5573d..e3cf82f73eb 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_claude_model, ) -from litellm.types.utils import ModelResponse - +from litellm.llms.snowflake.embedding.transformation import SnowflakeEmbeddingConfig +from litellm.types.utils import EmbeddingResponse, ModelResponse # ─── Fixtures ────────────────────────────────────────────────────────────── @@ -69,6 +69,7 @@ def _make_anthropic_response(content: str = "Hello!") -> httpx.Response: # ─── SnowflakeConfig (OpenAI-compatible) ─────────────────────────────────── + class TestSnowflakeConfigURL: def setup_method(self): self.cfg = SnowflakeConfig() @@ -266,6 +267,7 @@ class TestSnowflakeConfigResponse: # ─── SnowflakeConfig ──────────────────────────────────────── + class TestAnthropicConfigURL: def setup_method(self): self.cfg = SnowflakeConfig() @@ -533,6 +535,7 @@ class TestAnthropicConfigResponse: # ─── Model detection helper ──────────────────────────────────────────────── + class TestIsClaudeModel: def test_claude_model_detected(self): assert _is_claude_model("snowflake/claude-sonnet-4-5") is True @@ -548,6 +551,7 @@ class TestIsClaudeModel: # ─── Anthropic Tool Transformation Tests ────────────────────────────────── + class TestAnthropicToolTransformation: def setup_method(self): self.cfg = SnowflakeConfig() @@ -774,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" From e106c41ec947e3797ea0a4565ca9056c458c2a25 Mon Sep 17 00:00:00 2001 From: Raushankumar0720 Date: Thu, 1 Oct 2026 18:10:12 +0530 Subject: [PATCH 2/2] ci: trigger check re-run