diff --git a/chatgpt_pyapi/__init__.py b/chatgpt_pyapi/__init__.py index 566d26f..e69de29 100644 --- a/chatgpt_pyapi/__init__.py +++ b/chatgpt_pyapi/__init__.py @@ -1,91 +0,0 @@ -# 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 = [] diff --git a/chatgpt_pyapi/openai.py b/chatgpt_pyapi/openai.py new file mode 100644 index 0000000..566d26f --- /dev/null +++ b/chatgpt_pyapi/openai.py @@ -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 = [] diff --git a/cli.py b/cli.py new file mode 100755 index 0000000..2c452a1 --- /dev/null +++ b/cli.py @@ -0,0 +1,17 @@ +#!/usr/bin/env python3 + +# Copyright (c) 2023 Julian Müller (ChaoticByte) + +from os import environ +from sys import argv + +from chatgpt_pyapi.openai import ChatGPT, Message, Models + +# Read the API key from a environment variable +API_KEY = environ["OPENAI_API_KEY"] + +if __name__ == "__main__": + # Create ChatGPT API instance + cgpt = ChatGPT(API_KEY, model=Models.GPT_35_TURBO_0301) + user_input = " ".join(argv[1:]) + print(cgpt.chat(Message(user_input)).text) diff --git a/example.py b/example.py index 52cfba6..1e474b5 100755 --- a/example.py +++ b/example.py @@ -4,10 +4,7 @@ from os import environ -from chatgpt_pyapi import ChatGPT -from chatgpt_pyapi import Message -from chatgpt_pyapi import Models -from chatgpt_pyapi import Roles +from chatgpt_pyapi.openai import ChatGPT, Message, Models, Roles # Read the API key from a environment variable API_KEY = environ["OPENAI_API_KEY"]