RAG.py

import os
import warnings
import pickle
from together import Together
import faiss
from sentence_transformers import SentenceTransformer
from PyPDF2 import PdfReader
import glob

warnings.filterwarnings("ignore")

# Hardcoded Together API Key
TOGETHER_API_KEY = "81da53aa3044c7ebead342fb048f016a4e593a86928a783a6fdcc1e3883054e4"
client = Together(api_key=TOGETHER_API_KEY)

# Initialize embedding model
embedding_model = SentenceTransformer(
    "sentence-transformers/all-MiniLM-L6-v2",
    use_auth_token=os.environ.get("HUGGINGFACE_HUB_TOKEN"),
)

# Global variables for storing embeddings and documents
documents = []
filenames = []
index = None

def load_index():
    """Load the FAISS index and document metadata."""
    global index, documents, filenames
    
    if not os.path.exists("knowledge_base/faiss_index.bin") or not os.path.exists("knowledge_base/metadata.pkl"):
        raise ValueError("Index files not found. Please run preprocess.py first!")
    
    print("Loading index and metadata...")
    index = faiss.read_index("knowledge_base/faiss_index.bin")
    
    with open("knowledge_base/metadata.pkl", "rb") as f:
        metadata = pickle.load(f)
        documents = metadata["documents"]
        filenames = metadata["filenames"]
    
    print("Index and metadata loaded successfully!")

def answer_question(query: str) -> str:
    """
    Answer a question using the RAG system.
    
    Args:
        query (str): The user's question
        
    Returns:
        str: The generated answer
    """
    global index
    
    # Load index if not already loaded
    if index is None:
        load_index()
    
    # Get query embedding
    query_embedding = embedding_model.encode([query])
    faiss.normalize_L2(query_embedding)
    
    # Get top similar documents
    scores, indices = index.search(query_embedding, min(3, len(documents)))
    
    # Build context from retrieved documents
    context_parts = []
    relevant_docs = []
    
    for score, idx in zip(scores[0], indices[0]):
        if idx < len(documents):
            doc_info = {
                "content": documents[idx],
                "filename": filenames[idx],
                "score": float(score),
            }
            relevant_docs.append(doc_info)
            context_parts.append(f"[{doc_info['filename']}]\n{doc_info['content']}")
    
    if not relevant_docs:
        return "I couldn't find any relevant information to answer your question."
    
    # Combine context
    context = "\n\n".join(context_parts)
    
    # Create prompt for LLM
    llm_prompt = f"""Answer the question based on the provided context documents.

Context:
{context}

Question: {query}

Instructions:
- Answer based only on the information in the context
- If the context doesn't contain enough information, say so
- Mention which document(s) you're referencing
- Start with According to [document name]
- Add brackets to the document name

Answer:"""
    
    try:
        # Generate answer using Together AI
        response = client.chat.completions.create(
            model="meta-llama/Llama-3.3-70B-Instruct-Turbo",
            messages=[{"role": "user", "content": llm_prompt}],
            max_tokens=500,
            temperature=0.7,
        )
        answer = response.choices[0].message.content
        
        # Add source information
        sources_list = [doc["filename"] for doc in relevant_docs]
        sources_text = sources_list[0]
        full_answer = f"{answer}\n\n📄 Source Used: {sources_text}"
        
        return full_answer
        
    except Exception as e:
        return f"Error generating answer: {str(e)}"