diff --git a/bot/ai/chatgpt.py b/bot/ai/chatgpt.py index d232d57..9a51654 100644 --- a/bot/ai/chatgpt.py +++ b/bot/ai/chatgpt.py @@ -1,14 +1,33 @@ """ChatGPT (GPT-3.5+) language model from OpenAI.""" import logging -import openai +from openai import AsyncAzureOpenAI, AsyncOpenAI import tiktoken from bot.config import config -logger = logging.getLogger(__name__) +if config.openai.azure: + openai = AsyncAzureOpenAI( + api_key=config.openai.api_key, + api_version=config.openai.azure["version"], + azure_endpoint=config.openai.azure["endpoint"], + azure_deployment=config.openai.azure["deployment"], + ) +else: + openai = AsyncOpenAI(api_key=config.openai.api_key) -openai.api_key = config.openai.api_key encoding = tiktoken.get_encoding("cl100k_base") +logger = logging.getLogger(__name__) + +# Supported models and their context windows +MODELS = { + "gpt-4-1106-preview": 128000, + "gpt-4-vision-preview": 128000, + "gpt-4": 8192, + "gpt-4-32k": 32768, + "gpt-3.5-turbo-1106": 16385, + "gpt-3.5-turbo": 4096, + "gpt-3.5-turbo-16k": 16385, +} class Model: @@ -24,14 +43,20 @@ async def ask(self, question: str, history: list[tuple[str, str]]) -> str: n_input = _calc_n_input(self.name, n_output=config.openai.params["max_tokens"]) messages = self._generate_messages(question, history) messages = shorten(messages, length=n_input) - params = self._prepare_params() - resp = await openai.ChatCompletion.acreate( + params = config.openai.params + logger.debug( + f"> chat request: model=%s, params=%s, messages=%s", + self.name, + params, + messages, + ) + resp = await openai.chat.completions.create( model=self.name, messages=messages, **params, ) logger.debug( - "prompt_tokens=%s, completion_tokens=%s, total_tokens=%s", + "< chat response: prompt_tokens=%s, completion_tokens=%s, total_tokens=%s", resp.usage.prompt_tokens, resp.usage.completion_tokens, resp.usage.total_tokens, @@ -39,16 +64,9 @@ async def ask(self, question: str, history: list[tuple[str, str]]) -> str: answer = self._prepare_answer(resp) return answer - def _prepare_params(self) -> dict: - params = config.openai.params.copy() - if config.openai.azure: - params["api_type"] = "azure" - params["api_base"] = config.openai.azure["endpoint"] - params["api_version"] = config.openai.azure["version"] - params["deployment_id"] = config.openai.azure["deployment"] - return params - - def _generate_messages(self, question: str, history: list[tuple[str, str]]) -> list[dict]: + def _generate_messages( + self, question: str, history: list[tuple[str, str]] + ) -> list[dict]: """Builds message history to provide context for the language model.""" messages = [{"role": "system", "content": config.openai.prompt}] for prev_question, prev_answer in history: @@ -105,7 +123,5 @@ def _calc_n_input(name: str, n_output: int) -> int: """ # OpenAI counts length in tokens, not charactes. # We need to leave some tokens reserved for the output. - n_total = 4096 # max 4096 tokens total by default - if name == "gpt-4": - n_total = 8192 + n_total = MODELS.get(name, 4096) # max 4096 tokens total by default return n_total - n_output diff --git a/bot/ai/custom.py b/bot/ai/custom.py index aa73da2..6ef24f7 100644 --- a/bot/ai/custom.py +++ b/bot/ai/custom.py @@ -1,10 +1,10 @@ """Fine-tuned language model from OpenAI.""" import re -import openai +from openai import AsyncOpenAI from bot.config import config -openai.api_key = config.openai.api_key +openai = AsyncOpenAI(api_key=config.openai.api_key) DEFAULT_STOP = "###" PRE_RE = re.compile(r"<(/?pre)") @@ -29,7 +29,7 @@ async def ask(self, question, history=None) -> str: try: history = history or [] prompt = self._generate_prompt(question, history) - resp = await openai.Completion.acreate( + resp = await openai.completions.create( model=self.name, prompt=prompt, temperature=0.7, diff --git a/bot/ai/dalle.py b/bot/ai/dalle.py index 6ecc63a..2e5a7af 100644 --- a/bot/ai/dalle.py +++ b/bot/ai/dalle.py @@ -1,9 +1,9 @@ """DALL-E model from OpenAI.""" -import openai +from openai import AsyncOpenAI from bot.config import config -openai.api_key = config.openai.api_key +openai = AsyncOpenAI(api_key=config.openai.api_key) class Model: @@ -11,7 +11,9 @@ class Model: async def imagine(self, prompt: str, size: str) -> str: """Generates an image of the specified size according to the description.""" - resp = await openai.Image.acreate(prompt=prompt, size=size, n=1) + resp = await openai.images.generate( + model="dall-e-3", prompt=prompt, size=size, n=1 + ) if len(resp.data) == 0: raise ValueError("received an empty answer") return resp.data[0].url diff --git a/bot/ai/davinci.py b/bot/ai/davinci.py index 816205e..cd12a68 100644 --- a/bot/ai/davinci.py +++ b/bot/ai/davinci.py @@ -1,10 +1,10 @@ """DaVinci (GPT-3) language model from OpenAI.""" import re -import openai +from openai import AsyncOpenAI from bot.config import config -openai.api_key = config.openai.api_key +openai = AsyncOpenAI(api_key=config.openai.api_key) BASE_PROMPT = "Your primary goal is to answer my questions. This may involve writing code or providing helpful information. Be detailed and thorough in your responses. Write code inside 
,
tags." @@ -18,7 +18,7 @@ async def ask(self, question, history=None): """Asks the language model a question and returns an answer.""" history = history or [] prompt = self._generate_prompt(question, history) - resp = await openai.Completion.acreate( + resp = await openai.completions.create( model="text-davinci-003", prompt=prompt, temperature=0.7, diff --git a/bot/askers.py b/bot/askers.py index 0d34c27..ad050c0 100644 --- a/bot/askers.py +++ b/bot/askers.py @@ -24,7 +24,9 @@ async def ask(self, question: str, history: list[tuple[str, str]]) -> str: """Asks AI a question.""" pass - async def reply(self, message: Message, context: CallbackContext, answer: str) -> None: + async def reply( + self, message: Message, context: CallbackContext, answer: str + ) -> None: """Replies with an answer from AI.""" pass @@ -38,7 +40,9 @@ async def ask(self, question: str, history: list[tuple[str, str]]) -> str: """Asks AI a question.""" return await self.model.ask(question, history) - async def reply(self, message: Message, context: CallbackContext, answer: str) -> None: + async def reply( + self, message: Message, context: CallbackContext, answer: str + ) -> None: """Replies with an answer from AI.""" html_answer = markdown.to_html(answer) if len(html_answer) <= MessageLimit.MAX_TEXT_LENGTH: @@ -47,7 +51,8 @@ async def reply(self, message: Message, context: CallbackContext, answer: str) - doc = io.StringIO(answer) caption = ( - textwrap.shorten(answer, width=40, placeholder="...") + " (see attachment for the rest)" + textwrap.shorten(answer, width=40, placeholder="...") + + " (see attachment for the rest)" ) reply_to_message_id = message.id if message.chat.type != Chat.PRIVATE else None await context.bot.send_document( @@ -64,8 +69,13 @@ class ImagineAsker(Asker): model = ai.dalle.Model() size_re = re.compile(r"(256|512|1024)(?:x\1)?\s?(?:px)?") - sizes = {"256": "256x256", "512": "512x512", "1024": "1024x1024"} - default_size = "512x512" + sizes = { + "256": "256x256", + "512": "512x512", + "1024": "1024x1024", + "1792": "1792x1024", + } + default_size = "1024x1024" def __init__(self) -> None: self.caption = "" @@ -76,7 +86,9 @@ async def ask(self, question: str, history: list[tuple[str, str]]) -> str: self.caption = self._extract_caption(question) return await self.model.imagine(prompt=self.caption, size=size) - async def reply(self, message: Message, context: CallbackContext, answer: str) -> None: + async def reply( + self, message: Message, context: CallbackContext, answer: str + ) -> None: """Replies with an answer from AI.""" await message.reply_photo(answer, caption=self.caption) @@ -85,7 +97,7 @@ def _extract_size(self, question: str) -> str: if not match: return self.default_size width = match.group(1) - return self.sizes[width] + return self.sizes.get(width, width) def _extract_caption(self, question: str) -> str: caption = self.size_re.sub("", question).strip() diff --git a/bot/bot.py b/bot/bot.py index 308b39d..a960fd1 100644 --- a/bot/bot.py +++ b/bot/bot.py @@ -28,6 +28,7 @@ level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s", ) +logging.getLogger("httpx").setLevel(logging.WARNING) logging.getLogger("openai").setLevel(logging.WARNING) logging.getLogger("bot.ai.chatgpt").setLevel(logging.INFO) logging.getLogger("bot.commands").setLevel(logging.INFO) @@ -65,24 +66,36 @@ def add_handlers(application: Application): # info commands application.add_handler(CommandHandler("start", commands.Start())) - application.add_handler(CommandHandler("help", commands.Help(), filters=filters.users)) - application.add_handler(CommandHandler("version", commands.Version(), filters=filters.users)) + application.add_handler( + CommandHandler("help", commands.Help(), filters=filters.users) + ) + application.add_handler( + CommandHandler("version", commands.Version(), filters=filters.users) + ) # admin commands application.add_handler( - CommandHandler("config", commands.Config(filters), filters=filters.admins_private) + CommandHandler( + "config", commands.Config(filters), filters=filters.admins_private + ) ) # message-related commands application.add_handler( - CommandHandler("retry", commands.Retry(reply_to), filters=filters.users_or_chats) + CommandHandler( + "retry", commands.Retry(reply_to), filters=filters.users_or_chats + ) ) application.add_handler( - CommandHandler("imagine", commands.Imagine(reply_to), filters=filters.users_or_chats) + CommandHandler( + "imagine", commands.Imagine(reply_to), filters=filters.users_or_chats + ) ) # non-command handler: the default action is to reply to a message - application.add_handler(MessageHandler(filters.messages, commands.Message(reply_to))) + application.add_handler( + MessageHandler(filters.messages, commands.Message(reply_to)) + ) # generic error handler application.add_error_handler(commands.Error()) @@ -108,19 +121,25 @@ async def post_shutdown(application: Application) -> None: def with_message_limit(func): """Refuses to reply if the user has exceeded the message limit.""" - async def wrapper(message: Message, context: CallbackContext, question: str) -> None: + async def wrapper( + message: Message, context: CallbackContext, question: str + ) -> None: username = message.from_user.username user = UserData(context.user_data) # check if the message counter exceeds the message limit if ( not filters.is_known_user(username) - and user.message_counter.value >= config.conversation.message_limit.count > 0 + and user.message_counter.value + >= config.conversation.message_limit.count + > 0 and not user.message_counter.is_expired() ): # this is a group user and they have exceeded the message limit wait_for = models.format_timedelta(user.message_counter.expires_after()) - await message.reply_text(f"Please wait {wait_for} before asking a new question.") + await message.reply_text( + f"Please wait {wait_for} before asking a new question." + ) return # this is a known user or they have not exceeded the message limit, @@ -137,7 +156,9 @@ async def wrapper(message: Message, context: CallbackContext, question: str) -> @with_message_limit async def reply_to(message: Message, context: CallbackContext, question: str) -> None: """Replies to a specific question.""" - await message.chat.send_action(action="typing", message_thread_id=message.message_thread_id) + await message.chat.send_action( + action="typing", message_thread_id=message.message_thread_id + ) try: asker = askers.create(question) diff --git a/requirements.txt b/requirements.txt index 893e9cb..97d6840 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ -httpx==0.23.3 -openai==0.27.2 -beautifulsoup4==4.12.0 -python-telegram-bot==20.2 +httpcore==1.0.2 +httpx==0.25.1 +openai==1.2.3 +beautifulsoup4==4.12.2 +python-telegram-bot==20.6 PyYAML==6.0.1 -tiktoken==0.3.3 +tiktoken==0.5.1 diff --git a/tests/test_askers.py b/tests/test_askers.py index f6e00c2..3291a32 100644 --- a/tests/test_askers.py +++ b/tests/test_askers.py @@ -57,7 +57,7 @@ def test_extract_size(self): size = asker._extract_size(question="a cat 256px") self.assertEqual(size, "256x256") size = asker._extract_size(question="a cat 384") - self.assertEqual(size, "512x512") + self.assertEqual(size, "1024x1024") def test_extract_caption(self): asker = ImagineAsker()