Update src/rag.py

src/rag.py

@@ -466,21 +466,39 @@ class AACAssistant:
         response = self.chain.invoke(
             {"question": user_query, "emotion_analysis": emotion_analysis}
         )
-        raw_full_answer = response.get("answer", "")
-        assistant_marker = "</s> <|assistant|>" 
-        
-        # Find the last occurrence of the marker
-        marker_position = raw_full_answer.rfind(assistant_marker)
+        raw_chain_output_answer = response.get("answer", "")
+        prompt_end_marker = "Please generate your response as the AAC user, following the instructions above.</s>\n<|assistant|>"
+
+        # For debugging, let's print what we're searching for and a snippet of where we're searching
+        print(f"DEBUG: process_query - Attempting to find marker: [{prompt_end_marker}]")
+        # print(f"DEBUG: process_query - Last 200 chars of raw_chain_output_answer: [...{raw_chain_output_answer[-200:]}]")
+
+
+        marker_position = raw_chain_output_answer.rfind(prompt_end_marker)
         
+        actual_response = "" 
         if marker_position != -1:
-            # Extract the text *after* the marker
-            actual_response = raw_full_answer[marker_position + len(assistant_marker):].strip()
+            # If the marker is found, take everything AFTER it
+            actual_response = raw_chain_output_answer[marker_position + len(prompt_end_marker):].strip()
+            print(f"DEBUG: process_query - Marker found. Extracted response before cleaning EOS: '{actual_response}'")
+
+            # Llama 3 models often output an <|eot_id|> at the end of their turn.
+            # Let's remove this if present.
+            eot_marker = "<|eot_id|>" 
+            if actual_response.endswith(eot_marker):
+                actual_response = actual_response[:-len(eot_marker)].strip()
+                print(f"DEBUG: process_query - Cleaned <|eot_id|>, final response: '{actual_response}'")
+            
         else:
-            # Fallback if the marker is not found in the response.
-            # This might happen if the LLM's output is unexpected or if the prompt structure changed.
-            print(f"WARNING: Assistant marker '{assistant_marker}' not found in raw answer. Returning raw answer as fallback.")
-            actual_response = raw_full_answer.strip() # Or handle as an error
-
-        print(f"DEBUG: process_query - Extracted assistant response: '{actual_response}'")
+            # This block will be hit if the precise prompt_end_marker isn't found.
+            # This indicates a mismatch between your defined marker and the actual raw output.
+            print(f"ERROR: Precise marker [{prompt_end_marker}] NOT FOUND in raw answer.")
+            print(f"DEBUG: process_query - Raw full answer from chain (length {len(raw_chain_output_answer)}):")
+            print(f"'''{raw_chain_output_answer}'''") # Print the whole thing for analysis
+            actual_response = "Error: Could not parse the assistant's response correctly." # Or return raw_chain_output_answer for debugging in UI
+
+        # --- END OF CORRECTED PARSING LOGIC ---
+        
+        print(f"DEBUG: process_query - Final extracted assistant response: '{actual_response}'")
         return actual_response
         # return response["answer"]