审计结果文件下载

zhangqian

2024-11-19 f37670f13f8faf018a87d5b73b662bb1909ebe87

 app/api/files.py

@@ -1,19 +1,22 @@
import io
from typing import Optional

import requests
from fastapi import Depends, APIRouter, HTTPException, UploadFile, File, Query
from fastapi import Depends, APIRouter, HTTPException, UploadFile, File, Query, Form
from pydantic import BaseModel
from sqlalchemy.orm import Session
from starlette.responses import StreamingResponse
from werkzeug.utils import send_file

from app.api import Response, get_current_user, ResponseList
from app.config.config import settings
from app.models.agent_model import AgentType, AgentModel
from app.models.base_model import get_db
from app.models.user_model import UserModel
from app.service.basic import BasicService
from app.service.bisheng import BishengService
from app.service.ragflow import RagflowService
from app.service.token import get_ragflow_token, get_bisheng_token
from app.service.service_token import get_ragflow_token, get_bisheng_token
import urllib.parse

router = APIRouter()
@@ -58,6 +61,12 @@
            raise HTTPException(status_code=500, detail=str(e))
        result["file_name"] = file.filename
        return Response(code=200, msg="", data=result)
    elif agent.agent_type == AgentType.BASIC:
        if agent_id == "basic_excel_talk":
            service = BasicService(base_url=settings.basic_base_url)
            result = await service.excel_talk_upload(chat_id, file.filename, file_content)

            return Response(code=200, msg="", data=result)

    else:
        return Response(code=200, msg="Unsupported agent type")
@@ -69,6 +78,8 @@
        agent_id: str = Query(..., description="Agent ID"),
        doc_id: Optional[str] = Query(None, description="Optional doc id for ragflow agents"),
        doc_name: Optional[str] = Query(None, description="Optional doc name for ragflow agents"),
        file_id:  Optional[str] = Query(None, description="Optional file id for basic agents"),
        file_type:  Optional[str] = Query(None, description="Optional file type for basic agents"),
        db: Session = Depends(get_db)
):
    agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first()
@@ -86,6 +97,10 @@
            return Response(code=400, msg="doc_id is required")
        url = f"{settings.fwr_base_url}/v1/document/get/{doc_id}"
        filename = doc_name
    elif agent.agent_type == AgentType.BASIC:
        if agent_id == "basic_excel_talk":
            return await download_basic_file(file_id, file_type)

    else:
        return Response(code=400, msg="Unsupported agent type")

@@ -102,3 +117,25 @@
        )
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Error downloading file: {e}")


async def download_basic_file(file_id: str, file_type: str):
    service = BasicService(base_url=settings.basic_base_url)
    if not file_type or not file_id:
        return Response(code=400, msg="file_type and file_id is required")
    if file_type == "image":
        content, filename, mimetype = await service.excel_talk_image_download(file_id)
        return StreamingResponse(
                io.BytesIO(content),
                media_type=mimetype,
                headers={"Content-Disposition": f"attachment; filename={filename}"}
            )
    elif file_type == "excel":
        content, filename, mimetype = await service.excel_talk_excel_download(file_id)
        return StreamingResponse(
            io.BytesIO(content),
            media_type=mimetype,
            headers={"Content-Disposition": f"attachment; filename={filename}"}
        )
    else:
        return Response(code=400, msg="Unsupported file type")