phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 03:34:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(embeddings_handler.py): support async gemini embeddings

Browse source
This commit is contained in:
Krrish Dholakia 2024-08-27 18:31:57 -07:00
parent 6a483a1908
commit 4bb59b7b2c
3 changed files with 109 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

17
litellm/tests/test_embedding.py
View file

@ -686,14 +686,21 @@ async def test_triton_embeddings():
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_gemini_embeddings():
async def test_gemini_embeddings(sync_mode):
try:
litellm.set_verbose = True
response = await litellm.aembedding(
model="gemini/text-embedding-004",
input=["good morning from litellm"],
)
if sync_mode:
response = litellm.embedding(
model="gemini/text-embedding-004",
input=["good morning from litellm"],
)
else:
response = await litellm.aembedding(
model="gemini/text-embedding-004",
input=["good morning from litellm"],
)
print(f"response: {response}")
# stubbed endpoint is setup to return this