Add support for Gemini 2.0-flash and 2.5-pro models

Signed-off-by: JR <jrchang11015043@gmail.com>

Author: RayChang987

README.md

@@ -49,12 +49,11 @@ For more information, please refer to this [link](https://docs.docker.com/refere
 You can specify the Gemini model version using the environment variable `GOOGLE_GEMINI` in your `.env` file.
 The following models are supported:
 
-| Environment Value | Model Name          | Description |
-|------------------|------------------|-------------|
-| `1_pro`          | `gemini-pro`     | Legacy Gemini 1 Pro model (Vertex AI / Generative AI Studio). |
-| `1.5_flash`      | `gemini-1.5-flash` | Lightweight, faster model suitable for low-latency tasks. |
-| `1.5_pro`        | `gemini-1.5-pro` | More capable model for complex reasoning and higher-quality outputs. |
-| `2.5_flash`      | `gemini-2.5-flash` | Latest generation, faster and more accurate than 1.5_flash. |
+| Environment Value | Model Name | Description |
+| :--- | :--- | :--- |
+| `2.0_flash` | `gemini-2.0-flash` | Next-generation lightweight model with improved speed and efficiency. |
+| `2.5_flash` | `gemini-2.5-flash` | Latest generation, faster and more accurate than 1.5_flash. |
+| `2.5_pro` | `gemini-2.5-pro` | Latest generation, state-of-the-art model for the most demanding tasks. |
 
 Set the model by updating your `.env` file:
 ```bash

backend/src/api/routers/chains.py

@@ -54,14 +54,18 @@
     llm = ChatOllama(model=model_name, temperature=llm_temp)
 
 elif os.getenv("LLM_MODEL") == "gemini":
-    if os.getenv("GOOGLE_GEMINI") == "1_pro":
-        llm = ChatGoogleGenerativeAI(model="gemini-pro", temperature=llm_temp)
-    elif os.getenv("GOOGLE_GEMINI") == "1.5_flash":
-        llm = ChatVertexAI(model_name="gemini-1.5-flash", temperature=llm_temp)
-    elif os.getenv("GOOGLE_GEMINI") == "1.5_pro":
-        llm = ChatVertexAI(model_name="gemini-1.5-pro", temperature=llm_temp)
-    elif os.getenv("GOOGLE_GEMINI") == "2.5_flash":
+    gemini_model = os.getenv("GOOGLE_GEMINI")
+    if gemini_model in {"1_pro", "1.5_flash", "1.5_pro"}:
+        raise ValueError(
+            f"The selected Gemini model '{gemini_model}' (version 1.0–1.5) is disabled. "
+            "Please upgrade to version 2.0 or higher (e.g., 2.0_flash, 2.5_pro)."
+        )
+    elif gemini_model == "2.0_flash":
+        llm = ChatVertexAI(model_name="gemini-2.0-flash", temperature=llm_temp)
+    elif gemini_model == "2.5_flash":
         llm = ChatVertexAI(model_name="gemini-2.5-flash", temperature=llm_temp)
+    elif gemini_model == "2.5_pro":
+        llm = ChatVertexAI(model_name="gemini-2.5-pro", temperature=llm_temp)
     else:
         raise ValueError("GOOGLE_GEMINI environment variable not set to a valid value.")
 

backend/src/api/routers/graphs.py

@@ -75,14 +75,18 @@
     llm = ChatOllama(model=model_name, temperature=llm_temp)
 
 elif os.getenv("LLM_MODEL") == "gemini":
-    if os.getenv("GOOGLE_GEMINI") == "1_pro":
-        llm = ChatGoogleGenerativeAI(model="gemini-pro", temperature=llm_temp)
-    elif os.getenv("GOOGLE_GEMINI") == "1.5_flash":
-        llm = ChatVertexAI(model_name="gemini-1.5-flash", temperature=llm_temp)
-    elif os.getenv("GOOGLE_GEMINI") == "1.5_pro":
-        llm = ChatVertexAI(model_name="gemini-1.5-pro", temperature=llm_temp)
-    elif os.getenv("GOOGLE_GEMINI") == "2.5_flash":
+    gemini_model = os.getenv("GOOGLE_GEMINI")
+    if gemini_model in {"1_pro", "1.5_flash", "1.5_pro"}:
+        raise ValueError(
+            f"The selected Gemini model '{gemini_model}' (version 1.0–1.5) is disabled. "
+            "Please upgrade     to version 2.0 or higher (e.g., 2.0_flash, 2.5_pro)."
+        )
+    elif gemini_model == "2.0_flash":
+        llm = ChatVertexAI(model_name="gemini-2.0-flash", temperature=llm_temp)
+    elif gemini_model == "2.5_flash":
         llm = ChatVertexAI(model_name="gemini-2.5-flash", temperature=llm_temp)
+    elif gemini_model == "2.5_pro":
+        llm = ChatVertexAI(model_name="gemini-2.5-pro", temperature=llm_temp)
     else:
         raise ValueError("GOOGLE_GEMINI environment variable not set to a valid value.")