Update app.py

app.py

@@ -1,110 +1,203 @@
-import gradio as gr
 import os
-import datetime
+import gradio as gr
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
 from langchain.chains import RetrievalQA
-from langchain.vectorstores import Chroma
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.llms import OpenAI
 from langchain_groq import ChatGroq
-from langchain.text_splitter import CharacterTextSplitter
-from langchain.document_loaders import TextLoader, PyPDFLoader
-from langchain.prompts import PromptTemplate
+from langchain_community.document_loaders import TextLoader, PyPDFLoader
+from langchain.schema import Document
+from pathlib import Path
+from typing import List
+import logging
+import numpy as np
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
-from sentence_transformers import SentenceTransformer
-import numpy as np
-
-# Embedding Models
-hf_embed = HuggingFaceEmbeddings()
-fallback_model = SentenceTransformer('all-MiniLM-L6-v2')
-
-# Vector Store
-vector_store = Chroma(collection_name="ragstore", embedding_function=hf_embed)
-
-# LLM
-llm = ChatGroq(temperature=0, model_name="llama3-8b-8192")
-
-# Prompt Template
-prompt_template = PromptTemplate.from_template(
-    "Answer the following question using ONLY the context provided:\n\n{context}\n\nQuestion: {question}"
-)
-
-# RetrievalQA Chain
-qa_chain = RetrievalQA.from_chain_type(
-    llm=llm,
-    retriever=vector_store.as_retriever(search_kwargs={"k": 3}),
-    chain_type="stuff",
-    chain_type_kwargs={"prompt": prompt_template}
-)
-
-# TF-IDF Fallback
-def tfidf_fallback(query, documents):
-    texts = [doc.page_content for doc in documents]
-    vectorizer = TfidfVectorizer().fit(texts + [query])
-    vectors = vectorizer.transform(texts + [query])
-    cosine_sim = cosine_similarity(vectors[-1], vectors[:-1]).flatten()
-    top_idx = np.argmax(cosine_sim)
-    return texts[top_idx], cosine_sim[top_idx]
-
-# Ingestion
-def ingest_files(files):
-    for file in files:
-        if file.name.endswith(".pdf"):
-            loader = PyPDFLoader(file.name)
-        else:
-            loader = TextLoader(file.name)
-        docs = loader.load()
-        chunks = CharacterTextSplitter(chunk_size=1000, chunk_overlap=200).split_documents(docs)
-        vector_store.add_documents(chunks)
-
-# Evaluation Info
-def evaluate_retrieval(query):
-    docs = vector_store.similarity_search_with_score(query, k=3)
-    top_docs = [doc[0].page_content for doc in docs]
-    scores = [doc[1] for doc in docs]
-    similarities = [1 - s for s in scores]  # cosine similarity approximation
-    return top_docs, similarities
-
-# Final Response Generator
-def ask_question(query):
-    if not query.strip():
-        return "", "", "", "", ""
-
-    # Retrieve docs and similarities
-    docs, similarities = evaluate_retrieval(query)
-    formatted_docs = "\n\n".join([f"Doc {i+1} (Score: {similarities[i]*100:.2f}%)\n{docs[i]}" for i in range(len(docs))])
-    context_block = f"### Top Retrieved Documents:\n{formatted_docs}"
-
-    # Answer from RAG
-    answer = qa_chain.run(query)
-
-    # Baseline (Direct LLM, no context)
-    baseline = llm.invoke(query)
-
-    # Confidence score approximation
-    confidence = np.mean(similarities) * 100
-
-    return answer, context_block, f"{confidence:.2f}%", baseline, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-
-# Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 RAG-Based Learning & Code Assistant\nUpload docs, ask questions, get answers with confidence & evidence.")
-
-    with gr.Row():
-        with gr.Column():
-            file_input = gr.File(label="Upload PDF or TXT", file_types=[".pdf", ".txt"], file_count="multiple")
-            ingest_btn = gr.Button("Ingest Documents")
-            question_input = gr.Textbox(label="Ask a Question")
-            ask_btn = gr.Button("Ask")
-        with gr.Column():
-            answer_output = gr.Textbox(label="RAG Answer", lines=5)
-            retrieved_docs_output = gr.Textbox(label="Top 3 Retrieved Documents", lines=10)
-            confidence_output = gr.Textbox(label="Confidence (%)")
-            baseline_output = gr.Textbox(label="Baseline (Direct LLM)", lines=5)
-            timestamp_output = gr.Textbox(label="Timestamp")
-
-    ingest_btn.click(fn=ingest_files, inputs=file_input, outputs=[])
-    ask_btn.click(fn=ask_question, inputs=question_input,
-                  outputs=[answer_output, retrieved_docs_output, confidence_output, baseline_output, timestamp_output])
-
-demo.launch()
+from dotenv import load_dotenv
+
+# Configure 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 RetrieverEvaluator:
+    def __init__(self, retriever, ground_truth, k=3):
+        self.retriever = retriever
+        self.ground_truth = ground_truth
+        self.k = k
+
+    def recall_at_k(self):
+        correct = 0
+        for query, relevant_docs in self.ground_truth.items():
+            results = self.retriever.get_relevant_documents(query)
+            retrieved = [Path(doc.metadata.get("source", "")).name for doc in results]
+            if any(doc in retrieved[:self.k] for doc in relevant_docs):
+                correct += 1
+        recall = correct / len(self.ground_truth)
+        print(f"Recall@{self.k}: {recall:.2f}")
+        return recall
+
+    def mean_reciprocal_rank(self):
+        mrr_total = 0
+        for query, relevant_docs in self.ground_truth.items():
+            results = self.retriever.get_relevant_documents(query)
+            retrieved = [Path(doc.metadata.get("source", "")).name for doc in results]
+            for rank, doc in enumerate(retrieved[:self.k], 1):
+                if doc in relevant_docs:
+                    mrr_total += 1 / rank
+                    break
+        mrr = mrr_total / len(self.ground_truth)
+        print(f"MRR@{self.k}: {mrr:.2f}")
+        return mrr
+
+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)
+        self.learning_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._init_vector_store()
+
+    def _init_embeddings(self):
+        try:
+            from langchain_huggingface import HuggingFaceEmbeddings
+            for model_name in ["all-MiniLM-L6-v2", "paraphrase-MiniLM-L3-v2", "all-mpnet-base-v2"]:
+                try:
+                    return HuggingFaceEmbeddings(model_name=model_name, model_kwargs={'device': 'cpu'})
+                except:
+                    continue
+        except ImportError:
+            pass
+        return SimpleEmbeddings()
+
+    def _init_vector_store(self):
+        self.learning_vectorstore = Chroma(
+            persist_directory=self.learning_persist_dir,
+            embedding_function=self.embeddings,
+            collection_name="learning_materials"
+        )
+
+    def load_documents(self, files: List[str]) -> str:
+        documents = []
+        for file_path in files:
+            try:
+                loader = PyPDFLoader(file_path) if file_path.endswith(".pdf") else TextLoader(file_path, encoding="utf-8")
+                docs = loader.load()
+                documents.extend(docs)
+            except Exception as e:
+                print(f"Error loading {file_path}: {e}")
+        if not documents:
+            return "No valid documents found."
+        chunks = self.text_splitter.split_documents(documents)
+        for chunk in chunks:
+            chunk.metadata['source'] = chunk.metadata.get('source', 'unknown')
+        self.learning_vectorstore.add_documents(chunks)
+        self.learning_vectorstore.persist()
+        return f"Loaded {len(chunks)} document chunks."
+
+    def get_response(self, query: str) -> str:
+        if not self.learning_vectorstore:
+            return "Please upload learning materials first."
+        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
+        )
+        prompt = f"""
+        You are a helpful educational assistant.
+        Answer the student's question clearly and provide references if applicable.
+
+        Question: {query}
+        """
+        result = qa_chain({"query": prompt})
+        response = result['result']
+        if result.get("source_documents"):
+            response += "\n\n**Sources:**\n"
+            for doc in result["source_documents"]:
+                response += f"- {Path(doc.metadata.get('source', 'Unknown')).name}\n"
+        return response
+
+    def evaluate_retriever(self, user_queries: List[str], file_names: List[str]):
+        """Evaluate with user-provided queries and expected file names"""
+        ground_truth = dict(zip(user_queries, file_names))
+        retriever = self.learning_vectorstore.as_retriever(search_kwargs={"k": 3})
+        evaluator = RetrieverEvaluator(retriever, ground_truth, k=3)
+        recall = evaluator.recall_at_k()
+        mrr = evaluator.mean_reciprocal_rank()
+        return f"Recall@3: {recall:.2f}, MRR@3: {mrr:.2f}"
+
+def create_interface(assistant: RAGAssistant):
+    def upload_files(files):
+        file_paths = [f.name for f in files]
+        return assistant.load_documents(file_paths)
+
+    def chat_fn(message, history):
+        response = assistant.get_response(message)
+        history.append((message, response))
+        return history, ""
+
+    def evaluate_fn(queries, file_names):
+        query_list = [q.strip() for q in queries.split('\n') if q.strip()]
+        file_list = [f.strip() for f in file_names.split('\n') if f.strip()]
+        if len(query_list) != len(file_list):
+            return "Number of queries and expected file names must match."
+        return assistant.evaluate_retriever(query_list, file_list)
+
+    with gr.Blocks(title="RAG Assistant") as demo:
+        gr.Markdown("# 📘 RAG-Based Assistant")
+        with gr.Tab("📄 Upload & Chat"):
+            file_input = gr.File(label="Upload PDFs or Text Files", file_count="multiple", file_types=[".pdf", ".txt"])
+            upload_btn = gr.Button("Load Documents")
+            status = gr.Textbox(label="Status", interactive=False)
+            chatbot = gr.Chatbot()
+            user_input = gr.Textbox(label="Ask a question")
+            send_btn = gr.Button("Send")
+
+            upload_btn.click(fn=upload_files, inputs=[file_input], outputs=[status])
+            send_btn.click(fn=chat_fn, inputs=[user_input, chatbot], outputs=[chatbot, user_input])
+            user_input.submit(fn=chat_fn, inputs=[user_input, chatbot], outputs=[chatbot, user_input])
+
+        with gr.Tab("📊 Evaluate Retriever"):
+            gr.Markdown("Paste queries and expected file names (one per line).")
+            queries = gr.Textbox(lines=5, label="Queries")
+            filenames = gr.Textbox(lines=5, label="Expected File Names")
+            eval_btn = gr.Button("Run Evaluation")
+            eval_result = gr.Textbox(label="Evaluation Result")
+            eval_btn.click(fn=evaluate_fn, inputs=[queries, filenames], outputs=[eval_result])
+
+        gr.Markdown("---")
+        gr.Markdown("*Powered by LangChain, ChromaDB, and Groq API*")
+
+    return demo
+
+def main():
+    load_dotenv()
+    groq_api_key = os.getenv("GROQ_API_KEY")
+    if not groq_api_key:
+        print("Missing GROQ_API_KEY. Set it in your environment.")
+        return
+    assistant = RAGAssistant(groq_api_key)
+    app = create_interface(assistant)
+    app.launch(server_name="0.0.0.0", server_port=7860, share=True)
+
+if __name__ == "__main__":
+    main()