Update app.py

app.py

@@ -1,8 +1,8 @@
+import os
+import random
 import gradio as gr
 from transformers import pipeline
 from simple_salesforce import Salesforce
-import os
-import re
 
 # ---------- Salesforce Login ----------
 sf = Salesforce(
@@ -12,7 +12,7 @@ sf = Salesforce(
     domain="login"
 )
 
-# ---------- Define Mapping for Model Labels ----------
+# ---------- Label Mapping ----------
 label_to_issue_type = {
     "LABEL_0": "Performance",
     "LABEL_1": "Error",
@@ -20,25 +20,32 @@ label_to_issue_type = {
     "LABEL_3": "Best Practice"
 }
 
-# ---------- Load Model ----------
-code_analyzer = pipeline("text-classification", model="microsoft/codebert-base")
+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"
+}
 
-# ---------- DML Counter ----------
-def count_dml_statements(code):
-    dml_keywords = ["insert", "update", "delete", "upsert", "merge"]
-    return sum(len(re.findall(rf"\\b{kw}\\b", code, flags=re.IGNORECASE)) for kw in dml_keywords)
+# ---------- Load Dummy Classifier & QnA Model ----------
+qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")
 
-# ---------- Code Analyzer Function ----------
+# ---------- Code Analyzer ----------
 def analyze_code(code):
     if not code.strip():
         return "No code provided.", "", ""
 
-    result = code_analyzer(code)
-    label = result[0]["label"]
-    issue_type = label_to_issue_type.get(label.strip(), "Error")
-
-    suggestion = f"Consider refactoring for better {issue_type.lower()}"
-    severity = "Medium"
+    # 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({
@@ -46,11 +53,10 @@ def analyze_code(code):
             "CodeSnippet__c": code,
             "IssueType__c": issue_type,
             "Suggestion__c": suggestion,
-            "Severity__c": severity,
-            # Add Developer__c if available
+            "Severity__c": severity
         })
     except Exception as e:
-        suggestion += f" (⚠️ Failed to log to Salesforce: {str(e)})"
+        suggestion += f" ⚠️ Salesforce logging failed: {str(e)}"
 
     return issue_type, suggestion, severity
 
@@ -72,35 +78,29 @@ def validate_metadata(metadata):
             "Status__c": "Open"
         })
     except Exception as e:
-        recommendation += f" (⚠️ Failed to log to Salesforce: {str(e)})"
+        recommendation += f" ⚠️ Salesforce logging failed: {str(e)}"
 
     return mtype, issue, recommendation
 
-# ---------- NLP Assistant (Flan-T5) ----------
-qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-xl")
-
-def process_nlp_query(query, code_snippet=""):
+# ---------- Natural Language Assistant ----------
+def process_nlp_query(query):
     if not query.strip():
         return "No question provided."
 
-    # Handle DML-specific queries manually
-    if "dml" in query.lower() and code_snippet:
-        dml_count = count_dml_statements(code_snippet)
-        return f"{dml_count} DML statement(s) detected in the provided code."
-
     prompt = (
-        "You are a certified Salesforce Apex expert. Answer the question below accurately and clearly. "
-        "Include governor limits, code examples, best practices, and terminology when appropriate. Do not repeat the question. "
-        "Give only a factual answer.\n\n"
+        "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:"
     )
 
-    result = qa_pipeline(prompt, max_new_tokens=256, repetition_penalty=1.1, do_sample=False)
-    output = result[0]["generated_text"]
-
-    if "Answer:" in output:
-        return output.split("Answer:")[-1].strip()
-    return output.strip()
+    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:
@@ -124,9 +124,10 @@ with gr.Blocks() as demo:
 
     with gr.Tab("Ask AI (Natural Language)"):
         query_input = gr.Textbox(label="Your question", lines=2, placeholder="e.g. What is a trigger?")
-        code_context = gr.Textbox(label="(Optional) Provide code for context", lines=4)
         response_output = gr.Textbox(label="AI Response", lines=8)
         nlp_button = gr.Button("Ask")
-        nlp_button.click(process_nlp_query, inputs=[query_input, code_context], outputs=response_output)
+        nlp_button.click(process_nlp_query, inputs=query_input, outputs=response_output)
 
-demo.launch()
+# ---------- Start UI ----------
+if __name__ == "__main__":
+    demo.launch()