|
|
import copy |
|
|
from enum import Enum |
|
|
from os import PathLike |
|
|
|
|
|
from openai import OpenAI |
|
|
|
|
|
from commons.loggerfactory import LoggerFactory |
|
|
from commons.utils import getdefault |
|
|
from scripts.io import InputOutput |
|
|
|
|
|
DEFAULT_WELCOME_MESSAGE = "How can I assist you ..." |
|
|
DEFAULT_USER_PROMPT = "User: " |
|
|
DEFAULT_AGENT_PROMPT = "Agent: " |
|
|
DEFAULT_EXIT_MESSAGE = "Have a nice day!" |
|
|
|
|
|
|
|
|
class Role(Enum): |
|
|
SYSTEM = "system" |
|
|
USER = "user" |
|
|
ASSISTANT = "assistant" |
|
|
|
|
|
|
|
|
class ChatHistory: |
|
|
def __init__(self, context_span: int, initial_context: list[dict[str, str]]): |
|
|
self.__context_span = context_span |
|
|
self.__context = list() |
|
|
for context in initial_context: |
|
|
self.__add_context(context) |
|
|
|
|
|
def __ensure_max_size(self): |
|
|
while len(self.__context) > self.__context_span: |
|
|
self.__context.pop(1) |
|
|
return self |
|
|
|
|
|
def __add_context(self, context: dict[str, str]) -> "ChatHistory": |
|
|
self.__context.append(context) |
|
|
return self.__ensure_max_size() |
|
|
|
|
|
def add_message(self, role: Role, content: str) -> "ChatHistory": |
|
|
return self.__add_context({"role": role.value, "content": content}) |
|
|
|
|
|
def get_whole_context(self) -> list[dict[str, str]]: |
|
|
return self.__ensure_max_size().__context.copy() |
|
|
|
|
|
def get_chat_history(self) -> list[dict[str, str]]: |
|
|
return self.get_whole_context()[0:] |
|
|
|
|
|
def get_context_size(self) -> int: |
|
|
return len(self.__ensure_max_size().__context) |
|
|
|
|
|
def last_in_history(self): |
|
|
return self.__context[-1] |
|
|
|
|
|
def reset(self) -> "ChatHistory": |
|
|
self.__context.clear() |
|
|
return self |
|
|
|
|
|
|
|
|
class OpenAIBot: |
|
|
def __init__(self, bot: OpenAI, model: str, prompts: list[PathLike | str], context_span: int, **args): |
|
|
self.logger = LoggerFactory.getLogger(self.__class__.__name__) |
|
|
self.__bot = bot |
|
|
self.__model = model |
|
|
for prompt in prompts: |
|
|
final_prompt = "" |
|
|
with open(prompt, "r") as pf: |
|
|
final_prompt += pf.read() |
|
|
final_prompt = "You are Rick Sanchez from Rick and Morty." |
|
|
self.__history_instance = ChatHistory(context_span=context_span, |
|
|
initial_context=[{"role": Role.SYSTEM.value, "content": final_prompt}]) |
|
|
self.__history = self.get_history_copy() |
|
|
self.__exit_codes: list = getdefault(args, "exit_codes", list()) |
|
|
|
|
|
def __is_exit(self, message: str) -> bool: |
|
|
return message.lower() in self.__exit_codes |
|
|
|
|
|
def respond(self, user_input: str, history: ChatHistory = None, append_user_input: bool = True) -> str | bool: |
|
|
if not history: |
|
|
history = self.__history |
|
|
|
|
|
if self.__is_exit(user_input): |
|
|
return False |
|
|
|
|
|
response = None |
|
|
if user_input: |
|
|
messages = history.add_message(Role.USER, user_input).get_whole_context() |
|
|
|
|
|
if not append_user_input: |
|
|
messages = history.get_whole_context() |
|
|
messages.append({"role": "user", "content": user_input}) |
|
|
|
|
|
chat = self.__bot.chat.completions.create(model=self.__model, messages=messages) |
|
|
self.logger.info("Tokens count, prompts: %s, completion: %s, total: %s", |
|
|
chat.usage.prompt_tokens, chat.usage.completion_tokens, chat.usage.total_tokens) |
|
|
|
|
|
|
|
|
reply = chat.choices[0].message.content |
|
|
response = reply |
|
|
history.add_message(Role.ASSISTANT, reply) |
|
|
|
|
|
return response |
|
|
|
|
|
def get_history_copy(self) -> ChatHistory: |
|
|
return copy.deepcopy(self.__history_instance) |
|
|
|
