Crawl/backend/app/services/key_pool_pusher.py

"""
号池数据推送服务
定时将模型和 key 数据推送到 company_base_pool

兼容两种绑定方式：
1. 旧方式：模型直接绑定 key（models.api_key_id）
2. 新方式：模型通过分组绑定 key（models.group_id）
"""
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:8010")
SYNC_ENDPOINT = f"{COMPANY_BASE_POOL_URL}/api/crawler/sync"


async def push_model_api_keys_to_pool(pool):
    """
    推送模型和 API Key 绑定数据到 company_base_pool（旧方式兼容）

    逻辑：
    - 只推送有 api_key_id 的模型
    - 按 group_name 聚合，每个分组推送一个对象
    """
    try:
        async with pool.acquire() as conn:
            # 设置 search_path
            await conn.execute("SET search_path TO crawl")

            # 获取所有有 api_key_id 的模型
            models = await conn.fetch("""
                SELECT m.name, m.url, m.api_key_id, m.group_id,
                       g.name AS group_name,
                       ak.key_value AS api_key
                FROM models m
                LEFT JOIN model_groups g ON g.id = m.group_id
                LEFT JOIN api_keys ak ON ak.id = m.api_key_id
                WHERE m.api_key_id IS NOT NULL
            """)

            if not models:
                logger.info("没有模型绑定 API Key，跳过推送")
                return

            # 按 group_name 聚合
            groups_map = {}
            for m in models:
                group_name = m['group_name'] or '未分组'
                if group_name not in groups_map:
                    groups_map[group_name] = {
                        'group_name': group_name,
                        'provider': 'dashscope',
                        'base_url': 'https://dashscope.aliyuncs.com',
                        'model_name': None,
                        'api_keys': []
                    }
                if m['api_key'] and m['api_key'] not in groups_map[group_name]['api_keys']:
                    groups_map[group_name]['api_keys'].append(m['api_key'])

            # 推送到 company_base_pool
            data = list(groups_map.values())

            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"推送模型 API Key 到号池：{len(data)} 个分组，{sum(len(g['api_keys']) for g in data)} 个 key")
                else:
                    logger.error(f"推送失败：{response.status_code} {response.text}")

    except Exception as e:
        logger.error(f"推送模型 API Key 失败：{e}")


async def push_group_keys_to_pool(pool, group_id: int):
    """
    推送分组的 key 到 company_base_pool（新方式）
    """
    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]

            # 推送到 company_base_pool
            data = [{
                "group_name": group['name'],
                "provider": "dashscope",
                "base_url": "https://dashscope.aliyuncs.com",
                "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:
        # 方式 1：推送模型直接绑定的 API Key（旧方式）
        await push_model_api_keys_to_pool(pool)

        # 方式 2：推送分组和 key（新方式）
        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}")