Update app.py

app.py

@@ -1,38 +1,34 @@
 import os
+import json
 import time
 import threading
+import datetime
+import csv
+import docx
 import streamlit as st
-from twilio.rest import Client
-from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer
-import faiss
+from io import StringIO
 import numpy as np
-import docx
-from groq import Groq
 import requests
-from io import StringIO
+import pdfplumber
 from pdfminer.high_level import extract_text_to_fp
 from pdfminer.layout import LAParams
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer
+import faiss
+from twilio.rest import Client
 from twilio.base.exceptions import TwilioRestException
-import pdfplumber
-import datetime
-import csv
-import json
-import re
 
+# Start time for filtering incoming messages
 APP_START_TIME = datetime.datetime.now(datetime.timezone.utc)
 os.environ["PYTORCH_JIT"] = "0"
 
-# ---------------- PDF & DOCX & JSON Extraction ----------------
+# ---------------- PDF / DOCX / JSON LOADERS ----------------
 def _extract_tables_from_page(page):
     tables = page.extract_tables()
     formatted_tables = []
     for table in tables:
-        formatted_table = []
-        for row in table:
-            formatted_row = [cell if cell is not None else "" for cell in row]
-            formatted_table.append(formatted_row)
-        formatted_tables.append(formatted_table)
+        formatted_row = [[cell if cell else "" for cell in row] for row in table]
+        formatted_tables.append(formatted_row)
     return formatted_tables
 
 def extract_text_from_pdf(pdf_path):
@@ -46,55 +42,48 @@ def extract_text_from_pdf(pdf_path):
                 if text:
                     text_output.write(text + "\n\n")
     except Exception as e:
-        print(f"pdfplumber error: {e}")
+        print(f"[pdfplumber error] Falling back: {e}")
         with open(pdf_path, 'rb') as file:
             extract_text_to_fp(file, text_output, laparams=LAParams(), output_type='text')
     return text_output.getvalue(), all_tables
 
-def _format_tables_internal(tables):
-    formatted_tables_str = []
-    for table in tables:
-        with StringIO() as csvfile:
-            writer = csv.writer(csvfile)
-            writer.writerows(table)
-            formatted_tables_str.append(csvfile.getvalue())
-    return "\n\n".join(formatted_tables_str)
-
-def clean_extracted_text(text):
-    return '\n'.join(' '.join(line.strip().split()) for line in text.splitlines() if line.strip())
-
 def extract_text_from_docx(docx_path):
     try:
         doc = docx.Document(docx_path)
-        return '\n'.join(para.text for para in doc.paragraphs)
-    except:
+        return "\n".join(para.text for para in doc.paragraphs)
+    except Exception as e:
+        print(f"[DOCX error] {e}")
         return ""
 
 def load_json_data(json_path):
     try:
         with open(json_path, 'r', encoding='utf-8') as f:
             data = json.load(f)
+        if isinstance(data, list):
+            return "\n\n".join("\n".join(f"{k}: {v}" for k, v in d.items()) for d in data if isinstance(d, dict))
         if isinstance(data, dict):
-            # Flatten dictionary values (avoiding nested structures as strings)
-            return "\n".join(f"{key}: {value}" for key, value in data.items() if not isinstance(value, (dict, list)))
-        elif isinstance(data, list):
-            # Flatten list of dictionaries
-            all_items = []
-            for item in data:
-                if isinstance(item, dict):
-                    all_items.append("\n".join(f"{key}: {value}" for key, value in item.items() if not isinstance(value, (dict, list))))
-            return "\n\n".join(all_items)
-        else:
-            return json.dumps(data, ensure_ascii=False, indent=2)
+            return "\n".join(f"{k}: {v}" for k, v in data.items())
+        return str(data)
     except Exception as e:
-        print(f"JSON read error: {e}")
+        print(f"[JSON error] {e}")
         return ""
 
-# ---------------- Chunking ----------------
+def _format_tables_internal(tables):
+    formatted = []
+    for table in tables:
+        with StringIO() as csvfile:
+            writer = csv.writer(csvfile)
+            writer.writerows(table)
+            formatted.append(csvfile.getvalue())
+    return "\n\n".join(formatted)
+
+def clean_extracted_text(text):
+    return '\n'.join(' '.join(line.strip().split()) for line in text.splitlines() if line.strip())
+
+# ---------------- CHUNKING ----------------
 def chunk_text(text, tokenizer, chunk_size=128, chunk_overlap=32):
     tokens = tokenizer.tokenize(text)
-    chunks = []
-    start = 0
+    chunks, start = [], 0
     while start < len(tokens):
         end = min(start + chunk_size, len(tokens))
         chunk = tokens[start:end]
