Update app.py

app.py

@@ -1,204 +1,164 @@
 import logging
+import os
 import gradio as gr
 from transformers import pipeline
 from sentence_transformers import SentenceTransformer, util
 import PyPDF2
-import os
 
-# Set up logging with a dedicated file handler
-logger = logging.getLogger('SupportBot')
-logger.setLevel(logging.INFO)
-if logger.handlers:
-    logger.handlers.clear()
-
-# Define log file path in a writable directory (/tmp)
-log_file_path = '/tmp/support_bot_log.txt'
-
-# Create a file handler with append mode
-file_handler = logging.FileHandler(log_file_path, mode='a')
-file_handler.setLevel(logging.INFO)
-formatter = logging.Formatter('%(asctime)s - %(message)s')
-file_handler.setFormatter(formatter)
-logger.addHandler(file_handler)
-
-# Add a stream handler to output logs to the console as well
-stream_handler = logging.StreamHandler()
-stream_handler.setLevel(logging.INFO)
-stream_handler.setFormatter(formatter)
-logger.addHandler(stream_handler)
-
-# Load models
-qa_model = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")
-embedder = SentenceTransformer('all-MiniLM-L6-v2')
-
-# Helper function to extract text from a PDF
-def extract_text_from_pdf(file_path):
-    text = ""
-    with open(file_path, "rb") as file:
-        pdf_reader = PyPDF2.PdfReader(file)
-        for page in pdf_reader.pages:
-            extracted_text = page.extract_text()
-            if extracted_text:
-                text += extracted_text + "\n"
-    return text
-
-# Find the most relevant section in the document
-def find_relevant_section(query, sections, section_embeddings):
-    stopwords = {"and", "the", "is", "for", "to", "a", "an", "of", "in", "on", "at", "with", "by", "it", "as", "so", "what"}
-    
-    logger.info(f"Searching for relevant section for query: {query}")
-    query_embedding = embedder.encode(query, convert_to_tensor=True)
-    similarities = util.cos_sim(query_embedding, section_embeddings)[0]
-    best_idx = similarities.argmax().item()
-    best_section = sections[best_idx]
-    similarity_score = similarities[best_idx].item()
-    
-    SIMILARITY_THRESHOLD = 0.4
-    if similarity_score >= SIMILARITY_THRESHOLD:
-        logger.info(f"Found relevant section using embeddings (score: {similarity_score})")
-        file_handler.flush()  # Ensure log is written immediately
-        return best_section
-    
-    logger.info(f"Low similarity ({similarity_score}). Falling back to keyword search.")
-    query_words = {word for word in query.lower().split() if word not in stopwords}
-    for section in sections:
-        section_words = {word for word in section.lower().split() if word not in stopwords}
-        common_words = query_words.intersection(section_words)
-        if len(common_words) >= 2:
-            logger.info(f"Keyword match found with common words: {common_words}")
-            file_handler.flush()
-            return section
-    
-    logger.info("No good match found. Returning default response.")
-    file_handler.flush()
-    return "I don’t have enough information to answer that."
-
-# Process the uploaded file
-def process_file(file, state):
-    logger.info("Received file upload request")
-    if file is None:
-        logger.info("No file uploaded")
-        file_handler.flush()
-        return [("Bot", "Please upload a file.")], state
-
-    # Save the uploaded file to a temporary location
-    file_path = file.name
-    temp_file_path = os.path.join("/tmp", os.path.basename(file_path))
-    with open(temp_file_path, "wb") as f:
-        # Check if the file has a 'read' method; if not, assume it's already the content.
-        if hasattr(file, "read"):
-            content = file.read()
+# Set up logging: we write logs to /tmp so that it's writable on Spaces.
+log_file_path = "/tmp/support_bot_log.txt"
+logging.basicConfig(filename=log_file_path, level=logging.INFO, format='%(asctime)s - %(message)s')
+
+class SupportBotAgent:
+    def __init__(self, document_path):
+        # Load a pre-trained question-answering model
+        self.qa_model = pipeline("question-answering", model="distilbert-base-uncased-distilled-squad")
+        # Set up an embedding model for finding relevant sections
+        self.embedder = SentenceTransformer('all-MiniLM-L6-v2')
+        # Load the document text and split it into sections (by paragraphs)
+        self.document_text = self.load_document(document_path)
+        self.sections = self.document_text.split('\n\n')
+        self.section_embeddings = self.embedder.encode(self.sections, convert_to_tensor=True)
+        logging.info(f"Loaded document: {document_path}")
+
+    def load_document(self, path):
+        """Loads and extracts text from a TXT or PDF file."""
+        if path.lower().endswith(".txt"):
+            file_type = "Text File"
+            with open(path, 'r', encoding='utf-8') as file:
+                text = file.read()
+        elif path.lower().endswith(".pdf"):
+            file_type = "PDF File"
+            text = ""
+            with open(path, "rb") as file:
+                pdf_reader = PyPDF2.PdfReader(file)
+                for page in pdf_reader.pages:
+                    page_text = page.extract_text()
+                    if page_text:
+                        text += page_text + "\n"
         else:
