LWX/ai_teaching_system/scripts/generate_question_variants.py

# -*- coding: utf-8 -*-
"""
自动生成 question_variants：
- 依据 question_knowledge_points 的知识点映射，找“同知识点”的题作为变式题
- 为每道 base_question 最多生成 N 道 variant_question（默认 3）
- 使用 question_variants(base_question_id, variant_question_id) 的唯一约束去重

注意：
- 如果 question_knowledge_points 为空/不完整，会先把 questions.knowledge_point 同步补齐到 question_knowledge_points
  （避免因为历史未触发 lazy 建映射导致 variants 无法生成）
"""

import argparse
import random
from collections import defaultdict

import os
import sys

from sqlalchemy import and_

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from backend.config import get_settings
from backend.database import SessionLocal
from backend.models import KnowledgePoint, Question, QuestionKnowledgePoint, QuestionVariant


def _ensure_knowledge_points(session) -> None:
    """确保 knowledge_points 表至少包含 questions.knowledge_point 中出现的知识点名。"""
    kp_names = (
        session.query(Question.knowledge_point)
        .filter(and_(Question.knowledge_point.isnot(None), Question.knowledge_point != ""))
        .distinct()
        .all()
    )
    names = [x[0] for x in kp_names if x and x[0]]
    if not names:
        return

    existing = {
        n[0] for n in session.query(KnowledgePoint.name).filter(KnowledgePoint.name.in_(names)).all()
    }
    to_create = [n for n in names if n not in existing]
    for name in to_create:
        session.add(KnowledgePoint(name=name))
    if to_create:
        session.commit()


def _ensure_question_knowledge_points(session) -> None:
    """
    若 question_knowledge_points 缺少 (question_id, knowledge_point_id) 组合，
    则从 Question.knowledge_point 同步补齐 weight=1 映射。
    """
    # 预取 existing pair（question_id, knowledge_point_id）
    existing_pairs = set(
        session.query(QuestionKnowledgePoint.question_id, QuestionKnowledgePoint.knowledge_point_id).all()
    )

    # 分批拉取 questions，逐条补齐
    questions = (
        session.query(Question.id, Question.knowledge_point)
        .filter(and_(Question.knowledge_point.isnot(None), Question.knowledge_point != ""))
        .all()
    )
    if not questions:
        return

    # 缓存 knowledge_point name -> id
    kp_rows = session.query(KnowledgePoint.id, KnowledgePoint.name).all()
    kp_name_to_id = {name: kid for kid, name in [(r[0], r[1]) for r in kp_rows]}

    created = 0
    for qid, kp_name in questions:
        kpid = kp_name_to_id.get(kp_name)
        if not kpid:
            continue
        if (qid, kpid) in existing_pairs:
            continue
        session.add(
            QuestionKnowledgePoint(
                question_id=qid,
                knowledge_point_id=kpid,
                weight=1,
            )
        )
        existing_pairs.add((qid, kpid))
        created += 1
    if created:
        session.commit()


def generate_variants(
    session,
    max_variants_per_base: int = 3,
    max_total_bases: int | None = None,
    seed: int = 42,
) -> int:
    """
    生成 question_variants，并写入数据库。
    返回：插入的 rows 数
    """
    random.seed(seed)

    # 1) 拉取 qkp 映射
    qkps = (
        session.query(
            QuestionKnowledgePoint.question_id,
            QuestionKnowledgePoint.knowledge_point_id,
            QuestionKnowledgePoint.weight,
        )
        .all()
    )
    if not qkps:
        print("[WARN] question_knowledge_points 为空，无法生成 variants。请先补齐映射。")
        return 0

    question_kps: dict[int, list[tuple[int, int]]] = defaultdict(list)
    kp_q_weight: dict[int, dict[int, int]] = defaultdict(dict)

    for qid, kpid, w in qkps:
        if not kpid or not qid:
            continue
        w_int = int(w) if w is not None else 1
        question_kps[int(qid)].append((int(kpid), w_int))
        # 同一 (kp, q) 可能重复映射时取最大权重
        prev = kp_q_weight[int(kpid)].get(int(qid))
        if prev is None or w_int > prev:
            kp_q_weight[int(kpid)][int(qid)] = w_int

    base_qids = list(question_kps.keys())
    if not base_qids:
        return 0
    if max_total_bases is not None:
        base_qids = base_qids[: max_total_bases]

    # 2) 预取 existing variants 去重
    existing_pairs = set(
        session.query(QuestionVariant.base_question_id, QuestionVariant.variant_question_id).all()
    )

    # 3) 为每道 base_question 生成 candidates
    to_insert: list[QuestionVariant] = []
    inserted = 0

    # 增加一个轻量的随机打散，避免完全相同 overlap score 时固定选择
    random.shuffle(base_qids)

    for i, base_qid in enumerate(base_qids, start=1):
        base_kps = question_kps.get(base_qid) or []
        if len(base_kps) == 0:
            continue

        candidate_scores: dict[int, int] = defaultdict(int)
        candidate_best_kp: dict[int, int] = {}
        candidate_best_score: dict[int, int] = {}

        for kpid, w_base in base_kps:
            cand_map = kp_q_weight.get(kpid) or {}
            for cand_qid, w_cand in cand_map.items():
                if cand_qid == base_qid:
                    continue
                w_base = int(w_base)
                w_cand = int(w_cand)
                contrib = w_base * w_cand
                candidate_scores[cand_qid] += contrib
                if contrib > candidate_best_score.get(cand_qid, 0):
                    candidate_best_score[cand_qid] = contrib
                    candidate_best_kp[cand_qid] = kpid

        if not candidate_scores:
            continue

        # 候选排序：score desc + cand_qid 轻量打散
        sorted_cands = sorted(
            candidate_scores.items(),
            key=lambda kv: (kv[1], random.random()),
            reverse=True,
        )

        variants_for_base = 0
        for cand_qid, score in sorted_cands:
            pair = (base_qid, cand_qid)
            if pair in existing_pairs:
                continue
            vid_kp = candidate_best_kp.get(cand_qid)
            # weight：用相似度 score 作为排序依据；数值太大也没关系（这里只是降序）
            w = int(score) if score is not None else 1

            to_insert.append(
                QuestionVariant(
                    base_question_id=base_qid,
                    variant_question_id=cand_qid,
                    knowledge_point_id=vid_kp,
                    weight=w,
                )
            )
            existing_pairs.add(pair)
            variants_for_base += 1
            inserted += 1
            if variants_for_base >= max_variants_per_base:
                break

        # 每 500 条 flush，避免内存过大
        if inserted > 0 and inserted % 500 == 0:
            session.add_all(to_insert)
            session.commit()
            to_insert.clear()
            print(f"[INFO] 已插入 {inserted} rows（base 已处理到 {i}/{len(base_qids)}）")

    if to_insert:
        session.add_all(to_insert)
        session.commit()

    return inserted


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--max-variants-per-base", type=int, default=3)
    parser.add_argument("--max-total-bases", type=int, default=2000)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    settings = get_settings()
    print(f"[INFO] Using MYSQL_URI: {settings.MYSQL_URI}")

    session = SessionLocal()
    try:
        _ensure_knowledge_points(session)
        _ensure_question_knowledge_points(session)

        inserted = generate_variants(
            session,
            max_variants_per_base=args.max_variants_per_base,
            max_total_bases=args.max_total_bases,
            seed=args.seed,
        )
        print(f"[DONE] question_variants inserted rows: {inserted}")
    finally:
        session.close()


if __name__ == "__main__":
    main()