scripts/sync_bailian_models.py

#!/usr/bin/env python3
import os
import sys
import sqlite3
import json
import urllib.request
import subprocess
from datetime import datetime

DB_PATH = "/app/server/data/freeapi.db"
BAILIAN_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1/models"

def get_table_columns(conn, table):
    cursor = conn.execute(f"PRAGMA table_info({table})")
    return [row[1] for row in cursor.fetchall()]

def ensure_provider_table(conn):
    # 检查是否存在 providers 或 channels 表，如果没有则创建
    cursor = conn.cursor()
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name IN ('providers', 'channels')")
    tables = [row[0] for row in cursor.fetchall()]
    if 'providers' in tables:
        return 'providers'
    elif 'channels' in tables:
        return 'channels'
    else:
        cursor.execute("""
            CREATE TABLE providers (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT,
                base_url TEXT,
                api_key TEXT,
                created_at TEXT
            )
        """)
        conn.commit()
        return 'providers'

def get_or_create_bailian_provider(conn, provider_table):
    cursor = conn.cursor()
    cursor.execute(f"SELECT id FROM {provider_table} WHERE base_url LIKE '%dashscope.aliyuncs.com%' LIMIT 1")
    row = cursor.fetchone()
    if row:
        return row[0]
    cursor.execute(f"""
        INSERT INTO {provider_table} (name, base_url, api_key, created_at)
        VALUES (?, ?, ?, ?)
    """, ("Bailian", BAILIAN_URL, "", datetime.utcnow().isoformat()))
    conn.commit()
    return cursor.lastrowid

def main():
    api_key = os.getenv("BAILIAN_API_KEY")
    if not api_key:
        print("⚠️ BAILIAN_API_KEY not set, skipping model sync")
        return 0

    # 获取模型列表
    req = urllib.request.Request(BAILIAN_URL)
    req.add_header("Authorization", f"Bearer {api_key}")
    try:
        with urllib.request.urlopen(req, timeout=30) as response:
            data = json.loads(response.read().decode('utf-8'))
            models = [item["id"] for item in data.get("data", [])]
            if not models:
                print("⚠️ No models found")
                return 1
    except Exception as e:
        print(f"❌ Failed to fetch models: {e}")
        return 1

    conn = sqlite3.connect(DB_PATH)
    provider_table = ensure_provider_table(conn)
    provider_id = get_or_create_bailian_provider(conn, provider_table)

    # 获取 models 表的列
    model_columns = get_table_columns(conn, "models")
    print(f"models table columns: {model_columns}")

    # 根据实际列名动态构建插入语句
    # 我们期望的映射：模型ID -> 可能的列名
    id_col_candidates = ['model_id', 'id', 'name']
    id_col = None
    for col in id_col_candidates:
        if col in model_columns:
            id_col = col
            break
    if not id_col:
        print("❌ models table missing id column")
        return 1

    # 其他可用列
    insert_cols = [id_col]
    if 'platform' in model_columns:
        insert_cols.append('platform')
    if 'provider_id' in model_columns:
        insert_cols.append('provider_id')
    if 'enabled' in model_columns:
        insert_cols.append('enabled')
    if 'created_at' in model_columns:
        insert_cols.append('created_at')
    # 如果有 display_name 列，可以用作显示名称
    if 'display_name' in model_columns:
        insert_cols.append('display_name')

    if not insert_cols:
        print("❌ No columns to insert")
        return 1

    placeholders = ','.join(['?' for _ in insert_cols])
    cursor = conn.cursor()
    added = 0
    for model_id in models:
        # 检查是否存在
        cursor.execute(f"SELECT 1 FROM models WHERE {id_col} = ?", (model_id,))
        if cursor.fetchone():
            print(f"⏭️ Already exists: {model_id}")
            continue

        values = []
        for col in insert_cols:
            if col == id_col:
                values.append(model_id)
            elif col == 'platform':
                values.append('bailian')
            elif col == 'provider_id':
                values.append(provider_id)
            elif col == 'enabled':
                values.append(1)
            elif col == 'created_at':
                values.append(datetime.utcnow().isoformat())
            elif col == 'display_name':
                values.append(model_id)  # 或更友好的名称
        sql = f"INSERT INTO models ({','.join(insert_cols)}) VALUES ({placeholders})"
        cursor.execute(sql, values)
        added += 1
        print(f"➕ Added model: {model_id}")

    conn.commit()
    conn.close()
    print(f"✅ Sync completed. New models added: {added}")

    if added > 0:
        print("Triggering immediate backup to persist new models...")
        backup_script = "/app/scripts/backup_to_dataset.py"
        if os.path.exists(backup_script):
            subprocess.run(["python3", backup_script])
        else:
            print("Backup script not found, skipping immediate backup.")

    return 0

if __name__ == "__main__":
    sys.exit(main())