Rename

1 files changed, 23 insertions, 0 deletions

diff --git a/rag/db/embeddings.py b/rag/db/embeddings.py
new file mode 100644
index 0000000..f9d0735
--- /dev/null
+++ b/rag/db/embeddings.py
@@ -0,0 +1,23 @@
+import os
+from typing import Tuple
+
+import faiss
+import numpy as np
+
+
+# TODO: inner product distance metric?
+class Embeddings:
+    def __init__(self):
+        self.dim = os.environ["EMBEDDING_DIM"]
+        self.index = faiss.IndexFlatL2(self.dim)
+        # TODO: load from file
+
+    def add(self, embeddings: np.ndarray):
+        # TODO: save to file
+        self.index.add(embeddings)
+
+    def search(
+        self, query: np.ndarray, neighbors: int = 4
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        score, indices = self.index.search(query, neighbors)
+        return score, indices