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

feat(vertex_ai.py): adds support for gemini-pro on vertex ai

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-13 10:24:54 -08:00
parent 7e1a355b77
commit ef7a6e3ae1
6 changed files with 67 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

3
litellm/main.py
View file

@ -390,7 +390,6 @@ def completion(
model=deployment_id
custom_llm_provider="azure"
model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key)
### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ###
if input_cost_per_token is not None and output_cost_per_token is not None:
litellm.register_model({
@ -1136,7 +1135,7 @@ def completion(
)
return response
response = model_response
elif model in litellm.vertex_chat_models or model in litellm.vertex_code_chat_models or model in litellm.vertex_text_models or model in litellm.vertex_code_text_models or custom_llm_provider == "vertex_ai":
elif custom_llm_provider == "vertex_ai":
vertex_ai_project = (litellm.vertex_project
or get_secret("VERTEXAI_PROJECT"))
vertex_ai_location = (litellm.vertex_location