from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
import asyncio
from datetime import datetime, timedelta
import httpx
import os
import logging
import json
# Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("main")
app = FastAPI(
title="Vitizen WebSocket API",
description="API WebSocket pour l'assistant viticole",
version="1.0.0"
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Restrict in production
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
LANGSERVE_URL = os.getenv("LANGSERVE_URL", "https://thibautmodrin-vitizen-chat.hf.space/chat/invoke")
@app.get("/", response_class=HTMLResponse)
async def home():
return """
Bienvenue sur l'API WebSocket Vitizen
"""
@app.websocket("/ws/chat")
async def websocket_endpoint(websocket: WebSocket):
await websocket.accept()
logger.info("✅ Nouvelle connexion WebSocket établie")
last_call_time = datetime.utcnow() - timedelta(seconds=5)
async with httpx.AsyncClient(timeout=30) as client:
try:
while True:
data = await websocket.receive_json()
question = data.get("message", "").strip()
if not question:
await websocket.send_text("[Erreur] Le message est vide.")
continue
logger.info(f"💬 Message reçu: {question[:100]}...")
now = datetime.utcnow()
delta = (now - last_call_time).total_seconds()
if delta < 2:
wait_time = round(2 - delta, 2)
await websocket.send_text(f"⏳ Merci d'attendre {wait_time}s.")
continue
last_call_time = now
logger.info(f"📤 Envoi à LangServe: {LANGSERVE_URL}")
try:
response = await client.post(
LANGSERVE_URL,
json={"input": question}, # ✅ FORMAT CORRECT
headers={"User-Agent": "vitizen-client/1.0"},
timeout=30
)
response.raise_for_status()
logger.info(f"✅ Statut HTTP: {response.status_code}")
try:
json_data = response.json()
logger.info(f"🔍 Contenu JSON: {json_data}")
answer = json_data.get("output", "[Erreur: réponse vide]")
except Exception as e:
answer = f"[Erreur] Exception JSON: {str(e)}"
except httpx.HTTPStatusError as e:
answer = f"[Erreur] HTTP {e.response.status_code} - {e.response.reason_phrase}"
except httpx.RequestError as e:
answer = f"[Erreur] Échec réseau: {e.__class__.__name__} - {str(e)}"
except Exception as e:
answer = f"[Erreur] Exception durant la requête: {str(e)}"
try:
await websocket.send_text(answer)
logger.info("📨 Réponse envoyée")
except RuntimeError as e:
logger.warning(f"⛔ Connexion fermée avant l'envoi de la réponse: {str(e)}")
break
except WebSocketDisconnect:
logger.info("🔌 Connexion WebSocket fermée par le client")
except Exception as e:
logger.error(f"[Erreur] Exception inattendue: {str(e)}")
try:
await websocket.send_text(f"[Erreur] {str(e)}")
except:
pass
@app.get("/health")
async def health_check():
return {
"status": "ok",
"langserve_url": LANGSERVE_URL,
"version": "1.0.0"
}
if __name__ == "__main__":
import uvicorn
logger.info("🚀 Lancement du serveur sur le port 8000")
uvicorn.run(app, host="0.0.0.0", port=8000)