phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

test triton embeddings

Browse source
This commit is contained in:
Ishaan Jaff 2024-05-10 18:50:34 -07:00
parent 18ed87edc4
commit 0cde9473c9
1 changed files with 17 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

17
litellm/tests/test_embedding.py
View file

@ -516,6 +516,23 @@ def test_voyage_embeddings():
pytest.fail(f"Error occurred: {e}") 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() # test_voyage_embeddings()
# def test_xinference_embeddings(): # def test_xinference_embeddings():
# try: # try: