My chatgpt command line client
from setuptools import setup

with open("requirements.txt") as f:
    requirements = f.read().splitlines()

setup(
    name="gptline",
    version="1.0.7",
    packages=["src"],
    install_requires=requirements,
    entry_points={
        "console_scripts": [
            "gptline = src.main:main"
        ]
    },
)

setup.py
gnachman/gptline
import threading

class BackgroundTask:
    def __init__(self, func):
        self._thread = threading.Thread(target=self._run, args=(func,))
        self._result = None
        self.exception = None
        self._done = False
        self._thread.start()

    def _run(self, func):
        try:
            self._result = func()
        except Exception as e:
            print(e)
            self.exception = e
        self._done = True

    def done(self):
        return self._done

    def result(self):
        return self._result



src/background_task.py
gnachman/gptline
import simplejson
import enum
import inspect
import openai
import sys
import threading
import time
import typing
from src.spin import spin

def get_json_type_name(value):
    if isinstance(value, str):
        return "string"
    elif isinstance(value, int):
        return "integer"
    elif isinstance(value, float):
        return "number"
    elif isinstance(value, bool):
        return "boolean"
    else:
        return "null"

def _json_schema(func):
    api_info = {}

    # Get function name
    api_info["name"] = func.__name__

    # Get function description from docstring
    docstring = inspect.getdoc(func)
    if docstring:
        api_info["description"] = docstring.split("\n\n")[0]

    # Get function parameters
    parameters = {}
    parameters["type"] = "object"

    properties = {}

    signature = inspect.signature(func)
    required = []
    for param_name, param in signature.parameters.items():
        param_info = {}

        # Get parameter type from type hints
        if typing.get_origin(param.annotation) is typing.Union and type(None) in typing.get_args(param.annotation):
            # Handle Optional case
            inner_type = typing.get_args(param.annotation)[0]
            if issubclass(inner_type, enum.Enum):
                param_info["type"] = get_json_type_name(inner_type.__members__[list(inner_type.__members__)[0]].value)
                param_info["enum"] = [member.value for member in inner_type]
            else:
                param_info["type"] = get_json_type_name(inner_type.__name__)
        elif issubclass(param.annotation, enum.Enum):
            param_info["type"] = get_json_type_name(param.annotation.__members__[list(param.annotation.__members__)[0]].value)
            param_info["enum"] = [member.value for member in param.annotation]
        else:
            param_info["type"] = get_json_type_name(param.annotation.__name__)

        # Get parameter description from docstring
        if docstring and param_name in docstring:
            param_info["description"] = docstring.split(param_name + ":")[1].split("\n")[0].strip()

        # Check if parameter is required
        if param.default == inspect.Parameter.empty:
            required.append(param_name)

        # Add parameter info to parameters dict
        properties[param_name] = param_info

    # Add parameters to api_info
    parameters["properties"] = properties
    parameters["required"] = required
    api_info["parameters"] = parameters

    return api_info

def create_chat_with_spinner(messages, temperature, functions, model):
    return create_chat(messages, temperature, functions, model, True)

def create_chat(messages, temperature, functions, model, spinner=False):
    def create_chat_model():
        args = {
            "model": model,
            "messages": messages,
            "stream": True,
            "temperature": temperature,
        }
        if functions:
            args['functions'] = list(map(lambda f: _json_schema(f), functions))
        return openai.ChatCompletion.create(**args)

    if not spinner:
        chats = []
        return create_chat_model()
    else:
        return spin(create_chat_model)


def suggest_name(chat_id, message):
    try:
        chat_completion_resp = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You assign names to conversations based on the first message. Respond with only a short, descriptive title for a conversation."},
                {"role": "user", "content": message}
            ],
            temperature=0,
            max_tokens=10
        )
        name = chat_completion_resp.choices[0].message.content
        return (chat_id, name)
    except Exception as e:
        return None

def invoke(functions, name, args_str):
    try:
        args = simplejson.loads(args_str, strict=False)
    except Exception as e:
        return invoke_sloppy(functions, name, args_str)

    for func in functions:
        if func.__name__ == name:
            return func(**args)
    else:
        raise ValueError(f"Function '{name}' not found.")

def invoke_sloppy(functions, name, args_str):
    for func in functions:
        if func.__name__ == name and func.sloppy:
            args = [None] * len(inspect.signature(func).parameters)
            args[0] = args_str
            return func(*args)
    else:
        raise ValueError(f"Function '{name}' not found.")

src/chat.py
gnachman/gptline
import sqlite3
from typing import Optional
import fcntl
import json
import os
import sys

def fullpath(file):
    xdg_root = os.getenv("XDG_ROOT")
    if xdg_root:
        return os.path.join(xdg_root, file)
    else:
        home_dir = os.path.expanduser("~")
        return os.path.join(home_dir, file)

lock_fd = None

