Updated tests

This commit is contained in:
Prathamesh 2025-04-16 16:52:03 +05:30
parent 0a5cb5174c
commit 021fa9fae3

View file

@ -27,58 +27,58 @@ class TestVoyageAI(BaseLLMEmbeddingTest):
} }
# def test_voyage_ai_embedding_extra_params(): def test_voyage_ai_embedding_extra_params():
# try: try:
# client = HTTPHandler() client = HTTPHandler()
# litellm.set_verbose = True litellm.set_verbose = True
# with patch.object(client, "post") as mock_client: with patch.object(client, "post") as mock_client:
# response = litellm.embedding( response = litellm.embedding(
# model="voyage/voyage-3-lite", model="voyage/voyage-3-lite",
# input=["a"], input=["a"],
# dimensions=512, dimensions=512,
# input_type="document", input_type="document",
# client=client, client=client,
# ) )
# mock_client.assert_called_once() mock_client.assert_called_once()
# json_data = json.loads(mock_client.call_args.kwargs["data"]) json_data = json.loads(mock_client.call_args.kwargs["data"])
# print("request data to voyage ai", json.dumps(json_data, indent=4)) print("request data to voyage ai", json.dumps(json_data, indent=4))
# # Assert the request parameters # Assert the request parameters
# assert json_data["input"] == ["a"] assert json_data["input"] == ["a"]
# assert json_data["model"] == "voyage-3-lite" assert json_data["model"] == "voyage-3-lite"
# assert json_data["output_dimension"] == 512 assert json_data["output_dimension"] == 512
# assert json_data["input_type"] == "document" assert json_data["input_type"] == "document"
# except Exception as e: except Exception as e:
# pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
# def test_voyage_ai_embedding_prompt_token_mapping(): def test_voyage_ai_embedding_prompt_token_mapping():
# try: try:
# client = HTTPHandler() client = HTTPHandler()
# litellm.set_verbose = True litellm.set_verbose = True
# with patch.object(client, "post", return_value=MagicMock(status_code=200, json=lambda: {"usage": {"total_tokens": 120}})) as mock_client: with patch.object(client, "post", return_value=MagicMock(status_code=200, json=lambda: {"usage": {"total_tokens": 120}})) as mock_client:
# response = litellm.embedding( response = litellm.embedding(
# model="voyage/voyage-3-lite", model="voyage/voyage-3-lite",
# input=["a"], input=["a"],
# dimensions=512, dimensions=512,
# input_type="document", input_type="document",
# client=client, client=client,
# ) )
# mock_client.assert_called_once() mock_client.assert_called_once()
# # Assert the response # Assert the response
# assert response.usage.prompt_tokens == 120 assert response.usage.prompt_tokens == 120
# assert response.usage.total_tokens == 120 assert response.usage.total_tokens == 120
# except Exception as e: except Exception as e:
# pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
### Rerank Tests ### Rerank Tests