added openai mode for express mode models

app/model_loader.py

@@ -33,11 +33,9 @@ async def fetch_and_parse_models_config() -> Optional[Dict[str, List[str]]]:
                 print("Successfully fetched and parsed model configuration.")
                 
                 # Add [EXPRESS] prefix to express models
-                prefixed_express_models = [f"[EXPRESS] {model_name}" for model_name in data["vertex_express_models"]]
-                
                 return {
                     "vertex_models": data["vertex_models"],
-                    "vertex_express_models": prefixed_express_models
+                    "vertex_express_models": data["vertex_express_models"]
                 }
             else:
                 print(f"ERROR: Fetched model configuration has an invalid structure: {data}")

app/routes/models_api.py

@@ -1,10 +1,10 @@
 import time
-from fastapi import APIRouter, Depends, Request # Added Request
-from typing import List, Dict, Any
+from fastapi import APIRouter, Depends, Request
+from typing import List, Dict, Any, Set
 from auth import get_api_key
 from model_loader import get_vertex_models, get_vertex_express_models, refresh_models_config_cache
-import config as app_config # Import config
-from credentials_manager import CredentialManager # To check its type
+import config as app_config
+from credentials_manager import CredentialManager
 
 router = APIRouter()
 
@@ -12,10 +12,10 @@ router = APIRouter()
 async def list_models(fastapi_request: Request, api_key: str = Depends(get_api_key)):
     await refresh_models_config_cache()
     
-    OPENAI_DIRECT_SUFFIX = "-openai"
-    EXPERIMENTAL_MARKER = "-exp-"
     PAY_PREFIX = "[PAY]"
-    # Access credential_manager from app state
+    EXPRESS_PREFIX = "[EXPRESS] "
+    OPENAI_DIRECT_SUFFIX = "-openai"
+    
     credential_manager_instance: CredentialManager = fastapi_request.app.state.credential_manager
     express_key_manager_instance = fastapi_request.app.state.express_key_manager
 