class ChatDB:
    def __init__(self, db_file=".chatgpt.db"):
        db_path = fullpath(db_file)
        self.conn = sqlite3.connect(db_path)
        self.conn.execute("PRAGMA journal_mode=WAL")  # Enable WAL mode
        self.cursor = self.conn.cursor()
        self.create_schema()

    def create_schema(self):
        query = """
        CREATE TABLE IF NOT EXISTS chats (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            name TEXT,
            last_update TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        """
        self.conn.execute(query)

        query = """
        CREATE TABLE IF NOT EXISTS messages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            chat_id INTEGER,
            role TEXT,
            content TEXT NULL,
            time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            deleted INTEGER DEFAULT 0,
            function_call_name TEXT,
            function_call_arguments TEXT,
            FOREIGN KEY (chat_id) REFERENCES chats (id)
        )
        """
        self.conn.execute(query)

        query = """
        CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING FTS5 (
            message_id UNINDEXED,
            content,
            content_rowid,
        )
        """
        self.conn.execute(query)

        query = """
        CREATE VIRTUAL TABLE IF NOT EXISTS chat_fts USING FTS5 (
            chat_id UNINDEXED,
            content
        )
        """
        self.conn.execute(query)

        query = """
        CREATE TABLE IF NOT EXISTS settings (
            key TEXT NOT NULL UNIQUE,
            value TEXT NOT NULL
        )
        """
        self.conn.execute(query)

        query = """
        CREATE TABLE IF NOT EXISTS kvs (
            key TEXT NOT NULL UNIQUE,
            value TEXT NOT NULL
        )
        """
        self.conn.execute(query)

    def create_chat(self, name=None):
        query = "INSERT INTO chats (name) VALUES (?)"
        self.cursor.execute(query, (name,))
        chat_id = self.cursor.lastrowid  # Access lastrowid from the cursor
        self.conn.commit()
        return chat_id

    def add_message(self, chat_id: int, role: str, content: str, function_call_name: Optional[str], function_call_arguments: Optional[str]):
        query = "INSERT INTO messages (chat_id, role, content, function_call_name, function_call_arguments) VALUES (?, ?, ?, ?, ?)"
        with self.conn:
            self.cursor.execute(query, (chat_id, role, content, function_call_name, function_call_arguments))
            last_message_id = self.cursor.lastrowid

            if content is not None and role != "function":
                fts_query = "INSERT INTO messages_fts (message_id, content) VALUES (?, ?)"
                self.conn.execute(fts_query, (last_message_id, content.lower()))

            query = f"UPDATE chats SET last_update = CURRENT_TIMESTAMP WHERE id = ?"
            self.cursor.execute(query, (chat_id, ))

        return last_message_id

    def set_chat_fulltext(self, chat_id, content):
        # Remove existing chat_fts row with chat_id
        delete_query = "DELETE FROM chat_fts WHERE chat_id = ?"
        self.cursor.execute(delete_query, (chat_id,))

        # Insert new chat_fts row with chat_id and content
        insert_query = "INSERT INTO chat_fts (chat_id, content) VALUES (?, ?)"
        self.cursor.execute(insert_query, (chat_id, content))

        # Commit the changes
        self.conn.commit()


    def num_messages(self, chat_id: int) -> int:
        query = "SELECT COUNT(*) FROM messages WHERE chat_id = ?"
        result = self.conn.execute(query, (chat_id,)).fetchone()
        return result[0]

    def get_message_by_id(self, message_id: int):
        query = "SELECT role, content, time, id, deleted FROM messages WHERE id = ?"
        result = self.conn.execute(query, (message_id,)).fetchone()
        if result:
            return result
        else:
            raise IndexError("Index out of range")

    def get_message_by_index(self, chat_id: int, index: int):
        query = "SELECT role, content, time, id, deleted, function_call_name, function_call_arguments FROM messages WHERE chat_id = ? ORDER BY id LIMIT 1 OFFSET ?"
        result = self.conn.execute(query, (chat_id, index)).fetchone()
        if result:
            return result
        else:
            raise IndexError("Index out of range")

    def list_chats(self):
      query = """
      SELECT chats.id, chats.name, chats.last_update, COUNT(messages.id) as num_messages
      FROM chats
      LEFT JOIN messages ON chats.id = messages.chat_id
      GROUP BY chats.id
      HAVING COUNT(messages.id) > 1
      ORDER BY chats.id DESC
      """
      result = self.conn.execute(query).fetchall()
      return result

    def set_chat_name(self, chat_id: int, name: str):
        query = "UPDATE chats SET name = ? WHERE id = ?"
        self.conn.execute(query, (name, chat_id))
        self.conn.commit()

    def get_chat_name(self, chat_id):
        query = "SELECT name FROM chats WHERE id = ?"
        cursor = self.conn.execute(query, (chat_id,))
        result = cursor.fetchone()
        if result:
            return result[0]
        else:
            return None

    def delete_message(self, message_id: int):
        query = """
        UPDATE messages
        SET deleted = 1
        WHERE id = ?
        """
        self.conn.execute(query, (message_id,))
        self.conn.commit()

    def tokenize_fts(self, query):
        words = query.split(" ")
        escaped_words = [word.replace('"', '""') for word in words]
        quoted_words = [f'"{w}"' for w in escaped_words]
        return " ".join(quoted_words)

    def search_chats(self, query: str, pagination_token: int, limit: int):
        """
        Searches chats across the contents of messages.
        Returns ([(chat ID, snippet), ...], next pagination token)
        """

        query = self.tokenize_fts(query)
        fts_query = """
            SELECT chat_id, snippet(chat_fts, 1, '\ue000', '\ue001', '...', 16) AS snippet
            FROM chat_fts
            WHERE content MATCH ?
            LIMIT ?
            OFFSET ?
            """
        if pagination_token is None:
            offset = 0
        else:
            offset = pagination_token
        parameters = (query.lower(), limit, offset)
        print(fts_query)
        print(query)

        results = list(self.conn.execute(fts_query, parameters))
        return results, (offset + len(results))


    def search_messages_in_chat(self, query: str, chat_id: int):
        query = self.tokenize_fts(query)
        fts_query = """
            SELECT m.id
            FROM messages_fts
            JOIN messages m ON messages_fts.message_id = m.id
            WHERE messages_fts.content MATCH ? AND m.chat_id=?
            ORDER BY m.id ASC
        """

        parameters = (query.lower(), chat_id)
        return [row[0] for row in self.conn.execute(fts_query, parameters)]

    def search_messages(self, query: str, pagination_token: int, limit: int):
        query = self.tokenize_fts(query)
        fts_query = """
            SELECT m.id, m.chat_id, snippet(messages_fts, 1, '\ue000', '\ue001', '...', 16) AS snippet
            FROM messages_fts
            JOIN messages m ON messages_fts.message_id = m.id
            WHERE messages_fts.content MATCH ?
            LIMIT ?
            OFFSET ?
        """

        if pagination_token is None:
            offset = 0
        else:
            offset = pagination_token

        parameters = (query.lower(), limit, offset)
        result = self.conn.execute(fts_query, parameters)

        message_ids = []
        chat_ids = []
        snippets = {}
        for row in result:
            message_ids.append(int(row[0]))
            snippets[int(row[0])] = row[2]
            chat_ids.append(int(row[1]))

        # Sort chat IDs
        sorted_chat_ids = sorted(set(chat_ids))

        messages_by_chat = []
        for chat_id in sorted_chat_ids:
            message_ids_for_chat = [(message_ids[i], snippets[message_ids[i]]) for i in range(len(message_ids)) if chat_ids[i] == chat_id]
            messages_by_chat.append((chat_id, message_ids_for_chat))

        # Determine the pagination token
        next_offset = offset + limit
        has_more_results = len(message_ids) > next_offset
        pagination_token = next_offset if has_more_results else None

        return messages_by_chat, offset + len(message_ids)

    def get_setting(self, name: str, default_value=None):
        query = "SELECT value FROM settings WHERE key = ?"
        cursor = self.conn.execute(query, (name,))
        result = cursor.fetchone()
        if result:
            try:
                return json.loads(result[0])
            except json.JSONDecodeError:
                return default_value
        else:
            return default_value

    def set_setting(self, name: str, value):
        encoded_value = json.dumps(value)
        query = "INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)"
        self.conn.execute(query, (name, encoded_value))
        self.conn.commit()

    def get_kv(self, name: str, default_value=None):
        query = "SELECT value FROM kvs WHERE key = ?"
        cursor = self.conn.execute(query, (name,))
        result = cursor.fetchone()
        if result:
            try:
                return json.loads(result[0])
            except json.JSONDecodeError:
                return default_value
        else:
            return default_value

    def set_kv(self, name: str, value):
        encoded_value = json.dumps(value)
        query = "INSERT OR REPLACE INTO kvs (key, value) VALUES (?, ?)"
        self.conn.execute(query, (name, encoded_value))
        self.conn.commit()



src/db.py
gnachman/gptline
from prompt_toolkit import print_formatted_text
from prompt_toolkit.formatted_text import HTML
from src.highlight import SyntaxHighlighter
import datetime
import html
import os

