Update app.py

app.py

@@ -1,12 +1,11 @@
 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
+import xml.etree.ElementTree as ET  # NEW: for parsing metadata
 
 # ---------- Load Environment Variables ----------
 load_dotenv()
@@ -36,7 +35,7 @@ severities = {
     "Best Practice": "Low"
 }
 
-# ---------- Mock Knowledge Base ----------
+# ---------- 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.",
@@ -105,18 +104,7 @@ def analyze_code(code):
 
     return issue_type, suggestion, severity
 
-# ---------- 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 (handles non-XML gracefully) ----------
+# ---------- Metadata Validator (Updated with dynamic XML parsing) ----------
 def validate_metadata(metadata, admin_id=None):
     if not metadata.strip():
         return "No metadata provided.", "", ""
@@ -125,40 +113,22 @@ def validate_metadata(metadata, admin_id=None):
     issue = "Unknown"
     recommendation = "No recommendation found."
 
-    # Preview to help diagnose pastes
-    preview = metadata[:100].replace("\n", "\\n")
-    log_to_console({"preview": preview}, "Metadata Input Preview")
-
-    cleaned = _strip_code_fences(metadata).lstrip("\ufeff").strip()
+    try:
+        root = ET.fromstring(metadata)
+        # Detect missing <description> tag
+        description_found = any(elem.tag.endswith('description') for elem in root)
 
-    # Only treat as XML if the first non-space character is '<'
-    if not cleaned or not cleaned.lstrip().startswith("<"):
-        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:
-        cleaned = _escape_bare_ampersands(cleaned)
-        try:
-            root = ET.fromstring(cleaned)
-            # Detect <description> regardless of namespace
-            has_description = any(elem.tag.split('}')[-1].lower() == "description" for elem in root.iter())
+        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."
 
-            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)}"
+    except Exception as e:
+        issue = "Invalid XML"
+        recommendation = f"Could not parse metadata XML. Error: {str(e)}"
 
-    # Log to Salesforce
     log_data = {
         "Name": f"MetadataLog_{mtype}",
         "MetadataType__c": mtype,
@@ -166,6 +136,7 @@ def validate_metadata(metadata, admin_id=None):
         "Recommendation__c": recommendation,
         "Status__c": "Open"
     }
+
     if admin_id:
         log_data["Admin__c"] = admin_id
 
@@ -185,71 +156,50 @@ def validate_metadata(metadata, admin_id=None):
 
     return mtype, issue, recommendation
 
-# ---------- Salesforce Chatbot (messages format, no tip line) ----------
-conversation_history = []  # optional global memory
+# ---------- Salesforce Chatbot ----------
+conversation_history = []
 
-def salesforce_chatbot(query: str, chat_messages: list):
-    """
-    query: latest user question string
-    chat_messages: list[{"role":"user"|"assistant","content":str}] from the UI state
-    """
+def salesforce_chatbot(query, history=[]):
     global conversation_history
-    if not query or not query.strip():
+    if not query.strip():
         return "Please provide a valid Salesforce-related question."
 
     salesforce_keywords = [
         "apex", "soql", "trigger", "lwc", "aura", "visualforce", "salesforce", "governor limits",
-        "dml", "metadata", "batch apex", "queueable", "future method", "api", "sfdc", "heap", "limits"
+        "dml", "metadata", "batch apex", "queueable", "future method", "api", "sfdc"
     ]
-    if not any(k in query.lower() for k in salesforce_keywords):
-        return "Please ask a Salesforce-related question (e.g., Apex, SOQL, LWC, limits, etc)."
-
-    # Build history summary from messages (last few pairs)
-    chat_messages = chat_messages or []
-    history_summary = []
-    for m in chat_messages[-8:]:
-        prefix = "User" if m.get("role") == "user" else "Assistant"
-        history_summary.append(f"{prefix}: {m.get('content','')}")
-    history_summary = "\n".join(history_summary)
-
-    # Quick KB shortcut
+
+    if not any(keyword.lower() in query.lower() for keyword in salesforce_keywords):
+        return "Please ask a Salesforce-related question."
+
     query_key = query.lower().strip()
     for kb_key, kb_answer in salesforce_knowledge_base.items():
         if kb_key in query_key:
             conversation_history.append((query, kb_answer))
-            conversation_history[:] = conversation_history[-6:]
+            conversation_history = conversation_history[-6:]
             log_to_console({"Question": query, "Answer": kb_answer}, "Chatbot Query")
             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 and certified architect. Provide 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.
+You are an expert Salesforce developer...
 
 Conversation History:
 {history_summary}
 
-User: {query.strip()}
-Assistant:
+Question: {query.strip()}
+Answer:
 """
 
     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("Assistant:"):
-            output = output.replace("Assistant:", "").strip()
-
-        # NOTE: No extra "Tip:" or any appended lines here.
-        # We rely purely on the model output (prompt encourages 2+ lines).
-
+        if output.startswith("Answer:"):
+            output = output[7:].strip()
+        if len(output) < 20:
+            output = f"I'm sorry, I couldn't find a precise answer for '{query}'. Please refer to Salesforce docs."
         conversation_history.append((query, output))
-        conversation_history[:] = conversation_history[-6:]
+        conversation_history = conversation_history[-6:]
         log_to_console({"Question": query, "Answer": output}, "Chatbot Query")
         return output
     except Exception as e:
@@ -259,7 +209,6 @@ Assistant:
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🤖 Advanced Salesforce AI Code Review & Chatbot")
 
-    # Code Review
     with gr.Tab("Code Review"):
         code_input = gr.Textbox(label="Apex / LWC Code", lines=8, placeholder="Enter your Apex or LWC code here")
         issue_type = gr.Textbox(label="Issue Type")
@@ -268,7 +217,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         code_button = gr.Button("Analyze Code")
         code_button.click(analyze_code, inputs=code_input, outputs=[issue_type, suggestion, severity])
 
-    # Metadata Validation
     with gr.Tab("Metadata Validation"):
         metadata_input = gr.Textbox(label="Metadata XML", lines=8, placeholder="Enter your metadata XML here")
         mtype = gr.Textbox(label="Type")
@@ -277,30 +225,24 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         metadata_button = gr.Button("Validate Metadata")
         metadata_button.click(validate_metadata, inputs=metadata_input, outputs=[mtype, issue, recommendation])
 
-    # Salesforce Chatbot (messages format)
     with gr.Tab("Salesforce Chatbot"):
-        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?")
+        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?")
         with gr.Row():
             chatbot_button = gr.Button("Ask")
             clear_button = gr.Button("Clear Chat")
-
-        # Messages state: list of {"role","content"}
         chat_state = gr.State(value=[])
 
-        def update_chatbot(query, chat_messages):
-            if not query or not query.strip():
-                return chat_messages, "Please enter a valid question."
-            # Get model answer
-            answer = salesforce_chatbot(query, chat_messages)
-            # Append messages in OpenAI style
-            chat_messages = (chat_messages or []) + [
-                {"role": "user", "content": query},
-                {"role": "assistant", "content": answer},
-            ]
-            return chat_messages, ""
+        def update_chatbot(query, chat_history):
+            if not query.strip():
+                return chat_history, "Please enter a valid question."
+            response = salesforce_chatbot(query, chat_history)
+            chat_history.append((query, response))
+            return chat_history, ""
 
         def clear_chat():
+            global conversation_history
+            conversation_history = []
             return [], ""
 
         chatbot_button.click(fn=update_chatbot, inputs=[query_input, chat_state], outputs=[chatbot_output, query_input])