add_gemini_frontend

topoteretes · Jan 10, 2025 · 5712e6e · 5712e6e
1 parent 5ab023d
commit 5712e6e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/cognee/api/v1/settings/routers/get_settings_router.py b/cognee/api/v1/settings/routers/get_settings_router.py
@@ -21,7 +21,7 @@ class SettingsDTO(OutDTO):
 
 
 class LLMConfigInputDTO(InDTO):
-    provider: Union[Literal["openai"], Literal["ollama"], Literal["anthropic"]]
+    provider: Union[Literal["openai"], Literal["ollama"], Literal["anthropic"], Literal["gemini"]]
     model: str
     api_key: str
 

diff --git a/cognee/modules/settings/get_settings.py b/cognee/modules/settings/get_settings.py
@@ -13,6 +13,7 @@ class ModelName(Enum):
     openai = "openai"
     ollama = "ollama"
     anthropic = "anthropic"
+    gemini = "gemini"
 
 
 class LLMConfig(BaseModel):
@@ -72,6 +73,10 @@ def get_settings() -> SettingsDict:
             "value": "anthropic",
             "label": "Anthropic",
         },
+        {
+            "value": "gemini",
+            "label": "Gemini",
+        },
     ]
 
     return SettingsDict.model_validate(
@@ -136,6 +141,12 @@ def get_settings() -> SettingsDict:
                             "label": "Claude 3 Haiku",
                         },
                     ],
+                    "gemini": [
+                        {
+                            "value": "gemini-2.0-flash-exp",
+                            "label": "Gemini 2.0 Flash",
+                        },
+                    ],
                 },
             },
             vector_db={