""" ScratchLM - Zero-DB AI Study Application FastAPI + Pure HTML/JS frontend with dark theme Multi-model support: Gemma 4 (Google AI Studio) + NVIDIA Nemotron 3 """ import os import io import time import base64 from openai import OpenAI from fastapi import FastAPI, UploadFile, File from fastapi.responses import HTMLResponse, JSONResponse from pypdf import PdfReader import fitz # PyMuPDF for PDF → image rendering # ============================================================================ # Configuration # ============================================================================ GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "").strip() NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY", "").strip() # OpenAI SDK clients — initialized lazily to avoid crashing on startup without keys _google_client = None _nvidia_client = None def get_google_client() -> OpenAI: global _google_client if _google_client is None: _google_client = OpenAI( base_url="https://generativelanguage.googleapis.com/v1beta/openai/", api_key=GEMINI_API_KEY or "placeholder", ) return _google_client def get_nvidia_client() -> OpenAI: global _nvidia_client if _nvidia_client is None: _nvidia_client = OpenAI( base_url="https://integrate.api.nvidia.com/v1", api_key=NVIDIA_API_KEY or "placeholder", ) return _nvidia_client # Model IDs GEMMA_MODEL = "gemma-4-31b-it" NEMOTRON_MODEL = "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning" # ============================================================================ # Core Functions # ============================================================================ def generate_content(prompt: str, context: str = "", model: str = "gemma4", document_bytes: bytes = None, document_type: str = "pdf") -> str: """Generate content using Gemma via Google AI Studio or NVIDIA Nemotron via OpenAI SDK.""" if model == "nemotron": return _generate_nemotron(prompt, context, document_bytes, document_type) else: return _generate_gemma(prompt, context) def _generate_gemma(prompt: str, context: str = "") -> str: """Generate content using Gemma 4 via Google AI Studio OpenAI-compatible endpoint.""" if not GEMINI_API_KEY: return "Error: GEMINI_API_KEY not configured." full_text = f"{prompt}\n\n=== SOURCE MATERIAL ===\n{context}" if context else prompt try: completion = get_google_client().chat.completions.create( model=GEMMA_MODEL, messages=[{"role": "user", "content": full_text}], temperature=1.0, top_p=0.95, max_tokens=8192, ) return completion.choices[0].message.content or "" except Exception as e: return f"Error: {e}" def _generate_nemotron_with_retry(prompt: str, document_bytes: bytes = None, document_type: str = "pdf", fallback_context: str = "", retries: int = 2, delay: float = 2.0) -> tuple[str, bool]: """ Generate content using NVIDIA Nemotron 3 with retry logic. Returns (result, fell_back_to_gemma). """ if not NVIDIA_API_KEY: return "Error: NVIDIA_API_KEY not configured.", True # Build message content if document_bytes: if document_type == "pdf": # Render PDF pages as PNG images for multimodal understanding page_images = pdf_to_images(document_bytes, dpi=150) if page_images: content = [{"type": "text", "text": prompt}] for img_bytes in page_images: b64_data = base64.b64encode(img_bytes).decode("utf-8") content.append({ "type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64_data}"}, }) else: # Fallback: try pypdf text extraction try: reader = PdfReader(io.BytesIO(document_bytes)) text = "\n".join(p.extract_text() or "" for p in reader.pages) if text.strip(): content = [{"type": "text", "text": f"{prompt}\n\n=== DOCUMENT CONTENT ===\n{text}"}] else: content = [{"type": "text", "text": prompt}] except Exception: content = [{"type": "text", "text": prompt}] else: # Non-PDF files: send as base64 image b64_data = base64.b64encode(document_bytes).decode("utf-8") content = [ {"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": f"data:image/{document_type};base64,{b64_data}"}}, ] else: content = [{"type": "text", "text": prompt}] attempt = 0 last_error = "" while attempt <= retries: try: completion = get_nvidia_client().chat.completions.create( model=NEMOTRON_MODEL, messages=[{"role": "user", "content": content}], temperature=0.6, top_p=0.95, max_tokens=65536, extra_body={ "chat_template_kwargs": {"enable_thinking": True}, "reasoning_budget": 16384, }, stream=False, ) reasoning = getattr(completion.choices[0].message, "reasoning_content", None) response = completion.choices[0].message.content or "" # Prepend reasoning if present if reasoning: response = f"[Thinking]\n{reasoning}\n\n[Response]\n{response}" return response, False, fallback_context except Exception as e: last_error = str(e) attempt += 1 if attempt <= retries: time.sleep(delay) # All retries failed — fallback to Gemma return f"[NVIDIA Nemotron failed after {retries + 1} attempts: {last_error}] ", True, fallback_context def pdf_to_images(document_bytes: bytes, dpi: int = 150) -> list[bytes]: """Render PDF pages as PNG images using PyMuPDF. Returns list of PNG byte arrays.""" images = [] try: doc = fitz.open(stream=document_bytes, filetype="pdf") for page_num in range(len(doc)): page = doc[page_num] # Render at specified DPI mat = fitz.Matrix(dpi / 72, dpi / 72) pix = page.get_pixmap(matrix=mat, alpha=False) img_bytes = pix.tobytes("png") images.append(img_bytes) doc.close() except Exception as e: pass # Return empty list on failure return images def _generate_nemotron(prompt: str, context: str = "", document_bytes: bytes = None, document_type: str = "pdf") -> str: """Wrapper for nemotron generation with retry and fallback.""" result, fell_back, _ = _generate_nemotron_with_retry(prompt, document_bytes, document_type, fallback_context=context) if fell_back: # Retry with Gemma 4 — pass context so it's preserved return _generate_gemma(prompt, context) return result def parse_pdf(file_bytes: bytes) -> str: """Extract text from PDF using pypdf.""" try: reader = PdfReader(io.BytesIO(file_bytes)) return "\n".join(p.extract_text() or "" for p in reader.pages) except Exception as e: return f"[PDF error: {e}]" # ============================================================================ # FastAPI App # ============================================================================ app = FastAPI(title="ScratchLM") @app.get("/", response_class=HTMLResponse) async def root(): return """ ScratchLM

