Update app.py

app.py

@@ -1,9 +1,9 @@
 import random
 import gradio as gr
-from transformers import pipeline
 from datetime import datetime
+from transformers import pipeline
 
-# ---------- Label Mapping ----------
+# ---------- Label Mapping (Unchanged) ----------
 label_to_issue_type = {
     "LABEL_0": "Performance",
     "LABEL_1": "Error",
@@ -25,26 +25,88 @@ severities = {
     "Best Practice": "Low"
 }
 
+# ---------- Mock Salesforce 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 (e.g., Batch Apex, Queueable).",
+    "governor limits dml": "The governor limit for DML statements is 150 per synchronous or asynchronous transaction.",
+    "bulkify apex trigger": """
+To bulkify an Apex trigger, ensure it handles multiple records to stay within governor limits:
+- Use collections (e.g., Lists, Sets, Maps) to process records.
+- Perform SOQL queries outside loops.
+- Execute DML operations in bulk.
+Example:
+```apex
+trigger AccountTrigger on Account (before insert) {
+    Set<String> accountNames = new Set<String>();
+    for (Account acc : Trigger.new) {
+        accountNames.add(acc.Name);
+    }
+    List<Account> existingAccounts = [SELECT Name FROM Account WHERE Name IN :accountNames];
+    // Process records
+}
+```""",
+    "soql injection": """
+To prevent SOQL injection:
+- Use bind variables (e.g., `:variable`) instead of dynamic SOQL.
+- Escape single quotes with `String.escapeSingleQuotes()`.
+Example:
+```apex
+String userInput = 'Test';
+List<Account> accounts = [SELECT Name FROM Account WHERE Name = :userInput];
+```""",
+    "lwc best practices": """
+LWC best practices include:
+- Use `@api` decorators for public properties.
+- Leverage `@wire` for efficient data retrieval.
+- Avoid hardcoding IDs; use dynamic queries.
+- Follow SLDS for consistent UI.
+Example:
+```javascript
+import { LightningElement, api } from 'lwc';
+export default class MyComponent extends LightningElement {
+    @api recordId;
+}
+```""",
+    "batch apex": """
+Batch Apex processes records in chunks to handle large data volumes:
+- Implement `Database.Batchable<SObject>`.
+- Use `start`, `execute`, and `finish` methods.
+Example:
+```apex
+global class MyBatch implements Database.Batchable<SObject> {
+    global Database.QueryLocator start(Database.BatchableContext bc) {
+        return Database.getQueryLocator('SELECT Id FROM Account');
+    }
+    global void execute(Database.BatchableContext bc, List<SObject> scope) {
+        // Process records
+    }
+    global void finish(Database.BatchableContext bc) {}
+}
+```""",
+}
+
 # ---------- Load QnA Model ----------
-qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")
+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")
 
-# ---------- Local Logging (Replacement for Salesforce Logging) ----------
+# ---------- Local Logging (Unchanged) ----------
 def log_to_console(data, log_type):
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     print(f"[{timestamp}] {log_type} Log: {data}")
 
-# ---------- Code Analyzer ----------
+# ---------- Code Analyzer (Unchanged) ----------
 def analyze_code(code):
     if not code.strip():
         return "No code provided.", "", ""
 
-    # Simulate classification (replace with actual analysis for production)
     label = random.choice(list(label_to_issue_type.keys()))
     issue_type = label_to_issue_type[label]
     suggestion = suggestions[issue_type]
     severity = severities[issue_type]
 
