wehub-resource-sync
120 lines
4.1 KiB
Python
120 lines
4.1 KiB
Python
import threading
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
from langchain_core.callbacks import BaseCallbackHandler
|
|
from langchain_core.outputs import LLMResult
|
|
from langchain_core.messages import AIMessage
|
|
|
|
|
|
def _pick_int(d: Optional[Dict[str, Any]], *keys: str) -> int:
|
|
if not d:
|
|
return 0
|
|
for k in keys:
|
|
v = d.get(k)
|
|
if v is not None:
|
|
try:
|
|
return int(v)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
return 0
|
|
|
|
|
|
def _extract_tokens_from_usage_dict(u: Dict[str, Any]) -> Tuple[int, int]:
|
|
"""兼容 OpenAI / Anthropic / Gemini 等字段命名。"""
|
|
tin = _pick_int(u, "input_tokens", "prompt_tokens", "cache_read_input_tokens")
|
|
tout = _pick_int(u, "output_tokens", "completion_tokens", "candidates_token_count")
|
|
return tin, tout
|
|
|
|
|
|
class StatsCallbackHandler(BaseCallbackHandler):
|
|
"""Callback handler that tracks LLM calls, tool calls, and token usage."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self._lock = threading.Lock()
|
|
self.llm_calls = 0
|
|
self.tool_calls = 0
|
|
self.tokens_in = 0
|
|
self.tokens_out = 0
|
|
|
|
def on_llm_start(
|
|
self,
|
|
serialized: Dict[str, Any],
|
|
prompts: List[str],
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""Increment LLM call counter when an LLM starts."""
|
|
with self._lock:
|
|
self.llm_calls += 1
|
|
|
|
def on_chat_model_start(
|
|
self,
|
|
serialized: Dict[str, Any],
|
|
messages: List[List[Any]],
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""Increment LLM call counter when a chat model starts."""
|
|
with self._lock:
|
|
self.llm_calls += 1
|
|
|
|
def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
|
|
"""Extract token usage from LLM response."""
|
|
tin, tout = 0, 0
|
|
try:
|
|
generation = response.generations[0][0]
|
|
except (IndexError, TypeError):
|
|
generation = None
|
|
|
|
if generation is not None and hasattr(generation, "message"):
|
|
message = generation.message
|
|
if isinstance(message, AIMessage):
|
|
if hasattr(message, "usage_metadata") and message.usage_metadata:
|
|
um = message.usage_metadata
|
|
if isinstance(um, dict):
|
|
tin = _pick_int(um, "input_tokens", "prompt_tokens")
|
|
tout = _pick_int(um, "output_tokens", "completion_tokens")
|
|
if tin == 0 and tout == 0 and hasattr(message, "response_metadata"):
|
|
rm = message.response_metadata or {}
|
|
u = rm.get("token_usage") or rm.get("usage_metadata") or rm.get("usage")
|
|
if isinstance(u, dict):
|
|
tin, tout = _extract_tokens_from_usage_dict(u)
|
|
|
|
if tin == 0 and tout == 0 and generation is not None:
|
|
gen_info = getattr(generation, "generation_info", None) or {}
|
|
if isinstance(gen_info, dict):
|
|
u = gen_info.get("token_usage") or gen_info.get("usage")
|
|
if isinstance(u, dict):
|
|
tin, tout = _extract_tokens_from_usage_dict(u)
|
|
|
|
if tin == 0 and tout == 0:
|
|
llm_out = getattr(response, "llm_output", None)
|
|
if isinstance(llm_out, dict):
|
|
u = llm_out.get("token_usage") or llm_out.get("usage")
|
|
if isinstance(u, dict):
|
|
tin, tout = _extract_tokens_from_usage_dict(u)
|
|
|
|
if tin or tout:
|
|
with self._lock:
|
|
self.tokens_in += tin
|
|
self.tokens_out += tout
|
|
|
|
def on_tool_start(
|
|
self,
|
|
serialized: Dict[str, Any],
|
|
input_str: str,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""Increment tool call counter when a tool starts."""
|
|
with self._lock:
|
|
self.tool_calls += 1
|
|
|
|
def get_stats(self) -> Dict[str, Any]:
|
|
"""Return current statistics."""
|
|
with self._lock:
|
|
return {
|
|
"llm_calls": self.llm_calls,
|
|
"tool_calls": self.tool_calls,
|
|
"tokens_in": self.tokens_in,
|
|
"tokens_out": self.tokens_out,
|
|
}
|