def formattedTime(timestamp):
    dt = datetime.datetime.strptime(timestamp, "%Y-%m-%d %H:%M:%S")
    dt = dt.replace(tzinfo=datetime.timezone.utc)
    dt = dt.astimezone()
    return dt.strftime("%b %d, %Y at %I:%M %p")

def print_message(timestamp, role, content, deleted, prefix=""):
    s = prefix + f"[{html.escape(formattedTime(timestamp))}] <i>{html.escape(role)}</i>:"
    if deleted:
        print_formatted_text(HTML("<strike>" + s + "</strike>"))
    else:
        print_formatted_text(HTML(s))
    if role == "assistant":
        sh = SyntaxHighlighter()
        sh.put(content)
        sh.eof()
    else:
        print(content, end='')

def print_snippet(timestamp, role, content, deleted, prefix=""):
    s = prefix + f"[{html.escape(formattedTime(timestamp))}] <i>{html.escape(role)}</i>: " + content
    if deleted:
        print_formatted_text(HTML("<strike>" + s + "</strike>"))
    else:
        print_formatted_text(HTML(s))


def setMark():
    if 'TERM' in os.environ and os.environ['TERM'].startswith('screen'):
        osc = "\033Ptmux;\033\033]"
        st = "\a\033\\"
    else:
        osc = "\033]"
        st = "\a"

    print(f"{osc}1337;SetMark{st}", end='')


src/formatting.py
gnachman/gptline
from blessings import Terminal
from dataclasses import dataclass
import re
import sys
import termios
import tty
import uuid
from pygments.lexers import get_lexer_by_name
from pygments import highlight
from pygments.formatters import TerminalTrueColorFormatter
from pygments.style import Style
from pygments.token import Token
from pygments.styles import get_style_by_name


@dataclass
class Theme:
    code_block_banner_bg: str
    code_block_banner_fg: str
    code_block_bg: str
    inline_code: str

def foreground_color(r, g, b):
    return f"\033[38;2;{r};{g};{b}m"

def background_color(r, g, b):
    return f"\033[48;2;{r};{g};{b}m"

class SyntaxHighlighter:
    def __init__(self):
        self.term = Terminal()
        self.buffer = ""
        self.is_bold = False
        self.is_italic = False
        self.is_underline = False
        self.is_inline_code = False
        self.is_block_code = False
        self.block_id = None
        self.will_enter_block_code = False
        self.prev_attrs = []
        self.markup_eligible = True
        self.line = ""
        self.language = None
        darkTheme = Theme(background_color(55, 55, 55),
                          foreground_color(215, 215, 215),
                          background_color(15, 15, 15),
                          self.term.bold)

        lightTheme = Theme(background_color(215, 215, 215),
                           foreground_color(15, 15, 15),
                           background_color(240, 240, 240),
                           self.term.bold)

        if is_background_dark():
            self.theme = darkTheme
            self.formatter = TerminalTrueColorFormatter()
        else:
            self.theme = lightTheme
            self.formatter = TerminalTrueColorFormatter()
        self.lexer = None
        self.lex_buffer = ""

    def put(self, s):
        self._put(s)
        sys.stdout.flush()

    def _put(self, s):
        if len(s) == 0:
            return

        self.buffer += s
        while len(self.buffer) > 0:
            if not self.is_block_code and self.buffer.startswith("**"):
                self.is_bold = not self.is_bold
                self.buffer = self.buffer[2:]
            elif not self.is_block_code and self.buffer.startswith("*") and len(self.buffer) > 1 and self.buffer[1] != "*" and self.buffer[1] != " " and self.buffer[1] != "\n":
                self.is_italic = not self.is_italic
                self.buffer = self.buffer[1:]
            elif not self.is_block_code and self.buffer == "*":
                return
            elif not self.is_block_code and self.buffer.startswith("_") and self.markup_eligible:
                self.is_underline = not self.is_underline
                self.buffer = self.buffer[1:]
            elif not self.is_inline_code and self.buffer == "`":
                # A single backtick is ambiguous. Could be an inline block or a code block.
                return
            elif not self.is_block_code and not self.is_inline_code and self.buffer.startswith("`") and len(self.buffer) > 1 and self.buffer[1] != "`" and self.line and self.line[-1] == ' ':
                # Start an inline block
                self.is_inline_code = True
                self.buffer = self.buffer[1:]
            elif not self.is_block_code and self.is_inline_code and self.buffer.startswith("\n"):
                # End an inline block at newline to avoid it running away
                self.is_inline_code = False
            elif not self.is_block_code and self.is_inline_code and self.buffer.startswith("`"):
                # End an inline block
                self.is_inline_code = False
                self.buffer = self.buffer[1:]
            elif not self.is_block_code and not self.will_enter_block_code and re.match(r'^[ ]{0,8}$', self.line) is not None and self.buffer.startswith("```"):
                # Start a code block
                if '\n' not in self.buffer[1:]:
                    # Ignore language, which follows ```
                    return
                self.is_block_code = True
                self.block_id = uuid.uuid4()
                self.will_enter_block_code = True
                # Remove up to and including first newline
                first, second = self.buffer.split('\n', 1)
                if len(first) > 4:
                  # Print language
                  print(self.theme.code_block_banner_bg + self.theme.code_block_banner_fg, end='')
                  print(first[3:] + self.erase_line() + "    " + self.copyButton(self.block_id) + " " + self.hyperlink(f"iterm2:copy-block?block={self.block_id}", "Copy to clipboard") + self.term.normal)
                  self.language = first[3:]
                else:
                  self.language = "text/plain"
                try:
                    self.lexer = get_lexer_by_name(first[3:])
                except:
                    pass
                self.prev_attrs = None
                self.buffer = second
            elif self.is_block_code and not self.will_enter_block_code and self.buffer.startswith("```\n"):
                # End a code block
                if not self.will_enter_block_code:
                    self.emit(self.end_block(self.block_id) + self.term.normal)
                self.block_id = None
                self.will_enter_block_code = False
                self.buffer = self.buffer[4:]
                self.is_block_code = False
                self.lexer = None
            elif self.is_block_code and (self.buffer == "`" or self.buffer == "``"):
                # Too soon to tell
                return
            elif not self.is_block_code and re.match(r"\[.*\]\(.*\)", self.buffer):
                link_text, link_url = re.findall(r"\[(.*?)\]\((.*?)\)", self.buffer)[0]
                self.buffer = self.buffer[len(link_text) + len(link_url) + 4:]
                self.emit(self.hyperlink(link_url, link_text))
            elif not self.is_block_code and (self.buffer.startswith('[') or self.buffer == "``") and '\n' not in self.buffer:
                return
            else:
                self.printFirst()

    def copyButton(self, block_id):
        return self.osc(1337) + f'Button=type=copy;block={block_id}' + self.st()

    def erase_line(self):
        return '\033[K'

    def osc(self, n):
        return f'\033]{n};'

    def st(self):
        return f'\a'

    def hyperlink(self, href, anchor):
        return self.osc(8) + f';{href}' + self.st() + anchor + self.osc(8) + ';' + self.st()

    def start_block(self, id, type):
        return self.osc(1337) + f'Block=attr=start;id={id};type={type}' + self.st()

    def end_block(self, id):
        return self.osc(1337) + f'Block=attr=end;id={id};render=0' + self.st()

    def eof(self):
        while len(self.buffer) > 0:
            self.printFirst()
        if self.is_block_code:
            self.emit(self.end_block(self.block_id))

    def lex(self):
        highlighted_code = highlight(self.lex_buffer, self.lexer, self.formatter)
        highlighted_code = highlighted_code.replace("\n", self.theme.code_block_bg + self.erase_line() + "\n")
        print(highlighted_code, end='')
        self.lex_buffer = ""

    def printFirst(self):
        attrs = []
        if self.is_block_code:
            attrs.append(self.theme.code_block_bg)
        if self.will_enter_block_code:
            self.emit(self.theme.code_block_bg + self.erase_line())
            self.emit(self.start_block(self.block_id, self.language))
            self.will_enter_block_code = False
            attrs.append(self.theme.code_block_bg)
        if self.lexer:
            if self.buffer[0] == "\n":
                self.lex()
                self.line = ""
            else:
                self.lex_buffer += self.buffer[0]
            self.buffer = self.buffer[1:]
            return

        if self.is_bold:
            attrs.append(self.term.bold)
        elif self.is_italic:
            attrs.append(self.term.italic)
        elif self.is_underline:
            attrs.append(self.term.underline)
        elif self.is_inline_code:
            attrs.append(self.theme.inline_code)
        if self.prev_attrs != attrs:
            self.emit(self.term.normal + "".join(attrs) + self.buffer[0])
            self.prev_attrs = attrs
        else:
            self.emit(self.buffer[0])
        self.markup_eligible = self.buffer[0] == ' ' or self.buffer[0] == '\n'

        if self.buffer[0] == "\n":
            self.line = ""
            self.emit(self.erase_line())
        else:
            self.line += self.buffer[0]
        self.buffer = self.buffer[1:]

    def emit(self, chunk):
        print(chunk, end='')

def is_background_dark():
    # Save the current terminal settings
    old_settings = termios.tcgetattr(sys.stdin)

    try:
        # Set the terminal to raw mode
        tty.setraw(sys.stdin)

        # Write the query control sequence to the terminal
        sys.stdout.write('\033]11;?\033\\')
        sys.stdout.flush()

        # Read the response from the terminal
        response = ''
        while True:
            char = sys.stdin.read(1)
            response += char
            if response.endswith('\033\\'):
                break

        # Extract the background color information from the response
        match = re.search(r'rgb:([0-9a-fA-F]{4})/([0-9a-fA-F]{4})/([0-9a-fA-F]{4})', response)
        if match:
            r = int(match.group(1), 16) / 65535.0
            g = int(match.group(2), 16) / 65535.0
            b = int(match.group(3), 16) / 65535.0
            luminance = (0.299 * r + 0.587 * g + 0.114 * b)
            return luminance < 0.5

        return False

    finally:
        # Restore the terminal settings
        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings)

src/highlight.py
gnachman/gptline
from dataclasses import dataclass
from prompt_toolkit import PromptSession, print_formatted_text, Application
from prompt_toolkit.application.current import get_app
from prompt_toolkit.formatted_text import HTML, FormattedText
from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.keys import Keys
from prompt_toolkit.layout.containers import Window
from prompt_toolkit.layout.controls import BufferControl
from prompt_toolkit.layout.layout import Layout
from prompt_toolkit.widgets import Frame
from src.formatting import formattedTime
from src.ui_utils import draw_horizontal_line
from typing import Optional
import html
import tiktoken

@dataclass
class UserInput:
    text: Optional[str] = None
    # -1 to create a new chat.
    chat_identifier: Optional[int] = None
    query: Optional[str] = None
    message_query: Optional[str] = None
    regenerate = False
    edit = False
    allow_execution = False
    settings = False

@dataclass
class Chat:
    chat_identifier: int
    name: str
    last_update: str
    num_messages: int

def usage(model, text):
    try:
        encoding = tiktoken.encoding_for_model(model)
        return len(encoding.encode(text))
    except Exception as e:
        return 0


