Update app.py

app.py

@@ -1,11 +1,12 @@
 import os
 import random
+import re
 import gradio as gr
 from datetime import datetime
 from transformers import pipeline
 from simple_salesforce import Salesforce, SalesforceLogin
 from dotenv import load_dotenv
-import xml.etree.ElementTree as ET  # NEW: for parsing metadata
+import xml.etree.ElementTree as ET
 
 # ---------- Load Environment Variables ----------
 load_dotenv()
@@ -35,7 +36,7 @@ severities = {
     "Best Practice": "Low"
 }
 
-# ---------- Knowledge Base ----------
+# ---------- Mock Knowledge Base ----------
 salesforce_knowledge_base = {
     "governor limits soql": "In Salesforce, the governor limit for SOQL queries is 100 per synchronous transaction and 200 per asynchronous transaction.",
     "governor limits dml": "The governor limit for DML statements is 150 per transaction.",
@@ -104,7 +105,22 @@ def analyze_code(code):
 
     return issue_type, suggestion, severity
 
-# ---------- Metadata Validator (Updated with dynamic XML parsing) ----------
+# ---------- Helpers for Metadata XML ----------
+def _strip_code_fences(text: str) -> str:
+    # Remove ```xml ... ``` or plain ``` ... ```
+    t = re.sub(r'^\s*```(?:xml)?\s*', '', text, flags=re.IGNORECASE)
+    t = re.sub(r'\s*```\s*$', '', t)
+    return t
+
+def _escape_bare_ampersands(text: str) -> str:
+    # Replace & that aren't valid entities with &
+    return re.sub(
+        r'&(?!amp;|lt;|gt;|apos;|quot;|#\d+;|#x[0-9a-fA-F]+;)',
+        '&',
+        text
+    )
+
+# ---------- Metadata Validator ----------
 def validate_metadata(metadata, admin_id=None):
     if not metadata.strip():
         return "No metadata provided.", "", ""
@@ -113,22 +129,54 @@ def validate_metadata(metadata, admin_id=None):
     issue = "Unknown"
     recommendation = "No recommendation found."
 
-    try:
-        root = ET.fromstring(metadata)
-        # Detect missing <description> tag
-        description_found = any(elem.tag.endswith('description') for elem in root)
+    # Preview to logs for debugging
+    preview = metadata[:100].replace("\n", "\\n")
+    log_to_console({"preview": preview}, "Metadata Input Preview")
 
-        if not description_found:
-            issue = "Missing description"
-            recommendation = "Add a meaningful <description> to improve maintainability and clarity."
-        else:
-            issue = "Unused field detected"
-            recommendation = "Remove it to improve performance or document its purpose."
+    # Clean typical paste artefacts
+    cleaned = _strip_code_fences(metadata).lstrip("\ufeff").strip()
 
-    except Exception as e:
-        issue = "Invalid XML"
-        recommendation = f"Could not parse metadata XML. Error: {str(e)}"
+    # Quick non-XML guard
+    if "<" not in cleaned or ">" not in cleaned:
+        issue = "Invalid Input (Not XML)"
+        recommendation = (
+            "Please paste Salesforce metadata XML (starts with '<'). Example:\n"
+            "<CustomField xmlns=\"http://soap.sforce.com/2006/04/metadata\">...</CustomField>"
+        )
+    else:
+        # If leading junk before first '<', cut it
+        if cleaned and not cleaned.lstrip().startswith("<"):
+            idx = cleaned.find("<")
+            if idx != -1:
+                cleaned = cleaned[idx:].strip()
+
+        # Escape stray ampersands
+        cleaned = _escape_bare_ampersands(cleaned)
+
+        # Try to parse
+        try:
+            root = ET.fromstring(cleaned)
+            # Find <description> regardless of namespace
+            has_description = any(
+                elem.tag.split('}')[-1].lower() == "description"
+                for elem in root.iter()
+            )
+
+            if not has_description:
+                issue = "Missing description"
+                recommendation = "Add a meaningful <description> to improve maintainability and clarity."
+            else:
+                issue = "Unused field detected"
+                recommendation = "Remove it to improve performance or document its purpose."
+
+        except ET.ParseError as pe:
+            issue = "Invalid XML"
+            recommendation = f"Could not parse metadata XML. Error: {pe}."
+        except Exception as e:
+            issue = "Validation Error"
+            recommendation = f"Unexpected error while validating metadata: {str(e)}"
 
+    # Log to Salesforce
     log_data = {
         "Name": f"MetadataLog_{mtype}",
         "MetadataType__c": mtype,
@@ -156,7 +204,7 @@ def validate_metadata(metadata, admin_id=None):
 
     return mtype, issue, recommendation
 
-# ---------- Salesforce Chatbot ----------
+# ---------- Salesforce Chatbot (Updated Prompt) ----------
 conversation_history = []
 
 def salesforce_chatbot(query, history=[]):
@@ -166,11 +214,11 @@ def salesforce_chatbot(query, history=[]):
 
     salesforce_keywords = [
         "apex", "soql", "trigger", "lwc", "aura", "visualforce", "salesforce", "governor limits",
-        "dml", "metadata", "batch apex", "queueable", "future method", "api", "sfdc"
+        "dml", "metadata", "batch apex", "queueable", "future method", "api", "sfdc", "heap", "limits"
     ]
 
     if not any(keyword.lower() in query.lower() for keyword in salesforce_keywords):
-        return "Please ask a Salesforce-related question."
+        return "Please ask a Salesforce-related question (e.g., Apex, SOQL, LWC, limits, etc)."
 
     query_key = query.lower().strip()
     for kb_key, kb_answer in salesforce_knowledge_base.items():
@@ -181,23 +229,31 @@ def salesforce_chatbot(query, history=[]):
             return kb_answer
 
     history_summary = "\n".join([f"User: {q}\nAssistant: {a}" for q, a in conversation_history[-4:]])
-    prompt = f"""
-You are an expert Salesforce developer...
 
+    prompt = f"""
+You are an expert Salesforce developer and certified architect. You provide 100% accurate, clear, and practical answers for topics like Apex, SOQL, LWC, governor limits, triggers, metadata, and more.
+When answering:
+- ALWAYS give at least 2 lines of explanation.
+- Be clear, concise, and technically correct.
+- Mention official Salesforce limits if applicable.
+- Use bullet points or code snippets when helpful.
+- Avoid speculation — if unknown, say so and suggest Trailhead or official docs.
+- Examples must be realistic and follow best practices.
 Conversation History:
 {history_summary}
-
-Question: {query.strip()}
-Answer:
+User: {query.strip()}
+Assistant:
 """
 
     try:
         result = qa_pipeline(prompt, max_new_tokens=1024, do_sample=False, temperature=0.1, top_k=50)
         output = result[0]["generated_text"].strip()
-        if output.startswith("Answer:"):
-            output = output[7:].strip()
+        if output.startswith("Assistant:"):
+            output = output.replace("Assistant:", "").strip()
+
         if len(output) < 20:
-            output = f"I'm sorry, I couldn't find a precise answer for '{query}'. Please refer to Salesforce docs."
+            output = f"I'm sorry, I couldn't generate a detailed answer for '{query}'. You can also check https://developer.salesforce.com/docs."
+
         conversation_history.append((query, output))
         conversation_history = conversation_history[-6:]
         log_to_console({"Question": query, "Answer": output}, "Chatbot Query")
@@ -226,8 +282,8 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         metadata_button.click(validate_metadata, inputs=metadata_input, outputs=[mtype, issue, recommendation])
 
     with gr.Tab("Salesforce Chatbot"):
-        chatbot_output = gr.Chatbot(label="Conversation History", height=400)
-        query_input = gr.Textbox(label="Your Question", placeholder="e.g., How do I bulkify an Apex trigger?")
+        chatbot_output = gr.Chatbot(label="Conversation History", height=400, type="messages")
+        query_input = gr.Textbox(label="Your Question", placeholder="e.g., What is heap size in Apex?")
         with gr.Row():
             chatbot_button = gr.Button("Ask")
             clear_button = gr.Button("Clear Chat")