"""Library router -> per-session content management. Each session gets its own isolated upload folder (see app.services.project_files) -> a session's document set is exactly what's in its own folder, never a shared, rescanned "current uploads" directory. Drag-and-drop files in, get an instant curriculum tree, chunk the full text in the background. """ import asyncio import hashlib import json import os import uuid # Cap concurrent highlight-concepts LLM calls to avoid token-burst 429s _deployment_env = os.getenv("DEPLOYMENT_ENV", "desktop") _highlight_concurrency = 2 if _deployment_env == "demo" else 20 _highlight_sem = asyncio.Semaphore(_highlight_concurrency) from typing import List, Optional from fastapi import APIRouter, BackgroundTasks, File, Form, HTTPException, UploadFile from fastapi.responses import FileResponse from pydantic import BaseModel from app.agents.brain_agent import BrainAgent from app.rag.ingestion import LIBRARY_COLLECTION, ingest_file from app.routers.session import set_ingest_status from app.schemas.graph import NodeData from app.services.project_files import list_project_files, project_upload_dir from app.websockets.handlers import get_db, get_graph_manager router = APIRouter(prefix="/library", tags=["library"]) _SESSIONS_DIR = os.path.expanduser("~/.studybuddy/sessions") _PDF_CACHE_DIR = os.path.expanduser("~/.studybuddy/pdfs") _brain: BrainAgent | None = None def _get_brain() -> BrainAgent: global _brain if _brain is None: _brain = BrainAgent() return _brain # ------------------------------------------------------------------ # # Models # # ------------------------------------------------------------------ # class HighlightConceptsRequest(BaseModel): page_text: str familiarity: str = "high_school" # ------------------------------------------------------------------ # # Helpers # # ------------------------------------------------------------------ # def _combined_document_id(file_contents: List[bytes]) -> str: """Order-independent identity for a session's whole set of uploaded files. Sorting the per-file hashes before joining means the same SET of files always produces the same id regardless of upload order or directory listing order -> callers (upload, resume, status) can compute this independently and still land on the same cache key. """ per_file = sorted(hashlib.sha256(c).hexdigest() for c in file_contents) return hashlib.sha256("".join(per_file).encode()).hexdigest() # ------------------------------------------------------------------ # # Endpoints # # ------------------------------------------------------------------ # @router.post("/highlight-concepts") async def highlight_concepts(req: HighlightConceptsRequest): """Given a page's text, return key concept phrases to highlight.""" loop = asyncio.get_event_loop() try: async with _highlight_sem: concepts = await loop.run_in_executor( None, _get_brain().identify_concepts, req.page_text, req.intention, ) except Exception: concepts = [] return {"concepts": concepts} @router.post("/upload-and-start") async def upload_and_start( background_tasks: BackgroundTasks, project_id: str = Form(...), intention: str = Form("learn"), topic_hint: str = Form(""), files: List[UploadFile] = File(...), ): """Accept direct file uploads (drag-and-drop), generate tree instantly, chunk in background. Works in both browser (no Electron needed) and Electron. - Reads document structure (headings/first pages) immediately -> no embedding needed - Generates curriculum tree in ~2s - Parses full documents into structured academic evidence in background """ if not files: raise HTTPException(400, "No files uploaded") file_store: List[tuple] = [] # (bytes, filename) for background evidence ingestion for f in files[:5]: # cap at 5 to keep prompt manageable content = await f.read() filename = f.filename or "upload" file_store.append((content, filename)) # The curriculum tree is NOT generated here -> the client opens the WebSocket and # fires BUILD_GRAPH, which streams nodes in parallel ("fireworks"). This returns fast. # Each session gets its OWN upload folder (project_id is always a fresh UUID from # /session/create) -> no shared directory, so there's nothing to clear and nothing # another session's leftover files could contaminate this one with. upload_dir = project_upload_dir(project_id) for content, filename in file_store: dest = os.path.join(upload_dir, filename) with open(dest, "wb") as fh: fh.write(content) # Cache EVERY uploaded file for layout/regions segmentation under its document_id hash # (multi-paper: each paper must be resolvable, not just the first). if file_store: pdf_dir = os.path.expanduser("~/.studybuddy/pdfs") os.makedirs(pdf_dir, exist_ok=True) for content, _fn in file_store: dest = os.path.join(pdf_dir, f"{hashlib.sha256(content).hexdigest()}.pdf") if not os.path.exists(dest): with open(dest, "wb") as fh: fh.write(content) # Rebuild canonical evidence from the project's current source files. db = get_db() from app.rag.evidence_ingestion import EvidenceIngestionService EvidenceIngestionService(db=db).reset_project(project_id) set_ingest_status(project_id, "indexing") def _chunk_all(): try: total = 0 for content, filename in file_store: total += ingest_file( content, filename, LIBRARY_COLLECTION, db=db, skip_if_indexed=False, project_id=project_id, ) set_ingest_status(project_id, "ready" if total > 0 else "error") except Exception: import logging logging.getLogger(__name__).exception( "Background structured evidence indexing failed for project %s", project_id ) set_ingest_status(project_id, "error") background_tasks.add_task(_chunk_all) # Stable document id = order-independent hash of the WHOLE uploaded file set # (same set of files, any order = same id). document_id = _combined_document_id([c for c, _ in file_store]) if file_store else "" return { "status": "ready", "nodes": [], # streamed via BUILD_GRAPH over the WebSocket "edges": [], "filenames": [name for _, name in file_store], "document_id": document_id, } # ------------------------------------------------------------------ # # History -> past learning materials, resumable from the start window # # ------------------------------------------------------------------ # @router.get("/history") def list_history(): """One entry per document set the student has committed progress on, most recent first. A document set only shows up here if at least one of its source files is still cached under ~/.studybuddy/pdfs/{file_id}.* -> otherwise there'd be nothing to resume. document_id is a combined, order-independent hash of the whole file set (see _combined_document_id), so it never corresponds to a single cached PDF filename directly -> file_ids (the per-file hashes) is what's checked against the PDF cache. """ if not os.path.isdir(_SESSIONS_DIR): return {"items": []} items = [] for name in os.listdir(_SESSIONS_DIR): if not (name.startswith("doc_") and name.endswith(".json")): continue document_id = name[len("doc_"):-len(".json")] path = os.path.join(_SESSIONS_DIR, name) try: with open(path, encoding="utf-8") as f: saved = json.load(f) except Exception: continue file_ids = saved.get("file_ids", []) if not any(os.path.exists(os.path.join(_PDF_CACHE_DIR, f"{fid}.pdf")) for fid in file_ids): continue items.append({ "document_id": document_id, "title": saved.get("title") or saved.get("topic") or "Untitled", "topic": saved.get("topic") or "Untitled", "intention": saved.get("intention", "learn"), "content_files": saved.get("content_files", []), "node_count": len(saved.get("nodes", [])), "updated_at": os.path.getmtime(path), }) items.sort(key=lambda x: x["updated_at"], reverse=True) return {"items": items} @router.get("/history/{document_id}/resume-stats") async def get_resume_stats(document_id: str): doc_json_path = os.path.join(_SESSIONS_DIR, f"doc_{document_id}.json") if not os.path.exists(doc_json_path): raise HTTPException(404, "Session not found") with open(doc_json_path, encoding="utf-8") as f: saved = json.load(f) nodes = saved.get("nodes", []) mastered = sum(1 for n in nodes if n.get("status") == "completed") active = sum(1 for n in nodes if n.get("status") == "ongoing") from app.services.student_memory import StudentMemoryService memory = StudentMemoryService() topic = saved.get("topic") or "Study Session" project_id = saved.get("project_id") or document_id # We query prior knowledge to extract a quick summary prior_knowledge = await memory.query_prior_knowledge(topic, project_id=project_id) weak_areas = [] if prior_knowledge: # Just grab the first few lines of insights lines = [line.strip() for line in prior_knowledge.split("\n") if line.strip() and not line.startswith("Prior learning context")] weak_areas = lines[:3] return { "title": saved.get("title") or saved.get("topic") or "Untitled", "last_active": os.path.getmtime(doc_json_path), "total_nodes": len(nodes), "mastered_nodes": mastered, "active_nodes": active, "weak_areas": weak_areas, "tokens_used": "N/A" # We don't have token tracking stored per-project yet } class ResumeHistoryRequest(BaseModel): document_id: str @router.post("/history/resume") async def resume_history(req: ResumeHistoryRequest): """Restore a past session's document set into its own folder and hand back its already-known curriculum tree -> no LLM regeneration, no re-chunking, no re-indexing. Chunks already exist in ChromaDB tagged with the original project_id (reused here), so they're immediately queryable again. """ doc_json_path = os.path.join(_SESSIONS_DIR, f"doc_{req.document_id}.json") if not os.path.exists(doc_json_path): raise HTTPException(404, "That session's material is no longer available.") with open(doc_json_path, encoding="utf-8") as f: saved = json.load(f) content_files = saved.get("content_files") or [] file_ids = saved.get("file_ids") or [] project_id = saved.get("project_id") or str(uuid.uuid4()) # Restore every file in this document set from its own per-file PDF cache into the # session's own folder -> content_files/file_ids are parallel lists built together # at commit time. upload_dir = project_upload_dir(project_id) restored_any = False for filename, file_id in zip(content_files, file_ids): src = os.path.join(_PDF_CACHE_DIR, f"{file_id}.pdf") if not os.path.exists(src): continue with open(src, "rb") as f: content = f.read() with open(os.path.join(upload_dir, filename), "wb") as fh: fh.write(content) restored_any = True if not restored_any: raise HTTPException(404, "That session's material is no longer available.") # No re-chunking, no re-indexing: this project_id's chunks already exist in # ChromaDB from the original session and are never deleted except by an explicit # /session/clear or a fresh upload reusing the same project_id. set_ingest_status(project_id, "ready") # The committed snapshot is the source of truth for the node set -> it always # reflects this project's current state, including Graph-Curator "exploration" # nodes that were added after the base curriculum was cached. The doc-graph # cache is the base curriculum only (content-addressed, shared across sessions # with the same papers), so it holds the real AI-generated edge relationships # but not the session-specific exploration nodes. Use snapshot nodes, take real # edges from the cache where present, and synthesize edges from parent_id for # any parent->child pair the cache doesn't cover (e.g. exploration nodes). nodes = [NodeData(**n) for n in saved.get("nodes", [])] graph_cache = get_graph_manager().load_doc_graph(req.document_id) edges = list(graph_cache[1]) if graph_cache else [] covered = {(e.get("source"), e.get("target")) for e in edges} for n in nodes: if n.parent_id and (n.parent_id, n.id) not in covered: edges.append({ "source": n.parent_id, "target": n.id, "relationship": "related" if n.origin == "exploration" else "prerequisite", }) get_graph_manager().set_graph(project_id, nodes) return { "status": "ready", "project_id": project_id, "topic": saved.get("topic") or "Study Session", "intention": saved.get("intention", "learn"), "nodes": [n.model_dump() for n in nodes], "edges": edges, "filenames": content_files, "document_id": req.document_id, } @router.get("/file/{project_id}/{filename}") def serve_library_file(project_id: str, filename: str): """Serve a file from this session's own upload folder (for browser-mode PDF viewer).""" path = os.path.join(project_upload_dir(project_id), filename) if not os.path.isfile(path): raise HTTPException(404, f"File not found: {filename}") ext = os.path.splitext(filename)[1].lower() media_types = {".pdf": "application/pdf", ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", ".txt": "text/plain"} return FileResponse(path, media_type=media_types.get(ext, "application/octet-stream")) class _RefineNodeInfo(BaseModel): id: str label: str description: str = "" depth: int = 1 parent_id: Optional[str] = None class _RefineEdgeInfo(BaseModel): source: str target: str relationship: str = "prerequisite" class RefineTreeRequest(BaseModel): project_id: str user_feedback: str intention: str = "learn" document_id: str = "" current_nodes: Optional[List[_RefineNodeInfo]] = None current_edges: Optional[List[_RefineEdgeInfo]] = None @router.post("/refine-tree") async def refine_tree(req: RefineTreeRequest): """Refine the curriculum tree based on student feedback. When current_nodes/current_edges are provided, sends the existing graph to the LLM for context-aware refinement. Falls back to full regeneration if the current graph is not supplied. """ from app.rag.models import ConsumerType, EvidenceRequest from app.rag.retrieval_service import PaperEvidenceService loop = asyncio.get_event_loop() retrieval = await loop.run_in_executor( None, lambda: PaperEvidenceService().retrieve(EvidenceRequest( query="main topics overview summary table of contents", consumer=ConsumerType.PROJECT_GRAPH, project_id=req.project_id, token_budget=9000, project_memory_policy="none", student_memory_policy="none", )), ) chunks = [ { "text": item.evidence.index_text, "source": item.evidence.document_id, "document_id": item.evidence.document_id, "evidence_id": item.evidence.evidence_id, "page_start": item.evidence.page_start, "section_path": item.evidence.section_path, } for item in retrieval.evidence ] if not chunks: raise HTTPException( 400, "Content is still being indexed -> wait a moment, then try again.", ) brain = _get_brain() if req.current_nodes and len(req.current_nodes) > 0: # Context-aware refinement -> send existing tree structure to LLM nodes, edges = await loop.run_in_executor( None, brain.refine_curriculum, chunks, req.intention, req.user_feedback, [n.model_dump() for n in req.current_nodes], [e.model_dump() for e in req.current_edges or []], ) else: # Fallback: no existing graph context, regenerate from scratch guidance_note = f"Student feedback on the previous tree: {req.user_feedback}" nodes, edges = await loop.run_in_executor( None, brain.extract_curriculum, chunks, req.intention, guidance_note, ) get_graph_manager().set_graph(req.project_id, nodes) # Persist the refined graph under whatever document set is CURRENTLY uploaded, # not the client-supplied document_id -> which can be stale if a file was # added/removed since the tree was first built, and would otherwise silently # orphan the refined tree under the wrong cache key. current_files = list_project_files(req.project_id) current_id = req.document_id if current_files: try: contents = [] for p in current_files: with open(p, "rb") as f: contents.append(f.read()) current_id = _combined_document_id(contents) except Exception: pass get_graph_manager().save_doc_graph(current_id, nodes, edges) return { "status": "ready", "nodes": [n.model_dump() for n in nodes], "edges": edges, }