Spaces:

YADUUG
/

NPS-ChatAgentBot

Sleeping

Base 1.0

4e06012 9 months ago

17.1 kB

import os, logging
import glob
import logging
import re
import hashlib
import numpy as np
import asyncio
from typing import List, Optional, Any
import shutil
import secrets

# --- Pydantic & FastAPI ---
from pydantic import BaseModel, Field
import uvicorn
from fastapi import FastAPI, HTTPException, APIRouter, UploadFile, File, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.security import HTTPBasic, HTTPBasicCredentials

# --- LangChain & RAG Core ---
from PyPDF2 import PdfReader

try:
    from docx import Document as DocxDocument
except ImportError:
    DocxDocument = None
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain.retrievers import EnsembleRetriever
from langchain.chains.combine_documents import (
    create_stuff_documents_chain,
)  # <-- DEPRECATION FIX
from langchain.chains import create_retrieval_chain
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.llms.base import LLM
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain.docstore.document import Document
from langchain_core.retrievers import BaseRetriever
from langchain.prompts import PromptTemplate

# --- Scikit-Learn for TF-IDF ---
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# --- TruLens for Evaluation ---
# from trulens_eval import Feedback, Tru, TruChain
# from trulens_eval.feedback.provider import LiteLLM

# tru = Tru()

# --- Mistral AI ---
from mistralai import Mistral as MistralClient
from dotenv import load_dotenv

load_dotenv()

# --- Global Variables & Configuration ---
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# Define the original, read-only data folder from the repo
REPO_DATA_FOLDER = "data"

# Use persistent storage if available (on Hugging Face Spaces)
if os.path.exists("/data"):
    logger.info("Persistent storage '/data' found. Using it for PDFs and cache.")
    DATA_FOLDER = "/data/pdf_storage"
    CACHE_DIR = "/data/cache"
else:
    logger.info(
        "No persistent storage '/data' found. Using local 'data' folder and '/tmp/cache'."
    )
    DATA_FOLDER = REPO_DATA_FOLDER
    CACHE_DIR = "/tmp/cache"

os.makedirs(DATA_FOLDER, exist_ok=True)
os.makedirs(CACHE_DIR, exist_ok=True)
retrieval_chain = None

HF_CACHE = "/tmp/hf"
os.environ["HF_HOME"] = HF_CACHE
os.environ["HF_HOME"] = HF_CACHE
os.environ["HF_DATASETS_CACHE"] = HF_CACHE
os.environ["HF_MODULES_CACHE"] = HF_CACHE

retrieval_chain = None

db_path = "/data/default.sqlite" if os.path.exists("/data") else "default.sqlite"

# ✅ Only run the dashboard locally (not in Hugging Face Spaces)
if os.environ.get("SPACE_ID") is None:
    # tru.run_dashboard()
    pass


# --- Pydantic Models ---
class QuestionRequest(BaseModel):
    question: str = Field(..., min_length=1)


class AnswerResponse(BaseModel):
    answer: str
    sources: List[str]


# --- Greeting Handler ---
GREETING_RESPONSES = {
    "hi": "Hello! How can I assist you?",
    "hello": "Hello! How can I help you today?",
    "bye": "Bye-bye! See you soon.",
    "ty": "Your Welcome!",
}

# --- Enhanced Prompt ---
UNIVERSITY_ASSISTANT_PROMPT = PromptTemplate.from_template(
    """
You are a meticulous Guru Nanak Dev University Academic Assistant.
Your task is to provide accurate, specific, and concise answers based ONLY on the provided Context.
**Instructions:**
1.  **Analyze the Context:** Carefully read the user's question and the provided context to synthesize an answer.
2.  **Extract Specifics:** If asked for a fee, number, or date, provide the EXACT information.
3.  **Cite Sources:** After your answer, cite the source document and page number: `(Source: Prospectus.pdf, Page: 15)`.
4.  **Handle Missing Information:** If the info is not in the context, state: "Based on the provided documents, I could not find specific information about this topic."
---
**Context:**
{context}
---
**Question:**
{input}
**Answer (with citations):**
"""
)


