Merge pull request #216 from AlexHTW/fix-token-count-if-stream-false

Fix token counting if stream=false

bot/telegram_bot.py

@@ -455,7 +455,9 @@ async def prompt(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
                         total_tokens = int(tokens)
 
             else:
+                total_tokens = 0
                 async def _reply():
+                    nonlocal total_tokens
                     response, total_tokens = await self.openai.get_chat_response(chat_id=chat_id, query=prompt)
 
                     # Split into chunks of 4096 characters (Telegram's message limit)
@@ -488,7 +490,8 @@ async def _reply():
                 allowed_user_ids = self.config['allowed_user_ids'].split(',')
                 if str(user_id) not in allowed_user_ids and 'guests' in self.usage:
                     self.usage["guests"].add_chat_tokens(total_tokens, self.config['token_price'])
-            except:
+            except Exception as e:
+                logging.warning(f'Failed to add tokens to usage_logs: {str(e)}')
                 pass
 
         except Exception as e: