Spaces:
Sleeping
Sleeping
File size: 3,368 Bytes
e109700 48b9b47 e109700 db280f4 e109700 dfb1c84 e109700 0053356 48b9b47 0053356 db280f4 0053356 dfb1c84 0053356 c2cd706 0053356 |
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 |
import asyncpg
import logging
from typing import List, Dict, Any, Union
from config.settings import DATABASE_URL
logger = logging.getLogger(__name__)
async def fetch_models_for_group(resource_group: str) -> List[Dict[str, Any]]:
"""Récupérer les détails des modèles depuis la base de données pour un groupe de ressources spécifique.
Args:
resource_group: Identifiant du groupe de ressources
Returns:
Liste de dictionnaires contenant les informations des modèles
Raises:
Exception: Si une erreur se produit lors de la connexion ou de la requête
"""
conn = None
try:
# Désactiver le cache des prepared statements pour compatibilité avec pgbouncer
conn = await asyncpg.connect(DATABASE_URL, statement_cache_size=0)
logger.info(f"Successfully connected to database")
# Récupérer les modèles spécifiques à ce groupe
query = """
SELECT
model_id,
display_name,
hf_repo_id,
hf_subfolder,
hf_filename,
updated_at
FROM models
WHERE hf_resource_group = $1
"""
rows = await conn.fetch(query, resource_group)
logger.info(f"Found {len(rows)} models for group '{resource_group}'")
return [dict(row) for row in rows]
except Exception as e:
logger.error(f"Database error: {e}", exc_info=True)
raise
finally:
if conn and not conn.is_closed():
await conn.close()
logger.debug("Database connection closed")
async def fetch_model_by_id(model_id: str) -> Union[Dict[str, Any], None]:
"""Récupérer les détails d'un modèle spécifique par son ID de base de données.
Args:
model_id: L'ID du modèle dans la base de données (peut être int ou str selon le schéma).
Returns:
Un dictionnaire contenant les informations du modèle si trouvé, sinon None.
Raises:
Exception: Si une erreur se produit lors de la connexion ou de la requête.
"""
conn = None
try:
# Désactiver le cache des prepared statements pour compatibilité avec pgbouncer
conn = await asyncpg.connect(DATABASE_URL, statement_cache_size=0)
logger.debug(f"Successfully connected to database to fetch model ID: {model_id}")
# Récupérer le modèle spécifique par son ID
query = """
SELECT
model_id,
display_name,
hf_repo_id,
hf_subfolder,
hf_filename,
updated_at
FROM models
WHERE model_id = $1
"""
row = await conn.fetchrow(query, model_id)
if row:
logger.info(f"Found model with ID '{model_id}': {row['display_name']}")
return dict(row)
else:
logger.warning(f"No model found with ID '{model_id}'")
return None
except Exception as e:
logger.error(f"Database error fetching model ID {model_id}: {e}", exc_info=True)
raise
finally:
if conn and not conn.is_closed():
await conn.close()
logger.debug(f"Database connection closed after fetching model ID: {model_id}")
|