AstrBot 4.9.2__py3-none-any.whl → 4.10.0a1__py3-none-any.whl

astrbot/cli/__init__.py

@@ -1 +1 @@
-__version__ = "4.9.2"
+__version__ = "4.10.0-alpha.1"

astrbot/core/agent/message.py

@@ -3,7 +3,7 @@
 
 from typing import Any, ClassVar, Literal, cast
 
-from pydantic import BaseModel, GetCoreSchemaHandler, model_validator
+from pydantic import BaseModel, GetCoreSchemaHandler, model_serializer, model_validator
 from pydantic_core import core_schema
 
 
@@ -122,10 +122,12 @@ class ToolCall(BaseModel):
     extra_content: dict[str, Any] | None = None
     """Extra metadata for the tool call."""
 
-    def model_dump(self, **kwargs: Any) -> dict[str, Any]:
+    @model_serializer(mode="wrap")
+    def serialize(self, handler):
+        data = handler(self)
         if self.extra_content is None:
-            kwargs.setdefault("exclude", set()).add("extra_content")
-        return super().model_dump(**kwargs)
+            data.pop("extra_content", None)
+        return data
 
 
 class ToolCallPart(BaseModel):

astrbot/core/agent/response.py

@@ -1,7 +1,8 @@
 import typing as T
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from astrbot.core.message.message_event_result import MessageChain
+from astrbot.core.provider.entities import TokenUsage
 
 
 class AgentResponseData(T.TypedDict):
@@ -12,3 +13,23 @@ class AgentResponseData(T.TypedDict):
 class AgentResponse:
     type: str
     data: AgentResponseData
+
+
+@dataclass
+class AgentStats:
+    token_usage: TokenUsage = field(default_factory=TokenUsage)
+    start_time: float = 0.0
+    end_time: float = 0.0
+    time_to_first_token: float = 0.0
+
+    @property
+    def duration(self) -> float:
+        return self.end_time - self.start_time
+
+    def to_dict(self) -> dict:
+        return {
+            "token_usage": self.token_usage.__dict__,
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "time_to_first_token": self.time_to_first_token,
+        }

astrbot/core/agent/run_context.py

@@ -9,7 +9,7 @@ from .message import Message
 TContext = TypeVar("TContext", default=Any)
 
 
-@dataclass(config={"arbitrary_types_allowed": True})
+@dataclass
 class ContextWrapper(Generic[TContext]):
     """A context for running an agent, which can be used to pass additional data or state."""
 

astrbot/core/agent/runners/tool_loop_agent_runner.py

@@ -1,4 +1,5 @@
 import sys
+import time
 import traceback
 import typing as T
 
@@ -12,6 +13,7 @@ from mcp.types import (
 )
 
 from astrbot import logger
+from astrbot.core.message.components import Json
 from astrbot.core.message.message_event_result import (
     MessageChain,
 )
@@ -24,7 +26,7 @@ from astrbot.core.provider.provider import Provider
 
 from ..hooks import BaseAgentRunHooks
 from ..message import AssistantMessageSegment, Message, ToolCallMessageSegment
-from ..response import AgentResponseData
+from ..response import AgentResponseData, AgentStats
 from ..run_context import ContextWrapper, TContext
 from ..tool_executor import BaseFunctionToolExecutor
 from .base import AgentResponse, AgentState, BaseAgentRunner
@@ -69,6 +71,9 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
             )
         self.run_context.messages = messages
 
+        self.stats = AgentStats()
+        self.stats.start_time = time.time()
+
     async def _iter_llm_responses(self) -> T.AsyncGenerator[LLMResponse, None]:
         """Yields chunks *and* a final LLMResponse."""
         if self.streaming:
@@ -98,6 +103,10 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
 
         async for llm_response in self._iter_llm_responses():
             if llm_response.is_chunk:
+                # update ttft
+                if self.stats.time_to_first_token == 0:
+                    self.stats.time_to_first_token = time.time() - self.stats.start_time
+
                 if llm_response.result_chain:
                     yield AgentResponse(
                         type="streaming_delta",
@@ -121,6 +130,10 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
                     )
                 continue
             llm_resp_result = llm_response
+
+            if not llm_response.is_chunk and llm_response.usage:
+                # only count the token usage of the final response for computation purpose
+                self.stats.token_usage += llm_response.usage
             break  # got final response
 
         if not llm_resp_result:
@@ -132,6 +145,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
         if llm_resp.role == "err":
             # 如果 LLM 响应错误，转换到错误状态
             self.final_llm_resp = llm_resp
