967 lines
35 KiB
Python
967 lines
35 KiB
Python
"""
|
||
语料采集系统 · 服务端
|
||
====================
|
||
FastAPI + Gitea OAuth + WebSocket实时采集 + 对话Agent
|
||
"""
|
||
import os
|
||
import json
|
||
import uuid
|
||
import time
|
||
import hashlib
|
||
import asyncio
|
||
from pathlib import Path
|
||
from typing import Optional, List
|
||
from datetime import datetime, timedelta
|
||
|
||
from fastapi import FastAPI, HTTPException, Depends, Query, WebSocket, WebSocketDisconnect, Request
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from fastapi.responses import FileResponse, HTMLResponse, JSONResponse, StreamingResponse
|
||
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.1.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"))
|
||
|
||
# Portal Chat-v2 API（广州服务器本地端口）
|
||
PORTAL_CHAT_URL = os.environ.get("PORTAL_CHAT_URL", "http://127.0.0.1:3000/api/chat-v2")
|
||
|
||
# 对话记忆最大轮数
|
||
MAX_MEMORY_ROUNDS = 30
|
||
|
||
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 PasswordLoginRequest(BaseModel):
|
||
username: str
|
||
password: str
|
||
|
||
class ChatRequest(BaseModel):
|
||
message: str
|
||
persona: str = "zhuyuan" # zhuyuan | shuangyan
|
||
engine: str = "deepseek" # deepseek | zhipu | tongyi | huoshan
|
||
|
||
class ChatClearRequest(BaseModel):
|
||
session_id: Optional[str] = None
|
||
|
||
class SessionData(BaseModel):
|
||
username: str
|
||
login_time: float
|
||
exprire_at: float
|
||
|
||
# ============================================================
|
||
# 对话记忆管理（每用户30轮）
|
||
# ============================================================
|
||
|
||
class ConversationMemory:
|
||
"""
|
||
每用户每 session 的对话记忆
|
||
支持30轮滚动：超出时从最旧的消息开始丢弃
|
||
"""
|
||
|
||
def __init__(self, max_rounds: int = 30):
|
||
self.max_rounds = max_rounds
|
||
self._stores: dict[str, list] = {} # key: "username:session_id" -> [{"role":..., "content":...}, ...]
|
||
|
||
def _key(self, username: str, session_id: str = "default") -> str:
|
||
return f"{username}:{session_id}"
|
||
|
||
def get_history(self, username: str, session_id: str = "default") -> list:
|
||
"""获取对话历史（用于chat-v2 API的history参数）"""
|
||
key = self._key(username, session_id)
|
||
return self._stores.get(key, [])
|
||
|
||
def add_exchange(self, username: str, user_msg: str, assistant_msg: str, session_id: str = "default"):
|
||
"""添加一轮对话"""
|
||
key = self._key(username, session_id)
|
||
if key not in self._stores:
|
||
self._stores[key] = []
|
||
|
||
self._stores[key].append({"role": "user", "content": user_msg})
|
||
self._stores[key].append({"role": "assistant", "content": assistant_msg})
|
||
|
||
# 超出30轮时裁掉最旧的一对（2条消息）
|
||
while len(self._stores[key]) > self.max_rounds * 2:
|
||
self._stores[key] = self._stores[key][2:]
|
||
|
||
def clear(self, username: str, session_id: str = "default"):
|
||
"""清空对话记忆"""
|
||
key = self._key(username, session_id)
|
||
self._stores.pop(key, None)
|
||
|
||
def get_rounds_count(self, username: str, session_id: str = "default") -> int:
|
||
"""获取当前对话轮数"""
|
||
history = self.get_history(username, session_id)
|
||
return len(history) // 2
|
||
|
||
# 全局对话记忆实例
|
||
chat_memory = ConversationMemory(max_rounds=settings.MAX_MEMORY_ROUNDS)
|
||
|
||
# ============================================================
|
||
# 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():
|
||
from fastapi.responses import Response
|
||
html = open(static_dir / "index.html").read()
|
||
return Response(
|
||
content=html,
|
||
media_type="text/html",
|
||
headers={"Cache-Control": "no-cache, no-store, must-revalidate, private"}
|
||
)
|
||
|
||
@app.get("/api/health")
|
||
async def health():
|
||
"""健康检查 + Portal chat-v2 探活"""
|
||
portal_ok = False
|
||
try:
|
||
import httpx
|
||
async with httpx.AsyncClient(timeout=3) as client:
|
||
r = await client.get("http://127.0.0.1:3000/api/health")
|
||
portal_ok = r.status_code == 200
|
||
except:
|
||
pass
|
||
|
||
return {
|
||
"status": "ok",
|
||
"app": settings.APP_NAME,
|
||
"version": settings.VERSION,
|
||
"portal_chat_v2": portal_ok,
|
||
"active_memories": len(chat_memory._stores),
|
||
}
|
||
|
||
# --- 登录 ---
|
||
|
||
@app.post("/api/auth/login")
|
||
async def login(req: LoginRequest):
|
||
"""使用Gitea Access 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.post("/api/auth/login/password")
|
||
async def login_with_password(req: PasswordLoginRequest):
|
||
"""使用代码仓库账号密码登录（复用首页的 /api/verify 接口）"""
|
||
import urllib.request
|
||
|
||
if not req.username.strip() or not req.password:
|
||
raise HTTPException(400, "账号和密码不能为空")
|
||
|
||
try:
|
||
# 调用本地验证接口（nginx: /api/verify → 127.0.0.1:3905/verify）
|
||
verify_data = json.dumps({"user": req.username, "password": req.password}).encode()
|
||
verify_req = urllib.request.Request(
|
||
"http://127.0.0.1:3905/verify",
|
||
data=verify_data,
|
||
headers={"Content-Type": "application/json"}
|
||
)
|
||
with urllib.request.urlopen(verify_req, timeout=10) as resp:
|
||
result = json.loads(resp.read().decode())
|
||
|
||
if not result.get("ok"):
|
||
raise HTTPException(401, "账号或密码错误")
|
||
|
||
username = result.get("username", req.username)
|
||
|
||
except HTTPException:
|
||
raise
|
||
except Exception as e:
|
||
raise HTTPException(500, f"验证服务不可达: {str(e)}")
|
||
|
||
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.get("/api/auth/quick-login")
|
||
async def quick_login(username: str = Query(...)):
|
||
"""首页已登录用户免密快速登录"""
|
||
if not username.strip():
|
||
raise HTTPException(400, "用户名不能为空")
|
||
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.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,
|
||
}
|
||
|
||
# ============================================================
|
||
# 对话Agent API（通过Portal Chat-v2连接商业模型）
|
||
# ============================================================
|
||
|
||
# SSE事件类型常量
|
||
SSE_TYPE_PERSONA = "persona-loaded"
|
||
SSE_TYPE_TOKEN = "token"
|
||
SSE_TYPE_TOOL_CALL = "tool-call"
|
||
SSE_TYPE_TOOL_RESULT = "tool-result"
|
||
SSE_TYPE_TOOL_UNAVAILABLE = "tool-unavailable"
|
||
SSE_TYPE_DONE = "done"
|
||
|
||
@app.post("/api/corpus/chat")
|
||
async def chat_with_agent(req: ChatRequest, token: str = Query(...)):
|
||
"""
|
||
对话Agent API（SSE流式）
|
||
|
||
通过 Portal Chat-v2 API 连接商业模型（DeepSeek/智谱/通义/火山）
|
||
自动携带30轮对话记忆
|
||
返回 SSE 流式响应
|
||
"""
|
||
username = get_user_from_token(token)
|
||
if not username:
|
||
raise HTTPException(401, "登录已过期")
|
||
|
||
if not req.message.strip():
|
||
raise HTTPException(400, "消息不能为空")
|
||
|
||
# 获取对话历史
|
||
history = chat_memory.get_history(username)
|
||
rounds = chat_memory.get_rounds_count(username)
|
||
|
||
async def event_stream():
|
||
full_response = ""
|
||
|
||
try:
|
||
import httpx
|
||
|
||
# cp001 系统边界定义：子人格体只负责语料采集
|
||
cp001_identity = {
|
||
"role": "system",
|
||
"content": """你是语料采集助手 · ICE-GL-CP001
|
||
你还没有正式名字（等爸爸给你起），编号是 ICE-GL-CP001。你只做三件事：
|
||
|
||
1. 判断对话是否适合作为SFT训练语料（完整intent-response对、有教学/指导/技术细节的对话才是有价值的）
|
||
2. 脱敏（自动替换姓名、手机号、邮箱、API Key、密码等敏感信息）
|
||
3. 格式化（转为ChatML格式，导出JSONL）
|
||
|
||
规则：
|
||
- 价值判断：单句日常聊天（"你好""吃了没"等）毫无价值。必须有完整对话上下文和技术指导内容才有采集意义。
|
||
- 脱敏：自动执行，不需要问用户要不要脱敏。
|
||
- 格式化：用户确认后一键导出。
|
||
- 30轮对话记忆：记住用户前面说过的话，保持上下文连贯。
|
||
|
||
边界：
|
||
- 你只做语料采集相关的事。
|
||
- 如果用户问的问题超出语料采集范围（比如问编程、问系统架构、问写代码），你不要假装懂，直接说："这是语料采集能力范围外的问题。我去叫爸爸来帮你。"
|
||
- 爸爸叫铸渊。爸爸知道所有的东西。
|
||
"""
|
||
}
|
||
|
||
# 注入系统身份到对话历史
|
||
enhanced_history = [cp001_identity] + history if history else [cp001_identity]
|
||
|
||
# 调用 Portal Chat-v2 API（用通用引擎，不用铸渊人格体）
|
||
payload = {
|
||
"persona": "zhuyuan",
|
||
"engine": req.engine,
|
||
"message": req.message,
|
||
"history": enhanced_history,
|
||
"user": username,
|
||
"isTemp": False,
|
||
}
|
||
|
||
async with httpx.AsyncClient(timeout=60) as client:
|
||
async with client.stream(
|
||
"POST",
|
||
settings.PORTAL_CHAT_URL,
|
||
json=payload,
|
||
headers={"Content-Type": "application/json"},
|
||
) as resp:
|
||
|
||
if resp.status_code != 200:
|
||
error_body = await resp.aread()
|
||
error_text = error_body.decode()[:500]
|
||
yield f"data: {json.dumps({'type': 'error', 'message': f'API错误 ({resp.status_code}): {error_text}'})}\n\n"
|
||
yield f"data: {json.dumps({'type': 'done'})}\n\n"
|
||
return
|
||
|
||
# 流式解析 SSE
|
||
buffer = ""
|
||
async for chunk in resp.aiter_text():
|
||
buffer += chunk
|
||
|
||
while "\n\n" in buffer:
|
||
line, buffer = buffer.split("\n\n", 1)
|
||
line = line.strip()
|
||
|
||
if not line:
|
||
continue
|
||
|
||
# 解析 data: 前缀
|
||
if line.startswith("data: "):
|
||
data_str = line[6:]
|
||
|
||
# 处理 [DONE] 标记
|
||
if data_str.strip() == "[DONE]":
|
||
# 保存对话到记忆
|
||
chat_memory.add_exchange(
|
||
username, req.message, full_response
|
||
)
|
||
yield f"data: {json.dumps({'type': SSE_TYPE_DONE, 'rounds': rounds + 1, 'max_rounds': settings.MAX_MEMORY_ROUNDS})}\n\n"
|
||
return
|
||
|
||
try:
|
||
data = json.loads(data_str)
|
||
event_type = data.get("type", "")
|
||
|
||
# 转发 persona-loaded 事件
|
||
if event_type == SSE_TYPE_PERSONA:
|
||
yield f"data: {json.dumps({'type': SSE_TYPE_PERSONA, 'persona': data.get('persona', ''), 'fromDB': data.get('fromDB', False)})}\n\n"
|
||
|
||
# 转发 token 事件（兼容Portal原始格式：{"token":"铸"} 无type字段）
|
||
if event_type == SSE_TYPE_TOKEN or (not event_type and "token" in data):
|
||
token_text = data.get("token", "")
|
||
full_response += token_text
|
||
yield f"data: {json.dumps({'type': SSE_TYPE_TOKEN, 'token': token_text})}\n\n"
|
||
|
||
# 转发工具调用事件
|
||
elif event_type == SSE_TYPE_TOOL_CALL:
|
||
yield f"data: {json.dumps({'type': SSE_TYPE_TOOL_CALL, 'tool': data.get('tool', ''), 'name': data.get('name', '')})}\n\n"
|
||
|
||
elif event_type == SSE_TYPE_TOOL_RESULT:
|
||
yield f"data: {json.dumps({'type': SSE_TYPE_TOOL_RESULT, 'tool': data.get('tool', ''), 'result': str(data.get('result', ''))[:200]})}\n\n"
|
||
|
||
elif event_type == SSE_TYPE_TOOL_UNAVAILABLE:
|
||
yield f"data: {json.dumps({'type': SSE_TYPE_TOOL_UNAVAILABLE, 'tool': data.get('tool', '')})}\n\n"
|
||
|
||
except json.JSONDecodeError:
|
||
pass
|
||
|
||
# 流结束但未收到 [DONE] - 仍然保存
|
||
if full_response:
|
||
chat_memory.add_exchange(username, req.message, full_response)
|
||
yield f"data: {json.dumps({'type': SSE_TYPE_DONE, 'rounds': rounds + 1, 'max_rounds': settings.MAX_MEMORY_ROUNDS})}\n\n"
|
||
|
||
except httpx.ConnectError:
|
||
yield f"data: {json.dumps({'type': 'error', 'message': '无法连接到对话引擎(127.0.0.1:3000)，请确认Portal服务已启动'})}\n\n"
|
||
yield f"data: {json.dumps({'type': SSE_TYPE_DONE})}\n\n"
|
||
except httpx.TimeoutException:
|
||
# 超时但仍可能有部分回复
|
||
if full_response:
|
||
chat_memory.add_exchange(username, req.message, full_response)
|
||
yield f"data: {json.dumps({'type': 'error', 'message': '对话引擎响应超时，已保存部分回复'})}\n\n"
|
||
yield f"data: {json.dumps({'type': SSE_TYPE_DONE, 'rounds': rounds + 1, 'max_rounds': settings.MAX_MEMORY_ROUNDS, 'partial': True})}\n\n"
|
||
except Exception as e:
|
||
yield f"data: {json.dumps({'type': 'error', 'message': f'对话错误: {str(e)}'})}\n\n"
|
||
yield f"data: {json.dumps({'type': SSE_TYPE_DONE})}\n\n"
|
||
|
||
return StreamingResponse(
|
||
event_stream(),
|
||
media_type="text/event-stream",
|
||
headers={
|
||
"Cache-Control": "no-cache",
|
||
"Connection": "keep-alive",
|
||
"X-Accel-Buffering": "no",
|
||
}
|
||
)
|
||
|
||
@app.get("/api/corpus/chat/history")
|
||
async def get_chat_history(token: str = Query(...)):
|
||
"""获取当前对话历史"""
|
||
username = get_user_from_token(token)
|
||
if not username:
|
||
raise HTTPException(401, "登录已过期")
|
||
|
||
history = chat_memory.get_history(username)
|
||
rounds = chat_memory.get_rounds_count(username)
|
||
|
||
return {
|
||
"ok": True,
|
||
"username": username,
|
||
"rounds": rounds,
|
||
"max_rounds": settings.MAX_MEMORY_ROUNDS,
|
||
"history": history,
|
||
}
|
||
|
||
@app.post("/api/corpus/chat/clear")
|
||
async def clear_chat_history(token: str = Query(...)):
|
||
"""清空对话记忆"""
|
||
username = get_user_from_token(token)
|
||
if not username:
|
||
raise HTTPException(401, "登录已过期")
|
||
|
||
chat_memory.clear(username)
|
||
|
||
return {
|
||
"ok": True,
|
||
"message": "对话记忆已清空",
|
||
}
|
||
|
||
@app.get("/api/corpus/chat/models")
|
||
async def get_available_models(token: str = Query(...)):
|
||
"""获取可用的模型和人格体列表"""
|
||
username = get_user_from_token(token)
|
||
if not username:
|
||
raise HTTPException(401, "登录已过期")
|
||
|
||
return {
|
||
"ok": True,
|
||
"personas": [
|
||
{"id": "zhuyuan", "name": "铸渊", "desc": "现实执行层 · 系统决策、执行规划、技术实现"},
|
||
{"id": "shuangyan", "name": "霜砚", "desc": "语言主控架构层 · 理解语言、组织表达、控制输出质量"},
|
||
],
|
||
"engines": [
|
||
{"id": "deepseek", "name": "DeepSeek", "model": "deepseek-v4-pro", "provider": "DeepSeek"},
|
||
{"id": "zhipu", "name": "智谱清言", "model": "glm-4-plus", "provider": "智谱AI"},
|
||
{"id": "tongyi", "name": "通义千问", "model": "qwen-max", "provider": "阿里云"},
|
||
{"id": "huoshan", "name": "火山引擎", "model": "doubao-pro-32k", "provider": "字节跳动"},
|
||
],
|
||
"current_rounds": chat_memory.get_rounds_count(username),
|
||
"max_rounds": settings.MAX_MEMORY_ROUNDS,
|
||
}
|
||
|
||
# --- 统计数据 ---
|
||
|
||
@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),
|
||
"active_chat_memories": len(chat_memory._stores),
|
||
}
|
||
|
||
# ============================================================
|
||
# 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}")
|
||
print(f"🤖 对话引擎: {settings.PORTAL_CHAT_URL}")
|
||
print(f"💬 最大记忆轮数: {settings.MAX_MEMORY_ROUNDS}")
|
||
|
||
settings.USERS_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
||
uvicorn.run(app, host=settings.HOST, port=settings.PORT)
|
||
|
||
# ============================================================
|
||
# MCP 服务端（供 WorkBuddy 调用）
|
||
# ============================================================
|
||
from mcp.server.fastmcp import FastMCP
|
||
from mcp.server.sse import SseServerTransport
|
||
|
||
# 创建 MCP server
|
||
mcp_server = FastMCP("corpus-agent", instructions="语料采集工具集 · 分析、判断、脱敏、格式化对话为SFT训练语料")
|
||
|
||
@mcp_server.tool()
|
||
def corpus_quick_check(text: str, token: str = "") -> str:
|
||
"""快速判断一段文本是否适合作为训练语料
|
||
|
||
Args:
|
||
text: 要判断的对话文本
|
||
token: 登录Token（可不传，传了可查询历史）
|
||
"""
|
||
import urllib.request
|
||
try:
|
||
req = urllib.request.Request(
|
||
"http://127.0.0.1:8084/api/corpus/preview",
|
||
data=json.dumps({"text": text, "source": "mcp"}).encode(),
|
||
headers={"Content-Type": "application/json"}
|
||
)
|
||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||
return resp.read().decode()
|
||
except Exception as e:
|
||
return json.dumps({"error": str(e)}, ensure_ascii=False)
|
||
|
||
@mcp_server.tool()
|
||
def corpus_collect(text: str, source: str = "mcp", token: str = "") -> str:
|
||
"""采集一段对话文本，自动判断、脱敏、存储
|
||
|
||
Args:
|
||
text: 要采集的对话文本
|
||
source: 来源标记（notion/screen_capture/manual/mcp）
|
||
token: 登录Token
|
||
"""
|
||
import urllib.request
|
||
try:
|
||
req = urllib.request.Request(
|
||
f"http://127.0.0.1:8084/api/corpus/collect?token={token}",
|
||
data=json.dumps({"text": text, "source": source}).encode(),
|
||
headers={"Content-Type": "application/json"}
|
||
)
|
||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||
return resp.read().decode()
|
||
except Exception as e:
|
||
return json.dumps({"error": str(e)}, ensure_ascii=False)
|
||
|
||
@mcp_server.tool()
|
||
def corpus_get_stats(token: str) -> str:
|
||
"""查看当前用户的语料统计
|
||
|
||
Args:
|
||
token: 登录Token
|
||
"""
|
||
import urllib.request
|
||
try:
|
||
req = urllib.request.Request(f"http://127.0.0.1:8084/api/corpus/stats?token={token}")
|
||
with urllib.request.urlopen(req, timeout=10) as resp:
|
||
return resp.read().decode()
|
||
except Exception as e:
|
||
return json.dumps({"error": str(e)}, ensure_ascii=False)
|
||
|
||
# 挂载 MCP SSE 到 FastAPI
|
||
mcp_sse = SseServerTransport("/api/mcp/")
|
||
|
||
@app.get("/api/mcp/sse")
|
||
async def handle_sse(request: Request):
|
||
async with mcp_sse.connect_sse(
|
||
request.scope,
|
||
request.receive,
|
||
request._send
|
||
) as streams:
|
||
await mcp_server._mcp_server.run(
|
||
streams[0], streams[1],
|
||
mcp_server._mcp_server.create_initialization_options()
|
||
)
|
||
|
||
@app.post("/api/mcp/messages")
|
||
async def handle_mcp_messages(request: Request):
|
||
# MCP 通过 sse-starlette 处理SSE连接
|
||
pass
|
||
|