Update app.py

app.py

@@ -28,24 +28,29 @@ def analyze_code(code):
         return "No code provided.", "", ""
 
     result = code_analyzer(code)
-    label = result[0]["label"]  # e.g., LABEL_0
+    label = result[0]["label"].strip()
     score = result[0]["score"]
 
-    # Map raw label to issue type (no lower(), no mismatch)
-    issue_type = label_to_issue_type.get(label.strip(), "Error")
+    # Map raw label to Salesforce-valid picklist value
+    issue_type = label_to_issue_type.get(label, "Error")
+
+    suggestion = {
+        "Performance": "Consider refactoring for better performance",
+        "Error": "Fix syntax or logic error",
+        "Security": "Avoid dynamic SOQL and validate inputs",
+        "Best Practice": "Follow Apex coding conventions"
+    }.get(issue_type, "Consider improving code quality")
 
-    # Static suggestion and severity for now
-    suggestion = "Consider refactoring for better performance"
     severity = "Medium"
 
     try:
         sf.CodeReviewResult__c.create({
             "Name": f"Review_{issue_type}",
             "CodeSnippet__c": code,
-            "IssueType__c": issue_type,  # Must match picklist
+            "IssueType__c": issue_type,
             "Suggestion__c": suggestion,
             "Severity__c": severity,
-            # Optional: add Developer__c if needed
+            # Add Developer__c if needed
         })
     except Exception as e:
         suggestion += f" (⚠️ Failed to log to Salesforce: {str(e)})"
@@ -82,18 +87,16 @@ def process_nlp_query(query):
         return "No question provided."
 
     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 certified Salesforce Apex expert. "
+        "Answer accurately using Apex best practices, governor limits, SOQL/DML limits, and examples. "
+        "Avoid repeating the question. Be concise, correct, and helpful.\n\n"
         f"Question: {query.strip()}\n\nAnswer:"
     )
 
     result = qa_pipeline(prompt, max_length=512, temperature=0.3, top_p=0.9, 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()
+    return output.split("Answer:")[-1].strip() if "Answer:" in output else output.strip()
 
 # ---------- Gradio UI ----------
 with gr.Blocks() as demo: