app/vertex/routes/models_api.py

import time
from fastapi import APIRouter, Depends, Request, HTTPException
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from typing import List, Dict, Any, Optional
from app.vertex.auth import get_api_key, validate_api_key
from app.vertex.model_loader import get_vertex_models, get_vertex_express_models, refresh_models_config_cache
import app.vertex.config as app_config
from app.vertex.credentials_manager import CredentialManager
from app.utils.logging import vertex_log
from app.config import settings

router = APIRouter(prefix="/models")

# Setup security
security = HTTPBearer()

@router.get("/list")
async def list_models(credentials: HTTPAuthorizationCredentials = Depends(security)):
    """
    List available models for Vertex generation.
    
    Returns a list of models in OpenAI-compatible format.
    """
    # Validate API key
    api_key = validate_api_key(credentials)
    if not api_key:
        raise HTTPException(status_code=401, detail="Invalid API key")
    
    # Get available models
    vertex_log('info', "Retrieving list of available models")
    standard_models = await get_vertex_models()
    express_models = await get_vertex_express_models()
    
    # Combine and format
    all_models = []
    
    # Format standard models
    for model_name in standard_models:
        all_models.append({
            "id": model_name,
            "object": "model",
            "created": 1677610602,  # Placeholder timestamp
            "owned_by": "google",
            "permission": [],
            "root": model_name,
            "parent": None,
            "pricing": {
                "prompt": "0.0020",  # Placeholder pricing
                "completion": "0.0020"  # Placeholder pricing
            }
        })
    
    # Format express models
    for model_name in express_models:
        all_models.append({
            "id": model_name,
            "object": "model",
            "created": 1677610602,  # Placeholder timestamp
            "owned_by": "google-express",
            "permission": [],
            "root": model_name,
            "parent": None,
            "pricing": {
                "prompt": "0.0010",  # Lower placeholder pricing for express
                "completion": "0.0010"  # Lower placeholder pricing for express
            }
        })
    
    # Add encryption variants if available
    for base_model in standard_models:
        encrypt_model_name = f"{base_model}-encrypt-full"
        all_models.append({
            "id": encrypt_model_name,
            "object": "model",
            "created": 1677610602,  # Placeholder timestamp
            "owned_by": "google-secure",
            "permission": [],
            "root": base_model,
            "parent": base_model,
            "pricing": {
                "prompt": "0.0025",  # Higher placeholder pricing for encrypt
                "completion": "0.0025"  # Higher placeholder pricing for encrypt
            }
        })
    
    vertex_log('info', f"Found {len(all_models)} available models")
    
    return {
        "object": "list",
        "data": all_models
    }

@router.post("/refresh")
async def refresh_models(credentials: HTTPAuthorizationCredentials = Depends(security)):
    """
    Refresh the models configuration cache.
    
    Forces a refresh of the models available from the configuration source.
    """
    # Validate API key
    api_key = validate_api_key(credentials)
    if not api_key:
        raise HTTPException(status_code=401, detail="Invalid API key")
    
    vertex_log('info', "Attempting to refresh models configuration")
    
    # Attempt refresh
    success = await refresh_models_config_cache()
    
    if success:
        # Get updated counts
        standard_models = await get_vertex_models()
        express_models = await get_vertex_express_models()
        
        vertex_log('info', f"Models refresh successful. Standard: {len(standard_models)}, Express: {len(express_models)}")
        
        return {
            "success": True,
            "message": "Models configuration refreshed successfully",
            "models_count": {
                "standard": len(standard_models),
                "express": len(express_models),
                "total": len(standard_models) + len(express_models)
            }
        }
    else:
        vertex_log('error', "Failed to refresh models configuration")
        
        raise HTTPException(
            status_code=500,
            detail="Failed to refresh models configuration"
        )

