kg_engine.py

"""知识抽取 & 图谱管理引擎"""

from models import Concept, Relation
from state import AppState
from llm_client import chat_json

_EXTRACT_SYSTEM = """你是一个专业的知识图谱构建助手。
你的任务是从给定的学习笔记文本中，抽取出核心概念以及它们之间的关系。

对于每个概念，你需要提供：
- name: 概念名称（简洁，3-8字）
- description: 一句话解释这个概念
- category: 分类标签（如：统计学、机器学习、Python、数据库、业务分析等）
- difficulty: 难度1-5（1=入门概念，5=高级/数学密集型概念）

对于每组关系，你需要提供：
- src: 源概念名称
- dst: 目标概念名称
- rel_type: 关系类型，必须是以下之一：
  * prerequisite: 学src之前必须先学dst（dst是src的前置知识）
  * subclass: src是dst的一种/子类
  * application: dst是src的应用场景
  * comparison: src与dst是相对/对比概念
  * related: 一般相关
- description: 关系的一句话说明

注意：
1. 只抽取文本中明确提到的概念，不要臆造
2. 每个概念至少关联一个关系
3. 关系数量不要超过概念数量的2倍
4. 同名概念只出现一次"""


def extract_concepts_from_chunk(text: str, source_file: str = "") -> tuple[list[Concept], list[Relation]]:
    """从一段文本中抽取概念和关系"""
    prompt = f"""请从以下学习笔记中抽取概念和关系：

```
{text[:4000]}
```

返回JSON格式：
{{
  "concepts": [
    {{"name": "概念名", "description": "一句话解释", "category": "分类", "difficulty": 3}}
  ],
  "relations": [
    {{"src": "概念A", "dst": "概念B", "rel_type": "prerequisite", "description": "关系说明"}}
  ]
}}"""

    try:
        result = chat_json(prompt, _EXTRACT_SYSTEM, temperature=0.2)
    except Exception as e:
        print(f"[kg_engine] LLM 调用失败: {e}")
        return [], []

    state = AppState()
    concepts = []
    relations = []

    # 先创建所有概念
    name_to_id = {}
    for item in result.get("concepts", []):
        c = Concept(
            name=item.get("name", "").strip(),
            description=item.get("description", "").strip(),
            category=item.get("category", "未分类").strip(),
            source_file=source_file,
            difficulty=min(5, max(1, int(item.get("difficulty", 3)))),
        )
        if c.name:
            concepts.append(c)
            name_to_id[c.name] = c.id

    # 再创建关系
    for item in result.get("relations", []):
        src_name = item.get("src", "").strip()
        dst_name = item.get("dst", "").strip()
        if src_name in name_to_id and dst_name in name_to_id:
            r = Relation(
                src_id=name_to_id[src_name],
                dst_id=name_to_id[dst_name],
                rel_type=item.get("rel_type", "related").strip(),
                description=item.get("description", "").strip(),
            )
            relations.append(r)

    return concepts, relations


def build_graph_from_text(text: str, source_file: str = "", chunk_size: int = 3000) -> int:
    """从全文构建知识图谱，返回抽取的概念总数"""
    state = AppState()

    # 简单分块：按双换行分段落，合并短段落
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
    chunks = []
    current = ""
    for p in paragraphs:
        if len(current) + len(p) > chunk_size and current:
            chunks.append(current)
            current = p
        else:
            current = current + "\n" + p if current else p
    if current:
        chunks.append(current)

    total_concepts = 0
    for i, chunk in enumerate(chunks):
        concepts, relations = extract_concepts_from_chunk(chunk, source_file)
        for c in concepts:
            state.add_concept(c)
            total_concepts += 1
        for r in relations:
            state.add_relation(r)

    state.save()
    return total_concepts


def build_graph_from_file(file_path: str) -> dict:
    """从文件构建知识图谱，返回统计信息"""
    import os
    source_file = os.path.basename(file_path)

    if file_path.endswith('.md') or file_path.endswith('.txt'):
        with open(file_path, 'r', encoding='utf-8') as f:
            text = f.read()
    elif file_path.endswith('.pdf'):
        try:
            from pypdf import PdfReader
            reader = PdfReader(file_path)
            text = "\n\n".join(page.extract_text() or "" for page in reader.pages)
        except ImportError:
            return {"error": "需要安装 pypdf: pip install pypdf"}
    else:
        return {"error": f"不支持的文件格式: {file_path}"}

    if not text.strip():
        return {"error": "文件内容为空"}

    total = build_graph_from_text(text, source_file)
    state = AppState()
    return {
        "concepts": len(state.concepts),
        "relations": len(state.relations),
        "new_concepts": total,
    }