Update app.py

app.py

@@ -13,7 +13,7 @@ import json
 app = FastAPI()
 model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
 index = faiss.IndexFlatL2(384)  # 384 is the dimensionality of the MiniLM model
-
+documents = []
 
 templates = Jinja2Templates(directory=".")
 
@@ -35,6 +35,7 @@ def embed_strings(request: EmbedRequest):
     new_embeddings = model.encode(new_documents)
     index.add(np.array(new_embeddings))
     new_size = index.ntotal
+    documents.extend(new_documents)
     return {
         "message": f"{len(new_documents)} new strings embedded and added to FAISS database. New size of the database: {new_size}"
     }
@@ -44,11 +45,14 @@ def embed_strings(request: EmbedRequest):
 def search_string(request: SearchRequest):
     embedding = model.encode([request.text])
     distances, indices = index.search(np.array(embedding), request.n)
-    found_documents = index.reconstruct_n(indices[0], int(request.n))
+
+    # Get the documents associated with the returned indices
+    found_documents = [documents[i] for i in indices[0]]
+
     return {
         "distances": distances[0].tolist(),
         "indices": indices[0].tolist(),
-        "documents": found_documents.tolist()
+        "documents": found_documents
     }
 
 #########################
@@ -63,19 +67,19 @@ def reset_database():
     index.reset()
     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)
+@app.get("/admin/embeddings/download")
+def download_embeddings():
+    # Reconstruct the embeddings from the FAISS index
+    embeddings = index.reconstruct_n(0, index.ntotal)
 
-    # Convert the documents list to a JSON string
-    documents_json = json.dumps(documents.tolist())
+    # Convert the embeddings list to a JSON string
+    embeddings_json = json.dumps(embeddings.tolist())
 
     # Create a response with the JSON string as the content
-    response = Response(content=documents_json, media_type="application/json")
+    response = Response(content=embeddings_json, media_type="application/json")
 
     # Set the content disposition header to trigger a download
-    response.headers["Content-Disposition"] = "attachment; filename=documents.json"
+    response.headers["Content-Disposition"] = "attachment; filename=embeddings.json"
 
     return response