Upload 2 files

agent.py

@@ -340,7 +340,7 @@ class AgentRouter:
                     "tools_used": [],
                 }
 
-            final_text = accumulated.strip()
+            final_text = self.model.clean_response(accumulated)
             if not final_text:
                 final_text = self.model.generate(message=message, memory_context=memory_context, tool_context="")
 
@@ -417,7 +417,7 @@ class AgentRouter:
                 "tools_used": tools_used,
             }
 
-        final_text = accumulated.strip()
+        final_text = self.model.clean_response(accumulated)
         if not final_text:
             final_text = self.model.generate(message=message, memory_context=memory_context, tool_context=tool_context)
 

model.py

@@ -161,6 +161,9 @@ class ModelManager:
 
         return cleaned
 
+    def clean_response(self, text: str) -> str:
+        return self._clean_response(text)
+
     def generate(self, message: str, memory_context: str = "", tool_context: str = "") -> str:
         self.load()
         max_new_tokens = self.dynamic_token_budget(message)
@@ -233,30 +236,42 @@ class ModelManager:
         worker = threading.Thread(target=model.generate, kwargs=generation_kwargs, daemon=True)
         worker.start()
 
-        pieces = []
-        truncated = False
+        markers = ["\nUser:", "\nAssistant:", "\nSystem:", "User:", "Assistant:", "System:"]
+        buffer = ""
+        yielded_len = 0
+        stop_idx = -1
+
         for piece in streamer:
             if not piece:
                 continue
 
-            if truncated:
-                continue
+            buffer += piece
 
-            emitted = piece
-            for marker in ["\nUser:", "\nAssistant:", "\nSystem:"]:
-                idx = emitted.find(marker)
-                if idx != -1:
-                    emitted = emitted[:idx]
-                    truncated = True
-                    break
+            # Find earliest marker in accumulated text (handles marker split across chunks).
+            marker_positions = [buffer.find(m) for m in markers if buffer.find(m) != -1]
+            if marker_positions:
+                stop_idx = min(marker_positions)
 
-            if emitted:
-                pieces.append(emitted)
-                yield emitted
+            # Hold a short tail so markers crossing boundaries are still detected safely.
+            safe_upto = len(buffer) - 20 if stop_idx == -1 else stop_idx
+            if safe_upto > yielded_len:
+                out = buffer[yielded_len:safe_upto]
+                if out:
+                    yield out
+                yielded_len = safe_upto
+
+            if stop_idx != -1:
+                break
 
         worker.join(timeout=0.1)
 
-        final_text = self._clean_response("".join(pieces))
+        if stop_idx == -1 and yielded_len < len(buffer):
+            out = buffer[yielded_len:]
+            if out:
+                yield out
+
+        truncated_final = buffer[:stop_idx] if stop_idx != -1 else buffer
+        final_text = self._clean_response(truncated_final)
         if final_text:
             self._set_cached(key, final_text)