Update app.py

app.py

@@ -1,305 +1,203 @@
-# -*- coding: utf-8 -*-
-"""App
-
-Automatically generated by Colab.
-
-Original file is located at
-    https://colab.research.google.com/drive/1TdjbTSA8V5GUProQ3Bd-uYmTLXSInoWf
-"""
-
-import gradio as gr
-import numpy as np
-from transformers import pipeline
 import os
-import time
-import groq
-import uuid  # For generating unique filenames
-
-# Updated imports to address LangChain deprecation warnings:
+import re
+import uuid
+import tempfile
+import numpy as np
+import gradio as gr
+import chardet
+import fitz  # PyMuPDF
+import docx
+import gtts
+from pptx import Presentation
+from typing import TypedDict, List
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
+from langgraph.graph import StateGraph, END
 from langchain_groq import ChatGroq
-from langchain.schema import HumanMessage
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain.docstore.document import Document
-
-# Importing chardet (make sure to add chardet to your requirements.txt)
-import chardet
-
-import fitz  # PyMuPDF for PDFs
-import docx  # python-docx for Word files
-import gtts  # Google Text-to-Speech library
-from pptx import Presentation  # python-pptx for PowerPoint files
-import re
-
-# Initialize Whisper model for speech-to-text
-transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-base.en")
-
-# Set API Key (Ensure it's stored securely in an environment variable)
-groq.api_key = os.getenv("GROQ_API_KEY")  # Replace with a valid API key
+from langchain_community.vectorstores import Chroma
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
 
-# Initialize Chat Model
-chat_model = ChatGroq(model_name="llama-3.1-8b-instant", api_key=groq.api_key)  #llama-3.3-70b-versatile
+# --- 1. INITIALIZATION & CORE TOOLS ---
+groq_api_key = os.getenv("GROQ_API_KEY")
 
-# Initialize Embeddings and chromaDB
-embedding_model = HuggingFaceEmbeddings()
-vectorstore = Chroma(embedding_function=embedding_model)
+chat_model = ChatGroq(model_name="llama-3.3-70b-versatile", api_key=groq_api_key)
+web_search_tool = DuckDuckGoSearchRun()
+embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+vectorstore = Chroma(embedding_function=embedding_model, persist_directory="chroma_db")
 
-# Short-term memory for the LLM
-chat_memory = []
+# --- 2. HELPER FUNCTIONS ---
 
-# Prompt for quiz generation with added remark
-quiz_prompt = """
-You are an AI assistant specialized in education and assessment creation. Given an uploaded document or text, generate a quiz with a mix of multiple-choice questions (MCQs) and fill-in-the-blank questions. The quiz should be directly based on the key concepts, facts, and details from the provided material.
-Remove all unnecessary formatting generated by the LLM, including <think> tags, asterisks, markdown formatting, and any bold or italic text, as well as **, ###, ##, and # tags.
-For each question:
-- Provide 4 answer choices (for MCQs), with only one correct answer.
-- Ensure fill-in-the-blank questions focus on key terms, phrases, or concepts from the document.
-- Include an answer key for all questions.
-- Ensure questions vary in difficulty and encourage comprehension rather than memorization.
-- Additionally, implement an instant feedback mechanism:
-    - When a user selects an answer, indicate whether it is correct or incorrect.
-    - If incorrect, provide a brief explanation from the document to guide learning.
-    - Ensure responses are concise and educational to enhance understanding.
-Output Example:
-1. Fill in the blank: The LLM Agent framework has a central decision-making unit called the _______________________.
-Answer: Agent Core
-Feedback: The Agent Core is the central component of the LLM Agent framework, responsible for managing goals, tool instructions, planning modules, memory integration, and agent persona.
-2. What is the main limitation of LLM-based applications?
-a) Limited token capacity
-b) Lack of domain expertise
-c) Prone to hallucination
-d) All of the above
-Answer: d) All of the above
-Feedback: LLM-based applications have several limitations, including limited token capacity, lack of domain expertise, and being prone to hallucination, among others.
-"""
 
