107 lines
3.7 KiB
Python
107 lines
3.7 KiB
Python
"""
|
|
Implements the OpenAI client classes and functions.
|
|
"""
|
|
import openai
|
|
from typing import Optional, Union
|
|
from ..tags import Tag
|
|
from ..message import Message, Answer
|
|
from ..chat import Chat
|
|
from ..ai import AI, AIResponse, Tokens
|
|
from ..configuration import OpenAIConfig
|
|
|
|
ChatType = list[dict[str, str]]
|
|
|
|
|
|
class OpenAI(AI):
|
|
"""
|
|
The OpenAI AI client.
|
|
"""
|
|
|
|
def __init__(self, config: OpenAIConfig) -> None:
|
|
self.ID = config.ID
|
|
self.name = config.name
|
|
self.config = config
|
|
openai.api_key = config.api_key
|
|
|
|
def request(self,
|
|
question: Message,
|
|
chat: Chat,
|
|
num_answers: int = 1,
|
|
otags: Optional[set[Tag]] = None) -> AIResponse:
|
|
"""
|
|
Make an AI request, asking the given question with the given
|
|
chat history. The nr. of requested answers corresponds to the
|
|
nr. of messages in the 'AIResponse'.
|
|
"""
|
|
oai_chat = self.openai_chat(chat, self.config.system, question)
|
|
response = openai.ChatCompletion.create(
|
|
model=self.config.model,
|
|
messages=oai_chat,
|
|
temperature=self.config.temperature,
|
|
max_tokens=self.config.max_tokens,
|
|
top_p=self.config.top_p,
|
|
n=num_answers,
|
|
frequency_penalty=self.config.frequency_penalty,
|
|
presence_penalty=self.config.presence_penalty)
|
|
question.answer = Answer(response['choices'][0]['message']['content'])
|
|
question.tags = otags
|
|
question.ai = self.ID
|
|
question.model = self.config.model
|
|
answers: list[Message] = [question]
|
|
for choice in response['choices'][1:]: # type: ignore
|
|
answers.append(Message(question=question.question,
|
|
answer=Answer(choice['message']['content']),
|
|
tags=otags,
|
|
ai=self.ID,
|
|
model=self.config.model))
|
|
return AIResponse(answers, Tokens(response['usage']['prompt_tokens'],
|
|
response['usage']['completion_tokens'],
|
|
response['usage']['total_tokens']))
|
|
|
|
def models(self) -> list[str]:
|
|
"""
|
|
Return all models supported by this AI.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def print_models(self) -> None:
|
|
"""
|
|
Print all models supported by the current AI.
|
|
"""
|
|
not_ready = []
|
|
for engine in sorted(openai.Engine.list()['data'], key=lambda x: x['id']):
|
|
if engine['ready']:
|
|
print(engine['id'])
|
|
else:
|
|
not_ready.append(engine['id'])
|
|
if len(not_ready) > 0:
|
|
print('\nNot ready: ' + ', '.join(not_ready))
|
|
|
|
def openai_chat(self, chat: Chat, system: str,
|
|
question: Optional[Message] = None) -> ChatType:
|
|
"""
|
|
Create a chat history with system message in OpenAI format.
|
|
Optionally append a new question.
|
|
"""
|
|
oai_chat: ChatType = []
|
|
|
|
def append(role: str, content: str) -> None:
|
|
oai_chat.append({'role': role, 'content': content.replace("''", "'")})
|
|
|
|
append('system', system)
|
|
for message in chat.messages:
|
|
if message.answer:
|
|
append('user', message.question)
|
|
append('assistant', message.answer)
|
|
if question:
|
|
append('user', question.question)
|
|
return oai_chat
|
|
|
|
def tokens(self, data: Union[Message, Chat]) -> int:
|
|
raise NotImplementedError
|
|
|
|
def print(self) -> None:
|
|
print(f"MODEL: {self.config.model}")
|
|
print("=== SYSTEM ===")
|
|
print(self.config.system)
|