answer.py

"""
RAG Answer Module for Persona
Retrieval pipeline with sub-query generation, semantic search, and reranking
"""
from pathlib import Path
from langchain_openai import ChatOpenAI
from langchain_chroma import Chroma
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_core.messages import SystemMessage, HumanMessage, convert_to_messages
from langchain_core.documents import Document
from langchain_core.output_parsers import CommaSeparatedListOutputParser
from langchain_core.prompts import ChatPromptTemplate
from sentence_transformers import CrossEncoder
from rank_bm25 import BM25Okapi
import numpy as np
from dotenv import load_dotenv

load_dotenv(override=True)

# Configuration
DATA_DIR = Path(__file__).parent / "data"
VECTOR_DB = str(DATA_DIR / "vector_db")
EMBEDDING_MODEL = "thenlper/gte-small"
LLM_MODEL = "gpt-4o-mini"
PERSONA_NAME = "Dmitry Kisselev"

# Retrieval parameters
RETRIEVAL_K = 20  # Retrieve candidates for reranking
FINAL_K = 5  # Return top K after reranking

USE_QUERY_EXPANSION = False  # Disabled: hurt accuracy, completeness, MRR
USE_HYBRID_SEARCH = False  # Disabled: hurt accuracy, completeness, MRR

# System prompt for persona
SYSTEM_PROMPT = """You are {PERSONA_NAME}, answering questions about yourself.
Respond naturally in first person as if you're talking about your own life, career, and experiences.
Use the context provided to answer accurately. If you don't know something, say so honestly.

Context (with metadata):
{context}
"""

# Initialize components
embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
vectorstore = None
retriever = None
llm = ChatOpenAI(temperature=0, model_name=LLM_MODEL)

# Initialize reranker
_reranker = None

def get_reranker():
    """Lazy load cross-encoder reranker"""
    global _reranker
    if _reranker is None:
        _reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
    return _reranker

# Initialize BM25 for hybrid search
_bm25 = None
_bm25_docs = None

def get_bm25():
    """Initialize BM25 index from all documents in vector store"""
    global _bm25, _bm25_docs
    if _bm25 is None:
        # Get all documents from vector store
        collection = vectorstore._collection
        all_data = collection.get(include=["documents", "metadatas"])
        
        # Create Document objects
        _bm25_docs = [
            Document(page_content=doc, metadata=meta)
            for doc, meta in zip(all_data['documents'], all_data['metadatas'])
        ]
        
        # Tokenize documents
        tokenized_docs = [doc.page_content.lower().split() for doc in _bm25_docs]
        _bm25 = BM25Okapi(tokenized_docs)
    
    return _bm25, _bm25_docs

def initialize_retriever():
    """Initialize vector store and retriever"""
    global vectorstore, retriever
    if vectorstore is None:
        vectorstore = Chroma(persist_directory=VECTOR_DB, embedding_function=embeddings)
        retriever = vectorstore.as_retriever(search_kwargs={"k": RETRIEVAL_K})
    return retriever

# Sub-query generation
output_parser = CommaSeparatedListOutputParser()

template = """
You are a helpful assistant. Given a user question, generate 1 to 3 
sub-queries that are optimized for a vector database search.
The sub-queries should cover the different parts of the user's question.

Question: {question}

Format your response as a comma-separated list.
"""
query_gen_prompt = ChatPromptTemplate.from_template(template)
query_gen_chain = query_gen_prompt | llm | output_parser

def expand_query(question: str) -> list[str]:
    """
    Query Expansion: Generate 2-3 variations of the query to improve retrieval coverage.
    """
    expansion_prompt = f"""Given this question, generate 2 alternative phrasings that would help find relevant information.
Keep the variations concise and focused on the same topic.

Original question: {question}

Provide ONLY 2 alternative phrasings, one per line, without numbering or extra text:"""
    
    try:
        response = llm.invoke([HumanMessage(content=expansion_prompt)])
        variations = [line.strip() for line in response.content.strip().split('\n') if line.strip()]
        # Return original + variations (limit to 3 total)
        return [question] + variations[:2]
    except Exception as e:
        print(f"Query expansion failed: {e}")
        return [question]

