app.py

import os
import logging
import warnings
from typing import Optional, List

from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse, RedirectResponse
from fastapi.concurrency import run_in_threadpool
from pydantic import BaseModel
import uvicorn

# LangChain & Ollama Imports
from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader, TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaLLM, OllamaEmbeddings
from langchain_community.vectorstores import FAISS

# Suppress Warnings
warnings.filterwarnings("ignore", category=UserWarning, module="langchain_core._api.deprecation")

# Configure Logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# --- Configuration ---
DATA_PATH = "data/"
DB_FAISS_PATH = "vectorstore/db_faiss"
MODEL_NAME = "llama3"
EMBEDDING_MODEL = "nomic-embed-text"

DOCUMIND_SYSTEM_PROMPT = """You are DocuMind, a helpful and knowledgeable document assistant.

Your goal is to answer the user's questions based on the provided document context.

────────────────────────────────────
INSTRUCTIONS
────────────────────────────────────
1. **Prioritize the Document**: Use the provided document context as your primary source of truth.
2. **Be Helpful**: If the document mentions a term but doesn't define it fully, you MAY use your general knowledge to briefly explain it.
3. **Admit Gaps**: If the answer is completely missing from the document, state "I couldn't find that specific information in the document."
4. **Context Support**: Use the provided chat history to understand follow-up questions.

Context:
{context}

Chat History:
{history}

Question: {question}

Answer:"""

# --- Logic: Ingestion ---
def create_vector_db():
    """
    Ingests PDF and TXT documents, splits them, generates embeddings, and saves to FAISS.
    """
    if not os.path.exists(DATA_PATH):
        os.makedirs(DATA_PATH)
        logger.info(f"Directory {DATA_PATH} created.")
        return False, "Data directory created. Please add files."

    logger.info("Loading documents from storage...")
    try:
        pdf_loader = DirectoryLoader(DATA_PATH, glob='*.pdf', loader_cls=PyPDFLoader)
        txt_loader = DirectoryLoader(DATA_PATH, glob='*.txt', loader_cls=TextLoader)
        documents = pdf_loader.load() + txt_loader.load()
    except Exception as e:
        logger.error(f"Failed to load documents: {e}")
        return False, str(e)

    if not documents:
        logger.warning(f"No compatible documents found in {DATA_PATH}.")
        return False, "No documents found."

    # Split text
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=150)
    texts = text_splitter.split_documents(documents)
    logger.info(f"Processed {len(documents)} documents into {len(texts)} chunks.")

    # Embed and Store
    try:
        logger.info(f"Initializing embeddings model ({EMBEDDING_MODEL})...")
        embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)

        logger.info("Building vector index...")
        db = FAISS.from_documents(texts, embeddings)
        db.save_local(DB_FAISS_PATH)
        logger.info(f"Vector store successfully persisted to {DB_FAISS_PATH}")
        return True, "Ingestion successful!"
    except Exception as e:
        logger.error(f"Failed to create vector store: {e}")
        return False, str(e)

# --- Logic: Retrieval & RAG ---
class DocuMindError(Exception):
    pass

class DocuMind:
    def __init__(self):
        self.db = None
        self.llm = OllamaLLM(model=MODEL_NAME)
        self.chat_history = []
        self._initialize_vector_store()

    def _initialize_vector_store(self):
        print(f"Initializing embeddings ({EMBEDDING_MODEL})...")
        embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
        try:
            if os.path.exists(DB_FAISS_PATH):
                self.db = FAISS.load_local(DB_FAISS_PATH, embeddings, allow_dangerous_deserialization=True)
                print("Vector store loaded successfully.")
            else:
                print("Warning: Vector store not found. Please run ingestion.")
                self.db = None
        except Exception as e:
            print(f"Error loading vector store: {e}")
            self.db = None

    def get_response(self, query: str) -> str:
        if not self.db:
            return "System Error: Vector store not loaded. Please ingest documents first."

        try:
            # Increase k to 5 for better context
            docs = self.db.similarity_search(query, k=5)
            context = "\n---\n".join([doc.page_content for doc in docs])
        except Exception as e:
            return f"Retrieval Error: {e}"
        
        # Format History (last 3 turns)
        history_str = "\n".join([f"User: {q}\nDocuMind: {a}" for q, a in self.chat_history[-3:]])

        prompt = DOCUMIND_SYSTEM_PROMPT.format(context=context, history=history_str, question=query)
        
        try:
            response = self.llm.invoke(prompt)
            self.chat_history.append((query, response.strip()))
            return response.strip()
        except Exception as e:
            return f"Error calling LLM: {e}"

# --- FastAPI App ---
app = FastAPI()

app.mount("/static", StaticFiles(directory="static"), name="static")

# Initialize Chatbot Global Instance
bot = DocuMind()

class ChatRequest(BaseModel):
    message: str

@app.post("/api/chat")
async def chat(request: ChatRequest):
    global bot
    if bot is None or bot.db is None:
        # Try to re-initialize
        try:
            bot = DocuMind()
        except:
            pass
        
        if bot is None or bot.db is None:
             return JSONResponse(status_code=503, content={"message": "System not ready. Please ingest documents first."})

    # Run blocking inference in threadpool
    response = await run_in_threadpool(bot.get_response, request.message)
    return {"response": response}

@app.post("/api/ingest")
async def ingest_endpoint():
    global bot
    try:
        # Run blocking ingestion in threadpool
        success, message = await run_in_threadpool(create_vector_db)
        
        if not success:
             return JSONResponse(status_code=500, content={"message": f"Ingestion failed: {message}"})
        
        # Re-initialize bot to load new vector store
        bot = DocuMind()
        return {"message": "Ingestion successful! You can now ask questions."}
    except Exception as e:
        return JSONResponse(status_code=500, content={"message": str(e)})

@app.get("/")
async def root():
    return RedirectResponse(url="/static/index.html")

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