PlantAI-rag-new/agent/legacy-rag/scripts/utils/kg_sqlite.py

#!/usr/bin/env python3
"""SQLite版知识图谱（临时替代Neo4j）"""

import json
import sqlite3
from pathlib import Path

class GeneFamilyKG_SQLite:
    def __init__(self, db_path="/home/xzq/gene-family-mas-v2/kg.db"):
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path)
        self.init_schema()
        
    def init_schema(self):
        """初始化Schema"""
        cursor = self.conn.cursor()
        
        # 基因家族表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS gene_families (
                name TEXT PRIMARY KEY,
                full_name TEXT,
                description TEXT
            )
        ''')
        
        # 物种表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS species (
                name TEXT PRIMARY KEY,
                common_name TEXT
            )
        ''')
        
        # 基因表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS genes (
                id TEXT PRIMARY KEY,
                name TEXT,
                species TEXT,
                family TEXT,
                FOREIGN KEY (species) REFERENCES species(name),
                FOREIGN KEY (family) REFERENCES gene_families(name)
            )
        ''')
        
        # 系统发育分组表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS phylogeny_groups (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT,
                gene_family TEXT,
                UNIQUE(name, gene_family)
            )
        ''')
        
        # 基因-分组关系表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS gene_groups (
                gene_id TEXT,
                group_id INTEGER,
                FOREIGN KEY (gene_id) REFERENCES genes(id),
                FOREIGN KEY (group_id) REFERENCES phylogeny_groups(id),
                PRIMARY KEY (gene_id, group_id)
            )
        ''')
        
        self.conn.commit()
        print("SQLite Schema初始化完成")
        
    def import_phylogeny_data(self, json_path: Path):
        """导入v5提取数据"""
        with open(json_path) as f:
            data = json.load(f)
            
        gene_family = data.get("gene_family", "Unknown")
        species = data.get("species", "Unknown")
        groups = data.get("groups", {})
        
        cursor = self.conn.cursor()
        
        # 插入基因家族
        cursor.execute(
            "INSERT OR REPLACE INTO gene_families (name) VALUES (?)",
            (gene_family,)
        )
        
        # 插入物种
        cursor.execute(
            "INSERT OR REPLACE INTO species (name) VALUES (?)",
            (species,)
        )
        
        # 插入分组和基因
        for group_name, genes in groups.items():
            # 插入分组
            cursor.execute(
                "INSERT OR REPLACE INTO phylogeny_groups (name, gene_family) VALUES (?, ?)",
                (group_name, gene_family)
            )
            cursor.execute(
                "SELECT id FROM phylogeny_groups WHERE name=? AND gene_family=?",
                (group_name, gene_family)
            )
            group_id = cursor.fetchone()[0]
            
            # 插入基因并建立关系
            for gene_id in genes:
                cursor.execute(
                    "INSERT OR REPLACE INTO genes (id, name, species, family) VALUES (?, ?, ?, ?)",
                    (gene_id, gene_id, species, gene_family)
                )
                cursor.execute(
                    "INSERT OR REPLACE INTO gene_groups (gene_id, group_id) VALUES (?, ?)",
                    (gene_id, group_id)
                )
        
        self.conn.commit()
        print(f"导入完成: {gene_family}")
        print(f"  - 分组数: {len(groups)}")
        print(f"  - 基因数: {sum(len(g) for g in groups.values())}")
        
    def get_stats(self):
        """获取统计信息"""
        cursor = self.conn.cursor()
        stats = {}
        for table in ['gene_families', 'species', 'genes', 'phylogeny_groups']:
            cursor.execute(f"SELECT COUNT(*) FROM {table}")
            stats[table] = cursor.fetchone()[0]
        return stats
        
    def close(self):
        self.conn.close()

if __name__ == "__main__":
    kg = GeneFamilyKG_SQLite()
    
    # 导入番茄P450数据
    data_path = Path("/home/xzq/gene-family-mas-v2/evaluation/auto_results/番茄p450基因家族BMCgenomics2020_v5.json")
    if data_path.exists():
        kg.import_phylogeny_data(data_path)
        print("\n统计信息:", kg.get_stats())
    else:
        print(f"数据文件不存在: {data_path}")
    
    kg.close()