Create vector_store/vector_store.py

AI_Agent/vector_store/vector_store.py

@@ -0,0 +1,27 @@
+# AI_Agent/vector_store/vector_store.py
+from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
+
+class SimpleVectorStore:
+    def __init__(self, dim=384, model_name="all-MiniLM-L6-v2"):
+        self.embed_model = SentenceTransformer(model_name)
+        self.dim = dim
+        self.index = faiss.IndexFlatL2(dim)
+        self.docs = []
+
+    def add_documents(self, docs):
+        embeddings = self.embed_model.encode(docs, convert_to_numpy=True)
+        if embeddings.ndim == 1:
+            embeddings = embeddings.reshape(1, -1)
+        self.index.add(embeddings.astype("float32"))
+        self.docs.extend(docs)
+
+    def search(self, query, k=3):
+        query_emb = self.embed_model.encode([query], convert_to_numpy=True).astype("float32")
+        D, I = self.index.search(query_emb, k)
+        results = []
+        for idx, dist in zip(I[0], D[0]):
+            if idx < len(self.docs):
+                results.append((self.docs[idx], float(dist)))
+        return results