From 8dffc8620cc29d4a3518cfd6a674a01f223be6d5 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 12 Jun 2024 10:13:48 -0700 Subject: [PATCH] test test_vertexai_embedding_embedding_latest --- litellm/tests/test_amazing_vertex_completion.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index c7f3aa03fbe..5291f269160 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -816,10 +816,14 @@ def test_vertexai_embedding_embedding_latest(): litellm.set_verbose = True response = embedding( model="vertex_ai/text-embedding-004", - input=["hi"], + input=[ + {"text": "good morning from litellm", "task_type": "RETRIEVAL_DOCUMENT"} + ], dimensions=1, + auto_truncate=True, ) assert len(response.data[0]["embedding"]) == 1 + assert response.usage.prompt_tokens > 0 print(f"response:", response) except litellm.RateLimitError as e: pass