From 528faceb3597cfaa8d97a59513760a8ab9054799 Mon Sep 17 00:00:00 2001 From: sino Date: Fri, 17 May 2024 14:45:20 +0800 Subject: [PATCH] fix: cot agent token usage is empty (#4474) --- api/core/agent/cot_agent_runner.py | 2 +- api/core/agent/output_parser/cot_output_parser.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/api/core/agent/cot_agent_runner.py b/api/core/agent/cot_agent_runner.py index ff45474680..e180b9f71c 100644 --- a/api/core/agent/cot_agent_runner.py +++ b/api/core/agent/cot_agent_runner.py @@ -121,7 +121,7 @@ class CotAgentRunner(BaseAgentRunner, ABC): raise ValueError("failed to invoke llm") usage_dict = {} - react_chunks = CotAgentOutputParser.handle_react_stream_output(chunks) + react_chunks = CotAgentOutputParser.handle_react_stream_output(chunks, usage_dict) scratchpad = AgentScratchpadUnit( agent_response='', thought='', diff --git a/api/core/agent/output_parser/cot_output_parser.py b/api/core/agent/output_parser/cot_output_parser.py index 91ac41143b..5a445e9e59 100644 --- a/api/core/agent/output_parser/cot_output_parser.py +++ b/api/core/agent/output_parser/cot_output_parser.py @@ -9,7 +9,7 @@ from core.model_runtime.entities.llm_entities import LLMResultChunk class CotAgentOutputParser: @classmethod - def handle_react_stream_output(cls, llm_response: Generator[LLMResultChunk, None, None]) -> \ + def handle_react_stream_output(cls, llm_response: Generator[LLMResultChunk, None, None], usage_dict: dict) -> \ Generator[Union[str, AgentScratchpadUnit.Action], None, None]: def parse_action(json_str): try: @@ -58,6 +58,8 @@ class CotAgentOutputParser: thought_idx = 0 for response in llm_response: + if response.delta.usage: + usage_dict['usage'] = response.delta.usage response = response.delta.message.content if not isinstance(response, str): continue