TUI的LLM界面

Table of Contents

tui_llm.png

一个模仿aider界面的实验，以gemini为例子，在terminal中和其进行交互，给它取个名字叫tullm。

LLM对象

class GeminiChat:
    def __init__(self, api_key, initial_message, model="gemini-2.0-flash"):
        self.client = genai.Client(api_key=api_key)
        self.chat = self.client.chats.create(model=model)
        self.history = []
        self._add_to_history("user", initial_message)
        response = self.chat.send_message(initial_message)
        self._add_to_history("assistant", response.text)

    def send(self, message):
        self._add_to_history("user", message)
        response = self.chat.send_message(message)
        self._add_to_history("assistant", response.text)
        return response.text

    def _add_to_history(self, role, text):
        self.history.append({"role": role, "text": text})

    def print_history(self):
        for message in self.history:
            print(f'Role: {message["role"]}')
            print(f'Message: {message["text"]}')
            print("-" * 50)

这是用保存向gemini发送和接收信息的对象，主要关注几个函数和变量

  • initial_message ： 用来初始化一些prompt，你可以预先设置这次对话模型的角色信息之类的。
  • send: 很明显就是向模型发送信息的函数。
  • history ： 一个用来保存聊天内容的list。

终端交互

class TerminalChat:
    class CommandCompleter(Completer):
        def __init__(self, commands):
            self.commands = commands

        def get_completions(self, document, complete_event):
            text = document.text
            if text.startswith("/"):
                word = text[1:]
                for command in self.commands:
                    if command.startswith(word):
                        yield Completion(command, start_position=-len(word))

    def __init__(
        self,
        send_function=None,
        recv_function=None,
        banner_function=None,
        custom_commands=None,
    ):
        self.send_function = send_function or self._default_send_function
        self.banner_function = banner_function or self._default_banner_function
        self.recv_function = recv_function or self._default_recv_function

        self.commands = {
            "quit": {
                "function": self._quit_command,
                "description": "Exit the chat application",
            },
            "help": {
                "function": self._help_command,
                "description": "Display available commands and their descriptions",
            },
        }

        if custom_commands:
            self.commands.update(custom_commands)

        self.completer = self.CommandCompleter(self.commands.keys())

        self.kb = KeyBindings()

        @self.kb.add("c-j")
        def _(event):
            event.current_buffer.insert_text("\n")

    def _default_send_function(self, user_input):
        return f"Echo: {user_input}"

    def _default_banner_function(self):
        print(f"{Fore.GREEN}Welcome to Terminal Chat!{Style.RESET_ALL}")
        print(f"{Fore.GREEN}Type '/help' for available commands{Style.RESET_ALL}")

    def _default_recv_function(self, text):
        print(f"{Fore.GREEN}Assistant: {Style.RESET_ALL}{text}")

    def _quit_command(self, _):
        print(f"{Fore.GREEN}Goodbye!{Style.RESET_ALL}")
        return True

    def _help_command(self, _):
        print(f"{Fore.GREEN}Available commands:{Style.RESET_ALL}")
        for cmd, cmd_info in sorted(self.commands.items()):
            description = cmd_info.get("description", "No description available")
            print(f"{Fore.GREEN}  /{cmd} - {description}{Style.RESET_ALL}")
        return False

    def clear_screen(self):
        os.system("cls" if os.name == "nt" else "clear")

    def print_separator(self):
        print(f"{Fore.GREEN}{'─' * os.get_terminal_size().columns}{Style.RESET_ALL}")

    def print_message(self, role, text):
        print(f"{Fore.GREEN}{role.capitalize()}: {Style.RESET_ALL}{text}")
        print("-" * 50)

    def start(self):
        self.print_separator()
        self.banner_function()

        while True:
            self.print_separator()
            try:
                user_input = prompt(
                    HTML('<style fg="green">></style> '),
                    completer=self.completer,
                    key_bindings=self.kb,
                )

                if not user_input.strip():
                    continue

                if user_input.startswith("/"):
                    command = user_input[1:].lower()
                    if command in self.commands:
                        should_quit = self.commands[command]["function"](user_input)
                        if should_quit:
                            break
                    else:
                        print(
                            f"{Fore.RED}Unknown command: {user_input}{Style.RESET_ALL}"
                        )
                        print(
                            f"{Fore.GREEN}Type '/help' for available commands{Style.RESET_ALL}"
                        )
                else:
                    response = self.send_function(user_input)
                    self.recv_function(response)
            except KeyboardInterrupt:
                os._exit(0)
            except EOFError:
                os._exit(0)

