Update app.py

app.py

@@ -11,51 +11,52 @@ sf = Salesforce(
     domain="login"
 )
 
-# ---------- Load HuggingFace Pipelines ----------
+# ---------- Define Mapping for Model Labels ----------
+label_to_issue_type = {
+    "LABEL_0": "Performance",
+    "LABEL_1": "Error",
+    "LABEL_2": "Security",
+    "LABEL_3": "Best Practice"
+}
+
+# ---------- Load Model ----------
 code_analyzer = pipeline("text-classification", model="microsoft/codebert-base")
-qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-xl")
 
-# ---------- Code Review ----------
+# ---------- Code Analyzer Function ----------
 def analyze_code(code):
     if not code.strip():
         return "No code provided.", "", ""
 
     result = code_analyzer(code)
-    label = result[0]["label"].strip()
-
-    label_to_issue_type = {
-        "LABEL_0": "Performance",
-        "LABEL_1": "Error",
-        "LABEL_2": "Security",
-        "LABEL_3": "Best Practice"
-    }
+    label = result[0]["label"]  # e.g., LABEL_0
+    score = result[0]["score"]
 
-    if label not in label_to_issue_type:
-        raise ValueError(f"Unknown label from model: {label}")
+    # Map raw label to issue type (no lower(), no mismatch)
+    issue_type = label_to_issue_type.get(label.strip(), "Error")
 
-    issue_type = label_to_issue_type[label]
+    # Static suggestion and severity for now
     suggestion = "Consider refactoring for better performance"
     severity = "Medium"
 
-    # ✅ Use valid developer ID (REPLACE with real user ID)
-    developer_id = os.getenv("DEVELOPER_ID", "005NS00000Sn2q9")  # ← CHANGE THIS ID
-
     try:
         sf.CodeReviewResult__c.create({
             "Name": f"Review_{issue_type}",
             "CodeSnippet__c": code,
-            "IssueType__c": issue_type,
+            "IssueType__c": issue_type,  # Must match picklist
             "Suggestion__c": suggestion,
             "Severity__c": severity,
-            "Developer__c": developer_id
+            # Optional: add Developer__c if needed
         })
     except Exception as e:
         suggestion += f" (⚠️ Failed to log to Salesforce: {str(e)})"
 
     return issue_type, suggestion, severity
 
-# ---------- Metadata Validation ----------
+# ---------- 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"
@@ -73,35 +74,26 @@ def validate_metadata(metadata):
 
     return mtype, issue, recommendation
 
-# ---------- Natural Language AI Assistant ----------
+# ---------- NLP Assistant (Flan-T5) ----------
+qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-xl")
+
 def process_nlp_query(query):
-    try:
-        query = query.strip()
-        if not query:
-            return "Please enter a valid Salesforce-related question."
-
-        prompt = (
-            f"Question: {query}\n\n"
-            "You are a certified Salesforce developer. Answer this question in a clear, detailed, and educational way. "
-            "Include governor limits, examples, and best practices if relevant.\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()
+    if not query.strip():
+        return "No question provided."
 
-    except Exception as e:
-        return f"⚠️ AI Error: {str(e)}"
+    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"
+        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()
 
 # ---------- Gradio UI ----------
 with gr.Blocks() as demo:
@@ -124,7 +116,7 @@ with gr.Blocks() as demo:
         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 governor limit in Apex?")
+        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)