Update app.py

app.py

@@ -1,6 +1,18 @@
+import os
 import random
+import gradio as gr
+from transformers import pipeline
+from simple_salesforce import Salesforce
 
-# ---------- Define Mapping for Model Labels ----------
+# ---------- Salesforce Login ----------
+sf = Salesforce(
+    username=os.getenv("SF_USERNAME"),
+    password=os.getenv("SF_PASSWORD"),
+    security_token=os.getenv("SF_SECURITY_TOKEN"),
+    domain="login"
+)
+
+# ---------- Label Mapping ----------
 label_to_issue_type = {
     "LABEL_0": "Performance",
     "LABEL_1": "Error",
@@ -8,36 +20,33 @@ label_to_issue_type = {
     "LABEL_3": "Best Practice"
 }
 
+suggestions = {
+    "Performance": "Consider optimizing loops and database access.",
+    "Error": "Add proper error handling and null checks.",
+    "Security": "Avoid dynamic SOQL. Use binding variables.",
+    "Best Practice": "Refactor for readability and use bulk-safe patterns."
+}
+severities = {
+    "Performance": "Medium",
+    "Error": "High",
+    "Security": "High",
+    "Best Practice": "Low"
+}
+
+# ---------- Load Dummy Classifier & QnA Model ----------
+qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")
+
+# ---------- Code Analyzer ----------
 def analyze_code(code):
     if not code.strip():
         return "No code provided.", "", ""
 
-    # --- Simulate classification labels for demo/testing ---
-    simulated_labels = ["LABEL_0", "LABEL_1", "LABEL_2", "LABEL_3"]
-    label = random.choice(simulated_labels)
-    issue_type = label_to_issue_type.get(label, "Error")
-
-    # Simulated suggestion/severity based on issue type
-    suggestions = {
-        "Performance": "Consider optimizing loops and database access.",
-        "Error": "Add proper error handling and null checks.",
-        "Security": "Avoid dynamic SOQL. Use binding variables.",
-        "Best Practice": "Refactor for readability and use bulk-safe patterns."
-    }
-    severities = {
-        "Performance": "Medium",
-        "Error": "High",
-        "Security": "High",
-        "Best Practice": "Low"
-    }
-
-    suggestion = suggestions.get(issue_type)
-    severity = severities.get(issue_type)
-
-    # Optional: Logging for debugging
-    print(f"[DEBUG] Label: {label} → Issue: {issue_type} | Severity: {severity}")
-
-    # Log to Salesforce
+    # Simulate classification
+    label = random.choice(list(label_to_issue_type.keys()))
+    issue_type = label_to_issue_type[label]
+    suggestion = suggestions[issue_type]
+    severity = severities[issue_type]
+
     try:
         sf.CodeReviewResult__c.create({
             "Name": f"Review_{issue_type}",
@@ -45,9 +54,80 @@ def analyze_code(code):
             "IssueType__c": issue_type,
             "Suggestion__c": suggestion,
             "Severity__c": severity
-            # You can also add Developer__c here if needed
         })
     except Exception as e:
-        suggestion += f" (⚠️ Salesforce error: {str(e)})"
+        suggestion += f" ⚠️ Salesforce logging failed: {str(e)}"
 
     return issue_type, suggestion, severity
+
+# ---------- Metadata Validator ----------
+def validate_metadata(metadata):
+    if not metadata.strip():
+        return "No metadata provided.", "", ""
+
+    mtype = "Field"
+    issue = "Unused field detected"
+    recommendation = "Remove it to improve performance"
+
+    try:
+        sf.MetadataAuditLog__c.create({
+            "Name": f"MetadataLog_{mtype}",
+            "MetadataType__c": mtype,
+            "IssueDescription__c": issue,
+            "Recommendation__c": recommendation,
+            "Status__c": "Open"
+        })
+    except Exception as e:
+        recommendation += f" ⚠️ Salesforce logging failed: {str(e)}"
+
+    return mtype, issue, recommendation
+
+# ---------- Natural Language Assistant ----------
+def process_nlp_query(query):
+    if not query.strip():
+        return "No question provided."
+
+    prompt = (
+        "You are a Salesforce Apex expert. "
+        "Answer the question clearly and include governor limits, code best practices, and real examples.\n\n"
+        f"Question: {query.strip()}\n\nAnswer:"
+    )
+
+    try:
+        result = qa_pipeline(prompt, max_new_tokens=256, do_sample=False)
+        output = result[0]["generated_text"]
+        if "Answer:" in output:
+            return output.split("Answer:")[-1].strip()
+        return output.strip()
+    except Exception as e:
+        return f"⚠️ AI Response Error: {str(e)}"
+
+# ---------- Gradio UI ----------
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖 Salesforce AI Code Review & Q&A")
+
+    with gr.Tab("Code Review"):
+        code_input = gr.Textbox(label="Apex / LWC Code", lines=8)
+        issue_type = gr.Textbox(label="Issue Type")
+        suggestion = gr.Textbox(label="AI Suggestion")
+        severity = gr.Textbox(label="Severity")
+        code_button = gr.Button("Analyze Code")
+        code_button.click(analyze_code, inputs=code_input, outputs=[issue_type, suggestion, severity])
+
+    with gr.Tab("Metadata Validation"):
+        metadata_input = gr.Textbox(label="Metadata XML", lines=8)
+        mtype = gr.Textbox(label="Type")
+        issue = gr.Textbox(label="Issue")
+        recommendation = gr.Textbox(label="Recommendation")
+        metadata_button = gr.Button("Validate Metadata")
+        metadata_button.click(validate_metadata, inputs=metadata_input, outputs=[mtype, issue, recommendation])
+
+    with gr.Tab("Ask AI (Natural Language)"):
+        query_input = gr.Textbox(label="Your question", lines=2, placeholder="e.g. What is a trigger?")
+        response_output = gr.Textbox(label="AI Response", lines=8)
+        nlp_button = gr.Button("Ask")
+        nlp_button.click(process_nlp_query, inputs=query_input, outputs=response_output)
+
+# ---------- Start UI ----------
+if __name__ == "__main__":
+    demo.launch()