Merge pull request #9260 from Grizzly-jobs/fix/voyage-ai-token-usage-tracking

fix: VoyageAI `prompt_token` always empty

litellm/llms/voyage/embedding/transformation.py

@@ -132,7 +132,7 @@ class VoyageEmbeddingConfig(BaseEmbeddingConfig):
         model_response.object = raw_response_json.get("object")
 
         usage = Usage(
-            prompt_tokens=raw_response_json.get("usage", {}).get("prompt_tokens", 0),
+            prompt_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
             total_tokens=raw_response_json.get("usage", {}).get("total_tokens", 0),
         )
         model_response.usage = usage

tests/llm_translation/test_voyage_ai.py

@@ -54,3 +54,27 @@ def test_voyage_ai_embedding_extra_params():
 
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
+
+
+def test_voyage_ai_embedding_prompt_token_mapping():
+    try:
+
+        client = HTTPHandler()
+        litellm.set_verbose = True
+
+        with patch.object(client, "post", return_value=MagicMock(status_code=200, json=lambda: {"usage": {"total_tokens": 120}})) as mock_client:
+            response = litellm.embedding(
+                model="voyage/voyage-3-lite",
+                input=["a"],
+                dimensions=512,
+                input_type="document",
+                client=client,
+            )
+
+            mock_client.assert_called_once()
+            # Assert the response
+            assert response.usage.prompt_tokens == 120
+            assert response.usage.total_tokens == 120
+
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")