-# Function to clean AI response by removing unwanted formatting
 def clean_response(response):
-    """Removes <think> tags, asterisks, and markdown formatting."""
-    cleaned_text = re.sub(r"<think>.*?</think>", "", response, flags=re.DOTALL)
-    cleaned_text = re.sub(r"(\*\*|\*)", "", cleaned_text)
-    cleaned_text = re.sub(r"^#+\s*", "", cleaned_text, flags=re.MULTILINE)
-    cleaned_text = re.sub(r"\\", "", cleaned_text)
-    return cleaned_text.strip()
+    """Remove <think>...</think> blocks and common markdown artifacts."""
+    # Remove think tags and their content (greedily, case-insensitive)
+    cleaned = re.sub(r"<think>.*?(?:</think>|$)", "", response, flags=re.DOTALL | re.IGNORECASE)
+    # Remove stray closing tags and markdown symbols
+    cleaned = re.sub(r"</?think>|\*\*|\*|\[|\]|#", "", cleaned)
+    return cleaned.strip()
+    #return cleaned_text.strip()
 
-# Function to generate quiz based on content
-def generate_quiz(content):
-    prompt = f"{quiz_prompt}\n\nDocument content:\n{content}"
-    response = chat_model([HumanMessage(content=prompt)])
-    cleaned_response = clean_response(response.content)
-    return cleaned_response
 
-# Function to retrieve relevant documents from vectorstore based on user query
 def retrieve_documents(query):
     results = vectorstore.similarity_search(query, k=3)
     return [doc.page_content for doc in results]
 
-# Function to handle chatbot interactions with short-term memory
-def chat_with_groq(user_input):
-    try:
-        # Retrieve relevant documents for additional context
-        relevant_docs = retrieve_documents(user_input)
-        context = "\n".join(relevant_docs) if relevant_docs else "No relevant documents found."
-
-        # Construct proper prompting with conversation history
-        system_prompt = "You are a helpful AI assistant. Answer questions accurately and concisely."
-        conversation_history = "\n".join(chat_memory[-10:])  # Keep the last 10 exchanges
-        prompt = f"{system_prompt}\n\nConversation History:\n{conversation_history}\n\nUser Input: {user_input}\n\nContext:\n{context}"
-
-        # Call the chat model
-        response = chat_model([HumanMessage(content=prompt)])
-
-        # Clean response to remove any unwanted formatting
-        cleaned_response_text = clean_response(response.content)
-
-        # Append conversation history
-        chat_memory.append(f"User: {user_input}")
-        chat_memory.append(f"AI: {cleaned_response_text}")
-
-        # Convert response to speech
-        audio_file = speech_playback(cleaned_response_text)
-        
-        return cleaned_response_text, audio_file
-    except Exception as e:
-        return f"Error: {str(e)}", None
-
-# Function to play response as speech using gTTS
 def speech_playback(text):
     try:
-        # Generate a unique filename for each audio file
         unique_id = str(uuid.uuid4())
-        audio_file = f"output_audio_{unique_id}.mp3"
-        
-        # Convert text to speech
-        tts = gtts.gTTS(text, lang='en')
+        audio_file = f"/content/output_audio_{unique_id}.mp3"
+        tts = gtts.gTTS(text[:500], lang='en')
         tts.save(audio_file)
-        
-        # Return the path to the audio file
         return audio_file
     except Exception as e:
-        print(f"Error in speech_playback: {e}")
+        print(f"TTS error: {e}")
         return None
 
-# Function to detect encoding safely
-def detect_encoding(file_path):
-    try:
-        with open(file_path, "rb") as f:
-            raw_data = f.read(4096)
-            detected = chardet.detect(raw_data)
-            encoding = detected["encoding"]
-        return encoding if encoding else "utf-8"
-    except Exception:
-        return "utf-8"
+# --- 3. DOCUMENT INGESTION FUNCTION ---
+def extract_and_store_document(file_path: str):
+    text = ""
+    file_ext = os.path.splitext(file_path)[1].lower()
 
-# Function to extract text from PDF
-def extract_text_from_pdf(pdf_path):
     try:
