main.py

import json
from pathlib import Path

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from openai import OpenAI
from pydantic import BaseModel, ConfigDict, Field, field_validator

BASE_DIR = Path(__file__).resolve().parent

app = FastAPI(title="AI 科技树引擎 API V6.0")


app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


class TreeNode(BaseModel):
    id: str = Field(..., min_length=1, max_length=128)
    name: str = Field(..., min_length=1, max_length=200)
    mastery: int = Field(..., ge=0, le=100)

    @field_validator("id", "name", mode="before")
    @classmethod
    def strip_str(cls, v):
        if isinstance(v, str):
            return v.strip()
        return v


class TreeEdge(BaseModel):
    source: str = Field(..., min_length=1, max_length=128)
    target: str = Field(..., min_length=1, max_length=128)

    @field_validator("source", "target", mode="before")
    @classmethod
    def strip_str(cls, v):
        if isinstance(v, str):
            return v.strip()
        return v


class TreePayload(BaseModel):
    nodes: list[TreeNode] = Field(..., min_length=1)
    edges: list[TreeEdge] = Field(default_factory=list)


class LlmConfigBody(BaseModel):
    model_config = ConfigDict(populate_by_name=True)

    api_key: str = Field(..., min_length=1, max_length=2048)
    base_url: str = Field(..., min_length=1, max_length=512)
    llm_model: str = Field(..., min_length=1, max_length=128, alias="model")

    @field_validator("api_key", "base_url", "llm_model", mode="before")
    @classmethod
    def strip_fields(cls, v):
        if isinstance(v, str):
            return v.strip()
        return v


class GenerateTreeRequest(LlmConfigBody):
    text: str = Field(..., min_length=1, max_length=12000)


class ExpandNodeRequest(LlmConfigBody):
    node_id: str = Field(..., min_length=1, max_length=128)
    node_name: str = Field(..., min_length=1, max_length=200)

    @field_validator("node_id", "node_name", mode="before")
    @classmethod
    def strip_fields(cls, v):
        if isinstance(v, str):
            return v.strip()
        return v


def clean_json_string(raw: str) -> str:
    raw = raw.strip()
    if raw.startswith("```"):
        raw = raw.split("\n", 1)[-1]
        if raw.rstrip().endswith("```"):
            raw = raw.rstrip().rsplit("\n", 1)[0]
    return raw.strip()


def make_client(api_key: str, base_url: str) -> OpenAI:
    url = base_url.rstrip("/")
    return OpenAI(api_key=api_key, base_url=url)


def parse_and_validate_tree(content: str) -> dict:
    try:
        raw = json.loads(clean_json_string(content))
    except json.JSONDecodeError as e:
        raise HTTPException(
            status_code=502,
            detail={"message": "模型返回不是合法 JSON", "error": str(e)},
        ) from e
    try:
        payload = TreePayload.model_validate(raw)
    except Exception as e:
        raise HTTPException(
            status_code=502,
            detail={"message": "JSON 结构不符合技能树约定", "error": str(e)},
        ) from e
    return payload.model_dump()


@app.get("/health")
async def health():
    return {"status": "ok"}


@app.get("/")
async def serve_index():
    index = BASE_DIR / "index.html"
    if not index.is_file():
        raise HTTPException(status_code=404, detail="index.html 不存在")
    return FileResponse(index)


@app.post("/generate_tree")
async def generate_tree(req: GenerateTreeRequest):
    sys_prompt = """你是一个专业的技能树构建AI。你的任务是分析用户的学习经历和目标，提取出相关的技能节点，并构建它们之间的前置/后续关系（有向无环图）。
你必须且只能返回纯 JSON 格式的数据，不要包含任何 Markdown 标记（如 ```json）、解释说明或多余的废话。

JSON 结构必须严格如下：
{
  "nodes": [
    {"id": "唯一的英文ID", "name": "技能中文名", "mastery": 掌握度(0-100的整数)}
  ],
  "edges": [
    {"source": "前置技能的ID", "target": "后续技能的ID"}
  ]
}

规则：
1. mastery (掌握度)：根据用户描述推断。如果已经学完/熟练掌握，给 80-100；正在学给 40-70；仅仅是未来目标给 0-20。
2. edges：必须符合逻辑。比如“微积分”通常是“信号与系统”的 source。
"""
    client = make_client(req.api_key, req.base_url)
    try:
        response = client.chat.completions.create(
            model=req.llm_model,
            messages=[
                {"role": "system", "content": sys_prompt},
                {"role": "user", "content": req.text},
            ],
            temperature=0.1,
        )
        content = response.choices[0].message.content or ""
        data = parse_and_validate_tree(content)
        return JSONResponse(status_code=200, content={"status": "success", "data": data})
    except HTTPException:
        raise
    except Exception as e:
        return JSONResponse(
            status_code=502,
            content={"status": "error", "message": str(e)},
        )


@app.post("/expand_node")
async def expand_node(req: ExpandNodeRequest):
    nname = json.dumps(req.node_name, ensure_ascii=False)
    nid_lit = json.dumps(req.node_id, ensure_ascii=False)
    sys_prompt = f"""用户目前正在学习或已经掌握了技能：{nname}。
当前节点 id 为 {nid_lit}（JSON 字符串形式，edges 里 source 字段必须等于去掉引号后的该 id，与现有图一致）。
请推断 2 到 3 个逻辑上最紧密的进阶技能或衍生方向。
返回严格的 JSON 格式。新节点的 mastery 默认设为 10（刚起步）。
必须包含 edges：每条边的 source 必须等于上述节点 id，target 为新节点 id。
格式如下：
{{
  "nodes": [
    {{"id": "唯一的英文ID", "name": "新技能中文名", "mastery": 10}}
  ],
  "edges": [
    {{"source": "与当前节点 id 完全一致", "target": "唯一的英文ID"}}
  ]
}}
"""
    client = make_client(req.api_key, req.base_url)
    try:
        response = client.chat.completions.create(
            model=req.llm_model,
            messages=[{"role": "user", "content": sys_prompt}],
            temperature=0.6,
        )
        content = response.choices[0].message.content or ""
        data = parse_and_validate_tree(content)
        for e in data["edges"]:
            if e["source"] != req.node_id:
                e["source"] = req.node_id
        return JSONResponse(status_code=200, content={"status": "success", "data": data})
    except HTTPException:
        raise
    except Exception as e:
        return JSONResponse(
            status_code=502,
            content={"status": "error", "message": str(e)},
        )


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=7860)