main.py

from fastapi import FastAPI, File, UploadFile, HTTPException, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, HTMLResponse, FileResponse
from pydantic import BaseModel
from typing import Optional, List, Dict
import os
import tempfile
import shutil
from pathlib import Path
import uuid
from datetime import datetime

from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader, PyPDFLoader
from langchain.schema import Document
from langchain.chains import ConversationalRetrievalChain
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory, ConversationSummaryMemory
from langchain_core.messages import HumanMessage, AIMessage
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain.chains.question_answering import load_qa_chain
from langchain.chains import LLMChain
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

app = FastAPI(title="RAG System with Multiple Retrievers")

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global variables to store vectorstore and conversation sessions
vectorstore = None
documents = None
conversation_sessions = {}  # Store conversation memory per session

# Request/Response Models
class QueryRequest(BaseModel):
    question: str
    retriever_type: str = "similarity"  # similarity, mmr, mqr
    k: int = 4
    fetch_k: int = 20
    lambda_mult: float = 0.5
    session_id: Optional[str] = None  # Session ID for conversation tracking
    use_memory: bool = True  # Whether to use conversation memory

class QueryResponse(BaseModel):
    answer: str
    sources: List[dict]
    session_id: str
    conversation_history: Optional[List[dict]] = None

class ConversationHistory(BaseModel):
    session_id: str
    messages: List[dict]

# Initialize Google API
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")

if not GOOGLE_API_KEY:
    print("Warning: GOOGLE_API_KEY not set. Please set it in environment variables.")

embeddings = GoogleGenerativeAIEmbeddings(
    model="models/gemini-embedding-001",
    google_api_key=GOOGLE_API_KEY
)

llm = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash",
    google_api_key=GOOGLE_API_KEY,
    temperature=0.7
)

def get_or_create_session(session_id: Optional[str] = None) -> str:
    """Get existing session or create new one"""
    if session_id and session_id in conversation_sessions:
        return session_id
    
    # Create new session
    new_session_id = str(uuid.uuid4())
    conversation_sessions[new_session_id] = {
        "memory": ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True,
            output_key="answer"
        ),
        "created_at": datetime.now().isoformat(),
        "last_accessed": datetime.now().isoformat()
    }
    return new_session_id

def get_conversation_history(session_id: str) -> List[dict]:
    """Get conversation history for a session"""
    if session_id not in conversation_sessions:
        return []
    
    memory = conversation_sessions[session_id]["memory"]
    messages = memory.chat_memory.messages
    
    history = []
    for msg in messages:
        if isinstance(msg, HumanMessage):
            history.append({"role": "user", "content": msg.content})
        elif isinstance(msg, AIMessage):
            history.append({"role": "assistant", "content": msg.content})
    
    return history

def clear_session_memory(session_id: str):
    """Clear memory for a specific session"""
    if session_id in conversation_sessions:
        conversation_sessions[session_id]["memory"].clear()

def delete_session(session_id: str):
    """Delete a session completely"""
    if session_id in conversation_sessions:
        del conversation_sessions[session_id]

def load_document(file_path: str, file_type: str) -> List[Document]:
    """Load document based on file type"""
    try:
        if file_type == "pdf":
            loader = PyPDFLoader(file_path)
        else:  # txt
            loader = TextLoader(file_path)
        
        documents = loader.load()
        return documents
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error loading document: {str(e)}")

def chunk_documents(documents: List[Document], chunk_size: int, chunk_overlap: int) -> List[Document]:
    """Split documents into chunks"""
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
        separators=["\n\n", "\n", " ", ""]
    )
    
    chunks = text_splitter.split_documents(documents)
    return chunks

@app.get("/", response_class=HTMLResponse)
async def root():
    """Serve the index.html file"""
    html_path = Path(__file__).parent / "index.html"
    if html_path.exists():
        with open(html_path, "r", encoding="utf-8") as f:
            return f.read()
    return "<h1>RAG System API is running</h1><p>index.html not found</p>"