# --- Custom Scikit-Learn TF-IDF Retriever ---
class ScikitLearnTfidfRetriever(BaseRetriever):
    docs: List[Document]
    tfidf_vectorizer: TfidfVectorizer
    tfidf_matrix: Any
    k: int = 15

    class Config:
        arbitrary_types_allowed = True

    def _get_relevant_documents(self, query: str) -> List[Document]:
        query_vector = self.tfidf_vectorizer.transform([query])
        similarity_scores = cosine_similarity(query_vector, self.tfidf_matrix).flatten()
        top_k_indices = np.argsort(similarity_scores)[-self.k :][::-1]
        return [self.docs[i] for i in top_k_indices]


# --- MistralLLM Class ---
class MistralLLM(LLM):
    client: Any = None
    model: str = "mistral-small-latest"
    temperature: float = 0.0
    api_key: str = ""

    def __init__(
        self,
        api_key: str,
        model: str = "mistral-small-latest",
        temperature: float = 0.0,
    ):
        super().__init__()
        self.api_key = api_key
        self.model = model
        self.temperature = temperature
        self.client = MistralClient(api_key=self.api_key)

    @property
    def _llm_type(self) -> str:
        return "mistral"

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
        **kwargs: Any,
    ) -> str:
        messages = [{"role": "user", "content": prompt}]
        response = self.client.chat.complete(
            model=self.model, messages=messages, temperature=self.temperature
        )
        return response.choices[0].message.content.strip()


# --- Caching and Data Processing Functions ---
def get_data_hash(folder_path=DATA_FOLDER) -> str:
    """Generates a hash based on the filenames and modification times of supported files."""
    hasher = hashlib.md5()
    if not os.path.exists(folder_path):
        return "no_data"

    supported_extensions = [".pdf", ".txt", ".md", ".docx"]
    files_to_hash = []
    for ext in supported_extensions:
        files_to_hash.extend(glob.glob(os.path.join(folder_path, f"*{ext}")))

    for file_path in sorted(files_to_hash):
        try:
            stat = os.stat(file_path)
            file_info = f"{os.path.basename(file_path)}-{stat.st_mtime}-{stat.st_size}"
            hasher.update(file_info.encode("utf-8"))
        except FileNotFoundError:
            continue  # File might have been deleted during the process

    return hasher.hexdigest()


def get_documents_from_files(folder_path=DATA_FOLDER) -> List[Document]:
    docs = []
    supported_extensions = [".pdf", ".txt", ".md", ".docx"]

    for file_path in glob.glob(os.path.join(folder_path, "*")):
        file_name = os.path.basename(file_path)
        file_ext = os.path.splitext(file_name)[1].lower()

        if file_ext not in supported_extensions:
            continue

        if file_ext == ".pdf":
            pdf_reader = PdfReader(file_path)
            for i, page in enumerate(pdf_reader.pages):
                text = page.extract_text() or ""
                docs.append(
                    Document(
                        page_content=text, metadata={"source": file_name, "page": i + 1}
                    )
                )
        elif file_ext in [".txt", ".md"]:
            with open(file_path, "r", encoding="utf-8") as f:
                text = f.read()
                docs.append(
                    Document(
                        page_content=text, metadata={"source": file_name, "page": 1}
                    )
                )
        elif file_ext == ".docx":
            if DocxDocument is None:
                logger.warning(
                    "Skipping .docx file because 'python-docx' is not installed. Please install it with 'pip install python-docx'."
                )
                continue
            try:
                doc = DocxDocument(file_path)
                text = "\n".join([para.text for para in doc.paragraphs])
                docs.append(
                    Document(
                        page_content=text, metadata={"source": file_name, "page": 1}
                    )
                )
            except Exception as e:
                logger.error(f"Error processing .docx file {file_name}: {e}")

    return docs


