gpt.py

import sys
import json
import os
import time

import keyboard
from openai import OpenAI

chat_cache = "chat.json"


def pre_process(
    api_key,
    content,
    model,
    temprature,
    context_time,
    chat_number,
    clear_keyword,
    clipboard_keyword,
    roles,
    signature,
    print_input,
    api_base,
):
    # context_time
    if float(context_time) < 0:
        context_time = float("inf")
    if os.path.isfile(chat_cache):
        mod_time = os.path.getmtime(chat_cache)
        current_time = time.time()
        if (current_time - mod_time) / 60 > float(context_time):
            os.remove(chat_cache)

    # chat_number
    chat_number = int(int(chat_number) * 2)

    # clear_keyword
    if clear_keyword in content:
        content = content.replace(clear_keyword, "")
        if os.path.isfile(chat_cache):
            os.remove(chat_cache)

    # clipboard_keyword
    if clipboard_keyword in content:
        import pyperclip

        content = content.replace(clipboard_keyword, pyperclip.paste())

    # roles
    for role in roles.split("\n"):
        index = role.find(":")
        k = role[:index]
        v = role[index + 1 :].strip()
        if k in content:
            content = content.replace(k, "")
            content = v + content

    # signature
    signature = int(signature)

    # print_input
    print_input = int(print_input)

    messages = []
    if os.path.exists(chat_cache):
        with open(chat_cache, "r") as chat_json:
            messages = json.load(chat_json)
    messages += [{"role": "user", "content": content}]

    return (
        api_key,
        model,
        temprature,
        messages,
        signature,
        print_input,
        chat_number,
        api_base,
    )


def ask_gpt(api_key, model, temprature, messages, signature, print_input, api_base):
    client = OpenAI(api_key=api_key, base_url=api_base)
    temperature_fixed = int(temprature)
    if print_input:
        keyboard.write(">>> " + messages[-1]["content"] + "\n")
    chat_completion = client.chat.completions.create(
        messages=messages, model=model, stream=True, temperature=temperature_fixed
    )

    assistant_content = ""
    for chunk in chat_completion:
        if not chunk.choices[0].finish_reason == "stop":
            c = chunk.choices[0].delta.content
            if c:
                keyboard.write(c)
                assistant_content += c
        else:
            if signature == 1:
                from datetime import datetime

                current = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                keyboard.write("\n[Generated by ChatGPT, {}]\n\n".format(current))

    messages += [{"role": "assistant", "content": assistant_content}]
    return messages


def post_process(messages, chat_number):
    if len(messages) > chat_number:
        messages = messages[len(messages) - chat_number : len(messages)]
    with open(chat_cache, "w") as chat_json:
        json.dump(messages, chat_json, indent=4)


if __name__ == "__main__":
    params = sys.argv[1:]
    api_key, model, temprature, messages, signature, print_input, chat_number, api_base = pre_process(
        *params
    )
    if messages[-1].get("content", None):
        messages = ask_gpt(
            api_key, model, temprature, messages, signature, print_input, api_base
        )
        post_process(messages, chat_number)