Update app.py

app.py

@@ -5,7 +5,7 @@ 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()
@@ -35,8 +35,22 @@ severities = {
     "Best Practice": "Low"
 }
 
-# ---------- Load QnA Model (no fallback) ----------
-qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-large")
+# ---------- 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.",
+    "bulkify apex trigger": "...",
+    "soql injection": "...",
+    "lwc best practices": "...",
+    "batch apex": "..."
+}
+
+# ---------- Load QnA Model ----------
+try:
+    qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-large")
+except Exception as e:
+    print(f"Model loading error: {e}. Falling back to flan-t5-base.")
+    qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")
 
 # ---------- Logging ----------
 def log_to_console(data, log_type):
@@ -90,7 +104,7 @@ def analyze_code(code):
 
     return issue_type, suggestion, severity
 
-# ---------- Metadata Validator ----------
+# ---------- Metadata Validator (Updated with dynamic XML parsing) ----------
 def validate_metadata(metadata, admin_id=None):
     if not metadata.strip():
         return "No metadata provided.", "", ""
@@ -101,6 +115,7 @@ def validate_metadata(metadata, admin_id=None):
 
     try:
         root = ET.fromstring(metadata)
+        # Detect missing <description> tag
         description_found = any(elem.tag.endswith('description') for elem in root)
 
         if not description_found:
@@ -109,6 +124,7 @@ def validate_metadata(metadata, admin_id=None):
         else:
             issue = "Unused field detected"
             recommendation = "Remove it to improve performance or document its purpose."
+
     except Exception as e:
         issue = "Invalid XML"
         recommendation = f"Could not parse metadata XML. Error: {str(e)}"
@@ -140,7 +156,7 @@ def validate_metadata(metadata, admin_id=None):
 
     return mtype, issue, recommendation
 
-# ---------- Salesforce Chatbot (Improved Prompt) ----------
+# ---------- Salesforce Chatbot ----------
 conversation_history = []
 
 def salesforce_chatbot(query, history=[]):
@@ -150,41 +166,38 @@ 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", "heap", "limits"
+        "dml", "metadata", "batch apex", "queueable", "future method", "api", "sfdc"
     ]
 
     if not any(keyword.lower() in query.lower() for keyword in salesforce_keywords):
         return "Please ask a Salesforce-related question."
 
-    history_summary = "\n".join([f"User: {q}\nAssistant: {a}" for q, a in conversation_history[-4:]])
+    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:]
+            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 a certified Salesforce developer and architect. Your role is to answer with 100% accurate and detailed technical explanations, especially about limits, code, and platform best practices.
-
-Your answers MUST:
-- Always be at least two lines long.
-- Be correct, clear, and production-safe.
-- Include official Salesforce governor limits when applicable.
-- Use bullet points or code snippets when needed.
-- Recommend Trailhead or official docs if the answer isn't definitive.
-- Follow real-world practices (bulkification, error handling, etc).
+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()
-
-        if len(output.split()) < 15:
-            output += "\n\nRefer to: https://developer.salesforce.com/docs for more."
-
+        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:]
         log_to_console({"Question": query, "Answer": output}, "Chatbot Query")
@@ -214,7 +227,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     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 many DML operations are allowed in Apex?")
+        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")