refactor(models): update model fetching logic and dependencies

yufeikang · Dec 19, 2024 · c798e68 · c798e68
1 parent cff9e85
commit c798e68
Show file tree

Hide file tree

Showing 3 changed files with 86 additions and 25 deletions.
diff --git a/app/models.py b/app/models.py
@@ -442,30 +442,35 @@ def __generate_content(
             ),
         )
 
+    def __remove_model_name_prefix(self, model_name):
+        return model_name.replace("models/", "")
+
     async def get_models(self):
-        default_models = _get_default_model_dict("gemini-pro")
+        genai_models = genai.list_models()
+        # filter gemini models
+        genai_models = [
+            model
+            for model in genai_models
+            if model.name.endswith("latest") or model.name.endswith("exp")
+        ]
+        default_models = _get_default_model_dict(
+            self.__remove_model_name_prefix(genai_models[0].name)
+        )
+
         models = [
             {
-                "id": "gemini-pro",
-                "model": "gemini-pro",
-                "name": "Gemini Pro",
-                "provider": "google",
-                "provider_name": "Google",
-                "provider_brand": "google",
-                "context": 16,
-                **_get_model_extra_info(),
-            },
-            {
-                "id": "gemini-1.5-pro",
-                "model": "gemini-1.5-pro-latest",
-                "name": "Gemini 1.5 Pro",
+                "id": self.__remove_model_name_prefix(model.name),
+                "model": self.__remove_model_name_prefix(model.name),
+                "name": model.display_name,
                 "provider": "google",
                 "provider_name": "Google",
                 "provider_brand": "google",
                 "context": 16,
-                **_get_model_extra_info(),
-            },
+                **_get_model_extra_info(self.__remove_model_name_prefix(model.name)),
+            }
+            for model in genai_models
         ]
+
         return {"default_models": default_models, "models": models}
 
 

diff --git a/pdm.lock b/pdm.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -9,7 +9,7 @@ authors = [
 dependencies = [
     "fastapi>=0.95.1",
     "openai>=1.3.6",
-    "google-generativeai>=0.4.0",
+    "google-generativeai>=0.8.3",
     "httpx>=0.24.0",
     "uvicorn[standard]>=0.30.1",
     "pyyaml>=6.0.1",