dtyago commited on
Commit
97db660
1 Parent(s): 4f116c5

Implement Admin db clean up

Browse files
app/admin/admin_functions.py CHANGED
@@ -3,6 +3,14 @@ from typing import Optional
3
  import bcrypt
4
  import os
5
  import shutil
 
 
 
 
 
 
 
 
6
  from ..utils import get_user_cropped_image_from_photo
7
 
8
 
@@ -68,6 +76,7 @@ def verify_admin_password(submitted_user: str, submitted_password: str) -> bool:
68
 
69
  return False
70
 
 
71
  def get_disk_usage(path="/home/user/data"):
72
  total, used, free = shutil.disk_usage(path)
73
  # Convert bytes to MB by dividing by 2^20
@@ -83,3 +92,29 @@ def get_disk_usage(path="/home/user/data"):
83
  # - Moderating chat messages or viewing chat history.
84
  # - Managing system settings or configurations.
85
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3
  import bcrypt
4
  import os
5
  import shutil
6
+
7
+ # Import vector store for database operations
8
+ from langchain.vectorstores import Chroma
9
+ # Import embeddings module from langchain for vector representations of text
10
+ from langchain.embeddings import HuggingFaceEmbeddings
11
+ from app.main import CHROMADB_LOC
12
+
13
+ from app.utils.chat_rag import sanitize_collection_name
14
  from ..utils import get_user_cropped_image_from_photo
15
 
16
 
 
76
 
77
  return False
78
 
79
+ # Get disk usage
80
  def get_disk_usage(path="/home/user/data"):
81
  total, used, free = shutil.disk_usage(path)
82
  # Convert bytes to MB by dividing by 2^20
 
92
  # - Moderating chat messages or viewing chat history.
93
  # - Managing system settings or configurations.
94
 
