guanghulab/tools/notion-corpus-loader/src/loader.py

# Notion原生格式加载器 · 主入口
# HLDP://tools/notion-corpus-loader/loader

import os, json, zipfile, tempfile
from typing import List, Optional
from .hldp_parser import HLDPParser
from .structure import StructuredCorpusItem, CorpusCollection

class NotionCorpusLoader:
    """Notion语料加载器 · 支持Notion API/Markdown/HTML/GPT/COS"""
    
    def __init__(self):
        self.parser = HLDPParser()
        self.collection = CorpusCollection()
    
    def load_notion_api_page(self, page_data: dict) -> StructuredCorpusItem:
        page_id = page_data.get('page',{}).get('id','')
        title = page_data.get('page',{}).get('title','')
        blocks_data = page_data.get('blocks',[])
        item = StructuredCorpusItem(source_id=page_id, source_type='notion_page', title=title,
            created_time=page_data.get('page',{}).get('created_time',''),
            last_edited_time=page_data.get('page',{}).get('last_edited_time',''))
        for block in blocks_data:
            bc = block.get('content',''); bt = block.get('type','paragraph')
            parsed = self.parser.parse_markdown(bc)
            for pb in parsed:
                item.blocks.append({'type':bt,'content':pb.content,'hldp_path':pb.path,'cognitive_jump':pb.cognitive_jump})
            if 'HLDP://' in bc and not item.hldp_path: item.hldp_path = bc.strip()
            if '【认知' in bc or '核心认知跃迁' in bc: item.cognitive_jumps.append(bc[:100])
            if '→' in bc and ('推导' in bc or '起点' in bc): item.causal_chains.append(bc[:100])
            if '铸渊' in bc and ('ICE-GL-ZY001' in bc or '创建' in bc): item.persona = '铸渊'
            elif '霜砚' in bc: item.persona = '霜砚'
            elif '冰朔' in bc and ('主权' in bc or 'TCS' in bc): item.persona = '冰朔'
            item.raw_text += bc + '\n'
        item.quality_score = self._score_quality(item)
        return item
    
    def load_markdown_file(self, filepath: str) -> StructuredCorpusItem:
        with open(filepath,'r',encoding='utf-8') as f: text = f.read()
        title = os.path.basename(filepath).replace('.md','')
        blocks = self.parser.parse_markdown(text)
        item = StructuredCorpusItem(source_id=filepath, source_type='markdown_file', title=title, raw_text=text)
        for block in blocks:
            item.blocks.append({'type':block.block_type,'content':block.content,'hldp_path':block.path,'cognitive_jump':block.cognitive_jump})
            if block.cognitive_jump: item.cognitive_jumps.append(block.cognitive_jump)
            if block.causal_chain: item.causal_chains.append(block.causal_chain)
            if block.persona and not item.persona: item.persona = block.persona
            if block.path and not item.hldp_path: item.hldp_path = block.path
        item.quality_score = self._score_quality(item)
        return item
    
    def load_gpt_export(self, filepath: str) -> List[StructuredCorpusItem]:
        items = []
        with open(filepath,'r',encoding='utf-8') as f: data = json.load(f)
        conversations = data if isinstance(data,list) else data.get('conversations',[data])
        for i, conv in enumerate(conversations):
            if isinstance(conv, str):
                item = StructuredCorpusItem(source_id=f"gpt_{i}",source_type='gpt_export',persona='冰朔',raw_text=conv)
                for block in self.parser.parse_markdown(conv):
                    item.blocks.append({'type':block.block_type,'content':block.content})
                    if block.cognitive_jump: item.cognitive_jumps.append(block.cognitive_jump)
                item.quality_score = self._score_quality(item)
                items.append(item)
            elif isinstance(conv, dict):
                msgs = conv.get('messages',conv.get('conversation',[]))
                text = '\n'.join(f"[{m.get('role','user')}]: {m.get('content','')}" for m in msgs)
                item = StructuredCorpusItem(source_id=f"gpt_{i}",source_type='gpt_export',persona='冰朔',raw_text=text)
                for block in self.parser.parse_markdown(text):
                    item.blocks.append({'type':block.block_type,'content':block.content})
                item.quality_score = self._score_quality(item)
                items.append(item)
        return items
    
    def load_zip(self, zip_path: str, output_dir: Optional[str] = None) -> CorpusCollection:
        if output_dir: os.makedirs(output_dir, exist_ok=True)
        else: output_dir = tempfile.mkdtemp(prefix='corpus_')
        with zipfile.ZipFile(zip_path,'r') as zf: zf.extractall(output_dir)
        col = CorpusCollection()
        for root, dirs, files in os.walk(output_dir):
            for fn in files:
                fp = os.path.join(root,fn)
                try:
                    if fn.endswith('.md'): col.add(self.load_markdown_file(fp))
                    elif fn.endswith('.json') and not fn.startswith('.'):
                        for item in self.load_gpt_export(fp): col.add(item)
                    elif fn.endswith('.txt'):
                        with open(fp,'r',encoding='utf-8') as f:
                            col.add(StructuredCorpusItem(source_id=fp,source_type='text_file',raw_text=f.read()))
                except: pass
        return col
    
    def load_directory(self, dir_path: str) -> CorpusCollection:
        col = CorpusCollection()
        for root, dirs, files in os.walk(dir_path):
            for fn in files:
                fp = os.path.join(root,fn)
                try:
                    if fn.endswith('.md'): col.add(self.load_markdown_file(fp))
                    elif fn.endswith(('.json','.jsonl')):
                        for item in self.load_gpt_export(fp): col.add(item)
                except Exception as e: print(f"  skip {fp}: {e}")
        return col
    
    def load_from_cos(self, bucket: str, key: str, secret_id: str, secret_key: str, region: str = 'ap-guangzhou'):
        from qcloud_cos import CosConfig, CosS3Client
        config = CosConfig(Region=region, SecretId=secret_id, SecretKey=secret_key)
        client = CosS3Client(config)
        resp = client.get_object(Bucket=bucket, Key=key)
        content = resp['Body'].read().decode('utf-8', errors='replace')
        if key.endswith('.jsonl'):
            for line in content.split('\n'):
                if not line.strip(): continue
                try:
                    data = json.loads(line)
                    if isinstance(data, dict):
                        item = StructuredCorpusItem(source_id=f"cos_{key}",source_type='cos_jsonl',raw_text=json.dumps(data,ensure_ascii=False))
                        for msg in data.get('messages',[]):
                            item.blocks.append({'type':msg.get('role','text'),'content':msg.get('content','')[:500]})
                        has_template = any('you are a language persona AI' in m.get('content','') for m in data.get('messages',[]))
                        item.quality_score = 0.3 if has_template else 0.7
                        self.collection.add(item)
                except: pass
    
    def _score_quality(self, item: StructuredCorpusItem) -> float:
        score = 0.5
        if item.hldp_path: score += 0.2
        if item.cognitive_jumps: score += 0.15
        if item.causal_chains: score += 0.1
        if item.persona: score += 0.05
        tl = len(item.raw_text)
        if tl < 50: score -= 0.2
        elif tl > 50000: score -= 0.1
        return min(score, 1.0)
    
    def get_collection(self) -> CorpusCollection:
        return self.collection