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

Merge pull request #5393 from BerriAI/litellm_gemini_embedding_support

Browse source 
feat(vertex_ai_and_google_ai_studio): Support Google AI Studio Embedding Endpoint
This commit is contained in:
Krish Dholakia 2024-08-28 13:46:28 -07:00 • committed by GitHub
commit d928220ed2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 481 additions and 71 deletions
Download patch file Download diff file Expand all files Collapse all files

27
litellm/tests/test_embedding.py
View file

@ -695,6 +695,33 @@ async def test_triton_embeddings():
pytest.fail(f"Error occurred: {e}")
@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, input):
try:
litellm.set_verbose = True
if sync_mode:
response = litellm.embedding(
model="gemini/text-embedding-004",
input=input,
)
else:
response = await litellm.aembedding(
model="gemini/text-embedding-004",
input=input,
)
print(f"response: {response}")
# stubbed endpoint is setup to return this
assert isinstance(response.data[0]["embedding"], list)
assert response.usage.prompt_tokens > 0
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_databricks_embeddings(sync_mode):