autocomplete-llm.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "openai",
#     "rich",
#     "prompt_toolkit",
# ]
# ///

import os
from openai import OpenAI
from rich.console import Console
from prompt_toolkit import PromptSession
from prompt_toolkit.key_binding import KeyBindings

# Configure the client for standard completions
client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="not-needed"
)

MODEL = "allenai/Olmo-3-1025-7B"
console = Console()

# Set up custom keybindings for the prompt
bindings = KeyBindings()

@bindings.add('enter')
def handle_enter(event):
    """Pressing Enter submits the text to the model."""
    event.current_buffer.validate_and_handle()

@bindings.add('escape', 'enter')
def handle_alt_enter(event):
    """Pressing Alt+Enter (or Esc then Enter) inserts a newline instead of submitting."""
    event.current_buffer.insert_text('\n')

def main():
    console.print(f"[bold green]Connected to {MODEL} (Base Model Mode)[/bold green]")
    console.print("This interface acts like an advanced autocomplete. It maintains a continuous block of text.\n")
    console.print("[bold]Commands:[/bold]")
    console.print("  [blue][Enter][/blue]      : Generate more tokens based on the current context")
    console.print("  [blue][Alt+Enter][/blue]  : Insert a new line without submitting (keep typing!)")
    console.print("  [blue]/retry[/blue]       : Undo the last generated chunk and try again")
    console.print("  [blue]/pop [n][/blue]     : Delete the last 'n' characters from the context")
    console.print("  [blue]/clear[/blue]       : Wipe the context entirely and start over")
    console.print("  [blue]/quit[/blue]        : Exit the script\n")

    context = ""
    last_generated_chunk = ""
    
    # PromptSession manages the input state beautifully
    session = PromptSession()

    while True:
        try:
            # Determine the prompt indicator
            prompt_text = "[Start typing]📝: " if not context else "\n[Add text, Enter to continue, or Command]✍️: "

            # Grab input using prompt_toolkit with our custom keybindings
            user_input = session.prompt(
                prompt_text, 
                multiline=True, 
                key_bindings=bindings
            )

            # Process commands
            cmd_parts = user_input.strip().lower().split()
            cmd = cmd_parts[0] if cmd_parts else ""

            if cmd in ["/quit", "/exit"]:
                break
            elif cmd == "/clear":
                context = ""
                last_generated_chunk = ""
                console.print("[yellow]Context cleared.[/yellow]\n")
                continue
            elif cmd == "/retry":
                if last_generated_chunk:
                    context = context[:-len(last_generated_chunk)]
                    console.print("[yellow]Discarded last chunk. Regenerating...[/yellow]\n")
                    console.print(f"[dim]{context}[/dim]", end="")
                else:
                    console.print("[red]Nothing to retry yet.[/red]")
                    continue
            elif cmd == "/pop":
                try:
                    chars_to_pop = int(cmd_parts[1]) if len(cmd_parts) > 1 else 1
                    if chars_to_pop >= len(context):
                        context = ""
                        console.print("[yellow]Popped all characters. Context is empty.[/yellow]")
                    else:
                        context = context[:-chars_to_pop]
                        console.print(f"[yellow]Popped {chars_to_pop} characters.[/yellow]")
                        console.print(f"[dim]Current context:[/dim]\n{context}")
                    last_generated_chunk = ""  
                    continue
                except ValueError:
                    console.print("[red]Usage: /pop [number of characters][/red]")
                    continue
            else:
                # If there's input, append it exactly as formatted (including actual newlines)
                if user_input:
                    context += user_input
            
            if not context:
                continue

            # Generation Step
            console.print("\n[bold purple] OLMo 🤖:[/bold purple] ", end="")
            
            response = client.completions.create(
                model=MODEL,
                prompt=context,
                temperature=0.7,
                max_tokens=64, 
                stream=True
            )

            last_generated_chunk = ""
            for chunk in response:
                if chunk.choices[0].text:
                    text = chunk.choices[0].text
                    print(text, end="", flush=True)
                    last_generated_chunk += text
                    context += text
            print() 

        except KeyboardInterrupt:
            break
        except Exception as e:
            console.print(f"\n[bold red]Error:[/bold red] {e}\n")

if __name__ == "__main__":
    main()