Update api.py

api.py

@@ -1,27 +1,25 @@
 """
-api.py — FastAPI REST endpoint for DocMind AI
-Runs alongside the Streamlit app, exposed on port 7861.
-Add this to your HuggingFace Space and update your Dockerfile to run both.
+api.py — FastAPI REST endpoint for DocMind AI (Multimodal + Memory)
+Runs on port 7861 alongside the Streamlit app (port 7860).
 """
 
 from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-import tempfile
+from typing import List
 import os
 
-app = FastAPI(title="DocMind AI API")
+app = FastAPI(title="DocMind AI API", version="2.0")
 
-# Allow requests from your portfolio website
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Replace with "https://ryanfarahani.com" for production
+    allow_origins=["*"],  # Lock down to "https://rayanfarahani.com" in production
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Lazy-load the RAG engine (shared with Streamlit app)
+# Shared RAG engine instance
 _rag_engine = None
 
 def get_rag():
@@ -32,60 +30,117 @@ def get_rag():
     return _rag_engine
 
 
-# ── Models ─────────────────────────────────────────────────────────────────
+# ── Models ───────────────────────────────────────────────────────────────────
+
 class QueryRequest(BaseModel):
     question: str
 
 class QueryResponse(BaseModel):
-    answer: str
-    sources: list[str]
-    success: bool
-    error: str = ""
+    answer:       str
+    sources:      List[str]
+    success:      bool
+    error:        str = ""
+    memory_count: int = 0   # how many past exchanges the model remembers
+
+class UploadResponse(BaseModel):
+    success:   bool
+    filename:  str
+    chunks:    int
+    file_type: str
+    message:   str
+
+class MemoryResponse(BaseModel):
+    exchanges: int
+    messages:  List[dict]
+
 
+# ── Routes ───────────────────────────────────────────────────────────────────
 
-# ── Routes ─────────────────────────────────────────────────────────────────
 @app.get("/health")
 def health():
-    return {"status": "ok", "service": "DocMind AI API"}
+    return {"status": "ok", "service": "DocMind AI API", "version": "2.0"}
 
 
-@app.post("/upload")
+@app.post("/upload", response_model=UploadResponse)
 async def upload_document(file: UploadFile = File(...)):
-    """Upload and ingest a PDF or TXT document."""
-    try:
-        suffix = os.path.splitext(file.filename)[-1].lower() or ".txt"
-        if suffix not in [".pdf", ".txt"]:
-            raise HTTPException(status_code=400, detail="Only PDF and TXT files are supported.")
+    """
+    Upload and ingest a document.
+    Supported: PDF, TXT, DOCX, CSV, XLSX, JPG, PNG, WEBP
+    """
+    filename = file.filename
+    suffix   = os.path.splitext(filename)[-1].lower()
+
+    SUPPORTED = {".pdf", ".txt", ".docx", ".doc", ".csv", ".xlsx", ".xls",
+                 ".jpg", ".jpeg", ".png", ".webp"}
+    if suffix not in SUPPORTED:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Unsupported file type: {suffix}. Supported: {', '.join(sorted(SUPPORTED))}"
+        )
 
-        rag = get_rag()
+    try:
+        rag    = get_rag()
         chunks = rag.ingest_file(file)
-        return {
-            "success": True,
-            "filename": file.filename,
-            "chunks": chunks,
-            "message": f"Successfully indexed {chunks} chunks from {file.filename}"
+
+        type_labels = {
+            ".pdf": "PDF Document",
+            ".txt": "Text File",
+            ".docx": "Word Document", ".doc": "Word Document",
+            ".csv": "CSV Spreadsheet",
+            ".xlsx": "Excel Spreadsheet", ".xls": "Excel Spreadsheet",
+            ".jpg": "Image", ".jpeg": "Image", ".png": "Image", ".webp": "Image",
         }
+
+        return UploadResponse(
+            success=True,
+            filename=filename,
+            chunks=chunks,
+            file_type=type_labels.get(suffix, suffix),
+            message=f"Successfully indexed {chunks} chunks from {filename}"
+        )
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 
 @app.post("/query", response_model=QueryResponse)
 async def query_document(req: QueryRequest):
-    """Ask a question about the currently loaded document."""
+    """Ask a question. The model uses conversation memory for follow-ups."""
+    if not req.question.strip():
+        raise HTTPException(status_code=400, detail="Question cannot be empty.")
     try:
-        if not req.question.strip():
-            raise HTTPException(status_code=400, detail="Question cannot be empty.")
-
-        rag = get_rag()
+        rag            = get_rag()
         answer, sources = rag.query(req.question)
-        return QueryResponse(answer=answer, sources=sources, success=True)
+        return QueryResponse(
+            answer=answer,
+            sources=sources,
+            success=True,
+            memory_count=rag.get_memory_count()
+        )
     except Exception as e:
         return QueryResponse(answer="", sources=[], success=False, error=str(e))
 
 
+@app.get("/memory", response_model=MemoryResponse)
+def get_memory():
+    """Return current conversation history."""
+    rag = get_rag()
+    return MemoryResponse(
+        exchanges=rag.get_memory_count(),
+        messages=rag.get_memory_messages()
+    )
+
+
+@app.post("/memory/clear")
+def clear_memory():
+    """Clear conversation history without removing the document."""
+    rag = get_rag()
+    rag.clear_memory()
+    return {"success": True, "message": "Conversation memory cleared."}
+
+
 @app.post("/reset")
 def reset():
-    """Reset the loaded document."""
+    """Reset everything — document and memory."""
     global _rag_engine
     _rag_engine = None
-    return {"success": True, "message": "Document cleared."}
+    return {"success": True, "message": "Document and memory cleared."}