async def initialize_rag_system():
    global retrieval_chain
    logger.info("🚀 Initializing RAG system...")
    os.makedirs(CACHE_DIR, exist_ok=True)

    data_hash = get_data_hash()
    cache_path = os.path.join(CACHE_DIR, f"vectorstore_{data_hash}.faiss")

    embeddings = HuggingFaceEmbeddings(
        model_name="BAAI/bge-base-en-v1.5",
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True},
    )

    if os.path.exists(cache_path):
        logger.info(f"✅ Loading vector store from cache: {cache_path}")
        vectorstore = FAISS.load_local(
            cache_path, embeddings, allow_dangerous_deserialization=True
        )
        documents = get_documents_from_files()
        text_chunks = RecursiveCharacterTextSplitter(
            chunk_size=1000, chunk_overlap=150
        ).split_documents(documents)
    else:
        logger.warning(
            "No cache found or data has changed. Building new vector store..."
        )
        documents = get_documents_from_files()
        if not documents:
            logger.warning(
                "No documents found in /data folder. RAG system will be inactive."
            )
            retrieval_chain = None  # Ensure chain is None
            return

        text_chunks = RecursiveCharacterTextSplitter(
            chunk_size=1000, chunk_overlap=150
        ).split_documents(documents)

        logger.info(
            f"Embedding {len(text_chunks)} chunks... (This will be slow on the first run)"
        )
        vectorstore = FAISS.from_documents(documents=text_chunks, embedding=embeddings)

        logger.info(f"💾 Saving new vector store to cache: {cache_path}")
        vectorstore.save_local(cache_path)

    # --- Chain Setup ---
    mistral_api_key = os.getenv("MISTRAL_API_KEY")
    if not mistral_api_key:
        raise ValueError("MISTRAL_API_KEY not found")

    llm = MistralLLM(api_key=mistral_api_key)

    # --- DEPRECATION FIX: Use create_stuff_documents_chain ---
    question_answer_chain = create_stuff_documents_chain(
        llm, UNIVERSITY_ASSISTANT_PROMPT
    )

    faiss_retriever = vectorstore.as_retriever(
        search_type="similarity", search_kwargs={"k": 15}
    )

    tfidf_vectorizer = TfidfVectorizer(stop_words="english")

    # Ensure text_chunks is not empty before fitting
    if not text_chunks:
        logger.warning("No text chunks to process for TF-IDF.")
        # Handle case with no documents gracefully
        tfidf_retriever = ScikitLearnTfidfRetriever(
            docs=[],
            tfidf_vectorizer=tfidf_vectorizer,  # an empty vectorizer
            tfidf_matrix=tfidf_vectorizer.fit_transform([]),  # an empty matrix
            k=15,
        )
    else:
        tfidf_matrix = tfidf_vectorizer.fit_transform(
            [doc.page_content for doc in text_chunks]
        )
        tfidf_retriever = ScikitLearnTfidfRetriever(
            docs=text_chunks,
            tfidf_vectorizer=tfidf_vectorizer,
            tfidf_matrix=tfidf_matrix,
            k=15,
        )

    ensemble_retriever = EnsembleRetriever(
        retrievers=[tfidf_retriever, faiss_retriever], weights=[0.5, 0.5]
    )

    # Create the final retrieval chain that combines the retriever and the QA chain
    retrieval_chain = create_retrieval_chain(ensemble_retriever, question_answer_chain)

    logger.info("✅ RAG system initialized successfully.")


def extract_sources_from_docs(docs: List[Document]) -> List[str]:
    """Extracts unique source and page numbers from a list of documents."""
    return sorted(
        list(
            set(
                f"({doc.metadata.get('source', 'N/A')}, Page: {doc.metadata.get('page', 'N/A')})"
                for doc in docs
            )
        )
    )


# --- FastAPI Application ---
app = FastAPI(title="GNDU Assistant (Cached)", version="4.0.0")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

security = HTTPBasic()


def get_current_username(credentials: HTTPBasicCredentials = Depends(security)):
    correct_username = secrets.compare_digest(
        credentials.username, os.environ.get("ADMIN_USERNAME", "admin")
    )
    correct_password = secrets.compare_digest(
        credentials.password, os.environ.get("ADMIN_PASSWORD", "password")
    )
    if not (correct_username and correct_password):
        raise HTTPException(
            status_code=401,
            detail="Incorrect email or password",
            headers={"WWW-Authenticate": "Basic"},
        )
    return credentials.username


