fix(ai-proxy): cohere analytics

Kong · Apr 25, 2024 · ea61c0d · ea61c0d
1 parent 45a062a
commit ea61c0d
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/kong/llm/drivers/cohere.lua b/kong/llm/drivers/cohere.lua
@@ -250,9 +250,20 @@ local transformers_from = {
       messages.id = response_table.generation_id
 
       local stats = {
-        completion_tokens = response_table.token_count and response_table.token_count.response_tokens or nil,
-        prompt_tokens = response_table.token_count and response_table.token_count.prompt_tokens or nil,
-        total_tokens = response_table.token_count and response_table.token_count.total_tokens or nil,
+        completion_tokens = response_table.meta
+                        and response_table.meta.billed_units
+                        and response_table.meta.billed_units.output_tokens
+                        or nil,
+
+        prompt_tokens = response_table.meta
+                    and response_table.meta.billed_units
+                    and response_table.meta.billed_units.input_tokens
+                    or nil,
+
+        total_tokens = response_table.meta
+                  and response_table.meta.billed_units
+                  and (response_table.meta.billed_units.output_tokens + response_table.meta.billed_units.input_tokens)
+                  or nil,
       }
       messages.usage = stats