Complete while typing with prompt_toolkit

Replace readline with a prompt_toolkit session: the completion menu opens
as you type, history suggestions show in gray (accept with Right), and
history lives in .cache/history. Enter on a highlighted candidate accepts
it instead of submitting the line. Ctrl+C now exits cleanly.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

embedlab.py

@@ -2,12 +2,17 @@
 
 print("Loading torch/transformers, please wait...", flush=True)
 
-import atexit
 import os
-import readline
 
 import torch
 import torch.nn.functional as F
+from prompt_toolkit import PromptSession
+from prompt_toolkit.application import get_app
+from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
+from prompt_toolkit.completion import Completer, Completion
+from prompt_toolkit.filters import Condition
+from prompt_toolkit.history import FileHistory
+from prompt_toolkit.key_binding import KeyBindings
 from transformers import AutoModel, AutoTokenizer
 
 # Model -> default parameters. The bge family uses CLS pooling.
@@ -176,48 +181,56 @@ def encode(texts):
     return torch.cat(outs)
 
 
-def completer(text, state):
-    """Tab completion: commands at line start, model names after /select."""
-    line = readline.get_line_buffer().lstrip()
+def candidates(line, text):
+    """Completions for the word `text` at the end of `line`: commands at line start, model names after /select."""
+    line = line.lstrip()
     if line.startswith("/select "):
-        options = [m for m in MODELS if m.startswith(text)]
-    elif line.startswith("/set "):
+        return [m for m in MODELS if m.startswith(text)]
+    if line.startswith("/set "):
         words = line.split(" ")
         if len(words) <= 2:  # completing the key
-            options = [k + " " for k in PARAM_KEYS if k.startswith(text)]
-        elif len(words) == 3 and words[1] == "pooling":
-            options = [v for v in ("cls", "mean") if v.startswith(text)]
-        else:
-            options = []
-    elif " " not in line:
-        options = [c + " " if c in ("/select", "/sim", "/tokens", "/set") else c for c in COMMANDS if c.startswith(text)]
-    else:
-        options = []
-    return options[state] if state < len(options) else None
+            return [k + " " for k in PARAM_KEYS if k.startswith(text)]
+        if len(words) == 3 and words[1] == "pooling":
+            return [v for v in ("cls", "mean") if v.startswith(text)]
+        return []
+    if " " not in line:
+        return [c + " " if c in ("/select", "/sim", "/tokens", "/set") else c for c in COMMANDS if c.startswith(text)]
+    return []
 
 
-def setup_readline():
-    readline.set_completer_delims(" \t\n")  # do not treat / and - as delimiters
-    readline.set_completer(completer)
-    doc = readline.__doc__ or ""
-    if "libedit" in doc:  # macOS
-        readline.parse_and_bind("bind ^I rl_complete")
-    else:
-        readline.parse_and_bind("tab: complete")
-    try:
-        readline.read_history_file(HISTORY_FILE)
-    except OSError:
-        pass
-    readline.set_history_length(1000)
-    atexit.register(save_history)
+class InputCompleter(Completer):
+    """Completes while typing."""
+
+    def get_completions(self, document, complete_event):
+        text = document.get_word_before_cursor(WORD=True)
+        for option in candidates(document.text_before_cursor, text):
+            yield Completion(option, start_position=-len(text))
 
 
-def save_history():
+def make_session():
     try:
         os.makedirs(os.path.dirname(HISTORY_FILE), exist_ok=True)
-        readline.write_history_file(HISTORY_FILE)
     except OSError:
         pass
+    bindings = KeyBindings()
+
+    @Condition
+    def completion_selected():
+        state = get_app().current_buffer.complete_state
+        return state is not None and state.current_completion is not None
+
+    @bindings.add("enter", filter=completion_selected)
+    def _(event):
+        # Enter on a highlighted candidate accepts it (already inserted) instead of submitting
+        event.current_buffer.complete_state = None
+
+    return PromptSession(
+        key_bindings=bindings,
+        completer=InputCompleter(),
+        complete_while_typing=True,
+        auto_suggest=AutoSuggestFromHistory(),
+        history=FileHistory(HISTORY_FILE),
+    )
 
 
 HELP = """Commands:
@@ -241,13 +254,13 @@ def main():
     parser.add_argument("--full", action="store_true", help="print the full vector")
     args = parser.parse_args()
 
-    setup_readline()
+    session = make_session()
     select_model(load_last_model())
     print(HELP)
     while True:
         try:
-            text = input("> ").strip()
-        except EOFError:
+            text = session.prompt("> ").strip()
+        except (EOFError, KeyboardInterrupt):
             break
         if not text or text in ("/quit", "/exit"):
             break