ReOpenAI / app.py
YchKhan's picture
Update app.py
cc2d90e verified
raw
history blame
4.85 kB
import asyncio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from typing import List, Dict
import uvicorn
# Initialiser l'application FastAPI
app = FastAPI()
# Configurer CORS pour autoriser toutes les origines
# Permet à n'importe quel site web d'appeler votre API
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Autorise toutes les origines
allow_credentials=True,
allow_methods=["*"], # Autorise toutes les méthodes (GET, POST, etc.)
allow_headers=["*"], # Autorise tous les en-têtes
)
# Monter un répertoire statique pour servir le fichier index.html
app.mount("/static", StaticFiles(directory="static"), name="static")
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é, bien que simple, c'est plus robuste
client_id = str(id(websocket))
self.response_futures[client_id] = future
return future
return None
manager = ConnectionManager()
@app.post("/v1/mock")
async def mock_endpoint(request: Request):
"""
Endpoint API qui prend un string, le transmet via WebSocket,
attend une réponse et la retourne.
"""
try:
# Récupérer les données JSON du corps de la requête
data = await request.json()
input_string = data.get("parameter")
if input_string is None:
return {"error": "Le paramètre 'parameter' est manquant."}
print(f"Endpoint /v1/mock appelé avec: '{input_string}'")
if not manager.active_connections:
return {"error": "Aucun client WebSocket n'est connecté."}
# 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:
return {"error": "Échec de la diffusion du message."}
try:
# Attendre la réponse du client WebSocket avec un timeout de 10 secondes
websocket_response = await asyncio.wait_for(response_future, timeout=10.0)
print(f"Réponse reçue du WebSocket: '{websocket_response}'")
return {"response_from_client": websocket_response}
except asyncio.TimeoutError:
print("Timeout: Aucune réponse du client WebSocket.")
return {"error": "Timeout: Le client n'a pas répondu à temps."}
except Exception as e:
print(f"Erreur dans /v1/mock: {e}")
return {"error": 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)