BitBasher/EduConnect

1import os2import chromadb3 4from fastapi import FastAPI, Request, Form, File, UploadFile5from fastapi.middleware.cors import CORSMiddleware6from fastapi.staticfiles import StaticFiles7from fastapi.responses import HTMLResponse, RedirectResponse8from fastapi.templating import Jinja2Templates9 10from admin import admin_functions as admin11from utils.db import UserFaceEmbeddingFunction,ChromaDBFaceHelper12from api import userlogin, userlogout, userchat, userupload13from utils.db import ChromaDBFaceHelper14from utils.chat_rag import LlamaModelSingleton15 16app = FastAPI()17 18CHROMADB_LOC = os.getenv('CHROMADB_LOC')19 20# Add middleware21# Set all origins to wildcard for simplicity, but we should limit this in production22app.add_middleware(23    CORSMiddleware,24    allow_origins=["*"],25    allow_credentials=True,26    allow_methods=["*"],27    allow_headers=["*"],28)29 30# Persitent storage for chromadb setup in /data volume31ec_client = chromadb.PersistentClient(CHROMADB_LOC)32 33# The following collection reference is needed for admin function to register face34user_faces_db = ec_client.get_or_create_collection(name="user_faces_db", embedding_function=UserFaceEmbeddingFunction())35 36 37@app.on_event("startup")38async def startup_event():39    global chromadb_face_helper40    # Assuming chromadb persistent store client for APIs is in helper41    db_path = CHROMADB_LOC42    chromadb_face_helper = ChromaDBFaceHelper(db_path) # Used by APIs43    44    # Perform any other startup tasks here45    # Preload the LLM model46    await LlamaModelSingleton.get_instance()47    print("LLM model loaded and ready.")48 49    print(f"MODEL_PATH in main.py = {os.getenv('MODEL_PATH')} ")50 51# Mount static files52app.mount("/static", StaticFiles(directory="static"), name="static")53 54# Setup Jinja2Templates to point to the templates directory55templates = Jinja2Templates(directory="admin/templates")56 57@app.get("/")58async def get_admin_login(request: Request):59    return templates.TemplateResponse("admin_login.html", {"request": request})60 61# Admin Login Handler62@app.post("/admin/login", response_class=HTMLResponse)63async def handle_admin_login(request: Request, username: str = Form(...), password: str = Form(...)):64 65    if admin.verify_admin_password(username, password):66        # Redirect to user registration page upon successful login67        return RedirectResponse(url="/admin/register_user", status_code=303)68    else:69        # Reload login page with error message70        return templates.TemplateResponse("admin_login.html", {"request": request, "error": "Invalid password"})71 72# To display the register user page73@app.get("/admin/register_user", response_class=HTMLResponse)74async def get_user_registration(request: Request):75    # Render the registration form76    return templates.TemplateResponse("user_registration.html", {"request": request})77 78# User Registration Handler79@app.post("/admin/register_user", response_class=HTMLResponse)80async def handle_user_registration(request: Request, email: str = Form(...), name: str = Form(...), role: str = Form(...), file: UploadFile = File(...)):81    user_id = await admin.register_user(user_faces_db, email, name, role, file)82    if user_id:83        # Calculate disk usage84        disk_usage = admin.get_disk_usage("/home/user/data")85 86        # Redirect or display a success message87        return templates.TemplateResponse("registration_success.html", {88            "request": request,89            "disk_usage": disk_usage 90        })91    else:92        # Reload registration page with error message93        return templates.TemplateResponse("user_registration.html", {"request": request, "error": "Registration failed"})94 95# To display admin utilities96@app.get("/admin/data_management", response_class=HTMLResponse)97async def get_db_details(request: Request):98    # Render the Chroma DB details99    faces = admin.faces_count(ec_client, user_faces_db)100    return templates.TemplateResponse("data_management.html", {101        "request": request,102        "faces" : faces103    })104 105@app.post("/admin/delete_faces")106async def delete_faces(request: Request):107    try:108        # Call your function to remove all faces109        admin.remove_all_faces(ec_client)110        request.session['flash'] = "All user data successfully deleted."111    except Exception as e:112        request.session['flash'] = f"Failed to delete user data: {str(e)}"113    114    return RedirectResponse(url="/admin/data_management", status_code=303)115 116app.include_router(userlogin.router)117app.include_router(userlogout.router)118app.include_router(userchat.router)119app.include_router(userupload.router)120