# -*- coding: utf-8 -*- """ 种子数据执行入口：初始化 SQLite 表结构 → 灌入模拟数据 → 构建 Chroma 向量索引。 幂等设计：重复执行会自动清空重灌，方便数据更新后一键重建。 用法：python3 seed.py [--rebuild-vector] （--rebuild-vector 强制重建向量集合） """ import argparse import logging import sys import time logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") log = logging.getLogger("seed") from db import init_db, execute, executemany, query_one, table_count # noqa: E402 from seed_teams import SPORTS, LEAGUES, TEAMS, STANDINGS # noqa: E402 from seed_players_a import PLAYERS_A # noqa: E402 from seed_players_b import PLAYERS_B # noqa: E402 from seed_games import GAMES, BOX # noqa: E402 from seed_news import NEWS, WIKI # noqa: E402 from seed_persons import PERSONS # noqa: E402 import vector_store # noqa: E402 def wipe(): for t in ("game_player_stats", "games", "standings", "players", "teams", "leagues", "sports", "news", "persons"): execute(f"DELETE FROM {t}") execute(f"DELETE FROM sqlite_sequence WHERE name='{t}'") log.info("已清空旧数据") def seed_core(): """sports / leagues / teams / players / standings / games / persons""" for code, name, name_en in SPORTS: execute("INSERT INTO sports(code,name,name_en) VALUES(?,?,?)", (code, name, name_en)) sport_id = query_one("SELECT id FROM sports WHERE code='basketball'")["id"] for code, name, name_en, country, season in LEAGUES: execute("INSERT INTO leagues(sport_id,code,name,name_en,country,season) VALUES(?,?,?,?,?,?)", (sport_id, code, name, name_en, country, season)) league_id = query_one("SELECT id FROM leagues WHERE code='NBA'")["id"] # 球队 for code, name, en, city, arena, founded, champs, coach, intro in TEAMS: execute("""INSERT INTO teams(league_id,code,name,name_en,city,arena,founded,champion_count,head_coach,intro) VALUES(?,?,?,?,?,?,?,?,?,?)""", (league_id, code, name, en, city, arena, founded, champs, coach, intro)) team_id = {r["code"]: r["id"] for r in __import__("db").query("SELECT id,code FROM teams")} # 球员 p_rows = [] for p in PLAYERS_A + PLAYERS_B: (t, name, en, pos, num, h, w, country, dy, dp, sal, pts, reb, ast, stl, blk, mn, cpts, creb, cast, cg, awards, bio) = p p_rows.append((team_id[t], name, en, pos, num, h, w, country, dy, dp, sal, pts, reb, ast, stl, blk, mn, cpts, creb, cast, cg, awards, bio)) executemany("""INSERT INTO players(team_id,name,name_en,position,number,height_cm,weight_kg,country, draft_year,draft_pick,salary_m,season_pts,season_reb,season_ast,season_stl,season_blk,season_min, career_pts,career_reb,career_ast,career_games,awards,bio) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)""", p_rows) log.info("球员 %d 名", len(p_rows)) # 排名 s_rows = [(league_id, "2025-26", team_id[c], w, l, round(w / (w + l), 4), conf, rk) for c, conf, rk, w, l in STANDINGS] executemany("""INSERT INTO standings(league_id,season,team_id,wins,losses,win_pct,conference,rank) VALUES(?,?,?,?,?,?,?,?)""", s_rows) # 比赛 game_id = {} g_rows = [] for rnd, gt, status, home, away, hs, as_, venue, bcast in GAMES: g_rows.append((league_id, "2025-26", rnd, gt, status, team_id[home], team_id[away], hs, as_, venue, bcast)) executemany("""INSERT INTO games(league_id,season,round_name,game_time,status,home_team_id,away_team_id, home_score,away_score,venue,broadcast) VALUES(?,?,?,?,?,?,?,?,?,?,?)""", g_rows) for i, (rnd, gt, *_rest) in enumerate(GAMES): game_id[(rnd, gt)] = query_one( "SELECT id FROM games WHERE round_name=? AND game_time=?", (rnd, gt))["id"] # 球员技术统计（box score） box_rows = [] for (rnd, gt), lines in BOX.items(): gid = game_id.get((rnd, gt)) if not gid: continue for pname, pts, reb, ast, stl, blk, mn in lines: player = query_one("SELECT id,team_id FROM players WHERE name=?", (pname,)) if player: box_rows.append((gid, player["id"], player["team_id"], pts, reb, ast, stl, blk, mn)) executemany("""INSERT INTO game_player_stats(game_id,player_id,team_id,points,rebounds,assists,steals,blocks,minutes) VALUES(?,?,?,?,?,?,?,?,?)""", box_rows) log.info("比赛 %d 场，技术统计 %d 条", len(GAMES), len(box_rows)) # 人物 per_rows = [] for t, name, en, role, role_cn, title, bio, ach in PERSONS: per_rows.append((sport_id, league_id, team_id.get(t), name, en, role, role_cn, title, bio, ach)) executemany("""INSERT INTO persons(sport_id,league_id,team_id,name,name_en,role,role_cn,title,bio,achievements) VALUES(?,?,?,?,?,?,?,?,?,?)""", per_rows) log.info("人物 %d 名", len(per_rows)) # 新闻 / 百科 n_rows = [] for t, title, content, author, source, pt, tags, kind in NEWS + WIKI: n_rows.append((sport_id, league_id, team_id.get(t), title, content, author, source, pt, tags, kind)) executemany("""INSERT INTO news(sport_id,league_id,team_id,title,content,author,source,publish_time,tags,kind) VALUES(?,?,?,?,?,?,?,?,?,?)""", n_rows) log.info("新闻/百科 %d 篇", len(n_rows)) def seed_vector(rebuild=False): """把新闻/百科分块后写入 Chroma（title + 正文 按 400 字切块，重叠 60 字）""" from db import query as dbq rows = dbq("SELECT id,title,content,kind,team_id,publish_time,tags FROM news ORDER BY id") ids, docs, metas = [], [], [] for r in rows: text = f"{r['title']}\n{r['content']}" chunk_size, overlap = 400, 60 start = 0 while start < len(text): chunk = text[start:start + chunk_size] ids.append(f"news_{r['id']}_{start}") docs.append(chunk) metas.append({"news_id": r["id"], "title": r["title"], "kind": r["kind"], "team_id": r["team_id"] or 0, "publish_time": r["publish_time"], "tags": r["tags"] or ""}) start += chunk_size - overlap log.info("向量文档 %d 块（%d 篇）", len(ids), len(rows)) if rebuild: vector_store.reset_collection() # 若已有数据则跳过（增量） if vector_store.collection_count() >= len(ids): log.info("向量库已有 %d 条，跳过写入", vector_store.collection_count()) return if vector_store.collection_count() > 0: vector_store.reset_collection() for i in range(0, len(ids), 50): # 分批写入，避免单次请求过大 vector_store.add_documents(ids[i:i + 50], docs[i:i + 50], metas[i:i + 50]) log.info("向量索引完成，共 %d 条", vector_store.collection_count()) def main(): ap = argparse.ArgumentParser() ap.add_argument("--rebuild-vector", action="store_true", help="强制重建向量集合") ap.add_argument("--no-vector", action="store_true", help="跳过向量索引") args = ap.parse_args() t0 = time.time() init_db() wipe() seed_core() if not args.no_vector: seed_vector(rebuild=args.rebuild_vector) else: log.info("跳过向量索引") log.info("数据灌入完成，耗时 %.1fs", time.time() - t0) for t in ("sports", "leagues", "teams", "players", "games", "standings", "game_player_stats", "news", "persons"): log.info(" %-16s %d", t, table_count(t)) if __name__ == "__main__": sys.exit(main())