Update app.py

app.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI, Request, Query
 from fastapi.templating import Jinja2Templates
+from fastapi import File, UploadFile
 from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer
 import faiss
@@ -8,8 +9,7 @@ import numpy as np
 app = FastAPI()
 model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
 index = faiss.IndexFlatL2(384)  # 384 is the dimensionality of the MiniLM model
-# Create a list to store the documents
-documents = []
+
 
 templates = Jinja2Templates(directory=".")
 
@@ -27,33 +27,24 @@ def read_root(request: Request):
 
 @app.post("/embed")
 def embed_strings(request: EmbedRequest):
-    # Add the new texts to the documents list
     new_documents = request.texts
-    documents.extend(new_documents)
-
-    # Encode the new documents and add them to the FAISS database
     new_embeddings = model.encode(new_documents)
     index.add(np.array(new_embeddings))
-
-    # Get the new size of the FAISS database
-    new_size = len(documents)
-
+    new_size = index.ntotal
     return {
         "message": f"{len(new_documents)} new strings embedded and added to FAISS database. New size of the database: {new_size}"
     }
 
+
 @app.post("/search")
 def search_string(request: SearchRequest):
     embedding = model.encode([request.text])
     distances, indices = index.search(np.array(embedding), request.n)
-
-    # Get the documents associated with the returned indices
-    found_documents = [documents[i] for i in indices[0]]
-
+    found_documents = index.reconstruct_n(indices[0], request.n)
     return {
         "distances": distances[0].tolist(),
         "indices": indices[0].tolist(),
-        "documents": found_documents
+        "documents": found_documents.tolist()
     }
 
 #########################
@@ -61,18 +52,20 @@ def search_string(request: SearchRequest):
 #########################
 @app.get("/admin/database/length")
 def get_database_length():
-    return {"length": len(documents)}
+    return {"length": index.ntotal}
 
-@app.post("/admin/database/clear")
-def clear_database():
-    documents.clear()
+@app.post("/admin/database/reset")
+def reset_database():
     index.reset()
-    return {"message": "Database cleared"}
+    return {"message": "Database reset"}
 
 @app.get("/admin/documents/download")
 def download_documents():
+    # Reconstruct the documents from the FAISS index
+    documents = index.reconstruct_n(0, index.ntotal)
+
     # Convert the documents list to a JSON string
-    documents_json = json.dumps(documents)
+    documents_json = json.dumps(documents.tolist())
 
     # Create a response with the JSON string as the content
     response = Response(content=documents_json, media_type="application/json")