ColeMurray · September 8, 2024 01:22 · Sep 8, 2024
diff --git a/open-hands-demo-faiss-server.py b/open-hands-demo-faiss-server.py
@@ -0,0 +1,121 @@
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List, Optional, Dict
+import faiss
+import numpy as np
+import os
+
+app = FastAPI()
+
+class QueryRequest(BaseModel):
+    query: List[float]
+    namespace: str
+    identifier: str
+    num_results: Optional[int] = 10
+
+import pickle
+
+class VectorDatabase:
+    def __init__(self, base_path: str = "./vector_dbs"):
+        self.base_path = base_path
+        self.databases: Dict[str, faiss.IndexFlatL2] = {}
+        self.vectors: Dict[str, List[np.ndarray]] = {}
+        self.ids: Dict[str, List[int]] = {}
+        os.makedirs(base_path, exist_ok=True)
+        self.load_all_dbs()
+
+    def db_path(self, namespace: str, identifier: str) -> str:
+        return os.path.join(self.base_path, f"{namespace}_{identifier}.index")
+
+    def metadata_path(self, namespace: str, identifier: str) -> str:
+        return os.path.join(self.base_path, f"{namespace}_{identifier}.metadata")
+
+    def load_all_dbs(self):
+        for filename in os.listdir(self.base_path):
+            if filename.endswith(".index"):
+                namespace, identifier = filename[:-6].split("_", 1)
+                self.load_db(namespace, identifier)
+
+    def load_db(self, namespace: str, identifier: str):
+        db_key = f"{namespace}_{identifier}"
+        index_path = self.db_path(namespace, identifier)
+        metadata_path = self.metadata_path(namespace, identifier)
+        print(f"Checking index path: {index_path}")
+        print(f"Checking metadata path: {metadata_path}")
+        if os.path.exists(index_path):
+            self.databases[db_key] = faiss.read_index(index_path)
+            if os.path.exists(metadata_path):
+                with open(metadata_path, 'rb') as f:
+                    metadata = pickle.load(f)
+                    self.vectors[db_key] = metadata['vectors']
+                    self.ids[db_key] = metadata['ids']
+                print(f"Metadata file contents: {metadata}")
+            else:
+                print("Metadata file does not exist")
+                self.vectors[db_key] = []
+                self.ids[db_key] = []
+            print(f"Loaded database: {db_key}")
+            print(f"Number of vectors: {self.databases[db_key].ntotal}")
+            print(f"Number of IDs: {len(self.ids[db_key])}")
+            print(f"IDs: {self.ids[db_key]}")
+        else:
+            print(f"Index file does not exist: {index_path}")
+
+    def save_db(self, namespace: str, identifier: str):
+        db_key = f"{namespace}_{identifier}"
+        if db_key in self.databases:
+            index_path = self.db_path(namespace, identifier)
+            metadata_path = self.metadata_path(namespace, identifier)
+            faiss.write_index(self.databases[db_key], index_path)
+            with open(metadata_path, 'wb') as f:
+                pickle.dump({'vectors': self.vectors[db_key], 'ids': self.ids[db_key]}, f)
+            print(f"Saved database: {db_key}")
+
+    def create_or_get_db(self, namespace: str, identifier: str, dim: int):
+        db_key = f"{namespace}_{identifier}"
+        if db_key not in self.databases:
+            self.databases[db_key] = faiss.IndexFlatL2(dim)
+            self.vectors[db_key] = []
+            self.ids[db_key] = []
+        return self.databases[db_key]
+
+    def add_vector(self, namespace: str, identifier: str, vector: List[float], id: int):
+        db_key = f"{namespace}_{identifier}"
+        db = self.create_or_get_db(namespace, identifier, len(vector))
+        np_vector = np.array([vector], dtype=np.float32)
+        db.add(np_vector)
+        self.vectors[db_key].append(np_vector)
+        self.ids[db_key].append(id)
+        self.save_db(namespace, identifier)
+
+    def query(self, namespace: str, identifier: str, query_vector: List[float], k: int):
+        db_key = f"{namespace}_{identifier}"
+        if db_key not in self.databases:
+            raise HTTPException(status_code=404, detail="Database not found")
+        db = self.databases[db_key]
+        np_query = np.array([query_vector], dtype=np.float32)
+        distances, indices = db.search(np_query, k)
+        results = []
+        for idx, score in zip(indices[0], distances[0]):
+            if idx < len(self.ids[db_key]):
+                results.append({"id": self.ids[db_key][idx], "score": float(score)})
+            else:
+                results.append({"id": None, "score": float(score)})
+        return results
+
+vector_db = VectorDatabase()
+
+@app.post("/query")
+def query_vector_db(request: QueryRequest):
+    results = vector_db.query(request.namespace, request.identifier, request.query, request.num_results)
+    return {"results": results}
+
+@app.post("/add_vector")
+def add_vector(namespace: str, identifier: str, vector: List[float], id: int):
+    vector_db.add_vector(namespace, identifier, vector, id)
+    return {"message": "Vector added successfully"}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)
No results found