From 3071226e67ddb6b2d806af8778d94b38507a95ce Mon Sep 17 00:00:00 2001 From: suluyan Date: Fri, 26 Dec 2025 18:34:27 +0800 Subject: [PATCH 1/3] feat: current usage --- ms_agent/agent/llm_agent.py | 27 +++++++++++++++++++++++++-- ms_agent/llm/openai_llm.py | 4 ++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/ms_agent/agent/llm_agent.py b/ms_agent/agent/llm_agent.py index b47e3a65a..3128646b3 100644 --- a/ms_agent/agent/llm_agent.py +++ b/ms_agent/agent/llm_agent.py @@ -26,6 +26,11 @@ from ..config.config import Config, ConfigLifecycleHandler from .base import Agent +# Current process shared +TOTAL_PROMPT_TOKENS = 0 +TOTAL_COMPLETION_TOKENS = 0 +TOKEN_LOCK = asyncio.Lock() + class LLMAgent(Agent): """ @@ -467,9 +472,27 @@ async def step( messages = await self.parallel_tool_call(messages) await self.after_tool_call(messages) + + # usage + + prompt_tokens = _response_message.prompt_tokens + completion_tokens = _response_message.completion_tokens + + # 使用全局累积 + global TOTAL_PROMPT_TOKENS, TOTAL_COMPLETION_TOKENS, TOKEN_LOCK + async with TOKEN_LOCK: + TOTAL_PROMPT_TOKENS += prompt_tokens + TOTAL_COMPLETION_TOKENS += completion_tokens + + # tokens in the current step + self.log_output( + f'[usage] prompt_tokens: {prompt_tokens}, completion_tokens: {completion_tokens}' + ) + # total tokens for the process so far self.log_output( - f'[usage] prompt_tokens: {_response_message.prompt_tokens}, ' - f'completion_tokens: {_response_message.completion_tokens}') + f'[usage_total] total_prompt_tokens: {TOTAL_PROMPT_TOKENS}, ' + f'total_completion_tokens: {TOTAL_COMPLETION_TOKENS}') + yield messages def prepare_llm(self): diff --git a/ms_agent/llm/openai_llm.py b/ms_agent/llm/openai_llm.py index 3d012619d..ff5bb6744 100644 --- a/ms_agent/llm/openai_llm.py +++ b/ms_agent/llm/openai_llm.py @@ -132,6 +132,10 @@ def _call_llm(self, """ messages = self._format_input_message(messages) + if kwargs.get('stream', False) and self.args.get( + 'stream_options', {}).get('include_usage', True): + kwargs.setdefault('stream_options', {})['include_usage'] = True + return self.client.chat.completions.create( model=self.model, messages=messages, tools=tools, **kwargs) From 87d3099b61f0a9a803c3a5d42164337f2b633bce Mon Sep 17 00:00:00 2001 From: suluyan Date: Fri, 26 Dec 2025 18:44:44 +0800 Subject: [PATCH 2/3] fix typo --- ms_agent/agent/llm_agent.py | 1 - 1 file changed, 1 deletion(-) diff --git a/ms_agent/agent/llm_agent.py b/ms_agent/agent/llm_agent.py index 3128646b3..3f491b55a 100644 --- a/ms_agent/agent/llm_agent.py +++ b/ms_agent/agent/llm_agent.py @@ -478,7 +478,6 @@ async def step( prompt_tokens = _response_message.prompt_tokens completion_tokens = _response_message.completion_tokens - # 使用全局累积 global TOTAL_PROMPT_TOKENS, TOTAL_COMPLETION_TOKENS, TOKEN_LOCK async with TOKEN_LOCK: TOTAL_PROMPT_TOKENS += prompt_tokens From 99e1e22854ea39ae4ff70d59d3a8d8f2240993f6 Mon Sep 17 00:00:00 2001 From: suluyan Date: Mon, 29 Dec 2025 10:27:40 +0800 Subject: [PATCH 3/3] fix import --- ms_agent/agent/llm_agent.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ms_agent/agent/llm_agent.py b/ms_agent/agent/llm_agent.py index 3f491b55a..b6abfa3ac 100644 --- a/ms_agent/agent/llm_agent.py +++ b/ms_agent/agent/llm_agent.py @@ -1,4 +1,5 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +import asyncio import importlib import inspect import os.path