-        doc = fitz.open(pdf_path)
-        text = "\n".join([page.get_text("text") for page in doc])
-        return text if text.strip() else "No extractable text found."
-    except Exception as e:
-        return f"Error extracting text from PDF: {str(e)}"
-
-# Function to extract text from Word files (.docx)
-def extract_text_from_docx(docx_path):
-    try:
-        doc = docx.Document(docx_path)
-        text = "\n".join([para.text for para in doc.paragraphs])
-        return text if text.strip() else "No extractable text found."
-    except Exception as e:
-        return f"Error extracting text from Word document: {str(e)}"
-
-# Function to extract text from PowerPoint files (.pptx)
-def extract_text_from_pptx(pptx_path):
-    try:
-        presentation = Presentation(pptx_path)
-        text = ""
-        for slide in presentation.slides:
-            for shape in slide.shapes:
-                if hasattr(shape, "text"):
-                    text += shape.text + "\n"
-        return text if text.strip() else "No extractable text found."
-    except Exception as e:
-        return f"Error extracting text from PowerPoint: {str(e)}"
-
-# Function to process documents safely
-#def process_document(file):
-#    try:
-#        file_extension = os.path.splitext(file.name)[-1].lower()
-#        if file_extension in [".png", ".jpg", ".jpeg"]:
-#            return "Error: Images cannot be processed for text extraction."
-#        if file_extension == ".pdf":
-#            content = extract_text_from_pdf(file.name)
-#        elif file_extension == ".docx":
-#            content = extract_text_from_docx(file.name)
-#        elif file_extension == ".pptx":
-#            content = extract_text_from_pptx(file.name)
-#        else:
-#            encoding = detect_encoding(file.name)
-#            with open(file.name, "r", encoding=encoding, errors="replace") as f:
-#               content = f.read()
-#        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-#        documents = [Document(page_content=chunk) for chunk in text_splitter.split_text(content)]
-#        vectorstore.add_documents(documents)
-#        quiz = generate_quiz(content)
-#        return f"Document processed successfully (File Type: {file_extension}). Quiz generated:\n{quiz}"
-#    except Exception as e:
-#        return f"Error processing document: {str(e)}"
-
-def process_document(file):
-    try:
-        if not file or not hasattr(file, "name") or not isinstance(file.name, str):
-            return "Error: Invalid file uploaded."
-
-        file_extension = os.path.splitext(file.name)[-1].lower()
-        
-        if file_extension in [".png", ".jpg", ".jpeg"]:
-            return "Error: Images cannot be processed for text extraction."
-        
-        if file_extension == ".pdf":
-            content = extract_text_from_pdf(file.name)
-        elif file_extension == ".docx":
-            content = extract_text_from_docx(file.name)
-        elif file_extension == ".pptx":
-            content = extract_text_from_pptx(file.name)
+        if file_ext == ".pdf":
+            doc = fitz.open(file_path)
+            for page in doc:
+                text += page.get_text()
+            doc.close()
+        elif file_ext == ".docx":
+            doc = docx.Document(file_path)
+            text = "\n".join([para.text for para in doc.paragraphs])
+        elif file_ext == ".pptx":
+            prs = Presentation(file_path)
+            for slide in prs.slides:
+                for shape in slide.shapes:
+                    if hasattr(shape, "text"):
+                        text += shape.text + "\n"
         else:
-            encoding = detect_encoding(file.name)
-            with open(file.name, "r", encoding=encoding, errors="replace") as f:
-                content = f.read()
+            with open(file_path, 'rb') as f:
+                raw_data = f.read()
+                encoding = chardet.detect(raw_data)['encoding'] or 'utf-8'
+                text = raw_data.decode(encoding, errors='ignore')
+        
+        if not text.strip():
+            return False
 
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-        documents = [Document(page_content=chunk) for chunk in text_splitter.split_text(content)]
+        splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+        chunks = splitter.split_text(text)
+        documents = [Document(page_content=chunk, metadata={"source": os.path.basename(file_path)}) for chunk in chunks]
         vectorstore.add_documents(documents)
