ReOpenAI / app.py
YchKhan's picture
Update app.py
a413d29 verified
raw
history blame
5.06 kB
import asyncio
import uuid
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Dict, Any
import uvicorn
# 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
class ConnectionManager:
"""Gère les connexions WebSocket actives et les réponses en attente."""
def __init__(self):
self.active_connections: List[WebSocket] = []
# Utilise un ID de requête pour lier les requêtes aux réponses
self.response_futures: Dict[str, asyncio.Future] = {}
async def connect(self, websocket: WebSocket):
await websocket.accept()
self.active_connections.append(websocket)
print(f"Nouvelle connexion WebSocket. Total: {len(self.active_connections)}")
def disconnect(self, websocket: WebSocket):
# Annuler tous les futures en attente pour ce client déconnecté
keys_to_remove = [key for key, (ws, _) in self.response_futures.items() if ws == websocket]
for key in keys_to_remove:
self.response_futures[key][1].cancel()
del self.response_futures[key]
self.active_connections.remove(websocket)
print(f"Déconnexion WebSocket. Total: {len(self.active_connections)}")
async def send_action_and_wait(self, action: str, data: Any):
"""Envoie une action JSON au client et attend une réponse."""
if not self.active_connections:
return None
request_id = str(uuid.uuid4())
websocket = self.active_connections[0] # Simplification : on envoie au premier client
future = asyncio.get_event_loop().create_future()
self.response_futures[request_id] = future
message_to_send = {"request_id": request_id, "action": action, "data": data}
await websocket.send_json(message_to_send)
return await future
manager = ConnectionManager()
async def handle_api_request(action: str, payload: MockRequest):
"""Factorise la logique commune aux endpoints API."""
try:
input_string = payload.parameter
print(f"Endpoint pour l'action '{action}' appelé avec: '{input_string}'")
if not manager.active_connections:
return {"error": "Aucun client WebSocket n'est connecté."}
print(f"Envoi de l'action '{action}' au client WebSocket...")
websocket_response = await manager.send_action_and_wait(action, input_string)
if websocket_response is None:
return {"error": "Échec de la communication avec le client."}
print(f"Réponse reçue du WebSocket pour l'action '{action}': '{websocket_response}'")
return {"response_from_client": websocket_response}
except asyncio.CancelledError:
print("La tâche de réponse a été annulée (probablement déconnexion du client).")
return {"error": "La requête a été annulée car le client s'est déconnecté."}
except Exception as e:
print(f"Erreur dans handle_api_request: {e}")
return {"error": f"Une erreur interne est survenue: {str(e)}"}
@app.post("/v1/mock")
async def mock_endpoint(payload: MockRequest):
"""Demande au client une phrase prédéfinie."""
return await handle_api_request("get_sentence", payload)
@app.post("/v1/reverse")
async def reverse_endpoint(payload: MockRequest):
"""Demande au client d'inverser une chaîne de caractères."""
return await handle_api_request("reverse_string", payload)
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await manager.connect(websocket)
try:
while True:
# Attend une réponse JSON du client
response_data = await websocket.receive_json()
request_id = response_data.get("request_id")
response_payload = response_data.get("response")
if request_id and request_id in manager.response_futures:
# Marque le future comme terminé avec le résultat
manager.response_futures[request_id].set_result(response_payload)
del manager.response_futures[request_id]
else:
print(f"Réponse reçue avec un ID inconnu ou manquant: {request_id}")
except WebSocketDisconnect:
manager.disconnect(websocket)
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)