ReOpenAI / app.py
YchKhan's picture
Update app.py
262f775 verified
raw
history blame
5.39 kB
import asyncio
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
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
# 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 MockRequest(BaseModel):
"""Définit la structure attendue pour le corps de la requête POST."""
parameter: 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é, 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.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, 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 = payload.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 60 secondes
websocket_response = await asyncio.wait_for(response_future, timeout=60.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)