agstack 1.24.0__tar.gz → 1.25.0__tar.gz

{agstack-1.24.0 → agstack-1.25.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agstack
-Version: 1.24.0
+Version: 1.25.0
 Summary: Production-ready toolkit for building FastAPI and LLM applications
 Author-email: XtraVisions <gitadmin@xtravisions.com>, Chen Hao <chenhao@xtravisions.com>
 Maintainer-email: XtraVisions <gitadmin@xtravisions.com>, Chen Hao <chenhao@xtravisions.com>

{agstack-1.24.0 → agstack-1.25.0}/agstack/llm/client.py

@@ -2,7 +2,8 @@
 
 import logging
 import time
-from typing import TYPE_CHECKING, Any, AsyncIterator, Literal, overload
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Literal, overload
 
 import httpx
 from httpx import AsyncClient
@@ -24,6 +25,70 @@ from ..decorators import autoretry
 logger = logging.getLogger(__name__)
 
 
+@dataclass(frozen=True, slots=True)
+class UsageEvent:
+    """单次 LLM 调用的用量事件
+
+    :param model: 模型名称
+    :param kind: 调用类型（chat / chat_stream / vision / embedding / rerank）
+    :param prompt_tokens: 输入 token 数（后端未返回 usage 时为 0）
+    :param completion_tokens: 输出 token 数
+    :param total_tokens: 总 token 数
+    :param duration_ms: 调用耗时（毫秒）
+    """
+
+    model: str
+    kind: str
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+    duration_ms: int
+
+
+UsageCallback = Callable[[UsageEvent], None]
+"""usage 回调类型：必须同步、快速返回且自行捕获所有异常（不阻塞、不影响主调用链）"""
+
+_usage_callback: UsageCallback | None = None
+
+
+def set_usage_callback(callback: UsageCallback | None) -> None:
+    """注册全局 usage 回调（进程级单例，传 None 注销）
+
+    每次 LLM 调用（chat/chat_stream/vision/embed/rerank，含同步变体）完成后触发一次。
+    回调实现方应自行兜底异常；此处仍会捕获并告警，绝不向主调用链抛出。
+    """
+    global _usage_callback
+    _usage_callback = callback
+
+
+def _usage_field(usage: Any, key: str) -> int:
+    """从 usage（对象/字典/None）中容错提取整数字段"""
+    if usage is None:
+        return 0
+    value = usage.get(key) if isinstance(usage, dict) else getattr(usage, key, None)
+    return int(value or 0)
+
+
+def _emit_usage(model: str, kind: str, usage: Any, duration_ms: int) -> None:
+    """触发 usage 回调；usage 缺失时各分量记 0（保留 model/duration 供监控）"""
+    callback = _usage_callback
+    if callback is None:
+        return
+    try:
+        callback(
+            UsageEvent(
+                model=model,
+                kind=kind,
+                prompt_tokens=_usage_field(usage, "prompt_tokens"),
+                completion_tokens=_usage_field(usage, "completion_tokens"),
+                total_tokens=_usage_field(usage, "total_tokens"),
+                duration_ms=duration_ms,
+            )
+        )
+    except Exception:
+        logger.warning("Usage callback failed", exc_info=True)
+
+
 class LLMError(AppException):
     """LLM 调用错误基类"""
 
@@ -140,10 +205,15 @@ class LLMClient:
         """
         start = time.time()
         model_name = model
+        # 内部调用类型标记（vision 经由 chat 转发时传入，不透传给推理后端）
+        usage_kind = kwargs.pop("usage_kind", "chat")
 
         try:
             if stream:
-                return self._chat_stream(messages, model_name, temperature, max_tokens, start, **kwargs)
+                stream_kind = "chat_stream" if usage_kind == "chat" else usage_kind
+                return self._chat_stream(
+                    messages, model_name, temperature, max_tokens, start, usage_kind=stream_kind, **kwargs
+                )
 
             @autoretry(
                 logger,
@@ -171,6 +241,7 @@ class LLMClient:
                 logger.info(f"LLM: model={model_name}, tokens={usage.total_tokens}, duration={duration_ms}ms")
             else:
                 logger.info(f"LLM: model={model_name}, duration={duration_ms}ms")
+            _emit_usage(model_name, usage_kind, usage, duration_ms)
 
             return response
 
@@ -210,6 +281,7 @@ class LLMClient:
         client = self._get_sync_client()
         start = time.time()
         model_name = model
+        usage_kind = kwargs.pop("usage_kind", "chat")
 
         try:
             response = client.chat.completions.create(
@@ -226,6 +298,7 @@ class LLMClient:
                 logger.info(f"LLM (sync): model={model_name}, tokens={usage.total_tokens}, duration={duration_ms}ms")
             else:
                 logger.info(f"LLM (sync): model={model_name}, duration={duration_ms}ms")
+            _emit_usage(model_name, usage_kind, usage, duration_ms)
 
             return response
 
@@ -250,10 +323,11 @@ class LLMClient:
         temperature: float,
         max_tokens: int | None,
         start_time: float,
+        usage_kind: str = "chat_stream",
         **kwargs: Any,
     ) -> AsyncIterator["ChatCompletionChunk"]:
         """流式响应"""
-        total_tokens = 0
+        final_usage = None
 
         try:
             # noinspection PyTypeChecker
@@ -269,15 +343,17 @@ class LLMClient:
             )
 
             async for chunk in stream:
-                # 收集 token 统计
+                # 收集 token 统计（usage 通常在末尾 chunk 返回）
                 if chunk.usage:
-                    total_tokens = chunk.usage.total_tokens
+                    final_usage = chunk.usage
 
                 yield chunk
 
             # 记录指标
             duration_ms = int((time.time() - start_time) * 1000)
+            total_tokens = final_usage.total_tokens if final_usage else 0
             logger.info(f"LLM stream: model={model}, tokens={total_tokens}, duration={duration_ms}ms")
+            _emit_usage(model, usage_kind, final_usage, duration_ms)
 
         except APITimeoutError as e:
             logger.error(f"LLM stream timeout: {e}")
@@ -294,6 +370,7 @@ class LLMClient:
         :param model: 模型名称
         :return: 向量列表
         """
+        start = time.time()
 
         @autoretry(
             logger,
@@ -309,6 +386,7 @@ class LLMClient:
 
         try:
             response = await _call()
+            _emit_usage(model, "embedding", getattr(response, "usage", None), int((time.time() - start) * 1000))
             data = getattr(response, "data", None)
             if data:
                 return [item.embedding for item in data]
@@ -326,9 +404,11 @@ class LLMClient:
         :return: 向量列表
         """
         client = self._get_sync_client()
+        start = time.time()
 
         try:
             response = client.embeddings.create(model=model, input=texts)
+            _emit_usage(model, "embedding", getattr(response, "usage", None), int((time.time() - start) * 1000))
             data = getattr(response, "data", None)
             if data:
                 return [item.embedding for item in data]
@@ -395,9 +475,9 @@ class LLMClient:
 
         # bypass type check of Literal param `stream`
         if stream:
-            return await self.chat(messages, model=model, stream=True, **kwargs)
+            return await self.chat(messages, model=model, stream=True, usage_kind="vision", **kwargs)
 
-        return await self.chat(messages, model=model, stream=False, **kwargs)
+        return await self.chat(messages, model=model, stream=False, usage_kind="vision", **kwargs)
 
     def vision_sync(
         self,
@@ -431,7 +511,7 @@ class LLMClient:
             {"role": "user", "content": content}  # type: ignore[list-item]
         ]
 
-        return self.chat_sync(messages, model=model, **kwargs)
+        return self.chat_sync(messages, model=model, usage_kind="vision", **kwargs)
 
     async def rerank(
         self,
@@ -450,6 +530,8 @@ class LLMClient:
         """
         from httpx import ConnectTimeout, TimeoutException
 
+        start = time.time()
+
         @autoretry(
             logger,
             retries=3,
@@ -475,6 +557,10 @@ class LLMClient:
             response.raise_for_status()
             data = response.json()
 
+            # 上报用量（仅后端返回 usage 时）
+            if usage := data.get("usage"):
+                _emit_usage(model, "rerank", usage, int((time.time() - start) * 1000))
+
             # 解析响应
             results = []
             for item in data.get("results", []):
@@ -519,6 +605,7 @@ class LLMClient:
         import requests
 
         session = self._get_sync_http_session()
+        start = time.time()
 
         try:
             response = session.post(
@@ -535,6 +622,10 @@ class LLMClient:
             response.raise_for_status()
             data = response.json()
 
+            # 上报用量（仅后端返回 usage 时）
+            if usage := data.get("usage"):
+                _emit_usage(model, "rerank", usage, int((time.time() - start) * 1000))
+
             results = []
             for item in data.get("results", []):
                 index = item.get("index")

{agstack-1.24.0 → agstack-1.25.0}/agstack/llm/flow/agent.py

@@ -31,6 +31,7 @@ class Agent:
         max_tokens: int | None = None,
         max_turns: int = 10,
         *,
+        tool_choice: str = "auto",
         label: str | None = None,
         echo: bool = False,
     ):
@@ -53,6 +54,7 @@ class Agent:
         self.temperature = temperature
         self.max_tokens = max_tokens
         self.max_turns = max_turns
+        self.tool_choice = tool_choice
         self.label = label
         self.echo = echo
 
@@ -135,7 +137,7 @@ class Agent:
 
                 if tools_schema:
                     kwargs["tools"] = tools_schema
-                    kwargs["tool_choice"] = "auto"
+                    kwargs["tool_choice"] = self.tool_choice
 
                 stream = await client.chat(stream=True, **kwargs)
 

{agstack-1.24.0 → agstack-1.25.0}/agstack.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agstack
-Version: 1.24.0
+Version: 1.25.0
 Summary: Production-ready toolkit for building FastAPI and LLM applications
 Author-email: XtraVisions <gitadmin@xtravisions.com>, Chen Hao <chenhao@xtravisions.com>
 Maintainer-email: XtraVisions <gitadmin@xtravisions.com>, Chen Hao <chenhao@xtravisions.com>

{agstack-1.24.0 → agstack-1.25.0}/agstack.egg-info/SOURCES.txt

@@ -73,5 +73,6 @@ tests/test_cache_redis.py
 tests/test_flow_io.py
 tests/test_flow_iterator.py
 tests/test_flow_switch_subflow.py
+tests/test_llm_usage_callback.py
 tests/test_messagebus_memory.py
 tests/test_messagebus_redis.py

{agstack-1.24.0 → agstack-1.25.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "agstack"
-version = "1.24.0"
+version = "1.25.0"
 description = "Production-ready toolkit for building FastAPI and LLM applications"
 readme = "README.md"
 license = "MIT"

agstack-1.25.0/tests/test_llm_usage_callback.py

@@ -0,0 +1,83 @@
+#  Copyright (c) 2020-2026 XtraVisions, All rights reserved.
+
+"""LLM usage 回调钩子测试（纯内存，不依赖真实推理后端）"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from agstack.llm import client as llm_client
+from agstack.llm.client import UsageEvent, _emit_usage, set_usage_callback
+
+
+@pytest.fixture(autouse=True)
+def _reset_callback():
+    """每个用例结束后注销全局回调，避免用例间串扰"""
+    yield
+    set_usage_callback(None)
+
+
+def test_emit_with_object_usage():
+    """openai 响应对象形式的 usage 正常提取三个分量"""
+    events: list[UsageEvent] = []
+    set_usage_callback(events.append)
+
+    usage = SimpleNamespace(prompt_tokens=100, completion_tokens=50, total_tokens=150)
+    _emit_usage("qwen3", "chat", usage, 1200)
+
+    assert events == [
+        UsageEvent(
+            model="qwen3", kind="chat", prompt_tokens=100, completion_tokens=50, total_tokens=150, duration_ms=1200
+        )
+    ]
+
+
+def test_emit_with_dict_usage():
+    """rerank 的 JSON 字典形式 usage 同样可提取，缺失字段记 0"""
+    events: list[UsageEvent] = []
+    set_usage_callback(events.append)
+
+    _emit_usage("bge-reranker", "rerank", {"total_tokens": 320}, 80)
+
+    assert events[0].total_tokens == 320
+    assert events[0].prompt_tokens == 0
+    assert events[0].completion_tokens == 0
+
+
+def test_emit_with_none_usage():
+    """后端未返回 usage 时仍触发事件，各分量为 0（保留 model/duration 供监控）"""
+    events: list[UsageEvent] = []
+    set_usage_callback(events.append)
+
+    _emit_usage("qwen3", "chat_stream", None, 500)
+
+    assert events[0].total_tokens == 0
+    assert events[0].model == "qwen3"
+    assert events[0].duration_ms == 500
+
+
+def test_callback_exception_swallowed():
+    """回调抛异常不得影响主调用链"""
+
+    def _bad_callback(_evt: UsageEvent) -> None:
+        raise RuntimeError("boom")
+
+    set_usage_callback(_bad_callback)
+    _emit_usage("qwen3", "chat", None, 10)  # 不应抛出
+
+
+def test_no_callback_noop():
+    """未注册回调时静默跳过"""
+    assert llm_client._usage_callback is None
+    _emit_usage("qwen3", "chat", None, 10)  # 不应抛出
+
+
+def test_unregister_callback():
+    """传 None 注销回调后不再触发"""
+    events: list[UsageEvent] = []
+    set_usage_callback(events.append)
+    set_usage_callback(None)
+
+    _emit_usage("qwen3", "chat", None, 10)
+
+    assert events == []