import asyncio import os from fastapi import FastAPI, WebSocket, WebSocketDisconnect, HTTPException from fastapi.responses import HTMLResponse from fastapi.staticfiles import StaticFiles from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from typing import List, Dict import uvicorn import json import time # Initialiser l'application FastAPI app = FastAPI() # Configurer CORS pour autoriser toutes les origines app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Monter un répertoire statique pour servir le fichier index.html app.mount("/static", StaticFiles(directory="static"), name="static") class MockRequest(BaseModel): """Définit la structure attendue pour le corps de la requête POST.""" parameter: str model: str = None secret: str class ConnectionManager: """Gère les connexions WebSocket actives.""" def __init__(self): self.active_connections: List[WebSocket] = [] # Dictionnaire pour attendre les réponses des clients self.response_futures: Dict[str, asyncio.Future] = {} async def connect(self, websocket: WebSocket): """Accepte une nouvelle connexion WebSocket.""" await websocket.accept() self.active_connections.append(websocket) print(f"Nouvelle connexion WebSocket. Total: {len(self.active_connections)}") def disconnect(self, websocket: WebSocket): """Ferme une connexion WebSocket.""" self.active_connections.remove(websocket) print(f"Déconnexion WebSocket. Total: {len(self.active_connections)}") async def broadcast(self, message: str): """Envoie un message à tous les clients connectés.""" # Pour ce cas simple, nous n'envoyons qu'au premier client connecté if self.active_connections: websocket = self.active_connections[0] await websocket.send_text(message) # Créer un Future pour attendre la réponse future = asyncio.get_event_loop().create_future() # Utilise l'identifiant du client comme clé client_id = str(id(websocket)) self.response_futures[client_id] = future return future return None manager = ConnectionManager() def verify_secret(provided_secret: str) -> bool: """Vérifie si le secret fourni correspond à celui de la variable d'environnement.""" expected_secret = os.getenv("API_SECRET") if not expected_secret: print("ATTENTION: Variable d'environnement API_SECRET non définie!") return False return provided_secret == expected_secret @app.get("/", response_class=HTMLResponse) async def root(): """Serve the main HTML page.""" try: with open("static/index.html", "r", encoding="utf-8") as f: return HTMLResponse(content=f.read()) except FileNotFoundError: raise HTTPException(status_code=404, detail="index.html not found") @app.post("/v1/mock") async def mock_endpoint(payload: MockRequest): """ Endpoint API qui prend un string et un secret, vérifie le secret, puis transmet via WebSocket, attend une réponse et la retourne. """ start_time = time.monotonic() try: input_string = payload.parameter selected_model = payload.model provided_secret = payload.secret # Vérification du secret AVANT tout traitement if not verify_secret(provided_secret): print(f"Tentative d'accès avec un secret invalide: '{provided_secret[:10]}...'") raise HTTPException( status_code=401, detail="Secret invalide. Accès non autorisé." ) print(f"Secret vérifié avec succès. Endpoint /v1/mock appelé avec: '{input_string}'") if input_string is None: raise HTTPException(status_code=400, detail="Le paramètre 'parameter' est manquant.") if not manager.active_connections: raise HTTPException(status_code=503, detail="Aucun client WebSocket n'est connecté.") # Créer un dictionnaire avec les données à envoyer message_data = { "prompt": input_string, "model": selected_model } # Envoyer le message via WebSocket (sérialiser en JSON) response_future = await manager.broadcast(json.dumps(message_data)) # Envoyer le message via WebSocket et obtenir un "future" pour la réponse print("Envoi du message au client WebSocket...") response_future = await manager.broadcast(input_string) if response_future is None: raise HTTPException(status_code=500, detail="Échec de la diffusion du message.") try: # Attendre la réponse du client WebSocket avec un timeout de 60 secondes websocket_response = await asyncio.wait_for(response_future, timeout=60.0) print(f"Réponse reçue du WebSocket: '{websocket_response}'") end_time = time.monotonic() duration = end_time - start_time print(f"Requête complétée en {duration:.2f} secondes.") return { "response_from_client": websocket_response, "completion_time_in_seconds": round(duration, 2) } except asyncio.TimeoutError: print("Timeout: Aucune réponse du client WebSocket.") raise HTTPException(status_code=408, detail="Timeout: Le client n'a pas répondu à temps.") except HTTPException: # Re-lever les HTTPException sans les wrapper raise except Exception as e: print(f"Erreur dans /v1/mock: {e}") raise HTTPException(status_code=500, detail=f"Une erreur interne est survenue: {str(e)}") @app.websocket("/ws") async def websocket_endpoint(websocket: WebSocket): """Gère la communication WebSocket avec le client.""" await manager.connect(websocket) try: while True: # Attendre un message du client data = await websocket.receive_text() print(f"Message reçu du client: '{data}'") # Trouver le "future" correspondant et y mettre le résultat client_id = str(id(websocket)) if client_id in manager.response_futures: manager.response_futures[client_id].set_result(data) del manager.response_futures[client_id] # Nettoyer après utilisation except WebSocketDisconnect: manager.disconnect(websocket) print("Client déconnecté.") except Exception as e: print(f"Erreur dans le WebSocket: {e}") manager.disconnect(websocket) if __name__ == "__main__": uvicorn.run(app, host="0.0.0.0", port=7860)