Upload 4 files

app.py

@@ -97,7 +97,7 @@ class ResponseWrapper:
     def _extract_text(self) -> str:
         try:
             for part in self._data['candidates'][0]['content']['parts']:
-                if part.get('role') == 'model':
+                if 'thought' not in part:
                     return part['text']
             return ""
         except (KeyError, IndexError):
@@ -519,6 +519,9 @@ def chat_completions():
     else:
         try:
             text_content = response.text
+            prompt_tokens = response.prompt_token_count
+            completion_tokens = response.candidates_token_count
+            total_tokens = response.total_token_count
         except (AttributeError, IndexError, TypeError, ValueError) as e:
             if "response.candidates" in str(e) or "response.text" in str(e):
                 logger.error(f"用户输入被AI安全过滤器阻止")
@@ -551,9 +554,9 @@ def chat_completions():
                 'finish_reason': 'stop'
             }],
             'usage': {
-                'prompt_tokens': 0,
-                'completion_tokens': 0,
-                'total_tokens': 0
+                'prompt_tokens': prompt_tokens,
+                'completion_tokens': completion_tokens,
+                'total_tokens': total_tokens
             }
         }
         logger.info(f"200!")