Enable readline for the chat prompt and persist history

mlx_lm/cli_ui.py

@@ -1,11 +1,21 @@
 # Copyright © 2024 Apple Inc.
 
+import os
 import re
 import shutil
 import sys
 from contextlib import contextmanager
 from functools import lru_cache
 
+try:
+    # Importing readline is what makes input() honour the \x01/\x02
+    # non-printing markers corridor_input() emits, and what provides history
+    # and line editing. It is unavailable on some platforms (e.g. Windows
+    # without pyreadline), so the prompt must keep working without it.
+    import readline
+except ImportError:
+    readline = None
+
 import mlx.core as mx
 from rich.box import ROUNDED
 from rich.console import Console
@@ -249,6 +259,32 @@ def report_save(self, checkpoint):
         )
 
 
+CHAT_HISTORY_FILE = os.path.expanduser("~/.mlx_lm_chat_history")
+CHAT_HISTORY_LENGTH = 1000
+
+
+def load_chat_history(path: str = CHAT_HISTORY_FILE) -> None:
+    """Restore prompt history from a previous session, if there is one."""
+    if readline is None:
+        return
+    readline.set_history_length(CHAT_HISTORY_LENGTH)
+    try:
+        readline.read_history_file(path)
+    except (OSError, PermissionError):
+        # No history yet, or it is unreadable. Neither is worth failing over.
+        pass
+
+
+def save_chat_history(path: str = CHAT_HISTORY_FILE) -> None:
+    """Persist prompt history for the next session."""
+    if readline is None:
+        return
+    try:
+        readline.write_history_file(path)
+    except (OSError, PermissionError):
+        pass
+
+
 class ChatUI:
     """Helper class for rendering the chat UI and streaming responses."""
 
@@ -259,6 +295,7 @@ def __init__(self, args, rank: int = 0):
 
     def __enter__(self):
         if self._rank == 0:
+            load_chat_history()
             rows = [("model", str(self._args.model))]
             if self._args.adapter_path:
                 rows.append(("adapter", str(self._args.adapter_path)))
@@ -273,6 +310,8 @@ def __enter__(self):
         return self
 
     def __exit__(self, *exc):
+        if self._rank == 0:
+            save_chat_history()
         return False
 
     def prompt(self) -> str:

tests/test_chat.py

@@ -1,7 +1,10 @@
 import argparse
+import os
+import tempfile
 import unittest
 from unittest.mock import MagicMock, patch
 
+from mlx_lm import cli_ui
 from mlx_lm.chat import setup_arg_parser
 
 
@@ -171,5 +174,59 @@ def test_no_system_prompt_in_messages(
         self.assertEqual(call_args[0]["content"], "What is the weather?")
 
 
+class TestChatHistory(unittest.TestCase):
+    """Prompt history is persisted across sessions, on rank 0 only."""
+
+    def setUp(self):
+        self.tmp = tempfile.TemporaryDirectory()
+        self.path = os.path.join(self.tmp.name, "history")
+
+    def tearDown(self):
+        self.tmp.cleanup()
+
+    def test_round_trip(self):
+        if cli_ui.readline is None:
+            self.skipTest("readline unavailable on this platform")
+        cli_ui.readline.clear_history()
+        cli_ui.readline.add_history("hello there")
+        cli_ui.save_chat_history(self.path)
+
+        cli_ui.readline.clear_history()
+        self.assertEqual(cli_ui.readline.get_current_history_length(), 0)
+
+        cli_ui.load_chat_history(self.path)
+        self.assertEqual(cli_ui.readline.get_history_item(1), "hello there")
+
+    def test_missing_file_is_not_an_error(self):
+        cli_ui.load_chat_history(os.path.join(self.tmp.name, "does-not-exist"))
+
+    def test_unwritable_path_is_not_an_error(self):
+        cli_ui.save_chat_history(os.path.join(self.tmp.name, "no", "such", "dir"))
+
+    def test_history_is_skipped_without_readline(self):
+        with patch.object(cli_ui, "readline", None):
+            cli_ui.load_chat_history(self.path)
+            cli_ui.save_chat_history(self.path)
+        self.assertFalse(os.path.exists(self.path))
+
+    def test_only_rank_zero_touches_history(self):
+        args = MagicMock()
+        args.model, args.adapter_path = "m", None
+        args.max_tokens, args.system_prompt = 10, None
+
+        with patch.object(cli_ui, "load_chat_history") as load, patch.object(
+            cli_ui, "save_chat_history"
+        ) as save:
+            with cli_ui.ChatUI(args, rank=1):
+                pass
+            load.assert_not_called()
+            save.assert_not_called()
+
+            with cli_ui.ChatUI(args, rank=0):
+                pass
+            load.assert_called_once()
+            save.assert_called_once()
+
+
 if __name__ == "__main__":
     unittest.main()