添加获取token数和计费

models/generator.py

@@ -1,4 +1,5 @@
 from typing import Optional
+from decimal import Decimal
 from dashscope import MultiModalConversation
 from models.logger import setup_logger
 
@@ -10,8 +11,8 @@ class Generator:
         self.model = model
         self.text = text
 
-    def process_local_image(self, image_path: str) -> Optional[str]:
-        """处理本地单张图片，生成描述词"""
+    def process_local_image(self, image_path: str):
+        """处理本地单张图片"""
         try:
             image_url = f"file://{image_path}"
             messages = [
@@ -29,22 +30,14 @@ class Generator:
                 messages=messages
             )
 
-            if response and hasattr(response, 'output'):
-                if hasattr(response.output.choices[0].message.content[0], 'text'):
-                    return response.output.choices[0].message.content[0]["text"]
-                elif isinstance(response.output.choices[0].message.content[0], dict):
-                    return response.output.choices[0].message.content[0].get("text", "")
+            return response
 
-            return None
-
-        except AttributeError as e:
-            raise AttributeError(e)
         except Exception as e:
             logger.error(f"API调用失败: {e}")
             return None
 
-    def process_link_image(self, image_url: str) -> Optional[str]:
-        """处理链接图片，生成描述词"""
+    def process_link_image(self, image_url: str):
+        """处理链接图片"""
         try:
             messages = [
                 {
@@ -61,15 +54,22 @@ class Generator:
                 messages=messages
             )
 
-            if response and hasattr(response, 'output'):
-                if hasattr(response.output.choices[0].message.content[0], 'text'):
-                    return response.output.choices[0].message.content[0]["text"]
-                elif isinstance(response.output.choices[0].message.content[0], dict):
-                    return response.output.choices[0].message.content[0].get("text", "")
-            return None
+            return response
 
-        except AttributeError as e:
-            raise AttributeError(e)
         except Exception as e:
             logger.error(f"API调用失败: {e}")
             return None
+
+    def get_data(self, image_string: str) -> Optional[tuple[str, tuple]]:
+        """处理图片"""
+        if image_string.startswith(('http://', 'https://')):
+            response = self.process_link_image(image_string)
+        else:
+            response = self.process_local_image(image_string)
+
+        if isinstance(response.output.choices[0].message.content[0], dict):
+            text =  response.output.choices[0].message.content[0].get("text", "")
+            tokens = (response.usage.input_tokens, response.usage.output_tokens)
+            return text, tokens
+
+        return None