Refactor API endpoints and implement JWT authentication (#11)

- Removed the chat endpoint and integrated AG-UI WebSocket routing.
- Added JWT authentication for conversation initialization, allowing
secure WebSocket connections.
- Updated requirements to include PyJWT for token management.
- Cleaned up unused CopilotKit integration code and adjusted
conversation service to persist messages in the database.

This refactor enhances the security and organization of the API while
streamlining the conversation management process.

Author: EmiLabarthe

app/api/v1/agui_ws.py

@@ -0,0 +1,120 @@
+"""
+AG-UI WebSocket endpoint.
+
+Implements the Agent-User Interaction protocol over WebSockets.
+Each connection is authenticated via a JWT query-param that embeds the
+``conversationId``.  Messages are persisted through ``ChatService``.
+"""
+
+import json
+import logging
+import uuid
+
+from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect
+
+from app.security.auth import verify_conversation_token
+from app.services.chat_service import chat_service
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+
+def _agui_event(event_type: str, **fields) -> str:
+    """Serialise a single AG-UI event as a JSON text frame."""
+    return json.dumps({"type": event_type, **fields})
+
+
+@router.websocket("/ws")
+async def agui_websocket(
+    websocket: WebSocket,
+    token: str = Query(...),
+):
+    conversation_id = verify_conversation_token(token)
+    await websocket.accept()
+    logger.info(f"[AG-UI] WebSocket connected  conv_id={conversation_id}")
+
+    try:
+        while True:
+            raw = await websocket.receive_text()
+
+            try:
+                data = json.loads(raw)
+            except json.JSONDecodeError:
+                await websocket.send_text(
+                    _agui_event("RUN_ERROR", message="Invalid JSON", code="BAD_REQUEST")
+                )
+                continue
+
+            user_message = data.get("message", "").strip()
+            if not user_message:
+                await websocket.send_text(
+                    _agui_event("RUN_ERROR", message="Empty message", code="BAD_REQUEST")
+                )
+                continue
+
+            wizard_state = data.get("wizard_state")
+            run_id = str(uuid.uuid4())
+            thread_id = str(conversation_id)
+
+            await websocket.send_text(
+                _agui_event("RUN_STARTED", threadId=thread_id, runId=run_id)
+            )
+
+            try:
+                result = await chat_service.process_message(
+                    message=user_message,
+                    conversation_id=conversation_id,
+                    wizard_state=wizard_state,
+                )
+
+                response_text = result.get("response", "")
+                message_id = str(uuid.uuid4())
+
+                await websocket.send_text(
+                    _agui_event("TEXT_MESSAGE_START", messageId=message_id, role="assistant")
+                )
+                await websocket.send_text(
+                    _agui_event("TEXT_MESSAGE_CONTENT", messageId=message_id, delta=response_text)
+                )
+                await websocket.send_text(
+                    _agui_event("TEXT_MESSAGE_END", messageId=message_id)
+                )
+
+                state_snapshot = _build_state_snapshot(result)
+                if state_snapshot:
+                    await websocket.send_text(
+                        _agui_event("STATE_SNAPSHOT", snapshot=state_snapshot)
+                    )
+
+                await websocket.send_text(
+                    _agui_event("RUN_FINISHED", threadId=thread_id, runId=run_id)
+                )
+
+            except Exception as exc:
+                logger.error(f"[AG-UI] Error processing message: {exc}", exc_info=True)
+                await websocket.send_text(
+                    _agui_event("RUN_ERROR", message=str(exc), code="INTERNAL_ERROR")
+                )
+                await websocket.send_text(
+                    _agui_event("RUN_FINISHED", threadId=thread_id, runId=run_id)
+                )
+
+    except WebSocketDisconnect:
+        logger.info(f"[AG-UI] WebSocket disconnected  conv_id={conversation_id}")
+
+
+def _build_state_snapshot(result: dict) -> dict | None:
+    """Extract wizard / agent metadata into an AG-UI STATE_SNAPSHOT payload."""
+    snapshot: dict = {}
+
+    if result.get("wizard_state") and result["wizard_state"] != "INACTIVE":
+        snapshot["wizard_state"] = result.get("wizard_state")
+        snapshot["current_question"] = result.get("current_question")
+        snapshot["wizard_responses"] = result.get("wizard_responses", {})
+        snapshot["awaiting_answer"] = result.get("awaiting_answer", False)
+        snapshot["wizard_session_id"] = result.get("wizard_session_id")
+
+    snapshot["agent_used"] = result.get("agent_used", "unknown")
+
+    return snapshot if snapshot else None

app/api/v1/chat.py

@@ -8,6 +8,7 @@
 
 from app.db.config.database import get_async_session
 from app.db.models import Conversation
+from app.security.auth import create_conversation_token
 
 router = APIRouter()
 
@@ -22,6 +23,11 @@ class ConversationResponse(BaseModel):
     started_at: datetime
 
 
+class ConversationInitResponse(BaseModel):
+    token: str
+    conversationId: int
+
+
 @router.post("/conversations", response_model=ConversationResponse)
 async def create_conversation(
         conversation: ConversationCreate,
@@ -42,6 +48,25 @@ async def create_conversation(
         raise HTTPException(status_code=500, detail="Error creating conversation")
 
 
+@router.post("/conversations/init", response_model=ConversationInitResponse)
+async def init_conversation(
+        payload: ConversationCreate = ConversationCreate(),
+        session: AsyncSession = Depends(get_async_session),
+) -> ConversationInitResponse:
+    """Create a new conversation and return a signed JWT for the WebSocket."""
+    try:
+        new_conv = Conversation(email=payload.email)
+        session.add(new_conv)
+        await session.commit()
+        await session.refresh(new_conv)
+
+        token = create_conversation_token(new_conv.id)
+        return ConversationInitResponse(token=token, conversationId=new_conv.id)
+    except Exception:
+        await session.rollback()
+        raise HTTPException(status_code=500, detail="Error creating conversation")
+
+
 @router.get("/conversations")
 async def get_conversations(
         session: AsyncSession = Depends(get_async_session)

app/api/v1/conversations.py

@@ -12,31 +12,28 @@
     force=True,
 )
 
-from app.api.v1.chat import router as chat_router
+from app.api.v1.agui_ws import router as agui_ws_router
 from app.api.v1.conversations import router as conversations_router
-from app.api.v1.copilotkit_endpoint import router as copilotkit_router
 from app.api.v1.documents import router as documents_router
 from app.api.v1.scoring import router as scoring_router
 
 v1 = '/api/v1'
 
 app = FastAPI(title="Chatbot Backend", version="1.0.0")
 
-# Configurar CORS para permitir conexiones desde el frontend
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["http://localhost:3000", "http://127.0.0.1:3000",
-                   "http://localhost:3001"],  # Frontend URLs
+                   "http://localhost:3001"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
 app.include_router(conversations_router, prefix=v1, tags=["Conversations"])
-app.include_router(chat_router, prefix=v1, tags=["Chat"])
+app.include_router(agui_ws_router, prefix=v1, tags=["AG-UI WebSocket"])
 app.include_router(documents_router, prefix=v1, tags=["Documents"])
 app.include_router(scoring_router, prefix=v1, tags=["Scoring"])
-app.include_router(copilotkit_router, prefix=v1, tags=["CopilotKit"])
 
 
 @app.get("/")

app/api/v1/copilotkit_endpoint.py

@@ -1,12 +1,18 @@
 import os
 import secrets
 from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
 
+import jwt
 from fastapi import Depends, HTTPException, status
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
 bearer_scheme = HTTPBearer(auto_error=False)
 
+JWT_SECRET_KEY = os.getenv("JWT_SECRET_KEY", "change-me-in-production")
+JWT_ALGORITHM = os.getenv("JWT_ALGORITHM", "HS256")
+JWT_EXPIRATION_MINUTES = int(os.getenv("JWT_EXPIRATION_MINUTES", "60"))
+
 
 @dataclass(frozen=True)
 class AuthUser:
@@ -76,3 +82,37 @@ async def require_admin_user(
             detail="Permisos insuficientes. Se requiere rol admin",
         )
     return user
+
+
+# ---------------------------------------------------------------------------
+# JWT helpers for WebSocket conversation tokens
+# ---------------------------------------------------------------------------
+
+def create_conversation_token(conversation_id: int) -> str:
+    """Sign a short-lived JWT that embeds the conversation ID."""
+    payload = {
+        "sub": str(conversation_id),
+        "exp": datetime.now(timezone.utc) + timedelta(minutes=JWT_EXPIRATION_MINUTES),
+        "iat": datetime.now(timezone.utc),
+    }
+    return jwt.encode(payload, JWT_SECRET_KEY, algorithm=JWT_ALGORITHM)
+
+
+def verify_conversation_token(token: str) -> int:
+    """Validate a conversation JWT and return the integer conversation_id.
+
+    Raises ``HTTPException(401)`` on any validation failure.
+    """
+    try:
+        payload = jwt.decode(token, JWT_SECRET_KEY, algorithms=[JWT_ALGORITHM])
+        return int(payload["sub"])
+    except jwt.ExpiredSignatureError:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Token expirado",
+        )
+    except (jwt.InvalidTokenError, KeyError, ValueError) as exc:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=f"Token inválido: {exc}",
+        )