Update main.py

main.py

@@ -1,385 +1,386 @@
-# backend/main.py
-from urllib.parse import unquote
-from typing import List, Optional
-import uuid
-import auth
-# --- Standard Imports ---
-from fastapi import FastAPI, Depends, HTTPException, UploadFile, File, Form
-from pydantic import BaseModel
-from database import engine, get_db, SessionLocal
-from datetime import datetime
-
-from sqlalchemy.orm import Session
-from sqlalchemy.exc import OperationalError
-from sqlalchemy import text
-import os
-import shutil # Important for file operations
-
-from pipeline import highlight_text
-from ml_qna import qna as generate_ml_answer
-
-# from email_automation import download_attached_file
-# import imaplib
-from contextlib import asynccontextmanager
-from pipeline import pipeline_process_pdf, load_all_models
-from fastapi import BackgroundTasks
-
-# --- Middleware Import ---
-from fastapi.middleware.cors import CORSMiddleware
-
-# --- Local Module Imports ---
-import crud
-import models
-import schemas
-from database import engine, get_db
-from supabase_utils import upload_file_to_supabase
-
-# This creates/updates the database tables in your Neon database
-# based on your models.py file.
-models.Base.metadata.create_all(bind=engine)
-
-# --- (3) SETUP FOR LOADING MODELS ON STARTUP ---
-# This dictionary will hold our loaded models so we don't reload them on every request
-ml_models = {}
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    # This code runs ONCE when the server starts up
-    print("[INFO] Server starting up...")
-
-    # --- ADD THIS ENTIRE BLOCK ---
-    print("[INFO] Ensuring system 'automation_user' exists...")
-    db = SessionLocal()
-    try:
-        # Check if the user already exists
-        automation_user = crud.get_user(db, user_id="automation_user")
-        if not automation_user:
-            # If not, create it
-            print("[INFO] 'automation_user' not found. Creating it now...")
-            user_data = schemas.UserCreate(
-                id="automation_user",
-                name="Automation Service",
-                department="System",
-                role="system",
-                password="automation_pass" # A placeholder password
-            )
-            crud.create_user(db, user_data)
-            print("[INFO] 'automation_user' created successfully.")
-        else:
-            print("[INFO] 'automation_user' already exists.")
-    finally:
-        db.close() # Always close the database session
-    # --- END OF BLOCK ---
-
-    print("[INFO] Loading ML models...")
-    tokenizer, model, nlp_model = load_all_models()
-    ml_models["tokenizer"] = tokenizer
-    ml_models["model"] = model
-    ml_models["nlp_model"] = nlp_model
-    print("[INFO] ML models loaded successfully and are ready.")
-
-    yield
-
-    ml_models.clear()
-    print("[INFO] Server shutting down.")
-
-
-# @asynccontextmanager
-# async def lifespan(app: FastAPI):
-#     # This code runs ONCE when the server starts up
-#     print("[INFO] Server starting up. Loading ML models...")
-#     tokenizer, model, nlp_model = load_all_models()
-#     ml_models["tokenizer"] = tokenizer
-#     ml_models["model"] = model
-#     ml_models["nlp_model"] = nlp_model
-#     print("[INFO] ML models loaded successfully and are ready.")
-#     yield
-#     # This code runs when the server shuts down
-#     ml_models.clear()
-#     print("[INFO] Server shutting down.")
-
-
-app = FastAPI(lifespan=lifespan)
-
-# This list now includes the new port your frontend is using
-origins = [
-    "http://localhost:3000",
-    "http://127.0.0.1:3000",
-    "http://localhost:3003", # <-- ADD THIS LINE
-    "http://127.0.0.1:3003", # <-- And this one for good measure
-]
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=origins, # Use the updated list
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# --- LOCAL UPLOAD DIRECTORY for temporary storage ---
-UPLOAD_DIRECTORY = "uploads"
-os.makedirs(UPLOAD_DIRECTORY, exist_ok=True)
-
-# --- Diagnostic Endpoints ---
-@app.get("/")
-def read_root():
-    return {"status": "ok", "service": "kmrl-backend-service"}
-
-@app.get("/ping-db")
-def ping_db(db: Session = Depends(get_db)):
-    try:
-        db.execute(text("SELECT 1"))
-        return {"status": "ok", "message": "Database connection successful."}
-    except OperationalError as e:
-        raise HTTPException(
-            status_code=500,
-            detail=f"Database connection failed: {str(e)}"
-        )
-
-# --- User Management Endpoints ---
-@app.post("/users/", response_model=schemas.User)
-def create_user(user: schemas.UserCreate, db: Session = Depends(get_db)):
-    db_user = crud.get_user(db, user_id=user.id)
-    if db_user:
-        raise HTTPException(status_code=400, detail="User ID already registered")
-    return crud.create_user(db=db, user=user)
-
-@app.get("/users/{user_id}", response_model=schemas.User)
-def read_user(user_id: str, db: Session = Depends(get_db)):
-    db_user = crud.get_user(db, user_id=user_id)
-    if db_user is None:
-        raise HTTPException(status_code=404, detail="User not found")
-    return db_user
-
-# --- Document Management Endpoints ---
-
-@app.post("/documents/upload")
-def upload_document(
-        # Optional fields for email automation, but required for frontend
-        title: Optional[str] = Form(None),
-        department: Optional[str] = Form(None),
-        user_id: Optional[str] = Form(None),
-        # The file is always required
-        file: UploadFile = File(...),
-        db: Session = Depends(get_db)
-):
-    # --- 1. Set Default Values & Validate User ---
-    # If a title wasn't provided (from email), create a default one.
-    final_title = title or f"Email Attachment - {file.filename}"
-
-    # If a user_id wasn't provided, it MUST be the automation user.
-    final_user_id = user_id or "automation_user"
-
-    # If a department wasn't provided, set it to be auto-detected by the pipeline.
-    final_department = department or "auto-detected"
-
-    # Now, use these final variables to validate the user
-    user = crud.get_user(db, user_id=final_user_id)
-    if not user:
-        raise HTTPException(status_code=404, detail=f"Uploader '{final_user_id}' not found")
-
-    # --- 2. Upload Original File to Cloud ---
-    print("Uploading original file to cloud storage...")
-    public_url = upload_file_to_supabase(file.file, file.filename)
-    if not public_url:
-        raise HTTPException(status_code=500, detail="Could not upload file to cloud storage.")
-    print("File uploaded successfully. Public URL:", public_url)
-
-    file.file.seek(0) # Rewind file for local processing
-
-    # --- 3. Create Initial Database Record ---
-    # This now correctly matches the function in crud.py (which should not take highlighted_file_path)
-    document_data = schemas.DocumentCreate(title=final_title, department=final_department)
-    db_document = crud.create_document(db=db, document=document_data, file_path=public_url, user_id=final_user_id)
-    print(f"Initial document record created in DB with ID: {db_document.id}")
-
-    # --- 4. Save Local Copy & Run ML Pipeline ---
-    local_file_path = os.path.join(UPLOAD_DIRECTORY, file.filename)
-    with open(local_file_path, "wb") as buffer:
-        shutil.copyfileobj(file.file, buffer)
-
-    print("Starting ML pipeline processing...")
-    ml_results = pipeline_process_pdf(
-        pdf_path=local_file_path,
-        clf_tokenizer=ml_models["tokenizer"],
-        clf_model=ml_models["model"],
-        nlp_model=ml_models["nlp_model"]
-    )
-    print("ML pipeline processing complete.")
-
-    # --- 5. Upload Highlighted PDF (if created) ---
-    highlighted_pdf_path = ml_results.get("highlighted_pdf")
-    highlighted_public_url = None
-    if highlighted_pdf_path and os.path.exists(highlighted_pdf_path):
-        print("Uploading highlighted file to cloud storage...")
-        with open(highlighted_pdf_path, "rb") as f:
-            highlighted_filename = os.path.basename(highlighted_pdf_path)
-            highlighted_public_url = upload_file_to_supabase(f, highlighted_filename)
-        print("Highlighted PDF uploaded successfully.")
-
-    # ... (after the ML pipeline runs) ...
-
-    # --- (6) UPDATE THE DATABASE RECORD WITH ML RESULTS ---
-    print("Updating database record with ML results...")
-    final_document = crud.update_document_with_ml_results(
-        db,
-        document_id=db_document.id,
-        ml_results=ml_results,
-        highlighted_file_path=highlighted_public_url
-    )
-    print("Database record updated successfully.")
-
-    # --- (7) CREATE NOTIFICATION FOR THE DEPARTMENT ---
-    # The ML results contain the department the document was routed to.
-    routed_department = final_document.department
-    if routed_department and routed_department != "Unknown":
-        notification_message = f"New document '{final_document.title}' has been assigned to your department."
-        crud.create_notification(
-            db=db,
-            document_id=final_document.id,
-            department=routed_department,
-            message=notification_message
-        )
-        print(f"Notification created for department: {routed_department}")
-
-
-    # --- 8. Cleanup Local Files ---
-    try:
-        if os.path.exists(local_file_path):
-            os.remove(local_file_path)
-        if highlighted_pdf_path and os.path.exists(highlighted_pdf_path):
-            os.remove(highlighted_pdf_path)
-    except OSError as e:
-        print(f"Error during file cleanup: {e}")
-
-    # --- 9. Return Final Response ---
-    return {
-        "message": "Document processed and all data saved successfully.",
-        "document_info": schemas.Document.model_validate(final_document),
-        "highlighted_pdf_url": highlighted_public_url
-    }
-
-# --- Read Endpoints ---
-@app.get("/documents/", response_model=list[schemas.Document])
-def read_all_documents(skip: int = 0, limit: int = 100, db: Session = Depends(get_db)):
-    documents = crud.get_all_documents(db, skip=skip, limit=limit)
-    return documents
-
-@app.get("/documents/{department}", response_model=list[schemas.Document])
-def read_documents_for_department(department: str, skip: int = 0, limit: int = 100, db: Session = Depends(get_db)):
-    documents = crud.get_documents_by_department(db, department=department, skip=skip, limit=limit)
-    return documents
-
-# --- ADD THESE NEW ENDPOINTS FOR Q&A ---
-
-def run_ml_qna_in_background(question_id: uuid.UUID, pinecone_pdf_id: str, question_text: str):
-    print(f"[BACKGROUND TASK] Starting ML RAG pipeline for question ID: {question_id}")
-
-    # Call the ML function to get the answer
-    answer_text = generate_ml_answer(
-        pdf_id=pinecone_pdf_id, # <--- This now uses the correct filename ID
-        query=question_text
-    )
-    print(f"[BACKGROUND TASK] Answer generated: {answer_text[:100]}...")
-
-    # Use the CRUD function to save the answer to the database
-    db = SessionLocal()
-    try:
-        # Use the new 'db' session to update the database
-        crud.update_question_with_answer(
-            db=db,
-            question_id=question_id,
-            answer_text=answer_text
-        )
-        print(f"[BACKGROUND TASK] Answer saved to database for question ID: {question_id}")
-    finally:
-        db.close()
-
-@app.post("/documents/{document_id}/questions", response_model=schemas.Question)
-def ask_question_on_document(
-        document_id: uuid.UUID,
-        question: schemas.QuestionCreate,
-        background_tasks: BackgroundTasks,
-        db: Session = Depends(get_db)
-):
-    """
-    Endpoint for the frontend to submit a new question.
-    It saves the question, calls the ML RAG pipeline to generate a real answer,
-    and saves the answer to the database.
-    """
-    # First, fetch the document to get its uploader's ID
-    document = crud.get_document_by_id(db, document_id=document_id)
-    if not document:
-        raise HTTPException(status_code=404, detail="Document not found")
-
-    user_id_who_asked = document.uploader_id
-
-    # Create the question in the database with a NULL answer first
-    db_question = crud.create_question(
-        db=db,
-        document_id=document_id,
-        user_id=user_id_who_asked,
-        question=question
-    )
-    print(f"New question saved with ID: {db_question.id}. Triggering background ML task.")
-
-
-    pinecone_pdf_id = os.path.splitext(os.path.basename(unquote(document.file_path)))[0]
-
-    background_tasks.add_task(
-        run_ml_qna_in_background,
-        db_question.id,
-        pinecone_pdf_id, # <--- Pass the correct filename ID
-        question.question_text,
-    )
-    # --- END OF KEY CHANGE ---
-
-    # Return the new question object to the frontend immediately.
-    # The frontend will see that `answer_text` is still null.
-    return db_question
-
-
-@app.get("/documents/{document_id}/questions", response_model=List[schemas.Question])
-def get_document_questions(
-        document_id: uuid.UUID,
-        db: Session = Depends(get_db) # Ensure there are no typos like 'get_d b'
-):
-    """
-    Endpoint for the frontend to retrieve the full conversation history
-    (all questions and their answers) for a document.
-    """
-    return crud.get_questions_for_document(db=db, document_id=document_id)
-
-
-# --- ADD THIS NEW ENDPOINT FOR EMAIL AUTOMATION ---
-
-
-@app.patch("/questions/{question_id}/answer")
-def submit_answer(
-        question_id: uuid.UUID,
-        answer: schemas.Answer,
-        db: Session = Depends(get_db)
-):
-    """
-    INTERNAL ENDPOINT for the ML service to submit its generated answer
-    for a question that has already been created.
-    """
-    updated_question = crud.update_question_with_answer(
-        db=db,
-        question_id=question_id,
-        answer_text=answer.answer_text
-    )
-    if not updated_question:
-        raise HTTPException(status_code=404, detail="Question not found")
-
-    print(f"Answer submitted for question ID: {question_id}")
-    return {"status": "success", "question": updated_question}
-
-
-# --- NEW ENDPOINT FOR NOTIFICATIONS ---
-@app.get("/notifications/{department}", response_model=List[schemas.Notification])
-def read_notifications(department: str, db: Session = Depends(get_db)):
-    """Fetches unread notifications for a given department."""
-    notifications = crud.get_notifications_for_department(db, department=department)
-    return notifications
+# backend/main.py
+from urllib.parse import unquote
+from typing import List, Optional
+import uuid
+import auth
+# --- Standard Imports ---
+from fastapi import FastAPI, Depends, HTTPException, UploadFile, File, Form
+from pydantic import BaseModel
+from database import engine, get_db, SessionLocal
+from datetime import datetime
+
+from sqlalchemy.orm import Session
+from sqlalchemy.exc import OperationalError
+from sqlalchemy import text
+import os
+import shutil # Important for file operations
+
+from pipeline import highlight_text
+from ml_qna import qna as generate_ml_answer
+
+# from email_automation import download_attached_file
+# import imaplib
+from contextlib import asynccontextmanager
+from pipeline import pipeline_process_pdf, load_all_models
+from fastapi import BackgroundTasks
+
+# --- Middleware Import ---
+from fastapi.middleware.cors import CORSMiddleware
+
+# --- Local Module Imports ---
+import crud
+import models
+import schemas
+from database import engine, get_db
+from supabase_utils import upload_file_to_supabase
+
+# This creates/updates the database tables in your Neon database
+# based on your models.py file.
+models.Base.metadata.create_all(bind=engine)
+
+# --- (3) SETUP FOR LOADING MODELS ON STARTUP ---
+# This dictionary will hold our loaded models so we don't reload them on every request
+ml_models = {}
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # This code runs ONCE when the server starts up
+    print("[INFO] Server starting up...")
+
+    # --- ADD THIS ENTIRE BLOCK ---
+    print("[INFO] Ensuring system 'automation_user' exists...")
+    db = SessionLocal()
+    try:
+        # Check if the user already exists
+        automation_user = crud.get_user(db, user_id="automation_user")
+        if not automation_user:
+            # If not, create it
+            print("[INFO] 'automation_user' not found. Creating it now...")
+            user_data = schemas.UserCreate(
+                id="automation_user",
+                name="Automation Service",
+                department="System",
+                role="system",
+                password="automation_pass" # A placeholder password
+            )
+            crud.create_user(db, user_data)
+            print("[INFO] 'automation_user' created successfully.")
+        else:
+            print("[INFO] 'automation_user' already exists.")
+    finally:
+        db.close() # Always close the database session
+    # --- END OF BLOCK ---
+
+    print("[INFO] Loading ML models...")
+    tokenizer, model, nlp_model = load_all_models()
+    ml_models["tokenizer"] = tokenizer
+    ml_models["model"] = model
+    ml_models["nlp_model"] = nlp_model
+    print("[INFO] ML models loaded successfully and are ready.")
+
+    yield
+
+    ml_models.clear()
+    print("[INFO] Server shutting down.")
+
+
+# @asynccontextmanager
+# async def lifespan(app: FastAPI):
+#     # This code runs ONCE when the server starts up
+#     print("[INFO] Server starting up. Loading ML models...")
+#     tokenizer, model, nlp_model = load_all_models()
+#     ml_models["tokenizer"] = tokenizer
+#     ml_models["model"] = model
+#     ml_models["nlp_model"] = nlp_model
+#     print("[INFO] ML models loaded successfully and are ready.")
+#     yield
+#     # This code runs when the server shuts down
+#     ml_models.clear()
+#     print("[INFO] Server shutting down.")
+
+
+app = FastAPI(lifespan=lifespan)
+
+# This list now includes the new port your frontend is using
+origins = [
+    "http://localhost:3000",
+    "http://127.0.0.1:3000",
+    "https://kochi-metro-document.vercel.app",
+    "http://localhost:3003", # <-- ADD THIS LINE
+    "http://127.0.0.1:3003", # <-- And this one for good measure
+]
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins, # Use the updated list
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# --- LOCAL UPLOAD DIRECTORY for temporary storage ---
+UPLOAD_DIRECTORY = "uploads"
+os.makedirs(UPLOAD_DIRECTORY, exist_ok=True)
+
+# --- Diagnostic Endpoints ---
+@app.get("/")
+def read_root():
+    return {"status": "ok", "service": "kmrl-backend-service"}
+
+@app.get("/ping-db")
+def ping_db(db: Session = Depends(get_db)):
+    try:
+        db.execute(text("SELECT 1"))
+        return {"status": "ok", "message": "Database connection successful."}
+    except OperationalError as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Database connection failed: {str(e)}"
+        )
+
+# --- User Management Endpoints ---
+@app.post("/users/", response_model=schemas.User)
+def create_user(user: schemas.UserCreate, db: Session = Depends(get_db)):
+    db_user = crud.get_user(db, user_id=user.id)
+    if db_user:
+        raise HTTPException(status_code=400, detail="User ID already registered")
+    return crud.create_user(db=db, user=user)
+
+@app.get("/users/{user_id}", response_model=schemas.User)
+def read_user(user_id: str, db: Session = Depends(get_db)):
+    db_user = crud.get_user(db, user_id=user_id)
+    if db_user is None:
+        raise HTTPException(status_code=404, detail="User not found")
+    return db_user
+
+# --- Document Management Endpoints ---
+
+@app.post("/documents/upload")
+def upload_document(
+        # Optional fields for email automation, but required for frontend
+        title: Optional[str] = Form(None),
+        department: Optional[str] = Form(None),
+        user_id: Optional[str] = Form(None),
+        # The file is always required
+        file: UploadFile = File(...),
+        db: Session = Depends(get_db)
+):
+    # --- 1. Set Default Values & Validate User ---
+    # If a title wasn't provided (from email), create a default one.
+    final_title = title or f"Email Attachment - {file.filename}"
+
+    # If a user_id wasn't provided, it MUST be the automation user.
+    final_user_id = user_id or "automation_user"
+
+    # If a department wasn't provided, set it to be auto-detected by the pipeline.
+    final_department = department or "auto-detected"
+
+    # Now, use these final variables to validate the user
+    user = crud.get_user(db, user_id=final_user_id)
+    if not user:
+        raise HTTPException(status_code=404, detail=f"Uploader '{final_user_id}' not found")
+
+    # --- 2. Upload Original File to Cloud ---
+    print("Uploading original file to cloud storage...")
+    public_url = upload_file_to_supabase(file.file, file.filename)
+    if not public_url:
+        raise HTTPException(status_code=500, detail="Could not upload file to cloud storage.")
+    print("File uploaded successfully. Public URL:", public_url)
+
+    file.file.seek(0) # Rewind file for local processing
+
+    # --- 3. Create Initial Database Record ---
+    # This now correctly matches the function in crud.py (which should not take highlighted_file_path)
+    document_data = schemas.DocumentCreate(title=final_title, department=final_department)
+    db_document = crud.create_document(db=db, document=document_data, file_path=public_url, user_id=final_user_id)
+    print(f"Initial document record created in DB with ID: {db_document.id}")
+
+    # --- 4. Save Local Copy & Run ML Pipeline ---
+    local_file_path = os.path.join(UPLOAD_DIRECTORY, file.filename)
+    with open(local_file_path, "wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
+
+    print("Starting ML pipeline processing...")
+    ml_results = pipeline_process_pdf(
+        pdf_path=local_file_path,
+        clf_tokenizer=ml_models["tokenizer"],
+        clf_model=ml_models["model"],
+        nlp_model=ml_models["nlp_model"]
+    )
+    print("ML pipeline processing complete.")
+
+    # --- 5. Upload Highlighted PDF (if created) ---
+    highlighted_pdf_path = ml_results.get("highlighted_pdf")
+    highlighted_public_url = None
+    if highlighted_pdf_path and os.path.exists(highlighted_pdf_path):
+        print("Uploading highlighted file to cloud storage...")
+        with open(highlighted_pdf_path, "rb") as f:
+            highlighted_filename = os.path.basename(highlighted_pdf_path)
+            highlighted_public_url = upload_file_to_supabase(f, highlighted_filename)
+        print("Highlighted PDF uploaded successfully.")
+
+    # ... (after the ML pipeline runs) ...
+
+    # --- (6) UPDATE THE DATABASE RECORD WITH ML RESULTS ---
+    print("Updating database record with ML results...")
+    final_document = crud.update_document_with_ml_results(
+        db,
+        document_id=db_document.id,
+        ml_results=ml_results,
+        highlighted_file_path=highlighted_public_url
+    )
+    print("Database record updated successfully.")
+
+    # --- (7) CREATE NOTIFICATION FOR THE DEPARTMENT ---
+    # The ML results contain the department the document was routed to.
+    routed_department = final_document.department
+    if routed_department and routed_department != "Unknown":
+        notification_message = f"New document '{final_document.title}' has been assigned to your department."
+        crud.create_notification(
+            db=db,
+            document_id=final_document.id,
+            department=routed_department,
+            message=notification_message
+        )
+        print(f"Notification created for department: {routed_department}")
+
+
+    # --- 8. Cleanup Local Files ---
+    try:
+        if os.path.exists(local_file_path):
+            os.remove(local_file_path)
+        if highlighted_pdf_path and os.path.exists(highlighted_pdf_path):
+            os.remove(highlighted_pdf_path)
+    except OSError as e:
+        print(f"Error during file cleanup: {e}")
+
+    # --- 9. Return Final Response ---
+    return {
+        "message": "Document processed and all data saved successfully.",
+        "document_info": schemas.Document.model_validate(final_document),
+        "highlighted_pdf_url": highlighted_public_url
+    }
+
+# --- Read Endpoints ---
+@app.get("/documents/", response_model=list[schemas.Document])
+def read_all_documents(skip: int = 0, limit: int = 100, db: Session = Depends(get_db)):
+    documents = crud.get_all_documents(db, skip=skip, limit=limit)
+    return documents
+
+@app.get("/documents/{department}", response_model=list[schemas.Document])
+def read_documents_for_department(department: str, skip: int = 0, limit: int = 100, db: Session = Depends(get_db)):
+    documents = crud.get_documents_by_department(db, department=department, skip=skip, limit=limit)
+    return documents
+
+# --- ADD THESE NEW ENDPOINTS FOR Q&A ---
+
+def run_ml_qna_in_background(question_id: uuid.UUID, pinecone_pdf_id: str, question_text: str):
+    print(f"[BACKGROUND TASK] Starting ML RAG pipeline for question ID: {question_id}")
+
+    # Call the ML function to get the answer
+    answer_text = generate_ml_answer(
+        pdf_id=pinecone_pdf_id, # <--- This now uses the correct filename ID
+        query=question_text
+    )
+    print(f"[BACKGROUND TASK] Answer generated: {answer_text[:100]}...")
+
+    # Use the CRUD function to save the answer to the database
+    db = SessionLocal()
+    try:
+        # Use the new 'db' session to update the database
+        crud.update_question_with_answer(
+            db=db,
+            question_id=question_id,
+            answer_text=answer_text
+        )
+        print(f"[BACKGROUND TASK] Answer saved to database for question ID: {question_id}")
+    finally:
+        db.close()
+
+@app.post("/documents/{document_id}/questions", response_model=schemas.Question)
+def ask_question_on_document(
+        document_id: uuid.UUID,
+        question: schemas.QuestionCreate,
+        background_tasks: BackgroundTasks,
+        db: Session = Depends(get_db)
+):
+    """
+    Endpoint for the frontend to submit a new question.
+    It saves the question, calls the ML RAG pipeline to generate a real answer,
+    and saves the answer to the database.
+    """
+    # First, fetch the document to get its uploader's ID
+    document = crud.get_document_by_id(db, document_id=document_id)
+    if not document:
+        raise HTTPException(status_code=404, detail="Document not found")
+
+    user_id_who_asked = document.uploader_id
+
+    # Create the question in the database with a NULL answer first
+    db_question = crud.create_question(
+        db=db,
+        document_id=document_id,
+        user_id=user_id_who_asked,
+        question=question
+    )
+    print(f"New question saved with ID: {db_question.id}. Triggering background ML task.")
+
+
+    pinecone_pdf_id = os.path.splitext(os.path.basename(unquote(document.file_path)))[0]
+
+    background_tasks.add_task(
+        run_ml_qna_in_background,
+        db_question.id,
+        pinecone_pdf_id, # <--- Pass the correct filename ID
+        question.question_text,
+    )
+    # --- END OF KEY CHANGE ---
+
+    # Return the new question object to the frontend immediately.
+    # The frontend will see that `answer_text` is still null.
+    return db_question
+
+
+@app.get("/documents/{document_id}/questions", response_model=List[schemas.Question])
+def get_document_questions(
+        document_id: uuid.UUID,
+        db: Session = Depends(get_db) # Ensure there are no typos like 'get_d b'
+):
+    """
+    Endpoint for the frontend to retrieve the full conversation history
+    (all questions and their answers) for a document.
+    """
+    return crud.get_questions_for_document(db=db, document_id=document_id)
+
+
+# --- ADD THIS NEW ENDPOINT FOR EMAIL AUTOMATION ---
+
+
+@app.patch("/questions/{question_id}/answer")
+def submit_answer(
+        question_id: uuid.UUID,
+        answer: schemas.Answer,
+        db: Session = Depends(get_db)
+):
+    """
+    INTERNAL ENDPOINT for the ML service to submit its generated answer
+    for a question that has already been created.
+    """
+    updated_question = crud.update_question_with_answer(
+        db=db,
+        question_id=question_id,
+        answer_text=answer.answer_text
+    )
+    if not updated_question:
+        raise HTTPException(status_code=404, detail="Question not found")
+
+    print(f"Answer submitted for question ID: {question_id}")
+    return {"status": "success", "question": updated_question}
+
+
+# --- NEW ENDPOINT FOR NOTIFICATIONS ---
+@app.get("/notifications/{department}", response_model=List[schemas.Notification])
+def read_notifications(department: str, db: Session = Depends(get_db)):
+    """Fetches unread notifications for a given department."""
+    notifications = crud.get_notifications_for_department(db, department=department)
+    return notifications