Arag / app /api /chat.py
AuthorBot
Fix analytics accuracy, tenant isolation, and event-loop blocking.
ff0847b
Raw
History Blame Contribute Delete
21.1 kB
"""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]
@router.post("/{author_slug}/session/init", response_model=SessionInitResponse, status_code=201)
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")
@router.post("/{author_slug}", response_model=ChatResponse)
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).
@router.post("/{author_slug}/stream-start", response_model=StreamStartResponse)
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")
@router.get("/{author_slug}/poll/{job_id}", response_model=PollResponse)
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")
@router.post("/{author_slug}/session/farewell", response_model=ChatResponse)
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",
)
@router.get("/{author_slug}/session/history/{session_id}")
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
@router.post("/{author_slug}/session/reset")
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"}
@router.get("/{author_slug}/events")
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"},
)
@router.post("/{author_slug}/track-click")
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}
@router.post("/{author_slug}/feedback")
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}
@router.post("/{author_slug}/session/rate")
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}