Update main.py

main.py

@@ -54,6 +54,32 @@ model        = "HuggingFaceTB/SmolLM3-3B"
 temperature  = 0.1
 max_tokens   = 500
 
+import re
+
+def clean_response(answer: str) -> str:
+    """Clean up AI response by removing unwanted artifacts and formatting."""
+    if not answer:
+        return answer
+    
+    # Remove thinking tags and content
+    answer = re.sub(r'<think>.*?</think>', '', answer, flags=re.DOTALL)
+    answer = re.sub(r'<thinking>.*?</thinking>', '', answer, flags=re.DOTALL)
+    
+    # Remove other common AI response artifacts
+    answer = re.sub(r'\[.*?\]', '', answer, flags=re.DOTALL)  # Remove bracketed content
+    answer = re.sub(r'\{.*?\}', '', answer, flags=re.DOTALL)  # Remove curly bracketed content
+    answer = re.sub(r'```.*?```', '', answer, flags=re.DOTALL)  # Remove code blocks
+    answer = re.sub(r'---.*?---', '', answer, flags=re.DOTALL)  # Remove dashed sections
+    
+    # Remove excessive whitespace and newlines
+    answer = re.sub(r'\s+', ' ', answer).strip()
+    
+    # Remove common AI-generated prefixes/suffixes
+    answer = re.sub(r'^(Assistant:|AI:|Grok:)\s*', '', answer, flags=re.IGNORECASE)
+    answer = re.sub(r'\s*(Sincerely,.*|Best regards,.*|Regards,.*)$', '', answer, flags=re.IGNORECASE)
+    
+    return answer
+    
 def response_generator(query: str) -> str:
     """Ask the RAG chain to answer `query`, with JSON‑error fallback."""
     # log usage
@@ -109,6 +135,8 @@ def response_generator(query: str) -> str:
             "I’m sorry, I don’t have enough information to answer that. "
             "If you have a public data source to add, please email [email protected]."
         )
+
+    answer = clean_response(answer)
     return answer
 
 # ─────── Streamlit UI ──────────────────────────────────────────────────────────