Spaces:
Running
Running
File size: 3,672 Bytes
1a8333a 79d7ca9 dbc6c62 7832703 dbc6c62 0b77457 dbc6c62 bfa9638 7832703 bfa9638 79d7ca9 bfa9638 79d7ca9 7832703 10399f1 7832703 bfa9638 1a8333a bfa9638 7832703 dbc6c62 10f29f7 fb44d11 10f29f7 dbc6c62 79d7ca9 dbc6c62 79d7ca9 dbc6c62 a04e1ac dbc6c62 a04e1ac 9146dbe a04e1ac dbc6c62 bfa9638 a04e1ac bfa9638 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import os
import chromadb
from fastapi import FastAPI, Request, Form, File, UploadFile, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, HTMLResponse, RedirectResponse
from fastapi.templating import Jinja2Templates
from .admin import admin_functions as admin
from .utils.db import UserFaceEmbeddingFunction,ChromaDBFaceHelper
from .api import userlogin, userlogout, userchat, userupload
from .utils.db import tinydb_helper, ChromaDBFaceHelper
CHROMADB_LOC = "/home/user/data/chromadb"
app = FastAPI()
# Add middleware
# Set all origins to wildcard for simplicity, but we should limit this in production
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Persitent storage for chromadb setup in /data volume
ec_client = chromadb.PersistentClient(CHROMADB_LOC)
# The following collection reference is needed for admin function to register face
user_faces_db = ec_client.get_or_create_collection(name="user_faces_db", embedding_function=UserFaceEmbeddingFunction())
@app.on_event("startup")
async def startup_event():
global chromadb_face_helper
# Assuming chromadb persistent store client for APIs is in helper
db_path = CHROMADB_LOC
chromadb_face_helper = ChromaDBFaceHelper(db_path) # Used by APIs
# Perform any other startup tasks here
print(f"MODEL_PATH in main.py = {os.getenv('MODEL_PATH')} ")
# Mount static files
app.mount("/static", StaticFiles(directory="static"), name="static")
# Setup Jinja2Templates to point to the templates directory
templates = Jinja2Templates(directory="app/admin/templates")
@app.get("/")
async def get_admin_login(request: Request):
return templates.TemplateResponse("admin_login.html", {"request": request})
# Admin Login Handler
@app.post("/admin/login", response_class=HTMLResponse)
async def handle_admin_login(request: Request, username: str = Form(...), password: str = Form(...)):
if admin.verify_admin_password(username, password):
# Redirect to user registration page upon successful login
return RedirectResponse(url="/admin/register_user", status_code=303)
else:
# Reload login page with error message
return templates.TemplateResponse("admin_login.html", {"request": request, "error": "Invalid password"})
# To display the register user page
@app.get("/admin/register_user", response_class=HTMLResponse)
async def get_user_registration(request: Request):
# Render the registration form
return templates.TemplateResponse("user_registration.html", {"request": request})
# User Registration Handler
@app.post("/admin/register_user", response_class=HTMLResponse)
async def handle_user_registration(request: Request, email: str = Form(...), name: str = Form(...), role: str = Form(...), file: UploadFile = File(...)):
user_id = await admin.register_user(user_faces_db, email, name, role, file)
if user_id:
# Calculate disk usage
disk_usage = admin.get_disk_usage("/home/user/data")
# Redirect or display a success message
return templates.TemplateResponse("registration_success.html", {
"request": request,
"disk_usage": disk_usage
})
else:
# Reload registration page with error message
return templates.TemplateResponse("user_registration.html", {"request": request, "error": "Registration failed"})
app.include_router(userlogin.router)
app.include_router(userlogout.router)
app.include_router(userchat.router)
app.include_router(userupload.router)
|