Browse Source

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

tags/0.6.9
sino 1 year ago
parent
commit
528faceb35
No account linked to committer's email address

+ 1
- 1
api/core/agent/cot_agent_runner.py View File

@@ -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='',

+ 3
- 1
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:
@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

Loading…
Cancel
Save