Upload 4 files

app.py

@@ -90,10 +90,16 @@ class ResponseWrapper:
         self._data = data
         self._text = self._extract_text()
         self._finish_reason = self._extract_finish_reason()
+        self._prompt_token_count = self._extract_prompt_token_count()
+        self._candidates_token_count = self._extract_candidates_token_count()
+        self._total_token_count = self._extract_total_token_count()
 
     def _extract_text(self) -> str:
         try:
-            return str(self._data)
+            for part in self._data['candidates'][0]['content']['parts']:
+                if part.get('role') == 'model':
+                    return part['text']
+            return ""
         except (KeyError, IndexError):
             return ""
 
@@ -103,6 +109,24 @@ class ResponseWrapper:
         except (KeyError, IndexError):
             return None
 
+    def _extract_prompt_token_count(self) -> Optional[int]:
+        try:
+            return self._data['usageMetadata'].get('promptTokenCount')
+        except (KeyError):
+            return None
+
+    def _extract_candidates_token_count(self) -> Optional[int]:
+        try:
+            return self._data['usageMetadata'].get('candidatesTokenCount')
+        except (KeyError):
+            return None
+
+    def _extract_total_token_count(self) -> Optional[int]:
+        try:
+            return self._data['usageMetadata'].get('totalTokenCount')
+        except (KeyError):
+            return None
+
     @property
     def text(self) -> str:
         return self._text
@@ -111,6 +135,18 @@ class ResponseWrapper:
     def finish_reason(self) -> Optional[str]:
         return self._finish_reason
 
+    @property
+    def prompt_token_count(self) -> Optional[int]:
+        return self._prompt_token_count
+
+    @property
+    def candidates_token_count(self) -> Optional[int]:
+        return self._candidates_token_count
+
+    @property
+    def total_token_count(self) -> Optional[int]:
+        return self._total_token_count
+
 class APIKeyManager:
     def __init__(self):
         self.api_keys = re.findall(r"AIzaSy[a-zA-Z0-9_-]{33}", os.environ.get('KeyArray'))