-            content = file
-        if isinstance(content, str):
-            content = content.encode("utf-8")
-        f.write(content)
-    
-    if temp_file_path.lower().endswith(".pdf"):
-        logger.info(f"Processing PDF file: {temp_file_path}")
-        text = extract_text_from_pdf(temp_file_path)
-    elif temp_file_path.lower().endswith(".txt"):
-        logger.info(f"Processing TXT file: {temp_file_path}")
-        with open(temp_file_path, 'r', encoding='utf-8') as f:
-            text = f.read()
-    else:
-        logger.error(f"Unsupported file format: {temp_file_path}")
-        file_handler.flush()
-        return [("Bot", "Unsupported file format. Please upload a PDF or TXT file.")], state
-    
-    sections = text.split('\n\n')
-    section_embeddings = embedder.encode(sections, convert_to_tensor=True)
-    state['document_text'] = text
-    state['sections'] = sections
-    state['section_embeddings'] = section_embeddings
-    state['current_query'] = None
-    state['feedback_count'] = 0
-    state['mode'] = 'waiting_for_query'
-    state['chat_history'] = [("Bot", "File processed. You can now ask questions.")]
-    logger.info(f"File processed successfully: {temp_file_path}")
-    file_handler.flush()
-    return state['chat_history'], state
-
-# Handle user input (queries and feedback)
-def handle_input(user_input, state):
-    if state['mode'] == 'waiting_for_upload':
-        logger.info("User input received before file upload")
-        state['chat_history'].append(("Bot", "Please upload a file first."))
-        file_handler.flush()
-    elif state['mode'] == 'waiting_for_query':
-        query = user_input
-        logger.info(f"User query: {query}")
-        state['current_query'] = query
-        state['feedback_count'] = 0
-        context = find_relevant_section(query, state['sections'], state['section_embeddings'])
-        if context == "I don’t have enough information to answer that.":
-            answer = context
+            file_type = "Unsupported Format"
+            logging.error(f"Unsupported file format: {path}")
+            raise ValueError("Unsupported file format. Please provide a TXT or PDF file.")
+        logging.info(f"Loaded {file_type}: {path}")
+        return text
+
+    def find_relevant_section(self, query):
+        """
+        First uses semantic similarity. If similarity is too low, falls back to a keyword search.
+        """
+        stopwords = {"and", "the", "is", "for", "to", "a", "an", "of", "in", "on", "at", "with", "by", "it", "as", "so", "what"}
+        query_embedding = self.embedder.encode(query, convert_to_tensor=True)
+        similarities = util.cos_sim(query_embedding, self.section_embeddings)[0]
+        best_idx = similarities.argmax().item()
+        best_section = self.sections[best_idx]
+        similarity_score = similarities[best_idx].item()
+        SIMILARITY_THRESHOLD = 0.4
+
+        if similarity_score >= SIMILARITY_THRESHOLD:
+            logging.info(f"Found relevant section using embeddings for query: {query}")
+            return best_section
+
+        logging.info(f"Low similarity ({similarity_score}). Falling back to keyword search.")
+        query_words = {word for word in query.lower().split() if word not in stopwords}
+        for section in self.sections:
+            section_words = {word for word in section.lower().split() if word not in stopwords}
+            common_words = query_words.intersection(section_words)
+            if len(common_words) >= 2:
+                logging.info(f"Keyword match found for query: {query} with common words: {common_words}")
+                return section
+
+        logging.info(f"No good keyword match found. Returning default fallback response.")
+        return "I don’t have enough information to answer that."
+
+    def answer_query(self, query):
+        context = self.find_relevant_section(query)
+        if not context:
+            answer = "I don’t have enough information to answer that."
         else:
