Added cli and moved official API to chatgpt_pyapi.openai
This commit is contained in:
parent
4fd0969457
commit
308bf0730f
4 changed files with 109 additions and 95 deletions
91
chatgpt_pyapi/openai.py
Normal file
91
chatgpt_pyapi/openai.py
Normal file
|
@ -0,0 +1,91 @@
|
|||
# Copyright (c) 2023 Julian Müller (ChaoticByte)
|
||||
|
||||
from json import dumps, loads
|
||||
from urllib import request as http_request
|
||||
|
||||
|
||||
class Models:
|
||||
'''This class holds available models'''
|
||||
GPT_35_TURBO = "gpt-3.5-turbo"
|
||||
GPT_35_TURBO_0301 = "gpt-3.5-turbo-0301"
|
||||
|
||||
|
||||
class Roles:
|
||||
'''This class holds available roles to be used in Messages'''
|
||||
ASSISTANT = "assistant"
|
||||
SYSTEM = "system"
|
||||
USER = "user"
|
||||
|
||||
|
||||
class Message:
|
||||
'''Message type. Supports roles.'''
|
||||
|
||||
def __init__(self, text: str, role: str = Roles.USER):
|
||||
assert type(text) == str
|
||||
assert type(role) == str
|
||||
self.text = text
|
||||
self.role = role
|
||||
|
||||
@classmethod
|
||||
def from_api(cls, api_msg: dict):
|
||||
'''Create a Message object from API format'''
|
||||
assert type(api_msg) == dict
|
||||
msg = api_msg["choices"][0]["message"]
|
||||
msg["content"] = msg["content"].strip("\n")
|
||||
return cls(
|
||||
msg["content"],
|
||||
msg["role"])
|
||||
|
||||
def to_api(self):
|
||||
'''Convert to API format'''
|
||||
return {"role": self.role, "content": self.text}
|
||||
|
||||
|
||||
class ChatGPT:
|
||||
'''ChatGPT API'''
|
||||
|
||||
API_ENDPOINT = "https://api.openai.com/v1/chat/completions"
|
||||
|
||||
def __init__(self, api_key: str, model: str = Models.GPT_35_TURBO):
|
||||
assert type(api_key) == str
|
||||
assert type(model) == str
|
||||
# Create string used in header
|
||||
self.auth = f"Bearer {api_key}"
|
||||
# This list will contain all prior messages
|
||||
self._message_history = []
|
||||
self.model = model
|
||||
|
||||
def chat(self, message: Message) -> Message:
|
||||
'''Add a message to the message history & send it to ChatGPT. Returns the answer as a Message instance.'''
|
||||
self.add_to_chat(message)
|
||||
# Create api_input from message_history & encode it
|
||||
api_input = [m.to_api() for m in self._message_history]
|
||||
api_input_encoded = dumps(
|
||||
{"model": self.model, "messages": api_input},
|
||||
separators=(",", ":")).encode()
|
||||
# Create a Request object with the right url, data, headers and http method
|
||||
request = http_request.Request(
|
||||
self.API_ENDPOINT,
|
||||
data=api_input_encoded,
|
||||
headers={
|
||||
"Authorization": self.auth,
|
||||
"Content-Type": "application/json"
|
||||
},
|
||||
method="POST")
|
||||
# Send the request with r as the response
|
||||
with http_request.urlopen(request) as r:
|
||||
# Read response and parse json
|
||||
api_output = loads(r.read())
|
||||
# Convert to Message object
|
||||
response_message = Message.from_api(api_output)
|
||||
self._message_history.append(response_message)
|
||||
return response_message
|
||||
|
||||
def add_to_chat(self, message: Message) -> Message:
|
||||
'''Add a message to the message history without sending it to ChatGPT'''
|
||||
# Check if the message parameter is the correct type
|
||||
assert type(message) == Message, "message must be an instance of Message"
|
||||
self._message_history.append(message)
|
||||
|
||||
def clear_message_history(self):
|
||||
self._message_history = []
|
Reference in a new issue