AI/dify
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git synced 2025-08-12 19:49:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: cot agent token usage is empty (#4474)

Browse Source
This commit is contained in:
sino 2024-05-17 14:45:20 +08:00 committed by GitHub
parent c2a8fa91b1
commit 528faceb35
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
api/core/agent/cot_agent_runner.py
View File

@ -121,7 +121,7 @@ class CotAgentRunner(BaseAgentRunner, ABC):
raise ValueError("failed to invoke llm") raise ValueError("failed to invoke llm")
usage_dict = {} usage_dict = {}
react_chunks = CotAgentOutputParser.handle_react_stream_output(chunks) react_chunks = CotAgentOutputParser.handle_react_stream_output(chunks, usage_dict)
scratchpad = AgentScratchpadUnit( scratchpad = AgentScratchpadUnit(
agent_response='', agent_response='',
thought='', thought='',

4
api/core/agent/output_parser/cot_output_parser.py
View File

@ -9,7 +9,7 @@ from core.model_runtime.entities.llm_entities import LLMResultChunk
class CotAgentOutputParser: class CotAgentOutputParser:
@classmethod @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]: Generator[Union[str, AgentScratchpadUnit.Action], None, None]:
def parse_action(json_str): def parse_action(json_str):
try: try:
@ -58,6 +58,8 @@ class CotAgentOutputParser:
thought_idx = 0 thought_idx = 0
for response in llm_response: for response in llm_response:
if response.delta.usage:
usage_dict['usage'] = response.delta.usage
response = response.delta.message.content response = response.delta.message.content
if not isinstance(response, str): if not isinstance(response, str):
continue continue