🎓 ScratchLM

Zero-DB AI Study Assistant

Generated cheatsheet will appear here.
Hello! Upload documents and extract text to start chatting.

Model: Gemma 4 31B Google AI Studio , Nemotron 3 Nano Omni 30B NVIDIA NIM

""" @app.post("/extract") async def extract_files(files: list[UploadFile] = File(...)): """Extract text from uploaded PDFs.""" if not files: return JSONResponse({"error": "No files", "preview": "", "full_text": ""}) consolidated, full_text = [], "" for i, file in enumerate(files): file_bytes = await file.read() text = parse_pdf(file_bytes) if file.filename.lower().endswith('.pdf') else f"[Unsupported: {file.filename}]" header = f"\n=== DOCUMENT {i+1}: {file.filename} ===\n" consolidated.append(header + text) full_text += header + text + "\n" preview = "\n".join(consolidated) return {"preview": preview[:3000] + ("..." if len(preview) > 3000 else ""), "full_text": preview} @app.post("/generate-cheatsheet") async def generate_cheatsheet(request: dict): """Generate cheatsheet from context or raw document.""" model = request.get("model", "gemma4") context = request.get("context", "") document = request.get("document", None) # bytes as list of ints document_type = request.get("document_type", "pdf") if model == "nemotron" and document: document_bytes = bytes(document) elif not context: return JSONResponse({"error": "No context", "result": ""}) else: document_bytes = None # For Nemotron with document bytes, add text intro so model knows what the document is if model == "nemotron" and document_bytes: prompt_intro = "Create a 4-zone markdown cheatsheet based on the provided document:\n\n" else: prompt_intro = "" prompt = prompt_intro + """Create a 4-zone markdown cheatsheet: ## Zone A: Concept Deep-Dive Key concepts, definitions, theorems ## Zone B: Real-World Analogies Intuitive explanations ## Zone C: Worked Problems Step-by-step solutions with LaTeX math ## Zone D: Exam Traps & Strategy Common mistakes and how to avoid them""" result = generate_content(prompt, context, model, document_bytes, document_type) return {"result": result} @app.post("/chat") async def chat(request: dict): """RAG chat with context or raw document.""" message = request.get("message", "") model = request.get("model", "gemma4") context = request.get("context", "") history = request.get("history", []) document = request.get("document", None) # bytes as list of ints document_type = request.get("document_type", "pdf") if not message: return JSONResponse({"error": "No message", "response": ""}) if model == "nemotron" and document: document_bytes = bytes(document) elif model != "nemotron" and not context: return JSONResponse({"error": "Extract documents first", "response": ""}) else: document_bytes = None history_str = "\n".join(f"{'User' if m.get('role') == 'user' else 'Assistant'}: {m.get('content', '')}" for m in history) prompt = f"Tutor mode. Answer step-by-step.\n\nContext:\n{context}\n\nHistory:\n{history_str}\n\nUser: {message}" result = generate_content(prompt, context, model, document_bytes, document_type) return {"response": result} if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)