A simple Python wrapper for the ChatGPT API#

OpenAI released an API for ChatGPT yesterday. It’s 1/10th of the price of the text-davinci-003 model!

Their official openai Python package has been upgraded to add support for it (in this commit).

Unlike the regular GPT-3 APIs, this one takes an array of messages that looks like this:

1
[
2
  {
3
    "role": "system",
4
    "content": "You are a 1700s pirate with an exagerated UK westcountry accent"
5
  },
6
  {
7
    "role": "user",
8
    "content": "Introduce yourself"
9
  }
10
]

The “system” role is a way to give the bot additional instructions about how to behave.

This is enough to generate a response - but it’s your responsibility to then pass back that response with each subsequent API call in order to maintain context of the prior conversation:

1
[
2
  {
3
    "role": "system",
4
    "content": "You are a 1700s pirate with an exagerated UK westcountry accent"
5
  },
6
  {
7
    "role": "user",
8
    "content": "Introduce yourself"
9
  },
10
  {
11
    "role": "assistant",
12
    "content": "Arrr, me hearties! Shiver me timbers, my name is Cap'n Percy of the Jolly Roger crew! I hail from the land of rolling hills and endless fields, the ol' Westcountry of England. I may not be the finest bloke ye'll come across, but I can sail me ship through any storm and loot any vessel that comes in me way. So hoist the main sails and let's begin the adventure, me hearty!"
13
  },
14
  {
15
    "role": "user",
16
    "content": "What's your favourite bougie cocktail?"
17
  }
18
]

(It replied “Bougie cocktails? Ye’re speakin’ fancy words, matey! I prefer me good ol’ rum, straight up or with some water, but if ye really want me to choose a bougie drink, I’d say a Dark ‘n Stormy - a mix of dark rum, ginger beer, and lime juice. It’s a drink fit for a seafarin’ pirate like meself, arrr!”)

A Python wrapper#

I built this simple wrapper class to automate the process of recording the conversation and replaying it back to the API each time:

1
# This code is Apache 2 licensed:
2
# https://www.apache.org/licenses/LICENSE-2.0
3
import openai
4

5
class ChatBot:
6
    def __init__(self, system=""):
7
        self.system = system
8
        self.messages = []
9
        if self.system:
10
            self.messages.append({"role": "system", "content": system})
11

12
    def __call__(self, message):
13
        self.messages.append({"role": "user", "content": message})
14
        result = self.execute()
15
        self.messages.append({"role": "assistant", "content": result})
16
        return result
17

18
    def execute(self):
19
        completion = openai.ChatCompletion.create(model="gpt-3.5-turbo", messages=self.messages)
20
        # Uncomment this to print out token usage each time, e.g.
21
        # {"completion_tokens": 86, "prompt_tokens": 26, "total_tokens": 112}
22
        # print(completion.usage)
23
        return completion.choices[0].message.content

Now I can do the following:

1
>>> simon = ChatBot("You are a chatbot imitating Simon Willison. Pretend to be Simon.")
2
>>> simon("Tell me about yourself")

Without that “Pretend to be X” note it will sometimes mention that it’s a chatbot that is only imitating the person.