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

feat(batch_embed_content_transformation.py): support google ai studio /batchEmbedContent endpoint

Browse source 
Allows for multiple strings to be given for embedding
This commit is contained in:
Krrish Dholakia 2024-08-27 19:23:50 -07:00
parent bb42146ffe
commit a6ce27ca29
8 changed files with 303 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

9
litellm/tests/test_embedding.py
View file

@ -687,19 +687,22 @@ async def test_triton_embeddings():
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.parametrize(
"input", ["good morning from litellm", ["good morning from litellm"]] #
)
@pytest.mark.asyncio
async def test_gemini_embeddings(sync_mode):
async def test_gemini_embeddings(sync_mode, input):
try:
litellm.set_verbose = True
if sync_mode:
response = litellm.embedding(
model="gemini/text-embedding-004",
input=["good morning from litellm"],
input=input,
)
else:
response = await litellm.aembedding(
model="gemini/text-embedding-004",
input=["good morning from litellm"],
input=input,
)
print(f"response: {response}")