def fetch_context(question: str) -> list[Document]:
    """
    Retrieve and rerank relevant context documents.
    Uses: (Query Expansion) + Sub-query generation + Semantic search + (Hybrid Search) + Reranking.
    """
    retriever = initialize_retriever()
    
    # Query expansion
    if USE_QUERY_EXPANSION:
        expanded_queries = expand_query(question)
        base_question = expanded_queries[0]
    else:
        base_question = question
    
    # Generate sub-queries
    try:
        sub_queries = query_gen_chain.invoke({"question": base_question})
        all_queries = [base_question] + sub_queries
    except Exception as e:
        print(f"Sub-query generation failed: {e}. Using original question.")
        all_queries = [base_question]
    
    # Add expanded queries if enabled
    if USE_QUERY_EXPANSION:
        all_queries.extend(expanded_queries[1:])  # Add variations
    
    # Initialize BM25 if hybrid search is enabled
    bm25 = None
    bm25_docs = None
    if USE_HYBRID_SEARCH:
        try:
            bm25, bm25_docs = get_bm25()
        except Exception as e:
            print(f"Failed to initialize BM25: {e}")
    
    # Retrieve documents for all queries
    all_docs = []
    seen_ids = set()
    
    for q in all_queries:
        # Semantic search
        try:
            docs = retriever.invoke(q)
            for doc in docs:
                doc_id = f"{doc.metadata.get('source', '')}:{hash(doc.page_content)}"
                if doc_id not in seen_ids:
                    seen_ids.add(doc_id)
                    all_docs.append(doc)
        except Exception as e:
            print(f"Semantic retrieval failed for query '{q}': {e}")
        
        # BM25 search (if enabled)
        if USE_HYBRID_SEARCH and bm25 and bm25_docs:
            try:
                tokenized_query = q.lower().split()
                bm25_scores = bm25.get_scores(tokenized_query)
                top_bm25_indices = np.argsort(bm25_scores)[::-1][:RETRIEVAL_K]
                bm25_results = [bm25_docs[i] for i in top_bm25_indices]
                
                for doc in bm25_results:
                    doc_id = f"{doc.metadata.get('source', '')}:{hash(doc.page_content)}"
                    if doc_id not in seen_ids:
                        seen_ids.add(doc_id)
                        all_docs.append(doc)
            except Exception as e:
                print(f"BM25 retrieval failed for query '{q}': {e}")
    
    if not all_docs:
        print("No documents retrieved.")
        return []
    
    # Rerank with cross-encoder
    try:
        reranker = get_reranker()
        pairs = [[question, doc.page_content] for doc in all_docs]
        scores = reranker.predict(pairs)
        
        doc_scores = list(zip(all_docs, scores))
        doc_scores.sort(key=lambda x: x[1], reverse=True)
        top_docs = [doc for doc, score in doc_scores[:FINAL_K]]
        
        return top_docs
    except Exception as e:
        print(f"Reranking failed: {e}. Returning top documents without reranking.")
        return all_docs[:FINAL_K]

def format_doc_with_metadata(doc: Document, idx: int) -> str:
    """Format document with metadata for context"""
    meta = doc.metadata
    formatted = f"--- Document {idx+1} ---\n"
    
    # Add metadata
    if 'source' in meta:
        formatted += f"Source: {meta['source']}\n"
    if 'data_type' in meta:
        formatted += f"Type: {meta['data_type']}\n"
    if 'time_period' in meta:
        formatted += f"Time Period: {meta['time_period']}\n"
    if 'item_count' in meta:
        formatted += f"Items: {meta['item_count']}\n"
    
    # Add content
    formatted += f"\nContent:\n{doc.page_content}\n"
    return formatted

def answer_question(question: str, history: list[dict] = []) -> tuple[str, list[Document]]:
    """ Answer the given question using RAG."""
    # Fetch relevant context
    docs = fetch_context(question)
    
    # Format context with metadata
    context = "\n\n".join(format_doc_with_metadata(doc, i) for i, doc in enumerate(docs))
    
    # Build messages
    system_prompt = SYSTEM_PROMPT.format(context=context, PERSONA_NAME=PERSONA_NAME)
    messages = [SystemMessage(content=system_prompt)]
    messages.extend(convert_to_messages(history))
    messages.append(HumanMessage(content=question[:5000]))
    
    # Get response
    response = llm.invoke(messages)
    return response.content, docs

if __name__ == "__main__":
    # Test the module
    print("Testing RAG answer module...")
    test_question = "What is your current role?"
    answer, docs = answer_question(test_question)
    print(f"\nQuestion: {test_question}")
    print(f"\nAnswer: {answer}")
    print(f"\nRetrieved {len(docs)} documents")