From 8957770e68489ea0e36fda49cf0f4114753433f4 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 28 Oct 2025 18:14:34 -0700 Subject: [PATCH] test_vertex_ai_bge_embedding_with_custom_api_base --- .../llms/vertex_ai/test_bge_embedding.py | 21 +++++++------------ 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py index d7bd07f53ef..636df93b026 100644 --- a/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py +++ b/tests/test_litellm/llms/vertex_ai/test_bge_embedding.py @@ -39,21 +39,16 @@ def test_vertex_ai_bge_embedding_with_custom_api_base(): ): mock_response = MagicMock() mock_response.status_code = 200 + # BGE models return embeddings directly as arrays, not wrapped in objects mock_response.json.return_value = { "predictions": [ - { - "embeddings": { - "values": [0.1, 0.2, 0.3, 0.4, 0.5], - "statistics": {"token_count": 2} - } - }, - { - "embeddings": { - "values": [0.6, 0.7, 0.8, 0.9, 1.0], - "statistics": {"token_count": 2} - } - } - ] + [0.1, 0.2, 0.3, 0.4, 0.5], + [0.6, 0.7, 0.8, 0.9, 1.0] + ], + "deployedModelId": "849506872875548672", + "model": "projects/1060139831167/locations/us-central1/models/baai_bge-small-en-v1.5", + "modelDisplayName": "baai_bge-small-en-v1.5", + "modelVersionId": "1" } mock_post.return_value = mock_response