+            self.stats.end_time = time.time()
             self._transition_state(AgentState.ERROR)
             yield AgentResponse(
                 type="err",
@@ -146,6 +160,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
             # 如果没有工具调用，转换到完成状态
             self.final_llm_resp = llm_resp
             self._transition_state(AgentState.DONE)
+            self.stats.end_time = time.time()
             # record the final assistant message
             self.run_context.messages.append(
                 Message(
@@ -175,22 +190,19 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
         # 如果有工具调用，还需处理工具调用
         if llm_resp.tools_call_name:
             tool_call_result_blocks = []
-            for tool_call_name in llm_resp.tools_call_name:
-                yield AgentResponse(
-                    type="tool_call",
-                    data=AgentResponseData(
-                        chain=MessageChain(type="tool_call").message(
-                            f"🔨 调用工具: {tool_call_name}"
-                        ),
-                    ),
-                )
             async for result in self._handle_function_tools(self.req, llm_resp):
                 if isinstance(result, list):
                     tool_call_result_blocks = result
                 elif isinstance(result, MessageChain):
-                    result.type = "tool_call_result"
+                    if result.type is None:
+                        # should not happen
+                        continue
+                    if result.type == "tool_direct_result":
+                        ar_type = "tool_call_result"
+                    else:
+                        ar_type = result.type
                     yield AgentResponse(
-                        type="tool_call_result",
+                        type=ar_type,
                         data=AgentResponseData(chain=result),
                     )
             # 将结果添加到上下文中
@@ -233,6 +245,19 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
             llm_response.tools_call_args,
             llm_response.tools_call_ids,
         ):
+            yield MessageChain(
+                type="tool_call",
+                chain=[
+                    Json(
+                        data={
+                            "id": func_tool_id,
+                            "name": func_tool_name,
+                            "args": func_tool_args,
+                            "ts": time.time(),
+                        }
+                    )
+                ],
+            )
             try:
                 if not req.func_tool:
                     return
@@ -306,7 +331,6 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
                                     content=res.content[0].text,
                                 ),
                             )
-                            yield MessageChain().message(res.content[0].text)
                         elif isinstance(res.content[0], ImageContent):
                             tool_call_result_blocks.append(
                                 ToolCallMessageSegment(
@@ -328,7 +352,6 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
                                         content=resource.text,
                                     ),
                                 )
-                                yield MessageChain().message(resource.text)
                             elif (
                                 isinstance(resource, BlobResourceContents)
                                 and resource.mimeType
@@ -352,7 +375,22 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
                                         content="返回的数据类型不受支持",
                                     ),
                                 )
-                                yield MessageChain().message("返回的数据类型不受支持。")
+
+                        # yield the last tool call result
+                        if tool_call_result_blocks:
+                            last_tcr_content = str(tool_call_result_blocks[-1].content)
+                            yield MessageChain(
+                                type="tool_call_result",
+                                chain=[
+                                    Json(
+                                        data={
+                                            "id": func_tool_id,
+                                            "ts": time.time(),
+                                            "result": last_tcr_content,
+                                        }
+                                    )
+                                ],
+                            )
 
                     elif resp is None:
                         # Tool 直接请求发送消息给用户
@@ -362,6 +400,7 @@ class ToolLoopAgentRunner(BaseAgentRunner[TContext]):
                             f"{func_tool_name} 没有没有返回值或者将结果直接发送给用户，此工具调用不会被记录到历史中。"
                         )
                         self._transition_state(AgentState.DONE)
+                        self.stats.end_time = time.time()
                     else:
                         # 不应该出现其他类型
                         logger.warning(

astrbot/core/astr_agent_context.py

@@ -6,8 +6,10 @@ from astrbot.core.platform.astr_message_event import AstrMessageEvent
 from astrbot.core.star.context import Context
 
 
-@dataclass(config={"arbitrary_types_allowed": True})
+@dataclass
 class AstrAgentContext:
+    __pydantic_config__ = {"arbitrary_types_allowed": True}
+
     context: Context
     """The star context instance"""
     event: AstrMessageEvent

astrbot/core/astr_agent_run_util.py

@@ -4,6 +4,7 @@ from collections.abc import AsyncGenerator
 from astrbot.core import logger
 from astrbot.core.agent.runners.tool_loop_agent_runner import ToolLoopAgentRunner
 from astrbot.core.astr_agent_context import AstrAgentContext
+from astrbot.core.message.components import Json
 from astrbot.core.message.message_event_result import (
     MessageChain,
     MessageEventResult,
@@ -33,16 +34,27 @@ async def run_agent(
                     msg_chain = resp.data["chain"]
                     if msg_chain.type == "tool_direct_result":
                         # tool_direct_result 用于标记 llm tool 需要直接发送给用户的内容
-                        await astr_event.send(resp.data["chain"])
+                        await astr_event.send(msg_chain)
                         continue
+                    if astr_event.get_platform_id() == "webchat":
+                        await astr_event.send(msg_chain)
                     # 对于其他情况，暂时先不处理
                     continue
                 elif resp.type == "tool_call":
                     if agent_runner.streaming:
                         # 用来标记流式响应需要分节
                         yield MessageChain(chain=[], type="break")
-                    if show_tool_use:
+
+                    if astr_event.get_platform_name() == "webchat":
                         await astr_event.send(resp.data["chain"])
+                    elif show_tool_use:
+                        json_comp = resp.data["chain"].chain[0]
+                        if isinstance(json_comp, Json):
+                            m = f"🔨 调用工具: {json_comp.data.get('name')}"
+                        else:
+                            m = "🔨 调用工具..."
+                        chain = MessageChain(type="tool_call").message(m)
+                        await astr_event.send(chain)
                     continue
 
                 if stream_to_general and resp.type == "streaming_delta":
@@ -69,6 +81,15 @@ async def run_agent(
                         continue
                     yield resp.data["chain"]  # MessageChain
             if agent_runner.done():
+                # send agent stats to webchat
+                if astr_event.get_platform_name() == "webchat":
+                    await astr_event.send(
+                        MessageChain(
+                            type="agent_stats",
+                            chain=[Json(data=agent_runner.stats.to_dict())],
+                        )
+                    )
+
                 break
 
         except Exception as e: