288 lines
11 KiB
Python
288 lines
11 KiB
Python
import base64
|
||
import json
|
||
import uuid
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
from fastapi import APIRouter, Depends, HTTPException, UploadFile
|
||
from fastapi.responses import StreamingResponse
|
||
from sqlalchemy import delete, select
|
||
from sqlalchemy.orm import Session
|
||
|
||
from app.config import get_settings
|
||
from app.db import SessionLocal, get_db
|
||
from app.deps import get_current_user
|
||
from app.integrations.adp_client import AdpClient
|
||
from app.integrations.storage_client import StorageClient
|
||
from app.models import Attachment, Conversation, Feedback, Message, MessageReference, User
|
||
from app.schemas.chat import (
|
||
AttachmentResponse,
|
||
ChatStreamRequest,
|
||
ConversationCreate,
|
||
ConversationResponse,
|
||
FeedbackCreate,
|
||
MessageResponse,
|
||
)
|
||
from app.services.audit_service import write_audit
|
||
|
||
|
||
def _attachment_data_url(att: Attachment) -> str | None:
|
||
storage_dir = Path(get_settings().local_storage_dir)
|
||
p = storage_dir / att.storage_key
|
||
if not p.exists():
|
||
return None
|
||
try:
|
||
raw = p.read_bytes()
|
||
except OSError:
|
||
return None
|
||
mime = att.file_type or "image/png"
|
||
return f"data:{mime};base64,{base64.b64encode(raw).decode()}"
|
||
|
||
router = APIRouter()
|
||
|
||
|
||
@router.get("/sessions", response_model=list[ConversationResponse])
|
||
def list_sessions(current_user: User = Depends(get_current_user), db: Session = Depends(get_db)) -> list[Conversation]:
|
||
return db.scalars(
|
||
select(Conversation).where(Conversation.user_id == current_user.id).order_by(Conversation.id.desc())
|
||
).all()
|
||
|
||
|
||
@router.post("/sessions", response_model=ConversationResponse)
|
||
def create_session(
|
||
payload: ConversationCreate,
|
||
current_user: User = Depends(get_current_user),
|
||
db: Session = Depends(get_db),
|
||
) -> Conversation:
|
||
conversation = Conversation(user_id=current_user.id, title=payload.title)
|
||
db.add(conversation)
|
||
write_audit(db, "chat.session.create", actor_id=current_user.id, target_type="conversation")
|
||
db.commit()
|
||
db.refresh(conversation)
|
||
return conversation
|
||
|
||
|
||
@router.delete("/sessions/{session_id}")
|
||
def delete_session(
|
||
session_id: int,
|
||
current_user: User = Depends(get_current_user),
|
||
db: Session = Depends(get_db),
|
||
) -> dict:
|
||
conversation = db.get(Conversation, session_id)
|
||
if not conversation or conversation.user_id != current_user.id:
|
||
raise HTTPException(status_code=404, detail="会话不存在")
|
||
message_ids = list(db.scalars(select(Message.id).where(Message.conversation_id == session_id)))
|
||
if message_ids:
|
||
db.execute(delete(MessageReference).where(MessageReference.message_id.in_(message_ids)))
|
||
db.execute(delete(Feedback).where(Feedback.message_id.in_(message_ids)))
|
||
db.execute(delete(Message).where(Message.conversation_id == session_id))
|
||
db.delete(conversation)
|
||
write_audit(db, "chat.session.delete", actor_id=current_user.id, target_type="conversation", target_id=str(session_id))
|
||
db.commit()
|
||
return {"message": "会话已删除"}
|
||
|
||
|
||
@router.get("/sessions/{session_id}/messages", response_model=list[MessageResponse])
|
||
def list_messages(
|
||
session_id: int,
|
||
current_user: User = Depends(get_current_user),
|
||
db: Session = Depends(get_db),
|
||
) -> list[MessageResponse]:
|
||
conversation = db.get(Conversation, session_id)
|
||
if not conversation or conversation.user_id != current_user.id:
|
||
raise HTTPException(status_code=404, detail="会话不存在")
|
||
messages = db.scalars(select(Message).where(Message.conversation_id == session_id).order_by(Message.id)).all()
|
||
message_ids = [m.id for m in messages]
|
||
atts_by_msg: dict[int, list[Attachment]] = {}
|
||
if message_ids:
|
||
for att in db.scalars(select(Attachment).where(Attachment.message_id.in_(message_ids))):
|
||
atts_by_msg.setdefault(att.message_id, []).append(att)
|
||
return [
|
||
MessageResponse(
|
||
id=message.id,
|
||
conversation_id=message.conversation_id,
|
||
question=message.question,
|
||
answer=message.answer,
|
||
status=message.status,
|
||
token_usage=message.token_usage,
|
||
pu_usage=message.pu_usage,
|
||
created_at=message.created_at,
|
||
references=_dedup_refs(message.references),
|
||
attachments=[
|
||
AttachmentResponse(id=att.id, file_type=att.file_type, data_url=_attachment_data_url(att))
|
||
for att in atts_by_msg.get(message.id, [])
|
||
],
|
||
)
|
||
for message in messages
|
||
]
|
||
|
||
|
||
def _dedup_refs(refs) -> list[dict]:
|
||
seen: set[tuple[str, str]] = set()
|
||
out: list[dict] = []
|
||
for ref in refs:
|
||
title = (ref.source_title or "").strip()
|
||
if not title:
|
||
continue
|
||
key = (title, (ref.snippet or "").strip())
|
||
if key in seen:
|
||
continue
|
||
seen.add(key)
|
||
out.append({
|
||
"source_title": ref.source_title,
|
||
"snippet": ref.snippet,
|
||
"page_no": ref.page_no,
|
||
"section": ref.section,
|
||
"source_url": ref.source_url,
|
||
})
|
||
return out
|
||
|
||
|
||
@router.post("/upload-image")
|
||
def upload_image(file: UploadFile, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)) -> dict:
|
||
if file.content_type not in {"image/png", "image/jpeg", "image/jpg"}:
|
||
raise HTTPException(status_code=400, detail="仅支持 PNG/JPG 截图")
|
||
raw = file.file.read()
|
||
storage = StorageClient()
|
||
key = storage.save_bytes(raw, suffix=Path(file.filename or "upload.bin").suffix)
|
||
try:
|
||
remote_url = AdpClient().upload_chat_image(
|
||
raw, mime_type=file.content_type or "image/png", original_name=file.filename or ""
|
||
)
|
||
except Exception as exc:
|
||
raise HTTPException(status_code=502, detail=f"图片上传到 ADP 失败: {exc}") from exc
|
||
attachment = Attachment(
|
||
user_id=current_user.id,
|
||
file_type=file.content_type or "image",
|
||
storage_key=key,
|
||
remote_url=remote_url,
|
||
original_name=file.filename,
|
||
)
|
||
db.add(attachment)
|
||
db.flush()
|
||
write_audit(db, "chat.image.upload", actor_id=current_user.id, target_type="attachment", target_id=str(attachment.id))
|
||
db.commit()
|
||
return {"id": attachment.id, "url": storage.temporary_url(key)}
|
||
|
||
|
||
@router.post("/stream")
|
||
def stream_chat(
|
||
payload: ChatStreamRequest,
|
||
current_user: User = Depends(get_current_user),
|
||
db: Session = Depends(get_db),
|
||
) -> StreamingResponse:
|
||
current_user_id = current_user.id
|
||
allowed_tags = [] if current_user.is_superuser else [role.code for role in current_user.roles]
|
||
|
||
if not (payload.question or "").strip() and not payload.attachment_ids:
|
||
raise HTTPException(status_code=400, detail="问题文本和图片至少需要一项")
|
||
|
||
image_urls: list[str] = []
|
||
attachments: list[Attachment] = []
|
||
if payload.attachment_ids:
|
||
attachments = list(db.scalars(
|
||
select(Attachment).where(
|
||
Attachment.id.in_(payload.attachment_ids),
|
||
Attachment.user_id == current_user_id,
|
||
)
|
||
))
|
||
if len(attachments) != len(payload.attachment_ids):
|
||
raise HTTPException(status_code=404, detail="附件不存在或无权访问")
|
||
for att in attachments:
|
||
if not att.remote_url:
|
||
raise HTTPException(status_code=500, detail=f"附件缺少远端 URL，请重新上传: {att.id}")
|
||
image_urls.append(att.remote_url)
|
||
|
||
conversation = db.get(Conversation, payload.conversation_id) if payload.conversation_id else None
|
||
if conversation and conversation.user_id != current_user_id:
|
||
raise HTTPException(status_code=404, detail="会话不存在")
|
||
if not conversation:
|
||
title_seed = (payload.question or "").strip()[:40] or ("[图片] " + (attachments[0].original_name or "提问"))
|
||
conversation = Conversation(
|
||
user_id=current_user_id,
|
||
title=title_seed[:60],
|
||
adp_conversation_id=uuid.uuid4().hex,
|
||
)
|
||
db.add(conversation)
|
||
db.flush()
|
||
elif not conversation.adp_conversation_id:
|
||
conversation.adp_conversation_id = uuid.uuid4().hex
|
||
db.flush()
|
||
|
||
message = Message(conversation_id=conversation.id, user_id=current_user_id, question=payload.question, status="streaming")
|
||
db.add(message)
|
||
conversation.last_message_at = datetime.utcnow()
|
||
db.flush()
|
||
message_id = message.id
|
||
conversation_id = conversation.id
|
||
adp_conversation_id = conversation.adp_conversation_id
|
||
for att in attachments:
|
||
att.message_id = message_id
|
||
db.commit()
|
||
|
||
def event_stream():
|
||
answer_parts: list[str] = []
|
||
usage = {"token_usage": 0, "pu_usage": 0}
|
||
references: list[dict] = []
|
||
seen_refs: set[tuple[str, str]] = set()
|
||
for event in AdpClient().stream_answer(
|
||
payload.question,
|
||
user_id=str(current_user_id),
|
||
conversation_id=adp_conversation_id,
|
||
allowed_tags=allowed_tags,
|
||
image_urls=image_urls,
|
||
):
|
||
event_name = event["event"]
|
||
if event_name == "answer.delta":
|
||
answer_parts.append(event["data"])
|
||
elif event_name == "answer.reference":
|
||
title = (event["data"].get("source_title") or "").strip()
|
||
if not title:
|
||
continue
|
||
key = (title, (event["data"].get("snippet") or "").strip())
|
||
if key in seen_refs:
|
||
continue
|
||
seen_refs.add(key)
|
||
references.append(event["data"])
|
||
elif event_name == "answer.usage":
|
||
usage.update(event["data"])
|
||
yield f"event: {event_name}\ndata: {json.dumps(event['data'], ensure_ascii=False)}\n\n"
|
||
|
||
save_db = SessionLocal()
|
||
try:
|
||
saved = save_db.get(Message, message_id)
|
||
if saved:
|
||
saved.answer = "".join(answer_parts)
|
||
saved.status = "done"
|
||
saved.token_usage = usage["token_usage"]
|
||
saved.pu_usage = usage["pu_usage"]
|
||
for ref in references:
|
||
save_db.add(MessageReference(message_id=message_id, **ref))
|
||
write_audit(save_db, "chat.ask", actor_id=current_user_id, target_type="message", target_id=str(message_id))
|
||
save_db.commit()
|
||
finally:
|
||
save_db.close()
|
||
|
||
return StreamingResponse(
|
||
event_stream(),
|
||
media_type="text/event-stream",
|
||
headers={"X-Conversation-Id": str(conversation_id), "X-Message-Id": str(message_id)},
|
||
)
|
||
|
||
|
||
@router.post("/messages/{message_id}/feedback")
|
||
def create_feedback(
|
||
message_id: int,
|
||
payload: FeedbackCreate,
|
||
current_user: User = Depends(get_current_user),
|
||
db: Session = Depends(get_db),
|
||
) -> dict:
|
||
message = db.get(Message, message_id)
|
||
if not message or message.user_id != current_user.id:
|
||
raise HTTPException(status_code=404, detail="消息不存在")
|
||
feedback = Feedback(message_id=message_id, user_id=current_user.id, rating=payload.rating, comment=payload.comment)
|
||
db.add(feedback)
|
||
write_audit(db, "chat.feedback", actor_id=current_user.id, target_type="message", target_id=str(message_id))
|
||
db.commit()
|
||
return {"message": "反馈已提交"}
|