test_vertexai_embedding_embedding_latest

2026-08-13 02:22:50 +00:00 · 2024-06-12 11:00:27 -07:00
parent 9aa4b6e98c
commit 37d14a5e8c
1 changed files with 4 additions and 3 deletions
@@ -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)