From 37d14a5e8c91fa74cbc897b2f0f5d3676a98e7dc Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 12 Jun 2024 11:00:27 -0700 Subject: [PATCH] test_vertexai_embedding_embedding_latest --- litellm/tests/test_amazing_vertex_completion.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index 5291f26916..daa8756da8 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -814,14 +814,15 @@ def test_vertexai_embedding_embedding_latest(): try: load_vertex_ai_credentials() litellm.set_verbose = True + response = embedding( model="vertex_ai/text-embedding-004", - input=[ - {"text": "good morning from litellm", "task_type": "RETRIEVAL_DOCUMENT"} - ], + input=["hi"], dimensions=1, auto_truncate=True, + task_type="RETRIEVAL_QUERY", ) + assert len(response.data[0]["embedding"]) == 1 assert response.usage.prompt_tokens > 0 print(f"response:", response)