feat: add summary and search commands (#54)

* feat: add summary and search commands

Signed-off-by: Frost Ming <me@frostming.com>

* fix formats

Signed-off-by: Frost Ming <me@frostming.com>

* fix: clean up

Signed-off-by: Frost Ming <me@frostming.com>
This commit is contained in:
Frost Ming
2025-07-08 11:41:57 +08:00
committed by GitHub
parent e91862a548
commit 0b60ae2fbe
31 changed files with 1279 additions and 2588 deletions

View File

@ -1,20 +1,15 @@
from telebot import TeleBot
from telebot.types import Message
import requests
from openai import OpenAI
from os import environ
from . import *
import requests
from telebot import TeleBot
from telebot.types import Message
from config import settings
SD_API_KEY = environ.get("SD3_KEY")
# TODO refactor this shit to __init__
CHATGPT_API_KEY = environ.get("OPENAI_API_KEY")
CHATGPT_BASE_URL = environ.get("OPENAI_API_BASE") or "https://api.openai.com/v1"
CHATGPT_PRO_MODEL = "gpt-4o-2024-05-13"
client = OpenAI(api_key=CHATGPT_API_KEY, base_url=CHATGPT_BASE_URL)
CHATGPT_PRO_MODEL = settings.openai_model
def get_user_balance():
@ -33,7 +28,7 @@ def get_user_balance():
def generate_sd3_image(prompt):
response = requests.post(
f"https://api.stability.ai/v2beta/stable-image/generate/sd3",
"https://api.stability.ai/v2beta/stable-image/generate/sd3",
headers={"authorization": f"Bearer {SD_API_KEY}", "accept": "image/*"},
files={"none": ""},
data={
@ -61,18 +56,14 @@ def sd_handler(message: Message, bot: TeleBot):
)
m = message.text.strip()
prompt = m.strip()
try:
r = generate_sd3_image(prompt)
if r:
with open(f"sd3.jpeg", "rb") as photo:
bot.send_photo(
message.chat.id, photo, reply_to_message_id=message.message_id
)
else:
bot.reply_to(message, "prompt error")
except Exception as e:
print(e)
bot.reply_to(message, "sd3 error")
r = generate_sd3_image(prompt)
if r:
with open("sd3.jpeg", "rb") as photo:
bot.send_photo(
message.chat.id, photo, reply_to_message_id=message.message_id
)
else:
bot.reply_to(message, "prompt error")
def sd_pro_handler(message: Message, bot: TeleBot):
@ -83,7 +74,7 @@ def sd_pro_handler(message: Message, bot: TeleBot):
rewrite_prompt = (
f"revise `{prompt}` to a DALL-E prompt only return the prompt in English."
)
completion = client.chat.completions.create(
completion = settings.openai_client.chat.completions.create(
messages=[{"role": "user", "content": rewrite_prompt}],
max_tokens=2048,
model=CHATGPT_PRO_MODEL,
@ -95,21 +86,17 @@ def sd_pro_handler(message: Message, bot: TeleBot):
message,
f"Generating pretty sd3-turbo image may take some time please left credits {credits} every try will cost 4 criedits wait:\n the real prompt is: {sd_prompt}",
)
try:
r = generate_sd3_image(sd_prompt)
if r:
with open(f"sd3.jpeg", "rb") as photo:
bot.send_photo(
message.chat.id, photo, reply_to_message_id=message.message_id
)
else:
bot.reply_to(message, "prompt error")
except Exception as e:
print(e)
bot.reply_to(message, "sd3 error")
r = generate_sd3_image(sd_prompt)
if r:
with open("sd3.jpeg", "rb") as photo:
bot.send_photo(
message.chat.id, photo, reply_to_message_id=message.message_id
)
else:
bot.reply_to(message, "prompt error")
if SD_API_KEY and CHATGPT_API_KEY:
if SD_API_KEY and settings.openai_api_key:
def register(bot: TeleBot) -> None:
bot.register_message_handler(sd_handler, commands=["sd3"], pass_bot=True)