Merge pull request #158 from chriseckinger/patch-1

Added support for gemini-2.0-flash in _model_info.py
PySpur-Dev · Feb 13, 2025 · 00d8498 · 00d8498
2 parents 4606023 + 8bc7199
commit 00d8498
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/backend/app/nodes/llm/_model_info.py b/backend/app/nodes/llm/_model_info.py
@@ -98,6 +98,7 @@ class LLMModels(str, Enum):
 
     # Google Models
     GEMINI_2_0_FLASH_EXP = "gemini/gemini-2.0-flash-exp"
+    GEMINI_2_0_FLASH = "gemini/gemini-2.0-flash"
     GEMINI_1_5_PRO = "gemini/gemini-1.5-pro"
     GEMINI_1_5_FLASH = "gemini/gemini-1.5-flash"
     GEMINI_1_5_PRO_LATEST = "gemini/gemini-1.5-pro-latest"
@@ -353,6 +354,23 @@ def get_model_info(cls, model_id: str) -> LLMModel | None:
                     }
                 ),
             ),
+            cls.GEMINI_2_0_FLASH.value: LLMModel(  
+                id=cls.GEMINI_2_0_FLASH.value,
+                provider=LLMProvider.GEMINI,
+                name="Gemini 2.0 Flash",  
+                constraints=ModelConstraints(
+                    max_tokens=8192,  
+                    max_temperature=2.0, 
+                ).add_mime_categories(  
+                    {
+                        MimeCategory.IMAGES,
+                        MimeCategory.AUDIO,
+                        MimeCategory.VIDEO,
+                        MimeCategory.DOCUMENTS,
+                        MimeCategory.TEXT,
+                    }
+                ),
+            ),
             # Deepseek Models
             cls.DEEPSEEK_CHAT.value: LLMModel(
                 id=cls.DEEPSEEK_CHAT.value,