File size: 5,592 Bytes
7cc3183 |
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 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
import time
from fastapi import APIRouter, Depends, Request # Added Request
from typing import List, Dict, Any
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
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()
OPENAI_DIRECT_SUFFIX = "-openai"
EXPERIMENTAL_MARKER = "-exp-"
# Access credential_manager from app state
credential_manager_instance: CredentialManager = fastapi_request.app.state.credential_manager
has_sa_creds = credential_manager_instance.get_total_credentials() > 0
has_express_key = bool(app_config.VERTEX_EXPRESS_API_KEY_VAL)
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 model_id in sorted(list(all_model_ids)):
dynamic_models_data.append({
"id": model_id, "object": "model", "created": current_time, "owned_by": "google",
"permission": [], "root": model_id, "parent": None
})
# Conditionally add common variations (standard suffixes)
if not model_id.startswith("gemini-2.0"):
standard_suffixes = ["-search", "-encrypt", "-encrypt-full", "-auto"]
for suffix in standard_suffixes:
suffixed_id = f"{model_id}{suffix}"
# Check if this suffixed ID is already in all_model_ids (fetched from remote) or already added to dynamic_models_data
if suffixed_id not in all_model_ids and not any(m['id'] == suffixed_id for m in dynamic_models_data):
dynamic_models_data.append({
"id": suffixed_id, "object": "model", "created": current_time, "owned_by": "google",
"permission": [], "root": model_id, "parent": None
})
# Apply special suffixes for models starting with "gemini-2.5-flash"
if model_id.startswith("gemini-2.5-flash"):
special_flash_suffixes = ["-nothinking", "-max"]
for special_suffix in special_flash_suffixes:
suffixed_id = f"{model_id}{special_suffix}"
if suffixed_id not in all_model_ids and not any(m['id'] == suffixed_id for m in dynamic_models_data):
dynamic_models_data.append({
"id": suffixed_id, "object": "model", "created": current_time, "owned_by": "google",
"permission": [], "root": model_id, "parent": None
})
# Ensure uniqueness again after adding suffixes
# Add OpenAI direct variations for experimental models if SA creds are available
if has_sa_creds: # OpenAI direct mode only works with SA credentials
# We should iterate through the base models that could be experimental.
# `raw_vertex_models` should contain these.
for model_id in raw_vertex_models: # Iterate through the original list of base models
if EXPERIMENTAL_MARKER in model_id:
suffixed_id = f"{model_id}{OPENAI_DIRECT_SUFFIX}"
# Check if already added (e.g. if remote config somehow already listed it)
if not any(m['id'] == suffixed_id for m in dynamic_models_data):
dynamic_models_data.append({
"id": suffixed_id, "object": "model", "created": current_time, "owned_by": "google",
"permission": [], "root": model_id, "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'])} |