Update src/vectorstore.py

src/vectorstore.py

@@ -1,71 +1,84 @@
 import faiss
 import numpy as np
+import os
 
 # ==========================================================
 # BUILD FAISS INDEX (Cosine Similarity Safe Version)
 # ==========================================================
-def build_faiss_index(embeddings: list):
+def build_faiss_index(embeddings):
     """
     📘 Builds a FAISS index optimized for cosine similarity (float32-safe, dimension-aware).
 
     Args:
-        embeddings (list): List of embedding vectors (lists of floats).
+        embeddings (list | np.ndarray): List or array of embedding vectors.
     Returns:
         faiss.IndexFlatIP: FAISS index for cosine similarity search.
     """
-    if not embeddings or not isinstance(embeddings, list):
-        raise ValueError("❌ Invalid input: embeddings must be a non-empty list.")
-
-    # ✅ Convert to numpy float32
-    vectors = np.array(embeddings, dtype="float32")
-
-    # ✅ Normalize (so cosine == inner product)
+    # 🧩 Validation
+    if embeddings is None:
+        raise ValueError("❌ No embeddings provided to build FAISS index.")
+
+    # ✅ Convert ndarray to list if needed
+    if isinstance(embeddings, np.ndarray):
+        # Handle (n, d) NumPy embeddings
+        if embeddings.ndim == 1:
+            embeddings = embeddings.reshape(1, -1)
+        vectors = embeddings.astype("float32")
+    elif isinstance(embeddings, list):
+        vectors = np.array(embeddings, dtype="float32")
+    else:
+        raise TypeError(f"❌ Unexpected embeddings type: {type(embeddings)}")
+
+    # ✅ Ensure there are embeddings to index
+    if vectors.size == 0:
+        raise ValueError("❌ Empty embeddings array provided.")
+
+    # ✅ Normalize for cosine similarity
     faiss.normalize_L2(vectors)
 
+    # ✅ Build FAISS index (Inner Product = Cosine Similarity)
     dim = vectors.shape[1]
     index = faiss.IndexFlatIP(dim)
     index.add(vectors)
 
-    print(f"✅ FAISS index built with {index.ntotal} vectors of dimension {dim} (cosine similarity)")
+    print(f"✅ FAISS index built successfully — {index.ntotal} vectors | dim={dim} | cosine similarity mode.")
     return index
 
 
 # ==========================================================
 # SEARCH FAISS INDEX (with sanity checks)
 # ==========================================================
-def search_faiss(query_embedding: np.ndarray, index, chunks: list, top_k: int = 3):
+def search_faiss(query_embedding, index, chunks, top_k=3):
     """
-    🔍 Searches FAISS index for the most semantically similar chunks to query.
+    🔍 Searches FAISS index for semantically similar chunks.
 
     Args:
-        query_embedding (np.ndarray): Query embedding vector.
-        index (faiss.IndexFlatIP): Pre-built FAISS index.
-        chunks (list[str]): Text chunks used to build the index.
+        query_embedding (np.ndarray): Query vector (1D or 2D).
+        index (faiss.IndexFlatIP): Built FAISS index.
+        chunks (list[str]): Original document chunks.
         top_k (int): Number of results to return.
 
     Returns:
-        list[str]: Top matching text chunks.
+        list[str]: Top-matching chunks.
     """
     if index is None or index.ntotal == 0:
         raise ValueError("❌ FAISS index is empty or not initialized.")
 
-    # ✅ Ensure query vector is correct dtype and shape
-    query_vector = np.array([query_embedding], dtype="float32")
+    # ✅ Convert query to float32 and normalize
+    query_vector = np.array([query_embedding], dtype="float32") if query_embedding.ndim == 1 else query_embedding.astype("float32")
     faiss.normalize_L2(query_vector)
 
     # ✅ Dimension check
     if query_vector.shape[1] != index.d:
         raise ValueError(
             f"❌ Embedding dimension mismatch: query={query_vector.shape[1]}, index={index.d}. "
-            "Please rebuild the FAISS index with the current embedding model."
+            "Rebuild FAISS index with embeddings from the same model."
         )
 
-    # Search
+    # 🔍 Run search
     distances, indices = index.search(query_vector, top_k)
-
-    # ✅ Return sorted top-k chunks
     results = []
-    for i, idx in enumerate(indices[0]):
+    for idx in indices[0]:
         if 0 <= idx < len(chunks):
             results.append(chunks[idx])
 
@@ -73,17 +86,38 @@ def search_faiss(query_embedding: np.ndarray, index, chunks: list, top_k: int =
     return results
 
 
+# ==========================================================
+# SAVE / LOAD INDEX (Optional Utility)
+# ==========================================================
+def save_faiss_index(index, path="faiss_index.bin"):
+    """💾 Save FAISS index to disk."""
+    faiss.write_index(index, path)
+    print(f"💾 FAISS index saved to {path}")
+
+
+def load_faiss_index(path="faiss_index.bin"):
+    """📂 Load FAISS index from disk."""
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"❌ No FAISS index found at {path}")
+    index = faiss.read_index(path)
+    print(f"📂 FAISS index loaded from {path}")
+    return index
+
+
 # ==========================================================
 # DEBUG / LOCAL TEST
 # ==========================================================
 if __name__ == "__main__":
-    sample_embeddings = [
+    # Example embeddings (3 vectors, dim=3)
+    sample_embeddings = np.array([
         [0.1, 0.2, 0.3],
         [0.2, 0.1, 0.4],
         [0.9, 0.8, 0.7]
-    ]
-    query_vec = [0.15, 0.18, 0.35]
+    ], dtype="float32")
+
+    query_vec = np.array([0.15, 0.18, 0.35], dtype="float32")
 
+    # ✅ Build and search
     idx = build_faiss_index(sample_embeddings)
     results = search_faiss(query_vec, idx, ["Chunk A", "Chunk B", "Chunk C"], top_k=2)