# Returns UserInput
def read_input(chats, current_chat_name, have_any_messages, placeholder, allow_execution, used, max_tokens, model):
    result = UserInput()
    tokens_typed = [0]
    result.allow_execution = allow_execution

    session = PromptSession()

    def update_tokens_typed(_):
        buffer = session.default_buffer
        tokens_typed[0] = usage(model, buffer.text)
        session.app.invalidate()

    session.default_buffer.on_text_changed += update_tokens_typed
    kb = KeyBindings()

    # Add a binding for the 'Enter' key to insert a newline character
    @kb.add(Keys.Enter)
    def _(event):
        event.current_buffer.insert_text('\n')

    # Add a binding for ^D to accept the input.
    @kb.add(Keys.ControlD)
    def _(event):
        buffer = event.current_buffer
        if buffer.cursor_position < len(buffer.text):
            buffer.delete()
        else:
            buffer.validate_and_handle()

    SHOW_CHAT_LIST = "$$$SHOW_CHAT_LIST"
    NEW_CHAT = "$$$NEW_CHAT"
    SEARCH = "$$$SEARCH"
    SEARCH_THIS_CHAT = "$$$SEARCH_THIS_CHAT"
    REGENERATE = "$$$REGENERATE"
    EDIT = "$$$EDIT"
    TOGGLE_SETTING = "$$$TOGGLE_SETTING"
    SETTINGS = "$$$SETTINGS"

    @kb.add(Keys.F2)
    def _(event):
        app = get_app()
        app.exit(result=SHOW_CHAT_LIST)

    @kb.add(Keys.F3)
    def _(event):
        app = get_app()
        app.exit(result=NEW_CHAT)

    @kb.add(Keys.F4)
    def _(event):
        app = get_app()
        app.exit(result=SEARCH)

    if have_any_messages:
        @kb.add(Keys.F5)
        def _(event):
            app = get_app()
            app.exit(result=REGENERATE)
        @kb.add(Keys.F6)
        def _(event):
            app = get_app()
            app.exit(result=EDIT)

    @kb.add(Keys.F7)
    def _(event):
        result.allow_execution = not result.allow_execution
        if result.allow_execution:
            print("Command execution enabled.")
        else:
            print("Command execution disabled.")
        app = get_app()
        app.exit(result=TOGGLE_SETTING)

    @kb.add(Keys.F8)
    def _(event):
        app = get_app()
        app.exit(result=SEARCH_THIS_CHAT)

    @kb.add(Keys.F10)
    def _(event):
        app = get_app()
        app.exit(result=SETTINGS)

    def read_chat_search_query():
        search_session = PromptSession()
        return search_session.prompt(HTML("<b>Search this chat: </b>"))

    def read_search_query():
        search_session = PromptSession()
        return search_session.prompt(HTML("<b>Search: </b>"))

    def pick_chat():
        print_formatted_text(HTML(f'<b>Select a chat:</b>'))
        show_list = True
        base = 0
        PAGE_SIZE = Application().output.get_size().rows - 2
        while True:
            i = 0
            for chat in chats[base:]:
                i += 1
                if i == PAGE_SIZE + 1:
                    break
                if show_list:
                    time = formattedTime(chat.last_update)
                    n = chat.num_messages - 1
                    s = "s" if n > 1 else ""
                    print_formatted_text(HTML(f'{i}: <b>{html.escape(chat.name)}</b> ({time}, {n} message{s})'))
            show_list = True

            kb = KeyBindings()
            user_pressed_esc = [False]
            @kb.add(Keys.ControlD)
            def _(event):
                user_pressed_esc[0] = True
                event.app.exit()

            # Prompt the user for their selection
            picker_session = PromptSession(key_bindings=kb)
            selection = picker_session.prompt(HTML("<b>Enter chat number, press 'return' for more, + to create a new chat, or '^D' to cancel: </b>"))
            if user_pressed_esc[0]:
                return

            # Handle the user's selection
            if selection.isdigit():
                index = int(selection) + base - 1
                result.chat_identifier = chats[index].chat_identifier
                print_formatted_text(HTML(f'<b>Switched to: {html.escape(chats[index].name)}</b>'))
                return
            elif selection == '':
                if len(chats) > base + PAGE_SIZE:
                    base += PAGE_SIZE
                else:
                    print_formatted_text(HTML('<b>No more chats.</b>'))
                    show_list = False
            elif selection == '+':
                result.chat_identifier = -1
                return
            else:
                print_formatted_text(HTML('<b>Selection canceled or invalid choice.</b>'))


    # Create a frame around the default buffer
    frame = Frame(body=Window(content=BufferControl(buffer=session.default_buffer)))

    # Set the layout of the session to use the frame
    session.layout = Layout(container=frame)

    # Prompt the user for input
    draw_horizontal_line()
    try:
        while True:
            def bottom_toolbar():
                text = f'[{html.escape(current_chat_name)}] <b>M-⏎</b>: Send  <b>F2</b>: Switch chat  <b>F3</b>: New chat  <b>F4</b>: Search'
                if have_any_messages:
                    text += "  <b>F5</b>: Regenerate"
                    text += "  <b>F6</b>: Edit Last"
                if result.allow_execution:
                    text += "  <b>F7</b>: Disable Execution"
                else:
                    text += "  <b>F7</b>: Enable Execution"
                if have_any_messages:
                    text += "  <b>F8</b>: Search Current Chat"
                text += "  <b>F10</b>: Settings"
                text += "  "
                total_used = used + tokens_typed[0]
                if total_used >= max_tokens:
                  text += "<ansired>"
                text += f'{total_used}/{max_tokens} tokens'
                if total_used >= max_tokens:
                  text += "</ansired>"
                return HTML(text)
            value = session.prompt(
                    "",
                key_bindings=kb,
                multiline=True,
                bottom_toolbar=bottom_toolbar,
                default=placeholder
            )
            if value == SHOW_CHAT_LIST:
                pick_chat()
                if result.chat_identifier is not None:
                    return result
            elif value == NEW_CHAT:
                result.chat_identifier = -1
                return result
            elif value == REGENERATE:
                result.regenerate = True
                return result
            elif value == TOGGLE_SETTING:
              continue
            elif value == EDIT:
                result.edit = True
                return result
            elif value == SEARCH_THIS_CHAT:
                try:
                    result.message_query = read_chat_search_query()
                    if result.message_query is not None:
                        return result
                except EOFError:
                    pass
            elif value == SEARCH:
                try:
                    result.query = read_search_query()
                    if result.query is not None:
                        return result
                except EOFError:
                    pass
            elif value == SETTINGS:
                result.settings = True
                return result
            else:
                result.text = value
                return result
    except Exception as e:
        print(e)
        return None

src/input_reader.py
gnachman/gptline
#!/usr/bin/env python3
from dataclasses import dataclass
from src.background_task import BackgroundTask
from src.chat import create_chat, create_chat_with_spinner, suggest_name, invoke
from src.db import ChatDB
from src.formatting import print_message
from src.formatting import setMark
from src.highlight import SyntaxHighlighter
from src.input_reader import Chat, read_input, usage
from prompt_toolkit import print_formatted_text
from prompt_toolkit.formatted_text import HTML
from src.search import display_chat_search_results, display_search_results
from src.spin import Spinner
from src.ui_utils import draw_horizontal_line, draw_light_horizontal_line
import html
import json
import openai
import os
import subprocess
import sys
import tiktoken
import traceback
from typing import Optional, Any, Callable

openai.api_key = os.environ.get("OPENAI_KEY")
if not openai.api_key:
    print("Set the environment variable OPENAI_KEY to your api secret key")
    exit(1)

@dataclass
class Setting:
    name: str
    key: str
    # Take a decoded value and save it into App
    load: Callable
    default: Any
    # Returns a decoded value
    value: Callable
    # Take a string from the user and return a value that can be json encoded for storage. Throw if the value is bad.
    decode: Callable

def str_to_bool(s):
    if s.lower() == "y" or s.lower() == "t" or s.lower() == "true":
      return True
    if s.lower() == "n" or s.lower() == "f" or s.lower() == "false":
      return False
    raise Exception("Boolean values must be 'true' or 'false'")

