chore: import upstream snapshot with attribution

2026-07-13 12:36:27 +08:00
commit 05bc60394c
352 changed files with 344349 additions and 0 deletions
@@ -0,0 +1,119 @@
+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,
+            }