main.py

from fastapi import FastAPI, File, UploadFile, HTTPException, Form, Depends
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional, Dict, Tuple
import tempfile
import os
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from huggingface_hub import InferenceClient
from dotenv import load_dotenv

load_dotenv()

app = FastAPI()

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

class ChatRequest(BaseModel):
    question: str
    chat_history: List[Tuple[str, str]] = []

class ChatResponse(BaseModel):
    answer: str
    sources: List[str]
    retrieved_chunks: List[Dict[str, str]]

class CustomRetrieverChat:
    def __init__(self, vector_store: FAISS, hf_token: str):
        self.vector_store = vector_store
        self.client = InferenceClient(api_key=hf_token)
    
    def get_relevant_documents(self, query: str, k: int = 3):
        return self.vector_store.similarity_search(query, k=k)
    
    def format_prompt(self, question: str, context: str) -> str:
        return f"""You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.

        Context: {context}

        Question: {question}
        
        Answer:"""
    
    def generate_response(self, question: str) -> tuple[str, List[Dict]]:

        docs = self.get_relevant_documents(question)
        
        context = "\n\n".join([doc.page_content for doc in docs])
        
        prompt = self.format_prompt(question, context)
        
        messages = [{"role": "user", "content": prompt}]
        try:
            completion = self.client.chat.completions.create(
                model="meta-llama/Llama-3.2-3B-Instruct",
                messages=messages,
                max_tokens=500
            )
            answer = completion.choices[0].message.content
        except Exception as e:
            print(f"Error in LLM call: {e}")
            answer = "I apologize, but I encountered an error processing your request."
        
        retrieved_chunks = [
            {
                "content": doc.page_content,
                "source": os.path.basename(doc.metadata.get('source', 'Unknown Source'))
            }
            for doc in docs
        ]
        
        return answer, retrieved_chunks

vector_stores = {}
conversations = {}
hf_token = os.getenv('HF_TOKEN') 

def process_file(file_path: str) -> List:
    file_extension = os.path.splitext(file_path)[1].lower()

    if file_extension == '.pdf':
        loader = PyPDFLoader(file_path)
    elif file_extension == '.txt':
        loader = TextLoader(file_path)
    else:
        raise ValueError(f"Unsupported file type: {file_extension}")

    return loader.load()

def initialize_conversation(files: List[UploadFile]):
    all_documents = []
    temp_files = []
    
    try:
        for uploaded_file in files:
            suffix = os.path.splitext(uploaded_file.filename)[1].lower()
            with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:
                content = uploaded_file.file.read()
                tmp_file.write(content)
                temp_files.append(tmp_file.name)
                
                documents = process_file(tmp_file.name)
                all_documents.extend(documents)

        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len,
        )
        document_chunks = text_splitter.split_documents(all_documents)
        
        embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-mpnet-base-v2"
        )

        vector_store = FAISS.from_documents(document_chunks, embeddings)
        
        conversation = CustomRetrieverChat(vector_store, hf_token)
        
        return conversation
        
    finally:
        for temp_file in temp_files:
            try:
                os.unlink(temp_file)
            except Exception as e:
                print(f"Error removing temporary file: {e}")

@app.post("/upload")
async def upload_files(files: List[UploadFile] = File(...)):
    if len(files) > 5:
        raise HTTPException(status_code=400, detail="Maximum 5 files allowed")

    try:
        conversation = initialize_conversation(files)
        import uuid
        session_id = str(uuid.uuid4())
        conversations[session_id] = conversation

        return {"session_id": session_id, "message": "Files processed successfully"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/chat/{session_id}", response_model=ChatResponse)
async def chat(
    session_id: str,
    question: str = Form(None),
    chat_request: Optional[ChatRequest] = None
):
    if session_id not in conversations:
        raise HTTPException(status_code=404, detail="Session not found")

    try:
        conversation = conversations[session_id]
        
        if question is not None:
            query = question
        elif chat_request is not None:
            query = chat_request.question
        else:
            raise HTTPException(status_code=400, detail="No question provided")

        answer, retrieved_chunks = conversation.generate_response(query)
        
        sources = list(set(chunk["source"] for chunk in retrieved_chunks))

        return ChatResponse(
            answer=answer,
            sources=sources,
            retrieved_chunks=retrieved_chunks
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/session/{session_id}")
async def delete_session(session_id: str):
    if session_id in conversations:
        del conversations[session_id]
        return {"message": "Session deleted successfully"}
    raise HTTPException(status_code=404, detail="Session not found")

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