Update app.py

app.py

@@ -22,99 +22,92 @@ APP_START_TIME = datetime.datetime.now(datetime.timezone.utc)
 
 os.environ["PYTORCH_JIT"] = "0"
 
-# --- PDF Extraction ---
+# ---------------- PDF & DOCX & JSON Extraction ----------------
 def _extract_tables_from_page(page):
-    """Extracts tables from a single page of a PDF."""
-
     tables = page.extract_tables()
-    if not tables:
-        return []
-
     formatted_tables = []
     for table in tables:
         formatted_table = []
         for row in table:
-            if row:  # Filter out empty rows
-                formatted_row = [cell if cell is not None else "" for cell in row]  # Replace None with ""
-                formatted_table.append(formatted_row)
-            else:
-                formatted_table.append([""])  # Append an empty row if the row is None
+            formatted_row = [cell if cell is not None else "" for cell in row]
+            formatted_table.append(formatted_row)
         formatted_tables.append(formatted_table)
     return formatted_tables
-    
+
 def extract_text_from_pdf(pdf_path):
     text_output = StringIO()
     all_tables = []
     try:
         with pdfplumber.open(pdf_path) as pdf:
             for page in pdf.pages:
-                # Extract tables
-                page_tables = _extract_tables_from_page(page)
-                if page_tables:
-                    all_tables.extend(page_tables)
-                # Extract text
+                all_tables.extend(_extract_tables_from_page(page))
                 text = page.extract_text()
                 if text:
                     text_output.write(text + "\n\n")
     except Exception as e:
-        print(f"Error extracting with pdfplumber: {e}")
-        # Fallback to pdfminer if pdfplumber fails
+        print(f"pdfplumber error: {e}")
         with open(pdf_path, 'rb') as file:
-            extract_text_to_fp(file, text_output, laparams=LAParams(), output_type='text', codec=None)
-    extracted_text = text_output.getvalue()
-    return extracted_text, all_tables  # Return text and list of tables
-
-def clean_extracted_text(text):
-    lines = text.splitlines()
-    cleaned = []
-    for line in lines:
-        line = line.strip()
-        if line:
-            line = ' '.join(line.split())
-            cleaned.append(line)
-    return '\n'.join(cleaned)
+            extract_text_to_fp(file, text_output, laparams=LAParams(), output_type='text')
+    return text_output.getvalue(), all_tables
 
 def _format_tables_internal(tables):
-    """Formats extracted tables into a string representation."""
-
     formatted_tables_str = []
     for table in tables:
-        # Use csv writer to handle commas and quotes correctly
         with StringIO() as csvfile:
-            csvwriter = csv.writer(csvfile)
-            csvwriter.writerows(table)
+            writer = csv.writer(csvfile)
+            writer.writerows(table)
             formatted_tables_str.append(csvfile.getvalue())
     return "\n\n".join(formatted_tables_str)
 
-# --- DOCX Extraction ---
+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 Exception:
+    except:
         return ""
 
-# --- Chunking ---
-def chunk_text(text, tokenizer, chunk_size=128, chunk_overlap=32, max_tokens=512):
+def load_json_data(json_path):
+    try:
+        with open(json_path, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        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)
+    except Exception as e:
+        print(f"JSON read error: {e}")
+        return ""
+
+# ---------------- Chunking ----------------
+def chunk_text(text, tokenizer, chunk_size=128, chunk_overlap=32):
     tokens = tokenizer.tokenize(text)
     chunks = []
     start = 0
     while start < len(tokens):
         end = min(start + chunk_size, len(tokens))
-        chunk_tokens = tokens[start:end]
-        chunk_text = tokenizer.convert_tokens_to_string(chunk_tokens)
-        chunks.append(chunk_text)
-        if end == len(tokens):
-            break
+        chunk = tokens[start:end]
+        chunks.append(tokenizer.convert_tokens_to_string(chunk))
+        if end == len(tokens): break
         start += chunk_size - chunk_overlap
     return chunks
 
 def retrieve_chunks(question, index, embed_model, text_chunks, k=3):
-    question_embedding = embed_model.encode(question)
-    D, I = index.search(np.array([question_embedding]), k)
+    q_embedding = embed_model.encode(question)
+    D, I = index.search(np.array([q_embedding]), k)
     return [text_chunks[i] for i in I[0]]
 