这是聊天的TUI界面，就是通过————这样的横线分割每次对话，里面用一个>的输入提示符，通过向模型发送信息最后返回。主要关注的点就是init初始化的那几个参数。

  • send_function ： 就是chat消息发送函数，你可将LLM的send函数，包装成一个接收message的函数然后传递给它。
  • recv_function ：是接收消息的函数，你也可以包装成接收text的函数，然后将返回的内容自行处理。
  • banner_function ：显示一些提示信息相关的函数，可以写一些命令和快捷键相关的提示内容。
  • custom_commands ：这是一个命令字典，里面的键值对是"[command name]": {"function": [function name],"description": "[describe]",}这样的格式，你通过设置命令名、函数名、描述来声明命令。然后在输入中通过前缀/[command]这样的格式进行命令的调用。

main

main逻辑就是对这两个对象的组合，先通过apikey和初始化信息，定义好gemini的对象。然后将其包装在一些函数中，通过这些函数，传递给TerminalChat这个对象的初始化。

通过这种形式，你可以灵活得构建多种的LLM对象来接入聊天，只要符合TerminalChat初始化参数那些函数格式就好。 

if __name__ == "__main__":
    api_key = os.getenv("GEMINI_API_KEY")
    initial_message = "Your name is allen."

    if api_key is None:
        print(
            f"{Fore.RED}Please set the GEMINI_API_KEY environment variable.{Style.RESET_ALL}"
        )
    else:
        gemini_chat = GeminiChat(api_key, initial_message)

        def send_to_gemini(message):
            return gemini_chat.send(message)

        def receive_from_gemini(text):
            print(f"{Fore.GREEN}Assistant: {Style.RESET_ALL}{text}")

        def clear_command(_):
            print(f"{Fore.GREEN}Clearing screen...{Style.RESET_ALL}")
            os.system("cls" if os.name == "nt" else "clear")
            return False

        def quit_command(_):
            print(f"{Fore.GREEN}Custom quit!{Style.RESET_ALL}")
            return True

        def history_command(_):
            terminal_chat.clear_screen()
            print(f"{Fore.GREEN}Chat History:{Style.RESET_ALL}")
            for message in gemini_chat.history:
                terminal_chat.print_message(message["role"], message["text"])
            return False

        def gemini_banner_function():
            print(f"{Fore.GREEN}Welcome to Gemini Terminal Chat!{Style.RESET_ALL}")
            print(
                f"{Fore.GREEN}Type '/quit' to quit, '/history' to see chat history, '/help' for more commands{Style.RESET_ALL}"
            )
            print(
                f"{Fore.GREEN}Press Ctrl+C or Ctrl+D to quit directly{Style.RESET_ALL}"
            )

        custom_commands = {
            "clear": {
                "function": clear_command,
                "description": "Clear the terminal screen",
            },
            "quit": {"function": quit_command, "description": "Exit the application"},
            "history": {
                "function": history_command,
                "description": "Display the full conversation history",
            },
        }

        terminal_chat = TerminalChat(
            send_function=send_to_gemini,
            recv_function=receive_from_gemini,
            banner_function=gemini_banner_function,
            custom_commands=custom_commands,
        )

        terminal_chat.start()

完整代码

import os

from colorama import init, Fore, Style
from prompt_toolkit import prompt
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.formatted_text import HTML
from prompt_toolkit.completion import Completer, Completion
from google import genai

class GeminiChat:
    def __init__(self, api_key, initial_message, model="gemini-2.0-flash"):
        self.client = genai.Client(api_key=api_key)
        self.chat = self.client.chats.create(model=model)
        self.history = []
        self._add_to_history("user", initial_message)
        response = self.chat.send_message(initial_message)
        self._add_to_history("assistant", response.text)

    def send(self, message):
        self._add_to_history("user", message)
        response = self.chat.send_message(message)
        self._add_to_history("assistant", response.text)
        return response.text

    def _add_to_history(self, role, text):
        self.history.append({"role": role, "text": text})

    def print_history(self):
        for message in self.history:
            print(f'Role: {message["role"]}')
            print(f'Message: {message["text"]}')
            print("-" * 50)


