471 lines
26 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import argparse
import base64
import csv
import io
import json
import os
import secrets
import sqlite3
import tempfile
import threading
import urllib.error
import urllib.parse
import urllib.request
import webbrowser
from datetime import datetime, timezone
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from . import store
from .analysis import (SYSTEM_PROMPT, STRATEGIES, SIGNALS, Redactor, extract_events,
local_cards, make_card, prepare_plan, validate_model_cards)
from .importer import import_rows, read_rows, offset_zone
from .mac_connect import MacConnection
STATIC = Path(__file__).parent / 'static'
MAX_UPLOAD = 64 * 1024 * 1024
def now():
return datetime.now(timezone.utc).isoformat()
def validate_api(config):
endpoint = str(config.get('endpoint', '')).strip()
parsed = urllib.parse.urlsplit(endpoint)
if parsed.username or parsed.password or parsed.query or parsed.fragment:
raise ValueError('API 地址不能包含账号、密码、查询参数或片段')
if parsed.scheme != 'https' and not (parsed.scheme == 'http' and parsed.hostname in ('127.0.0.1', 'localhost', '::1')):
raise ValueError('远程模型 API 必须使用 HTTPS,本地模型可使用 localhost HTTP')
if not parsed.hostname or not parsed.path.endswith('/chat/completions'):
raise ValueError('请填写完整的兼容接口地址,以 /chat/completions 结尾')
model = str(config.get('model', '')).strip()
if not model or len(model) > 200:
raise ValueError('请填写模型名称')
return {'endpoint': endpoint, 'model': model}
class NoRedirect(urllib.request.HTTPRedirectHandler):
def redirect_request(self, *args, **kwargs):
raise ValueError('API 重定向已拒绝,避免脱敏文本和密钥被转发到其他地址')
def call_model(config, key, body):
if key and (not key.isascii() or any(ord(char) < 33 or ord(char) > 126 for char in key)):
raise ValueError('API 密钥含无效字符,请重新填写(不包含空格或换行)')
headers = {'Content-Type': 'application/json', 'Accept': 'application/json'}
if key:
headers['Authorization'] = 'Bearer ' + key
req = urllib.request.Request(config['endpoint'], data=json.dumps(body, ensure_ascii=False).encode(), headers=headers)
try:
# Do not implicitly route chat text through system HTTP proxies.
opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect())
with opener.open(req, timeout=120) as response:
data = response.read(8 * 1024 * 1024 + 1)
if len(data) > 8 * 1024 * 1024:
raise ValueError('模型响应过大')
result = json.loads(data)
if not isinstance(result, dict):
raise ValueError('模型接口返回格式错误')
choices = result.get('choices')
if not isinstance(choices, list) or not choices:
raise ValueError('模型接口未返回 choices')
text = choices[0].get('message', {}).get('content')
if not isinstance(text, str):
raise ValueError('模型没有返回文本结果')
text = text.strip()
if text.startswith('```'):
text = text.split('\n', 1)[1].rsplit('```', 1)[0].strip()
return json.loads(text)
except urllib.error.HTTPError as exc:
# Never persist provider response bodies: they can echo the prompt/key.
raise ValueError(f'模型 API 返回 HTTP {exc.code},请检查地址、密钥、模型和额度') from None
except (urllib.error.URLError, TimeoutError, OSError):
raise ValueError('模型 API 连接失败或超时,请检查网络与接口地址') from None
except (json.JSONDecodeError, IndexError, KeyError, AttributeError):
raise ValueError('模型响应不是可用的 JSON,未保存不完整分析') from None
def sqlite_rows(data, table=None, schema_only=False):
if not data.startswith(b'SQLite format 3\x00'):
raise ValueError('这不是明文 SQLite 数据库。微信加密数据库需先在 Mac 上取得可读导出;本工具不冒充已完成解密。')
with tempfile.TemporaryDirectory(prefix='wechat-import-') as directory:
path = Path(directory) / 'input.sqlite3'
path.write_bytes(data)
db = sqlite3.connect(path.as_uri() + '?mode=ro&immutable=1', uri=True)
db.row_factory = sqlite3.Row
try:
db.execute('PRAGMA query_only=ON')
tables = [r[0] for r in db.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'")]
def quoted(name):
return '"' + name.replace('"', '""') + '"'
if schema_only:
return [{'table': name, 'columns': [r[1] for r in db.execute('PRAGMA table_info(' + quoted(name) + ')')]} for name in tables]
if table not in tables:
raise ValueError('请选择数据库中存在的消息表')
rows = db.execute('SELECT * FROM ' + quoted(table) + ' LIMIT 500001').fetchall()
if len(rows) > 500000:
raise ValueError('单次最多导入 50 万条,请拆分文件')
return [dict(r) for r in rows]
finally:
db.close()
class Workbench(ThreadingHTTPServer):
daemon_threads = True
def __init__(self, address, data):
super().__init__(address, Handler)
self.data = Path(data).resolve()
self.token = secrets.token_urlsafe(32)
self.lock = threading.RLock()
self.plans = {}
self.api_key = os.environ.get('WECHAT_MODEL_API_KEY', '')
self.mac = MacConnection(self.data)
db = store.connect(self.data)
if not store.setting(db, 'salt'):
store.set_setting(db, 'salt', secrets.token_hex(32))
# Interrupted API jobs cannot be silently resumed or automatically re-sent.
db.execute("UPDATE runs SET status='failed',error='服务重启中断了任务;不会自动重新发送,请重新预览后执行' WHERE status='running'")
db.commit()
db.close()
def db(self):
return store.connect(self.data)
class Handler(BaseHTTPRequestHandler):
server_version = 'WeChatWorkbench/0.1'
def log_message(self, *args):
pass # Neither chat payloads nor credentials belong in request logs.
def respond(self, value, status=200, mime='application/json; charset=utf-8', filename=None):
data = json.dumps(value, ensure_ascii=False).encode() if mime.startswith('application/json') else value.encode() if isinstance(value, str) else value
self.send_response(status)
self.send_header('Content-Type', mime)
self.send_header('Content-Length', str(len(data)))
self.send_header('Cache-Control', 'no-store')
self.send_header('X-Content-Type-Options', 'nosniff')
self.send_header('Referrer-Policy', 'no-referrer')
self.send_header('Content-Security-Policy', "default-src 'self'; script-src 'self'; style-src 'self'; connect-src 'self'; img-src 'self' data:; frame-ancestors 'none'; base-uri 'none'; form-action 'self'")
if filename:
self.send_header('Content-Disposition', 'attachment; filename="' + filename + '"')
self.end_headers()
self.wfile.write(data)
def check_request(self, mutation=False):
expected = f'127.0.0.1:{self.server.server_port}'
if self.headers.get('Host') != expected:
raise PermissionError('只允许通过本机 127.0.0.1 访问')
origin = self.headers.get('Origin')
if origin and origin != 'http://' + expected:
raise PermissionError('拒绝跨站请求')
if self.headers.get('Sec-Fetch-Site') in ('cross-site', 'same-site'):
raise PermissionError('拒绝跨站请求')
if mutation and not secrets.compare_digest(self.headers.get('X-Workbench-Token', ''), self.server.token):
raise PermissionError('页面会话失效,请刷新')
def do_GET(self):
self.dispatch(False)
def do_POST(self):
self.dispatch(True)
def dispatch(self, post):
db = None
try:
self.check_request(post)
parts = urllib.parse.urlsplit(self.path)
path = parts.path
if not post and path in ('/', '/app.js', '/style.css'):
file = STATIC / ('index.html' if path == '/' else path[1:])
text = file.read_text(encoding='utf-8').replace('__SESSION_TOKEN__', self.server.token)
return self.respond(text, mime={'/': 'text/html; charset=utf-8', '/app.js': 'text/javascript; charset=utf-8', '/style.css': 'text/css; charset=utf-8'}[path])
body = {}
if post:
length = int(self.headers.get('Content-Length', '0'))
if not 0 < length <= MAX_UPLOAD:
raise ValueError('请求为空或超过 64 MB(文件 base64 后约增大三分之一)')
if 'application/json' not in self.headers.get('Content-Type', ''):
raise ValueError('请求必须是 JSON')
body = json.loads(self.rfile.read(length))
if not isinstance(body, dict):
raise ValueError('请求必须是对象')
query = urllib.parse.parse_qs(parts.query)
db = self.server.db()
if not post:
return self.get_route(db, path, query)
return self.post_route(db, path, body)
except PermissionError as exc:
self.respond({'error': str(exc)}, 403)
except (ValueError, TypeError, sqlite3.Error, KeyError, OverflowError) as exc:
# Type/key/sql errors do not contain user payloads in supported routes.
error = str(exc) if isinstance(exc, ValueError) else '请求格式或数据库结构不符合要求,请检查字段和输入'
self.respond({'error': error}, 400)
except (BrokenPipeError, ConnectionResetError):
pass
except Exception:
self.respond({'error': '内部错误,数据未发送到新的地址。请重启后重试或运行测试诊断。'}, 500)
finally:
if db:
db.close()
def get_route(self, db, path, query):
if path == '/api/mac/status':
return self.respond(self.server.mac.status())
if path == '/api/overview':
result = store.overview(db)
result['runs'] = [dict(r) for r in db.execute('SELECT id,created,mode,status,error,coverage FROM runs ORDER BY created DESC')]
for r in result['runs']:
r['coverage'] = json.loads(r['coverage'])
return self.respond(result)
if path == '/api/settings':
return self.respond({'api': store.setting(db, 'api', {}), 'custom': store.setting(db, 'custom', []),
'utc_offset': store.setting(db, 'utc_offset', '+08:00'), 'has_key': bool(self.server.api_key)})
if path == '/api/run':
run = self.get_run(db, query.get('id', [''])[0])
return self.respond(run)
if path == '/api/evidence':
ids = [int(i) for i in query.get('ids', [''])[0].split(',') if i]
if len(ids) > 1000:
raise ValueError('一次最多读取 1000 条证据')
rows = [dict(r) for r in db.execute('SELECT * FROM messages WHERE id IN (' + ','.join('?' for _ in ids) + ') ORDER BY timestamp,id', ids)] if ids else []
return self.respond(rows)
if path == '/api/context':
run = self.get_run(db, query.get('run_id', [''])[0])
card = next((c for c in run['cards'] if c['id'] == query.get('card_id', [''])[0]), None)
if not card:
raise ValueError('卡片不存在')
events = self.run_events(db, run['id'])
return self.respond([e for e in events if e['id'] in card['event_ids']])
if path == '/api/export':
run = self.get_run(db, query.get('id', [''])[0])
fmt = query.get('format', ['md'])[0]
if fmt == 'csv':
out = io.StringIO(newline='')
writer = csv.writer(out)
writer.writerow(['主题', '摘要', '事件数', '优先级(启发式)', '信号', '状态', '策略', '审核', '证据ID', '会话', '答复候选', '适用边界', '备注', '经验提供者候选'])
for c in run['cards']:
vals = [c['title'], c['summary'], c['event_count'], c['priority'], '、'.join(SIGNALS[s] for s in c['signals']), c['status'], c['strategy'], '已审核' if c['reviewed'] else '待审核', ','.join(map(str, c['evidence_ids'])), '、'.join(c['conversations']), c['answer'], c['boundaries'], c.get('notes', ''), '、'.join(c['responders'])]
# Neutralize spreadsheet formula injection from messages/model output.
writer.writerow(["'" + str(v) if str(v).lstrip().startswith(('=', '+', '-', '@', '\t', '\r')) else v for v in vals])
return self.respond('\ufeff' + out.getvalue(), mime='text/csv; charset=utf-8', filename='knowledge-needs.csv')
lines = ['# 微信知识需求清单', '', '生成时间:' + run['created'], '', '覆盖范围:' + json.dumps(run['coverage'], ensure_ascii=False), '', '模型结果和本地规则均为候选判断;人工审核不等于实践验证。', '']
for c in run['cards']:
lines.extend(['## ' + c['title'].replace('\n', ' '), '', c['summary'], '', f"- 独立事件候选数:{c['event_count']};优先级(启发式):{c['priority']}", '- 信号:' + '、'.join(SIGNALS[s] for s in c['signals']), '- 状态:' + c['status'], '- 处理策略:' + c['strategy'], '- 人工审核:' + ('已审核' if c['reviewed'] else '待审核'), '- 证据消息 ID:' + ', '.join(map(str, c['evidence_ids'])), '', '### 经验候选', '', c['answer'], '', '### 边界与缺口', '', c['boundaries'], '', c.get('notes', ''), ''])
lines.extend(['经验提供者候选(需核实):' + '、'.join(c['responders']), ''])
return self.respond('\n'.join(lines), mime='text/markdown; charset=utf-8', filename='knowledge-needs.md')
return self.respond({'error': '不存在的接口'}, 404)
def get_run(self, db, run_id):
row = db.execute('SELECT * FROM runs WHERE id=?', (run_id,)).fetchone()
if not row:
raise ValueError('分析任务不存在')
result = dict(row)
result.pop('events', None)
for field in ('cards', 'filters', 'coverage'):
result[field] = json.loads(result[field])
return result
def run_events(self, db, run_id):
row = db.execute('SELECT events FROM runs WHERE id=?', (run_id,)).fetchone()
if not row:
raise ValueError('分析任务不存在')
return json.loads(row[0])
def selected(self, db, filters):
rows = store.messages(db, filters)
if not rows:
raise ValueError('筛选范围内没有记录')
if len(rows) > 200000:
raise ValueError('一次最多分析 20 万条,请缩小时间范围')
events = extract_events(rows)
coverage = {'messages': len(rows), 'events': len(events), 'conversations': len({m['conversation_id'] for m in rows}),
'first': rows[0]['timestamp'], 'last': rows[-1]['timestamp'], 'media': sum(m['message_type'] != 'text' for m in rows),
'source': '仅本机已导入记录;时间窗口为讨论候选,不代表完整问答或因果关系'}
return rows, events, coverage
def save_run(self, db, mode, filters, coverage, cards, events, status='complete'):
rid = secrets.token_hex(12)
db.execute('INSERT INTO runs (id,created,mode,filters,cards,coverage,status,error,events) VALUES (?,?,?,?,?,?,?,?,?)', (rid, now(), mode, json.dumps(filters, ensure_ascii=False), json.dumps(cards, ensure_ascii=False), json.dumps(coverage, ensure_ascii=False), status, '', json.dumps(events, ensure_ascii=False)))
db.commit()
return rid
def post_route(self, db, path, body):
if path == '/api/mac/scan':
return self.respond(self.server.mac.scan(bool(body.get('choose_folder'))))
if path == '/api/mac/keys':
if len(json.dumps(body)) > 1024*1024:
raise ValueError('密钥文件最多 1 MB')
return self.respond(self.server.mac.load_keys(body.get('keys')))
if path == '/api/mac/start':
return self.respond(self.server.mac.start(body.get('action'), body))
if path == '/api/mac/recovery-plan':
return self.respond(self.server.mac.recovery_plan(body.get('account')))
if path == '/api/import':
data = base64.b64decode(body['data'], validate=True)
options = body.get('options') or {}
offset_zone(options.get('utc_offset', '+08:00'))
filename = Path(body['filename']).name
if filename.lower().endswith(('.db', '.sqlite', '.sqlite3')):
rows = sqlite_rows(data, options.get('table'))
else:
rows = read_rows(data, filename)
if len(rows) > 500000:
raise ValueError('单次最多导入 50 万条')
return self.respond(import_rows(db, rows, filename, options))
if path == '/api/schema':
return self.respond(sqlite_rows(base64.b64decode(body['data'], validate=True), schema_only=True))
if path == '/api/demo':
rows = json.loads((Path(__file__).parent.parent / 'examples' / 'demo.json').read_text(encoding='utf-8'))
return self.respond(import_rows(db, rows, 'demo.json(虚构样例)', {'utc_offset': '+08:00'}))
if path == '/api/settings':
api = validate_api(body['api']) if body.get('api', {}).get('endpoint') else {}
custom = body.get('custom', [])
if not isinstance(custom, list) or any(not isinstance(x, str) for x in custom) or len(custom) > 2000:
raise ValueError('自定义敏感词必须是字符串数组,最多 2000 项')
offset_zone(body.get('utc_offset', '+08:00'))
new_key = str(body.get('key', '')).strip()
if new_key and (not new_key.isascii() or any(ord(c) < 33 or ord(c) > 126 for c in new_key)):
raise ValueError('API 密钥含无效字符,请重新填写(不包含空格或换行)')
store.set_setting(db, 'api', api)
store.set_setting(db, 'custom', custom)
store.set_setting(db, 'utc_offset', body.get('utc_offset', '+08:00'))
with self.server.lock:
if body.get('clear_key'):
self.server.api_key = ''
elif body.get('key'):
self.server.api_key = new_key
self.server.plans.clear()
return self.respond({'saved': True})
if path in ('/api/analyze', '/api/preview'):
filters = body.get('filters') or {}
rows, events, coverage = self.selected(db, filters)
cross = bool(body.get('cross', True))
filters = dict(filters, cross=cross)
if path == '/api/analyze':
rid = self.save_run(db, 'local', filters, coverage, local_cards(events, cross), events)
return self.respond({'id': rid})
if not events:
raise ValueError('当前规则未找到问题或复用线索;可查看原文或调整范围')
config = validate_api(store.setting(db, 'api', {}))
# Collect names from all imported conversations, including excluded ones.
contacts = [dict(r) for r in db.execute('SELECT DISTINCT sender,sender_id,conversation,conversation_id FROM messages')]
redactor = Redactor(contacts, store.setting(db, 'salt'), store.setting(db, 'custom', []))
batches = prepare_plan(events, redactor, cross)
bodies = [{'model': config['model'], 'messages': [{'role': 'system', 'content': SYSTEM_PROMPT}, {'role': 'user', 'content': json.dumps(batch, ensure_ascii=False)}], 'temperature': 0.2} for batch in batches]
plan_id = secrets.token_hex(16)
plan = dict(config=config, bodies=bodies, events=events, filters=filters, coverage=coverage, created=datetime.now().timestamp())
with self.server.lock:
# Preview plaintext is redacted. Original event snapshot stays in memory.
if len(self.server.plans) >= 10:
self.server.plans.clear()
self.server.plans[plan_id] = plan
return self.respond({'plan_id': plan_id, 'endpoint': config['endpoint'], 'coverage': coverage, 'batches': len(bodies), 'characters': sum(len(json.dumps(b, ensure_ascii=False)) for b in bodies), 'requests': bodies})
if path == '/api/send':
with self.server.lock:
plan = self.server.plans.pop(body.get('plan_id', ''), None)
key = self.server.api_key
if not plan or datetime.now().timestamp() - plan['created'] > 1800:
raise ValueError('预览已失效或已发送,请重新生成预览')
coverage = dict(plan['coverage'], batches=len(plan['bodies']), cross_batch_merge=False)
rid = self.save_run(db, 'api', plan['filters'], coverage, [], plan['events'], 'running')
threading.Thread(target=self.run_ai, args=(rid, plan, key), daemon=True).start()
return self.respond({'id': rid})
if path == '/api/card':
with self.server.lock:
run = self.get_run(db, body['run_id'])
if run['status'] != 'complete':
raise ValueError('任务尚未完成')
cards = run['cards']
action = body.get('action', 'edit')
indices = [i for i, c in enumerate(cards) if c['id'] in body.get('ids', [])]
if not indices:
raise ValueError('请选择卡片')
if action == 'edit':
if len(indices) != 1:
raise ValueError('一次编辑一张卡片')
card = cards[indices[0]]
for field in ('title', 'summary', 'answer', 'boundaries', 'notes'):
if field in body:
card[field] = str(body[field])[:10000]
if body.get('strategy') in STRATEGIES:
card['strategy'] = body['strategy']
card['reviewed'] = bool(body.get('reviewed'))
elif action in ('merge', 'split'):
event_map = {e['id']: e for e in self.run_events(db, run['id'])}
if action == 'merge':
if len(indices) < 2:
raise ValueError('合并至少需要两张卡片')
events = [event_map[eid] for i in indices for eid in cards[i]['event_ids']]
if not run['filters'].get('cross', True) and len({e['conversation_id'] for e in events}) > 1:
raise ValueError('该任务按会话分析,不可跨会话合并')
new = make_card(events, body.get('title') or cards[indices[0]]['title'])
new['notes'] = '\n\n'.join(c['title'] + '\n' + c['summary'] + '\n' + c.get('notes', '') for i,c in enumerate(cards) if i in indices)
additions = [new]
else:
if len(indices) != 1:
raise ValueError('一次拆分一张卡片')
original = cards[indices[0]]
ids = original['event_ids']
chosen = body.get('event_ids', [])
if not chosen or not set(chosen) < set(ids):
raise ValueError('请选择要拆出的部分事件,不能选空或全选')
additions = [make_card([event_map[e] for e in ids if (e in chosen) == included]) for included in (True, False)]
for new in additions:
new['notes'] = '拆分前备注:' + original.get('notes', '')
cards = [c for i,c in enumerate(cards) if i not in indices] + additions
else:
raise ValueError('未知编辑操作')
db.execute('UPDATE runs SET cards=? WHERE id=?', (json.dumps(cards, ensure_ascii=False), run['id']))
db.commit()
return self.respond({'saved': True})
return self.respond({'error': '不存在的接口'}, 404)
def run_ai(self, rid, plan, key):
db = self.server.db()
try:
cards = []
event_map = {e['id']: e for e in plan['events']}
for body in plan['bodies']:
batch = json.loads(body['messages'][1]['content'])
events = [event_map[e['id']] for e in batch['events']]
result = call_model(plan['config'], key, body)
cards.extend(validate_model_cards(result, events, plan['filters']['cross']))
db.execute('UPDATE runs SET cards=?,status=? WHERE id=?', (json.dumps(cards, ensure_ascii=False), 'complete', rid))
db.commit()
except Exception as exc:
reason = str(exc) if isinstance(exc, ValueError) else '分析失败;未保存不完整卡片。请检查接口后重新预览执行。'
db.execute('UPDATE runs SET status=?,error=? WHERE id=?', ('failed', reason[:500], rid))
db.commit()
finally:
db.close()
def main():
parser = argparse.ArgumentParser(description='微信知识需求分析工具(仅绑定本机)')
parser.add_argument('--port', type=int, default=8765)
parser.add_argument('--data-dir', type=Path, default=store.DEFAULT_DATA)
parser.add_argument('--no-browser', action='store_true')
args = parser.parse_args()
try:
server = Workbench(('127.0.0.1', args.port), args.data_dir)
except OSError:
parser.exit(1, '端口被占用。请关闭旧进程,或使用 --port 8766。\n')
url = f'http://127.0.0.1:{server.server_port}'
print(f'WeChat knowledge workbench: {url}', flush=True)
print(f'Local data: {server.data}', flush=True)
if not args.no_browser:
webbrowser.open(url)
try:
server.serve_forever()
except KeyboardInterrupt:
print('\nStopped.')
finally:
server.server_close()
if __name__ == '__main__':
main()