Update app.py

app.py

@@ -79,26 +79,34 @@ def generate_response(
 
     context = search_relevant_text(message)  # Get relevant content from PDF
 
-    messages = [{"role": "system", "content": system_message}]
-    for user_msg, bot_msg in history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if bot_msg:
-            messages.append({"role": "assistant", "content": bot_msg})
-
-    messages.append({"role": "user", "content": f"Context: {context}\nQuestion: {message}"})
-
-    response = ""
-    for chunk in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = chunk.choices[0].delta.content or ""
-        response += token
-    return response
+    # Format the prompt for Mistral
+    prompt = f"""<s>[INST] {system_message}
+
+Context from the PDF:
+{context}
+
+User Question: {message} [/INST]"""
+
+    # Add conversation history if it exists
+    for prev_msg, prev_response in history:
+        prompt += f" {prev_response} </s>[INST] {prev_msg} [/INST]"
+
+    try:
+        response = client.text_generation(
+            prompt,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            stream=True
+        )
+        
+        full_response = ""
+        for chunk in response:
+            full_response += chunk
+        return full_response
+    except Exception as e:
+        print(f"Error generating response: {str(e)}")
+        return "I apologize, but I encountered an error while generating the response. Please try again."
 
 @app.route('/')
 def index():