-# --- 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")
@@ -124,8 +117,9 @@ def generate_answer_with_groq(question, context):
     }
     prompt = (
         f"Customer asked: '{question}'\n\n"
-        f"Here is the relevant product or policy info to help:\n{context}\n\n"
-        f"Respond in a friendly and helpful tone as a toy shop support agent."
+        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."
     )
     payload = {
         "model": "llama3-8b-8192",
@@ -133,9 +127,11 @@ def generate_answer_with_groq(question, context):
             {
                 "role": "system",
                 "content": (
-                    "You are ToyBot, a friendly and helpful WhatsApp assistant for an online toy shop. "
-                    "Your goal is to politely answer customer questions, help them choose the right toys, "
-                    "provide order or delivery information, explain return policies, and guide them through purchases."
+                    "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."
                 )
             },
             {"role": "user", "content": prompt},
@@ -175,46 +171,37 @@ def send_twilio_message(client, conversation_sid, body):
         author="system", body=body
     )
 
-# --- Load Knowledge Base ---
+# ---------------- Knowledge Base Setup ----------------
 def setup_knowledge_base():
     folder_path = "docs"
     all_text = ""
 
-    # Process PDFs
-    for filename in ["FAQ.pdf", "ProductReturnPolicy.pdf"]:
-        pdf_path = os.path.join(folder_path, filename)
-        text, tables = extract_text_from_pdf(pdf_path)
-        all_text += clean_extracted_text(text) + "\n"
-        all_text += _format_tables_internal(tables) + "\n"
-
-    # Process CSVs
-    for filename in ["CustomerOrders.csv"]:
-        csv_path = os.path.join(folder_path, filename)
-        try:
-            with open(csv_path, newline='', encoding='utf-8') as csvfile:
-                reader = csv.DictReader(csvfile)
-                for row in reader:
-                    line = f"Order ID: {row.get('OrderID')} | Customer Name: {row.get('CustomerName')} | Order Date: {row.get('OrderDate')} | ProductID: {row.get('ProductID')} | Date: {row.get('OrderDate')} | Quantity: {row.get('Quantity')} | UnitPrice(USD): {row.get('UnitPrice(USD)')} | TotalPrice(USD): {row.get('TotalPrice(USD)')} | ShippingAddress: {row.get('ShippingAddress')} | OrderStatus: {row.get('OrderStatus')}"
-                    all_text += line + "\n"
-        except Exception as e:
-            print(f"❌ Error reading {filename}: {e}")
-
-    for filename in ["Products.csv"]:
-        csv_path = os.path.join(folder_path, filename)
-        try:
-            with open(csv_path, newline='', encoding='utf-8') as csvfile:
-                reader = csv.DictReader(csvfile)
-                for row in reader:
-                    line = f"Product ID: {row.get('ProductID')} | Toy Name: {row.get('ToyName')} | Category: {row.get('Category')} | Price(USD): {row.get('Price(USD)')} | Stock Quantity: {row.get('StockQuantity')} | Description: {row.get('Description')}"
-                    all_text += line + "\n"
-        except Exception as e:
-            print(f"❌ Error reading {filename}: {e}")
+    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}")
 
-    # Tokenization & chunking
     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, truncation=True, max_length=512)
+    embeddings = model.encode(chunks, show_progress_bar=False)
     dim = embeddings[0].shape[0]
     index = faiss.IndexFlatL2(dim)
     index.add(np.array(embeddings).astype('float32'))
@@ -274,7 +261,10 @@ def start_conversation_monitor(client, index, embed_model, text_chunks):
 # --- Streamlit UI ---
 st.set_page_config(page_title="Quasa – Al-Powered WhatsApp Chatbot", layout="wide")
 st.title("📱 Quasa – Al-Powered WhatsApp Chatbot")
+index, model, chunks = setup_knowledge_base()
 
+st.success("Knowledge base loaded.")
+#st.write("Waiting for WhatsApp messages...")
 account_sid = st.secrets.get("TWILIO_SID")
 auth_token = st.secrets.get("TWILIO_TOKEN")
 GROQ_API_KEY = st.secrets.get("GROQ_API_KEY")