improve multi-line editing with prompt_toolkit

assistant.py

@@ -4,6 +4,9 @@ from ollama import Client
 import re
 import pyperclip
 import sys
+import tty
+import termios
+import signal
 import argparse
 import pygments
 from pygments.lexers import get_lexer_by_name, guess_lexer
@@ -11,6 +14,9 @@ from pygments.formatters import TerminalFormatter
 import os
 import json
 
+from prompt_toolkit.key_binding import KeyBindings
+from prompt_toolkit import PromptSession
+
 server = 'localhost:11434'
 model = 'gemma3:12b'
 reasoning_model='deepseek-r1:14b'
@@ -342,37 +348,24 @@ def handle_piped_input(args):
     if args.copy and len(blocks):
         copy_string_to_clipboard(blocks[0])
 
-def improved_input():
+kb = KeyBindings()
+
+@kb.add('c-d')
+def _(event):
+    event.current_buffer.validate_and_handle()
+
+session = PromptSession(multiline=True, prompt_continuation='', key_bindings=kb)
+def improved_input(prompt="> "):
     """
-    Handles multi-line input and prevents accidental sending of pasted content.
-    Returns the complete input text when the user indicates they're done.
+    Returns the full text (including embedded newlines) when you press Ctrl-D.
+    Arrow keys edit within or across lines automatically.
     """
-    lines = []
-    print("> ", end='', flush=True)
-    while True:
-        try:
-            line = sys.stdin.readline()
-            if not line:
-                print("", flush=True) # Some visual feedback that user input has ended
-                break  # EOF
-            if line.strip() == '':
-                continue  # ignore empty lines
-            lines.append(line)
-
-            # Check for termination signal (Ctrl+D or two consecutive empty lines)
-            if len(lines) > 1 and lines[-2] == '\n' and lines[-1] == '\n':
-                break
-
-        except KeyboardInterrupt:
-            print("\nUser aborted input")
-            return None
-        except Exception as e:
-            print(f"Error reading input: {e}")
-            return None
-
-    # Join the lines and strip any trailing newlines
-    full_input = ''.join(lines).rstrip('\n')
-    return full_input
+    try:
+        text = session.prompt(prompt)
+        return text
+    except KeyboardInterrupt:
+        print("\nUser aborted input")
+        return None
 
 def handle_non_piped_input(args):
     if args.shell: