message assistent content

app.py

@@ -61,9 +61,15 @@ def refine_prompt(prompt):
             {"role": "user", "content": prompt},
         ]
         refined_prompt = text_gen(messages)
+        
         # Extract just the assistant's content from the response
         try:
-            assistant_content = refined_prompt[0]['generated_text'][-1]['content']
+            messages = refined_prompt[0]['generated_text']
+            # Find the last message with role 'assistant'
+            assistant_messages = [msg for msg in messages if msg['role'] == 'assistant']
+            if not assistant_messages:
+                return "Error: No assistant response found"
+            assistant_content = assistant_messages[-1]['content']
             return assistant_content
         except (KeyError, IndexError):
             return "Error: Unexpected response format from the model"