diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index 69ddb39ff..8da847b64 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -516,6 +516,23 @@ def test_voyage_embeddings(): pytest.fail(f"Error occurred: {e}") +@pytest.mark.asyncio +async def test_triton_embeddings(): + try: + litellm.set_verbose = True + response = await litellm.aembedding( + model="triton/my-triton-model", + api_base="https://exampleopenaiendpoint-production.up.railway.app/triton/embeddings", + input=["good morning from litellm"], + ) + print(f"response: {response}") + + # stubbed endpoint is setup to return this + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + except Exception as e: + pytest.fail(f"Error occurred: {e}") + + # test_voyage_embeddings() # def test_xinference_embeddings(): # try: