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  # NEW: for parsing metadata
+import xml.etree.ElementTree as ET
 
 # ---------- Load Environment Variables ----------
 load_dotenv()
@@ -35,22 +35,8 @@ severities = {
     "Best Practice": "Low"
 }
 
-# ---------- 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")
+qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-large")
 
 # ---------- Logging ----------
 def log_to_console(data, log_type):
@@ -104,7 +90,7 @@ def analyze_code(code):
 
     return issue_type, suggestion, severity
 
-# ---------- Metadata Validator (Updated with dynamic XML parsing) ----------
+# ---------- Metadata Validator ----------
 def validate_metadata(metadata, admin_id=None):
     if not metadata.strip():
         return "No metadata provided.", "", ""
@@ -115,7 +101,6 @@ 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:
@@ -124,7 +109,6 @@ 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)}"
@@ -156,7 +140,7 @@ def validate_metadata(metadata, admin_id=None):
 
     return mtype, issue, recommendation
 
-# ---------- Salesforce Chatbot ----------
+# ---------- Salesforce Chatbot (Final Improved Prompt) ----------
 conversation_history = []
 
 def salesforce_chatbot(query, history=[]):
@@ -166,38 +150,42 @@ 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."
 
-    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 an expert Salesforce developer...
+You are a certified Salesforce developer and architect. You answer Salesforce questions with 100% accuracy and clarity, using official platform knowledge.
+
+Guidelines:
+- Always give answers that are at least 2 full lines long.
+- When asked about Apex limits (e.g., DML, SOQL, heap), always:
+  - Mention the numeric limit (e.g., 150 DML statements)
+  - Include the transaction context (sync/async)
+  - Suggest using monitoring methods like Limits.getDMLStatements()
+- Avoid making up numbers — if unsure, suggest Trailhead or official docs.
+- Format clearly, using bullet points or code if helpful.
 
 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 len(output) < 20:
-            output = f"I'm sorry, I couldn't find a precise answer for '{query}'. Please refer to Salesforce docs."
+        if output.startswith("Assistant:"):
+            output = output.replace("Assistant:", "").strip()
+
+        if len(output.split()) < 20:
+            output += "\n\nPlease refer to: https://developer.salesforce.com/docs"
+
         conversation_history.append((query, output))
         conversation_history = conversation_history[-6:]
         log_to_console({"Question": query, "Answer": output}, "Chatbot Query")
@@ -227,7 +215,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 do I bulkify an Apex trigger?")
+        query_input = gr.Textbox(label="Your Question", placeholder="e.g., How many DML operations are allowed in Apex?")
         with gr.Row():
             chatbot_button = gr.Button("Ask")
             clear_button = gr.Button("Clear Chat")