add logging to identify tokens used

microsoft · Sep 20, 2024 · 3c07881 · 3c07881
1 parent 4058e42
commit 3c07881
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/Server/modules/open_ai.py b/Server/modules/open_ai.py
@@ -37,13 +37,17 @@ def make_request(
     ):
         for run in range(retries):
             try:
-                response = self.client.chat.completions.create(
-                    model=self.model,
-                    messages=messages,
-                    max_tokens=self.max_tokens,
-                    stream=self.stream,
-                    response_model=response_model,
+                response, completion = (
+                    self.client.chat.completions.create_with_completion(
+                        model=self.model,
+                        messages=messages,
+                        max_tokens=self.max_tokens,
+                        stream=self.stream,
+                        response_model=response_model,
+                    )
                 )
+
+                print("Tokens used:", completion.usage)
                 return response
             except Exception as e:
                 error = str(e)