feat：飞书文档补全
This commit is contained in:
Evilom
@@ -0,0 +1,197 @@
|
||||
"""修复 wiki 解析 + 重拉失败的文档"""
|
||||
import json, os, subprocess, sys, time
|
||||
from datetime import datetime
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from paths import LINKS_DIR, DOCS_DIR, OUTPUT_DIR
|
||||
|
||||
LARK = r'C:\Users\admin\AppData\Roaming\npm\node_modules\@larksuite\cli\bin\lark-cli.exe'
|
||||
IMAGES_DIR = os.path.join(OUTPUT_DIR, 'feishu-images')
|
||||
|
||||
def lark_api(path, params=None, timeout=30):
|
||||
cmd = [LARK, 'api', 'GET', path, '--format', 'json']
|
||||
if params: cmd += ['--params', json.dumps(params)]
|
||||
r = subprocess.run(cmd, capture_output=True, timeout=timeout)
|
||||
out = r.stdout.decode('utf-8', errors='replace')
|
||||
try: return json.loads(out)
|
||||
except: return None
|
||||
|
||||
def download_media(token, local_path, timeout=30):
|
||||
if os.path.isfile(local_path): return True
|
||||
os.makedirs(os.path.dirname(local_path), exist_ok=True)
|
||||
filename = os.path.basename(local_path)
|
||||
try:
|
||||
r = subprocess.run(
|
||||
[LARK, 'api', 'GET', f'/open-apis/drive/v1/medias/{token}/download', '-o', filename],
|
||||
capture_output=True, timeout=timeout, cwd=os.path.dirname(local_path)
|
||||
)
|
||||
return r.returncode == 0 and os.path.isfile(local_path)
|
||||
except: return False
|
||||
|
||||
def resolve_wiki(node_id):
|
||||
data = lark_api('/open-apis/wiki/v2/spaces/get_node', params={'token': node_id})
|
||||
if data and data.get('code') == 0:
|
||||
node = data.get('data', {}).get('node', {})
|
||||
return node.get('obj_token', ''), node.get('title', '')
|
||||
return '', ''
|
||||
|
||||
def extract_text(text_obj):
|
||||
if not text_obj: return ''
|
||||
parts = []
|
||||
for elem in text_obj.get('elements', []):
|
||||
if 'text_run' in elem:
|
||||
content = elem['text_run'].get('content', '')
|
||||
style = elem['text_run'].get('text_element_style', {})
|
||||
if style.get('bold'): content = f'**{content}**'
|
||||
if style.get('italic'): content = f'*{content}*'
|
||||
parts.append(content)
|
||||
elif 'mention_user' in elem:
|
||||
parts.append(f'@{elem["mention_user"].get("user_id", "?")}')
|
||||
return ''.join(parts).strip()
|
||||
|
||||
def blocks_to_markdown(blocks, doc_id, doc_url):
|
||||
lines = []
|
||||
image_map = {}
|
||||
block_map = {b['block_id']: b for b in blocks}
|
||||
root_blocks = [b for b in blocks if b.get('parent_id') == doc_id or not b.get('parent_id')]
|
||||
|
||||
def process(b, depth=0):
|
||||
bt = b.get('block_type', 0)
|
||||
children = b.get('children', [])
|
||||
|
||||
if bt == 2:
|
||||
t = extract_text(b.get('text', {}))
|
||||
if t: lines.append(t)
|
||||
elif 3 <= bt <= 9:
|
||||
t = extract_text(b.get(f'heading{bt-2}', {}) or b.get('text', {}))
|
||||
if t: lines.append(f'{"#"*(bt-2)} {t}')
|
||||
elif bt == 10:
|
||||
t = extract_text(b.get('bullet', {}) or b.get('text', {}))
|
||||
if t: lines.append(f'{" "*depth}- {t}')
|
||||
elif bt == 11:
|
||||
t = extract_text(b.get('ordered', {}) or b.get('text', {}))
|
||||
if t: lines.append(f'{" "*depth}1. {t}')
|
||||
elif bt == 12:
|
||||
if 'bullet' in b:
|
||||
t = extract_text(b.get('bullet', {}))
|
||||
if t: lines.append(f'{" "*depth}- {t}')
|
||||
else:
|
||||
t = extract_text(b.get('code', {}) or b.get('text', {}))
|
||||
if t: lines.append(f'```\n{t}\n```')
|
||||
elif bt == 13:
|
||||
t = extract_text(b.get('quote', {}) or b.get('text', {}))
|
||||
if t: lines.append(f'> {t}')
|
||||
elif bt == 14:
|
||||
t = extract_text(b.get('todo', {}) or b.get('text', {}))
|
||||
if t: lines.append(f'- [ ] {t}')
|
||||
elif bt == 19:
|
||||
if 'callout' in b:
|
||||
emoji = b.get('callout', {}).get('emoji_id', '')
|
||||
if emoji: lines.append(f'\n> :{emoji}:')
|
||||
else: lines.append('---')
|
||||
elif bt == 27:
|
||||
img = b.get('image', {})
|
||||
token = img.get('token', '')
|
||||
if token:
|
||||
local_name = f'{doc_id}_{token[:12]}.png'
|
||||
local_path = os.path.join(IMAGES_DIR, local_name)
|
||||
image_map[token] = local_path
|
||||
lines.append(f'')
|
||||
elif bt == 30:
|
||||
sheet_token = b.get('sheet', {}).get('token', '')
|
||||
if sheet_token: lines.append(f'\n> [嵌入电子表格: {sheet_token[:20]}...]\n')
|
||||
elif bt == 28:
|
||||
lines.append('\n| (嵌入表格) |')
|
||||
elif bt in (21, 22, 24, 25, 31, 34): pass
|
||||
else:
|
||||
for key in ['text', 'bullet', 'heading1', 'heading2', 'heading3', 'quote', 'code', 'todo']:
|
||||
if key in b:
|
||||
t = extract_text(b[key])
|
||||
if t: lines.append(t)
|
||||
break
|
||||
|
||||
for cid in children:
|
||||
if cid in block_map: process(block_map[cid], depth + 1)
|
||||
|
||||
for b in root_blocks: process(b, depth=0)
|
||||
header = [f'# {doc_id}\n', f'Source: {doc_url}\n', '']
|
||||
return '\n'.join(header + lines), image_map
|
||||
|
||||
# 加载已有索引
|
||||
content_path = os.path.join(LINKS_DIR, 'all_feishu_content.json')
|
||||
with open(content_path, 'r', encoding='utf-8') as f:
|
||||
content_data = json.load(f)
|
||||
existing = {d['doc_id']: d for d in content_data.get('documents', [])}
|
||||
|
||||
# 找出需要重试的文档（fetched=False 或不存在的）
|
||||
links_path = os.path.join(LINKS_DIR, 'all_feishu_links.json')
|
||||
with open(links_path, 'r', encoding='utf-8') as f:
|
||||
all_links = json.load(f).get('links', [])
|
||||
|
||||
retry = []
|
||||
for link in all_links:
|
||||
did = link.get('doc_id', '')
|
||||
if not did: continue
|
||||
old = existing.get(did, {})
|
||||
if not old.get('fetched', False) or old.get('content_length', 0) <= 200:
|
||||
retry.append(link)
|
||||
|
||||
print(f'需要重试: {len(retry)} 篇\n')
|
||||
|
||||
new = err = 0
|
||||
for i, link in enumerate(retry):
|
||||
doc_id = link.get('doc_id', '')
|
||||
doc_type = link.get('type', 'docx')
|
||||
url = link.get('url', '')
|
||||
sender = link.get('sender', '?')
|
||||
|
||||
print(f'[{i+1}/{len(retry)}] {doc_id} ({sender})...', end=' ', flush=True)
|
||||
|
||||
real_id = doc_id
|
||||
title = ''
|
||||
if doc_type == 'wiki':
|
||||
real_id, wiki_title = resolve_wiki(doc_id)
|
||||
if not real_id:
|
||||
print('wiki FAIL'); err += 1; continue
|
||||
title = wiki_title
|
||||
|
||||
data = lark_api(f'/open-apis/docx/v1/documents/{real_id}/blocks')
|
||||
if not data or data.get('code') != 0:
|
||||
print('API FAIL'); err += 1; continue
|
||||
|
||||
blocks = data.get('data', {}).get('items', [])
|
||||
if not blocks:
|
||||
print('空'); err += 1; continue
|
||||
|
||||
md, image_map = blocks_to_markdown(blocks, doc_id, url)
|
||||
|
||||
downloaded = 0
|
||||
for token, local_path in image_map.items():
|
||||
if download_media(token, local_path): downloaded += 1
|
||||
time.sleep(0.1)
|
||||
|
||||
with open(os.path.join(DOCS_DIR, f'{doc_id}.md'), 'w', encoding='utf-8') as f: f.write(md)
|
||||
|
||||
if not title:
|
||||
for line in md.split('\n'):
|
||||
if line.startswith('# ') and len(line) > 2: title = line[2:].strip(); break
|
||||
|
||||
existing[doc_id] = {
|
||||
'url': url, 'doc_id': doc_id, 'title': title,
|
||||
'sender': sender, 'group': link.get('group', ''), 'time': link.get('time', ''),
|
||||
'content_length': len(md), 'image_blocks': len(image_map),
|
||||
'images_downloaded': downloaded, 'fetched': True,
|
||||
'fetched_at': datetime.now().isoformat(),
|
||||
}
|
||||
new += 1
|
||||
print(f'OK ({len(md)}c, {len(image_map)}img, {downloaded}dl)')
|
||||
time.sleep(0.15)
|
||||
|
||||
# 保存索引
|
||||
content_data['documents'] = list(existing.values())
|
||||
content_data['total'] = len(existing)
|
||||
with open(content_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(content_data, f, ensure_ascii=False, indent=2)
|
||||
|
||||
imgs = len([f for f in os.listdir(IMAGES_DIR) if f.endswith('.png')]) if os.path.isdir(IMAGES_DIR) else 0
|
||||
print(f'\n完成! 新拉取:{new} 失败:{err} 总计:{len(existing)} 图片:{imgs}')
|
||||
Reference in New Issue
Block a user