-        quiz = generate_quiz(content)
-        return f"Document processed successfully (File Type: {file_extension}). Quiz generated:\n{quiz}"
-    
+        vectorstore.persist()
+        return True
+
     except Exception as e:
-        return f"Error processing document: {str(e)}"
+        print(f"Error processing {file_path}: {e}")
+        return False
 
+# --- 4. REFRAG MULTI-AGENT LOGIC (LangGraph) ---
 
+class AgentState(TypedDict):
+    messages: List[BaseMessage]
+    context: str
+    decision: str 
+    source: str
 
-# Function to handle speech-to-text conversion
-def transcribe_audio(audio):
-    sr, y = audio
-    if y.ndim > 1:
-        y = y.mean(axis=1)
-    y = y.astype(np.float32)
-    y /= np.max(np.abs(y))
-    return transcriber({"sampling_rate": sr, "raw": y})["text"]
+def sensing_node(state: AgentState):
+    user_query = state["messages"][-1].content
+    relevant_docs = retrieve_documents(user_query) 
+    context = "\n".join(relevant_docs) if relevant_docs else ""
+    
+    prompt = f"Docs: {context}\nQuery: {user_query}\nIf docs answer this, reply 'RAG'. Else reply 'WEB'."
+    decision = chat_model.invoke([HumanMessage(content=prompt)]).content.strip().upper()
+    return {"context": context, "decision": "RAG" if "RAG" in decision else "WEB"}
+
+def expansion_node(state: AgentState):
+    if state["decision"] == "WEB":
+        user_query = state["messages"][-1].content
+        web_data = web_search_tool.run(user_query)
+        return {"context": f"WEB INFO: {web_data}\nLOCAL: {state['context']}", "source": "Web + Local Documents"}
+    return {"source": "Local Documents Only"}
+
+def generation_node(state: AgentState):
+    system_msg = f"You are a Tutor AI. Use this context: {state['context']}"
+    response = chat_model.invoke([SystemMessage(content=system_msg)] + state["messages"])
+    cleaned = clean_response(response.content)
+    return {"messages": [AIMessage(content=f"{cleaned}\n\n*(Verified via: {state['source']})*")]}
+
+workflow = StateGraph(AgentState)
+workflow.add_node("sense", sensing_node)
+workflow.add_node("expand", expansion_node)
+workflow.add_node("generate", generation_node)
+workflow.set_entry_point("sense")
+workflow.add_edge("sense", "expand")
+workflow.add_edge("expand", "generate")
+workflow.add_edge("generate", END)
+app_agent = workflow.compile()
+
+# --- 5. GRADIO APP WITH MANUAL AUDIO ---
+
+# Store last assistant response globally (simple approach for demo)
+last_assistant_response = ""
+
+def chat_handler(user_input, chat_history):
+    global last_assistant_response
+    if not user_input: 
+        return chat_history, "", None
+    
+    inputs = {"messages": [HumanMessage(content=user_input)], "context": "", "decision": "", "source": ""}
+    result = app_agent.invoke(inputs)
+    final_msg = result["messages"][-1].content
+    
+    chat_history.append({"role": "user", "content": user_input})
+    chat_history.append({"role": "assistant", "content": final_msg})
+    
+    # Save clean text for later TTS (without source note)
+    last_assistant_response = final_msg.split("*(Verified")[0].strip()
+    
+    # Return chat history and clear audio (no autoplay)
+    return chat_history, "", None
 
-# Your cleanup function
-def cleanup_old_files(directory=".", age_limit=60):
-    """Delete files older than `age_limit` seconds."""
-    current_time = time.time()
-    for filename in os.listdir(directory):
-        file_path = os.path.join(directory, filename)
-        if os.path.isfile(file_path) and filename.startswith("output_audio_"):
-            file_age = current_time - os.path.getmtime(file_path)
-            if file_age > age_limit:
-                os.remove(file_path)
+def generate_audio():
+    global last_assistant_response
+    if not last_assistant_response:
+        return None
+    return speech_playback(last_assistant_response)
 
