Update scripts/sync_bailian_models.py

scripts/sync_bailian_models.py

@@ -4,27 +4,45 @@ import sys
 import sqlite3
 import json
 import urllib.request
-import urllib.error
 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_name):
+    """返回表的列名列表"""
+    cursor = conn.execute(f"PRAGMA table_info({table_name})")
+    return [row[1] for row in cursor.fetchall()]
+
+def get_or_create_bailian_provider(conn):
+    """获取或创建百炼 provider"""
+    cursor = conn.cursor()
+    # 尝试查找已有的百炼 provider（根据 base_url 包含 dashscope）
+    cursor.execute("SELECT id FROM providers WHERE base_url LIKE '%dashscope.aliyuncs.com%' LIMIT 1")
+    row = cursor.fetchone()
+    if row:
+        return row[0]
+    # 如果没有，创建一个新的
+    cursor.execute("INSERT INTO providers (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
 
-    # 请求百炼 API
+    # 获取模型列表
     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_list = [item["id"] for item in data.get("data", [])]
-            if not models_list:
-                print("⚠️ No models found from API")
+            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}")
@@ -33,27 +51,54 @@ def main():
     # 连接数据库
     try:
         conn = sqlite3.connect(DB_PATH)
-        cursor = conn.cursor()
+        conn.row_factory = sqlite3.Row
     except Exception as e:
         print(f"❌ Cannot connect to database: {e}")
         return 1
 
-    added_count = 0
-    for model_id in models_list:
+    # 获取 provider_id
+    provider_id = get_or_create_bailian_provider(conn)
+    print(f"Using provider_id: {provider_id}")
+
+    # 获取 models 表列名
+    columns = get_table_columns(conn, "models")
+    print(f"models table columns: {columns}")
+
+    # 确定插入时使用的列
+    insert_cols = [col for col in ['id', 'provider_id', 'enabled', 'created_at', 'name'] if col in columns]
+    # 如果 'name' 列存在，则 name 使用 model_id；否则忽略
+    placeholders = ','.join(['?' for _ in insert_cols])
+
+    cursor = conn.cursor()
+    added = 0
+    for model_id in models:
+        # 检查是否已存在
         cursor.execute("SELECT 1 FROM models WHERE id = ?", (model_id,))
-        if not cursor.fetchone():
-            # 根据你的数据库实际结构调整 provider_id（可能需要查询 provider 表）
-            # 这里假设百炼渠道的 provider_id 为 1，请根据实际情况修改
-            cursor.execute("INSERT INTO models (id, name, provider_id, enabled, created_at) VALUES (?, ?, 1, 1, ?)",
-                           (model_id, model_id, datetime.utcnow().isoformat()))
-            added_count += 1
-            print(f"➕ Added model: {model_id}")
-        else:
+        if cursor.fetchone():
             print(f"⏭️ Already exists: {model_id}")
+            continue
+
+        # 准备插入数据
+        values = []
+        for col in insert_cols:
+            if col == 'id':
+                values.append(model_id)
+            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 == 'name':
+                values.append(model_id)  # name 与 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_count}")
+    print(f"✅ Sync completed. New models added: {added}")
     return 0
 
 if __name__ == "__main__":