297 lines
12 KiB
Python
297 lines
12 KiB
Python
import secrets
|
||
from pathlib import Path
|
||
from urllib.parse import quote
|
||
|
||
from flask import Flask, g, jsonify, redirect, request, send_file
|
||
|
||
from auth_service import (
|
||
build_dingtalk_login_url,
|
||
exchange_dingtalk_user,
|
||
get_user_by_session_token,
|
||
revoke_session,
|
||
upsert_user_and_session,
|
||
)
|
||
from config import Config
|
||
from database import init_database
|
||
from rule_service import (
|
||
delete_rule_document,
|
||
get_rule_document,
|
||
get_rule_html_path,
|
||
save_rule_document,
|
||
)
|
||
from task_service import (
|
||
TaskError,
|
||
cancel_task,
|
||
create_task,
|
||
get_file,
|
||
get_task,
|
||
get_task_extracted,
|
||
list_tasks,
|
||
recollect_task,
|
||
resume_task,
|
||
resume_unfinished_tasks,
|
||
)
|
||
|
||
|
||
def create_app() -> Flask:
|
||
app = Flask(__name__)
|
||
init_database()
|
||
|
||
@app.before_request
|
||
def assign_request_id():
|
||
g.request_id = request.headers.get("X-Request-Id") or f"req_{secrets.token_hex(10)}"
|
||
if request.method == "OPTIONS":
|
||
return None
|
||
open_paths = ("/health", "/api/v1/auth/", "/api/v1/skill/")
|
||
if request.path == "/health" or any(request.path.startswith(path) for path in open_paths):
|
||
return None
|
||
user = get_user_by_session_token(request.cookies.get(Config.SESSION_COOKIE_NAME))
|
||
if user is None and request.path.startswith("/api/"):
|
||
return jsonify({
|
||
"error": {"code": "UNAUTHORIZED", "message": "请先通过钉钉授权登录"},
|
||
"request_id": g.request_id,
|
||
}), 401
|
||
g.current_user = user
|
||
|
||
@app.after_request
|
||
def add_headers(response):
|
||
response.headers["X-Request-Id"] = g.get("request_id", "")
|
||
response.headers["Access-Control-Allow-Origin"] = Config.CORS_ORIGIN
|
||
response.headers["Access-Control-Allow-Headers"] = "Content-Type,Authorization,Idempotency-Key,X-Request-Id"
|
||
response.headers["Access-Control-Allow-Methods"] = "GET,POST,DELETE,OPTIONS"
|
||
response.headers["Access-Control-Allow-Credentials"] = "true"
|
||
return response
|
||
|
||
@app.errorhandler(TaskError)
|
||
def handle_task_error(error: TaskError):
|
||
return jsonify({
|
||
"error": {"code": error.code, "message": error.message, "details": error.details},
|
||
"request_id": g.request_id,
|
||
}), error.status_code
|
||
|
||
@app.errorhandler(Exception)
|
||
def handle_unexpected(error: Exception):
|
||
app.logger.exception("未处理异常")
|
||
return jsonify({
|
||
"error": {"code": "INTERNAL_ERROR", "message": "服务器内部错误"},
|
||
"request_id": g.get("request_id"),
|
||
}), 500
|
||
|
||
@app.route("/health", methods=["GET"])
|
||
def health():
|
||
return jsonify({"status": "ok", "request_id": g.request_id})
|
||
|
||
@app.route("/api/v1/auth/dingtalk/login", methods=["GET"])
|
||
def dingtalk_login():
|
||
return jsonify({
|
||
"data": build_dingtalk_login_url(request.args.get("redirect", "/tasks")),
|
||
"request_id": g.request_id,
|
||
})
|
||
|
||
@app.route("/api/v1/auth/dingtalk/callback", methods=["GET"])
|
||
def dingtalk_callback():
|
||
code = request.args.get("authCode") or request.args.get("code")
|
||
state = request.args.get("state", "")
|
||
frontend_redirect = "/tasks"
|
||
if ":" in state:
|
||
frontend_redirect = state.split(":", 1)[1] or frontend_redirect
|
||
if not code:
|
||
return redirect(f"{Config.FRONTEND_BASE_URL}/login?error={quote('钉钉授权缺少 code')}")
|
||
user_payload = exchange_dingtalk_user(code)
|
||
user, token, expires_at = upsert_user_and_session(user_payload)
|
||
response = redirect(f"{Config.FRONTEND_BASE_URL}{frontend_redirect}")
|
||
response.set_cookie(
|
||
Config.SESSION_COOKIE_NAME,
|
||
token,
|
||
expires=expires_at,
|
||
httponly=True,
|
||
secure=Config.SESSION_COOKIE_SECURE,
|
||
samesite=Config.SESSION_COOKIE_SAMESITE,
|
||
path="/",
|
||
)
|
||
# 给前端“Link AI”复制提示词使用。后端鉴权仍读取 hc_session；
|
||
# 前端只读取这个辅助 Cookie 的 value，并在提示词中拼成 hc_session=<value>。
|
||
response.set_cookie(
|
||
"hc_session_agent",
|
||
token,
|
||
expires=expires_at,
|
||
httponly=False,
|
||
secure=Config.SESSION_COOKIE_SECURE,
|
||
samesite=Config.SESSION_COOKIE_SAMESITE,
|
||
path="/",
|
||
)
|
||
app.logger.info("钉钉用户登录成功 user_id=%s name=%s", user["id"], user["name"])
|
||
return response
|
||
|
||
@app.route("/api/v1/auth/direct-token", methods=["GET"])
|
||
def auth_direct_token():
|
||
"""直接用钉钉 authCode 换取 session token，AI 无需浏览器 Cookie。
|
||
|
||
调用方式：
|
||
GET /api/v1/auth/direct-token?authCode=钉钉返回的code
|
||
|
||
返回：
|
||
{"data": {"token": "session_token", "user": {...}, "expires_at": "..."}}
|
||
|
||
AI 拿到 token 后，后续请求带上：
|
||
Cookie: hc_session=<token>
|
||
"""
|
||
code = request.args.get("authCode") or request.args.get("code")
|
||
if not code:
|
||
return jsonify({"code": "MISSING_AUTH_CODE", "msg": "缺少 authCode 参数"}), 400
|
||
try:
|
||
dingtalk_user = exchange_dingtalk_user(code)
|
||
user, token, expires_at_dt = upsert_user_and_session(dingtalk_user)
|
||
return jsonify({
|
||
"data": {
|
||
"token": token,
|
||
"user": user,
|
||
"expires_at": expires_at_dt.isoformat(timespec="seconds"),
|
||
},
|
||
"request_id": g.request_id,
|
||
})
|
||
except TaskError as e:
|
||
return jsonify({"code": e.code, "msg": e.message, "data": e.data}), e.http_status
|
||
|
||
# ── Skill 版本管理 ────────────────────────────────────────────
|
||
SKILL_BASE = Path(__file__).parent / "skills"
|
||
|
||
@app.route("/api/v1/skill/holy-crab/version", methods=["GET"])
|
||
def skill_holy_crab_version():
|
||
"""返回当前 Holy Crab Skill 的版本信息。"""
|
||
import json
|
||
skill_dir = SKILL_BASE / "holy-crab"
|
||
version_file = skill_dir / "VERSION"
|
||
if not version_file.exists():
|
||
return jsonify({"code": "SKILL_NOT_FOUND", "msg": "Skill 文件不存在"}), 404
|
||
info = json.loads(version_file.read_text(encoding="utf-8"))
|
||
return jsonify({"data": {"version": info}, "request_id": g.request_id})
|
||
|
||
@app.route("/api/v1/skill/holy-crab/download", methods=["GET"])
|
||
def skill_holy_crab_download():
|
||
"""下载 Holy Crab Skill 完整 zip 包（含所有文件）。"""
|
||
import hashlib, io, zipfile
|
||
skill_dir = SKILL_BASE / "holy-crab"
|
||
if not skill_dir.exists():
|
||
return jsonify({"code": "SKILL_NOT_FOUND", "msg": "Skill 文件不存在"}), 404
|
||
|
||
buf = io.BytesIO()
|
||
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
|
||
for fpath in skill_dir.rglob("*"):
|
||
if fpath.is_file():
|
||
arcname = str(fpath.relative_to(skill_dir))
|
||
zf.writestr(arcname, fpath.read_bytes())
|
||
buf.seek(0)
|
||
return buf.getvalue(), 200, {
|
||
"Content-Type": "application/zip",
|
||
"Content-Disposition": "attachment; filename=holy-crab-skill.zip",
|
||
"X-Skill-Version": json.loads((skill_dir / "VERSION").read_text(encoding="utf-8")).get("version", "unknown"),
|
||
}
|
||
|
||
@app.route("/api/v1/auth/me", methods=["GET"])
|
||
def auth_me():
|
||
user = get_user_by_session_token(request.cookies.get(Config.SESSION_COOKIE_NAME))
|
||
if not user:
|
||
raise TaskError("UNAUTHORIZED", "请先通过钉钉授权登录", 401)
|
||
return jsonify({"data": {"user": user}, "request_id": g.request_id})
|
||
|
||
@app.route("/api/v1/auth/logout", methods=["POST"])
|
||
def auth_logout():
|
||
revoke_session(request.cookies.get(Config.SESSION_COOKIE_NAME))
|
||
response = jsonify({"data": {"ok": True}, "request_id": g.request_id})
|
||
response.delete_cookie(Config.SESSION_COOKIE_NAME, path="/")
|
||
response.delete_cookie("hc_session_agent", path="/")
|
||
return response
|
||
|
||
@app.route("/api/v1/tasks", methods=["POST"])
|
||
def create_research_task():
|
||
payload = request.get_json(silent=True) or {}
|
||
current_user = getattr(g, "current_user", None) or {}
|
||
payload["created_by_id"] = current_user.get("id") or payload.get("created_by_id")
|
||
payload["created_by_name"] = current_user.get("name") or payload.get("created_by_name")
|
||
task = create_task(payload)
|
||
return jsonify({"data": task, "request_id": g.request_id}), 201
|
||
|
||
@app.route("/api/v1/tasks", methods=["GET"])
|
||
def query_research_tasks():
|
||
result = list_tasks(
|
||
q=request.args.get("q", "").strip(),
|
||
status=request.args.get("status", "").strip(),
|
||
page=int(request.args.get("page", "1")),
|
||
page_size=int(request.args.get("page_size", "10")),
|
||
)
|
||
return jsonify({"data": result, "request_id": g.request_id})
|
||
|
||
@app.route("/api/v1/tasks/<task_id>", methods=["GET"])
|
||
def get_research_task(task_id: str):
|
||
if request.args.get("view") in {"extracted", "cleaned", "parsed"}:
|
||
return jsonify({"data": get_task_extracted(task_id), "request_id": g.request_id})
|
||
return jsonify({"data": get_task(task_id), "request_id": g.request_id})
|
||
|
||
@app.route("/api/v1/tasks/<task_id>/extracted", methods=["GET"])
|
||
def get_research_task_extracted(task_id: str):
|
||
return jsonify({"data": get_task_extracted(task_id), "request_id": g.request_id})
|
||
|
||
@app.route("/api/v1/tasks/<task_id>/cancel", methods=["POST"])
|
||
def cancel_research_task(task_id: str):
|
||
body = request.get_json(silent=True) or {}
|
||
return jsonify({
|
||
"data": cancel_task(task_id, str(body.get("reason") or "")),
|
||
"request_id": g.request_id,
|
||
})
|
||
|
||
@app.route("/api/v1/tasks/<task_id>/retry", methods=["POST"])
|
||
@app.route("/api/v1/tasks/<task_id>/resume", methods=["POST"])
|
||
def resume_research_task(task_id: str):
|
||
return jsonify({"data": resume_task(task_id), "request_id": g.request_id})
|
||
|
||
@app.route("/api/v1/tasks/<task_id>/recollect", methods=["POST"])
|
||
def recollect_research_task(task_id: str):
|
||
return jsonify({
|
||
"data": recollect_task(task_id),
|
||
"request_id": g.request_id,
|
||
}), 201
|
||
|
||
@app.route("/api/v1/tasks/<task_id>/files", methods=["GET"])
|
||
def list_task_files(task_id: str):
|
||
task = get_task(task_id)
|
||
return jsonify({"data": {"items": task["files"]}, "request_id": g.request_id})
|
||
|
||
@app.route("/api/v1/tasks/<task_id>/files/<file_id>/preview", methods=["GET"])
|
||
def preview_task_file(task_id: str, file_id: str):
|
||
metadata, path = get_file(task_id, file_id)
|
||
if not metadata["previewable"]:
|
||
raise TaskError("FILE_TOO_LARGE", "文件超过预览大小限制", 413)
|
||
return send_file(path, as_attachment=False)
|
||
|
||
@app.route("/api/v1/tasks/<task_id>/files/<file_id>/download", methods=["GET"])
|
||
def download_task_file(task_id: str, file_id: str):
|
||
_, path = get_file(task_id, file_id)
|
||
return send_file(path, as_attachment=True, download_name=Path(path).name)
|
||
|
||
@app.route("/api/v1/rules", methods=["GET"])
|
||
def get_current_rule():
|
||
return jsonify({"data": get_rule_document(), "request_id": g.request_id})
|
||
|
||
@app.route("/api/v1/rules/current.html", methods=["GET"])
|
||
def preview_current_rule_html():
|
||
return send_file(get_rule_html_path(), mimetype="text/html; charset=utf-8")
|
||
|
||
@app.route("/api/v1/rules/html", methods=["POST"])
|
||
def upload_rule_html():
|
||
return jsonify({"data": save_rule_document(request), "request_id": g.request_id})
|
||
|
||
@app.route("/api/v1/rules/html", methods=["DELETE"])
|
||
def remove_rule_html():
|
||
return jsonify({"data": delete_rule_document(), "request_id": g.request_id})
|
||
|
||
resume_unfinished_tasks()
|
||
return app
|
||
|
||
|
||
app = create_app()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
app.run(host=Config.HOST, port=Config.PORT, threaded=True)
|