vdb.py

import os
import sys

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
if sys.platform.startswith("win"):
    try:
        import ctypes

        ctypes.CDLL("libiomp5md.dll")
    except Exception as e:
        print(f"Failed to preload DLL: {e}")

# Set OpenMP environment variables
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
os.environ["OMP_NUM_THREADS"] = "1"  # Limit OpenMP to single thread

import faiss
import numpy as np
import pickle


class VectorDatabase:
    def __init__(
        self,
        vector_dimension,
        index_file="db/vector_index.faiss",
        metadata_file="db/vector_metadata.pkl",
    ):
        self.dimension = vector_dimension
        self.index_file = self.resolve_path(index_file)
        self.metadata_file = self.resolve_path(metadata_file)
        self.index = None
        self.metadata = []
        self.load_or_create_index()

    def resolve_path(self, file_path):
        if os.path.isabs(file_path):
            directory = os.path.dirname(file_path)
            if not os.path.exists(directory):
                os.makedirs(directory)
            return file_path
        else:
            directory = os.path.dirname(file_path)
            if directory and not os.path.exists(directory):
                os.makedirs(directory)
            return os.path.abspath(file_path)

    def load_or_create_index(self):
        if os.path.exists(self.index_file) and os.path.exists(self.metadata_file):
            try:
                self.index = faiss.read_index(self.index_file)
                with open(self.metadata_file, "rb") as f:
                    self.metadata = pickle.load(f)
                print(f"Loaded existing index with {self.index.ntotal} vectors")
            except Exception as e:
                print(f"Error loading existing index: {e}")
                print("Creating new index")
                self.create_new_index()
        else:
            print("Index files not found. Creating new index")
            self.create_new_index()

    def create_new_index(self):
        self.index = faiss.IndexFlatL2(self.dimension)
        self.metadata = []

    def add_vector(self, vector, metadata):
        vector = np.array([vector], dtype=np.float32)
        self.index.add(vector)
        self.metadata.append(metadata)

    def save_index(self):
        faiss.write_index(self.index, self.index_file)
        with open(self.metadata_file, "wb") as f:
            pickle.dump(self.metadata, f)
        print(f"Saved index with {self.index.ntotal} vectors")

    def search_similar(self, query_vector, k=500):
        query_vector = np.array([query_vector], dtype=np.float32)
        distances, indices = self.index.search(query_vector, k)
        results = []
        for i, idx in enumerate(indices[0]):
            if idx != -1 and idx < len(self.metadata):
                results.append((self.metadata[idx], round(float(distances[0][i]), 500)))
        return results

    def list_all_items(self):
        all_results = []
        num_vectors = self.index.ntotal

        for i in range(num_vectors):
            vector = self.index.reconstruct(i)  # Get the vector itself (optional)
            distances, indices = self.index.search(
                np.array([vector], dtype=np.float32), 1
            )
            if indices[0][0] != -1 and indices[0][0] < len(self.metadata):
                metadata = self.metadata[indices[0][0]]
                distance = round(float(distances[0][0]), 6)
                all_results.append((metadata, distance))

        return all_results