Update app.py

app.py

@@ -137,55 +137,87 @@ def remove_repeated_sentences(text: str) -> str:
         unique_sentences.append(sentence)
     return " ".join(unique_sentences)
 
-def clean_agri_output(raw_text: str) -> str:
+def clean_agri_output(raw_text: str, user_message: str = "") -> str:
+    """
+    Cleans model output. 
+    - For Pest/Disease queries: Displays Problem, Possible Cause, Recommended Action, and Explanation.
+    - For other queries: Displays Recommended Action and Explanation.
+    """
     if not raw_text:
         return ""
 
     text = raw_text.strip()
 
+    # 1. Strip raw metadata key=value pairs and structural tags
     text = re.sub(
-        r'(?:Context\s*:?\s*)?(?:Region|Owner|Farm|Season|Soil|Severity|Source|Problem|Possible Cause|Recommended Action|Reason)\s*=\s*.*?(?=\.|\n|\Z)',
+        r'(?:Context\s*:?\s*)?(?:Region|Owner|Farm|Season|Soil|Severity|Source)\s*=\s*.*?(?=\.|\n|\Z)',
         '', text, flags=re.IGNORECASE
     )
     text = re.sub(r'\bContext\s*:?\s*', '', text, flags=re.IGNORECASE)
     text = re.sub(r'https?://\S+|\[.*?\]\(.*?\)', '', text)
     text = re.sub(r"<\|.*?\|>", "", text)
 
+    # 2. Check if the query is related to Pests or Diseases
+    pest_disease_keywords = [
+        "pest", "pests", "insect", "insects", "aphid", "thrips", "whitefly", 
+        "caterpillar", "weevil", "beetle", "borer", "worm", "disease", 
+        "diseases", "fungus", "fungal", "bacteria", "bacterial", "virus", 
+        "viral", "blight", "rust", "mildew", "rot", "leaf spot", "spots"
+    ]
+    
+    combined_check_text = (user_message + " " + text).lower()
+    is_pest_or_disease = any(k in combined_check_text for k in pest_disease_keywords)
+
+    # 3. Extract sections using Regex
+    problem_match = re.search(
+        r"Problem\s*:\s*(.*?)(?=\n\s*(?:Possible Cause|Recommended Action|Recommendation|Explanation)\s*:|\Z)",
+        text, flags=re.IGNORECASE | re.DOTALL
+    )
+    cause_match = re.search(
+        r"Possible Cause\s*:\s*(.*?)(?=\n\s*(?:Problem|Recommended Action|Recommendation|Explanation)\s*:|\Z)",
+        text, flags=re.IGNORECASE | re.DOTALL
+    )
     recommendation_match = re.search(
-        r"(?:Recommended Action|Recommendation)\s*:\s*(.*?)(?=\n\s*(?:Explanation|Problem|Possible Cause)\s*:|\Z)",
-        text,
-        flags=re.IGNORECASE | re.DOTALL,
+        r"(?:Recommended Action|Recommendation)\s*:\s*(.*?)(?=\n\s*(?:Problem|Possible Cause|Explanation)\s*:|\Z)",
+        text, flags=re.IGNORECASE | re.DOTALL
     )
-    recommendation = recommendation_match.group(1).strip() if recommendation_match else ""
-
     explanation_match = re.search(
         r"Explanation\s*:\s*(.*?)(?=\n\s*(?:Problem|Possible Cause|Recommended Action|Recommendation)\s*:|\Z)",
-        text,
-        flags=re.IGNORECASE | re.DOTALL,
+        text, flags=re.IGNORECASE | re.DOTALL
     )
-    explanation = explanation_match.group(1).strip() if explanation_match else ""
 
-    if not recommendation:
-        cleaned = re.sub(
-            r"^(#+\s*)?(Response|Instruction|Recommendation)\s*:\s*",
-            "",
-            text,
-            flags=re.IGNORECASE,
-        )
-        recommendation = re.split(r"\n#+\s*|\nContext:|\nExplanation:|\bRegion=", cleaned)[0].strip()
+    problem = remove_repeated_sentences(problem_match.group(1).strip()) if problem_match else ""
+    cause = remove_repeated_sentences(cause_match.group(1).strip()) if cause_match else ""
+    recommendation = remove_repeated_sentences(recommendation_match.group(1).strip()) if recommendation_match else ""
+    explanation = remove_repeated_sentences(explanation_match.group(1).strip()) if explanation_match else ""
 
-    recommendation = re.sub(r"^(Problem|Possible Cause)\s*:.*?\n", "", recommendation, flags=re.IGNORECASE)
-    recommendation = remove_repeated_sentences(recommendation)
-    explanation = remove_repeated_sentences(explanation)
+    # Fallback cleanup for unstructured generation
+    if not recommendation:
+        cleaned = re.sub(r"^(#+\s*)?(Response|Instruction|Recommendation)\s*:\s*", "", text, flags=re.IGNORECASE)
+        recommendation = remove_repeated_sentences(re.split(r"\n#+\s*|\nContext:|\nExplanation:|\bRegion=", cleaned)[0].strip())
 
     recommendation = re.sub(r"\s+", " ", recommendation).strip()
     explanation = re.sub(r"\s+", " ", explanation).strip()
+    problem = re.sub(r"\s+", " ", problem).strip()
+    cause = re.sub(r"\s+", " ", cause).strip()
 
     output = []
-    if recommendation and recommendation != ".":
-        output.append(f"### 🌱 Recommended Action\n\n{recommendation}")
-    if explanation and explanation != ".":
-        output.append(f"### 💡 Explanation\n\n{explanation}")
+
+    # 4. Construct response according to query type
+    if is_pest_or_disease:
+        if problem and problem != ".":
+            output.append(f"### ⚠️ Problem\n\n{problem}")
+        if cause and cause != ".":
+            output.append(f"### 🔍 Possible Cause\n\n{cause}")
+        if recommendation and recommendation != ".":
+            output.append(f"### 🌱 Recommended Action\n\n{recommendation}")
+        if explanation and explanation != ".":
+            output.append(f"### 💡 Explanation\n\n{explanation}")
+    else:
+        if recommendation and recommendation != ".":
+            output.append(f"### 🌱 Recommended Action\n\n{recommendation}")
+        if explanation and explanation != ".":
+            output.append(f"### 💡 Explanation\n\n{explanation}")
 
     if not output:
         return "I was unable to generate a clear agricultural recommendation. Please try rephrasing your question."
@@ -266,13 +298,13 @@ def generate_response(message, history):
     partial_text = ""
     for new_text in streamer:
         partial_text += new_text
-        cleaned = clean_agri_output(partial_text)
+        cleaned = clean_agri_output(partial_text, user_message=user_message)
         if cleaned:
             history[-1]["content"] = cleaned
             yield history
 
     thread.join()
-    final_response = clean_agri_output(partial_text)
+    final_response = clean_agri_output(partial_text, user_message=user_message)
     if final_response:
         history[-1]["content"] = final_response
     else: