import csv import hmac import json import logging from dataclasses import dataclass, field from functools import partial from io import StringIO from typing import Annotated, Any import anyio from fastapi import ( APIRouter, Depends, Header, HTTPException, Response, WebSocket, WebSocketDisconnect, ) from sqlalchemy import desc from sqlmodel import Session as DBSession from sqlmodel import select from app import crud from app.config import Settings, get_settings from app.db import get_db from app.providers import get_providers from app.schemas import ( ConfirmTurnRequest, FeedbackCreateRequest, SessionCreateRequest, SessionCreateResponse, ) from app.session import process_audio_turn, process_text_turn api_router = APIRouter(prefix="/api") ws_router = APIRouter() DBDep = Annotated[DBSession, Depends(get_db)] AdminToken = Annotated[str | None, Header(alias="X-Admin-Token")] logger = logging.getLogger(__name__) MAX_TYPED_TEXT_CHARS = 2000 @dataclass(slots=True) class InFlightTurn: speaker: str lang: str chunks: list[bytes] = field(default_factory=list) audio_bytes: int = 0 # ponytail: process-local turn lock; use DB/Redis if multi-worker deploy matters. _in_flight: dict[str, InFlightTurn] = {} def _risk_spans(turn: crud.TurnRecord) -> list[dict[str, Any]]: return json.loads(turn.risk_spans_json or "[]") def turn_payload(turn: crud.TurnRecord) -> dict[str, Any]: return { "id": turn.id, "session_id": turn.session_id, "seq": turn.seq, "speaker": turn.speaker, "src_lang": turn.src_lang, "tgt_lang": turn.tgt_lang, "source_text": turn.source_text, "normalized_text": turn.normalized_text, "translation": turn.translation, "asr_confidence": turn.asr_confidence, "mt_confidence": turn.mt_confidence, "risk_tier": turn.risk_tier, "risk_spans": _risk_spans(turn), "readback": json.loads(turn.readback_json) if turn.readback_json else None, "status": turn.status, "corrected_text": turn.corrected_text, "created_at": turn.created_at.isoformat(), } def feedback_payload(feedback: crud.FeedbackRecord) -> dict[str, Any]: return { "id": feedback.id, "turn_id": feedback.turn_id, "reason": feedback.reason, "comment": feedback.comment, "created_at": feedback.created_at.isoformat(), } def turn_result_payload(turn: crud.TurnRecord, settings: Settings) -> dict[str, Any]: low_confidence = ( turn.asr_confidence < settings.confidence_threshold or turn.mt_confidence < settings.confidence_threshold ) return { "type": "turn_result", "turn": turn_payload(turn), "requires_confirmation": turn.risk_tier in {"high", "critical"} or turn.status in {"awaiting_confirm", "blocked"}, "low_confidence": low_confidence, } def log_turn_processed(turn: crud.TurnRecord) -> None: logger.info( "turn_processed", extra={ "turn_id": turn.id, "session_id": turn.session_id, "risk_tier": turn.risk_tier, "status": turn.status, }, ) async def _send_pipeline_error( websocket: WebSocket, session_id: str, exc: Exception, *, speaker: str, lang: str, source_text: str | None, ) -> None: logger.warning( "turn_processing_failed session_id=%s error_class=%s", session_id, exc.__class__.__name__, ) _in_flight.pop(session_id, None) await websocket.send_json( { "type": "turn_error", "message": "translation failed — retry or use typed fallback", "retryable": True, "failure_context": { "speaker": speaker, "src_lang": lang, "tgt_lang": "en" if lang.lower().startswith("vi") else "vi", "source_text": source_text, "translation": None, }, } ) def _session_is_ended(db: DBSession, session_id: str) -> bool: return crud.require_session(db, session_id).status == "ended" async def _send_session_ended(websocket: WebSocket, session_id: str) -> None: _in_flight.pop(session_id, None) await websocket.send_json( {"type": "turn_error", "message": "session ended", "retryable": False} ) @api_router.get("/health") def health() -> dict[str, str]: settings = get_settings() return {"status": "ok", "provider_mode": settings.provider_mode} @api_router.post("/sessions", status_code=201) def create_session( payload: SessionCreateRequest, db: DBDep, ) -> SessionCreateResponse: session = crud.create_session(db, payload.consent) return SessionCreateResponse(session_id=session.id) @api_router.post("/sessions/{session_id}/end", status_code=204) def end_session(session_id: str, db: DBDep) -> Response: crud.end_session(db, session_id) return Response(status_code=204) @api_router.post("/sessions/{session_id}/escalate", status_code=204) def escalate_session(session_id: str, db: DBDep) -> Response: crud.escalate_session(db, session_id) return Response(status_code=204) @api_router.delete("/sessions/{session_id}", status_code=204) def delete_session(session_id: str, db: DBDep) -> Response: crud.hard_delete_session(db, session_id) return Response(status_code=204) @api_router.get("/sessions/{session_id}/transcript") def get_transcript(session_id: str, db: DBDep) -> list[dict[str, Any]]: if crud.get_session(db, session_id) is None: raise HTTPException(status_code=404, detail="session not found") return [turn_payload(turn) for turn in crud.list_turns(db, session_id)] @api_router.post("/turns/{turn_id}/confirm") def confirm_turn( turn_id: str, payload: ConfirmTurnRequest, db: DBDep, ) -> dict[str, Any]: return turn_payload(crud.confirm_turn(db, turn_id, payload.edited_translation)) @api_router.post("/turns/{turn_id}/feedback", status_code=201) def create_feedback( turn_id: str, payload: FeedbackCreateRequest, db: DBDep, ) -> dict[str, str]: feedback = crud.create_feedback(db, turn_id, payload.reason, payload.comment) return {"id": feedback.id} @api_router.get("/admin/review", response_model=None) def admin_review( db: DBDep, x_admin_token: AdminToken = None, risk: str | None = None, flagged: bool = False, low_confidence: bool = False, escalated: bool = False, limit: int = 50, offset: int = 0, format: str = "json", ) -> dict[str, Any] | Response: settings = get_settings() if not hmac.compare_digest(x_admin_token or "", settings.admin_token): raise HTTPException(status_code=401, detail="invalid admin token") sessions = {session.id: session for session in db.exec(select(crud.SessionRecord))} feedback: dict[str, list[crud.FeedbackRecord]] = {} for item in db.exec(select(crud.FeedbackRecord)): feedback.setdefault(item.turn_id, []).append(item) # ponytail: Python scan is enough for single-clinic SQLite MVP. rows: list[dict[str, Any]] = [] for turn in db.exec(select(crud.TurnRecord).order_by(desc(crud.TurnRecord.created_at))): turn_feedback = feedback.get(turn.id, []) session = sessions.get(turn.session_id) if risk and turn.risk_tier != risk: continue if flagged and not turn_feedback: continue if low_confidence and ( turn.asr_confidence >= settings.confidence_threshold and turn.mt_confidence >= settings.confidence_threshold ): continue if escalated and session and session.status != "escalated": continue payload = turn_payload(turn) payload["session_status"] = session.status if session else "unknown" payload["feedback"] = [feedback_payload(item) for item in turn_feedback] rows.append(payload) total = len(rows) page = rows[offset : offset + limit] if format == "csv": return review_csv(page) return {"items": page, "total": total, "limit": limit, "offset": offset} def review_csv(rows: list[dict[str, Any]]) -> Response: handle = StringIO() fieldnames = [ "id", "session_id", "speaker", "source_text", "translation", "corrected_text", "risk_tier", "status", "session_status", "feedback_reasons", ] writer = csv.DictWriter(handle, fieldnames=fieldnames) writer.writeheader() for row in rows: writer.writerow( { key: row.get(key, "") for key in fieldnames if key != "feedback_reasons" } | {"feedback_reasons": ";".join(item["reason"] for item in row["feedback"])} ) return Response(content=handle.getvalue(), media_type="text/csv") @ws_router.websocket("/ws/sessions/{session_id}") async def websocket_session( websocket: WebSocket, session_id: str, db: DBDep, ) -> None: try: await websocket.accept() except RuntimeError: return settings = get_settings() providers = get_providers(settings) try: if crud.get_session(db, session_id) is None: await websocket.close(code=1008, reason="session not found") return await websocket.send_json( { "type": "session_state", "turns": [turn_payload(turn) for turn in crud.list_turns(db, session_id)], } ) while True: message = await websocket.receive() if message["type"] == "websocket.disconnect": break if message.get("bytes") is not None: current = _in_flight.get(session_id) if current is None: await websocket.send_json( {"type": "turn_error", "message": "no turn in progress", "retryable": True} ) continue chunk = message["bytes"] current.audio_bytes += len(chunk) if current.audio_bytes > settings.max_turn_audio_bytes: _in_flight.pop(session_id, None) await websocket.send_json( { "type": "turn_error", "message": "audio turn too large", "retryable": True, } ) continue current.chunks.append(chunk) continue event = json.loads(message.get("text") or "{}") event_type = event.get("type") if event_type in {"start_turn", "end_turn", "text_turn"} and _session_is_ended( db, session_id ): await _send_session_ended(websocket, session_id) continue if event_type == "start_turn": if session_id in _in_flight: await websocket.send_json( { "type": "turn_error", "message": "turn already in progress", "retryable": True, } ) continue _in_flight[session_id] = InFlightTurn( speaker=event["speaker"], lang=event["lang"], ) elif event_type == "end_turn": current = _in_flight.pop(session_id, None) if current is None: await websocket.send_json( {"type": "turn_error", "message": "no turn in progress", "retryable": True} ) continue try: turn = await anyio.to_thread.run_sync( partial( process_audio_turn, db, providers, settings, session_id=session_id, speaker=current.speaker, lang=current.lang, audio=b"".join(current.chunks), ) ) except Exception as exc: await _send_pipeline_error( websocket, session_id, exc, speaker=current.speaker, lang=current.lang, source_text=None, ) continue log_turn_processed(turn) await websocket.send_json(turn_result_payload(turn, settings)) elif event_type == "text_turn": if len(event.get("text", "")) > MAX_TYPED_TEXT_CHARS: await websocket.send_json( { "type": "turn_error", "message": "typed turn too long", "retryable": True, } ) continue try: turn = await anyio.to_thread.run_sync( partial( process_text_turn, db, providers, settings, session_id=session_id, speaker=event["speaker"], lang=event["lang"], text=event["text"], asr_confidence=1.0, ) ) except Exception as exc: await _send_pipeline_error( websocket, session_id, exc, speaker=event["speaker"], lang=event["lang"], source_text=event["text"], ) continue log_turn_processed(turn) await websocket.send_json(turn_result_payload(turn, settings)) else: await websocket.send_json( {"type": "turn_error", "message": "unknown event type", "retryable": False} ) except (RuntimeError, WebSocketDisconnect): pass finally: _in_flight.pop(session_id, None)