reduce code duplication and reuse image_to_data_uri

This commit is contained in:
F4ria 2024-04-17 14:03:26 +08:00
parent 19bb441533
commit 2abcc90a32
3 changed files with 16 additions and 18 deletions

View File

@ -1,5 +1,6 @@
from __future__ import annotations
import base64
import importlib
import re
import traceback
@ -186,5 +187,16 @@ def enrich_text_with_urls(text: str) -> str:
return text
def image_to_data_uri(file_path):
with open(file_path, "rb") as image_file:
encoded_image = base64.b64encode(image_file.read()).decode("utf-8")
return f"data:image/png;base64,{encoded_image}"
# `import *` will give you these
__all__ = ["bot_reply_first", "bot_reply_markdown", "enrich_text_with_urls"]
__all__ = [
"bot_reply_first",
"bot_reply_markdown",
"enrich_text_with_urls",
"image_to_data_uri",
]

View File

@ -1,6 +1,4 @@
import base64
from os import environ
from pathlib import Path
import time
from openai import OpenAI
@ -172,9 +170,7 @@ def chatgpt_photo_handler(message: Message, bot: TeleBot) -> None:
downloaded_file = bot.download_file(file_path)
with open("chatgpt_temp.jpg", "wb") as temp_file:
temp_file.write(downloaded_file)
with open("chatgpt_temp.jpg", "rb") as image_file:
image_data = image_file.read()
base64_image_data = base64.b64encode(image_data).decode("utf-8")
try:
r = client.chat.completions.create(
max_tokens=1024,
@ -185,9 +181,7 @@ def chatgpt_photo_handler(message: Message, bot: TeleBot) -> None:
{"type": "text", "text": prompt},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{base64_image_data}"
},
"image_url": {"url": image_to_data_uri("chatgpt_temp.jpg")},
},
],
}

View File

@ -1,6 +1,4 @@
import base64
from os import environ
from pathlib import Path
from openai import OpenAI
import requests
@ -88,12 +86,6 @@ def yi_handler(message: Message, bot: TeleBot) -> None:
bot_reply_markdown(reply_id, who, yi_reply_text, bot)
def _image_to_data_uri(file_path):
with open(file_path, "rb") as image_file:
encoded_image = base64.b64encode(image_file.read()).decode("utf-8")
return f"data:image/png;base64,{encoded_image}"
def yi_photo_handler(message: Message, bot: TeleBot) -> None:
s = message.caption
prompt = s.strip()
@ -121,7 +113,7 @@ def yi_photo_handler(message: Message, bot: TeleBot) -> None:
{"type": "text", "text": prompt},
{
"type": "image_url",
"image_url": {"url": _image_to_data_uri("yi_temp.jpg")},
"image_url": {"url": image_to_data_uri("yi_temp.jpg")},
},
],
}