506 lines
16 KiB
Python
506 lines
16 KiB
Python
"""
|
||
语料采集系统 · 服务端
|
||
====================
|
||
FastAPI + Gitea OAuth + WebSocket实时采集
|
||
"""
|
||
import os
|
||
import json
|
||
import uuid
|
||
import time
|
||
import hashlib
|
||
from pathlib import Path
|
||
from typing import Optional, List
|
||
from datetime import datetime, timedelta
|
||
|
||
from fastapi import FastAPI, HTTPException, Depends, Query, WebSocket, WebSocketDisconnect
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
|
||
from fastapi.staticfiles import StaticFiles
|
||
from pydantic import BaseModel
|
||
import uvicorn
|
||
|
||
# 导入引擎
|
||
from engine import (
|
||
process_text_chunk, process_dialog_stream,
|
||
desensitize, is_valuable, classify_content,
|
||
to_chatml, format_sft_jsonl, generate_corpus_id
|
||
)
|
||
|
||
# ============================================================
|
||
# 配置
|
||
# ============================================================
|
||
|
||
class Settings:
|
||
APP_NAME = "语料采集系统 Corpus Agent"
|
||
VERSION = "1.0.0"
|
||
|
||
# 存储路径
|
||
DATA_DIR = Path(os.environ.get("CORPUS_DATA_DIR", "./data"))
|
||
USERS_DIR = DATA_DIR / "users"
|
||
|
||
# Gitea OAuth
|
||
GITEA_URL = os.environ.get("GITEA_URL", "https://guanghulab.com")
|
||
GITEA_CLIENT_ID = os.environ.get("GITEA_CLIENT_ID", "")
|
||
GITEA_CLIENT_SECRET = os.environ.get("GITEA_CLIENT_SECRET", "")
|
||
|
||
# Session
|
||
SESSION_SECRET = os.environ.get("SESSION_SECRET", "corpus-agent-secret-key")
|
||
SESSION_EXPIRE_HOURS = 24
|
||
|
||
# 服务器
|
||
HOST = os.environ.get("CORPUS_HOST", "0.0.0.0")
|
||
PORT = int(os.environ.get("CORPUS_PORT", "8084"))
|
||
|
||
settings = Settings()
|
||
|
||
# ============================================================
|
||
# 数据模型
|
||
# ============================================================
|
||
|
||
class TextChunk(BaseModel):
|
||
text: str
|
||
source: str = "manual"
|
||
session_id: Optional[str] = None
|
||
|
||
class DialogMessage(BaseModel):
|
||
role: str
|
||
content: str
|
||
source: str = "unknown"
|
||
|
||
class DialogBatch(BaseModel):
|
||
messages: List[DialogMessage]
|
||
session_id: Optional[str] = None
|
||
|
||
class CorpusSave(BaseModel):
|
||
samples: List[dict]
|
||
filename: str = "corpus"
|
||
|
||
class LoginRequest(BaseModel):
|
||
gitea_token: str
|
||
|
||
class SessionData(BaseModel):
|
||
username: str
|
||
login_time: float
|
||
exprire_at: float
|
||
|
||
# ============================================================
|
||
# App
|
||
# ============================================================
|
||
|
||
app = FastAPI(title=settings.APP_NAME, version=settings.VERSION)
|
||
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
# 静态文件
|
||
static_dir = Path(__file__).parent / "static"
|
||
static_dir.mkdir(exist_ok=True)
|
||
app.mount("/static", StaticFiles(directory=str(static_dir)), name="static")
|
||
|
||
# ============================================================
|
||
# Session管理（简化版，生产环境改为Redis）
|
||
# ============================================================
|
||
|
||
sessions = {} # token -> SessionData
|
||
|
||
def create_session(username: str) -> str:
|
||
token = hashlib.sha256(f"{username}{time.time()}{uuid.uuid4()}".encode()).hexdigest()[:32]
|
||
sessions[token] = SessionData(
|
||
username=username,
|
||
login_time=time.time(),
|
||
exprire_at=time.time() + settings.SESSION_EXPIRE_HOURS * 3600,
|
||
)
|
||
return token
|
||
|
||
def get_user_from_token(token: str) -> Optional[str]:
|
||
if token in sessions:
|
||
sd = sessions[token]
|
||
if time.time() < sd.exprire_at:
|
||
return sd.username
|
||
del sessions[token]
|
||
return None
|
||
|
||
# ============================================================
|
||
# 用户数据管理
|
||
# ============================================================
|
||
|
||
def get_user_dir(username: str) -> Path:
|
||
path = settings.USERS_DIR / username
|
||
path.mkdir(parents=True, exist_ok=True)
|
||
return path
|
||
|
||
def load_corpus(username: str) -> list:
|
||
path = get_user_dir(username) / "corpus.jsonl"
|
||
if not path.exists():
|
||
return []
|
||
samples = []
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if line:
|
||
try:
|
||
samples.append(json.loads(line))
|
||
except:
|
||
pass
|
||
return samples
|
||
|
||
def append_corpus(username: str, samples: list):
|
||
path = get_user_dir(username) / "corpus.jsonl"
|
||
with open(path, "a", encoding="utf-8") as f:
|
||
for s in samples:
|
||
f.write(json.dumps(s, ensure_ascii=False) + "\n")
|
||
|
||
def save_corpus_snapshot(username: str, filename: str) -> Path:
|
||
"""保存快照文件"""
|
||
samples = load_corpus(username)
|
||
snapshot_dir = get_user_dir(username) / "snapshots"
|
||
snapshot_dir.mkdir(exist_ok=True)
|
||
|
||
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
path = snapshot_dir / f"{filename}_{ts}.jsonl"
|
||
|
||
with open(path, "w", encoding="utf-8") as f:
|
||
for s in samples:
|
||
f.write(json.dumps(s, ensure_ascii=False) + "\n")
|
||
|
||
return path
|
||
|
||
# ============================================================
|
||
# API路由
|
||
# ============================================================
|
||
|
||
@app.get("/")
|
||
async def root():
|
||
return HTMLResponse(open(static_dir / "index.html").read())
|
||
|
||
@app.get("/api/health")
|
||
async def health():
|
||
return {"status": "ok", "app": settings.APP_NAME, "version": settings.VERSION}
|
||
|
||
# --- 登录 ---
|
||
|
||
@app.post("/api/auth/login")
|
||
async def login(req: LoginRequest):
|
||
"""使用Gitea Access Token登录"""
|
||
# 通过Gitea API验证token
|
||
import urllib.request
|
||
try:
|
||
gitea_req = urllib.request.Request(
|
||
f"{settings.GITEA_URL}/api/v1/user",
|
||
headers={"Authorization": f"token {req.gitea_token}"}
|
||
)
|
||
with urllib.request.urlopen(gitea_req, timeout=10) as resp:
|
||
user_data = json.loads(resp.read().decode())
|
||
username = user_data.get("login", "unknown")
|
||
except Exception as e:
|
||
# 降级：允许直接使用用户名登录（开发模式）
|
||
username = req.gitea_token.split(":")[0] if ":" in req.gitea_token else req.gitea_token
|
||
|
||
token = create_session(username)
|
||
|
||
# 统计现有语料
|
||
corpus = load_corpus(username)
|
||
|
||
return {
|
||
"ok": True,
|
||
"token": token,
|
||
"username": username,
|
||
"stats": {
|
||
"total_samples": len(corpus),
|
||
"total_chars": sum(len(json.dumps(s, ensure_ascii=False)) for s in corpus),
|
||
}
|
||
}
|
||
|
||
@app.get("/api/auth/check")
|
||
async def check_session(token: str = Query(...)):
|
||
username = get_user_from_token(token)
|
||
if not username:
|
||
raise HTTPException(401, "登录已过期")
|
||
return {"ok": True, "username": username}
|
||
|
||
# --- 语料操作 ---
|
||
|
||
@app.post("/api/corpus/collect")
|
||
async def collect_chunk(chunk: TextChunk, token: str = Query(...)):
|
||
"""单段文本采集"""
|
||
username = get_user_from_token(token)
|
||
if not username:
|
||
raise HTTPException(401, "登录已过期")
|
||
|
||
# 引擎处理
|
||
results = process_text_chunk(chunk.text, source=chunk.source)
|
||
if not results:
|
||
return {"ok": True, "collected": 0, "reason": "no_valuable_content"}
|
||
|
||
# 添加时间戳和归属
|
||
ts = datetime.now().isoformat()
|
||
for r in results:
|
||
r["timestamp"] = ts
|
||
r["user"] = username
|
||
|
||
# 转为ChatML格式存储
|
||
sft_samples = []
|
||
for r in results:
|
||
if "user" in r and "assistant" in r:
|
||
sft_samples.append(to_chatml(r["user"], r["assistant"]))
|
||
else:
|
||
# 单条文本也存起来
|
||
sft_samples.append({
|
||
"messages": [{"role": "user", "content": r["text"]}],
|
||
"source": r["source"],
|
||
"tags": r.get("tags", []),
|
||
"corpus_id": r.get("corpus_id", generate_corpus_id(r["text"])),
|
||
"collected_at": ts,
|
||
})
|
||
|
||
append_corpus(username, sft_samples)
|
||
|
||
return {
|
||
"ok": True,
|
||
"collected": len(sft_samples),
|
||
"preview": sft_samples[:3],
|
||
"stats": {"total": len(load_corpus(username))},
|
||
}
|
||
|
||
@app.post("/api/corpus/batch")
|
||
async def collect_batch(batch: DialogBatch, token: str = Query(...)):
|
||
"""批量对话采集"""
|
||
username = get_user_from_token(token)
|
||
if not username:
|
||
raise HTTPException(401, "登录已过期")
|
||
|
||
msgs = [{"role": m.role, "content": m.content, "source": m.source} for m in batch.messages]
|
||
result = process_dialog_stream(msgs)
|
||
|
||
if result["stats"]["valuable_pairs"] > 0:
|
||
# 添加用户归属
|
||
ts = datetime.now().isoformat()
|
||
for p in result["pairs"]:
|
||
p["collected_at"] = ts
|
||
p["user"] = username
|
||
|
||
append_corpus(username, result["pairs"])
|
||
|
||
return {
|
||
"ok": True,
|
||
**result["stats"],
|
||
}
|
||
|
||
@app.get("/api/corpus/list")
|
||
async def list_corpus(token: str = Query(...), page: int = Query(1), size: int = Query(50)):
|
||
"""查看已采集语料"""
|
||
username = get_user_from_token(token)
|
||
if not username:
|
||
raise HTTPException(401, "登录已过期")
|
||
|
||
all_samples = load_corpus(username)
|
||
total = len(all_samples)
|
||
|
||
# 分页
|
||
start = (page - 1) * size
|
||
end = start + size
|
||
page_samples = all_samples[start:end]
|
||
|
||
return {
|
||
"ok": True,
|
||
"total": total,
|
||
"page": page,
|
||
"size": size,
|
||
"samples": page_samples,
|
||
}
|
||
|
||
@app.get("/api/corpus/export")
|
||
async def export_corpus(token: str = Query(...), format: str = "jsonl"):
|
||
"""导出语料文件"""
|
||
username = get_user_from_token(token)
|
||
if not username:
|
||
raise HTTPException(401, "登录已过期")
|
||
|
||
if format == "jsonl":
|
||
path = save_corpus_snapshot(username, "export")
|
||
return FileResponse(
|
||
path,
|
||
media_type="application/octet-stream",
|
||
filename=f"corpus_{username}_{datetime.now().strftime('%Y%m%d')}.jsonl"
|
||
)
|
||
|
||
raise HTTPException(400, f"不支持的格式: {format}")
|
||
|
||
@app.post("/api/corpus/preview")
|
||
async def preview_text(chunk: TextChunk):
|
||
"""预览一段文本的处理结果（无需登录）"""
|
||
from engine import preview as engine_preview
|
||
return engine_preview(chunk.text)
|
||
|
||
@app.get("/api/corpus/stats")
|
||
async def corpus_stats(token: str = Query(...)):
|
||
"""语料统计"""
|
||
username = get_user_from_token(token)
|
||
if not username:
|
||
raise HTTPException(401, "登录已过期")
|
||
|
||
samples = load_corpus(username)
|
||
|
||
# 按来源统计
|
||
by_source = {}
|
||
# 按标签统计
|
||
by_tag = {}
|
||
total_chars = 0
|
||
|
||
for s in samples:
|
||
msgs = s.get("messages", [])
|
||
for m in msgs:
|
||
total_chars += len(m.get("content", ""))
|
||
|
||
source = s.get("source", "unknown")
|
||
by_source[source] = by_source.get(source, 0) + 1
|
||
|
||
for tag in s.get("tags", []):
|
||
by_tag[tag] = by_tag.get(tag, 0) + 1
|
||
|
||
return {
|
||
"ok": True,
|
||
"total_samples": len(samples),
|
||
"total_chars": total_chars,
|
||
"by_source": by_source,
|
||
"by_tag": by_tag,
|
||
}
|
||
|
||
# --- 统计数据 ---
|
||
|
||
@app.get("/api/system/stats")
|
||
async def system_stats():
|
||
"""系统统计（公开）"""
|
||
total_users = len(list(settings.USERS_DIR.iterdir())) if settings.USERS_DIR.exists() else 0
|
||
total_samples = 0
|
||
for user_dir in settings.USERS_DIR.iterdir() if settings.USERS_DIR.exists() else []:
|
||
if user_dir.is_dir():
|
||
total_samples += len(load_corpus(user_dir.name))
|
||
|
||
return {
|
||
"total_users": total_users,
|
||
"total_samples": total_samples,
|
||
"active_sessions": len(sessions),
|
||
}
|
||
|
||
# ============================================================
|
||
# WebSocket实时采集（Mac客户端用）
|
||
# ============================================================
|
||
|
||
class ConnectionManager:
|
||
def __init__(self):
|
||
self.active_connections: dict[str, list[WebSocket]] = {}
|
||
|
||
async def connect(self, websocket: WebSocket, username: str):
|
||
await websocket.accept()
|
||
if username not in self.active_connections:
|
||
self.active_connections[username] = []
|
||
self.active_connections[username].append(websocket)
|
||
|
||
def disconnect(self, websocket: WebSocket, username: str):
|
||
if username in self.active_connections:
|
||
self.active_connections[username] = [ws for ws in self.active_connections[username] if ws != websocket]
|
||
if not self.active_connections[username]:
|
||
del self.active_connections[username]
|
||
|
||
async def send_to_user(self, username: str, message: dict):
|
||
if username in self.active_connections:
|
||
for ws in self.active_connections[username]:
|
||
try:
|
||
await ws.send_json(message)
|
||
except:
|
||
pass
|
||
|
||
manager = ConnectionManager()
|
||
|
||
@app.websocket("/ws/collect")
|
||
async def websocket_collect(websocket: WebSocket, token: str = Query(...)):
|
||
username = get_user_from_token(token)
|
||
if not username:
|
||
await websocket.close(code=4001, reason="未登录")
|
||
return
|
||
|
||
await manager.connect(websocket, username)
|
||
|
||
try:
|
||
while True:
|
||
data = await websocket.receive_json()
|
||
|
||
# 支持多种消息类型
|
||
msg_type = data.get("type", "text")
|
||
|
||
if msg_type == "text":
|
||
chunk = TextChunk(text=data.get("text", ""), source=data.get("source", "screen_capture"))
|
||
results = process_text_chunk(chunk.text, source=chunk.source)
|
||
if results:
|
||
ts = datetime.now().isoformat()
|
||
for r in results:
|
||
r["timestamp"] = ts
|
||
r["user"] = username
|
||
|
||
sft_samples = []
|
||
for r in results:
|
||
if "user" in r and "assistant" in r:
|
||
sft_samples.append(to_chatml(r["user"], r["assistant"]))
|
||
else:
|
||
sft_samples.append({
|
||
"messages": [{"role": "user", "content": r["text"]}],
|
||
"source": r["source"],
|
||
"tags": r.get("tags", []),
|
||
"corpus_id": r.get("corpus_id", ""),
|
||
"collected_at": ts,
|
||
})
|
||
|
||
if sft_samples:
|
||
append_corpus(username, sft_samples)
|
||
|
||
await websocket.send_json({
|
||
"type": "result",
|
||
"collected": len(sft_samples),
|
||
"valuable": len(results) > 0,
|
||
"preview": results[:2],
|
||
})
|
||
else:
|
||
await websocket.send_json({"type": "result", "collected": 0, "valuable": False})
|
||
|
||
elif msg_type == "batch":
|
||
messages = data.get("messages", [])
|
||
result = process_dialog_stream(messages)
|
||
if result["stats"]["valuable_pairs"] > 0:
|
||
ts = datetime.now().isoformat()
|
||
for p in result["pairs"]:
|
||
p["collected_at"] = ts
|
||
p["user"] = username
|
||
append_corpus(username, result["pairs"])
|
||
|
||
await websocket.send_json({
|
||
"type": "batch_result",
|
||
**result["stats"],
|
||
})
|
||
|
||
elif msg_type == "ping":
|
||
await websocket.send_json({"type": "pong"})
|
||
|
||
except WebSocketDisconnect:
|
||
manager.disconnect(websocket, username)
|
||
except Exception as e:
|
||
manager.disconnect(websocket, username)
|
||
|
||
# ============================================================
|
||
# 启动
|
||
# ============================================================
|
||
|
||
if __name__ == "__main__":
|
||
print(f"🚀 {settings.APP_NAME} v{settings.VERSION}")
|
||
print(f"📂 数据目录: {settings.DATA_DIR.absolute()}")
|
||
print(f"🌐 服务地址: http://{settings.HOST}:{settings.PORT}")
|
||
|
||
settings.USERS_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
||
uvicorn.run(app, host=settings.HOST, port=settings.PORT)
|