Merge pull request #366 from majorcheng/main

fix: add estimated openai usage stats

Author: chenyme

app/services/grok/services/chat.py

@@ -32,6 +32,7 @@
     parse_tool_call_block,
     format_tool_history,
 )
+from app.services.grok.utils.usage import estimate_chat_usage, estimate_prompt_tokens
 from app.services.token import get_token_manager, EffortType
 
 
@@ -373,7 +374,8 @@ async def chat_openai(
             model_config_override=model_config_override,
         )
 
-        return response, stream, model
+        prompt_tokens = estimate_prompt_tokens(message)
+        return response, stream, model, prompt_tokens
 
 
 class ChatService:
@@ -426,7 +428,7 @@ async def completions(
             try:
                 # 请求 Grok
                 service = GrokChatService()
-                response, _, model_name = await service.chat_openai(
+                response, _, model_name, prompt_tokens = await service.chat_openai(
                     token,
                     model,
                     messages,
@@ -442,14 +444,27 @@ async def completions(
                 # 处理响应
                 if is_stream:
                     logger.debug(f"Processing stream response: model={model}")
-                    processor = StreamProcessor(model_name, token, show_think, tools=tools, tool_choice=tool_choice)
+                    processor = StreamProcessor(
+                        model_name,
+                        token,
+                        show_think,
+                        tools=tools,
+                        tool_choice=tool_choice,
+                        prompt_tokens=prompt_tokens,
+                    )
                     return wrap_stream_with_usage(
                         processor.process(response), token_mgr, token, model
                     )
 
                 # 非流式
                 logger.debug(f"Processing non-stream response: model={model}")
-                result = await CollectProcessor(model_name, token, tools=tools, tool_choice=tool_choice).process(response)
+                result = await CollectProcessor(
+                    model_name,
+                    token,
+                    tools=tools,
+                    tool_choice=tool_choice,
+                    prompt_tokens=prompt_tokens,
+                ).process(response)
                 try:
                     model_info = ModelService.get(model)
                     effort = (
@@ -506,7 +521,15 @@ async def completions(
 class StreamProcessor(proc_base.BaseProcessor):
     """Stream response processor."""
 
-    def __init__(self, model: str, token: str = "", show_think: bool = None, tools: List[Dict[str, Any]] = None, tool_choice: Any = None):
+    def __init__(
+        self,
+        model: str,
+        token: str = "",
+        show_think: bool = None,
+        tools: List[Dict[str, Any]] = None,
+        tool_choice: Any = None,
+        prompt_tokens: int = 0,
+    ):
         super().__init__(model, token)
         self.response_id: str = None
         self.fingerprint: str = ""
@@ -531,6 +554,17 @@ def __init__(self, model: str, token: str = "", show_think: bool = None, tools:
         self._tool_partial = ""
         self._tool_calls_seen = False
         self._tool_call_index = 0
+        self.prompt_tokens = max(0, int(prompt_tokens or 0))
+        self._completion_parts: list[str] = []
+        self._completion_tool_calls: list[dict[str, Any]] = []
+
+    def _record_content(self, content: str) -> None:
+        if content:
+            self._completion_parts.append(content)
+
+    def _record_tool_call(self, tool_call: Any) -> None:
+        if isinstance(tool_call, dict):
+            self._completion_tool_calls.append(tool_call)
 
     def _with_tool_index(self, tool_call: Any) -> Any:
         if not isinstance(tool_call, dict):
@@ -691,7 +725,14 @@ def _flush_tool_stream(self) -> list[tuple[str, Any]]:
         self._tool_state = "text"
         return events
 
-    def _sse(self, content: str = "", role: str = None, finish: str = None, tool_calls: list = None) -> str:
+    def _sse(
+        self,
+        content: str = "",
+        role: str = None,
+        finish: str = None,
+        tool_calls: list = None,
+        usage: dict | None = None,
+    ) -> str:
         """Build SSE response."""
         delta = {}
         if role:
@@ -712,6 +753,8 @@ def _sse(self, content: str = "", role: str = None, finish: str = None, tool_cal
                 {"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish}
             ],
         }
+        if usage is not None:
+            chunk["usage"] = usage
         return f"data: {orjson.dumps(chunk).decode()}\n\n"
 
     async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, None]:
@@ -780,6 +823,7 @@ async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, N
                         rendered = await dl_service.render_image(
                             url, self.token, img_id
                         )
+                        self._record_content(f"{rendered}\n")
                         yield self._sse(f"{rendered}\n")
 
                     if (
@@ -804,8 +848,10 @@ async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, N
                             if original:
                                 title_safe = title.replace("\n", " ").strip()
                                 if title_safe:
+                                    self._record_content(f"![{title_safe}]({original})\n")
                                     yield self._sse(f"![{title_safe}]({original})\n")
                                 else:
+                                    self._record_content(f"![image]({original})\n")
                                     yield self._sse(f"![image]({original})\n")
                     continue
 
@@ -834,17 +880,21 @@ async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, N
                             self.think_closed_once = True
 
                     if in_think:
+                        self._record_content(filtered)
                         yield self._sse(filtered)
                         continue
 
                     if self._tool_stream_enabled:
                         for kind, payload in self._handle_tool_stream(filtered):
                             if kind == "text":
+                                self._record_content(payload)
                                 yield self._sse(payload)
                             elif kind == "tool":
+                                self._record_tool_call(payload)
                                 yield self._sse(tool_calls=[payload])
                         continue
 
+                    self._record_content(filtered)
                     yield self._sse(filtered)
 
             if self.think_opened:
@@ -854,13 +904,29 @@ async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, N
             if self._tool_stream_enabled:
                 for kind, payload in self._flush_tool_stream():
                     if kind == "text":
+                        self._record_content(payload)
                         yield self._sse(payload)
                     elif kind == "tool":
+                        self._record_tool_call(payload)
                         yield self._sse(tool_calls=[payload])
                 finish_reason = "tool_calls" if self._tool_calls_seen else "stop"
-                yield self._sse(finish=finish_reason)
+                yield self._sse(
+                    finish=finish_reason,
+                    usage=estimate_chat_usage(
+                        prompt_tokens=self.prompt_tokens,
+                        content="".join(self._completion_parts),
+                        tool_calls=self._completion_tool_calls or None,
+                    ),
+                )
             else:
-                yield self._sse(finish="stop")
+                yield self._sse(
+                    finish="stop",
+                    usage=estimate_chat_usage(
+                        prompt_tokens=self.prompt_tokens,
+                        content="".join(self._completion_parts),
+                        tool_calls=self._completion_tool_calls or None,
+                    ),
+                )
 
             yield "data: [DONE]\n\n"
         except asyncio.CancelledError:
@@ -902,11 +968,19 @@ async def process(self, response: AsyncIterable[bytes]) -> AsyncGenerator[str, N
 class CollectProcessor(proc_base.BaseProcessor):
     """Non-stream response processor."""
 
-    def __init__(self, model: str, token: str = "", tools: List[Dict[str, Any]] = None, tool_choice: Any = None):
+    def __init__(
+        self,
+        model: str,
+        token: str = "",
+        tools: List[Dict[str, Any]] = None,
+        tool_choice: Any = None,
+        prompt_tokens: int = 0,
+    ):
         super().__init__(model, token)
         self.filter_tags = get_config("app.filter_tags")
         self.tools = tools
         self.tool_choice = tool_choice
+        self.prompt_tokens = max(0, int(prompt_tokens or 0))
 
     def _filter_content(self, content: str) -> str:
         """Filter special tags in content."""
@@ -1098,22 +1172,11 @@ def _render_card(match: re.Match) -> str:
                     "finish_reason": finish_reason,
                 }
             ],
-            "usage": {
-                "prompt_tokens": 0,
-                "completion_tokens": 0,
-                "total_tokens": 0,
-                "prompt_tokens_details": {
-                    "cached_tokens": 0,
-                    "text_tokens": 0,
-                    "audio_tokens": 0,
-                    "image_tokens": 0,
-                },
-                "completion_tokens_details": {
-                    "text_tokens": 0,
-                    "audio_tokens": 0,
-                    "reasoning_tokens": 0,
-                },
-            },
+            "usage": estimate_chat_usage(
+                prompt_tokens=self.prompt_tokens,
+                content=content,
+                tool_calls=tool_calls_result,
+            ),
         }
 
 

app/services/grok/services/responses.py

@@ -9,6 +9,7 @@
 import orjson
 
 from app.services.grok.services.chat import ChatService
+from app.services.grok.utils.usage import to_responses_usage
 from app.services.grok.utils import process as proc_base
 
 
@@ -725,8 +726,7 @@ async def create(
                 model=model,
                 output_text=content,
                 tool_calls=tool_calls,
-                usage=result.get("usage")
-                or {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0},
+                usage=to_responses_usage(result.get("usage")),
                 status="completed",
                 instructions=instructions,
                 max_output_tokens=max_output_tokens,
@@ -768,6 +768,7 @@ async def create(
         )
 
         async def _stream() -> AsyncGenerator[str, None]:
+            final_usage: Optional[Dict[str, Any]] = None
             yield adapter.created_event()
             yield adapter.in_progress_event()
             async for chunk in result:
@@ -780,6 +781,8 @@ async def _stream() -> AsyncGenerator[str, None]:
                     continue
 
                 if data.get("object") == "chat.completion.chunk":
+                    if data.get("usage"):
+                        final_usage = to_responses_usage(data.get("usage"))
                     delta = (data.get("choices") or [{}])[0].get("delta") or {}
                     if "content" in delta and delta["content"]:
                         for event in adapter.ensure_message_started():
@@ -815,7 +818,7 @@ async def _stream() -> AsyncGenerator[str, None]:
                     yield event
             for event in adapter.tool_arguments_done_events():
                 yield event
-            yield adapter.completed_event()
+            yield adapter.completed_event(final_usage)
 
         return _stream()