app.py

import os
import gradio as gr
import json
import logging
from pathlib import Path
from typing import List
from dotenv import load_dotenv

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.schema import Document
from langchain_community.document_loaders import TextLoader, PyPDFLoader
from langchain_groq import ChatGroq

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer

# Logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SimpleEmbeddings:
    def __init__(self):
        self.vectorizer = TfidfVectorizer(max_features=384, stop_words='english')
        self.fitted = False

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        if not self.fitted:
            self.vectorizer.fit(texts)
            self.fitted = True
        embeddings = self.vectorizer.transform(texts)
        return embeddings.toarray().tolist()

    def embed_query(self, text: str) -> List[float]:
        if not self.fitted:
            return [0.0] * 384
        embedding = self.vectorizer.transform([text])
        return embedding.toarray()[0].tolist()


class RAGAssistant:
    def __init__(self, groq_api_key: str):
        self.groq_api_key = groq_api_key
        self.embeddings = self._init_embeddings()

        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len
        )

        self.learning_vectorstore = None
        self.code_vectorstore = None

        self.llm = ChatGroq(
            groq_api_key=groq_api_key,
            model_name="llama3-70b-8192",
            temperature=0.1
        )

        self.learning_persist_dir = "./chroma_learning_db"
        self.code_persist_dir = "./chroma_code_db"

        self._init_vector_stores()

    def _init_embeddings(self):
        try:
            from langchain_huggingface import HuggingFaceEmbeddings
            models_to_try = ["all-MiniLM-L6-v2", "paraphrase-MiniLM-L3-v2", "all-mpnet-base-v2"]
            for model_name in models_to_try:
                try:
                    embeddings = HuggingFaceEmbeddings(
                        model_name=model_name,
                        model_kwargs={'device': 'cpu'},
                        encode_kwargs={'normalize_embeddings': False}
                    )
                    return embeddings
                except Exception:
                    continue
        except ImportError:
            pass
        return SimpleEmbeddings()

    def _init_vector_stores(self):
        self.learning_vectorstore = Chroma(
            persist_directory=self.learning_persist_dir,
            embedding_function=self.embeddings,
            collection_name="learning_materials"
        )
        self.code_vectorstore = Chroma(
            persist_directory=self.code_persist_dir,
            embedding_function=self.embeddings,
            collection_name="code_documentation"
        )

    def load_documents(self, files: List[str], assistant_type: str) -> str:
        try:
            documents = []
            for file_path in files:
                if file_path.endswith('.pdf'):
                    loader = PyPDFLoader(file_path)
                else:
                    loader = TextLoader(file_path, encoding='utf-8')
                docs = loader.load()
                documents.extend(docs)

            if not documents:
                return "No documents could be loaded."

            chunks = self.text_splitter.split_documents(documents)
            for chunk in chunks:
                chunk.metadata['assistant_type'] = assistant_type

            if assistant_type == "learning":
                self.learning_vectorstore.add_documents(chunks)
                self.learning_vectorstore.persist()
            elif assistant_type == "code":
                self.code_vectorstore.add_documents(chunks)
                self.code_vectorstore.persist()

            return f"Loaded {len(chunks)} chunks from {len(documents)} documents."
        except Exception as e:
            return f"Error loading documents: {str(e)}"

    def get_learning_tutor_response(self, question: str) -> str:
        if not self.learning_vectorstore:
            return "Please upload some learning materials."
        qa_chain = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=self.learning_vectorstore.as_retriever(search_kwargs={"k": 3}),
            return_source_documents=True
        )
        result = qa_chain({"query": question})
        response = result['result']
        if result.get('source_documents'):
            response += "\n\n**Sources:**\n"
            for doc in result['source_documents'][:3]:
                response += f"- {Path(doc.metadata.get('source', 'Unknown')).name}\n"
        return response

    def get_code_helper_response(self, question: str) -> str:
        if not self.code_vectorstore:
            return "Please upload some code documentation."
        qa_chain = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=self.code_vectorstore.as_retriever(search_kwargs={"k": 3}),
            return_source_documents=True
        )
        result = qa_chain({"query": question})
        response = result['result']
        if result.get('source_documents'):
            response += "\n\n**Documentation Sources:**\n"
            for doc in result['source_documents'][:3]:
                response += f"- {Path(doc.metadata.get('source', 'Unknown')).name}\n"
        return response


