diff --git a/mlx_lm/cli_ui.py b/mlx_lm/cli_ui.py index 133d1190e..1fe613fa3 100644 --- a/mlx_lm/cli_ui.py +++ b/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: diff --git a/tests/test_chat.py b/tests/test_chat.py index b4ddcf3ae..794ffff01 100644 --- a/tests/test_chat.py +++ b/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()