"""FastAPI entrypoint — runs internally on :8000, fronted by Next.js on :7860.""" from __future__ import annotations # ── Load .env FIRST so all module-level os.environ.get() calls in local # modules (auth, rag, hf_sync) pick up the values when they are imported. import asyncio import json import os from pathlib import Path from dotenv import load_dotenv _env_file = Path(__file__).parent / ".env" if not _env_file.exists(): _env_file = Path(__file__).parent.parent / ".env" load_dotenv(_env_file) import logging import shutil import uuid from contextlib import asynccontextmanager from typing import Any, AsyncIterator, List, Literal, Optional from fastapi import Depends, FastAPI, File, Form, HTTPException, Response, UploadFile from fastapi.responses import StreamingResponse from huggingface_hub import AsyncInferenceClient from pydantic import BaseModel, Field # Local (imported AFTER load_dotenv so their module-level env reads are correct) import auth import config import hf_sync import personas_store from rag import RagEngine logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") logger = logging.getLogger("backend") TMP_UPLOAD_DIR = config.TMP_UPLOAD_DIR LANCEDB_DIR = config.LANCEDB_DIR LLM_MODEL = config.LLM_MODEL LLM_MODEL_LOWER = config.LLM_MODEL_LOWER EMBED_MODEL = config.EMBED_MODEL HF_TOKEN = config.HF_TOKEN LLM_EXTRA_BODY_JSON = config.LLM_EXTRA_BODY_JSON LLM_FINAL_ANSWER_EXTRA_BODY_JSON = config.LLM_FINAL_ANSWER_EXTRA_BODY_JSON SHOW_LLM_REASONING = config.SHOW_LLM_REASONING LLM_MAX_TOKENS = config.LLM_MAX_TOKENS LLM_FINAL_ANSWER_MAX_TOKENS = config.LLM_FINAL_ANSWER_MAX_TOKENS # Globals populated in lifespan. rag: Optional[RagEngine] = None llm_client: Optional[AsyncInferenceClient] = None def _load_json_object_env(raw: str, env_name: str) -> Optional[dict[str, Any]]: if not raw: return None try: loaded = json.loads(raw) except json.JSONDecodeError as exc: logger.warning("Ignoring invalid %s: %s", env_name, exc) return None if not isinstance(loaded, dict): logger.warning("Ignoring %s because it is not a JSON object", env_name) return None return loaded def _is_qwen_thinking_model() -> bool: return "qwen/qwen3" in LLM_MODEL_LOWER or "qwen3" in LLM_MODEL_LOWER def _non_thinking_extra_body() -> dict[str, Any]: return { "chat_template_kwargs": {"enable_thinking": False}, # Some HF router/provider combinations look for this at the top level. "enable_thinking": False, } def _chat_extra_body() -> Optional[dict[str, Any]]: configured = _load_json_object_env(LLM_EXTRA_BODY_JSON, "LLM_EXTRA_BODY_JSON") if configured is not None: return configured if _is_qwen_thinking_model(): return _non_thinking_extra_body() return None def _final_answer_extra_body() -> Optional[dict[str, Any]]: configured = _load_json_object_env( LLM_FINAL_ANSWER_EXTRA_BODY_JSON, "LLM_FINAL_ANSWER_EXTRA_BODY_JSON", ) if configured is not None: return configured if _is_qwen_thinking_model(): return _non_thinking_extra_body() return None CHAT_EXTRA_BODY = _chat_extra_body() FINAL_ANSWER_EXTRA_BODY = _final_answer_extra_body() def _sync_vector_cache_sync() -> bool: if not hf_sync.is_configured(): return False try: return hf_sync.upload_vector_store(LANCEDB_DIR, hf_sync.list_remote_pdfs(), EMBED_MODEL) except Exception as exc: # noqa: BLE001 logger.warning("Vector cache sync failed: %s", exc) return False def _prepare_local_vector_cache_sync() -> tuple[bool, List[str]]: """Restore the local LanceDB folder from the dataset when a valid cache exists.""" if not hf_sync.is_configured(): return False, [] remote_pdfs = hf_sync.list_remote_pdfs() cache_restored = hf_sync.download_vector_store(LANCEDB_DIR, remote_pdfs, EMBED_MODEL) if not cache_restored: shutil.rmtree(LANCEDB_DIR, ignore_errors=True) return cache_restored, remote_pdfs def _cold_start_index_sync() -> None: if rag is None: return sync_dir = TMP_UPLOAD_DIR / "_sync" sync_dir.mkdir(parents=True, exist_ok=True) paths = hf_sync.download_all_pdfs(sync_dir) if paths: logger.info("Cold-start: indexing %d PDF(s) from dataset", len(paths)) added = rag.index_many((p, p.name) for p in paths) logger.info("Cold-start: %d nodes indexed", added) _sync_vector_cache_sync() else: logger.info("Cold-start: dataset has no PDFs yet") _sync_vector_cache_sync() def _upload_and_index_sync(tmp_path: Path, safe_name: str) -> tuple[int, bool]: if rag is None: return 0, False hf_sync.upload_pdf(tmp_path, safe_name) nodes_added = rag.index_pdf(tmp_path, safe_name) return nodes_added, _sync_vector_cache_sync() def _delete_and_sync_sync(filename: str) -> tuple[bool, int, bool]: removed = hf_sync.delete_pdf(filename) nodes_removed = rag.delete_by_filename(filename) if rag else 0 cache_synced = _sync_vector_cache_sync() if removed or nodes_removed > 0 else False return removed, nodes_removed, cache_synced @asynccontextmanager async def lifespan(app: FastAPI): # noqa: ARG001 global rag, llm_client TMP_UPLOAD_DIR.mkdir(parents=True, exist_ok=True) if not HF_TOKEN: logger.warning("HF_TOKEN is not set — chat & embeddings will fail until configured.") cache_restored = False remote_pdfs: List[str] = [] if hf_sync.is_configured(): try: cache_restored, remote_pdfs = await asyncio.to_thread(_prepare_local_vector_cache_sync) except Exception as exc: # noqa: BLE001 logger.exception("Vector cache restore failed: %s", exc) rag = RagEngine() llm_client = AsyncInferenceClient(model=LLM_MODEL, token=HF_TOKEN) # Load personas (HF Dataset → local file → bundled defaults). personas_store.load() # Cold-start sync: pull all PDFs from the linked dataset and rebuild the index. try: if cache_restored: logger.info("Cold-start: restored vector cache from dataset for %d PDF(s)", len(remote_pdfs)) elif hf_sync.is_configured(): await asyncio.to_thread(_cold_start_index_sync) except Exception as exc: # noqa: BLE001 logger.exception("Cold-start sync failed: %s", exc) yield app = FastAPI(title="IamEarthDev RAG", lifespan=lifespan) # ---------------- Models ---------------- class ChatMessage(BaseModel): role: Literal["system", "user", "assistant"] content: str class ChatRequest(BaseModel): messages: List[ChatMessage] persona_prompt: str = Field(default="", description="Persona-specific system prompt") temperature: float = Field(default=0.4, ge=0.0, le=1.5) top_k: int = Field(default=4, ge=1, le=10) class LoginRequest(BaseModel): passcode: str class DeleteRequest(BaseModel): filename: str # ---------------- Health & files ---------------- @app.get("/health") async def health() -> dict: return { "ok": True, "model": LLM_MODEL, "embed_model": EMBED_MODEL, "dataset_id": config.DATASET_ID, "admin_passcode_defaulted": config.ADMIN_PASSCODE == "password" and "ADMIN_PASSCODE" not in os.environ, "hf_token_configured": bool(HF_TOKEN), "chat_extra_body": CHAT_EXTRA_BODY, "final_answer_extra_body": FINAL_ANSWER_EXTRA_BODY, "show_reasoning": SHOW_LLM_REASONING, } @app.get("/files") async def list_files(_: None = Depends(auth.require_admin)) -> dict: return {"files": hf_sync.list_remote_pdfs()} @app.get("/documents") async def list_documents() -> dict: """Public read-only list of uploaded course materials.""" if not hf_sync.is_configured(): return {"files": []} return {"files": hf_sync.list_remote_pdfs()} # ---------------- Admin auth ---------------- @app.post("/admin/login") async def admin_login(body: LoginRequest, response: Response) -> dict: if not auth.verify_passcode(body.passcode): raise HTTPException(status_code=401, detail="Invalid passcode") response.set_cookie( key=auth.COOKIE_NAME, value=auth.issue_token(), httponly=True, samesite="lax", secure=False, # HF Spaces serves over https at the edge; cookie still works. max_age=60 * 60 * 12, path="/", ) return {"ok": True} @app.post("/admin/logout") async def admin_logout(response: Response) -> dict: response.delete_cookie(auth.COOKIE_NAME, path="/") return {"ok": True} @app.get("/admin/me") async def admin_me(_: None = Depends(auth.require_admin)) -> dict: return {"authenticated": True} # ---------------- Personas ---------------- class PersonaItem(BaseModel): id: str name: str description: str prompt: str @app.get("/personas") async def get_personas() -> dict: """Public — students and the chat UI need this.""" return {"personas": personas_store.get_all()} @app.put("/personas") async def put_personas( personas: List[PersonaItem], _: None = Depends(auth.require_admin), ) -> dict: saved = personas_store.save([p.model_dump() for p in personas]) return {"personas": saved} @app.post("/personas/reset") async def reset_personas(_: None = Depends(auth.require_admin)) -> dict: restored = personas_store.reset() return {"personas": restored} # ---------------- Upload / Delete ---------------- @app.post("/upload") async def upload( file: UploadFile = File(...), _: None = Depends(auth.require_admin), ) -> dict: if not file.filename or not file.filename.lower().endswith(".pdf"): raise HTTPException(status_code=400, detail="Only .pdf files are accepted.") safe_name = Path(file.filename).name # strip any path components tmp_path = TMP_UPLOAD_DIR / f"{uuid.uuid4().hex}_{safe_name}" try: with tmp_path.open("wb") as out: shutil.copyfileobj(file.file, out) try: nodes_added, vector_cache_synced = await asyncio.to_thread( _upload_and_index_sync, tmp_path, safe_name ) except RuntimeError as exc: raise HTTPException(status_code=502, detail=str(exc)) from exc except Exception as exc: # noqa: BLE001 logger.exception("Indexing failed for %s", safe_name) raise HTTPException( status_code=500, detail=f"Uploaded {safe_name} to the dataset, but indexing failed: {exc}", ) from exc finally: tmp_path.unlink(missing_ok=True) return { "ok": True, "filename": safe_name, "nodes_added": nodes_added, "vector_cache_synced": vector_cache_synced, } @app.post("/delete") async def delete(body: DeleteRequest, _: None = Depends(auth.require_admin)) -> dict: safe_name = Path(body.filename).name try: removed, nodes_removed, vector_cache_synced = await asyncio.to_thread( _delete_and_sync_sync, safe_name ) except RuntimeError as exc: raise HTTPException(status_code=502, detail=str(exc)) from exc if not removed and nodes_removed == 0: raise HTTPException(status_code=404, detail=f"{safe_name} not found") return { "ok": True, "filename": safe_name, "nodes_removed": nodes_removed, "vector_cache_synced": vector_cache_synced, } # ---------------- Chat (streaming) ---------------- SYSTEM_BASE = ( "You are a helpful course assistant for students. Use the selected persona " "instructions to decide the tone, teaching style, level of directness, and " "structure of your reply. Use the provided course material excerpts as your " "primary source of truth. If the material does not contain the answer, say so " "honestly. Always cite source filenames in square brackets when you use them." ) def _build_messages(req: ChatRequest, context: str) -> List[dict]: persona = req.persona_prompt.strip() sys_parts = [SYSTEM_BASE] if persona: sys_parts.append( "Selected persona instructions. Follow these for the response style; " "only the grounding and citation rules above take priority:\n" + persona ) if context: sys_parts.append( "Course material excerpts (use these as your primary source of truth):\n\n" + context ) msgs: List[dict] = [{"role": "system", "content": "\n\n".join(sys_parts)}] for m in req.messages: if m.role == "system": continue # we control the system prompt server-side msgs.append({"role": m.role, "content": m.content}) return msgs @app.post("/chat") async def chat(req: ChatRequest) -> StreamingResponse: if rag is None or llm_client is None: raise HTTPException(status_code=503, detail="Backend not ready.") if not req.messages: raise HTTPException(status_code=400, detail="messages must not be empty.") last_user = next((m.content for m in reversed(req.messages) if m.role == "user"), "") nodes = await asyncio.to_thread(rag.retrieve, last_user, req.top_k) if last_user else [] context = rag.format_context(nodes) sources = sorted({n.node.metadata.get("source_filename", "unknown") for n in nodes}) messages = _build_messages(req, context) final_answer_messages = [ *messages, { "role": "user", "content": ( "You have already reasoned about this request. Now produce the " "student-facing reply in the selected persona style, based on the " "same course material excerpts. Do not repeat your thinking process. " "Do not produce a thinking process, analysis section, or hidden reasoning. " "Answer immediately and cite source filenames in square brackets." ), }, ] async def stream() -> AsyncIterator[bytes]: # Emit sources first as a single SSE-style event so the UI can render them. yield (json.dumps({"type": "sources", "sources": sources}) + "\n").encode() saw_content = False saw_reasoning = False finish_reason = None try: primary_stream = await llm_client.chat_completion( messages=messages, max_tokens=LLM_MAX_TOKENS, temperature=req.temperature, stream=True, extra_body=CHAT_EXTRA_BODY, ) async for chunk in primary_stream: delta = "" reasoning = "" try: finish_reason = chunk.choices[0].finish_reason or finish_reason delta = chunk.choices[0].delta.content or "" reasoning = getattr(chunk.choices[0].delta, "reasoning", "") or "" except (AttributeError, IndexError): pass if delta: saw_content = True yield (json.dumps({"type": "delta", "text": delta}) + "\n").encode() if reasoning: saw_reasoning = True if SHOW_LLM_REASONING: yield (json.dumps({"type": "reasoning", "text": reasoning}) + "\n").encode() elif not saw_content: logger.info( "Model emitted hidden reasoning before visible content; " "switching to final-answer pass early." ) close_stream = getattr(primary_stream, "close", None) if callable(close_stream): close_stream() break if not saw_content and saw_reasoning: yield ( json.dumps( { "type": "phase", "value": "finalizing", "message": "Thinking completed. Generating the final answer.", } ) + "\n" ).encode() final_stream = await llm_client.chat_completion( messages=final_answer_messages, max_tokens=LLM_FINAL_ANSWER_MAX_TOKENS, temperature=req.temperature, stream=True, extra_body=FINAL_ANSWER_EXTRA_BODY, ) async for chunk in final_stream: delta = "" try: delta = chunk.choices[0].delta.content or "" except (AttributeError, IndexError): pass if delta: saw_content = True yield (json.dumps({"type": "delta", "text": delta}) + "\n").encode() if not saw_content and saw_reasoning: detail = "The model finished its reasoning but did not emit a final answer." if finish_reason == "length": detail += ( " It likely exhausted the completion budget while thinking. " "Increase LLM_MAX_TOKENS if you want a longer first-pass reasoning budget." ) yield (json.dumps({"type": "error", "message": detail}) + "\n").encode() yield (json.dumps({"type": "done"}) + "\n").encode() except Exception as exc: # noqa: BLE001 logger.exception("LLM stream failed") yield (json.dumps({"type": "error", "message": str(exc)}) + "\n").encode() return StreamingResponse( stream(), media_type="application/x-ndjson", headers={ "X-Accel-Buffering": "no", "Cache-Control": "no-cache", "Connection": "keep-alive", }, )