Add token usage to response

uwcirg · Nov 12, 2024 · af9b3c5 · af9b3c5
1 parent dc1c506
commit af9b3c5
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/app.py b/app.py
@@ -19,6 +19,8 @@ def chat():
         messages=data['messages']
     )
     message_dict = response.choices[0].message.to_dict()
+    message_dict['prompt_tokens'] = response.usage.prompt_tokens
+    message_dict['completion_tokens'] = response.usage.completion_tokens
     return jsonify(message_dict)
 
 @app.route('/api/go', methods=['GET'])