-            result = qa_model(question=query, context=context)
+            result = self.qa_model(question=query, context=context, max_answer_len=50)
             answer = result["answer"]
-        state['last_answer'] = answer
-        state['mode'] = 'waiting_for_feedback'
-        state['chat_history'].append(("User", query))
-        state['chat_history'].append(("Bot", f"Answer: {answer}\nPlease provide feedback: good, too vague, not helpful."))
-        logger.info(f"Generated answer: {answer}")
-        file_handler.flush()
-    elif state['mode'] == 'waiting_for_feedback':
+        logging.info(f"Answer for query '{query}': {answer}")
+        return answer
+
+    def adjust_response(self, query, response, feedback):
+        """Modify the response based on user feedback."""
+        if feedback == "too vague":
+            context = self.find_relevant_section(query)
+            adjusted_response = f"{response}\n\n(More details:\n{context[:500]}...)"
+        elif feedback == "not helpful":
+            adjusted_response = self.answer_query(query + " Please provide more detailed information with examples.")
+        else:
+            adjusted_response = response
+        logging.info(f"Adjusted answer for query '{query}': {adjusted_response}")
+        return adjusted_response
+
+# --- Gradio Functions and App Workflow ---
+
+def process_file(file, state):
+    """Handles the file upload and initializes the SupportBotAgent."""
+    if file is None:
+        logging.info("No file uploaded")
+        return [("Bot", "Please upload a TXT or PDF file.")], state
+    # Save the uploaded file to /tmp
+    temp_path = os.path.join("/tmp", file.name)
+    with open(temp_path, "wb") as f:
+        f.write(file.read())
+    try:
+        state["agent"] = SupportBotAgent(temp_path)
+    except Exception as e:
+        return [("Bot", f"Error processing file: {str(e)}")], state
+    state["chat_history"] = [("Bot", "File loaded successfully. Enter your query (or type 'exit' to end):")]
+    state["mode"] = "query"
+    state["last_query"] = ""
+    state["last_answer"] = ""
+    state["feedback_count"] = 0
+    return state["chat_history"], state
+
+def process_input(user_input, state):
+    """
+    Processes user input as either a query or feedback based on the current mode.
+    Typing 'exit' stops the session.
+    """
+    if state.get("mode", "query") == "ended":
+        return state["chat_history"], state
+    if user_input.lower() == "exit":
+        state["chat_history"].append(("Bot", "Session ended. You may now download the log file."))
+        state["mode"] = "ended"
+        return state["chat_history"], state
+    if state["mode"] == "query":
+        state["last_query"] = user_input
+        answer = state["agent"].answer_query(user_input)
+        state["last_answer"] = answer
+        state["feedback_count"] = 0
+        state["chat_history"].append(("User", user_input))
+        state["chat_history"].append(("Bot", f"Answer: {answer}\nPlease provide feedback (good, too vague, not helpful):"))
+        state["mode"] = "feedback"
+    elif state["mode"] == "feedback":
         feedback = user_input.lower()
