mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-22 16:06:50 +00:00
rayhpeng
9d0a42c1fb
Major refactoring of deerflow/runtime/: - runs/callbacks/ - new callback system (builder, events, title, tokens) - runs/internal/ - execution internals (executor, supervisor, stream_logic, registry) - runs/internal/execution/ - execution artifacts and events handling - runs/facade.py - high-level run facade - runs/observer.py - run observation protocol - runs/types.py - type definitions - runs/store/ - simplified store interfaces (create, delete, query, event) Refactor stream_bridge/: - Replace old providers with contract.py and exceptions.py - Remove async_provider.py, base.py, memory.py Add documentation: - README.md and README_zh.md for runtime module Remove deprecated: - manager.py moved to internal/ - worker.py, schemas.py - user_context.py Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
123 lines
4.4 KiB
Python
123 lines
4.4 KiB
Python
"""Token and message summary callback for runs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
from uuid import UUID
|
|
|
|
from langchain_core.callbacks import BaseCallbackHandler
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RunCompletionData:
|
|
total_input_tokens: int = 0
|
|
total_output_tokens: int = 0
|
|
total_tokens: int = 0
|
|
llm_call_count: int = 0
|
|
lead_agent_tokens: int = 0
|
|
subagent_tokens: int = 0
|
|
middleware_tokens: int = 0
|
|
message_count: int = 0
|
|
last_ai_message: str | None = None
|
|
first_human_message: str | None = None
|
|
|
|
def to_dict(self) -> dict[str, object]:
|
|
return {
|
|
"total_input_tokens": self.total_input_tokens,
|
|
"total_output_tokens": self.total_output_tokens,
|
|
"total_tokens": self.total_tokens,
|
|
"llm_call_count": self.llm_call_count,
|
|
"lead_agent_tokens": self.lead_agent_tokens,
|
|
"subagent_tokens": self.subagent_tokens,
|
|
"middleware_tokens": self.middleware_tokens,
|
|
"message_count": self.message_count,
|
|
"last_ai_message": self.last_ai_message,
|
|
"first_human_message": self.first_human_message,
|
|
}
|
|
|
|
|
|
class RunTokenCallback(BaseCallbackHandler):
|
|
"""Aggregate token and message summary data for one run."""
|
|
|
|
def __init__(self, *, track_token_usage: bool = True) -> None:
|
|
super().__init__()
|
|
self._track_token_usage = track_token_usage
|
|
self._total_input_tokens = 0
|
|
self._total_output_tokens = 0
|
|
self._total_tokens = 0
|
|
self._llm_call_count = 0
|
|
self._lead_agent_tokens = 0
|
|
self._subagent_tokens = 0
|
|
self._middleware_tokens = 0
|
|
self._message_count = 0
|
|
self._last_ai_message: str | None = None
|
|
self._first_human_message: str | None = None
|
|
|
|
def set_first_human_message(self, content: str) -> None:
|
|
self._first_human_message = content[:2000] if content else None
|
|
|
|
def on_llm_end(self, response: Any, *, run_id: UUID, **kwargs: Any) -> None:
|
|
try:
|
|
message = response.generations[0][0].message
|
|
except (IndexError, AttributeError):
|
|
return
|
|
|
|
self._record_ai_message(message, kwargs)
|
|
if not self._track_token_usage:
|
|
return
|
|
|
|
usage = dict(getattr(message, "usage_metadata", None) or {})
|
|
input_tk = usage.get("input_tokens", 0) or 0
|
|
output_tk = usage.get("output_tokens", 0) or 0
|
|
total_tk = usage.get("total_tokens", 0) or input_tk + output_tk
|
|
if total_tk <= 0:
|
|
return
|
|
|
|
self._total_input_tokens += input_tk
|
|
self._total_output_tokens += output_tk
|
|
self._total_tokens += total_tk
|
|
self._llm_call_count += 1
|
|
|
|
caller = self._identify_caller(kwargs)
|
|
if caller.startswith("subagent:"):
|
|
self._subagent_tokens += total_tk
|
|
elif caller.startswith("middleware:"):
|
|
self._middleware_tokens += total_tk
|
|
else:
|
|
self._lead_agent_tokens += total_tk
|
|
|
|
def completion_data(self) -> RunCompletionData:
|
|
return RunCompletionData(
|
|
total_input_tokens=self._total_input_tokens,
|
|
total_output_tokens=self._total_output_tokens,
|
|
total_tokens=self._total_tokens,
|
|
llm_call_count=self._llm_call_count,
|
|
lead_agent_tokens=self._lead_agent_tokens,
|
|
subagent_tokens=self._subagent_tokens,
|
|
middleware_tokens=self._middleware_tokens,
|
|
message_count=self._message_count,
|
|
last_ai_message=self._last_ai_message,
|
|
first_human_message=self._first_human_message,
|
|
)
|
|
|
|
def _record_ai_message(self, message: Any, kwargs: dict[str, Any]) -> None:
|
|
if self._identify_caller(kwargs) != "lead_agent":
|
|
return
|
|
if getattr(message, "tool_calls", None):
|
|
return
|
|
content = getattr(message, "content", "")
|
|
if isinstance(content, str) and content:
|
|
self._last_ai_message = content[:2000]
|
|
self._message_count += 1
|
|
|
|
def _identify_caller(self, kwargs: dict[str, Any]) -> str:
|
|
for tag in kwargs.get("tags") or []:
|
|
if isinstance(tag, str) and (
|
|
tag.startswith("subagent:")
|
|
or tag.startswith("middleware:")
|
|
or tag == "lead_agent"
|
|
):
|
|
return tag
|
|
return "lead_agent"
|