class App:
    def __init__(self):
        self.tasks = []
        self.chat_db = ChatDB()
        self.current_chat_id = None
        self.messages = []
        self.placeholder = ""
        self.allow_execution = False
        self.temperature = 0
        self.chats = []
        self.chat = None
        self.content = None
        self.max_tokens = 8000
        self.load_settings()
        self.print_settings()
        self.index_all_chats()

    def new_chat(self):
        self.current_chat_id = self.chat_db.create_chat("New chat")
        self.messages.clear()
        self.messages.append({
            "role": "system",
            "content": "You assist a user in a terminal emulator."})
        self.chat_db.add_message(self.current_chat_id, self.messages[0]["role"],
                self.messages[0]["content"], None, None)
        return self.current_chat_id

    def run_forever(self):
        while True:
            self.iterate()

    def delete_until_user_message(self):
         while self.messages[-1]["role"] == "assistant" or self.messages[-1]["role"] == "function" or "function_call" in self.messages[-1]:
             self.chat_db.delete_message(self.messages[-1]["id"])
             self.messages = self.messages[:-1]

    def iterate(self):
        self.temperature = 0
        self.chats = list(map(lambda x: Chat(*x), self.chat_db.list_chats()))
        setMark()
        current_chat_name = self.get_chat_name()
        user_input = self.read(current_chat_name)
        self.placeholder = ""
        if not user_input:
            # Empty query means to quit.
            return False
        if user_input.settings:
            self.do_settings()
            return
        if user_input.regenerate:
            # Delete the last response and preceding prompt, and act as though
            # the user retyped it.
            # prompt.
            user_input.text = self.regenerate()
        elif user_input.edit:
            # Remove the last message, set placeholder, and ask for input again.
            self.edit()
            return True
        message = user_input.text
        if user_input.query:
            new_i = self.search(user_input.query)
            if new_i is not None:
                user_input.chat_identifier = new_i
            else:
                return True
        if user_input.message_query:
            self.search_messages(user_input.message_query)
        if user_input.chat_identifier:
            if user_input.chat_identifier < 0:
                print_formatted_text(HTML(f'<b>Starting a new chat</b>'))
                self.current_chat_id = self.new_chat()
            else:
                self.switch(user_input.chat_identifier)
            if not message:
                return True
        elif self.current_chat_id is None:
            self.current_chat_id = self.new_chat()
        if not message:
            return False
        if len(self.messages) == 1:
            self.assign_name(message)

        self.send_message(message)
        return True

    def send_message(self, message):
        draw_light_horizontal_line()
        setMark()

        self.messages.append(
            {"role": "user", "content": message},
        )
        message_id = self.chat_db.add_message(self.current_chat_id, "user", message, None, None)
        self.update_chat_fulltext()
        self.messages[-1]["id"] = message_id
        sanitized = [
                {k: v for k, v in message.items() if k != 'id'} 
                for message in self.messages]
        if self.auto_truncate:
            truncated = 0
            while len(sanitized) > 1 and self.usage(sanitized) + 500 > self.max_tokens:
                del sanitized[0]
                truncated += 1
            if truncated:
                s = "s" if truncated != 1 else ""
                print_formatted_text(HTML(f'<em>Warning: token limit exceeded. {truncated} message{s} dropped.</em>'))

        if self.allow_execution:
            execute_command.sloppy = False
            create_file.sloppy = False
            execute_python.sloppy = True
            functions = [execute_command, create_file, execute_python]
        else:
            functions = []
        try:
            self.chat = create_chat_with_spinner(sanitized, self.temperature, functions, self.model)
        except Exception as e:
            print(f"Failed to create chat: {e}")
            return
        self.content = ""
        sh = SyntaxHighlighter()
        fspinner = None
        try:
            # There can be more than one chat when there's a function call.
            while self.read_response(functions, sh):
                pass
            sh.eof()
            print("")
            print("")
        except KeyboardInterrupt:
            sh.eof()
            self.chat.close()
            print("")
        self.commit_ordinary()

        # Check if any tasks are done.
        self.check_tasks()

    def set_model(self, model):
        self.model = model

    def validate_model(self, model):
        try:
            tiktoken.encoding_for_model(model)
        except KeyError:
            raise Exception("Unsupported model")
        return model

    def settings(self):
        return [
                Setting(
                    "Model",
                    "model",
                    lambda s: self.set_model(s),
                    "gpt-3.5-turbo",
                    lambda: self.model,
                    lambda s: self.validate_model(s)),
                Setting(
                    "Auto-truncate",
                    "auto-truncate",
                    lambda s: self.set_auto_truncate(s),
                    True,
                    lambda: self.auto_truncate,
                    lambda s: str_to_bool(s))]

    def set_auto_truncate(self, value):
        self.auto_truncate = value

    def load_settings(self):
        for setting in self.settings():
            setting.load(self.chat_db.get_setting(setting.key, setting.default))

    def print_settings(self): 
        for setting in self.settings():
            print(f'{setting.name}: {setting.value()}')

    def do_settings(self):
        self.load_settings()
        i = 1
        settings = self.settings()
        for setting in settings:
            print(f'({i}) {setting.name}: {setting.value()}')
            i += 1
        num = input("Enter the number of the setting to enter or press return to cancel: ")
        try:
            num = int(num) - 1
            if num < 0 or num >= len(settings):
                print("Invalid number")
                return
            setting = settings[num]
            value = input(f"Enter new value for {setting.name}: ")
            if not value:
                print("Cancel")
                return
            try:
                self.chat_db.set_setting(setting.key, setting.decode(value))
            except Exception as e:
                print(f"Bad value: {e}")
                return
            self.load_settings()
        except ValueError:
            print("Cancel")
            return

    def check_tasks(self):
        for task in self.tasks:
            if task.done():
                if not task.exception:
                    (chat_id, name) = task.result()
                    self.chat_db.set_chat_name(chat_id, name)
            self.tasks = [task for task in self.tasks if not task.done()]

    def get_chat_name(self):
        if self.current_chat_id is None:
            return "New chat"
        else:
            return self.chat_db.get_chat_name(self.current_chat_id)

    def read(self, current_chat_name):
        try:
            user_input = read_input(
                    self.chats,
                    current_chat_name,
                    len(self.messages) > 1,
                    self.placeholder,
                    self.allow_execution,
                    self.usage(self.messages),
                    self.max_tokens,
                    self.model)
            self.allow_execution = user_input.allow_execution
            return user_input 
        except KeyboardInterrupt:
            print("^C")
            exit(0)

    def message_content(self, message):
        c = message["content"]
        if c is not None:
            return c
        f = message["function_call"]
        if f is not None:
            return json.dumps(f)
        return 0

    def usage(self, messages):
        text = ""
        contents = [self.message_content(message) for message in messages]
        text = "\n".join(contents)
        return usage(self.model, text)

    def regenerate(self):
        print("Regenerating response...")
        self.delete_until_user_message()
        text = self.messages[-1]["content"]
        self.temperature = 0.5
        self.chat_db.delete_message(self.messages[-1]["id"])
        self.messages = self.messages[:-1]
        return text

    def edit(self):
        print("Editing previous message...")
        self.delete_until_user_message()
        self.placeholder = self.messages[-1]["content"]
        self.chat_db.delete_message(self.messages[-1]["id"])
        self.messages = self.messages[:-1]

    def search_messages(self, query):
        """Search messages in the current chat"""
        results = self.chat_db.search_messages_in_chat(query, self.current_chat_id)
        i = 1
        for message_id in results:
            draw_light_horizontal_line()
            print_formatted_text(HTML(f'<b>Message {i} of {len(results)}</b>'))
            i += 1
            role, content, timestamp, mid, deleted = self.chat_db.get_message_by_id(message_id)
            print_message(timestamp, role, content, deleted)
            print("")

    def search(self, query):
        """Returns id of chat to switch to or else None.
        A negative id means to create a new chat."""
        i = None
        cursor = None
        PAGE_SIZE = 10
        while True:
            search_results, cursor = self.chat_db.search_chats(query, cursor, PAGE_SIZE)
            if not len(search_results):
                print("No results")
                return None
            i = display_chat_search_results(query, search_results, self.chat_db)
            if i is None:
                cursor += PAGE_SIZE
                continue
            if i < 0:
                return -1
            return i


    def search_by_message(self, query):
        i = None
        cursor = None
        PAGE_SIZE = 10
        while True:
            search_results, cursor = self.chat_db.search_messages(query, cursor, PAGE_SIZE)
            if not len(search_results):
                print("No results")
                return None
            i = display_search_results(query, search_results, self.chat_db)
            if i is None:
                cursor += PAGE_SIZE
                continue
            if i < 0:
                return -1
            return i

    def switch(self, chat_id, show=True):
        """Replace self.messages with contents of chat_id and also print the
        messages to the screen. As a side-effect, set self.current_chat_id."""
        self.current_chat_id = chat_id
        n = self.chat_db.num_messages(self.current_chat_id)
        self.messages = []
        for i in range(1, n):
            (role, self.content, time, message_id, deleted, fname, fargs) = self.chat_db.get_message_by_index(self.current_chat_id, i)
            if self.content is not None:
                m = {
                    "id": message_id,
                    "role": role,
                    "content": self.content
                    }
                if role == "function":
                    m["name"] = fname
                else:
                    if show:
                        if role == "user":
                            draw_horizontal_line()
                        else:
                            draw_light_horizontal_line()
                        print_message(time, role, self.content.rstrip(), deleted)
                        print("")
                self.messages.append(m)
            elif fname and fargs:
                self.messages.append({
                    "id": message_id,
                    "role": "assistant",
                    "content": None,
                    "function_call": {
                        "name": fname,
                        "arguments": fargs}})

    def assign_name(self, message):
        # The chat needs a name
        self.tasks.append(BackgroundTask(lambda: suggest_name(self.current_chat_id, message)))

    def handle_function_call(self, functions, fspinner, call_name, call_args):
        try:
            if fspinner:
                fspinner.stop()
            return (None, invoke(functions, call_name, call_args), None)
        except Exception as e:
            return (None, None, str(e))

    def stop_unexpectedly(self, finish_reason):
        print("")
        print(f"Stopping because {finish_reason}")

    def read_response(self, functions, sh):
        """Read an entire response and handle it. Return True to call this again."""
        fspinner = None
        try:
            call_name = None
            call_args = None
            function_output = None
            error_output = None
            for resp in self.chat:
                if resp.choices[0].finish_reason:
                    finish_reason = resp.choices[0].finish_reason
                    if finish_reason == "function_call" and call_name and call_args:
                        fspinner, function_output, error_output = self.handle_function_call(
                                functions,
                                fspinner,
                                call_name,
                                call_args)
                    elif finish_reason != "stop":
                        self.stop_unexpectedly(finish_reason)
                    break
                elif "content" in resp.choices[0].delta and resp.choices[0].delta.content:
                    self.content = self.handle_content(resp, self.content, sh)
                elif "function_call" in resp.choices[0].delta and resp.choices[0].delta.function_call and self.allow_execution:
                    call_name, call_args, fspinner = self.accrue_function_call(
                            resp, call_name, call_args, fspinner)


            return self.commit_special(
                    call_name,
                    call_args,
                    function_output,
                    error_output,
                    functions)
        finally:
            if fspinner:
                fspinner.stop()
                fspinner = None

    def handle_content(self, resp, content, sh):
        chunk = resp.choices[0].delta.content
        content += chunk
        sh.put(chunk)
        return content

    def accrue_function_call(self, resp, call_name, call_args, fspinner):
        if "name" in resp.choices[0].delta.function_call:
            if call_name:
                call_name += resp.choices[0].delta.function_call.name
            else:
                fspinner = Spinner()
                call_name = resp.choices[0].delta.function_call.name
        if "arguments" in resp.choices[0].delta.function_call:
            if call_args:
                call_args += resp.choices[0].delta.function_call.arguments
            else:
                call_args = resp.choices[0].delta.function_call.arguments
        return (call_name, call_args, fspinner)

    def index_all_chats(self):
        if self.chat_db.get_kv("chat_fts_migration") == "done":
            return
        print("Re-indexing all chats for better full text search. This could take a second.")
        saved = self.current_chat_id
        for info in self.chat_db.list_chats():
            cid = info[0]
            self.switch(cid, False)
            self.update_chat_fulltext()
        self.current_chat_id = saved
        print("Done")
        self.chat_db.set_kv("chat_fts_migration", "done")

    def update_chat_fulltext(self):
        content = self.get_chat_name() + "\n" + "\n".join([self.plaintext_message(m) for m in self.messages])
        self.chat_db.set_chat_fulltext(self.current_chat_id, content)

    def plaintext_message(self, m):
        if m["role"] == "assistant":
            role = "Assistant"
        elif m["role"] == "user":
            role = "User"
        else:
            return ""
        if "content" not in m:
            return ""
        content = m["content"]
        return f'{role}: {content}'

    def commit_special(self, call_name, call_args, function_output,
            error_output, functions):
        if call_name and call_args:
            self.commit_function_call_request(call_name, call_args);
        if function_output:
            self.commit_function_output(functions, call_name, function_output)
            return True
        if error_output:
            return self.commit_error(error_output, functions)
        return False

    def commit_ordinary(self):
        if self.content is not None:
            self.messages.append({"role": "assistant", "content": self.content})
        message_id = self.chat_db.add_message(
                self.current_chat_id,
                self.messages[-1]["role"],
                self.messages[-1]["content"],
                None,
                None)
        self.messages[-1]["id"] = message_id
        self.update_chat_fulltext()

    def commit_function_call_request(self, call_name, call_args):
        # Record that a function call was requested
        self.messages.append({
            "role": "assistant",
            "content": None,
            "function_call": {
                "name": call_name,
                "arguments": call_args}})
        message_id = self.chat_db.add_message(self.current_chat_id,
                self.messages[-1]["role"], self.messages[-1]["content"], call_name, call_args)
        self.messages[-1]["id"] = message_id
        self.update_chat_fulltext()

    def commit_function_output(self, functions, call_name, function_output):
        # Record the output of the function call
        self.messages.append({
            "role": "function",
            "name": call_name,
            "content": function_output})
        message_id = self.chat_db.add_message(
                self.current_chat_id,
                self.messages[-1]["role"],
                self.messages[-1]["content"],
                call_name,
                None)
        self.messages[-1]["id"] = message_id
        sanitized = [
                {k: v for k, v in message.items() if k != 'id'} 
                for message in self.messages]
        try:
            self.chat = create_chat(sanitized, self.temperature, functions, self.model)
        except Exception as e:
            print(f"Failed to create chat: {e}")

    def commit_error(self, error_output, functions):
        # Something went wrong
        print(f'Error: {error_output}')
        auto_retry = False
        self.messages.append({
            "role": "system",
            "content": error_output})
        message_id = self.chat_db.add_message(
                self.current_chat_id,
                self.messages[-1]["role"],
                self.messages[-1]["content"],
                None,
                None)
        self.messages[-1]["id"] = message_id
        if auto_retry:
            sanitized = [
                    {k: v for k, v in message.items() if k != 'id'} 
                    for message in self.messages]
            print(sanitized)
            try:
                self.chat = create_chat(sanitized, self.temperature, functions, self.model)
            except Exception as e:
                print(f"Failed to create chat: {e}")
                return False
            return True
        else:
            print(f"An error was encountered while executing a function: {error_output}")
            self.content = None
            return False

