guanghulab/server/corpus-agent/server.py

"""
语料采集系统 · 服务端
====================
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)