@@ -106,9 +95,9 @@ def chunk_text(text, tokenizer, chunk_size=128, chunk_overlap=32):
 def retrieve_chunks(question, index, embed_model, text_chunks, k=3):
     q_embedding = embed_model.encode(question)
     D, I = index.search(np.array([q_embedding]), k)
-    return [text_chunks[i] for i in I[0]]
+    return [text_chunks[i] for i in I[0] if i < len(text_chunks)]
 
-# ---------------- Groq Answer Generator ----------------
+# ---------------- GROQ ANSWER GENERATOR ----------------
 def generate_answer_with_groq(question, context):
     url = "https://api.groq.com/openai/v1/chat/completions"
     api_key = os.environ.get("GROQ_API_KEY")
@@ -116,131 +105,86 @@ def generate_answer_with_groq(question, context):
         "Authorization": f"Bearer {api_key}",
         "Content-Type": "application/json",
     }
+
     prompt = (
         f"Customer asked: '{question}'\n\n"
-        f"Here is the relevant information to help:\n{context}\n\n"
-        f"Respond in a friendly and helpful tone as a toy shop support agent, "
-        f"addressing the customer by their name if it's available in the context."
+        f"Here is the relevant information:\n{context}\n\n"
+        "Reply as a friendly toy shop assistant, include customer name from the context if possible."
     )
+
     payload = {
         "model": "llama3-8b-8192",
         "messages": [
             {
                 "role": "system",
-                "content": (
-                    "You are ToyBot, a friendly WhatsApp assistant for an online toy shop. "
-                    "Help customers with toys, delivery, and returns in a helpful tone. "
-                    "When responding, try to find the customer's name in the provided context "
-                    "and address them directly. If the context contains order details and status, "
-                    "include that information in your response."
-                )
+                "content": "You are ToyBot, a helpful assistant for an online toy shop."
             },
-            {"role": "user", "content": prompt},
-        ],
-        "temperature": 0.5,
-        "max_tokens": 300,
+            {"role": "user", "content": prompt}
+        ]
     }
-    response = requests.post(url, headers=headers, json=payload)
-    response.raise_for_status()
-    return response.json()['choices'][0]['message']['content'].strip()
 
-# ---------------- Twilio Integration ----------------
-def fetch_latest_incoming_message(client, conversation_sid):
-    try:
-        messages = client.conversations.v1.conversations(conversation_sid).messages.list()
-        for msg in reversed(messages):
-            if msg.author.startswith("whatsapp:"):
-                return {
-                    "sid": msg.sid,
-                    "body": msg.body,
-                    "author": msg.author,
-                    "timestamp": msg.date_created,
-                }
-    except TwilioRestException as e:
-        print(f"Twilio error: {e}")
-    return None
-
-def send_twilio_message(client, conversation_sid, body):
-    return client.conversations.v1.conversations(conversation_sid).messages.create(
-        author="system", body=body
-    )
+    response = requests.post(url, headers=headers, json=payload)
+    return response.json()["choices"][0]["message"]["content"]
+
+# ---------------- TWILIO MONITOR ----------------
+def handle_incoming_messages(index, embed_model, tokenizer, text_chunks):
+    client = Client(os.environ["TWILIO_ACCOUNT_SID"], os.environ["TWILIO_AUTH_TOKEN"])
+    conversation_sid = os.environ.get("TWILIO_CONVERSATION_SID")
+
+    if not conversation_sid:
+        print("❌ TWILIO_CONVERSATION_SID not set")
+        return
+
+    last_check_time = APP_START_TIME
+
+    while True:
+        try:
+            messages = client.conversations.conversations(conversation_sid).messages.list(order="asc")
+            for msg in messages:
+                msg_time = msg.date_created.replace(tzinfo=datetime.timezone.utc)
+                if msg_time > last_check_time:
+                    print(f"📩 New message from {msg.author}: {msg.body}")
+                    answer = generate_answer_with_groq(msg.body, "\n".join(retrieve_chunks(msg.body, index, embed_model, text_chunks)))
+                    client.conversations.conversations(conversation_sid).messages.create(author="ToyBot", body=answer)
+            last_check_time = datetime.datetime.now(datetime.timezone.utc)
+        except TwilioRestException as e:
+            print(f"[Twilio error] {e}")
+        time.sleep(10)
+
+# ---------------- STREAMLIT UI ----------------
+st.title("🎁 ToyShop Assistant – RAG WhatsApp Bot")
+
+uploaded_files = st.file_uploader("📎 Upload your documents (PDF, DOCX, JSON)", accept_multiple_files=True)
+if uploaded_files:
+    full_text = ""
+    all_tables = []
 
