src/rag/retriever.py

import os
import pickle
from pathlib import Path
from src.config.settings import FAQS_TXT, VECTORS_DIR

_chunks = []
_index = None


def _chunk_text(text: str, chunk_size: int = 400, overlap: int = 80) -> list[str]:
    words = text.split()
    chunks = []
    i = 0
    while i < len(words):
        chunk = " ".join(words[i: i + chunk_size])
        chunks.append(chunk)
        i += chunk_size - overlap
    return chunks


def build_vector_store():
    try:
        import faiss
        import numpy as np
        import google.generativeai as genai
        from src.config.settings import GEMINI_API_KEY

        genai.configure(api_key=GEMINI_API_KEY)
        text = FAQS_TXT.read_text()
        chunks = _chunk_text(text)

        embeddings = []
        for chunk in chunks:
            result = genai.embed_content(
                model="models/embedding-001",
                content=chunk,
                task_type="retrieval_document",
            )
            embeddings.append(result["embedding"])

        emb_array = np.array(embeddings, dtype="float32")
        index = faiss.IndexFlatL2(emb_array.shape[1])
        index.add(emb_array)

        VECTORS_DIR.mkdir(parents=True, exist_ok=True)
        faiss.write_index(index, str(VECTORS_DIR / "faqs.index"))
        with open(VECTORS_DIR / "chunks.pkl", "wb") as f:
            pickle.dump(chunks, f)

        print(f"Vector store built with {len(chunks)} chunks.")
    except Exception as e:
        print(f"Vector store build failed (using keyword fallback): {e}")


def load_vector_store():
    global _chunks, _index
    chunks_path = VECTORS_DIR / "chunks.pkl"
    index_path = VECTORS_DIR / "faqs.index"

    if chunks_path.exists():
        with open(chunks_path, "rb") as f:
            _chunks = pickle.load(f)

    if index_path.exists():
        try:
            import faiss
            _index = faiss.read_index(str(index_path))
        except Exception:
            _index = None

    # Always load raw text as fallback
    if not _chunks and FAQS_TXT.exists():
        text = FAQS_TXT.read_text()
        _chunks = _chunk_text(text)


def retrieve_faq(query: str, top_k: int = 3) -> str:
    global _chunks, _index

    if not _chunks:
        load_vector_store()

    # Try FAISS semantic search first
    if _index is not None and _chunks:
        try:
            import numpy as np
            import google.generativeai as genai
            from src.config.settings import GEMINI_API_KEY

            genai.configure(api_key=GEMINI_API_KEY)
            result = genai.embed_content(
                model="models/embedding-001",
                content=query,
                task_type="retrieval_query",
            )
            q_emb = np.array([result["embedding"]], dtype="float32")
            distances, indices = _index.search(q_emb, top_k)
            retrieved = [_chunks[i] for i in indices[0] if i < len(_chunks)]
            return "\n\n---\n\n".join(retrieved)
        except Exception:
            pass

    # Keyword fallback
    return _keyword_search(query)


def _keyword_search(query: str) -> str:
    if not _chunks:
        load_vector_store()
    q_lower = query.lower()
    scored = []
    for chunk in _chunks:
        score = sum(1 for word in q_lower.split() if word in chunk.lower())
        scored.append((score, chunk))
    scored.sort(key=lambda x: -x[0])
    top = [c for _, c in scored[:3] if _ > 0]
    return "\n\n---\n\n".join(top) if top else ""