rag.py

import os
from dotenv import load_dotenv
import re
import pickle
import faiss
import numpy as np
from typing import List, Dict
from sentence_transformers import SentenceTransformer, CrossEncoder, util
from rank_bm25 import BM25Okapi
import nltk
from nltk.corpus import stopwords
import requests
import json
from openai import OpenAI
import logging
#import generate_indexes

load_dotenv()
#generate_indexes.main()

# ---------------- Logging Setup ----------------
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s %(levelname)s %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

nltk.download("stopwords")
STOPWORDS = set(stopwords.words("english"))

os.environ["TOKENIZERS_PARALLELISM"] = "false"

# ---------------- Paths & Models ----------------
EMBED_MODEL   = "sentence-transformers/all-MiniLM-L6-v2"
CROSS_ENCODER = "cross-encoder/ms-marco-MiniLM-L-6-v2"
OUT_DIR       = "data/index_merged"

FAISS_PATH = os.path.join(OUT_DIR, "faiss_merged.index")
BM25_PATH  = os.path.join(OUT_DIR, "bm25_merged.pkl")
META_PATH  = os.path.join(OUT_DIR, "meta_merged.pkl")

# ---------------- Load Indexes ----------------
logger.info("Loading FAISS, BM25, metadata, and models...")
try:
    faiss_index = faiss.read_index(FAISS_PATH)
    with open(BM25_PATH, "rb") as f:
        bm25_obj = pickle.load(f)
    bm25 = bm25_obj["bm25"]
    with open(META_PATH, "rb") as f:
        meta: List[Dict] = pickle.load(f)
    embed_model = SentenceTransformer(EMBED_MODEL)
    reranker = CrossEncoder(CROSS_ENCODER)
    api_key = os.getenv("HF_API_KEY")
    if not api_key:
        logger.error("HF_API_KEY environment variable not set. Please check your .env file or environment.")
        raise ValueError("HF_API_KEY environment variable not set.")
    client = OpenAI(
        base_url="https://router.huggingface.co/v1",
        api_key=api_key
    )
except Exception as e:
    logger.error(f"Error loading models or indexes: {e}")
    raise

def get_mistral_answer(query: str, context: str) -> str:
    """
    Calls Mistral 7B Instruct API via Hugging Face Inference API.
    Adds error handling and logging.
    """
    prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer in full sentences using context."
    try:
        logger.info(f"Calling Mistral API for query: {query}")
        completion = client.chat.completions.create(
            model="dphn/Dolphin-Mistral-24B-Venice-Edition:featherless-ai",
            messages=[
                {
                    "role": "user",
                    "content": prompt
                }
            ]
        )
        answer = str(completion.choices[0].message.content)
        logger.info(f"Mistral API response: {answer}")
        return answer
    except Exception as e:
        logger.error(f"Error in Mistral API call: {e}")
        return f"Error fetching answer from LLM: {e}"

# ---------------- Guardrails ----------------
BLOCKED_TERMS = ["weather", "cricket", "movie", "song", "football", "holiday",
                 "travel", "recipe", "music", "game", "sports", "politics", "election"]

FINANCE_DOMAINS = [
    "financial reporting", "balance sheet", "income statement",
    "assets and liabilities", "equity", "revenue", "profit and loss",
    "goodwill impairment", "cash flow", "dividends", "taxation",
    "investment", "valuation", "capital structure", "ownership interests",
    "subsidiaries", "shareholders equity", "expenses", "earnings",
    "debt", "amortization", "depreciation"
]
finance_embeds = embed_model.encode(FINANCE_DOMAINS, convert_to_tensor=True)

def validate_query(query: str, threshold: float = 0.5) -> bool:
    q_lower = query.lower()
    if any(bad in q_lower for bad in BLOCKED_TERMS):
        print("[Guardrail] Rejected by blocklist.")
        return False
    q_emb = embed_model.encode(query, convert_to_tensor=True)
    sim_scores = util.cos_sim(q_emb, finance_embeds)
    max_score = float(sim_scores.max())
    if max_score > threshold:
        print(f"[Guardrail] Accepted (semantic match {max_score:.2f})")
        return True
    else:
        print(f"[Guardrail] Rejected (low semantic score {max_score:.2f})")
        return False

# ---------------- Preprocess ----------------
def preprocess_query(query: str, remove_stopwords: bool = True) -> str:
    query = query.lower()
    query = re.sub(r"[^a-z0-9\s]", " ", query)
    tokens = query.split()
    if remove_stopwords:
        tokens = [t for t in tokens if t not in STOPWORDS]
    return " ".join(tokens)

