Added the initial version of the library source code
This commit is contained in:
parent
645101cc46
commit
b5a5920722
3 changed files with 103 additions and 0 deletions
|
@ -1,2 +1,3 @@
|
|||
# ChatGPT-PyAPI
|
||||
|
||||
A minimalistic Python API for OpenAI's ChatGPT
|
||||
|
|
81
chatgpt_pyapi/__init__.py
Normal file
81
chatgpt_pyapi/__init__.py
Normal file
|
@ -0,0 +1,81 @@
|
|||
|
||||
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):
|
||||
self.text = text
|
||||
self.role = role
|
||||
|
||||
@classmethod
|
||||
def from_api(cls, message_dict:str):
|
||||
'''Create a Message object from API format'''
|
||||
return cls(
|
||||
message_dict["content"],
|
||||
message_dict["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):
|
||||
# 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:
|
||||
# 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)
|
||||
# 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())
|
||||
api_output_answer = api_output["choices"][0]["message"]
|
||||
# remove leading and trailing newlines
|
||||
api_output_answer["content"] = api_output_answer["content"].strip("\n")
|
||||
# convert to Message object
|
||||
response_message = Message.from_api(api_output_answer)
|
||||
self.message_history.append(response_message)
|
||||
return response_message
|
||||
|
||||
def clear_message_history(self):
|
||||
self.message_history = []
|
21
example.py
Executable file
21
example.py
Executable file
|
@ -0,0 +1,21 @@
|
|||
#!/usr/bin/env python3
|
||||
|
||||
from os import environ
|
||||
|
||||
from chatgpt_pyapi import ChatGPT
|
||||
from chatgpt_pyapi import Models
|
||||
from chatgpt_pyapi import Message
|
||||
|
||||
|
||||
API_KEY = environ["OPENAI_API_KEY"]
|
||||
|
||||
if __name__ == "__main__":
|
||||
cgpt = ChatGPT(API_KEY, model=Models.GPT_35_TURBO_0301)
|
||||
user_in = "Who are you?"
|
||||
print(f"USER: {user_in}")
|
||||
print(cgpt.chat(Message(user_in)).text)
|
||||
user_in = "Could you please elaborate?"
|
||||
print(f"\nUSER: {user_in}")
|
||||
print(cgpt.chat(Message(user_in)).text)
|
||||
print("\nMessage History:")
|
||||
[print(m.to_api()) for m in cgpt.message_history]
|
Reference in a new issue