From 88798d4b3db4d79a767b702820c9ca3bf74fa8b6 Mon Sep 17 00:00:00 2001 From: Lucas Raschek Date: Fri, 14 Mar 2025 18:14:41 +0100 Subject: [PATCH] Add unit test --- tests/llm_translation/test_voyage_ai.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/llm_translation/test_voyage_ai.py b/tests/llm_translation/test_voyage_ai.py index ca49d53c393..d1c01024a05 100644 --- a/tests/llm_translation/test_voyage_ai.py +++ b/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}") \ No newline at end of file