@router.get("/v1/models")
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]"
    
    # 获取credential_manager，如果不存在则创建一个新的
    try:
        credential_manager_instance = fastapi_request.app.state.credential_manager
        vertex_log('info', "Using existing credential manager from app state")
    except AttributeError:
        # 如果app.state中没有credential_manager，则创建一个新的
        vertex_log('warning', "No credential_manager found in app.state, creating a new one")
        credential_manager_instance = CredentialManager()

    # 检查是否有SA凭证和Express Key
    has_sa_creds = credential_manager_instance.get_total_credentials() > 0
    
    # 优先从settings中获取EXPRESS API KEY
    has_express_key = False
    if hasattr(settings, 'VERTEX_EXPRESS_API_KEY') and settings.VERTEX_EXPRESS_API_KEY:
        vertex_log('info', "使用settings中的VERTEX_EXPRESS_API_KEY")
        has_express_key = True
    elif app_config.VERTEX_EXPRESS_API_KEY_VAL:
        vertex_log('info', "使用app_config中的VERTEX_EXPRESS_API_KEY_VAL")
        has_express_key = True

    raw_vertex_models = await get_vertex_models()
    raw_express_models = await get_vertex_express_models()
    
    candidate_model_ids = set()

    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]] = []
    current_time = int(time.time())

    # Add base models and their variations
    for original_model_id in sorted(list(all_model_ids)):
        current_display_prefix = ""
        # 检查是否为[EXPRESS]模型，保留原始前缀
        is_express_model = original_model_id.startswith("[EXPRESS]")
        base_model_without_prefix = original_model_id
        
        # 如果是EXPRESS模型，移除前缀用于后续处理，但在显示时会保留
        if is_express_model:
            # 从原始ID中提取不带[EXPRESS]前缀的模型名
            base_model_without_prefix = original_model_id[len("[EXPRESS] "):]
            vertex_log('info', f"处理EXPRESS模型: {original_model_id}, 基础名称: {base_model_without_prefix}")
        
        # 只有非EXPRESS模型才考虑添加PAY_PREFIX
        if not is_express_model and has_sa_creds and EXPERIMENTAL_MARKER not in original_model_id:
            # 只要有SA凭证就应该显示PAY前缀，不管是否有Express Key
            current_display_prefix = PAY_PREFIX
            vertex_log('info', f"为非EXPRESS模型添加PAY前缀: {current_display_prefix}{original_model_id}")
        
        # 构建显示ID，如果是EXPRESS模型，确保保留[EXPRESS]前缀
        base_display_id = original_model_id if is_express_model else f"{current_display_prefix}{original_model_id}"
        
        vertex_log('info', f"添加模型到列表: {base_display_id}")
        dynamic_models_data.append({
            "id": base_display_id, "object": "model", "created": current_time, "owned_by": "google",
            "permission": [], "root": base_model_without_prefix, "parent": None
        })
        
        # Conditionally add common variations (standard suffixes)
        if not base_model_without_prefix.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 (without EXPRESS prefix)
                suffixed_model_part = f"{base_model_without_prefix}{suffix}"
                # Then the whole thing is prefixed - retain EXPRESS prefix if it was there
                final_suffixed_display_id = ""
                if is_express_model:
                    final_suffixed_display_id = f"[EXPRESS] {suffixed_model_part}"
                else:
                    # 非EXPRESS模型的后缀版本也需要正确显示PAY前缀
                    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):
                    vertex_log('info', f"添加后缀模型到列表: {final_suffixed_display_id}")
                    dynamic_models_data.append({
                        "id": final_suffixed_display_id, "object": "model", "created": current_time, "owned_by": "google",
                        "permission": [], "root": base_model_without_prefix, "parent": None
                    })
        
        # Apply special suffixes for models starting with "gemini-2.5-flash"
        if base_model_without_prefix.startswith("gemini-2.5-flash"): # Suffix rules based on original_model_id
            special_flash_suffixes = ["-nothinking", "-max"]
            for special_suffix in special_flash_suffixes:
                suffixed_model_part = f"{base_model_without_prefix}{special_suffix}"
                # Retain EXPRESS prefix if original model had it
                if is_express_model:
                    final_special_suffixed_display_id = f"[EXPRESS] {suffixed_model_part}"
                else:
                    # 非EXPRESS模型的特殊后缀版本也需要正确显示PAY前缀
                    final_special_suffixed_display_id = f"{current_display_prefix}{suffixed_model_part}"

                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):
                    vertex_log('info', f"添加特殊后缀模型到列表: {final_special_suffixed_display_id}")
                    dynamic_models_data.append({
                        "id": final_special_suffixed_display_id, "object": "model", "created": current_time, "owned_by": "google",
                        "permission": [], "root": base_model_without_prefix, "parent": None
                    })

        # Apply special suffixes for gemini-2.5-pro-preview-06-05
        if base_model_without_prefix == "gemini-2.5-pro-preview-06-05":
            special_pro_suffixes = ["-nothinking", "-max"]
            for special_suffix in special_pro_suffixes:
                suffixed_model_part = f"{base_model_without_prefix}{special_suffix}"
                # Retain EXPRESS prefix if original model had it
                if is_express_model:
                    final_special_suffixed_display_id = f"[EXPRESS] {suffixed_model_part}"
                else:
                    # 非EXPRESS模型的特殊后缀版本也需要正确显示PAY前缀
                    final_special_suffixed_display_id = f"{current_display_prefix}{suffixed_model_part}"

                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):
                    vertex_log('info', f"添加gemini-2.5-pro-preview-06-05特殊后缀模型到列表: {final_special_suffixed_display_id}")
                    dynamic_models_data.append({
                        "id": final_special_suffixed_display_id, "object": "model", "created": current_time, "owned_by": "google",
                        "permission": [], "root": base_model_without_prefix, "parent": None
                    })

        # Ensure uniqueness again after adding suffixes
        # Add OpenAI direct variations if SA creds are available
        if has_sa_creds: # OpenAI direct mode only works with SA credentials
            # `all_model_ids` contains the comprehensive list of base models that are eligible based on current credentials
            # We iterate through this to determine which ones get an -openai variation.
            # `raw_vertex_models` is used here to ensure we only add -openai suffix to models that are
            # fundamentally Vertex models, not just any model that might appear in `all_model_ids` (e.g. from Express list exclusively)
            # if express only key is provided.
            # We iterate through the base models from the main Vertex list.
            for base_model_id_for_openai in raw_vertex_models: # Iterate through original list of GAIA/Vertex base 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:
                    # OpenAI直接模式下也应该保持PAY前缀
                    display_model_id = f"{PAY_PREFIX}{base_model_id_for_openai}{OPENAI_DIRECT_SUFFIX}"
                
                # Check if already added (e.g. if remote config somehow already listed it or added as a base model)
                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
                    })
    
    return {"object": "list", "data": sorted(dynamic_models_data, key=lambda x: x['id'])}