# --- Evaluation Utilities ---
def evaluate_retrieval(assistant, assistant_type: str, eval_file: str, k=3):
    try:
        if not os.path.exists(eval_file):
            return f"Evaluation file {eval_file} not found."

        with open(eval_file, 'r', encoding='utf-8') as f:
            eval_data = json.load(f)

        total = len(eval_data)
        hits = 0
        mrr = 0

        for idx, item in enumerate(eval_data):
            question = item.get("question", "")
            keywords = item.get("keywords", [])
            result = assistant.get_learning_tutor_response(question) if assistant_type == "learning" else assistant.get_code_helper_response(question)
            hit = any(kw.lower() in result.lower() for kw in keywords)
            if hit:
                hits += 1
                mrr += 1 / (idx + 1)

        precision = hits / total if total else 0
        recall = precision
        mean_rr = mrr / total if total else 0

        return f"""
📊 Evaluation Summary ({assistant_type.title()} Assistant):
- Total Queries: {total}
- Precision@{k}: {precision:.2f}
- Recall@{k}: {recall:.2f}
- MRR: {mean_rr:.2f}

⚙️ Config:
- Retriever Top-K: {k}
- Embedding Model: {getattr(assistant.embeddings, 'model_name', 'TF-IDF (fallback)')}
- Vector Store: ChromaDB
        """
    except Exception as e:
        return f"Evaluation error: {str(e)}"


def create_gradio_interface(assistant: RAGAssistant):
    def upload_learning_files(files):
        return assistant.load_documents([f.name for f in files], "learning") if files else "No files uploaded."

    def upload_code_files(files):
        return assistant.load_documents([f.name for f in files], "code") if files else "No files uploaded."

    def learning_chat(message, history):
        if not message.strip():
            return history, ""
        response = assistant.get_learning_tutor_response(message)
        history.append((message, response))
        return history, ""

    def code_chat(message, history):
        if not message.strip():
            return history, ""
        response = assistant.get_code_helper_response(message)
        history.append((message, response))
        return history, ""

    with gr.Blocks(title="RAG-Based Learning & Code Assistant", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🎓 RAG-Based Learning & Code Assistant")

        with gr.Tabs():
            with gr.TabItem("📚 Learning Tutor"):
                learning_files = gr.File(file_types=[".pdf", ".txt", ".md"], file_count="multiple")
                learning_upload_btn = gr.Button("Upload Materials")
                learning_status = gr.Textbox()
                learning_chatbot = gr.Chatbot(height=400)
                learning_input = gr.Textbox(placeholder="Ask about your learning materials")
                learning_submit = gr.Button("Ask")

                learning_upload_btn.click(upload_learning_files, [learning_files], [learning_status])
                learning_submit.click(learning_chat, [learning_input, learning_chatbot], [learning_chatbot, learning_input])
                learning_input.submit(learning_chat, [learning_input, learning_chatbot], [learning_chatbot, learning_input])

            with gr.TabItem("💻 Code Documentation Helper"):
                code_files = gr.File(file_types=[".pdf", ".txt", ".md", ".py", ".js", ".json"], file_count="multiple")
                code_upload_btn = gr.Button("Upload Documentation")
                code_status = gr.Textbox()
                code_chatbot = gr.Chatbot(height=400)
                code_input = gr.Textbox(placeholder="Ask about your codebase or APIs")
                code_submit = gr.Button("Ask")

                code_upload_btn.click(upload_code_files, [code_files], [code_status])
                code_submit.click(code_chat, [code_input, code_chatbot], [code_chatbot, code_input])
                code_input.submit(code_chat, [code_input, code_chatbot], [code_chatbot, code_input])

            with gr.TabItem("📈 Evaluation Bench"):
                gr.Markdown("Upload a supported document (.json) containing evaluation questions and expected keywords.")
                eval_file_input = gr.File(file_types=[".json"], file_count="single")
                eval_assistant_choice = gr.Radio(["learning", "code"], label="Assistant", value="learning")
                eval_button = gr.Button("Run Evaluation")
                eval_output = gr.Textbox(lines=10)

                def run_eval(file, assistant_type):
                    if file is None:
                        return "Please upload a valid evaluation file."
                    return evaluate_retrieval(assistant, assistant_type, file.name)

                eval_button.click(run_eval, [eval_file_input, eval_assistant_choice], [eval_output])

    return demo


def main():
    load_dotenv()
    groq_api_key = os.getenv("GROQ_API_KEY")
    if not groq_api_key:
        print("Set your GROQ_API_KEY in .env file")
        return
    assistant = RAGAssistant(groq_api_key)
    demo = create_gradio_interface(assistant)
    demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)


if __name__ == "__main__":
    main()