def main():
    print("Welcome to gptline! Enter a question and press option-Enter to send it.")
    app = App()
    app.run_forever()

def execute_command(command_line: str, input_string: str):
    """
    Executes a unix command at the shell. Note that it does not run in a TTY so interactive commands will not work.

    Args:
        command_line: Unix command to execute. For example, "ls -l". This string will be passed to system().
        input_string: A string containing input to send to stdin.
    """
    print("")
    print("ChatGPT wants to run the following command:")
    print(command_line)
    print("")
    ok = input(f"OK to run (y/n)? ")
    if ok != "y":
        print("Not running it")
        return "Error: user denied permission to execute function call"
    print("Executing...")

    try:
        process = subprocess.Popen(command_line, shell=True, cwd="/tmp", stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True)
        output, _ = process.communicate(input=input_string)
        print("Program output:")
        print(output)
        print("")
        return output.strip()
    except Exception as e:
        print(f'Exception while executing provided code: {e}')
        return str(e)

def create_file(name: str, content: str):
    """
    Creates a file with the given name and contents.

    Args:
        name: A file name. This will be relative to a private temporary directory.
        content: A string to write to the file.
    """
    print(f"ChatGPT wants to create a file called {name} with the following contents:")
    sh = SyntaxHighlighter()
    sh.put("```")
    sh.put(content)
    sh.put("")
    sh.put("```")
    sh.eof()
    ok = input(f"OK to create (y/n)? ")
    if ok != "y":
        print("Not running it")
        return "Error: user denied permission to create file"
    file_path = os.path.join('/tmp', name)

    try:
        # Open the file in write mode and write the content
        with open(file_path, 'w') as file:
            file.write(content)
        return f"File {name} created successfully."
    except Exception as e:
        return f"Error creating file: {str(e)}"

