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

feat(vertex_ai.py): vertex ai gecko text embedding support

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-03 09:48:29 -08:00
parent 6cdb9aede0
commit d9ba8668f4
6 changed files with 154 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

24
litellm/main.py
View file

@ -2486,7 +2486,7 @@ def embedding(
client=client,
aembedding=aembedding,
)
elif model in litellm.cohere_embedding_models:
elif custom_llm_provider == "cohere":
cohere_key = (
api_key
or litellm.cohere_key
@ -2528,6 +2528,28 @@ def embedding(
optional_params=optional_params,
model_response=EmbeddingResponse(),
)
elif custom_llm_provider == "vertex_ai":
vertex_ai_project = (
optional_params.pop("vertex_ai_project", None)
or litellm.vertex_project
or get_secret("VERTEXAI_PROJECT")
)
vertex_ai_location = (
optional_params.pop("vertex_ai_location", None)
or litellm.vertex_location
or get_secret("VERTEXAI_LOCATION")
)
response = vertex_ai.embedding(
model=model,
input=input,
encoding=encoding,
logging_obj=logging,
optional_params=optional_params,
model_response=EmbeddingResponse(),
vertex_project=vertex_ai_project,
vertex_location=vertex_ai_location,
)
elif custom_llm_provider == "oobabooga":
response = oobabooga.embedding(
model=model,