-# ---------------- Knowledge Base Setup ----------------
-def setup_knowledge_base():
-    folder_path = "docs"
-    all_text = ""
-
-    for filename in os.listdir(folder_path):
-        file_path = os.path.join(folder_path, filename)
-        if filename.endswith(".pdf"):
-            text, tables = extract_text_from_pdf(file_path)
-            all_text += clean_extracted_text(text) + "\n"
-            all_text += _format_tables_internal(tables) + "\n"
-        elif filename.endswith(".docx"):
-            text = extract_text_from_docx(file_path)
-            all_text += clean_extracted_text(text) + "\n"
-        elif filename.endswith(".json"):
-            text = load_json_data(file_path)
-            all_text += text + "\n"
-        elif filename.endswith(".csv"):
-            try:
-                with open(file_path, newline='', encoding='utf-8') as csvfile:
-                    reader = csv.DictReader(csvfile)
-                    for row in reader:
-                        line = ' | '.join(f"{k}: {v}" for k, v in row.items())
-                        all_text += line + "\n"
-            except Exception as e:
-                print(f"CSV read error: {e}")
-
-    tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
-    chunks = chunk_text(all_text, tokenizer)
-    model = SentenceTransformer('all-mpnet-base-v2')
-    embeddings = model.encode(chunks, show_progress_bar=False)
-    dim = embeddings[0].shape[0]
-    index = faiss.IndexFlatL2(dim)
-    index.add(np.array(embeddings).astype('float32'))
-    return index, model, chunks
-
-# ---------------- Monitor Twilio Conversations ----------------
-def start_conversation_monitor(client, index, embed_model, text_chunks):
-    processed_convos = set()
-    last_processed_timestamp = {}
-
-    def poll_convo(convo_sid):
-        while True:
-            latest_msg = fetch_latest_incoming_message(client, convo_sid)
-            if latest_msg:
-                msg_time = latest_msg["timestamp"]
-                if convo_sid not in last_processed_timestamp or msg_time > last_processed_timestamp[convo_sid]:
-                    last_processed_timestamp[convo_sid] = msg_time
-                    question = latest_msg["body"]
-                    sender = latest_msg["author"]
-                    print(f"📩 New message from {sender}: {question}")
-                    context = "\n\n".join(retrieve_chunks(question, index, embed_model, text_chunks))
-                    answer = generate_answer_with_groq(question, context)
-                    send_twilio_message(client, convo_sid, answer)
-            time.sleep(5)
-
-    for convo in client.conversations.v1.conversations.list():
-        if convo.sid not in processed_convos:
-            processed_convos.add(convo.sid)
-            threading.Thread(target=poll_convo, args=(convo.sid,), daemon=True).start()
-
-# ---------------- Main Entry ----------------
-if __name__ == "__main__":
-    st.title("🤖 ToyBot WhatsApp Assistant")
-    st.write("Initializing knowledge base...")
-
-    index, model, chunks = setup_knowledge_base()
-
-    st.success("Knowledge base loaded.")
-    st.write("Waiting for WhatsApp messages...")
-
-    account_sid = os.environ.get("TWILIO_ACCOUNT_SID")
-    auth_token = os.environ.get("TWILIO_AUTH_TOKEN")
-    if not account_sid or not auth_token:
-        st.error("❌ Twilio credentials not set.")
-    else:
-        client = Client(account_sid, auth_token)
-        start_conversation_monitor(client, index, model, chunks)
-        st.info("✅ Bot is now monitoring Twilio conversations.")
+    for file in uploaded_files:
+        ext = file.name.lower().split(".")[-1]
+        if ext == "pdf":
+            text, tables = extract_text_from_pdf(file)
+            all_tables.extend(tables)
+        elif ext == "docx":
+            text = extract_text_from_docx(file)
+        elif ext == "json":
+            text = load_json_data(file)
+        else:
+            text = file.read().decode("utf-8")
+        full_text += clean_extracted_text(text) + "\n\n"
+
+    # Load models
+    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+    embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+    chunks = chunk_text(full_text, tokenizer)
+    embeddings = embed_model.encode(chunks)
+    
+    index = faiss.IndexFlatL2(embeddings.shape[1])
+    index.add(np.array(embeddings))
+
+    # Start listener thread
+    if "listener_started" not in st.session_state:
+        threading.Thread(target=handle_incoming_messages, args=(index, embed_model, tokenizer, chunks), daemon=True).start()
+        st.session_state.listener_started = True
+        st.success("✅ WhatsApp listener started.")
+
+    st.success(f"📚 Knowledge base built with {len(chunks)} chunks")