Updated

main.py

@@ -8,6 +8,7 @@ from pydantic import BaseModel
 import torch
 import logging
 import os
+import re
 
 from huggingface_hub import login
 from langchain.llms.huggingface_pipeline import HuggingFacePipeline
@@ -102,7 +103,7 @@ Conversation so far:
 {conversation_history}
 
 User: {query}
-Tech Disciples AI (respond with warmth, depth, and Biblical understanding):
+Tech Disciples AI:
 """
 
 prompt = PromptTemplate(
@@ -112,6 +113,25 @@ prompt = PromptTemplate(
 
 chain = LLMChain(prompt=prompt, llm=llm, memory=memory) if llm else None
 
+# =====================================================
+# HELPER FUNCTION TO CLEAN RESPONSE
+# =====================================================
+def clean_response(text):
+    """Remove unwanted markers and clean up the response"""
+    # Remove common end markers
+    text = re.sub(r'\[End of conversation\]', '', text, flags=re.IGNORECASE)
+    text = re.sub(r'\[END\]', '', text, flags=re.IGNORECASE)
+    text = re.sub(r'<\|endoftext\|>', '', text)
+    text = re.sub(r'</s>', '', text)
+    
+    # Remove excessive newlines
+    text = re.sub(r'\n{3,}', '\n\n', text)
+    
+    # Strip whitespace
+    text = text.strip()
+    
+    return text
+
 # =====================================================
 # REQUEST MODEL
 # =====================================================
@@ -136,7 +156,11 @@ async def ai_chat(data: QueryInput, x_api_key: str = Header(None)):
 
     try:
         response = chain.run(query=data.query.strip())
-        return {"reply": response.strip()}
+        
+        # Clean the response before sending
+        cleaned_response = clean_response(response)
+        
+        return {"reply": cleaned_response}
     except Exception as e:
         logger.error(f"⚠️ Model runtime error: {e}")
-        raise HTTPException(status_code=500, detail=f"Model failed to respond — {e}")
+        raise HTTPException(status_code=500, detail=f"Model failed to respond — {e}")