95
+ # Display all faces in collection
96
+ def faces_count(client, db):
97
+ return {
98
+ "face_count" : db.count(),
99
+ "all_faces" : db.get(),
100
+ "all_collections" : client.list_collections() # List all collections at this location
101
+ }
102
+
103
+ # Delete all faces in collection
104
+ def remove_all_faces(client, user_faces_collection="user_faces_db"):
105
+ # Fetch all user IDs from the user_faces_db collection
106
+ all_user_ids = client.get_all_ids(collection_name=user_faces_collection)
107
+
108
+ # Loop through all user IDs and delete associated collections
109
+ for user_id in all_user_ids:
110
+ sanitized_collection_name = sanitize_collection_name(user_id)
111
+ vectordb = Chroma(
112
+ collection_name=sanitized_collection_name,
113
+ embedding_function=HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2'),
114
+ persist_directory=f"{CHROMADB_LOC}/{sanitized_collection_name}", # Optional: Separate directory for each user's data
115
+ )
116
+ all_ids = vectordb._collection.get()
117
+ vectordb._collection.delete(ids=all_ids)
118
+ # Finally, delete the user_faces_db collection itself
119
+ client.delete_collection(user_faces_collection)
120
+ print(f"All user collections and {user_faces_collection} have been removed.")
app/admin/templates/data_management.html ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ <!DOCTYPE html>
2
+ <html>
3
+ <head>
4
+ <title>Econnect Data Management</title>
5
+ <meta name="viewport" content="width=device-width, initial-scale=1.0">
6
+ <link href="static/css/mvp.css" rel="stylesheet" media="screen">
7
+ <!-- Bootstrap CSS -->
8
+ <link rel="stylesheet" href="https://stackpath.bootstrapcdn.com/bootstrap/4.3.1/css/bootstrap.min.css">
9
+ </head>
10
+ <body>
11
+ <h1>Face DB</h1>
12
+ <div>
13
+ <p>Total Faces: {{ faces.face_count }}</p>
14
+ <h2>Registered faces:</h2>
15
+ <p>{{ faces.all_faces }}</p>
16
+ <hr>
17
+ <h2>All collections:</h2>
18
+ <p>{{ faces.all_collections }}</p>
19
+ </div>
20
+ <hr>
21
+ <div class="container">
22
+ <h2>User Registration</h2>
23
+ <form id="deleteForm" action="/admin/delete_faces" method="post" enctype="multipart/form-data">
24
+ <h2>Delete and clean vectordb:</h2>
25
+ <!-- Updated button to only trigger the modal -->
26
+ <button type="button" class="btn btn-default" data-toggle="modal" data-target="#confirmationModal">Delete ALL</button>
27
+ </form>
28
+ {% if error %}
29
+ <p class="error"><strong>Error:</strong> {{ error }}</p>
30
+ {% endif %}
31
+ </div>
32
+
33
+ <!-- Modal -->
34
+ <div class="modal fade" id="confirmationModal" tabindex="-1" role="dialog" aria-labelledby="confirmationModalLabel" aria-hidden="true">
35
+ <div class="modal-dialog" role="document">
36
+ <div class="modal-content">
37
+ <div class="modal-header">
38
+ <h5 class="modal-title" id="confirmationModalLabel">Confirm Deletion</h5>
39
+ <button type="button" class="close" data-dismiss="modal" aria-label="Close">
40
+ <span aria-hidden="true">&times;</span>
41
+ </button>
42
+ </div>
43
+ <div class="modal-body">
44
+ Are you sure you want to delete ALL faces?
45
+ </div>
46
+ <div class="modal-footer">
47
+ <button type="button" class="btn btn-secondary" data-dismiss="modal">Cancel</button>
48
+ <button type="button" class="btn btn-danger" id="confirmDelete">Delete All</button>
49
+ </div>
50
+ </div>
51
+ </div>
52
+ </div>
53
+
54
+ <!-- Bootstrap and jQuery scripts -->
55
+ <script src="https://code.jquery.com/jquery-3.3.1.slim.min.js"></script>
56
+ <script src="https://cdnjs.cloudflare.com/ajax/libs/popper.js/1.14.7/umd/popper.min.js"></script>
57
+ <script src="https://stackpath.bootstrapcdn.com/bootstrap/4.3.1/js/bootstrap.min.js"></script>
58
+ <script>
59
+ document.addEventListener('DOMContentLoaded', (event) => {
60
+ document.getElementById('confirmDelete').addEventListener('click', function() {
61
+ document.getElementById('deleteForm').submit();
62
+ });
63
+ });
64
+ </script>
65
+ </body>
66
+ </html>
app/admin/templates/user_registration.html CHANGED
@@ -30,5 +30,9 @@
30
  <p class="error"><strong>Error:</strong> {{ error }}
31
  {% endif %}
32
  </div>
 
 
 
 
33
  </body>
34
  </html>
 
30
  <p class="error"><strong>Error:</strong> {{ error }}
31
  {% endif %}
32
  </div>
33
+ <hr>
34
+ <div>
35
+ <a href="/admin/data_management">Data Management</a>
36
+ </div>
37
  </body>
38
  </html>
app/main.py CHANGED
@@ -1,16 +1,16 @@
1
  import os
2
  import chromadb
3
 
4
- from fastapi import FastAPI, Request, Form, File, UploadFile, Depends
5
  from fastapi.middleware.cors import CORSMiddleware
6
  from fastapi.staticfiles import StaticFiles
7
- from fastapi.responses import FileResponse, HTMLResponse, RedirectResponse
8
  from fastapi.templating import Jinja2Templates
9
 
10
  from .admin import admin_functions as admin
11
  from .utils.db import UserFaceEmbeddingFunction,ChromaDBFaceHelper
12
  from .api import userlogin, userlogout, userchat, userupload
13
- from .utils.db import tinydb_helper, ChromaDBFaceHelper
14
 
15
  CHROMADB_LOC = "/home/user/data/chromadb"
16
 