@@ -25,116 +25,49 @@ async def list_models(fastapi_request: Request, api_key: str = Depends(get_api_k
     raw_vertex_models = await get_vertex_models()
     raw_express_models = await get_vertex_express_models()
     
-    candidate_model_ids = set()
-    raw_vertex_models_set = set(raw_vertex_models)  # For checking origin during prefixing
-
-    if has_express_key:
-        candidate_model_ids.update(raw_express_models)
-        # If *only* express key is available, only express models (and their variants) should be listed.
-        # The current `vertex_model_ids` from remote config might contain non-express models.
-        # The `get_vertex_express_models()` should be the source of truth for express-eligible base models.
-        if not has_sa_creds:
-            # Only list models that are explicitly in the express list.
-            # Suffix generation will apply only to these if they are not gemini-2.0
-            all_model_ids = set(raw_express_models)
-        else:
-            # Both SA and Express are available, combine all known models
-            all_model_ids = set(raw_vertex_models + raw_express_models)
-    elif has_sa_creds:
-        # Only SA creds available, use all vertex_models (which might include express-eligible ones)
-        all_model_ids = set(raw_vertex_models)
-    else:
-        # No credentials available
-        all_model_ids = set()
-    
-    # Create extended model list with variations (search, encrypt, auto etc.)
-    # This logic might need to be more sophisticated based on actual supported features per base model.
-    # For now, let's assume for each base model, we might have these variations.
-    # A better approach would be if the remote config specified these variations.
-    
-    dynamic_models_data: List[Dict[str, Any]] = []
+    final_model_list: List[Dict[str, Any]] = []
+    processed_ids: Set[str] = set()
     current_time = int(time.time())
 
-    # Add base models and their variations
-    for original_model_id in sorted(list(all_model_ids)):
-        current_display_prefix = ""
-        # Only add PAY_PREFIX if the model is not already an EXPRESS model (which has its own prefix)
-        # Apply PAY_PREFIX if SA creds are present, it's a model from raw_vertex_models,
-        # it's not experimental, and not already an EXPRESS model.
-        if has_sa_creds and \
-           original_model_id in raw_vertex_models_set and \
-           EXPERIMENTAL_MARKER not in original_model_id and \
-           not original_model_id.startswith("[EXPRESS]"):
-            current_display_prefix = PAY_PREFIX
+    def add_model_and_variants(base_id: str, prefix: str):
+        """Adds a model and its variants to the list if not already present."""
         
-        base_display_id = f"{current_display_prefix}{original_model_id}"
+        # Define all possible suffixes for a given model
+        suffixes = [""] # For the base model itself
+        if not base_id.startswith("gemini-2.0"):
+            suffixes.extend(["-search", "-encrypt", "-encrypt-full", "-auto"])
+        if "gemini-2.5-flash" in base_id or "gemini-2.5-pro" in base_id:
+            suffixes.extend(["-nothinking", "-max"])
         
-        dynamic_models_data.append({
-            "id": base_display_id, "object": "model", "created": current_time, "owned_by": "google",
-            "permission": [], "root": original_model_id, "parent": None
-        })
-        
-        # Conditionally add common variations (standard suffixes)
-        if not original_model_id.startswith("gemini-2.0"): # Suffix rules based on original_model_id
-            standard_suffixes = ["-search", "-encrypt", "-encrypt-full", "-auto"]
-            for suffix in standard_suffixes:
-                # Suffix is applied to the original model ID part
-                suffixed_model_part = f"{original_model_id}{suffix}"
-                # Then the whole thing is prefixed
-                final_suffixed_display_id = f"{current_display_prefix}{suffixed_model_part}"
-                
-                # Check if this suffixed ID is already in all_model_ids (unlikely with prefix) or already added
-                if final_suffixed_display_id not in all_model_ids and not any(m['id'] == final_suffixed_display_id for m in dynamic_models_data):
-                    dynamic_models_data.append({
-                        "id": final_suffixed_display_id, "object": "model", "created": current_time, "owned_by": "google",
-                        "permission": [], "root": original_model_id, "parent": None
-                    })
-        
-        # Apply special suffixes for models starting with "gemini-2.5-flash" or containing "gemini-2.5-pro"
-        # This includes both regular and EXPRESS versions
-        if "gemini-2.5-flash" in original_model_id or "gemini-2.5-pro" in original_model_id: # Suffix rules based on original_model_id
-            special_thinking_suffixes = ["-nothinking", "-max"]
-            for special_suffix in special_thinking_suffixes:
-                suffixed_model_part = f"{original_model_id}{special_suffix}"
-                final_special_suffixed_display_id = f"{current_display_prefix}{suffixed_model_part}"
+        # Add the openai variant for all models
+        suffixes.append(OPENAI_DIRECT_SUFFIX)
 
-                if final_special_suffixed_display_id not in all_model_ids and not any(m['id'] == final_special_suffixed_display_id for m in dynamic_models_data):
-                    dynamic_models_data.append({
-                        "id": final_special_suffixed_display_id, "object": "model", "created": current_time, "owned_by": "google",
-                        "permission": [], "root": original_model_id, "parent": None
-                    })
+        for suffix in suffixes:
+            model_id_with_suffix = f"{base_id}{suffix}"
+            
+            # Experimental models have no prefix
+            final_id = f"{prefix}{model_id_with_suffix}" if "-exp-" not in base_id else model_id_with_suffix
 
-        # Ensure uniqueness again after adding suffixes
-        # Add OpenAI direct variations if SA creds are available
-        # Add OpenAI direct variations for SA credentials
-        if has_sa_creds:
-            for base_model_id_for_openai in raw_vertex_models:
-                display_model_id = ""
-                if EXPERIMENTAL_MARKER in base_model_id_for_openai:
-                    display_model_id = f"{base_model_id_for_openai}{OPENAI_DIRECT_SUFFIX}"
-                else:
-                    display_model_id = f"{PAY_PREFIX}{base_model_id_for_openai}{OPENAI_DIRECT_SUFFIX}"
-                
-                if display_model_id and not any(m['id'] == display_model_id for m in dynamic_models_data):
-                    dynamic_models_data.append({
-                        "id": display_model_id, "object": "model", "created": current_time, "owned_by": "google",
-                        "permission": [], "root": base_model_id_for_openai, "parent": None
-                    })
+            if final_id not in processed_ids:
+                final_model_list.append({
+                    "id": final_id,
+                    "object": "model",
+                    "created": current_time,
+                    "owned_by": "google",
+                    "permission": [],
+                    "root": base_id,
+                    "parent": None
+                })
+                processed_ids.add(final_id)
 
-        # Add OpenAI direct variations for Express keys
-        if has_express_key:
-            EXPRESS_PREFIX = "[EXPRESS] "
-            for base_model_id_for_express_openai in raw_express_models:
-                # Express models are prefixed with [EXPRESS]
-                display_model_id = f"{EXPRESS_PREFIX}{base_model_id_for_express_openai}{OPENAI_DIRECT_SUFFIX}"
-                
-                if display_model_id and not any(m['id'] == display_model_id for m in dynamic_models_data):
-                    dynamic_models_data.append({
-                        "id": display_model_id, "object": "model", "created": current_time, "owned_by": "google",
-                        "permission": [], "root": base_model_id_for_express_openai, "parent": None
-                    })
-    # final_models_data_map = {m["id"]: m for m in dynamic_models_data}
-    # model_list = list(final_models_data_map.values())
-    # model_list.sort()
-    
-    return {"object": "list", "data": sorted(dynamic_models_data, key=lambda x: x['id'])}
+    # Process Express Key models first
+    if has_express_key:
+        for model_id in raw_express_models:
+            add_model_and_variants(model_id, EXPRESS_PREFIX)
+
+    # Process Service Account (PAY) models, they have lower priority
+    if has_sa_creds:
+        for model_id in raw_vertex_models:
+            add_model_and_variants(model_id, PAY_PREFIX)
+
+    return {"object": "list", "data": sorted(final_model_list, key=lambda x: x['id'])}