From 6d5661fad992d7ae43ddcaa4b339f2097418e25a Mon Sep 17 00:00:00 2001 From: shixiao-coder Date: Tue, 21 Apr 2026 10:24:57 -0400 Subject: [PATCH] Modify the code to only keep Gemini 3 related setup and remove all Gemini 2.5 setup --- server/__init__.py | 6 +----- server/lib/feature_flags.py | 1 - server/lib/nl/detection/llm_api.py | 16 ++-------------- server/lib/nl/explore/overview.py | 10 ++-------- server/lib/nl/explore/related.py | 10 ++-------- server/lib/utils/gemini_utils.py | 14 +++++--------- 6 files changed, 12 insertions(+), 45 deletions(-) diff --git a/server/__init__.py b/server/__init__.py index 71ffe261b4..f4a9bc7715 100644 --- a/server/__init__.py +++ b/server/__init__.py @@ -33,7 +33,6 @@ from server.lib.disaster_dashboard import get_disaster_dashboard_data from server.lib.feature_flags import BIOMED_NL_FEATURE_FLAG from server.lib.feature_flags import DATA_OVERVIEW_FEATURE_FLAG -from server.lib.feature_flags import ENABLE_GEMINI_3_FLASH from server.lib.feature_flags import ENABLE_NL_AGENT_DETECTOR from server.lib.feature_flags import is_feature_enabled import server.lib.i18n as i18n @@ -435,10 +434,7 @@ def create_app(nl_root=DEFAULT_NL_ROOT): 'palm-api-key') if is_feature_enabled(ENABLE_NL_AGENT_DETECTOR, app): os.environ['GEMINI_API_KEY'] = app.config['LLM_API_KEY'] - if is_feature_enabled(ENABLE_GEMINI_3_FLASH, app): - default_model = "gemini-3-flash-preview" - else: - default_model = "gemini-2.5-flash" + default_model = "gemini-3-flash-preview" app.config['NL_DETECTION_AGENT'] = create_detection_agent( os.environ.get("AGENT_MODEL", default_model), os.environ.get("DC_MCP_URL")) diff --git a/server/lib/feature_flags.py b/server/lib/feature_flags.py index aa34f499a5..7445113794 100644 --- a/server/lib/feature_flags.py +++ b/server/lib/feature_flags.py @@ -34,7 +34,6 @@ ENABLE_STAT_VAR_AUTOCOMPLETE = 'enable_stat_var_autocomplete' ENABLE_NL_AGENT_DETECTOR = 'enable_nl_agent_detector' NEW_RANKING_PAGE = 'new_ranking_page' -ENABLE_GEMINI_3_FLASH = 'enable_gemini_3_flash' USE_V2_API = 'use_v2_api' # This flag controls the switching of detect-and-fulfill API to use v2/resolve from current nl search vars USE_V2_RESOLVE_FOR_NL_SEARCH_VARS = 'use_v2_resolve_for_nl_search_vars' diff --git a/server/lib/nl/detection/llm_api.py b/server/lib/nl/detection/llm_api.py index d2d8964856..7ac39428c1 100644 --- a/server/lib/nl/detection/llm_api.py +++ b/server/lib/nl/detection/llm_api.py @@ -24,14 +24,10 @@ from google.genai import types import json5 -from server.lib.feature_flags import ENABLE_GEMINI_3_FLASH -from server.lib.feature_flags import is_feature_enabled from server.lib.nl.common import counters _GEMINI_3_0_FLASH = 'gemini-3-flash-preview' -_GEMINI_2_5_FLASH = 'gemini-2.5-flash' _API_VERSION_3 = 'v1beta' -_API_VERSION_2 = 'v1' # TODO: Consider tweaking this. And maybe consider passing as url param. _TEMPERATURE = 0.1 @@ -55,13 +51,7 @@ "category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE" }, - ]) - -_GEMINI_2_5_CONFIG = _GEMINI_CONFIG - -_GEMINI_3_CONFIG = types.GenerateContentConfig( - temperature=_TEMPERATURE, - safety_settings=_GEMINI_CONFIG.safety_settings, + ], thinking_config=types.ThinkingConfig(thinking_level="low")) _SKIP_BEGIN_CHARS = ['`', '*'] @@ -197,6 +187,4 @@ def _extract_answer(resp: str) -> str: def detect_model_name() -> tuple[str, str, types.GenerateContentConfig]: - if is_feature_enabled(ENABLE_GEMINI_3_FLASH): - return _GEMINI_3_0_FLASH, _API_VERSION_3, _GEMINI_3_CONFIG - return _GEMINI_2_5_FLASH, _API_VERSION_2, _GEMINI_2_5_CONFIG + return _GEMINI_3_0_FLASH, _API_VERSION_3, _GEMINI_CONFIG diff --git a/server/lib/nl/explore/overview.py b/server/lib/nl/explore/overview.py index 9da2e767c6..b5e1a89684 100644 --- a/server/lib/nl/explore/overview.py +++ b/server/lib/nl/explore/overview.py @@ -20,8 +20,6 @@ from pydantic import ConfigDict from pydantic.alias_generators import to_camel -from server.lib.feature_flags import ENABLE_GEMINI_3_FLASH -from server.lib.feature_flags import is_feature_enabled from server.lib.nl.explore.gemini_prompts import PAGE_OVERVIEW_PROMPT from server.lib.utils.gemini_utils import call_gemini @@ -81,12 +79,8 @@ def generate_page_overview( formatted_page_overview_prompt = PAGE_OVERVIEW_PROMPT.format( initial_query=query, stat_var_titles=stat_var_titles) - if is_feature_enabled(ENABLE_GEMINI_3_FLASH): - overview_gemini_model = _OVERVIEW_GEMINI_3_1_LITE - use_thinking_config = True - else: - overview_gemini_model = _OVERVIEW_GEMINI_2_5_LITE - use_thinking_config = False + overview_gemini_model = _OVERVIEW_GEMINI_3_1_LITE + use_thinking_config = True page_overview = call_gemini(api_key=gemini_api_key, formatted_prompt=formatted_page_overview_prompt, diff --git a/server/lib/nl/explore/related.py b/server/lib/nl/explore/related.py index eb76f7c007..ebc8cc10ec 100644 --- a/server/lib/nl/explore/related.py +++ b/server/lib/nl/explore/related.py @@ -21,8 +21,6 @@ from flask import current_app from pydantic import BaseModel -from server.lib.feature_flags import ENABLE_GEMINI_3_FLASH -from server.lib.feature_flags import is_feature_enabled import server.lib.nl.common.topic as topic import server.lib.nl.common.utils as utils import server.lib.nl.detection.types as dtypes @@ -316,12 +314,8 @@ def generate_follow_up_questions(query: str, formatted_follow_up_questions_prompt = FOLLOW_UP_QUESTIONS_PROMPT.format( initial_query=query, related_topics=related_topics) - if is_feature_enabled(ENABLE_GEMINI_3_FLASH): - gemini_model = _QUESTIONS_GEMINI_3 - use_thinking_config = True - else: - gemini_model = _QUESTIONS_GEMINI_2 - use_thinking_config = False + gemini_model = _QUESTIONS_GEMINI_3 + use_thinking_config = True follow_up_questions = call_gemini( api_key=gemini_api_key, diff --git a/server/lib/utils/gemini_utils.py b/server/lib/utils/gemini_utils.py index c40a97f0c3..91747303a4 100644 --- a/server/lib/utils/gemini_utils.py +++ b/server/lib/utils/gemini_utils.py @@ -19,14 +19,12 @@ from pydantic import BaseModel -def get_gemini_config(schema: Optional[BaseModel] = None, - use_thinking_config: bool = False) -> dict: +def get_gemini_config(schema: Optional[BaseModel] = None) -> dict: config = { "response_mime_type": "application/json", - "response_schema": schema + "response_schema": schema, + "thinking_config": genai.types.ThinkingConfig(thinking_level="low") } if schema else {} - if use_thinking_config: - config["thinking_config"] = genai.types.ThinkingConfig(thinking_level="low") return config @@ -34,15 +32,13 @@ def call_gemini( api_key: str, formatted_prompt: str, gemini_model: str, - schema: Optional[BaseModel] = None, - use_thinking_config: bool = False) -> Optional[Union[BaseModel, str]]: + schema: Optional[BaseModel] = None) -> Optional[Union[BaseModel, str]]: """A helper for all Gemini generations through the Python Gen AI client. Args: api_key: A string representing the API key required for authentication with the Gemini service. formatted_prompt: A string containing the structured prompt or input to be sent to the Gemini model for generation. schema: A Pydantic BaseModel class that defines the expected model's JSON response. gemini_model: A string specifying the name of the Gemini model to utilize. - use_thinking_config: Boolean advising whether to use thinking configuration for Gemini 3. Returns: The output of the call. @@ -50,7 +46,7 @@ def call_gemini( if not api_key or not formatted_prompt: return None - generate_content_config = get_gemini_config(schema, use_thinking_config) + generate_content_config = get_gemini_config(schema) gemini = genai.Client(api_key=api_key) try: