src/search.py

# File: src/search.py
import faiss
import numpy as np
import pickle
import os
from config.settings import Settings
from src.logger import logger

class Search:
    def __init__(self):
        self.index = faiss.IndexFlatL2(384)
        self.texts = []
        if os.path.exists(Settings.FAISS_INDEX_PATH):
            self.index = faiss.read_index(Settings.FAISS_INDEX_PATH)
            if os.path.exists(Settings.TEXTS_PATH):
                with open(Settings.TEXTS_PATH, "rb") as f:
                    self.texts = pickle.load(f)
            logger.info("Đã tải FAISS index và texts")

    def add(self, text, vector):
        self.index.add(np.array([vector], dtype=np.float32))
        self.texts.append(text)
        faiss.write_index(self.index, Settings.FAISS_INDEX_PATH)
        with open(Settings.TEXTS_PATH, "wb") as f:
            pickle.dump(self.texts, f)
        logger.info(f"Đã thêm vào FAISS: {text[:50]}...")

    def search(self, vector, k=1):
        distances, indices = self.index.search(np.array([vector], dtype=np.float32), k)
        return [(self.texts[i], distances[0][j]) for j, i in enumerate(indices[0]) if i != -1]