-    # Log to console instead of Salesforce
     log_data = {
         "Name": f"Review_{issue_type}",
         "CodeSnippet": code,
@@ -56,7 +118,7 @@ def analyze_code(code):
 
     return issue_type, suggestion, severity
 
-# ---------- Metadata Validator ----------
+# ---------- Metadata Validator (Unchanged) ----------
 def validate_metadata(metadata):
     if not metadata.strip():
         return "No metadata provided.", "", ""
@@ -65,7 +127,6 @@ def validate_metadata(metadata):
     issue = "Unused field detected"
     recommendation = "Remove it to improve performance or document its purpose."
 
-    # Log to console instead of Salesforce
     log_data = {
         "Name": f"MetadataLog_{mtype}",
         "MetadataType": mtype,
@@ -83,23 +144,44 @@ conversation_history = []
 def salesforce_chatbot(query, history=[]):
     global conversation_history
     if not query.strip():
-        return "No question provided."
+        return "Please provide a valid Salesforce-related question."
 
-    # Expanded Salesforce-related keywords
+    # Expanded Salesforce keywords
     salesforce_keywords = [
         "apex", "soql", "sosl", "trigger", "lwc", "aura", "visualforce", "salesforce",
         "governor limits", "bulkification", "dml", "query", "metadata", "flow",
         "process builder", "sobject", "schema", "lightning", "custom object",
         "validation rule", "workflow", "platform event", "batch apex", "queueable",
-        "future method", "lightning web component", "api", "rest", "soap"
+        "future method", "lightning web component", "api", "rest", "soap", "integration",
+        "trigger", "profile", "permission set", "sharing rule", "field", "record type",
+        "crm", "sfdc", "force.com"
     ]
+
+    # Check if query is Salesforce-related
     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., about Apex, SOQL, LWC, or Salesforce platform features)."
+
+    # Check knowledge base for exact matches
+    query_key = query.lower().strip()
+    for kb_key, kb_answer in salesforce_knowledge_base.items():
+        if kb_key in query_key:
+            # Store in conversation history
+            conversation_history.append((query, kb_answer))
+            conversation_history = conversation_history[-6:]  # Keep last 6 exchanges
+            log_to_console({"Question": query, "Answer": kb_answer}, "Chatbot Query")
+            return kb_answer
 
     # Build conversation context
-    history_summary = "\n".join([f"Q: {q} A: {a}" for q, a in conversation_history[-3:]])  # Last 3 exchanges
+    history_summary = "\n".join([f"User: {q}\nAssistant: {a}" for q, a in conversation_history[-4:]])
     prompt = f"""
-You are an expert Salesforce developer with deep knowledge of Apex, SOQL, LWC, and Salesforce platform features. Answer the following Salesforce-related question accurately, concisely, and based strictly on Salesforce best practices and official documentation (e.g., governor limits: 100 SOQL queries, 150 DML statements per transaction). Provide code examples where relevant, use correct terminology, and address the question directly. If the question builds on prior context, consider the conversation history.
+You are an expert Salesforce developer with deep knowledge of Apex, SOQL, SOSL, LWC, Aura, Visualforce, and Salesforce platform features. Your role is to provide 100% accurate answers based strictly on Salesforce official documentation and best practices (e.g., governor limits: 100 SOQL queries, 150 DML statements per transaction). Follow these guidelines:
+- Provide concise, accurate answers with no speculation.
+- Include code snippets in ```apex``` or ```javascript``` blocks for technical questions.
+- Reference governor limits or best practices explicitly.
+- If the question is ambiguous, ask for clarification within the response.
+- Use bullet points for lists or steps.
+- Leverage conversation history for context.
+- If unsure, admit limitations and suggest checking Salesforce documentation or Trailhead.
 
 Conversation History:
 {history_summary}
@@ -109,24 +191,33 @@ Answer:
 """
 
     try:
-        result = qa_pipeline(prompt, max_new_tokens=512, do_sample=False)
+        # Generate response with strict parameters
+        result = qa_pipeline(prompt, max_new_tokens=1024, do_sample=False, temperature=0.1, top_k=50)
         output = result[0]["generated_text"].strip()
 
+        # Clean up response
+        if output.startswith("Answer:"):
+            output = output[7:].strip()
+
+        # Validate response quality
+        if len(output) < 20 or output.lower() in ["unknown", "i don't know", "not sure"]:
+            output = f"I'm sorry, I couldn't find a precise answer for '{query}'. Please clarify or refer to Salesforce documentation at https://developer.salesforce.com/docs."
+
         # Store in conversation history
         conversation_history.append((query, output))
-        conversation_history = conversation_history[-5:]  # Keep last 5 exchanges
+        conversation_history = conversation_history[-6:]  # Keep last 6 exchanges
 
         # Log question and answer
         log_to_console({"Question": query, "Answer": output}, "Chatbot Query")
 
         return output
     except Exception as e:
-        error_msg = f"⚠️ AI Response Error: {str(e)}"
+        error_msg = f"⚠️ Error generating response: {str(e)}. Please try rephrasing your question or check Salesforce documentation at https://developer.salesforce.com/docs."
         log_to_console({"Question": query, "Error": error_msg}, "Chatbot Error")
         return error_msg
 
 # ---------- Gradio UI ----------
-with gr.Blocks() as demo:
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🤖 Advanced Salesforce AI Code Review & Chatbot")
 
     with gr.Tab("Code Review"):
@@ -146,10 +237,38 @@ with gr.Blocks() as demo:
         metadata_button.click(validate_metadata, inputs=metadata_input, outputs=[mtype, issue, recommendation])
 
     with gr.Tab("Salesforce Chatbot"):
-        query_input = gr.Textbox(label="Ask a Salesforce question", lines=2, placeholder="e.g., How do I bulkify an Apex trigger?")
-        response_output = gr.Textbox(label="AI Response", lines=8)
-        chatbot_button = gr.Button("Ask")
-        chatbot_button.click(salesforce_chatbot, inputs=query_input, outputs=response_output)
+        gr.Markdown("### Ask a Salesforce Question\nGet expert answers on Apex, SOQL, LWC, and more!")
+        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")
+        
+        # State to manage chat history
+        chat_state = gr.State(value=[])
+
+        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]
+        )
+        clear_button.click(
+            fn=clear_chat,
+            inputs=None,
+            outputs=[chatbot_output, query_input]
+        )
 
 # ---------- Start UI ----------
 if __name__ == "__main__":