@app.on_event("startup")
async def startup_event():
    # On Hugging Face, copy initial files from the read-only repo folder to persistent storage
    if DATA_FOLDER != REPO_DATA_FOLDER and os.path.exists(REPO_DATA_FOLDER):
        supported_extensions = [".pdf", ".txt", ".md", ".docx"]
        # Check if the destination is empty of supported files before copying
        if not any(
            any(f.endswith(ext) for ext in supported_extensions)
            for f in os.listdir(DATA_FOLDER)
        ):
            logger.info(
                f"Copying initial files from '{REPO_DATA_FOLDER}' to '{DATA_FOLDER}'..."
            )
            for filename in os.listdir(REPO_DATA_FOLDER):
                if any(filename.endswith(ext) for ext in supported_extensions):
                    src_path = os.path.join(REPO_DATA_FOLDER, filename)
                    dest_path = os.path.join(DATA_FOLDER, filename)
                    if not os.path.exists(dest_path):
                        shutil.copy(src_path, dest_path)

    await initialize_rag_system()


admin_router = APIRouter()


@admin_router.get("/files")
async def get_files(_: str = Depends(get_current_username)):
    """Lists all supported files in the data directory."""
    if not os.path.exists(DATA_FOLDER):
        return []
    supported_extensions = [".pdf", ".txt", ".md", ".docx"]
    files = sorted(
        [
            f
            for f in os.listdir(DATA_FOLDER)
            if any(f.endswith(ext) for ext in supported_extensions)
        ]
    )
    return files


@admin_router.post("/files")
async def upload_file(
    file: UploadFile = File(...), _: str = Depends(get_current_username)
):
    """Uploads a new PDF file."""
    file_path = os.path.join(DATA_FOLDER, file.filename)
    if os.path.exists(file_path):
        raise HTTPException(
            status_code=409, detail="File with this name already exists."
        )

    try:
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
    finally:
        file.file.close()

    await initialize_rag_system()

    return {"filename": file.filename}


@admin_router.delete("/files/{filename}")
async def delete_file(filename: str, _: str = Depends(get_current_username)):
    """Deletes a PDF file."""
    # Basic security check
    if ".." in filename:
        raise HTTPException(status_code=400, detail="Invalid filename.")

    file_path = os.path.join(DATA_FOLDER, filename)

    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="File not found.")

    os.remove(file_path)

    await initialize_rag_system()

    return {"message": f"File '{filename}' deleted successfully."}


app.include_router(admin_router, prefix="/admin")


