openai non-stream metrics

open-webui · Sep 27, 2024 · bd4d824 · bd4d824
1 parent 68938c0
commit bd4d824
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/examples/filters/langfuse_filter_pipeline.py b/examples/filters/langfuse_filter_pipeline.py
@@ -128,12 +128,12 @@ async def outlet(self, body: dict, user: Optional[dict] = None) -> dict:
         if assistant_message_obj:
             info = assistant_message_obj.get("info", {})
             if isinstance(info, dict):
-                prompt_eval_count = info.get("prompt_eval_count")
-                eval_count = info.get("eval_count")
-                if prompt_eval_count is not None and eval_count is not None:
+                input_tokens = info.get("prompt_eval_count") or info.get("prompt_tokens")
+                output_tokens = info.get("eval_count") or info.get("completion_tokens")
+                if input_tokens is not None and output_tokens is not None:
                     usage = {
-                        "input": prompt_eval_count,
-                        "output": eval_count,
+                        "input": input_tokens,
+                        "output": output_tokens,
                         "unit": "TOKENS",
                     }