Merge pull request #1313 from yu-iskw/support-google-ai-studio

Support Google AI Studio

pr_agent/algo/__init__.py

@@ -38,6 +38,8 @@ MAX_TOKENS = {
     'vertex_ai/gemini-1.5-pro': 1048576,
     'vertex_ai/gemini-1.5-flash': 1048576,
     'vertex_ai/gemma2': 8200,
+    'gemini/gemini-1.5-pro': 1048576,
+    'gemini/gemini-1.5-flash': 1048576,
     'codechat-bison': 6144,
     'codechat-bison-32k': 32000,
     'anthropic.claude-instant-v1': 100000,

pr_agent/algo/ai_handlers/litellm_ai_handler.py

@@ -83,6 +83,11 @@ class LiteLLMAIHandler(BaseAiHandler):
             litellm.vertex_location = get_settings().get(
                 "VERTEXAI.VERTEX_LOCATION", None
             )
+        # Google AI Studio
+        # SEE https://docs.litellm.ai/docs/providers/gemini
+        if get_settings().get("GOOGLE_AI_STUDIO.GEMINI_API_KEY", None):
+          os.environ["GEMINI_API_KEY"] = get_settings().google_ai_studio.gemini_api_key
+
     def prepare_logs(self, response, system, user, resp, finish_reason):
         response_log = response.dict().copy()
         response_log['system'] = system

pr_agent/settings/.secrets_template.toml

@@ -43,6 +43,9 @@ api_base = "" # the base url for your local Llama 2, Code Llama, and other model
 vertex_project = "" # the google cloud platform project name for your vertexai deployment
 vertex_location = "" # the google cloud platform location for your vertexai deployment
 
+[google_ai_studio]
+gemini_api_key = "" # the google AI Studio API key
+
 [github]
 # ---- Set the following only for deployment type == "user"
 user_token = ""  # A GitHub personal access token with 'repo' scope.