@app.post("/ask", response_model=AnswerResponse)
async def ask_question(request: QuestionRequest):
    normalized_question = request.question.strip().lower()
    if normalized_question in GREETING_RESPONSES:
        return AnswerResponse(
            answer=GREETING_RESPONSES[normalized_question],
            sources=["Conversational Greeting"],
        )

    if not retrieval_chain:
        raise HTTPException(
            status_code=503,
            detail="RAG system is not initialized. Please add PDF files to the /data folder and restart.",
        )

    try:
        # The new chain structure returns a dictionary with 'answer' and 'context'
        result = await retrieval_chain.ainvoke({"input": request.question})

        answer = result.get("answer", "No answer could be generated.")
        sources = extract_sources_from_docs(result.get("context", []))

        return AnswerResponse(answer=answer, sources=sources)

    except Exception as e:
        logger.error(f"Error during RAG chain invocation: {e}")
        raise HTTPException(status_code=500, detail=str(e))


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

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

	import os, logging
	import glob
	import logging
	import re
	import hashlib
	import numpy as np
	import asyncio
	from typing import List, Optional, Any
	import shutil
	import secrets

	# --- Pydantic & FastAPI ---
	from pydantic import BaseModel, Field
	import uvicorn
	from fastapi import FastAPI, HTTPException, APIRouter, UploadFile, File, Depends
	from fastapi.middleware.cors import CORSMiddleware
	from fastapi.staticfiles import StaticFiles
	from fastapi.security import HTTPBasic, HTTPBasicCredentials

	# --- LangChain & RAG Core ---
	from PyPDF2 import PdfReader

	try:
	from docx import Document as DocxDocument
	except ImportError:
	DocxDocument = None
	from langchain.text_splitter import RecursiveCharacterTextSplitter
	from langchain_community.vectorstores import FAISS
	from langchain.retrievers import EnsembleRetriever
	from langchain.chains.combine_documents import (
	create_stuff_documents_chain,
	) # <-- DEPRECATION FIX
	from langchain.chains import create_retrieval_chain
	from langchain_huggingface import HuggingFaceEmbeddings
	from langchain.llms.base import LLM
	from langchain_core.callbacks.manager import CallbackManagerForLLMRun
	from langchain.docstore.document import Document
	from langchain_core.retrievers import BaseRetriever
	from langchain.prompts import PromptTemplate

	# --- Scikit-Learn for TF-IDF ---
	from sklearn.feature_extraction.text import TfidfVectorizer
	from sklearn.metrics.pairwise import cosine_similarity

	# --- TruLens for Evaluation ---
	# from trulens_eval import Feedback, Tru, TruChain
	# from trulens_eval.feedback.provider import LiteLLM

	# tru = Tru()

	# --- Mistral AI ---
	from mistralai import Mistral as MistralClient
	from dotenv import load_dotenv

	load_dotenv()

	# --- Global Variables & Configuration ---
	logging.basicConfig(
	level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
	)
	logger = logging.getLogger(__name__)

	# Define the original, read-only data folder from the repo
	REPO_DATA_FOLDER = "data"

	# Use persistent storage if available (on Hugging Face Spaces)
	if os.path.exists("/data"):
	logger.info("Persistent storage '/data' found. Using it for PDFs and cache.")
	DATA_FOLDER = "/data/pdf_storage"
	CACHE_DIR = "/data/cache"
	else:
	logger.info(
	"No persistent storage '/data' found. Using local 'data' folder and '/tmp/cache'."
	)
	DATA_FOLDER = REPO_DATA_FOLDER
	CACHE_DIR = "/tmp/cache"

	os.makedirs(DATA_FOLDER, exist_ok=True)
	os.makedirs(CACHE_DIR, exist_ok=True)
	retrieval_chain = None

	HF_CACHE = "/tmp/hf"
	os.environ["HF_HOME"] = HF_CACHE
	os.environ["HF_HOME"] = HF_CACHE
	os.environ["HF_DATASETS_CACHE"] = HF_CACHE
	os.environ["HF_MODULES_CACHE"] = HF_CACHE

	retrieval_chain = None

	db_path = "/data/default.sqlite" if os.path.exists("/data") else "default.sqlite"

	# ✅ Only run the dashboard locally (not in Hugging Face Spaces)
	if os.environ.get("SPACE_ID") is None:
	# tru.run_dashboard()
	pass


	# --- Pydantic Models ---
	class QuestionRequest(BaseModel):
	question: str = Field(..., min_length=1)


	class AnswerResponse(BaseModel):
	answer: str
	sources: List[str]


	# --- Greeting Handler ---
	GREETING_RESPONSES = {
	"hi": "Hello! How can I assist you?",
	"hello": "Hello! How can I help you today?",
	"bye": "Bye-bye! See you soon.",
	"ty": "Your Welcome!",
	}

	# --- Enhanced Prompt ---
	UNIVERSITY_ASSISTANT_PROMPT = PromptTemplate.from_template(
	"""
	You are a meticulous Guru Nanak Dev University Academic Assistant.
	Your task is to provide accurate, specific, and concise answers based ONLY on the provided Context.
	Instructions:
	1. Analyze the Context: Carefully read the user's question and the provided context to synthesize an answer.
	2. Extract Specifics: If asked for a fee, number, or date, provide the EXACT information.
	3. Cite Sources: After your answer, cite the source document and page number: `(Source: Prospectus.pdf, Page: 15)`.
	4. Handle Missing Information: If the info is not in the context, state: "Based on the provided documents, I could not find specific information about this topic."
	---
	Context:
	{context}
	---
	Question:
	{input}
	Answer (with citations):
	"""
	)


	# --- Custom Scikit-Learn TF-IDF Retriever ---
	class ScikitLearnTfidfRetriever(BaseRetriever):
	docs: List[Document]
	tfidf_vectorizer: TfidfVectorizer
	tfidf_matrix: Any
	k: int = 15

	class Config:
	arbitrary_types_allowed = True

	def _get_relevant_documents(self, query: str) -> List[Document]:
	query_vector = self.tfidf_vectorizer.transform([query])
	similarity_scores = cosine_similarity(query_vector, self.tfidf_matrix).flatten()
	top_k_indices = np.argsort(similarity_scores)[-self.k :][::-1]
	return [self.docs[i] for i in top_k_indices]


	# --- MistralLLM Class ---
	class MistralLLM(LLM):
	client: Any = None
	model: str = "mistral-small-latest"
	temperature: float = 0.0
	api_key: str = ""

	def __init__(
	self,
	api_key: str,
	model: str = "mistral-small-latest",
	temperature: float = 0.0,
	):
	super().__init__()
	self.api_key = api_key
	self.model = model
	self.temperature = temperature
	self.client = MistralClient(api_key=self.api_key)

	@property
	def _llm_type(self) -> str:
	return "mistral"

	def _call(
	self,
	prompt: str,
	stop: Optional[List[str]] = None,
	run_manager: Optional[CallbackManagerForLLMRun] = None,
	**kwargs: Any,
	) -> str:
	messages = [{"role": "user", "content": prompt}]
	response = self.client.chat.complete(
	model=self.model, messages=messages, temperature=self.temperature
	)
	return response.choices[0].message.content.strip()


	# --- Caching and Data Processing Functions ---
	def get_data_hash(folder_path=DATA_FOLDER) -> str:
	"""Generates a hash based on the filenames and modification times of supported files."""
	hasher = hashlib.md5()
	if not os.path.exists(folder_path):
	return "no_data"

	supported_extensions = [".pdf", ".txt", ".md", ".docx"]
	files_to_hash = []
	for ext in supported_extensions:
	files_to_hash.extend(glob.glob(os.path.join(folder_path, f"*{ext}")))

	for file_path in sorted(files_to_hash):
	try:
	stat = os.stat(file_path)
	file_info = f"{os.path.basename(file_path)}-{stat.st_mtime}-{stat.st_size}"
	hasher.update(file_info.encode("utf-8"))
	except FileNotFoundError:
	continue # File might have been deleted during the process

	return hasher.hexdigest()


	def get_documents_from_files(folder_path=DATA_FOLDER) -> List[Document]:
	docs = []
	supported_extensions = [".pdf", ".txt", ".md", ".docx"]

	for file_path in glob.glob(os.path.join(folder_path, "*")):
	file_name = os.path.basename(file_path)
	file_ext = os.path.splitext(file_name)[1].lower()

	if file_ext not in supported_extensions:
	continue

	if file_ext == ".pdf":
	pdf_reader = PdfReader(file_path)
	for i, page in enumerate(pdf_reader.pages):
	text = page.extract_text() or ""
	docs.append(
	Document(
	page_content=text, metadata={"source": file_name, "page": i + 1}
	)
	)
	elif file_ext in [".txt", ".md"]:
	with open(file_path, "r", encoding="utf-8") as f:
	text = f.read()
	docs.append(
	Document(
	page_content=text, metadata={"source": file_name, "page": 1}
	)
	)
	elif file_ext == ".docx":
	if DocxDocument is None:
	logger.warning(
	"Skipping .docx file because 'python-docx' is not installed. Please install it with 'pip install python-docx'."
	)
	continue
	try:
	doc = DocxDocument(file_path)
	text = "\n".join([para.text for para in doc.paragraphs])
	docs.append(
	Document(
	page_content=text, metadata={"source": file_name, "page": 1}
	)
	)
	except Exception as e:
	logger.error(f"Error processing .docx file {file_name}: {e}")

	return docs


	async def initialize_rag_system():
	global retrieval_chain
	logger.info("🚀 Initializing RAG system...")
	os.makedirs(CACHE_DIR, exist_ok=True)

	data_hash = get_data_hash()
	cache_path = os.path.join(CACHE_DIR, f"vectorstore_{data_hash}.faiss")

	embeddings = HuggingFaceEmbeddings(
	model_name="BAAI/bge-base-en-v1.5",
	model_kwargs={"device": "cpu"},
	encode_kwargs={"normalize_embeddings": True},
	)

	if os.path.exists(cache_path):
	logger.info(f"✅ Loading vector store from cache: {cache_path}")
	vectorstore = FAISS.load_local(
	cache_path, embeddings, allow_dangerous_deserialization=True
	)
	documents = get_documents_from_files()
	text_chunks = RecursiveCharacterTextSplitter(
	chunk_size=1000, chunk_overlap=150
	).split_documents(documents)
	else:
	logger.warning(
	"No cache found or data has changed. Building new vector store..."
	)
	documents = get_documents_from_files()
	if not documents:
	logger.warning(
	"No documents found in /data folder. RAG system will be inactive."
	)
	retrieval_chain = None # Ensure chain is None
	return

	text_chunks = RecursiveCharacterTextSplitter(
	chunk_size=1000, chunk_overlap=150
	).split_documents(documents)

	logger.info(
	f"Embedding {len(text_chunks)} chunks... (This will be slow on the first run)"
	)
	vectorstore = FAISS.from_documents(documents=text_chunks, embedding=embeddings)

	logger.info(f"💾 Saving new vector store to cache: {cache_path}")
	vectorstore.save_local(cache_path)

	# --- Chain Setup ---
	mistral_api_key = os.getenv("MISTRAL_API_KEY")
	if not mistral_api_key:
	raise ValueError("MISTRAL_API_KEY not found")

	llm = MistralLLM(api_key=mistral_api_key)

	# --- DEPRECATION FIX: Use create_stuff_documents_chain ---
	question_answer_chain = create_stuff_documents_chain(
	llm, UNIVERSITY_ASSISTANT_PROMPT
	)

	faiss_retriever = vectorstore.as_retriever(
	search_type="similarity", search_kwargs={"k": 15}
	)

	tfidf_vectorizer = TfidfVectorizer(stop_words="english")

	# Ensure text_chunks is not empty before fitting
	if not text_chunks:
	logger.warning("No text chunks to process for TF-IDF.")
	# Handle case with no documents gracefully
	tfidf_retriever = ScikitLearnTfidfRetriever(
	docs=[],
	tfidf_vectorizer=tfidf_vectorizer, # an empty vectorizer
	tfidf_matrix=tfidf_vectorizer.fit_transform([]), # an empty matrix
	k=15,
	)
	else:
	tfidf_matrix = tfidf_vectorizer.fit_transform(
	[doc.page_content for doc in text_chunks]
	)
	tfidf_retriever = ScikitLearnTfidfRetriever(
	docs=text_chunks,
	tfidf_vectorizer=tfidf_vectorizer,
	tfidf_matrix=tfidf_matrix,
	k=15,
	)

	ensemble_retriever = EnsembleRetriever(
	retrievers=[tfidf_retriever, faiss_retriever], weights=[0.5, 0.5]
	)

	# Create the final retrieval chain that combines the retriever and the QA chain
	retrieval_chain = create_retrieval_chain(ensemble_retriever, question_answer_chain)

	logger.info("✅ RAG system initialized successfully.")


	def extract_sources_from_docs(docs: List[Document]) -> List[str]:
	"""Extracts unique source and page numbers from a list of documents."""
	return sorted(
	list(
	set(
	f"({doc.metadata.get('source', 'N/A')}, Page: {doc.metadata.get('page', 'N/A')})"
	for doc in docs
	)
	)
	)


	# --- FastAPI Application ---
	app = FastAPI(title="GNDU Assistant (Cached)", version="4.0.0")
	app.add_middleware(
	CORSMiddleware,
	allow_origins=["*"],
	allow_credentials=True,
	allow_methods=["*"],
	allow_headers=["*"],
	)

	security = HTTPBasic()


	def get_current_username(credentials: HTTPBasicCredentials = Depends(security)):
	correct_username = secrets.compare_digest(
	credentials.username, os.environ.get("ADMIN_USERNAME", "admin")
	)
	correct_password = secrets.compare_digest(
	credentials.password, os.environ.get("ADMIN_PASSWORD", "password")
	)
	if not (correct_username and correct_password):
	raise HTTPException(
	status_code=401,
	detail="Incorrect email or password",
	headers={"WWW-Authenticate": "Basic"},
	)
	return credentials.username


	@app.on_event("startup")
	async def startup_event():
	# On Hugging Face, copy initial files from the read-only repo folder to persistent storage
	if DATA_FOLDER != REPO_DATA_FOLDER and os.path.exists(REPO_DATA_FOLDER):
	supported_extensions = [".pdf", ".txt", ".md", ".docx"]
	# Check if the destination is empty of supported files before copying
	if not any(
	any(f.endswith(ext) for ext in supported_extensions)
	for f in os.listdir(DATA_FOLDER)
	):
	logger.info(
	f"Copying initial files from '{REPO_DATA_FOLDER}' to '{DATA_FOLDER}'..."
	)
	for filename in os.listdir(REPO_DATA_FOLDER):
	if any(filename.endswith(ext) for ext in supported_extensions):
	src_path = os.path.join(REPO_DATA_FOLDER, filename)
	dest_path = os.path.join(DATA_FOLDER, filename)
	if not os.path.exists(dest_path):
	shutil.copy(src_path, dest_path)

	await initialize_rag_system()


	admin_router = APIRouter()


	@admin_router.get("/files")
	async def get_files(_: str = Depends(get_current_username)):
	"""Lists all supported files in the data directory."""
	if not os.path.exists(DATA_FOLDER):
	return []
	supported_extensions = [".pdf", ".txt", ".md", ".docx"]
	files = sorted(
	[
	f
	for f in os.listdir(DATA_FOLDER)
	if any(f.endswith(ext) for ext in supported_extensions)
	]
	)
	return files


	@admin_router.post("/files")
	async def upload_file(
	file: UploadFile = File(...), _: str = Depends(get_current_username)
	):
	"""Uploads a new PDF file."""
	file_path = os.path.join(DATA_FOLDER, file.filename)
	if os.path.exists(file_path):
	raise HTTPException(
	status_code=409, detail="File with this name already exists."
	)

	try:
	with open(file_path, "wb") as buffer:
	shutil.copyfileobj(file.file, buffer)
	finally:
	file.file.close()

	await initialize_rag_system()

	return {"filename": file.filename}


	@admin_router.delete("/files/{filename}")
	async def delete_file(filename: str, _: str = Depends(get_current_username)):
	"""Deletes a PDF file."""
	# Basic security check
	if ".." in filename:
	raise HTTPException(status_code=400, detail="Invalid filename.")

	file_path = os.path.join(DATA_FOLDER, filename)

	if not os.path.exists(file_path):
	raise HTTPException(status_code=404, detail="File not found.")

	os.remove(file_path)

	await initialize_rag_system()

	return {"message": f"File '{filename}' deleted successfully."}


	app.include_router(admin_router, prefix="/admin")


	@app.post("/ask", response_model=AnswerResponse)
	async def ask_question(request: QuestionRequest):
	normalized_question = request.question.strip().lower()
	if normalized_question in GREETING_RESPONSES:
	return AnswerResponse(
	answer=GREETING_RESPONSES[normalized_question],
	sources=["Conversational Greeting"],
	)

	if not retrieval_chain:
	raise HTTPException(
	status_code=503,
	detail="RAG system is not initialized. Please add PDF files to the /data folder and restart.",
	)

	try:
	# The new chain structure returns a dictionary with 'answer' and 'context'
	result = await retrieval_chain.ainvoke({"input": request.question})

	answer = result.get("answer", "No answer could be generated.")
	sources = extract_sources_from_docs(result.get("context", []))

	return AnswerResponse(answer=answer, sources=sources)

	except Exception as e:
	logger.error(f"Error during RAG chain invocation: {e}")
	raise HTTPException(status_code=500, detail=str(e))


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

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