from __future__ import annotations import io from urllib.parse import quote from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile from fastapi.responses import StreamingResponse from pydantic import BaseModel from server.utils.auth_middleware import get_required_user from yuxi import knowledge_base from yuxi.services.workspace_service import ( create_workspace_directory, delete_workspace_path, download_workspace_file, list_workspace_tree, read_workspace_file_content, upload_workspace_file, write_workspace_file_content, ) from yuxi.storage.postgres.models_business import User workspace = APIRouter(prefix="/workspace", tags=["workspace"]) class CreateWorkspaceDirectoryRequest(BaseModel): parent_path: str name: str class UpdateWorkspaceFileContentRequest(BaseModel): path: str content: str async def _ensure_knowledge_read_access(current_user: User, kb_id: str) -> None: allowed = await knowledge_base.check_accessible( { "uid": current_user.uid, "role": current_user.role, "department_id": current_user.department_id, }, kb_id, ) if not allowed: raise HTTPException(status_code=403, detail="Access denied") def _raise_knowledge_read_error(error: ValueError) -> None: message = str(error) or "知识库文件读取失败" if message.startswith("Dify 知识库不支持"): raise HTTPException(status_code=501, detail=message) from error raise HTTPException(status_code=400, detail=message) from error @workspace.get("/tree", response_model=dict) async def get_workspace_tree( path: str = Query("/", description="工作区目录路径"), recursive: bool = Query(False, description="是否递归返回子目录文件"), files_only: bool = Query(False, description="是否仅返回文件"), current_user: User = Depends(get_required_user), ): return await list_workspace_tree( path=path, recursive=recursive, files_only=files_only, current_user=current_user, ) @workspace.get("/file", response_model=dict) async def get_workspace_file( path: str = Query(..., description="工作区文件路径"), current_user: User = Depends(get_required_user), ): return await read_workspace_file_content(path=path, current_user=current_user) @workspace.get("/knowledge/tree", response_model=dict) async def get_workspace_knowledge_tree( kb_id: str = Query(..., description="知识库 ID"), parent_id: str | None = Query(None, description="父文件夹 ID"), recursive: bool = Query(False, description="是否递归返回子目录文件"), files_only: bool = Query(False, description="是否仅返回文件"), current_user: User = Depends(get_required_user), ): await _ensure_knowledge_read_access(current_user, kb_id) try: return await knowledge_base.list_file_tree( kb_id=kb_id, parent_id=parent_id, recursive=recursive, files_only=files_only, ) except ValueError as error: _raise_knowledge_read_error(error) @workspace.get("/knowledge/file", response_model=dict) async def get_workspace_knowledge_file( kb_id: str = Query(..., description="知识库 ID"), file_id: str = Query(..., description="知识库文件 ID"), variant: str = Query("parsed", description="预览模式:parsed 或 original"), current_user: User = Depends(get_required_user), ): await _ensure_knowledge_read_access(current_user, kb_id) try: return await knowledge_base.read_file_preview(kb_id=kb_id, file_id=file_id, variant=variant) except ValueError as error: _raise_knowledge_read_error(error) @workspace.get("/knowledge/download") async def download_workspace_knowledge_file( kb_id: str = Query(..., description="知识库 ID"), file_id: str = Query(..., description="知识库文件 ID"), variant: str = Query("original", description="下载模式:original 或 parsed"), current_user: User = Depends(get_required_user), ): await _ensure_knowledge_read_access(current_user, kb_id) try: data = await knowledge_base.get_file_download(kb_id=kb_id, file_id=file_id, variant=variant) except ValueError as error: _raise_knowledge_read_error(error) filename = data["filename"] return StreamingResponse( io.BytesIO(data["content"]), media_type=data["media_type"], headers={"Content-Disposition": f"attachment; filename*=UTF-8''{quote(filename)}"}, ) @workspace.put("/file", response_model=dict) async def update_workspace_file( payload: UpdateWorkspaceFileContentRequest, current_user: User = Depends(get_required_user), ): return await write_workspace_file_content( path=payload.path, content=payload.content, current_user=current_user, ) @workspace.delete("/file", response_model=dict) async def delete_workspace_file_route( path: str = Query(..., description="工作区文件或目录路径"), current_user: User = Depends(get_required_user), ): return await delete_workspace_path(path=path, current_user=current_user) @workspace.post("/directory", response_model=dict) async def create_workspace_directory_route( payload: CreateWorkspaceDirectoryRequest, current_user: User = Depends(get_required_user), ): return await create_workspace_directory( parent_path=payload.parent_path, name=payload.name, current_user=current_user, ) @workspace.post("/upload", response_model=dict) async def upload_workspace_file_route( parent_path: str = Form(..., description="父目录路径"), file: UploadFile = File(..., description="上传文件"), current_user: User = Depends(get_required_user), ): return await upload_workspace_file(parent_path=parent_path, file=file, current_user=current_user) @workspace.get("/download") async def download_workspace( path: str = Query(..., description="工作区文件路径"), current_user: User = Depends(get_required_user), ): return await download_workspace_file(path=path, current_user=current_user)