import os import io import re import json import tempfile from typing import Optional import uvicorn from fastapi import FastAPI, File, UploadFile, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from pydantic import BaseModel import google.generativeai as genai # ── Document parsers ──────────────────────────────────────────────────────── import PyPDF2 import docx # ── App setup ─────────────────────────────────────────────────────────────── app = FastAPI( title="DocAI – Document Analyzer", description="Upload PDF, DOCX, or TXT files and get AI-generated title, summary, and keywords via Gemini 2.5 Flash.", version="1.0.0", ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ── Gemini client ──────────────────────────────────────────────────────────── GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "") if GEMINI_API_KEY: genai.configure(api_key=GEMINI_API_KEY) GEMINI_MODEL = "gemini-2.5-flash" # ── Pydantic models ────────────────────────────────────────────────────────── class AnalysisResult(BaseModel): title: str summary: str keywords: list[str] word_count: int char_count: int file_name: str file_type: str # ── Helpers ────────────────────────────────────────────────────────────────── ALLOWED_TYPES = { "application/pdf": "pdf", "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx", "application/msword": "doc", "text/plain": "txt", } MAX_FILE_SIZE_MB = 10 MAX_FILE_SIZE_BYTES = MAX_FILE_SIZE_MB * 1024 * 1024 def extract_text_from_pdf(file_bytes: bytes) -> str: reader = PyPDF2.PdfReader(io.BytesIO(file_bytes)) pages = [] for page in reader.pages: text = page.extract_text() if text: pages.append(text.strip()) return "\n\n".join(pages) def extract_text_from_docx(file_bytes: bytes) -> str: doc = docx.Document(io.BytesIO(file_bytes)) paragraphs = [p.text for p in doc.paragraphs if p.text.strip()] return "\n\n".join(paragraphs) def extract_text_from_txt(file_bytes: bytes) -> str: for enc in ("utf-8", "latin-1", "cp1252"): try: return file_bytes.decode(enc) except UnicodeDecodeError: continue raise ValueError("Cannot decode text file with supported encodings.") def extract_text(file_bytes: bytes, file_type: str) -> str: if file_type == "pdf": return extract_text_from_pdf(file_bytes) elif file_type in ("docx", "doc"): return extract_text_from_docx(file_bytes) elif file_type == "txt": return extract_text_from_txt(file_bytes) raise ValueError(f"Unsupported file type: {file_type}") def truncate_text(text: str, max_chars: int = 30_000) -> str: """Gemini has a large context window; we still cap to keep costs low.""" if len(text) <= max_chars: return text return text[:max_chars] + "\n\n[... document truncated for analysis ...]" def build_prompt(text: str) -> str: return f"""You are an expert document analyst. Analyze the document text below and respond ONLY with valid JSON — no markdown, no code fences, no extra text. Return exactly this structure: {{ "title": "A professional, descriptive title for this document (max 15 words)", "summary": "A concise, coherent summary preserving the key information (150-250 words)", "keywords": ["keyword1", "keyword2", "keyword3", "keyword4", "keyword5", "keyword6", "keyword7", "keyword8", "keyword9", "keyword10"] }} Rules: - title: Meaningful, professional, specific to the content. Do NOT use generic titles. - summary: Preserve facts, figures, names, and key arguments. Write in third-person prose. - keywords: Exactly 10 keywords. Order by importance (most important first). Single words or short phrases only. DOCUMENT TEXT: \"\"\" {text} \"\"\" """ def call_gemini(prompt: str) -> dict: if not GEMINI_API_KEY: raise HTTPException( status_code=500, detail="GEMINI_API_KEY environment variable is not set.", ) model = genai.GenerativeModel( model_name=GEMINI_MODEL, generation_config=genai.GenerationConfig( temperature=0.3, max_output_tokens=8192, ), ) response = model.generate_content(prompt) raw = response.text.strip() # Strip markdown fences if model adds them anyway raw = re.sub(r"^```(?:json)?\s*", "", raw, flags=re.IGNORECASE | re.MULTILINE) raw = re.sub(r"\s*```\s*$", "", raw, flags=re.MULTILINE) raw = raw.strip() # Extract the first JSON object if there is surrounding text match = re.search(r"\{.*\}", raw, re.DOTALL) if match: raw = match.group(0) return json.loads(raw) # ── Routes ─────────────────────────────────────────────────────────────────── @app.get("/", tags=["Health"]) def root(): return { "service": "DocAI – Document Analyzer", "model": GEMINI_MODEL, "status": "running", "endpoints": { "analyze": "POST /analyze", "health": "GET /health", "docs": "GET /docs", }, } @app.get("/health", tags=["Health"]) def health(): return {"status": "ok", "model": GEMINI_MODEL} @app.post("/analyze", response_model=AnalysisResult, tags=["Analysis"]) async def analyze_document(file: UploadFile = File(...)): """ Upload a PDF, DOCX, or TXT file. Returns an AI-generated title, summary, and 10 keywords. """ # ── Validate content type ──────────────────────────────────────────────── content_type = file.content_type or "" file_type = ALLOWED_TYPES.get(content_type) # Fallback: guess from extension if not file_type and file.filename: ext = file.filename.rsplit(".", 1)[-1].lower() if ext in ("pdf", "docx", "doc", "txt"): file_type = ext if not file_type: raise HTTPException( status_code=415, detail=f"Unsupported file type '{content_type}'. Allowed: PDF, DOCX, TXT.", ) # ── Read & size-check ──────────────────────────────────────────────────── file_bytes = await file.read() if len(file_bytes) > MAX_FILE_SIZE_BYTES: raise HTTPException( status_code=413, detail=f"File too large. Maximum size is {MAX_FILE_SIZE_MB} MB.", ) if len(file_bytes) == 0: raise HTTPException(status_code=400, detail="Uploaded file is empty.") # ── Extract text ───────────────────────────────────────────────────────── try: raw_text = extract_text(file_bytes, file_type) except Exception as exc: raise HTTPException( status_code=422, detail=f"Could not extract text from document: {exc}", ) raw_text = raw_text.strip() if not raw_text: raise HTTPException( status_code=422, detail="No readable text found in the document.", ) # ── Analyze with Gemini ────────────────────────────────────────────────── truncated = truncate_text(raw_text) prompt = build_prompt(truncated) try: result = call_gemini(prompt) except json.JSONDecodeError as exc: raise HTTPException( status_code=502, detail=f"Gemini returned non-JSON response: {exc}", ) except Exception as exc: raise HTTPException( status_code=502, detail=f"Gemini API error: {exc}", ) # ── Validate and normalize result ──────────────────────────────────────── title = str(result.get("title", "Untitled Document")).strip() summary = str(result.get("summary", "")).strip() keywords = result.get("keywords", []) if not isinstance(keywords, list): keywords = [] keywords = [str(k).strip() for k in keywords if k][:10] return AnalysisResult( title=title, summary=summary, keywords=keywords, word_count=len(raw_text.split()), char_count=len(raw_text), file_name=file.filename or "unknown", file_type=file_type, ) # ── Entry point (local dev) ────────────────────────────────────────────────── if __name__ == "__main__": uvicorn.run("main:app", host="0.0.0.0", port=7860, reload=True)