refac

backend/open_webui/utils/response.py

@@ -6,6 +6,47 @@ from open_webui.utils.misc import (
 )
 
 
+def normalize_usage(usage: dict) -> dict:
+    """
+    Normalize usage statistics to standard format.
+    Handles OpenAI, Ollama, and llama.cpp formats.
+
+    Adds standardized token fields to the original data:
+    - input_tokens: Number of tokens in the prompt
+    - output_tokens: Number of tokens generated
+    - total_tokens: Sum of input and output tokens
+    """
+    if not usage:
+        return {}
+
+    # Map various field names to standard names
+    input_tokens = (
+        usage.get("input_tokens")  # Already standard
+        or usage.get("prompt_tokens")  # OpenAI
+        or usage.get("prompt_eval_count")  # Ollama
+        or usage.get("prompt_n")  # llama.cpp
+        or 0
+    )
+
+    output_tokens = (
+        usage.get("output_tokens")  # Already standard
+        or usage.get("completion_tokens")  # OpenAI
+        or usage.get("eval_count")  # Ollama
+        or usage.get("predicted_n")  # llama.cpp
+        or 0
+    )
+
+    total_tokens = usage.get("total_tokens") or (input_tokens + output_tokens)
+
+    # Add standardized fields to original data
+    result = dict(usage)
+    result["input_tokens"] = int(input_tokens)
+    result["output_tokens"] = int(output_tokens)
+    result["total_tokens"] = int(total_tokens)
+
+    return result
+
+
 def convert_ollama_tool_call_to_openai(tool_calls: list) -> list:
     openai_tool_calls = []
     for tool_call in tool_calls: