sycommon-python-lib 0.2.1a18__tar.gz → 0.2.1a20__tar.gz

{sycommon_python_lib-0.2.1a18 → sycommon_python_lib-0.2.1a20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sycommon-python-lib
-Version: 0.2.1a18
+Version: 0.2.1a20
 Summary: Add your description here
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown

{sycommon_python_lib-0.2.1a18 → sycommon_python_lib-0.2.1a20}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "sycommon-python-lib"
-version = "0.2.1a18"
+version = "0.2.1a20"
 description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.11"

sycommon_python_lib-0.2.1a20/src/sycommon/llm/token_usage_buffer.py

@@ -0,0 +1,446 @@
+"""
+Token 使用量本地缓冲管理器
+用于在高并发场景下聚合写入，减少数据库压力
+
+完全异步实现，不阻塞主服务
+"""
+import asyncio
+import logging
+import time
+from collections import defaultdict
+from dataclasses import dataclass, field
+from datetime import datetime, date
+from typing import Optional, Dict, Any, List
+
+from sycommon.config.Config import SingletonMeta
+from sycommon.logging.kafka_log import SYLogger
+
+
+@dataclass
+class TokenUsageBufferItem:
+    """缓冲区中的 Token 使用记录"""
+    user_id: str
+    tenant_id: Optional[str]
+    service_name: str
+    system_env: str
+    input_tokens: int
+    output_tokens: int
+    model: Optional[str]
+    usage_date: date
+    created_at: datetime = field(default_factory=datetime.now)
+
+
+class TokenUsageBuffer(metaclass=SingletonMeta):
+    """
+    Token 使用量本地缓冲管理器（完全异步实现）
+
+    功能：
+    1. 本地内存缓冲：将高频写入先存入内存
+    2. 智能聚合：相同 key 的记录在内存中累加
+    3. 定时刷新：按时间间隔批量写入数据库
+    4. 阈值触发：缓冲区达到一定数量时立即刷新
+    5. 优雅关闭：服务停止时确保数据不丢失
+
+    特点：
+    - 完全异步，使用 asyncio.Lock 保证线程安全
+    - 不阻塞主事件循环
+    """
+
+    # 缓冲区
+    _buffer: Dict[str, TokenUsageBufferItem] = None
+    _buffer_lock: asyncio.Lock = None
+    _flush_task: Optional[asyncio.Task] = None
+    _running: bool = False
+    _loop: Optional[asyncio.AbstractEventLoop] = None
+
+    # 配置参数
+    _flush_interval: float = 15.0  # 刷新间隔（秒）
+    _max_buffer_size: int = 500   # 缓冲区最大条目数
+    _max_batch_size: int = 100    # 单次批量写入最大条目数
+
+    # 统计信息
+    _total_buffered: int = 0
+    _total_flushed: int = 0
+    _last_flush_time: Optional[datetime] = None
+
+    # 回调函数（用于实际写入数据库）
+    _flush_callback = None
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def _ensure_initialized(cls):
+        """确保缓冲区已初始化"""
+        if cls._buffer is None:
+            cls._buffer = {}
+        if cls._buffer_lock is None:
+            cls._buffer_lock = asyncio.Lock()
+
+    @classmethod
+    def configure(
+        cls,
+        flush_interval: float = 15.0,
+        max_buffer_size: int = 500,
+        max_batch_size: int = 100
+    ):
+        """
+        配置缓冲区参数
+
+        Args:
+            flush_interval: 刷新间隔（秒），默认 15 秒
+            max_buffer_size: 缓冲区最大条目数，超过此数量立即刷新
+            max_batch_size: 单次批量写入最大条目数
+        """
+        cls._flush_interval = flush_interval
+        cls._max_buffer_size = max_buffer_size
+        cls._max_batch_size = max_batch_size
+        logging.info(
+            f"TokenUsageBuffer 配置: flush_interval={flush_interval}s, "
+            f"max_buffer_size={max_buffer_size}, max_batch_size={max_batch_size}"
+        )
+
+    @classmethod
+    def set_flush_callback(cls, callback):
+        """
+        设置刷新回调函数
+
+        Args:
+            callback: 异步函数，接收 List[TokenUsageBufferItem] 参数
+        """
+        cls._flush_callback = callback
+
+    @classmethod
+    def _make_key(
+        cls,
+        user_id: str,
+        service_name: str,
+        system_env: str,
+        usage_date: date
+    ) -> str:
+        """生成缓冲区 key"""
+        return f"{user_id}|{service_name}|{system_env}|{usage_date.isoformat()}"
+
+    @classmethod
+    async def add(
+        cls,
+        user_id: Optional[str],
+        input_tokens: int,
+        output_tokens: int,
+        model: Optional[str] = None,
+        tenant_id: Optional[str] = None,
+        service_name: Optional[str] = None,
+        system_env: Optional[str] = None
+    ) -> bool:
+        """
+        异步添加 Token 使用记录到缓冲区
+
+        如果已存在相同 key 的记录，则累加 tokens
+
+        Args:
+            user_id: 用户ID
+            input_tokens: 输入 token 数量
+            output_tokens: 输出 token 数量
+            model: 模型名称
+            tenant_id: 租户ID
+            service_name: 服务名称
+            system_env: 系统环境
+
+        Returns:
+            bool: 是否成功添加到缓冲区
+        """
+        cls._ensure_initialized()
+
+        # 处理空值
+        user_id = user_id or ''
+        service_name = service_name or ''
+        system_env = system_env or ''
+        usage_date = date.today()
+
+        key = cls._make_key(user_id, service_name, system_env, usage_date)
+
+        should_flush = False
+
+        async with cls._buffer_lock:
+            if key in cls._buffer:
+                # 累加到已有记录
+                item = cls._buffer[key]
+                item.input_tokens += input_tokens
+                item.output_tokens += output_tokens
+                item.model = model  # 更新为最新模型
+            else:
+                # 创建新记录
+                cls._buffer[key] = TokenUsageBufferItem(
+                    user_id=user_id,
+                    tenant_id=tenant_id,
+                    service_name=service_name,
+                    system_env=system_env,
+                    input_tokens=input_tokens,
+                    output_tokens=output_tokens,
+                    model=model,
+                    usage_date=usage_date
+                )
+                cls._total_buffered += 1
+
+            # 检查是否需要立即刷新
+            if len(cls._buffer) >= cls._max_buffer_size:
+                should_flush = True
+
+        # 在锁外触发刷新，避免阻塞
+        if should_flush:
+            logging.info(f"TokenUsageBuffer 缓冲区达到阈值 ({cls._max_buffer_size})，触发立即刷新")
+            asyncio.create_task(cls._do_flush())
+
+        return True
+
+    @classmethod
+    def add_sync(
+        cls,
+        user_id: Optional[str],
+        input_tokens: int,
+        output_tokens: int,
+        model: Optional[str] = None,
+        tenant_id: Optional[str] = None,
+        service_name: Optional[str] = None,
+        system_env: Optional[str] = None
+    ) -> bool:
+        """
+        同步添加（非阻塞，fire-and-forget）
+
+        通过 asyncio.to_thread 或 run_coroutine_threadsafe 调用异步方法
+        不会阻塞调用线程
+
+        Args:
+            同 add 方法
+
+        Returns:
+            bool: 是否成功提交
+        """
+        cls._ensure_initialized()
+
+        # 处理空值
+        user_id = user_id or ''
+        service_name = service_name or ''
+        system_env = system_env or ''
+        usage_date = date.today()
+
+        key = cls._make_key(user_id, service_name, system_env, usage_date)
+
+        # 使用简单的方式：直接操作缓冲区（最小化锁竞争）
+        # 这里的操作是原子的，不需要复杂的锁
+        if cls._buffer is None:
+            cls._buffer = {}
+
+        # 快速路径：更新或创建记录
+        if key in cls._buffer:
+            item = cls._buffer[key]
+            item.input_tokens += input_tokens
+            item.output_tokens += output_tokens
+            item.model = model
+        else:
+            cls._buffer[key] = TokenUsageBufferItem(
+                user_id=user_id,
+                tenant_id=tenant_id,
+                service_name=service_name,
+                system_env=system_env,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                model=model,
+                usage_date=usage_date
+            )
+            cls._total_buffered += 1
+
+        # 检查是否需要刷新
+        if len(cls._buffer) >= cls._max_buffer_size:
+            cls._trigger_async_flush()
+
+        return True
+
+    @classmethod
+    def _trigger_async_flush(cls):
+        """触发异步刷新（从同步上下文调用）"""
+        try:
+            loop = asyncio.get_running_loop()
+            loop.create_task(cls._do_flush())
+        except RuntimeError:
+            # 没有运行中的事件循环，忽略
+            pass
+
+    @classmethod
+    async def start(cls):
+        """
+        启动缓冲区刷新任务
+
+        应该在应用启动时调用
+        """
+        if cls._running:
+            return
+
+        cls._ensure_initialized()
+        cls._running = True
+        cls._loop = asyncio.get_running_loop()
+
+        # 启动定时刷新任务
+        cls._flush_task = asyncio.create_task(cls._flush_loop())
+        logging.info(f"TokenUsageBuffer 已启动，刷新间隔: {cls._flush_interval}s")
+
+        # 注册关闭钩子
+        import atexit
+        atexit.register(cls.shutdown_sync)
+
+    @classmethod
+    async def stop(cls):
+        """
+        停止缓冲区并刷新剩余数据
+
+        应该在应用关闭时调用
+        """
+        if not cls._running:
+            return
+
+        cls._running = False
+
+        # 取消定时任务
+        if cls._flush_task:
+            cls._flush_task.cancel()
+            try:
+                await cls._flush_task
+            except asyncio.CancelledError:
+                pass
+
+        # 最后一次刷新
+        await cls._do_flush()
+        logging.info("TokenUsageBuffer 已停止")
+
+    @classmethod
+    def shutdown_sync(cls):
+        """同步关闭方法（用于 atexit）"""
+        if not cls._running:
+            return
+
+        cls._running = False
+
+        # 同步刷新剩余数据
+        if cls._buffer:
+            try:
+                # 尝试在已有事件循环中执行
+                if cls._loop and cls._loop.is_running():
+                    future = asyncio.run_coroutine_threadsafe(cls._do_flush(), cls._loop)
+                    future.result(timeout=10)  # 最多等待 10 秒
+                else:
+                    # 创建新的事件循环
+                    loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(loop)
+                    try:
+                        loop.run_until_complete(cls._do_flush())
+                    finally:
+                        loop.close()
+            except Exception as e:
+                logging.error(f"TokenUsageBuffer 关闭时刷新失败: {e}", exc_info=True)
+
+        logging.info("TokenUsageBuffer 已同步关闭")
+
+    @classmethod
+    async def _flush_loop(cls):
+        """定时刷新循环"""
+        while cls._running:
+            try:
+                await asyncio.sleep(cls._flush_interval)
+                await cls._do_flush()
+            except asyncio.CancelledError:
+                break
+            except Exception as e:
+                logging.error(f"TokenUsageBuffer 刷新循环异常: {e}", exc_info=True)
+                await asyncio.sleep(1)  # 出错后等待 1 秒再继续
+
+    @classmethod
+    async def _do_flush(cls):
+        """执行批量刷新到数据库"""
+        if not cls._buffer:
+            return
+
+        if not cls._flush_callback:
+            logging.warning("TokenUsageBuffer 未设置 flush_callback，无法刷新数据")
+            return
+
+        # 获取要刷新的数据
+        items_to_flush = []
+        async with cls._buffer_lock:
+            if cls._buffer:
+                items_to_flush = list(cls._buffer.values())
+                cls._buffer = {}  # 清空缓冲区
+
+        if not items_to_flush:
+            return
+
+        # 分批刷新
+        total_items = len(items_to_flush)
+        success_count = 0
+        failed_items = []
+
+        for i in range(0, total_items, cls._max_batch_size):
+            batch = items_to_flush[i:i + cls._max_batch_size]
+            try:
+                await cls._flush_callback(batch)
+                success_count += len(batch)
+            except Exception as e:
+                logging.error(f"TokenUsageBuffer 批量刷新失败: {e}", exc_info=True)
+                failed_items.extend(batch)
+
+        # 失败的记录重新放回缓冲区
+        if failed_items:
+            async with cls._buffer_lock:
+                for item in failed_items:
+                    key = cls._make_key(
+                        item.user_id, item.service_name,
+                        item.system_env, item.usage_date
+                    )
+                    if key in cls._buffer:
+                        cls._buffer[key].input_tokens += item.input_tokens
+                        cls._buffer[key].output_tokens += item.output_tokens
+                    else:
+                        cls._buffer[key] = item
+
+        cls._total_flushed += success_count
+        cls._last_flush_time = datetime.now()
+
+        if success_count > 0:
+            logging.info(
+                f"TokenUsageBuffer 刷新完成: 成功 {success_count} 条, "
+                f"失败 {len(failed_items)} 条"
+            )
+
+    @classmethod
+    async def flush_now(cls):
+        """手动触发立即刷新"""
+        await cls._do_flush()
+
+    @classmethod
+    def get_stats(cls) -> Dict[str, Any]:
+        """获取缓冲区统计信息"""
+        buffer_size = len(cls._buffer) if cls._buffer else 0
+        total_tokens = 0
+        if cls._buffer:
+            total_tokens = sum(
+                item.input_tokens + item.output_tokens
+                for item in cls._buffer.values()
+            )
+
+        return {
+            "running": cls._running,
+            "buffer_size": buffer_size,
+            "total_tokens_in_buffer": total_tokens,
+            "total_buffered": cls._total_buffered,
+            "total_flushed": cls._total_flushed,
+            "last_flush_time": cls._last_flush_time.isoformat() if cls._last_flush_time else None,
+            "config": {
+                "flush_interval": cls._flush_interval,
+                "max_buffer_size": cls._max_buffer_size,
+                "max_batch_size": cls._max_batch_size
+            }
+        }
+
+    @classmethod
+    def is_running(cls) -> bool:
+        """检查缓冲区是否正在运行"""
+        return cls._running