Update app.py

app.py

@@ -55,14 +55,28 @@ SYSTEM_PROMPT = (
 )
 
 # --- Prompt Builder ---
-def build_prompt(messages, context=""):
-    prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}\n\nContext:\n{context}<|im_end|>\n"
-    for msg in messages:
-        role = msg["role"]
-        prompt += f"<|im_start|>{role}\n{msg['content']}<|im_end|>\n"
-    prompt += "<|im_start|>assistant\n"
+def build_prompt(messages, context="", model_name="Qwen"):
+    if "Mistral" in model_name:
+        # Alpaca-style prompt
+        prompt = f"You are DigiTwin, an expert in offshore inspection, maintenance, and asset integrity.\n"
+        if context:
+            prompt += f"Here is relevant context:\n{context}\n\n"
+        for msg in messages:
+            if msg["role"] == "user":
+                prompt += f"### Instruction:\n{msg['content'].strip()}\n"
+            elif msg["role"] == "assistant":
+                prompt += f"### Response:\n{msg['content'].strip()}\n"
+        prompt += "### Response:\n"
+    else:
+        # Qwen-style
+        prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}\n\nContext:\n{context}<|im_end|>\n"
+        for msg in messages:
+            role = msg["role"]
+            prompt += f"<|im_start|>{role}\n{msg['content']}<|im_end|>\n"
+        prompt += "<|im_start|>assistant\n"
     return prompt
 
+
 # --- Embed Uploaded Documents ---
 @st.cache_resource
 def embed_uploaded_files(files):
@@ -125,7 +139,7 @@ if prompt := st.chat_input("Ask something based on uploaded documents..."):
 
     # Limit to last 6 messages for memory
     recent_messages = st.session_state.messages[-6:]
-    full_prompt = build_prompt(recent_messages, context)
+    full_prompt = build_prompt(recent_messages, context, model_name=model_id)
 
     with st.chat_message("assistant", avatar=BOT_AVATAR):
         start = time.time()