ChatCLI

背景

自从官方开放API开始,OpenAI就被墙了,国内用户访问ChatGPT越来越麻烦了,然而官方的API却出乎意料地,没有封禁国内用户。对ChatGPT重度用户,用API请求ChatGPT可以作为有时登不上官网的后备策略。

市面上已经有大量ChatGPT的GUI了,有WEB端,也有原生APP,但是大多用的是一次性得到response并显示的方式。事实上,官方API中,有一个 stream 参数(默认值为 False ),若为 True 就能实现像 ChatGPT 官方网页版那样,一个一个字吐出来的效果。这样相比一次请求获得结果的方式,具有以下好处:

  • 响应更快,可以逐步得到结果
  • 避免因为内容过长,等待时间过久造成 timeout
  • 可以实现动画

本文教大家用60行python代码实现一个在终端中运行的ChatGPT终端,具有以下功能:

  • 连续对话(可重置)
  • 多行编辑
  • markdown渲染
  • 一键复制响应/回答

实现这些功能,需要用到以下第三方库:

  • requests 用于发起HTTP请求
  • pyperclip 用于将内容复制到剪贴板
  • questionary 用来实现终端中支持多行编辑的UI
  • rich 用来渲染markdown

实现

对OpenAI的请求很简单,格式如下:

requests.post("https://api.openai.com/v1/chat/completions",
    json={
        "model": "gpt-3.5-turbo",
        "messages": context,
        "stream": True
    }, 
    headers={
        "Authorization": f"Bearer sk-***",
        "Content-Type": "application/json"
    },
    stream=True
)

其中,在 headersAuthorization 中填写你自己的 API key ,而 message 中的 messages 字段的格式如下:

[
    { "role": "system", "content": "你是ChatGPT" },
    { "role": "user", "content": "你好" },
    { "role": "assistant", "content": "你好,我是一个大型语言模型....."},
    ...
]

其中,role 标志了该消息是用户发给ChatGPT的,还是ChatGPT回答用户的,抑或是系统发给ChatGPT的……(详见官方文档

"stream": True 是关键的参数,这意味着OpenAI会用 Event stream format 流式传输返回的回答。而 stream=True 是传给 requests.post 的参数,这意味着 requests 会以流的方式读取response

我们看看response的格式:

response = requests.post(...)
for chunk in response.iter_lines():
    print(chunk)

结果是

...

b'data: {"id":"chatcmpl-6rL6EQEKUZiUhtot9uqm7ZuSKw4fC","object":"chat.completion.chunk","created":1678170510,"model":"gpt-3.5-turbo-0301","choices":[{"delta":{"content":"\xe6\x88\x91"},"index":0,"finish_reason":null}]}'
b''
b'data: {"id":"chatcmpl-6rL6EQEKUZiUhtot9uqm7ZuSKw4fC","object":"chat.completion.chunk","created":1678170510,"model":"gpt-3.5-turbo-0301","choices":[{"delta":{"content":"\xe3\x80\x82"},"index":0,"finish_reason":null}]}'
b''
b'data: {"id":"chatcmpl-6rL6EQEKUZiUhtot9uqm7ZuSKw4fC","object":"chat.completion.chunk","created":1678170510,"model":"gpt-3.5-turbo-0301","choices":[{"delta":{},"index":0,"finish_reason":"stop"}]}'
b''
b'data: [DONE]'
b''

打印出了满屏幕的数据,观察发现,数据是用空字符串隔开的,最后以 data: [DONE] 结束。

对每一个有数据的行,可以新增的字符串可以用如下方式获取:

delta = json.loads(chunk.strip()[6:].decode())["choices"][0]["delta"].get("content", "")

delta 连续不断地打印出来,就能获得ChatGPT官网UI挤牙膏的效果了:

print(delta, end="", flush=True)  # 这样去掉行尾的空行,并强制刷新终端

附加功能

实现多轮对话最简单的方法就是维护一个全局的 context 列表,每次用户输入/ChatGPT回答都加进去,这里篇幅所限就不详解了

比如终端获得用户的多行输入,并添加进 context 中:

user_input = questionary.text("enter your prompt:", multiline=True).ask()
# 用户可以进行多行输入,在IDE中按Alt+Enter退出,在Windows Terminal中按Esc后按Enter退出
context.append({"role": "user", "content": user_input})

那么复制到剪贴板也很好写了:

if questionary.select("copy what", ["input", "result"]).ask() == "input":
    pyperclip.copy(context[-2]["content"])
else:
    pyperclip.copy(context[-1]["content"])

questionary 的具体用法见官方文档


因为用ChatGPT写代码是常见的需求,因此代码高亮也很重要。常见的终端高亮办法是采用 rich.syntax 库高亮代码块,我们这里直接rich.markdown 显示markdown文档(ChatGPT返回的格式就是markdown)。

有个问题是,流式显示ChatGPT的回答与语法高亮不可兼得,我们这里采用一个折衷方案,就是先流式打印纯文本,响应完毕后清空终端并输出完整的markdown高亮的回答

from rich.console import Console
from rich.markdown import Markdown

console = Console()
console.print(Markdown("........"))

最后把以上所有部分完成,以及加上错误处理机制(尤其是http请求的部分,因为代理一般不会太稳定,经常请求不到,需要实现简答的重试机制),完整的代码如下:

from rich.markdown import Markdown
from rich.console import Console
import questionary
import pyperclip
import requests
import json

console = Console()
context = [{"role": "system", "content": "语言是中文。若无提示,编程语言是python。请详细回答用户提问的每个问题。"}]
headers = {
    "Authorization": f"Bearer sk-***", "Content-Type": "application/json"
}

def invoke_chat():
    user_input = questionary.text("enter your prompt:", multiline=True).ask()
    context.append({"role": "user", "content": user_input})
    while True:
        try:
            response = requests.post("https://api.openai.com/v1/chat/completions", json={
                "stream": True, "model": "gpt-3.5-turbo", "messages": context
            }, headers=headers, stream=True)
            parts = []
            for chunk in response.iter_lines():
                if chunk.strip() and not chunk.startswith(b"data: [DONE]"):
                    result = json.loads(chunk.strip()[6:].decode())["choices"][0]
                    if result["finish_reason"] is not None:
                        continue

                    delta = result["delta"].get("content", "")
                    parts.append(delta)
                    print(delta, end="", flush=True)

            whole = "".join(parts)
            context.append({"role": "assistant", "content": whole})
            console.clear()
            console.print(Markdown("---\n" + whole))
            break

        except requests.RequestException:
            import time
            time.sleep(0.5)

def main():
    while True:
        choice = questionary.select("what to do next", ["continue", "restart", "copy", "exit"]).ask()
        if choice == "exit":
            break
        if choice == "copy":
            if questionary.select("copy what", ["input", "result"]).ask() == "input":
                pyperclip.copy(context[-2]["content"])
            else:
                pyperclip.copy(context[-1]["content"])
            continue

        if choice == "restart":
            del context[1:]
        invoke_chat()


if __name__ == '__main__':
    main()

要运行以上代码,只需要安装3个依赖:

$ python -m pip install -U requests pyperclip questionary rich

本文最后编辑于2023年3月7日

2
1