-        logger.info(f"User feedback: {feedback}")
-        state['chat_history'].append(("User", feedback))
-        if feedback == "good" or state['feedback_count'] >= 2:
-            state['mode'] = 'waiting_for_query'
-            if feedback == "good":
-                state['chat_history'].append(("Bot", "Thank you for your feedback. You can ask another question."))
-                logger.info("Feedback 'good' received. Ready for next query.")
-            else:
-                state['chat_history'].append(("Bot", "Maximum feedback iterations reached. You can ask another question."))
-                logger.info("Max feedback iterations (2) reached. Ready for next query.")
-            file_handler.flush()
+        state["chat_history"].append(("User", feedback))
+        if feedback == "good" or state["feedback_count"] >= 1:
+            state["chat_history"].append(("Bot", "Thank you for your feedback. Enter your next query (or type 'exit' to end):"))
+            state["mode"] = "query"
         else:
-            query = state['current_query']
-            context = find_relevant_section(query, state['sections'], state['section_embeddings'])
-            if feedback == "too vague":
-                adjusted_answer = f"{state['last_answer']}\n\n(More details:\n{context[:500]}...)"
-                logger.info("Feedback 'too vague'. Providing context.")
-            elif feedback == "not helpful":
-                adjusted_answer = qa_model(question=query + " Please provide more detailed information with examples.", context=context)['answer']
-                logger.info("Feedback 'not helpful'. Re-searching with modified query.")
-            else:
-                state['chat_history'].append(("Bot", "Please provide valid feedback: good, too vague, not helpful."))
-                logger.info(f"Invalid feedback received: {feedback}")
-                file_handler.flush()
-                return state['chat_history'], state
-            state['last_answer'] = adjusted_answer
-            state['feedback_count'] += 1
-            state['chat_history'].append(("Bot", f"Updated answer: {adjusted_answer}\nPlease provide feedback: good, too vague, not helpful."))
-            logger.info(f"Updated answer: {adjusted_answer}")
-            file_handler.flush()
-    return state['chat_history'], state
-
-# Initial state
-initial_state = {
-    'document_text': None,
-    'sections': None,
-    'section_embeddings': None,
-    'current_query': None,
-    'feedback_count': 0,
-    'mode': 'waiting_for_upload',
-    'chat_history': [("Bot", "Please upload a PDF or TXT file to start.")],
-    'last_answer': None
-}
-
-# Gradio interface
+            new_answer = state["agent"].adjust_response(state["last_query"], state["last_answer"], feedback)
+            state["last_answer"] = new_answer
+            state["feedback_count"] += 1
+            state["chat_history"].append(("Bot", f"Updated Answer: {new_answer}\nPlease provide feedback (good, too vague, not helpful):"))
+    return state["chat_history"], state
+
+# --- Gradio UI Setup ---
+
 with gr.Blocks() as demo:
-    state = gr.State(initial_state)
-    file_upload = gr.File(label="Upload PDF or TXT file")
+    state = gr.State({"mode": "idle"})
+    gr.Markdown("## Customer Support Bot with Document Training")
+    file_upload = gr.File(label="Upload TXT or PDF file")
     chat = gr.Chatbot()
-    user_input = gr.Textbox(label="Your query or feedback")
+    user_input = gr.Textbox(label="Enter your query or feedback")
     submit_btn = gr.Button("Submit")
-    # Point the log file download to the writable log file path
-    log_file = gr.File(label="Download Log File", value=log_file_path)
+    log_file = gr.File(label="Download Log File", file_count="single", interactive=False, value=log_file_path)
 
     file_upload.upload(process_file, inputs=[file_upload, state], outputs=[chat, state])
-    submit_btn.click(handle_input, inputs=[user_input, state], outputs=[chat, state]).then(lambda: "", None, user_input)
+    submit_btn.click(process_input, inputs=[user_input, state], outputs=[chat, state]).then(lambda: "", None, user_input)
 
-demo.launch(share=True)
+demo.launch(share=True)