def execute_python(code: str, input_string: Optional[str]):
    """
    Executes python code. Outputs values from stdout. If you need a computed value make sure to print() it.

    Args:
        code: Python code to execute.
        input_string: A string containing input to send to stdin.
    """
    print("ChatGPT wants to run the following Python program:")
    sh = SyntaxHighlighter()
    sh.put("```python\n")
    sh.put(code)
    if not code.endswith("\n"):
        sh.put("\n")
    sh.put("```")
    sh.eof()
    print("")
    if input_string:
        print("With this supplied as input:")
        print(input_string)
        print("")
    ok = input(f"OK to run (y/n)? ")
    if ok != "y":
        print("Not running it")
        return "Error: user denied permission to execute function call"
    print("Executing...")
    try:
        process = subprocess.Popen(["python3", "-c", code], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        stdout, stderr = process.communicate(input=input_string if input_string else "")
        print("Program output:")
        print(stdout + stderr)
        if stderr:
            return f"Error: {stderr}"
        return stdout
    except Exception as e:
        return f"Error: {str(e)}"

src/main.py
gnachman/gptline
from prompt_toolkit import PromptSession, print_formatted_text
from prompt_toolkit.formatted_text import HTML
from src.formatting import print_snippet
import html

def display_chat_search_results(query, search_results, chat_db):
    session = PromptSession()

    while True:
        for index, (chat_id, snippet) in enumerate(search_results, start=1):
            chat_name = chat_db.get_chat_name(chat_id)
            print_formatted_text(HTML(f'<b>{index}. {html.escape(chat_name)}</b>'))
            escaped_snippet = html.escape(snippet).replace('\ue000', '<b>').replace('\ue001', '</b>').replace('\n', ' ')
            print_formatted_text(HTML("    " + escaped_snippet))

        # Prompt the user to either select a chat or request more results
        try:
            user_input = session.prompt("\nEnter a number to select a chat, press Enter for more results, or Ctrl+D to cancel: ")
        except EOFError:
            print("\nSearch canceled.")
            return -1
        except KeyboardInterrupt:
            print("\nSearch canceled.")
            return -1

        # Return None to request more results
        if user_input == "":
            return None

        # Return chat ID as int if selected
        try:
            selected_index = int(user_input) - 1
            if 0 <= selected_index < len(search_results):
                selected_chat_id = search_results[selected_index][0]
                return selected_chat_id
            else:
                print("Invalid input. Please enter a valid number.")
        except ValueError:
            print("Invalid input. Please enter a valid number.")



def display_search_results(query, messages_by_chat, chat_db):
    session = PromptSession()

    while True:
        for index, (chat_id, message_ids) in enumerate(messages_by_chat, start=1):
            chat_name = chat_db.get_chat_name(chat_id)
            print_formatted_text(HTML(f'<b>{index}. {html.escape(chat_name)}</b>'))
            for (message_id, snippet) in message_ids:
                role, content, timestamp, _id, deleted = chat_db.get_message_by_id(message_id)
                escaped_snippet = html.escape(snippet).replace('\ue000', '<b>').replace('\ue001', '</b>').replace('\n', ' ')
                print_snippet(timestamp, role, escaped_snippet, deleted, "    ")

        # Prompt the user to either select a chat or request more results
        try:
            user_input = session.prompt("\nEnter a number to select a chat, press Enter for more results, or Ctrl+D to cancel: ")
        except EOFError:
            print("\nSearch canceled.")
            return -1
        except KeyboardInterrupt:
            print("\nSearch canceled.")
            return -1

        # Return None to request more results
        if user_input == "":
            return None

        # Return chat ID as int if selected
        try:
            selected_index = int(user_input) - 1
            if 0 <= selected_index < len(messages_by_chat):
                selected_chat_id = messages_by_chat[selected_index][0]
                return selected_chat_id
            else:
                print("Invalid input. Please enter a valid number.")
        except ValueError:
            print("Invalid input. Please enter a valid number.")



src/search.py
gnachman/gptline
import sys
import threading
import time

class Spinner:
    def __init__(self):
        self.i = 0
        self.running = True
        self.thread = threading.Thread(target=self._run)
        self.thread.start()
        self.wrote = False

    def stop(self):
        self.running = False
        self.thread.join()
        if self.wrote:
            sys.stdout.write(chr(8) + " " + chr(8))
            sys.stdout.flush()

    def _run(self):
        bs = ""
        spinner = ["|", "/", "-", "\\"]
        i = 0
        while self.running:
            sys.stdout.write(bs + spinner[i % 4])
            bs = chr(8)
            sys.stdout.flush()
            self.wrote = True
            time.sleep(0.1)
            i += 1

def spin(f):
    s = Spinner()
    exc = None
    try:
        result = f()
    except Exception as e:
        exc = e
    s.stop()
    if exc:
        raise exc
    return result


src/spin.py
gnachman/gptline
import os
from prompt_toolkit.formatted_text import HTML
from prompt_toolkit import PromptSession, print_formatted_text

def draw_light_horizontal_line():
    draw_horizontal_line("<style fg='gray'>{}</style>")

def draw_horizontal_line(style=None):
    terminal_width = os.get_terminal_size().columns
    if style:
        line = HTML(style.format('\u2500' * terminal_width))
    else:
        line = HTML('\u2500' * terminal_width)
    print_formatted_text(line)



src/ui_utils.py
gnachman/gptline