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(gemini.py): add support for completion calls for gemini-pro (google ai studio)

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-24 09:42:45 +05:30
parent 61f41f7b72
commit 1262d89ab3
5 changed files with 272 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

25
litellm/main.py
View file

@ -54,6 +54,7 @@ from .llms import (
oobabooga,
openrouter,
palm,
gemini,
vertex_ai,
maritalk)
from .llms.openai import OpenAIChatCompletion, OpenAITextCompletion
@ -1137,6 +1138,30 @@ def completion(
)
return response
response = model_response
elif custom_llm_provider == "gemini":
gemini_api_key = (
api_key
or get_secret("GEMINI_API_KEY")
or get_secret("PALM_API_KEY") # older palm api key should also work
or litellm.api_key
)
# palm does not support streaming as yet :(
model_response = gemini.completion(
model=model,
messages=messages,
model_response=model_response,
print_verbose=print_verbose,
optional_params=optional_params,
litellm_params=litellm_params,
logger_fn=logger_fn,
encoding=encoding,
api_key=gemini_api_key,
logging_obj=logging,
acompletion=acompletion,
custom_prompt_dict=custom_prompt_dict
)
response = model_response
elif custom_llm_provider == "vertex_ai":
vertex_ai_project = (litellm.vertex_project
or get_secret("VERTEXAI_PROJECT"))