# ---------------- Hybrid Retrieval ----------------
def hybrid_candidates(query: str, candidate_k: int = 50, alpha: float = 0.5) -> List[int]:
    q_emb = embed_model.encode([preprocess_query(query, remove_stopwords=False)], convert_to_numpy=True, normalize_embeddings=True)
    faiss_scores, faiss_ids = faiss_index.search(q_emb, max(candidate_k, 50))
    faiss_ids = faiss_ids[0]
    faiss_scores = faiss_scores[0]

    tokenized_query = preprocess_query(query).split()
    bm25_scores = bm25.get_scores(tokenized_query)

    topN = max(candidate_k, 50)
    bm25_top = np.argsort(bm25_scores)[::-1][:topN]
    faiss_top = faiss_ids[:topN]
    union_ids = np.unique(np.concatenate([bm25_top, faiss_top]))

    faiss_score_map = {int(i): float(s) for i, s in zip(faiss_ids, faiss_scores)}
    f_arr = np.array([faiss_score_map.get(int(i), -1.0) for i in union_ids], dtype=float)
    f_min = np.min(f_arr)
    if np.any(f_arr < 0):
        f_arr = np.where(f_arr < 0, f_min, f_arr)
    b_arr = np.array([bm25_scores[int(i)] for i in union_ids], dtype=float)

    def _norm(x): return (x - np.min(x)) / (np.ptp(x) + 1e-9)
    combined = alpha * _norm(f_arr) + (1 - alpha) * _norm(b_arr)
    order = np.argsort(combined)[::-1]
    return union_ids[order][:candidate_k].tolist()

# ---------------- Cross-Encoder Rerank ----------------
def rerank_cross_encoder(query: str, cand_ids: List[int], top_k: int = 10) -> List[Dict]:
    pairs = [(query, meta[i]["content"]) for i in cand_ids]
    scores = reranker.predict(pairs)
    order = np.argsort(scores)[::-1][:top_k]
    return [{"id": cand_ids[i], "chunk_size": meta[cand_ids[i]]["chunk_size"], "content": meta[cand_ids[i]]["content"], "rerank_score": float(scores[i])} for i in order]

# ---------------- Extract Numeric ----------------
def extract_value_for_year_and_concept(year: str, concept: str, context_docs: List[Dict]) -> str:
    target_year = str(year)
    concept_lower = concept.lower()
    for doc in context_docs:
        text = doc.get("content", "")
        lines = [line for line in text.split("\n") if line.strip() and any(c.isdigit() for c in line)]
        header_idx = None
        year_to_col = {}
        for idx, line in enumerate(lines):
            years_in_line = re.findall(r"20\d{2}", line)
            if years_in_line:
                for col_idx, y in enumerate(years_in_line):
                    year_to_col[y] = col_idx
                header_idx = idx
                break
        if target_year not in year_to_col or header_idx is None:
            continue
        for line in lines[header_idx+1:]:
            if concept_lower in line.lower():
                cols = re.split(r"\s{2,}|\t", line)
                col_idx = year_to_col[target_year]
                if col_idx < len(cols):
                    return cols[col_idx].replace(",", "")
    return ""

# ---------------- RAG Pipeline ----------------
def generate_answer(query: str, top_k: int = 5, candidate_k: int = 50, alpha: float = 0.6):
    logger.info(f"Received query: {query}")
    try:
        if not validate_query(query):
            logger.warning("Query rejected: Not finance-related.")
            return "Query rejected: Please ask finance-related questions."

        cand_ids = hybrid_candidates(query, candidate_k=candidate_k, alpha=alpha)
        logger.info(f"Hybrid candidates retrieved: {cand_ids}")
        reranked = rerank_cross_encoder(query, cand_ids, top_k=top_k)
        logger.info(f"Reranked top docs: {[d['id'] for d in reranked]}")

        year_match = re.search(r"(20\d{2})", query)
        year = year_match.group(0) if year_match else None
        concept = re.sub(r"for the year 20\d{2}", "", query, flags=re.IGNORECASE).strip()

        year_specific_answer = None
        if year and concept:
            year_specific_answer = extract_value_for_year_and_concept(year, concept, reranked)
            logger.info(f"Year-specific answer: {year_specific_answer}")

        if year_specific_answer:
            answer = year_specific_answer
        else:
            # Pass top 5 chunks as context
            context_text = "\n".join([d["content"] for d in reranked])
            answer = get_mistral_answer(query, context_text)
        final_answer = answer
        logger.info(f"Final Answer: {final_answer}")
        return final_answer
    except Exception as e:
        logger.error(f"Error in RAG pipeline: {e}")
        return f"Error in RAG pipeline: {e}"