472 lines
17 KiB
Python
472 lines
17 KiB
Python
import os
|
|
import json
|
|
import shutil
|
|
import time
|
|
import uuid
|
|
|
|
from fastapi import APIRouter, UploadFile, File, Depends, Header, HTTPException
|
|
from pydantic import BaseModel
|
|
from sqlalchemy.orm import Session
|
|
|
|
from database import get_db
|
|
from models import KnowledgeDoc, QAPair, KnowledgeChunk, Avatar, User
|
|
from responses import ok, fail
|
|
import embeddings
|
|
from services.knowledge_vectorizer import knowledge_vectorizer
|
|
|
|
router = APIRouter()
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
UPLOAD_DIR = os.path.abspath(os.getenv("UPLOAD_DIR", os.path.join(BASE_DIR, "uploads")))
|
|
os.makedirs(UPLOAD_DIR, exist_ok=True)
|
|
|
|
ALLOWED_EXT = {".md", ".txt", ".pdf", ".doc", ".docx", ".xlsx"}
|
|
MAX_UPLOAD_BYTES = 50 * 1024 * 1024
|
|
UPLOAD_CHUNK_BYTES = 1024 * 1024
|
|
MULTIPART_CHUNK_BYTES = 5 * 1024 * 1024
|
|
MULTIPART_ROOT = ".multipart"
|
|
MULTIPART_TTL_SECONDS = 24 * 60 * 60
|
|
|
|
|
|
class QAIn(BaseModel):
|
|
question: str = ""
|
|
answer: str = ""
|
|
enabled: bool = True
|
|
|
|
|
|
class EnabledIn(BaseModel):
|
|
enabled: bool = True
|
|
|
|
|
|
class MultipartUploadIn(BaseModel):
|
|
filename: str
|
|
fileSize: int
|
|
totalChunks: int
|
|
|
|
|
|
def _validate_document(filename: str, file_size: int):
|
|
ext = os.path.splitext(filename or "")[1].lower()
|
|
if ext not in ALLOWED_EXT:
|
|
return None, f"不支持的文件类型：{ext or '空'}，仅支持 md/txt/pdf/doc/docx/xlsx"
|
|
if file_size <= 0:
|
|
return None, "文件内容不能为空"
|
|
if file_size > MAX_UPLOAD_BYTES:
|
|
return None, "文件不能超过 50MB"
|
|
return ext, ""
|
|
|
|
|
|
def _multipart_dir(avatar_id: str, upload_id: str) -> str:
|
|
safe_avatar_id = os.path.basename(avatar_id)
|
|
safe_upload_id = os.path.basename(upload_id)
|
|
if (
|
|
safe_avatar_id != avatar_id
|
|
or safe_upload_id != upload_id
|
|
or len(upload_id) != 32
|
|
or any(character not in "0123456789abcdef" for character in upload_id)
|
|
):
|
|
raise HTTPException(status_code=400, detail="上传标识无效")
|
|
return os.path.join(UPLOAD_DIR, MULTIPART_ROOT, safe_avatar_id, safe_upload_id)
|
|
|
|
|
|
def _purge_stale_multipart_uploads(avatar_id: str):
|
|
avatar_upload_root = os.path.join(UPLOAD_DIR, MULTIPART_ROOT, os.path.basename(avatar_id))
|
|
if not os.path.isdir(avatar_upload_root):
|
|
return
|
|
cutoff = time.time() - MULTIPART_TTL_SECONDS
|
|
for entry in os.scandir(avatar_upload_root):
|
|
if entry.is_dir(follow_symlinks=False) and entry.stat(follow_symlinks=False).st_mtime < cutoff:
|
|
shutil.rmtree(entry.path, ignore_errors=True)
|
|
|
|
|
|
def _read_multipart_metadata(avatar_id: str, upload_id: str) -> tuple[str, dict]:
|
|
upload_dir = _multipart_dir(avatar_id, upload_id)
|
|
metadata_path = os.path.join(upload_dir, "metadata.json")
|
|
if not os.path.isfile(metadata_path):
|
|
raise HTTPException(status_code=404, detail="上传任务不存在或已过期")
|
|
with open(metadata_path, "r", encoding="utf-8") as stream:
|
|
return upload_dir, json.load(stream)
|
|
|
|
|
|
def _create_knowledge_doc(db: Session, avatar_id: str, filename: str, ext: str, file_size: int, stored: str):
|
|
doc = KnowledgeDoc(
|
|
id=uuid.uuid4().hex,
|
|
avatar_id=avatar_id,
|
|
filename=filename,
|
|
file_type=ext.lstrip("."),
|
|
file_size=file_size,
|
|
file_url=f"/api/files/{avatar_id}/{stored}",
|
|
status="parsing",
|
|
index_stage="queued",
|
|
index_progress=0,
|
|
)
|
|
db.add(doc)
|
|
db.commit()
|
|
db.refresh(doc)
|
|
knowledge_vectorizer.enqueue(doc.id)
|
|
return doc
|
|
|
|
|
|
def _doc_payload(doc: KnowledgeDoc) -> dict:
|
|
payload = doc.to_dict()
|
|
stored_name = os.path.basename(doc.file_url or "")
|
|
stored_path = os.path.join(UPLOAD_DIR, doc.avatar_id, stored_name)
|
|
payload["filePresent"] = bool(stored_name and os.path.isfile(stored_path))
|
|
return payload
|
|
|
|
|
|
def _resolve_user(authorization: str | None, db: Session):
|
|
if not authorization:
|
|
return None
|
|
token = authorization.replace("Bearer ", "", 1).replace("bearer ", "", 1).strip()
|
|
return db.query(User).filter(User.app_token == token).first()
|
|
|
|
|
|
def _require_owned_avatar(db: Session, avatar_id: str, authorization: str | None):
|
|
avatar = db.query(Avatar).filter(Avatar.id == avatar_id).first()
|
|
if not avatar:
|
|
raise HTTPException(status_code=404, detail="avatar not found")
|
|
user = _resolve_user(authorization, db)
|
|
if not user:
|
|
raise HTTPException(status_code=401, detail="未登录")
|
|
if avatar.owner_id and avatar.owner_id != user.huihui_user_id:
|
|
raise HTTPException(status_code=403, detail="无权访问该分身")
|
|
return avatar
|
|
|
|
|
|
# ---------------- Documents ----------------
|
|
@router.get("/avatar/{avatar_id}/knowledge/docs")
|
|
def list_docs(avatar_id: str, authorization: str = Header(None), db: Session = Depends(get_db)):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
docs = (
|
|
db.query(KnowledgeDoc)
|
|
.filter(KnowledgeDoc.avatar_id == avatar_id)
|
|
.order_by(KnowledgeDoc.created_at.desc())
|
|
.all()
|
|
)
|
|
return ok([_doc_payload(d) for d in docs])
|
|
|
|
|
|
@router.post("/avatar/{avatar_id}/knowledge/docs")
|
|
async def upload_doc(avatar_id: str, file: UploadFile = File(...), authorization: str = Header(None), db: Session = Depends(get_db)):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
ext, validation_error = _validate_document(file.filename or "", 1)
|
|
if validation_error:
|
|
return fail(validation_error, code=400)
|
|
avatar_dir = os.path.join(UPLOAD_DIR, avatar_id)
|
|
os.makedirs(avatar_dir, exist_ok=True)
|
|
stored = f"{uuid.uuid4().hex}{ext}"
|
|
path = os.path.join(avatar_dir, stored)
|
|
file_size = 0
|
|
try:
|
|
# Stream large files to disk so a 100MB upload does not occupy 100MB RAM.
|
|
with open(path, "wb") as f:
|
|
while chunk := await file.read(UPLOAD_CHUNK_BYTES):
|
|
file_size += len(chunk)
|
|
if file_size > MAX_UPLOAD_BYTES:
|
|
raise ValueError("文件不能超过 50MB")
|
|
f.write(chunk)
|
|
except ValueError as exc:
|
|
if os.path.exists(path):
|
|
os.remove(path)
|
|
return fail(str(exc), code=400)
|
|
if file_size == 0:
|
|
if os.path.exists(path):
|
|
os.remove(path)
|
|
return fail("文件内容不能为空", code=400)
|
|
|
|
doc = _create_knowledge_doc(db, avatar_id, file.filename or stored, ext, file_size, stored)
|
|
return ok(_doc_payload(doc))
|
|
|
|
|
|
@router.post("/avatar/{avatar_id}/knowledge/uploads")
|
|
def create_multipart_upload(
|
|
avatar_id: str,
|
|
body: MultipartUploadIn,
|
|
authorization: str = Header(None),
|
|
db: Session = Depends(get_db),
|
|
):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
ext, validation_error = _validate_document(body.filename, body.fileSize)
|
|
if validation_error:
|
|
return fail(validation_error, code=400)
|
|
expected_chunks = (body.fileSize + MULTIPART_CHUNK_BYTES - 1) // MULTIPART_CHUNK_BYTES
|
|
if body.totalChunks != expected_chunks:
|
|
return fail("文件分片数量不正确", code=400)
|
|
|
|
_purge_stale_multipart_uploads(avatar_id)
|
|
upload_id = uuid.uuid4().hex
|
|
upload_dir = _multipart_dir(avatar_id, upload_id)
|
|
os.makedirs(upload_dir, exist_ok=False)
|
|
metadata = {
|
|
"filename": body.filename,
|
|
"fileSize": body.fileSize,
|
|
"totalChunks": body.totalChunks,
|
|
"extension": ext,
|
|
}
|
|
with open(os.path.join(upload_dir, "metadata.json"), "w", encoding="utf-8") as stream:
|
|
json.dump(metadata, stream, ensure_ascii=False)
|
|
return ok({"uploadId": upload_id, "chunkSize": MULTIPART_CHUNK_BYTES})
|
|
|
|
|
|
@router.post("/avatar/{avatar_id}/knowledge/uploads/{upload_id}/chunks/{chunk_index}")
|
|
async def upload_multipart_chunk(
|
|
avatar_id: str,
|
|
upload_id: str,
|
|
chunk_index: int,
|
|
file: UploadFile = File(...),
|
|
authorization: str = Header(None),
|
|
db: Session = Depends(get_db),
|
|
):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
upload_dir, metadata = _read_multipart_metadata(avatar_id, upload_id)
|
|
total_chunks = int(metadata["totalChunks"])
|
|
if chunk_index < 0 or chunk_index >= total_chunks:
|
|
return fail("文件分片序号不正确", code=400)
|
|
|
|
expected_size = min(
|
|
MULTIPART_CHUNK_BYTES,
|
|
int(metadata["fileSize"]) - chunk_index * MULTIPART_CHUNK_BYTES,
|
|
)
|
|
part_path = os.path.join(upload_dir, f"{chunk_index}.part")
|
|
temporary_path = f"{part_path}.uploading"
|
|
received = 0
|
|
try:
|
|
with open(temporary_path, "wb") as stream:
|
|
while chunk := await file.read(UPLOAD_CHUNK_BYTES):
|
|
received += len(chunk)
|
|
if received > expected_size:
|
|
raise ValueError("文件分片大小不正确")
|
|
stream.write(chunk)
|
|
if received != expected_size:
|
|
raise ValueError("文件分片大小不正确")
|
|
os.replace(temporary_path, part_path)
|
|
except ValueError as exc:
|
|
if os.path.exists(temporary_path):
|
|
os.remove(temporary_path)
|
|
return fail(str(exc), code=400)
|
|
return ok({"chunkIndex": chunk_index, "uploadedBytes": received})
|
|
|
|
|
|
@router.post("/avatar/{avatar_id}/knowledge/uploads/{upload_id}/complete")
|
|
def complete_multipart_upload(
|
|
avatar_id: str,
|
|
upload_id: str,
|
|
authorization: str = Header(None),
|
|
db: Session = Depends(get_db),
|
|
):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
upload_dir, metadata = _read_multipart_metadata(avatar_id, upload_id)
|
|
total_chunks = int(metadata["totalChunks"])
|
|
part_paths = [os.path.join(upload_dir, f"{index}.part") for index in range(total_chunks)]
|
|
if not all(os.path.isfile(path) for path in part_paths):
|
|
return fail("文件分片尚未上传完整", code=400)
|
|
if sum(os.path.getsize(path) for path in part_paths) != int(metadata["fileSize"]):
|
|
return fail("文件分片总大小不正确", code=400)
|
|
|
|
avatar_dir = os.path.join(UPLOAD_DIR, avatar_id)
|
|
os.makedirs(avatar_dir, exist_ok=True)
|
|
stored = f"{uuid.uuid4().hex}{metadata['extension']}"
|
|
final_path = os.path.join(avatar_dir, stored)
|
|
temporary_path = f"{final_path}.assembling"
|
|
try:
|
|
with open(temporary_path, "wb") as output:
|
|
for part_path in part_paths:
|
|
with open(part_path, "rb") as source:
|
|
shutil.copyfileobj(source, output, UPLOAD_CHUNK_BYTES)
|
|
os.replace(temporary_path, final_path)
|
|
doc = _create_knowledge_doc(
|
|
db,
|
|
avatar_id,
|
|
metadata["filename"],
|
|
metadata["extension"],
|
|
int(metadata["fileSize"]),
|
|
stored,
|
|
)
|
|
except Exception:
|
|
if os.path.exists(temporary_path):
|
|
os.remove(temporary_path)
|
|
raise
|
|
shutil.rmtree(upload_dir, ignore_errors=True)
|
|
return ok(_doc_payload(doc))
|
|
|
|
|
|
@router.post("/avatar/{avatar_id}/knowledge/docs/{doc_id}/retry")
|
|
def retry_doc(avatar_id: str, doc_id: str, authorization: str = Header(None), db: Session = Depends(get_db)):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
doc = db.query(KnowledgeDoc).filter(
|
|
KnowledgeDoc.id == doc_id, KnowledgeDoc.avatar_id == avatar_id
|
|
).first()
|
|
if not doc:
|
|
return fail("文档不存在", code=404)
|
|
if doc.vectorized and doc.status == "ready":
|
|
return ok(_doc_payload(doc))
|
|
stored_name = os.path.basename(doc.file_url or "")
|
|
if not stored_name or not os.path.isfile(os.path.join(UPLOAD_DIR, avatar_id, stored_name)):
|
|
return fail("原文件不可用，请重新上传", code=400)
|
|
db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == doc.id).delete()
|
|
doc.status = "parsing"
|
|
doc.vectorized = False
|
|
doc.embedding_model = ""
|
|
doc.chunk_count = 0
|
|
doc.vectorized_at = None
|
|
doc.error_message = ""
|
|
doc.index_stage = "queued"
|
|
doc.index_progress = 0
|
|
db.commit()
|
|
db.refresh(doc)
|
|
knowledge_vectorizer.enqueue(doc.id)
|
|
return ok(_doc_payload(doc))
|
|
|
|
|
|
@router.delete("/avatar/{avatar_id}/knowledge/docs/{doc_id}")
|
|
def delete_doc(avatar_id: str, doc_id: str, authorization: str = Header(None), db: Session = Depends(get_db)):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
doc = (
|
|
db.query(KnowledgeDoc)
|
|
.filter(KnowledgeDoc.id == doc_id, KnowledgeDoc.avatar_id == avatar_id)
|
|
.first()
|
|
)
|
|
if not doc:
|
|
return fail("文档不存在", code=404)
|
|
# 级联删除切片
|
|
db.query(KnowledgeChunk).filter(KnowledgeChunk.doc_id == doc_id).delete()
|
|
try:
|
|
fp = os.path.join(UPLOAD_DIR, avatar_id, os.path.basename(doc.file_url))
|
|
if os.path.exists(fp):
|
|
os.remove(fp)
|
|
except Exception:
|
|
pass
|
|
db.delete(doc)
|
|
db.commit()
|
|
return ok({"id": doc_id})
|
|
|
|
|
|
# ---------------- 向量检索 ----------------
|
|
@router.get("/avatar/{avatar_id}/knowledge/search")
|
|
def search_knowledge(avatar_id: str, q: str = "", top_k: int = 5, authorization: str = Header(None), db: Session = Depends(get_db)):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
q = (q or "").strip()
|
|
if not q:
|
|
return ok([])
|
|
chunks = (
|
|
db.query(KnowledgeChunk)
|
|
.filter(KnowledgeChunk.avatar_id == avatar_id)
|
|
.all()
|
|
)
|
|
if not chunks:
|
|
return ok([])
|
|
qvec = embeddings.embed([q])[0]
|
|
scored = []
|
|
for c in chunks:
|
|
try:
|
|
vec = json.loads(c.vector)
|
|
except Exception:
|
|
continue
|
|
scored.append((embeddings.cosine(qvec, vec), c))
|
|
scored.sort(key=lambda x: x[0], reverse=True)
|
|
results = []
|
|
for score, c in scored[: max(1, top_k)]:
|
|
doc = db.query(KnowledgeDoc).filter(KnowledgeDoc.id == c.doc_id).first()
|
|
snippet = c.content[:120] + ("…" if len(c.content) > 120 else "")
|
|
results.append(
|
|
{
|
|
"docId": c.doc_id,
|
|
"filename": doc.filename if doc else "",
|
|
"fileType": doc.file_type if doc else "",
|
|
"snippet": snippet,
|
|
"score": round(score, 4),
|
|
}
|
|
)
|
|
return ok(results)
|
|
|
|
|
|
# ---------------- Standard Q&A pairs ----------------
|
|
@router.get("/avatar/{avatar_id}/knowledge/qa")
|
|
def list_qa(avatar_id: str, authorization: str = Header(None), db: Session = Depends(get_db)):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
items = (
|
|
db.query(QAPair)
|
|
.filter(QAPair.avatar_id == avatar_id)
|
|
.order_by(QAPair.created_at.desc())
|
|
.all()
|
|
)
|
|
return ok([q.to_dict() for q in items])
|
|
|
|
|
|
@router.post("/avatar/{avatar_id}/knowledge/qa")
|
|
def create_qa(avatar_id: str, body: QAIn, authorization: str = Header(None), db: Session = Depends(get_db)):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
q = QAPair(
|
|
avatar_id=avatar_id,
|
|
question=body.question,
|
|
answer=body.answer,
|
|
enabled=body.enabled,
|
|
)
|
|
db.add(q)
|
|
db.commit()
|
|
db.refresh(q)
|
|
return ok(q.to_dict())
|
|
|
|
|
|
@router.put("/avatar/{avatar_id}/knowledge/qa/{qa_id}")
|
|
def update_qa(avatar_id: str, qa_id: str, body: QAIn, authorization: str = Header(None), db: Session = Depends(get_db)):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
q = (
|
|
db.query(QAPair)
|
|
.filter(QAPair.id == qa_id, QAPair.avatar_id == avatar_id)
|
|
.first()
|
|
)
|
|
if not q:
|
|
return fail("问答对不存在", code=404)
|
|
q.question = body.question
|
|
q.answer = body.answer
|
|
q.enabled = body.enabled
|
|
db.commit()
|
|
db.refresh(q)
|
|
return ok(q.to_dict())
|
|
|
|
|
|
@router.put("/avatar/{avatar_id}/knowledge/qa/{qa_id}/enabled")
|
|
def set_qa_enabled(avatar_id: str, qa_id: str, body: EnabledIn, authorization: str = Header(None), db: Session = Depends(get_db)):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
q = (
|
|
db.query(QAPair)
|
|
.filter(QAPair.id == qa_id, QAPair.avatar_id == avatar_id)
|
|
.first()
|
|
)
|
|
if not q:
|
|
return fail("问答对不存在", code=404)
|
|
q.enabled = bool(body.enabled)
|
|
db.commit()
|
|
db.refresh(q)
|
|
return ok(q.to_dict())
|
|
|
|
|
|
@router.delete("/avatar/{avatar_id}/knowledge/qa/{qa_id}")
|
|
def delete_qa(avatar_id: str, qa_id: str, authorization: str = Header(None), db: Session = Depends(get_db)):
|
|
_require_owned_avatar(db, avatar_id, authorization)
|
|
q = (
|
|
db.query(QAPair)
|
|
.filter(QAPair.id == qa_id, QAPair.avatar_id == avatar_id)
|
|
.first()
|
|
)
|
|
if not q:
|
|
return fail("问答对不存在", code=404)
|
|
db.delete(q)
|
|
db.commit()
|
|
return ok({"id": qa_id})
|
|
|
|
|
|
# ---------------- HuiHui user profile (mock; plug real interface via HUIHUI_USER_API) ----------------
|
|
@router.get("/user/profile")
|
|
def user_profile():
|
|
# 接入真实会会接口：设置环境变量 HUIHUI_USER_API 后在此请求并映射字段
|
|
api = os.getenv("HUIHUI_USER_API")
|
|
if api:
|
|
# TODO: 调用会会用户接口，返回 { userId, nickname, avatarUrl }
|
|
pass
|
|
return ok({
|
|
"userId": "hh_10001",
|
|
"nickname": "会会用户",
|
|
"avatarUrl": "https://api.dicebear.com/7.x/initials/svg?seed=HuiHui&backgroundColor=F97316",
|
|
})
|