Update app.py

app.py

@@ -1,212 +1,110 @@
-import os
 import gradio as gr
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import Chroma
+import os
+import datetime
 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_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 langchain.text_splitter import CharacterTextSplitter
+from langchain.document_loaders import TextLoader, PyPDFLoader
+from langchain.prompts import PromptTemplate
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
-import pickle
-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 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:
-                    continue
-        except ImportError:
-            pass
-        return SimpleEmbeddings()
-
-    def _init_vector_stores(self):
-        try:
-            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"
-            )
-        except Exception as e:
-            logger.error(f"Error initializing vector stores: {str(e)}")
-            raise
-
-    def load_documents(self, files: List[str], assistant_type: str) -> str:
-        try:
-            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}")
-                    continue
-            if not documents:
-                return "No documents could be loaded. Please check your files."
-            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"Successfully loaded {len(chunks)} chunks from {len(documents)} documents into {assistant_type} assistant."
-        except Exception as e:
-            logger.error(f"Error loading documents: {str(e)}")
-            return f"Error loading documents: {str(e)}"
-
-    def get_learning_tutor_response(self, question: str) -> str:
-        try:
-            if not self.learning_vectorstore:
-                return "Please upload some 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 an AI learning assistant. Answer the following student question based on uploaded course materials.
-            Question: {question}
-            """
-            result = qa_chain({"query": prompt})
-            response = result['result']
-            if result.get('source_documents'):
-                response += "\n\n**Sources:**\n"
-                for i, doc in enumerate(result['source_documents'][:3]):
-                    source = doc.metadata.get('source', 'Unknown')
-                    response += f"- {Path(source).name}\n"
-            return response
-        except Exception as e:
-            logger.error(f"Error in learning tutor: {str(e)}")
-            return f"Error generating response: {str(e)}"
-
-    def get_code_helper_response(self, question: str) -> str:
-        try:
-            if not self.code_vectorstore:
-                return "Please upload some code documentation first."
-            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
-            )
-            prompt = f"""
-            You are a code assistant. Answer the following developer question based on uploaded technical documentation.
-            Question: {question}
-            """
-            result = qa_chain({"query": prompt})
-            response = result['result']
-            if result.get('source_documents'):
-                response += "\n\n**Documentation Sources:**\n"
-                for i, doc in enumerate(result['source_documents'][:3]):
-                    source = doc.metadata.get('source', 'Unknown')
-                    response += f"- {Path(source).name}\n"
-            return response
-        except Exception as e:
-            logger.error(f"Error in code helper: {str(e)}")
-            return f"Error generating response: {str(e)}"
-
-def evaluate_retrieval(query: str, ground_truth_docs: List[str], retriever, k: int = 5):
-    try:
-        retrieved_docs = retriever.get_relevant_documents(query)
-        top_k = [doc.page_content for doc in retrieved_docs[:k]]
-        hits = sum([1 for doc in top_k if any(gt.lower() in doc.lower() for gt in ground_truth_docs)])
-        precision = hits / k
-        recall = hits / len(ground_truth_docs) if ground_truth_docs else 0.0
-        print("\n Query:", query)
-        print(" Top-K Retrieved Documents:")
-        for i, doc in enumerate(top_k, 1):
-            print(f"{i}. {doc[:200]}...")
-        print(f"\n Evaluation Results:")
-        print(f" Precision@{k}: {precision:.2f}")
-        print(f" Recall@{k}: {recall:.2f}")
-        return {
-            f"Precision@{k}": precision,
-            f"Recall@{k}": recall,
-            "Hits": hits,
-            "Retrieved": top_k
-        }
-    except Exception as e:
-        logger.error(f"❌ Error during evaluation: {str(e)}")
-        return {
-            f"Precision@{k}": 0.0,
-            f"Recall@{k}": 0.0,
-            "Hits": 0,
-            "Retrieved": []
-        }
-
-def main():
-    load_dotenv()
-    groq_api_key = os.getenv("GROQ_API_KEY")
-    if not groq_api_key:
-        print("Please set your GROQ_API_KEY environment variable")
-        return
-    assistant = RAGAssistant(groq_api_key)
-
-    # Example Evaluation
-    query = "What is supervised learning?"
-    ground_truth_docs = ["Supervised learning is a type of machine learning where the model learns from labeled data."]
-    evaluate_retrieval(
-        query=query,
-        ground_truth_docs=ground_truth_docs,
-        retriever=assistant.learning_vectorstore.as_retriever(search_kwargs={"k": 5})
-    )
-
-if __name__ == "__main__":
-    main()
+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()