@@ -28,6 +28,7 @@ app.add_middleware(
28
 
29
  # Persitent storage for chromadb setup in /data volume
30
  ec_client = chromadb.PersistentClient(CHROMADB_LOC)
 
31
  # The following collection reference is needed for admin function to register face
32
  user_faces_db = ec_client.get_or_create_collection(name="user_faces_db", embedding_function=UserFaceEmbeddingFunction())
33
 
@@ -85,6 +86,27 @@ async def handle_user_registration(request: Request, email: str = Form(...), nam
85
  # Reload registration page with error message
86
  return templates.TemplateResponse("user_registration.html", {"request": request, "error": "Registration failed"})
87
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
88
  app.include_router(userlogin.router)
89
  app.include_router(userlogout.router)
90
  app.include_router(userchat.router)
 
1
  import os
2
  import chromadb
3
 
4
+ from fastapi import FastAPI, Request, Form, File, UploadFile
5
  from fastapi.middleware.cors import CORSMiddleware
6
  from fastapi.staticfiles import StaticFiles
7
+ from fastapi.responses import HTMLResponse, RedirectResponse
8
  from fastapi.templating import Jinja2Templates
9
 
10
  from .admin import admin_functions as admin
11
  from .utils.db import UserFaceEmbeddingFunction,ChromaDBFaceHelper
12
  from .api import userlogin, userlogout, userchat, userupload
13
+ from .utils.db import ChromaDBFaceHelper
14
 
15
  CHROMADB_LOC = "/home/user/data/chromadb"
16
 
 
28
 
29
  # Persitent storage for chromadb setup in /data volume
30
  ec_client = chromadb.PersistentClient(CHROMADB_LOC)
31
+
32
  # The following collection reference is needed for admin function to register face
33
  user_faces_db = ec_client.get_or_create_collection(name="user_faces_db", embedding_function=UserFaceEmbeddingFunction())
34
 
 
86
  # Reload registration page with error message
87
  return templates.TemplateResponse("user_registration.html", {"request": request, "error": "Registration failed"})
88
 
89
+ # To display admin utilities
90
+ @app.get("/admin/data_management", response_class=HTMLResponse)
91
+ async def get_db_details(request: Request):
92
+ # Render the Chroma DB details
93
+ faces = admin.faces_count(user_faces_db)
94
+ return templates.TemplateResponse("data_management.html", {
95
+ "request": request,
96
+ "faces" : faces
97
+ })
98
+
99
+ @app.post("/admin/delete_faces")
100
+ async def delete_faces(request: Request):
101
+ try:
102
+ # Call your function to remove all faces
103
+ admin.remove_all_faces(ec_client)
104
+ request.session['flash'] = "All user data successfully deleted."
105
+ except Exception as e:
106
+ request.session['flash'] = f"Failed to delete user data: {str(e)}"
107
+
108
+ return RedirectResponse(url="/admin/data_management", status_code=303)
109
+
110
  app.include_router(userlogin.router)
111
  app.include_router(userlogout.router)
112
  app.include_router(userchat.router)
app/utils/chat_rag.py CHANGED
@@ -51,7 +51,7 @@ def get_vectordb_for_user(user_collection_name):
51
  vectordb = Chroma(
52
  collection_name=user_collection_name,
53
  embedding_function=HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2'),
54
- persist_directory=f"{CHROMADB_LOC}/{user_collection_name}", # Optional: Separate directory for each user's data
55
  )
56
  return vectordb
57
 
 
51
  vectordb = Chroma(
52
  collection_name=user_collection_name,
53
  embedding_function=HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2'),
54
+ persist_directory=f"{CHROMADB_LOC}" # Optional: Separate directory for each user's data
55
  )
56
  return vectordb
57
 
requirements.txt CHANGED
@@ -18,4 +18,4 @@ tinydb # The in memory database for storing JWT tokens
18
  langchain # Langgchain for RAG
19
  llama-cpp-python # To load the model
20
  sentence-transformers # For text embeddings
21
- pypdf # Handling PDF files
 
18
  langchain # Langgchain for RAG
19
  llama-cpp-python # To load the model
20
  sentence-transformers # For text embeddings
21
+ pypdf # Handling PDF files