Crawl/backend/app/services/key_pool_pusher.py

"""
号池数据推送服务
定时将分组和key数据推送到company_base_pool
"""
import logging
import asyncio
from datetime import datetime
import httpx
import os

logger = logging.getLogger(__name__)

# company_base_pool的API地址
COMPANY_BASE_POOL_URL = os.getenv("COMPANY_BASE_POOL_URL", "http://localhost:8011")
SYNC_ENDPOINT = f"{COMPANY_BASE_POOL_URL}/api/admin/key-pool/sync"


async def push_group_keys_to_pool(pool, group_id: int, api_key_ids: list):
    """
    推送分组的key到company_base_pool

    Args:
        pool: 数据库连接池
        group_id: 分组ID
        api_key_ids: key ID列表
    """
    try:
        async with pool.acquire() as conn:
            # 设置search_path
            await conn.execute("SET search_path TO crawl")

            # 获取分组信息
            group = await conn.fetchrow("SELECT name FROM model_groups WHERE id = $1", group_id)
            if not group:
                return

            # 获取keys
            keys = await conn.fetch("""
                SELECT ak.key_value
                FROM api_keys ak
                JOIN group_keys gk ON gk.api_key_id = ak.id
                WHERE gk.group_id = $1 AND gk.is_active = TRUE
            """, group_id)

            if not keys:
                return

            api_keys = [k['key_value'] for k in keys]

            # 获取分组的provider（从绑定的模型推断）
            provider_row = await conn.fetchrow("""
                SELECT DISTINCT m.provider
                FROM models m
                WHERE m.group_id = $1 AND m.provider IS NOT NULL
                LIMIT 1
            """, group_id)
            provider = provider_row['provider'] if provider_row else "dashscope"

            # 推送到company_base_pool
            data = [{
                "group_name": group['name'],
                "provider": provider,
                "model_name": None,
                "api_keys": api_keys
            }]

            async with httpx.AsyncClient(timeout=30) as client:
                response = await client.post(SYNC_ENDPOINT, json=data)

                if response.status_code == 200:
                    result = response.json()
                    logger.info(f"推送分组 {group['name']} 到号池: {result}")
                else:
                    logger.error(f"推送失败: {response.status_code} {response.text}")

    except Exception as e:
        logger.error(f"推送分组 {group_id} 失败: {e}")


async def push_all_groups_to_pool(pool):
    """推送所有分组到company_base_pool"""
    try:
        async with pool.acquire() as conn:
            # 设置search_path
            await conn.execute("SET search_path TO crawl")

            # 获取所有有key的分组
            groups = await conn.fetch("""
                SELECT DISTINCT g.id, g.name
                FROM model_groups g
                JOIN group_keys gk ON gk.group_id = g.id
                WHERE gk.is_active = TRUE
            """)

            logger.info(f"开始推送 {len(groups)} 个分组到号池")

            for group in groups:
                await push_group_keys_to_pool(pool, group['id'], [])

            logger.info("推送完成")

    except Exception as e:
        logger.error(f"推送所有分组失败: {e}")