from __future__ import annotations import json from collections.abc import Iterator from typing import Annotated from fastapi import APIRouter, Query, Request from fastapi.responses import StreamingResponse from backend.data.gateway import MarketDataUnavailable from backend.features.accounts.auth import ( AuthenticatedPrincipal, SmartAccessPrincipal, SmartWritePrincipal, ) from backend.features.mentor.schemas import ( CountResponse, MentorChatInput, MentorPreferencesInput, ) from backend.features.mentor.service import MentorError from backend.http.errors import AppError from backend.llm.gateway import LLMGatewayError router = APIRouter(prefix="/mentors", tags=["mentors"]) @router.get("/setup", response_model=dict) def setup( request: Request, principal: AuthenticatedPrincipal, trade_date: Annotated[str, Query(alias="date", min_length=10, max_length=10)], ) -> dict: return _call(request, "setup", principal, trade_date) @router.put("/preferences", response_model=dict) def save_preferences( payload: MentorPreferencesInput, request: Request, principal: SmartWritePrincipal, ) -> dict: return _call(request, "save_preferences", principal, payload.order, payload.pinned) @router.get("/messages", response_model=list[dict]) def messages( request: Request, principal: SmartAccessPrincipal, mentor_id: Annotated[str, Query(min_length=1, max_length=100)], trade_date: Annotated[str, Query(alias="date", min_length=10, max_length=10)], ) -> list[dict]: return _call(request, "messages", principal, mentor_id, trade_date) @router.delete("/messages", response_model=CountResponse) def clear_messages( request: Request, principal: SmartWritePrincipal, mentor_id: Annotated[str, Query(min_length=1, max_length=100)], trade_date: Annotated[str, Query(alias="date", min_length=10, max_length=10)], ) -> CountResponse: deleted = _call(request, "clear", principal, mentor_id, trade_date) return CountResponse(deleted=deleted) @router.post("/chat") def chat( payload: MentorChatInput, request: Request, principal: SmartWritePrincipal, ) -> StreamingResponse: prepared = _call( request, "prepare_chat", principal, payload.mentor_id, payload.trade_date, payload.question, ) def body() -> Iterator[bytes]: for event in request.app.state.container.mentor.stream_chat(prepared): yield (json.dumps(event, ensure_ascii=False, separators=(",", ":")) + "\n").encode( "utf-8" ) return StreamingResponse( body(), media_type="application/x-ndjson", headers={"Cache-Control": "no-cache, no-transform", "X-Accel-Buffering": "no"}, ) def _call(request: Request, method: str, *args): try: return getattr(request.app.state.container.mentor, method)(*args) except MentorError as exc: raise AppError("mentor_unavailable", str(exc), 409) from exc except MarketDataUnavailable as exc: raise AppError("market_data_unavailable", str(exc), 503) from exc except LLMGatewayError as exc: status = 403 if exc.code in {"membership_required", "quota_exhausted"} else 503 raise AppError(exc.code, str(exc), status) from exc