sana0721/vertex

1import httpx2import asyncio3import json4from typing import List, Dict, Optional, Any5 6# Assuming config.py is in the same directory level for Docker execution7import config as app_config 8 9_model_cache: Optional[Dict[str, List[str]]] = None10_cache_lock = asyncio.Lock()11 12async def fetch_and_parse_models_config() -> Optional[Dict[str, List[str]]]:13    """14    Fetches the model configuration JSON from the URL specified in app_config.15    Parses it and returns a dictionary with 'vertex_models' and 'vertex_express_models'.16    Returns None if fetching or parsing fails.17    """18    if not app_config.MODELS_CONFIG_URL:19        print("ERROR: MODELS_CONFIG_URL is not set in the environment/config.")20        return None21 22    print(f"Fetching model configuration from: {app_config.MODELS_CONFIG_URL}")23    try:24        async with httpx.AsyncClient() as client:25            response = await client.get(app_config.MODELS_CONFIG_URL)26            response.raise_for_status() # Raise an exception for HTTP errors (4xx or 5xx)27            data = response.json()28            29            # Basic validation of the fetched data structure30            if isinstance(data, dict) and \31               "vertex_models" in data and isinstance(data["vertex_models"], list) and \32               "vertex_express_models" in data and isinstance(data["vertex_express_models"], list):33                print("Successfully fetched and parsed model configuration.")34                35                # Add [EXPRESS] prefix to express models36                prefixed_express_models = [f"[EXPRESS] {model_name}" for model_name in data["vertex_express_models"]]37                38                return {39                    "vertex_models": data["vertex_models"],40                    "vertex_express_models": prefixed_express_models41                }42            else:43                print(f"ERROR: Fetched model configuration has an invalid structure: {data}")44                return None45    except httpx.RequestError as e:46        print(f"ERROR: HTTP request failed while fetching model configuration: {e}")47        return None48    except json.JSONDecodeError as e:49        print(f"ERROR: Failed to decode JSON from model configuration: {e}")50        return None51    except Exception as e:52        print(f"ERROR: An unexpected error occurred while fetching/parsing model configuration: {e}")53        return None54 55async def get_models_config() -> Dict[str, List[str]]:56    """57    Returns the cached model configuration.58    If not cached, fetches and caches it.59    Returns a default empty structure if fetching fails.60    """61    global _model_cache62    async with _cache_lock:63        if _model_cache is None:64            print("Model cache is empty. Fetching configuration...")65            _model_cache = await fetch_and_parse_models_config()66            if _model_cache is None: # If fetching failed, use a default empty structure67                print("WARNING: Using default empty model configuration due to fetch/parse failure.")68                _model_cache = {"vertex_models": [], "vertex_express_models": []}69    return _model_cache70 71async def get_vertex_models() -> List[str]:72    config = await get_models_config()73    return config.get("vertex_models", [])74 75async def get_vertex_express_models() -> List[str]:76    config = await get_models_config()77    return config.get("vertex_express_models", [])78 79async def refresh_models_config_cache() -> bool:80    """81    Forces a refresh of the model configuration cache.82    Returns True if successful, False otherwise.83    """84    global _model_cache85    print("Attempting to refresh model configuration cache...")86    async with _cache_lock:87        new_config = await fetch_and_parse_models_config()88        if new_config is not None:89            _model_cache = new_config90            print("Model configuration cache refreshed successfully.")91            return True92        else:93            print("ERROR: Failed to refresh model configuration cache.")94            # Optionally, decide if we want to clear the old cache or keep it95            # _model_cache = {"vertex_models": [], "vertex_express_models": []} # To clear96            return False