rag.py

import faiss
import numpy as np

from embeddings import get_embeddings


class RAG:

    def __init__(self):
        self.index = None
        self.chunks = []

    def chunk_text(self, text, chunk_size=500):

        words = text.split()

        chunks = []

        for i in range(0, len(words), chunk_size):
            chunks.append(" ".join(words[i:i + chunk_size]))

        return chunks

    def create_index(self, text):

        self.chunks = self.chunk_text(text)
        if not self.chunks:
            return

        embeddings = get_embeddings(self.chunks)

        embeddings = np.atleast_2d(np.array(embeddings)).astype("float32")

        dimension = embeddings.shape[1]

        self.index = faiss.IndexFlatL2(dimension)

        self.index.add(embeddings)

    def search(self, query, top_k=3):
        if self.index is None or not self.chunks:
            return []

        query_embedding = get_embeddings([query])
        query_embedding = np.atleast_2d(np.array(query_embedding)).astype("float32")

        top_k = min(top_k, len(self.chunks))
        distances, indices = self.index.search(query_embedding, top_k)

        results = []

        for score, idx in zip(distances[0], indices[0]):

            if idx == -1:
                continue

            results.append(
                {
                    "chunk": self.chunks[idx],
                    "score": float(score),
                    "chunk_id": int(idx)
                }
            )

        return results