rag/retriever.py

from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
from config import EMBEDDING_MODEL


class SimpleVectorStore:
    
    def __init__(self, embeddings, documents):
        self.embeddings = embeddings
        self.documents = documents
        self.index = None
        self._build_index()
    
    def _build_index(self):
        texts = [doc.page_content for doc in self.documents]
        vectors = self.embeddings.encode(texts)
        
        dimension = vectors.shape[1]
        self.index = faiss.IndexFlatL2(dimension)
        self.index.add(np.array(vectors).astype('float32'))
    
    def similarity_search(self, query, k=3):
        query_vector = self.embeddings.encode([query])
        distances, indices = self.index.search(
            np.array(query_vector).astype('float32'), 
            k
        )
        
        return [self.documents[i] for i in indices[0]]


_embeddings_model = SentenceTransformer(EMBEDDING_MODEL)


def build_vectorstore(documents):
    """Constrói vectorstore a partir de documentos"""
    return SimpleVectorStore(_embeddings_model, documents)