kg_agent/utils/llm_account_helper.py

import os
import sys
import random
import logging

current_dir = os.path.dirname(os.path.abspath(__file__))
root_dir = os.path.dirname(current_dir)
sys.path.insert(0, root_dir)
from utils.mysql import mysql_db


logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)


def get_api_key(llm_type: str) -> str:

    """
    获取指定类型的LLM API密钥。
    从数据库中随机选择一个状态为正常的API密钥。

    Args:
        llm_type (str): LLM类型，例如"openai"、"openrouter"等

    Returns:
        str: 有效的API密钥
    """
    try:
        # 查询所有状态为正常（status=1）的API密钥
        api_keys = mysql_db.select(
            table='llm_account',
            columns='api_key',
            where='status = %s and type = %s',
            where_params=(1, llm_type)
        )

        if not api_keys:
            raise ValueError(f"数据库中没有找到有效的{llm_type} API密钥")

        # 随机选择一个API密钥
        selected_key = random.choice(api_keys)['api_key']
        return selected_key

    except Exception as e:
        logging.error(f"获取{llm_type} API密钥失败: {e}")
        # 如果数据库查询失败，抛出异常
        raise