+def upload_file(file):
+    if file is None:
+        return "❌ No file uploaded."
+    try:
+        success = extract_and_store_document(file.name)
+        if success:
+            return f"✅ **{os.path.basename(file.name)}** successfully parsed and added to knowledge base!"
+        else:
+            return f"⚠️ Failed to extract text from **{os.path.basename(file.name)}**."
+    except Exception as e:
+        return f"❌ Error: {str(e)}"
 
-# Gradio UI with Video Clip
 with gr.Blocks() as demo:
-    gr.HTML("<h2 style='text-align: center;'>AI Tutor - We.</h2>")
+    gr.Markdown("# 🎓 Tutor AI (single agent with tool-routing capability)")
     
-    # Align image and video side by side
-    with gr.Row():
-        with gr.Column(scale=1):  # Adjust scale to control width ratio
-            gr.HTML("""
-                <div style="text-align: center; margin-bottom: 20px;">
-                    <img src="https://img.freepik.com/premium-photo/little-girl-is-seen-sitting-front-laptop-computer-engaged-with-nearby-robot-robot-assistant-helping-child-with-homework-ai-generated_585735-12266.jpg" 
-                    style="width: 100%; height: auto; border-radius: 10px; box-shadow: 0 4px 8px rgba(0,0,0,0.2);" />
-                </div>
-            """)
-        
-        #with gr.Column(scale=1):  # Adjust scale for equal width
-            gr.Video("https://github.com/lesterchia1/AI_tutor/raw/main/We%20not%20me%20video.mp4", label="Introduction Video")
+    with gr.Tab("AI Chatbot"):
+        chatbot = gr.Chatbot(type="messages", height=400)
+        with gr.Row():
+            msg = gr.Textbox(placeholder="Ask your tutor...", scale=4)
+            submit = gr.Button("Send", variant="primary")
+        # Manual audio control
+        with gr.Row():
+            play_audio_btn = gr.Button("🔊 Play Audio Response", variant="secondary")
+            audio_out = gr.Audio(label="Audio Response", autoplay=False)  # autoplay=False
         
-
-    # Add other UI elements below
-    with gr.Row():
-        with gr.Column():
-            audio_input = gr.Audio(type="numpy", label="Record Audio")
-            transcription_output = gr.Textbox(label="Transcription")
-            user_input = gr.Textbox(label="Ask a question")
-            chat_output = gr.Textbox(label="Response")
-            audio_output = gr.Audio(label="Audio Playback")
-            submit_btn = gr.Button("Ask")
-        with gr.Column():
-            file_upload = gr.File(label="Upload a document")
-            process_status = gr.Textbox(label="Processing Status")
-    
-    # Define button actions
-    submit_btn.click(chat_with_groq, inputs=user_input, outputs=[chat_output, audio_output])
-    audio_input.change(transcribe_audio, inputs=audio_input, outputs=transcription_output)
-    transcription_output.change(fn=lambda x: x, inputs=transcription_output, outputs=user_input)
-    file_upload.change(process_document, inputs=file_upload, outputs=process_status)
-
-    # Add cleanup function to be triggered periodically (e.g., every time a button is clicked or after certain actions)
-    #demo.load(lambda: cleanup_old_files(directory="./", age_limit=60), inputs=[], outputs=[])
-    demo.load(lambda: [], inputs=[], outputs=[])
-
-
-demo.launch()
+        # Chat submission
+        submit.click(chat_handler, [msg, chatbot], [chatbot, msg, audio_out])
+        msg.submit(chat_handler, [msg, chatbot], [chatbot, msg, audio_out])
+        # Manual audio generation
+        play_audio_btn.click(generate_audio, None, audio_out)
+
+    with gr.Tab("Upload Notes"):
+        file_input = gr.File(label="Upload PDF / DOCX / PPTX / TXT", file_types=[".pdf", ".docx", ".pptx", ".txt"])
+        upload_status = gr.Markdown()
+        file_input.change(upload_file, file_input, upload_status)
+
+demo.launch(share=True, debug=True)