@app.post("/upload")
async def upload_document(
    file: UploadFile = File(...),
    chunk_size: int = Form(1000),
    chunk_overlap: int = Form(200)
):
    """Upload and process document"""
    global vectorstore, documents
    
    try:
        # Validate file type
        file_extension = file.filename.split(".")[-1].lower()
        if file_extension not in ["txt", "pdf"]:
            raise HTTPException(status_code=400, detail="Only .txt and .pdf files are supported")
        
        # Save uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_extension}") as tmp_file:
            shutil.copyfileobj(file.file, tmp_file)
            tmp_file_path = tmp_file.name
        
        # Load document
        documents = load_document(tmp_file_path, file_extension)
        
        # Chunk documents
        chunks = chunk_documents(documents, chunk_size, chunk_overlap)
        
        # Create vectorstore
        vectorstore = FAISS.from_documents(chunks, embeddings)
        
        # Clean up temporary file
        os.unlink(tmp_file_path)
        
        return {
            "message": "Document uploaded and processed successfully",
            "filename": file.filename,
            "num_chunks": len(chunks),
            "chunk_size": chunk_size,
            "chunk_overlap": chunk_overlap
        }
    
    except Exception as e:
        # Clean up on error
        if 'tmp_file_path' in locals():
            try:
                os.unlink(tmp_file_path)
            except:
                pass
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/query", response_model=QueryResponse)
async def query_documents(request: QueryRequest):
    """Query documents using selected retriever with conversation memory"""
    global vectorstore
    
    if vectorstore is None:
        raise HTTPException(status_code=400, detail="No document uploaded. Please upload a document first.")
    
    try:
        # Get or create session
        session_id = get_or_create_session(request.session_id)
        
        # Get retriever based on type
        if request.retriever_type == "similarity":
            retriever = vectorstore.as_retriever(
                search_type="similarity",
                search_kwargs={"k": request.k}
            )
        elif request.retriever_type == "mmr":
            # MMR (Maximal Marginal Relevance)
            retriever = vectorstore.as_retriever(
                search_type="mmr",
                search_kwargs={
                    "k": request.k,
                    "fetch_k": request.fetch_k,
                    "lambda_mult": request.lambda_mult
                }
            )
        elif request.retriever_type == "mqr":
            # Multi-Query Retriever - generates multiple query variations
            base_retriever = vectorstore.as_retriever(
                search_type="similarity",
                search_kwargs={"k": request.k}
            )
            retriever = MultiQueryRetriever.from_llm(
                retriever=base_retriever,
                llm=llm
            )
        else:
            raise HTTPException(status_code=400, detail="Invalid retriever type")
        
        if request.use_memory:
            # Use conversational retrieval chain with memory
            memory = conversation_sessions[session_id]["memory"]
            
            # Update last accessed time
            conversation_sessions[session_id]["last_accessed"] = datetime.now().isoformat()
            
            # Get chat history
            chat_history = memory.chat_memory.messages
            
            # Format chat history as string
            formatted_history = ""
            for msg in chat_history:
                if isinstance(msg, HumanMessage):
                    formatted_history += f"Human: {msg.content}\n"
                elif isinstance(msg, AIMessage):
                    formatted_history += f"Assistant: {msg.content}\n"
            
            # Retrieve relevant documents
            try:
                docs = retriever.get_relevant_documents(request.question)
                source_docs = docs[:request.k]
            except Exception as e:
                source_docs = []
            
            # Format document context
            doc_context = "\n\n".join([doc.page_content for doc in source_docs]) if source_docs else "No relevant documents found."
            
            # Create prompt that includes BOTH chat history and document context
            final_prompt = f"""You are a helpful AI assistant. Answer the question using:
1. Information from the conversation history (if relevant)
2. Information from the provided document context (if relevant)

If the question can be answered from the conversation history, prioritize that.
If you need information from the documents, use the document context.
If you don't know the answer from either source, say so.

Conversation History:
{formatted_history}

Document Context:
{doc_context}

Current Question: {request.question}

Answer (be concise and helpful):"""
            
            # Get answer from LLM
            response = llm.invoke(final_prompt)
            answer = response.content
            
            # Save to memory
            memory.save_context(
                {"question": request.question},
                {"answer": answer}
            )
            
        else:
            # Use without memory (stateless)
            from langchain.chains import RetrievalQA
            
            template = """Use the following pieces of context to answer the question at the end. 
If you don't know the answer, just say that you don't know, don't try to make up an answer.

Context: {context}

Question: {question}

Answer: """
            
            QA_CHAIN_PROMPT = PromptTemplate(
                input_variables=["context", "question"],
                template=template,
            )
            
            qa_chain = RetrievalQA.from_chain_type(
                llm=llm,
                chain_type="stuff",
                retriever=retriever,
                return_source_documents=True,
                chain_type_kwargs={"prompt": QA_CHAIN_PROMPT}
            )
            
            result = qa_chain({"query": request.question})
            answer = result["result"]
            source_docs = result["source_documents"][:request.k]
        
        # Format sources
        sources = []
        for i, doc in enumerate(source_docs):
            sources.append({
                "content": doc.page_content[:200] + "..." if len(doc.page_content) > 200 else doc.page_content,
                "metadata": doc.metadata,
                "index": i + 1
            })
        
        # Get conversation history
        history = get_conversation_history(session_id) if request.use_memory else None
        
        return QueryResponse(
            answer=answer,
            sources=sources,
            session_id=session_id,
            conversation_history=history
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error processing query: {str(e)}")

@app.delete("/clear")
async def clear_documents():
    """Clear uploaded documents"""
    global vectorstore, documents
    
    vectorstore = None
    documents = None
    
    return {"message": "Documents cleared successfully"}

@app.get("/sessions")
async def list_sessions():
    """List all active conversation sessions"""
    sessions = []
    for session_id, data in conversation_sessions.items():
        sessions.append({
            "session_id": session_id,
            "created_at": data["created_at"],
            "last_accessed": data["last_accessed"],
            "message_count": len(data["memory"].chat_memory.messages)
        })
    return {"sessions": sessions}

@app.get("/sessions/{session_id}/history")
async def get_session_history(session_id: str):
    """Get conversation history for a specific session"""
    if session_id not in conversation_sessions:
        raise HTTPException(status_code=404, detail="Session not found")
    
    history = get_conversation_history(session_id)
    return {
        "session_id": session_id,
        "history": history,
        "message_count": len(history)
    }

@app.delete("/sessions/{session_id}/clear")
async def clear_session_history(session_id: str):
    """Clear conversation history for a specific session"""
    if session_id not in conversation_sessions:
        raise HTTPException(status_code=404, detail="Session not found")
    
    clear_session_memory(session_id)
    return {"message": f"Session {session_id} history cleared"}

@app.delete("/sessions/{session_id}")
async def delete_session_endpoint(session_id: str):
    """Delete a conversation session"""
    if session_id not in conversation_sessions:
        raise HTTPException(status_code=404, detail="Session not found")
    
    delete_session(session_id)
    return {"message": f"Session {session_id} deleted"}

@app.delete("/sessions")
async def clear_all_sessions():
    """Clear all conversation sessions"""
    global conversation_sessions
    conversation_sessions = {}
    return {"message": "All sessions cleared"}

@app.get("/status")
async def get_status():
    """Get system status"""
    return {
        "documents_loaded": vectorstore is not None,
        "api_key_set": bool(GOOGLE_API_KEY),
        "active_sessions": len(conversation_sessions),
        "total_messages": sum(len(s["memory"].chat_memory.messages) for s in conversation_sessions.values())
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)