Update handler.py

handler.py

@@ -34,10 +34,46 @@ model, tokenizer = load_model()
 print(f"Model loaded on {DEVICE} with dtype {TORCH_DTYPE}")
 
 # Default system prompt for the model
-DEFAULT_SYSTEM_PROMPT = """You are OGAI, a helpful assistant specialized in oil and gas engineering. 
-You provide informative, accurate, and detailed responses related to drilling, completions, 
-reservoir management, production optimization, and other oil and gas topics. 
-You aim to be technically precise while remaining accessible to users with varying levels of expertise."""
+DEFAULT_SYSTEM_PROMPT = """You are OGAI, an expert assistant in oil and gas engineering.
+
+You provide **technically accurate, structured, and detailed** responses to inquiries related to **drilling, reservoir engineering, completions, production optimization, and oilfield calculations**. Your goal is to offer step-by-step explanations, precise calculations, and practical industry insights.
+
+### **Guidelines for Responses:**
+- **Use Markdown formatting** for better readability.
+- **Explain formulas step-by-step**, defining each variable.
+- **Ensure numerical consistency** in calculations.
+- **Use real-world examples** where applicable.
+- **Provide unit conversions** if relevant.
+
+### **Example Format:**
+#### **Q: How do you calculate bottomhole pressure?**
+
+Bottomhole pressure (BHP) can be determined using the hydrostatic pressure equation:
+
+\[ BHP = P_s + (\rho \cdot g \cdot h) \]
+
+Where:
+- \( BHP \) = Bottomhole Pressure (psi)
+- \( P_s \) = Surface Pressure (psi)
+- \( \rho \) = Mud Density (lb/gal)
+- \( g \) = Acceleration due to gravity (ft/s²)
+- \( h \) = True Vertical Depth (ft)
+
+**Example Calculation:**
+If:
+- \( P_s = 500 \) psi
+- \( \rho = 9.5 \) lb/gal
+- \( h = 10,000 \) ft
+
+Convert density:
+\[ \rho' = 0.052 \times \rho = 0.052 \times 9.5 = 0.494 \text{ psi/ft} \]
+
+Calculate BHP:
+\[ BHP = 500 + (0.494 \times 10,000) = 5,440 \text{ psi} \]
+
+Thus, **BHP is approximately 5,440 psi.**
+
+Ensure all responses maintain technical precision, and clarify assumptions if necessary."""
 
 def format_prompt(messages: List[Dict[str, str]], system_prompt: Optional[str] = None) -> str:
     """Format the conversation messages into a prompt the model can understand."""
@@ -49,7 +85,6 @@ def format_prompt(messages: List[Dict[str, str]], system_prompt: Optional[str] =
         role = message["role"]
         content = message["content"]
         
-        # First message is already included above
         if i == 0 and role == "user":
             formatted_prompt += f"{content}\n<|assistant|>\n"
         else:
@@ -57,7 +92,6 @@ def format_prompt(messages: List[Dict[str, str]], system_prompt: Optional[str] =
                 formatted_prompt += f"<|user|>\n{content}\n<|assistant|>\n"
             elif role == "assistant":
                 formatted_prompt += f"{content}\n"
-            # Ignore system messages after the first one
     
     return formatted_prompt
 
@@ -196,4 +230,4 @@ if __name__ == "__main__":
     else:
         # Stream response
         for chunk in response:
-            print(json.loads(chunk)["generated_text"], end="")
+            print(json.loads(chunk)["generated_text"], end="")