mirror of
https://github.com/agent0ai/agent-zero.git
synced 2026-08-25 08:14:19 +00:00
Guard partially parsed tool arguments while streaming and render both text and legacy message response fields. Keep live response logs synchronized without allowing incomplete code arguments to raise.
74 lines
2.9 KiB
Python
74 lines
2.9 KiB
Python
from helpers import persist_chat, tokens
|
|
from helpers.extension import Extension
|
|
from agent import LoopData
|
|
import asyncio
|
|
from helpers.log import LogItem
|
|
from helpers import log
|
|
import math
|
|
from extensions.python.before_main_llm_call._10_log_for_stream import build_heading, build_default_heading
|
|
|
|
|
|
class LogFromStream(Extension):
|
|
|
|
async def execute(
|
|
self,
|
|
loop_data: LoopData = LoopData(),
|
|
text: str = "",
|
|
parsed: dict = {},
|
|
**kwargs,
|
|
):
|
|
if not self.agent:
|
|
return
|
|
|
|
heading = build_default_heading(self.agent)
|
|
if "headline" in parsed:
|
|
heading = build_heading(self.agent, parsed['headline'])
|
|
elif "tool_name" in parsed:
|
|
heading = build_heading(self.agent, f"Using {parsed['tool_name']}") # if the llm skipped headline
|
|
elif "thoughts" in parsed:
|
|
# thought length indicator
|
|
length = "|" * math.ceil(math.sqrt(len(text))/2)
|
|
heading = build_heading(self.agent, f"Thinking... {length}")
|
|
else:
|
|
heading = build_heading(self.agent, "Receiving...")
|
|
|
|
# create log message and store it in loop data temporary params
|
|
if "log_item_generating" not in loop_data.params_temporary:
|
|
loop_data.params_temporary["log_item_generating"] = (
|
|
self.agent.context.log.log(
|
|
type="agent",
|
|
heading=heading,
|
|
)
|
|
)
|
|
|
|
# update log message
|
|
log_item = loop_data.params_temporary["log_item_generating"]
|
|
|
|
# keep reasoning from previous logs in kvps
|
|
kvps = {}
|
|
if log_item.kvps is not None and "reasoning" in log_item.kvps:
|
|
kvps["reasoning"] = log_item.kvps["reasoning"]
|
|
|
|
# step description for UI - using tool XY, writing Python code, etc.
|
|
if parsed is not None and "tool_name" in parsed and parsed["tool_name"]:
|
|
kvps["step"] = f"Using {parsed['tool_name']}..." # using tool XY
|
|
if parsed["tool_name"]=="code_execution_tool":
|
|
tool_args = parsed.get("tool_args")
|
|
if isinstance(tool_args, dict) and "runtime" in tool_args:
|
|
length = ""
|
|
code = tool_args.get("code")
|
|
if isinstance(code, str):
|
|
length = f"({len(code)})"
|
|
kvps["step"] = f"Writing code... {length}"
|
|
if tool_args["runtime"] == "python":
|
|
kvps["step"] = f"Writing Python code... {length}"
|
|
elif tool_args["runtime"] == "nodejs":
|
|
kvps["step"] = f"Writing Node.js code... {length}"
|
|
elif tool_args["runtime"] == "terminal":
|
|
kvps["step"] = f"Writing terminal command... {length}"
|
|
kvps.update(parsed)
|
|
|
|
|
|
|
|
# update the log item
|
|
log_item.update(heading=heading, content=text, kvps=kvps)
|