From 3eddb75c81f436ed9a8c45e710437a849c6adab3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=A4=9A=E7=BE=85=E7=8B=BC?= <Dorowolf314@gmail.com>
Date: Sun, 5 Nov 2023 16:41:30 +0800
Subject: [PATCH] move petal to core

---
 core/petal.py               | 16 ++++++++++++++++
 modules/ask/__init__.py     | 15 ++-------------
 modules/summary/__init__.py | 16 ++++++++--------
 3 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/core/petal.py b/core/petal.py
index b5d8c075..fc1b9e9e 100644
--- a/core/petal.py
+++ b/core/petal.py
@@ -4,6 +4,22 @@ from config import Config
 from core.builtins import Bot
 from core.utils.storedata import get_stored_list, update_stored_list
 
+ONE_K = Decimal('1000')
+# https://openai.com/pricing
+BASE_COST_GPT_3_5 = Decimal('0.002')  # gpt-3.5-turbo： $0.002 / 1K tokens
+# We are not tracking specific tool usage like searches b/c I'm too lazy, use a universal multiplier
+THIRD_PARTY_MULTIPLIER = Decimal('1.5')
+PROFIT_MULTIPLIER = Decimal('1.1')  # At the time we are really just trying to break even
+PRICE_PER_1K_TOKEN = BASE_COST_GPT_3_5 * THIRD_PARTY_MULTIPLIER * PROFIT_MULTIPLIER
+
+
+async def count_petal(tokens)
+    USD_TO_CNY = Decimal('7.3')
+    CNY_TO_PETAL = 100  # Assuming 1 USD = 7.3 CNY, 100 petal = 1 CNY
+
+    price = tokens / ONE_K * PRICE_PER_1K_TOKEN
+    petal = price * USD_TO_CNY * CNY_TO_PETAL
+    return petal
 
 def gained_petal(msg: Bot.MessageSession, amount):
     if Config('openai_api_key') and Config('enable_get_petal'):
diff --git a/modules/ask/__init__.py b/modules/ask/__init__.py
index 6176e78e..ae7c85ca 100644
--- a/modules/ask/__init__.py
+++ b/modules/ask/__init__.py
@@ -9,6 +9,7 @@ from config import Config
 from core.builtins import Bot, Plain, Image
 from core.component import module
 from core.dirty_check import check_bool, rickroll
+from core.petal import count_petal
 from core.utils.cooldown import CoolDown
 
 os.environ['LANGCHAIN_TRACING_V2'] = "true"
@@ -20,17 +21,6 @@ from langchain.callbacks import get_openai_callback  # noqa: E402
 from .agent import agent_executor  # noqa: E402
 from .formatting import generate_latex, generate_code_snippet  # noqa: E402
 
-ONE_K = Decimal('1000')
-# https://openai.com/pricing
-BASE_COST_GPT_3_5 = Decimal('0.002')  # gpt-3.5-turbo： $0.002 / 1K tokens
-# We are not tracking specific tool usage like searches b/c I'm too lazy, use a universal multiplier
-THIRD_PARTY_MULTIPLIER = Decimal('1.5')
-PROFIT_MULTIPLIER = Decimal('1.1')  # At the time we are really just trying to break even
-PRICE_PER_1K_TOKEN = BASE_COST_GPT_3_5 * THIRD_PARTY_MULTIPLIER * PROFIT_MULTIPLIER
-# Assuming 1 USD = 7.3 CNY, 100 petal = 1 CNY
-USD_TO_CNY = Decimal('7.3')
-CNY_TO_PETAL = 100
-
 a = module('ask', developers=['Dianliang233'], desc='{ask.help.desc}')
 
 
@@ -56,8 +46,7 @@ async def _(msg: Bot.MessageSession):
             res = await agent_executor.arun(question)
             tokens = cb.total_tokens
         if not is_superuser:
-            price = tokens / ONE_K * PRICE_PER_1K_TOKEN
-            petal = price * USD_TO_CNY * CNY_TO_PETAL
+            petal = await count_petal(tokens)
             msg.data.modify_petal(-petal)
         else:
             petal = 0
diff --git a/modules/summary/__init__.py b/modules/summary/__init__.py
index a6a58ade..0e851e87 100644
--- a/modules/summary/__init__.py
+++ b/modules/summary/__init__.py
@@ -10,19 +10,17 @@ from core.logger import Logger
 
 openai.api_key = Config('openai_api_key')
 
-s = module(
-    'summary',
-    developers=[
-        'Dianliang233',
-        'OasisAkari'],
+s = module('summary', 
+    developers=['Dianliang233', 'OasisAkari'],
     desc='{summary.help.desc}',
-    available_for=[
-        'QQ',
-        'QQ|Group'])
+    available_for=['QQ', 'QQ|Group'])
 
 
 @s.handle('{{summary.help}}')
 async def _(msg: Bot.MessageSession):
+    if not Config('openai_api_key'):
+        raise Exception(msg.locale.t('error.config.secret.not_found'))
+
     f_msg = await msg.wait_next_message(msg.locale.t('summary.message'), append_instruction=False)
     try:
         f = re.search(r'\[Ke:forward,id=(.*?)\]', f_msg.as_display()).group(1)
@@ -32,6 +30,8 @@ async def _(msg: Bot.MessageSession):
     data = await f_msg.call_api('get_forward_msg', message_id=f)
     msgs = data['messages']
     texts = [f'\n{m["sender"]["nickname"]}：{m["content"]}' for m in msgs]
+    if await check_bool(texts):
+        rickroll(msg)
 
     char_count = sum([len(i) for i in texts])
     wait_msg = await msg.send_message(