class TerminalChat:
    class CommandCompleter(Completer):
        def __init__(self, commands):
            self.commands = commands

        def get_completions(self, document, complete_event):
            text = document.text
            if text.startswith("/"):
                word = text[1:]
                for command in self.commands:
                    if command.startswith(word):
                        yield Completion(command, start_position=-len(word))

    def __init__(
        self,
        send_function=None,
        recv_function=None,
        banner_function=None,
        custom_commands=None,
    ):
        self.send_function = send_function or self._default_send_function
        self.banner_function = banner_function or self._default_banner_function
        self.recv_function = recv_function or self._default_recv_function

        self.commands = {
            "quit": {
                "function": self._quit_command,
                "description": "Exit the chat application",
            },
            "help": {
                "function": self._help_command,
                "description": "Display available commands and their descriptions",
            },
        }

        if custom_commands:
            self.commands.update(custom_commands)

        self.completer = self.CommandCompleter(self.commands.keys())

        self.kb = KeyBindings()

        @self.kb.add("c-j")
        def _(event):
            event.current_buffer.insert_text("\n")

    def _default_send_function(self, user_input):
        return f"Echo: {user_input}"

    def _default_banner_function(self):
        print(f"{Fore.GREEN}Welcome to Terminal Chat!{Style.RESET_ALL}")
        print(f"{Fore.GREEN}Type '/help' for available commands{Style.RESET_ALL}")

    def _default_recv_function(self, text):
        print(f"{Fore.GREEN}Assistant: {Style.RESET_ALL}{text}")

    def _quit_command(self, _):
        print(f"{Fore.GREEN}Goodbye!{Style.RESET_ALL}")
        return True

    def _help_command(self, _):
        print(f"{Fore.GREEN}Available commands:{Style.RESET_ALL}")
        for cmd, cmd_info in sorted(self.commands.items()):
            description = cmd_info.get("description", "No description available")
            print(f"{Fore.GREEN}  /{cmd} - {description}{Style.RESET_ALL}")
        return False

    def clear_screen(self):
        os.system("cls" if os.name == "nt" else "clear")

    def print_separator(self):
        print(f"{Fore.GREEN}{'─' * os.get_terminal_size().columns}{Style.RESET_ALL}")

    def print_message(self, role, text):
        print(f"{Fore.GREEN}{role.capitalize()}: {Style.RESET_ALL}{text}")
        print("-" * 50)

    def start(self):
        self.print_separator()
        self.banner_function()

        while True:
            self.print_separator()
            try:
                user_input = prompt(
                    HTML('<style fg="green">></style> '),
                    completer=self.completer,
                    key_bindings=self.kb,
                )

                if not user_input.strip():
                    continue

                if user_input.startswith("/"):
                    command = user_input[1:].lower()
                    if command in self.commands:
                        should_quit = self.commands[command]["function"](user_input)
                        if should_quit:
                            break
                    else:
                        print(
                            f"{Fore.RED}Unknown command: {user_input}{Style.RESET_ALL}"
                        )
                        print(
                            f"{Fore.GREEN}Type '/help' for available commands{Style.RESET_ALL}"
                        )
                else:
                    response = self.send_function(user_input)
                    self.recv_function(response)
            except KeyboardInterrupt:
                os._exit(0)
            except EOFError:
                os._exit(0)


if __name__ == "__main__":
    api_key = os.getenv("GEMINI_API_KEY")
    initial_message = "Your name is allen."

    if api_key is None:
        print(
            f"{Fore.RED}Please set the GEMINI_API_KEY environment variable.{Style.RESET_ALL}"
        )
    else:
        gemini_chat = GeminiChat(api_key, initial_message)

        def send_to_gemini(message):
            return gemini_chat.send(message)

        def receive_from_gemini(text):
            print(f"{Fore.GREEN}Assistant: {Style.RESET_ALL}{text}")

        def clear_command(_):
            print(f"{Fore.GREEN}Clearing screen...{Style.RESET_ALL}")
            os.system("cls" if os.name == "nt" else "clear")
            return False

        def quit_command(_):
            print(f"{Fore.GREEN}Custom quit!{Style.RESET_ALL}")
            return True

        def history_command(_):
            terminal_chat.clear_screen()
            print(f"{Fore.GREEN}Chat History:{Style.RESET_ALL}")
            for message in gemini_chat.history:
                terminal_chat.print_message(message["role"], message["text"])
            return False

        def gemini_banner_function():
            print(f"{Fore.GREEN}Welcome to Gemini Terminal Chat!{Style.RESET_ALL}")
            print(
                f"{Fore.GREEN}Type '/quit' to quit, '/history' to see chat history, '/help' for more commands{Style.RESET_ALL}"
            )
            print(
                f"{Fore.GREEN}Press Ctrl+C or Ctrl+D to quit directly{Style.RESET_ALL}"
            )

        custom_commands = {
            "clear": {
                "function": clear_command,
                "description": "Clear the terminal screen",
            },
            "quit": {"function": quit_command, "description": "Exit the application"},
            "history": {
                "function": history_command,
                "description": "Display the full conversation history",
            },
        }

        terminal_chat = TerminalChat(
            send_function=send_to_gemini,
            recv_function=receive_from_gemini,
            banner_function=gemini_banner_function,
            custom_commands=custom_commands,
        )

        terminal_chat.start()