File size: 2,896 Bytes
7cc3183 d342ca5 7cc3183 d342ca5 7cc3183 9fde8ed d342ca5 7cc3183 da7a18e 7cc3183 da7a18e 7cc3183 d342ca5 7cc3183 d342ca5 2a81a94 d342ca5 2a81a94 d342ca5 2a81a94 d342ca5 7cc3183 d342ca5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
import time
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
from credentials_manager import CredentialManager
router = APIRouter()
@router.get("/v1/models")
async def list_models(fastapi_request: Request, api_key: str = Depends(get_api_key)):
await refresh_models_config_cache()
PAY_PREFIX = "[PAY]"
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
has_sa_creds = credential_manager_instance.get_total_credentials() > 0
has_express_key = express_key_manager_instance.get_total_keys() > 0
raw_vertex_models = await get_vertex_models()
raw_express_models = await get_vertex_express_models()
final_model_list: List[Dict[str, Any]] = []
processed_ids: Set[str] = set()
current_time = int(time.time())
def add_model_and_variants(base_id: str, prefix: str):
"""Adds a model and its variants to the list if not already present."""
# 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"])
# Add the openai variant for all models
suffixes.append(OPENAI_DIRECT_SUFFIX)
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
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)
# 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'])}
|