Spaces:
Running
Running
| """Author RAG — Chat API Router. | |
| Routes (per implementation plan): | |
| POST /api/chat/{author_slug} — Main chat endpoint (synchronous) | |
| POST /api/chat/{author_slug}/stream-start — Phase 7: Queue job, return job_id | |
| GET /api/chat/{author_slug}/poll/{job_id} — Phase 7: Poll for result | |
| POST /api/chat/{author_slug}/session/init — Initialize session | |
| GET /api/chat/{author_slug}/session/history/{session_id} | |
| POST /api/chat/{author_slug}/session/rate — Star rating | |
| POST /api/chat/{author_slug}/session/reset — Reset session | |
| GET /api/chat/{author_slug}/events — SSE for ingestion progress | |
| RULE: Requires X-Subscription-Token header on all routes. | |
| RULE: Visitor fingerprint is generated server-side. | |
| """ | |
| import asyncio | |
| import hashlib | |
| import json | |
| import uuid | |
| import structlog | |
| from fastapi import APIRouter, Depends, Header, HTTPException, Request | |
| from fastapi.responses import StreamingResponse | |
| from sqlalchemy import select, update | |
| from sqlalchemy.ext.asyncio import AsyncSession | |
| from app.dependencies import get_db, get_redis, get_subscription_author, _get_session_factory | |
| from app.models.chat_session import ChatMessage, ChatSession | |
| from app.services.analytics_core.geo import get_geo_info, get_real_ip | |
| from app.services.analytics_core.tracker import parse_device_info, track_turn | |
| from app.services.analytics_core.visitor_tracker import record_visitor | |
| from app.services.formatter import ResponseFormatter | |
| from app.services.prompter import FAREWELL_RESPONSE | |
| from app.services.rag_pipeline import run_pipeline | |
| from app.services.session_core.manager import SessionManager | |
| from app.services.token_budget import check_and_warn_budget | |
| from app.repositories.book_repo import BookRepository | |
| from app.repositories.link_repo import LinkRepository | |
| from app.schemas.chatbot import ( | |
| ChatRequest, | |
| ChatResponse, | |
| FarewellRequest, | |
| FeedbackRequest, | |
| PollResponse, | |
| RateSessionRequest, | |
| SessionInitRequest, | |
| SessionInitResponse, | |
| StreamStartResponse, | |
| TrackClickRequest, | |
| ) | |
| router = APIRouter() | |
| logger = structlog.get_logger(__name__) | |
| # Redis TTL for poll job results (10 minutes — generous for slow mobile connections) | |
| _JOB_TTL = 600 | |
| _JOB_PREFIX = "chat_job:" | |
| def _fingerprint(request: Request) -> str: | |
| """SHA-256 fingerprint from IP + UA (no PII stored). | |
| B1 fix: request.client can be None in Docker/proxy environments. | |
| Now reads X-Forwarded-For as fallback before defaulting to 'unknown'. | |
| """ | |
| ip = ( | |
| request.client.host | |
| if request.client | |
| else request.headers.get("X-Forwarded-For", "unknown").split(",")[0].strip() | |
| ) | |
| components = [ | |
| ip, | |
| request.headers.get("User-Agent", ""), | |
| request.headers.get("Accept-Language", ""), | |
| ] | |
| return hashlib.sha256("|".join(components).encode()).hexdigest()[:32] | |
| async def init_session( | |
| author_slug: str, | |
| request: Request, | |
| body: SessionInitRequest = SessionInitRequest(), | |
| author=Depends(get_subscription_author), | |
| db: AsyncSession = Depends(get_db), | |
| redis=Depends(get_redis), | |
| ): | |
| """Initialize a new chat session for a visitor. | |
| Accepts an optional visitor_uid from the widget's localStorage. | |
| Returns visitor_uid so the widget can store/update it in localStorage. | |
| This approach is HF-Spaces-safe — no cookies needed (iframe-safe). | |
| """ | |
| try: | |
| session_id = str(uuid.uuid4()) | |
| visitor_fp = _fingerprint(request) | |
| # ── Visitor deduplication ────────────────────────────────────────────── | |
| # record_visitor upserts the Visitor table: | |
| # - New visitor → creates row, returns (new_uid, True) | |
| # - Known visitor → updates lastSeen + pageViews, returns (uid, False) | |
| # Failures are caught internally — visitor_uid always returns a valid UUID. | |
| visitor_uid, _is_new = await record_visitor( | |
| request=request, | |
| author_id=author.id, | |
| db=db, | |
| redis=redis, | |
| visitor_uid=body.visitor_uid, | |
| ) | |
| book_repo = BookRepository(db) | |
| link_repo = LinkRepository(db) | |
| active_books = await book_repo.list_active_for_author(author.id) | |
| books_payload = [] | |
| for b in active_books: | |
| link = await link_repo.get_for_book(b.id, author.id) | |
| books_payload.append({ | |
| "id": b.id, | |
| "title": b.title, | |
| "tagline": b.tagline, | |
| "cover_path": b.cover_path, | |
| "ai_summary": b.ai_summary, | |
| "purchase_url": link.purchase_url if link else None, | |
| }) | |
| base_welcome = (author.welcome_message or "Hello!").strip() | |
| if "select a book" not in base_welcome.lower(): | |
| welcome = f"{base_welcome}\n\nSelect a book below to ask about it." | |
| else: | |
| welcome = base_welcome | |
| geo = await get_geo_info(request, redis=redis) | |
| device = parse_device_info(request) | |
| session = ChatSession( | |
| id=session_id, | |
| author_id=author.id, | |
| visitor_fingerprint=visitor_fp, | |
| visitor_uid=visitor_uid, | |
| ip_address=get_real_ip(request), | |
| country_code=geo.get("country_code") or None, | |
| country_name=geo.get("country_name") or None, | |
| city=geo.get("city") or None, | |
| device_type=device.get("device_type"), | |
| browser=device.get("browser"), | |
| os=device.get("os"), | |
| ) | |
| db.add(session) | |
| await db.commit() | |
| return SessionInitResponse( | |
| session_id=session_id, | |
| bot_name=author.bot_name or "Book Advisor", | |
| welcome_message=welcome, | |
| widget_theme=author.widget_theme or "midnight", | |
| books=books_payload, | |
| show_book_selector=len(active_books) > 0, | |
| visitor_uid=visitor_uid, | |
| ) | |
| except HTTPException: | |
| raise | |
| except Exception as e: | |
| logger.error("session init failed", author_id=author.id, error=str(e), exc_info=True) | |
| raise HTTPException(status_code=500, detail="Could not start chat session") | |
| async def chat( | |
| author_slug: str, | |
| payload: ChatRequest, | |
| author=Depends(get_subscription_author), | |
| db: AsyncSession = Depends(get_db), | |
| redis=Depends(get_redis), | |
| ): | |
| """Process one chat message through the 12-step RAG pipeline.""" | |
| session_mgr = SessionManager(redis) | |
| session_ctx = await session_mgr.load(payload.session_id, author.id) | |
| if payload.ask_all_books: | |
| session_ctx.is_cross_book = True | |
| session_ctx.selected_book_id = None | |
| await session_mgr.set_selected_book( | |
| payload.session_id, author.id, None, is_cross=True, | |
| ) | |
| elif payload.selected_book_id is not None: | |
| session_ctx.is_cross_book = False | |
| session_ctx.selected_book_id = payload.selected_book_id | |
| await session_mgr.set_selected_book( | |
| payload.session_id, author.id, payload.selected_book_id, | |
| ) | |
| result = await run_pipeline( | |
| query=payload.message, | |
| author=author, | |
| session_context=session_ctx, | |
| db=db, | |
| ) | |
| await session_mgr.save( | |
| context=session_ctx, | |
| user_message=payload.message, | |
| assistant_message=result.response["text"], | |
| ) | |
| await track_turn( | |
| db=db, | |
| redis=redis, | |
| session_id=payload.session_id, | |
| author_id=author.id, | |
| book_id=session_ctx.selected_book_id, | |
| user_message=payload.message, | |
| result=result, | |
| # DESIGN-5 fix: session_mgr.save() has already incremented turn_count above, | |
| # so this is the correct completed-turn number for this event row. | |
| turn_number=session_ctx.turn_count, | |
| ) | |
| # Phase 5A (B8 fix): Warn author by email when token budget hits 80%. | |
| # Non-blocking: gracefully skips if SMTP not configured or email already sent. | |
| if hasattr(author, "_access_grant") and author._access_grant: | |
| await check_and_warn_budget(redis, author._access_grant, author) | |
| return ChatResponse( | |
| text=result.response["text"], | |
| links=result.response["links"], | |
| has_links=result.response["has_links"], | |
| book_selector=result.response.get("book_selector"), | |
| session_id=payload.session_id, | |
| intent=result.intent, | |
| ) | |
| # ── Phase 7: Polling Streaming Endpoints ───────────────────────────────────── | |
| # Problem: nginx on HuggingFace Spaces buffers SSE responses, breaking streaming. | |
| # Solution: job_id + polling — no buffering issues, works with standard nginx. | |
| # | |
| # Flow: | |
| # 1. Widget POSTs to /stream-start → gets job_id immediately (< 50ms) | |
| # 2. Widget shows "thinking..." spinner | |
| # 3. Widget polls GET /poll/{job_id} every 800ms | |
| # 4. When status='done', widget renders the response and stops polling | |
| # | |
| # The pipeline runs as an asyncio.create_task() — non-blocking on the event loop. | |
| # Result stored in Redis with 10min TTL (generous for slow mobile connections). | |
| async def stream_start( | |
| author_slug: str, | |
| payload: ChatRequest, | |
| request: Request, | |
| author=Depends(get_subscription_author), | |
| redis=Depends(get_redis), | |
| ): | |
| """Phase 7: Queue a chat pipeline job, return job_id immediately. | |
| The pipeline runs in the background via asyncio.create_task(). | |
| Widget polls GET /poll/{job_id} every 800ms until status='done'. | |
| This is HF-Spaces-safe: no SSE, no streaming response, no nginx buffering. | |
| IMPORTANT: The background task opens its own DB session via _get_session_factory(). | |
| The request-scoped session from Depends(get_db) closes when this route returns, | |
| which would cause 'Session is closed' errors in the task if we shared it. | |
| """ | |
| job_id = str(uuid.uuid4()) | |
| job_key = f"{_JOB_PREFIX}{job_id}" | |
| # Store initial 'pending' status immediately so poll() returns something | |
| await redis.setex(job_key, _JOB_TTL, json.dumps({"status": "pending"})) | |
| async def _run_pipeline_task() -> None: | |
| """Inner task: opens its own DB session, runs pipeline, writes to Redis.""" | |
| async with _get_session_factory()() as db: | |
| try: | |
| session_mgr = SessionManager(redis) | |
| session_ctx = await session_mgr.load(payload.session_id, author.id) | |
| if payload.ask_all_books: | |
| session_ctx.is_cross_book = True | |
| session_ctx.selected_book_id = None | |
| await session_mgr.set_selected_book( | |
| payload.session_id, author.id, None, is_cross=True, | |
| ) | |
| elif payload.selected_book_id is not None: | |
| session_ctx.is_cross_book = False | |
| session_ctx.selected_book_id = payload.selected_book_id | |
| await session_mgr.set_selected_book( | |
| payload.session_id, author.id, payload.selected_book_id, | |
| ) | |
| pipeline_result = await run_pipeline( | |
| query=payload.message, | |
| author=author, | |
| session_context=session_ctx, | |
| db=db, | |
| ) | |
| await session_mgr.save( | |
| context=session_ctx, | |
| user_message=payload.message, | |
| assistant_message=pipeline_result.response["text"], | |
| ) | |
| await track_turn( | |
| db=db, | |
| redis=redis, | |
| session_id=payload.session_id, | |
| author_id=author.id, | |
| book_id=session_ctx.selected_book_id, | |
| user_message=payload.message, | |
| result=pipeline_result, | |
| # DESIGN-5 fix: same as sync path — turn_count already incremented. | |
| turn_number=session_ctx.turn_count, | |
| ) | |
| if hasattr(author, "_access_grant") and author._access_grant: | |
| await check_and_warn_budget(redis, author._access_grant, author) | |
| response_payload = ChatResponse( | |
| text=pipeline_result.response["text"], | |
| links=pipeline_result.response["links"], | |
| has_links=pipeline_result.response["has_links"], | |
| book_selector=pipeline_result.response.get("book_selector"), | |
| session_id=payload.session_id, | |
| intent=pipeline_result.intent, | |
| ) | |
| await redis.setex( | |
| job_key, | |
| _JOB_TTL, | |
| json.dumps({"status": "done", "result": response_payload.model_dump()}), | |
| ) | |
| except Exception as exc: | |
| logger.error("stream pipeline task failed", job_id=job_id, error=str(exc)) | |
| await redis.setex( | |
| job_key, | |
| _JOB_TTL, | |
| json.dumps({"status": "error", "error": str(exc)[:200]}), | |
| ) | |
| # Schedule pipeline — returns immediately, task runs concurrently | |
| asyncio.create_task(_run_pipeline_task()) | |
| return StreamStartResponse(job_id=job_id, status="queued") | |
| async def poll_job( | |
| author_slug: str, | |
| job_id: str, | |
| author=Depends(get_subscription_author), | |
| redis=Depends(get_redis), | |
| ): | |
| """Phase 7: Poll for the result of a stream-start job. | |
| Returns: | |
| {status: 'pending'} while pipeline is still running | |
| {status: 'done', result: ChatResponse} when complete | |
| {status: 'error', error: str} on pipeline failure | |
| {status: 'expired'} if job_id is no longer in Redis (10min TTL) | |
| Widget should: | |
| - Poll every 800ms | |
| - Stop polling when status is 'done', 'error', or 'expired' | |
| - Show fallback error UI on 'error' or 'expired' | |
| """ | |
| job_key = f"{_JOB_PREFIX}{job_id}" | |
| raw = await redis.get(job_key) | |
| if not raw: | |
| return PollResponse(job_id=job_id, status="expired") | |
| try: | |
| data = json.loads(raw) | |
| except Exception: | |
| return PollResponse(job_id=job_id, status="error", error="Corrupt job data") | |
| status = data.get("status", "pending") | |
| if status == "done" and "result" in data: | |
| return PollResponse( | |
| job_id=job_id, | |
| status="done", | |
| result=ChatResponse(**data["result"]), | |
| ) | |
| if status == "error": | |
| return PollResponse( | |
| job_id=job_id, | |
| status="error", | |
| error=data.get("error", "Unknown error"), | |
| ) | |
| return PollResponse(job_id=job_id, status="pending") | |
| async def session_farewell( | |
| author_slug: str, | |
| payload: FarewellRequest, | |
| author=Depends(get_subscription_author), | |
| db: AsyncSession = Depends(get_db), | |
| ): | |
| """Return a short farewell with Buy Book button when the visitor closes chat.""" | |
| book_repo = BookRepository(db) | |
| link_repo = LinkRepository(db) | |
| formatter = ResponseFormatter() | |
| book_title = "this book" | |
| purchase_url = None | |
| preview_url = None | |
| if payload.selected_book_id: | |
| book = await book_repo.get_by_id_for_author(payload.selected_book_id, author.id) | |
| if book: | |
| book_title = book.title | |
| link = await link_repo.get_for_book(book.id, author.id) | |
| if link: | |
| purchase_url = link.purchase_url | |
| preview_url = link.preview_url | |
| text = FAREWELL_RESPONSE.format(book_title=book_title) | |
| formatted = formatter.purchase_only(text, purchase_url, preview_url) | |
| return ChatResponse( | |
| text=formatted["text"], | |
| links=formatted["links"], | |
| has_links=formatted["has_links"], | |
| session_id=payload.session_id, | |
| intent="farewell", | |
| ) | |
| async def get_history( | |
| author_slug: str, | |
| session_id: str, | |
| author=Depends(get_subscription_author), | |
| redis=Depends(get_redis), | |
| ): | |
| """Return full chat history for a session.""" | |
| session_mgr = SessionManager(redis) | |
| ctx = await session_mgr.load(session_id, author.id) | |
| return {"session_id": session_id, "history": ctx.history, "turn_count": ctx.turn_count} | |
| # R-046: First rate_session definition removed — was duplicate of L346 which has proper body validation | |
| async def reset_session( | |
| author_slug: str, | |
| session_id: str, | |
| author=Depends(get_subscription_author), | |
| redis=Depends(get_redis), | |
| ): | |
| """Clear the session context from Redis (start fresh).""" | |
| session_mgr = SessionManager(redis) | |
| await session_mgr.delete(session_id, author.id) | |
| return {"message": "Session reset"} | |
| async def ingestion_events( | |
| author_slug: str, | |
| author=Depends(get_subscription_author), | |
| redis=Depends(get_redis), | |
| ): | |
| """SSE stream for real-time ingestion progress updates.""" | |
| async def event_stream(): | |
| pubsub = redis.pubsub() | |
| channel = f"ingestion:{author.id}" | |
| await pubsub.subscribe(channel) | |
| try: | |
| # R-023: Max connection timeout to prevent connection exhaustion | |
| import time as _time | |
| deadline = _time.monotonic() + 300 # 5 minute max | |
| heartbeat_interval = 30 | |
| last_heartbeat = _time.monotonic() | |
| async for message in pubsub.listen(): | |
| now = _time.monotonic() | |
| if now >= deadline: | |
| yield 'data: {"type": "timeout"}\n\n' | |
| break | |
| if now - last_heartbeat > heartbeat_interval: | |
| yield ": heartbeat\n\n" | |
| last_heartbeat = now | |
| if message["type"] == "message": | |
| yield f"data: {message['data']}\n\n" | |
| finally: | |
| await pubsub.unsubscribe(channel) | |
| return StreamingResponse( | |
| event_stream(), | |
| media_type="text/event-stream", | |
| headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, | |
| ) | |
| async def track_click( | |
| author_slug: str, | |
| body: TrackClickRequest, | |
| author=Depends(get_subscription_author), | |
| db: AsyncSession = Depends(get_db), | |
| ): | |
| """Track when a visitor clicks a purchase/preview link. R-029: Schema validated.""" | |
| if body.session_id: | |
| await db.execute( | |
| update(ChatSession) | |
| .where(ChatSession.id == body.session_id, ChatSession.author_id == author.id) | |
| .values(link_clicked=True) | |
| ) | |
| from app.services.analytics_core.tracker import track_link_click | |
| await track_link_click(db, body.session_id, author.id, body.book_id) | |
| await db.commit() | |
| return {"tracked": True, "link_type": body.link_type} | |
| async def message_feedback( | |
| author_slug: str, | |
| body: FeedbackRequest, | |
| author=Depends(get_subscription_author), | |
| db: AsyncSession = Depends(get_db), | |
| ): | |
| """Record thumbs up/down feedback for a bot message. R-029: Schema validated.""" | |
| if body.session_id and body.message_index is not None: | |
| result = await db.execute( | |
| select(ChatSession).where( | |
| ChatSession.id == body.session_id, | |
| ChatSession.author_id == author.id, | |
| ) | |
| ) | |
| session = result.scalar_one_or_none() | |
| if not session: | |
| raise HTTPException(404, "Session not found") | |
| msg_result = await db.execute( | |
| select(ChatMessage) | |
| .where(ChatMessage.session_id == body.session_id, ChatMessage.role == "assistant") | |
| .order_by(ChatMessage.created_at.asc()) | |
| ) | |
| messages = msg_result.scalars().all() | |
| if 0 <= body.message_index < len(messages): | |
| messages[body.message_index].user_feedback = body.reaction | |
| await db.commit() | |
| return {"status": "recorded", "reaction": body.reaction} | |
| async def rate_session( | |
| author_slug: str, | |
| body: RateSessionRequest, | |
| author=Depends(get_subscription_author), | |
| db: AsyncSession = Depends(get_db), | |
| ): | |
| """Record star rating (1-5) for a session. R-029: Schema validated.""" | |
| if body.session_id: | |
| await db.execute( | |
| update(ChatSession) | |
| .where(ChatSession.id == body.session_id, ChatSession.author_id == author.id) | |
| .values(rating=body.rating) | |
| ) | |
| await db.commit() | |
| return {"status": "rated", "rating": body.rating} | |