Sfoglia il codice sorgente

Merge remote-tracking branch 'origin/master'

zhangxiao 1 anno fa
parent
commit
81d420d88e
6 ha cambiato i file con 117 aggiunte e 8 eliminazioni
  1. 45 1
      app/api/__init__.py
  2. 13 2
      app/api/auth.py
  3. 48 1
      app/api/files.py
  4. 2 1
      app/api/report.py
  5. 6 1
      app/models/token_model.py
  6. 3 2
      pip_install.sh

+ 45 - 1
app/api/__init__.py

@@ -1,3 +1,6 @@
+import urllib
+from urllib.parse import urlencode
+
 import jwt
 from fastapi import FastAPI, Depends, HTTPException
 from fastapi.security import OAuth2PasswordBearer
@@ -72,4 +75,45 @@ async def get_current_user_websocket(websocket: WebSocket):
     except jwt.PyJWTError as e:
         print(e)
         await websocket.close(code=1008)
-        raise WebSocketDisconnect(code=status.WS_1008_POLICY_VIOLATION)
+        raise WebSocketDisconnect(code=status.WS_1008_POLICY_VIOLATION)
+
+
+def format_file_url(agent_id: str, file_url: str, doc_id: str = None, doc_name: str = None) -> str:
+    if file_url:
+        # 对 file_url 进行 URL 编码
+        encoded_file_url = urllib.parse.quote(file_url, safe=':/')
+        return f"./api/files/download/?url={encoded_file_url}&agent_id={agent_id}"
+
+    if doc_id:
+        # 对 doc_id 和 doc_name 进行 URL 编码
+        encoded_doc_id = urllib.parse.quote(doc_id, safe='')
+        encoded_doc_name = urllib.parse.quote(doc_name, safe='')
+        return f"./api/files/download/?doc_id={encoded_doc_id}&doc_name={encoded_doc_name}&agent_id={agent_id}"
+
+    return file_url
+
+
+def process_files(files, agent_id):
+    """
+    处理文件列表,格式化每个文件的 URL。
+
+    :param files: 文件列表,每个文件是一个字典
+    :param agent_id: 代理 ID
+    """
+    if not files:
+        return  # 如果文件列表为空,直接返回
+
+    for file in files:
+        if "file_url" in file and file["file_url"]:
+            try:
+                file["file_url"] = format_file_url(agent_id, file["file_url"])
+            except Exception as e:
+                # 记录异常信息,但继续处理其他文件
+                print(f"Error processing file URL: {e}")
+if __name__=="__main__":
+
+    files1 = [{"file_url": "aaa.com"}, {"file_url":"bbb.com"}]
+    print(files1)
+
+    process_files(files1,11111)
+    print(files1)

+ 13 - 2
app/api/auth.py

@@ -1,10 +1,10 @@
 from fastapi import APIRouter, Depends
 from sqlalchemy.orm import Session
 
-from app.api import Response, pwd_context
+from app.api import Response, pwd_context, get_current_user
 from app.config.config import settings
 from app.models.base_model import get_db
-from app.models.token_model import upsert_token
+from app.models.token_model import upsert_token, get_token
 from app.models.user import UserCreate, LoginData
 from app.models.user_model import UserModel
 from app.service.auth import authenticate_user, create_access_token
@@ -76,3 +76,14 @@ async def login(login_data: LoginData, db: Session = Depends(get_db)):
         "username": user.username,
         "nickname": "",
     })
+
+
+@router.get("/token", response_model=Response)
+async def token_api(db: Session = Depends(get_db), current_user: UserModel = Depends(get_current_user)):
+    # 查询现有记录
+    token = get_token(db, current_user.id)
+    if token is None:
+        return Response(code=400, msg="token not found")
+    return Response(code=200, msg="success", data={
+        "ragflow_token": token.ragflow_token,
+    })

+ 48 - 1
app/api/files.py

@@ -1,5 +1,10 @@
-from fastapi import Depends, APIRouter, HTTPException, UploadFile, File, requests, Query
+from typing import Optional
+
+import requests
+from fastapi import Depends, APIRouter, HTTPException, UploadFile, File, Query
+from pydantic import BaseModel
 from sqlalchemy.orm import Session
+from starlette.responses import StreamingResponse
 
 from app.api import Response, get_current_user, ResponseList
 from app.config.config import settings
@@ -9,6 +14,7 @@ from app.models.user_model import UserModel
 from app.service.bisheng import BishengService
 from app.service.ragflow import RagflowService
 from app.service.token import get_ragflow_token, get_bisheng_token
+import urllib.parse
 
 router = APIRouter()
 
@@ -55,3 +61,44 @@ async def upload_file(agent_id: str,
 
     else:
         return Response(code=200, msg="Unsupported agent type")
+
+
+@router.get("/download/", response_model=Response)
+async def download_file(
+        url: Optional[str] = Query(None, description="URL of the file to download for bisheng"),
+        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"),
+        db: Session = Depends(get_db)
+):
+    agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first()
+    if not agent:
+        return Response(code=404, msg="Agent not found")
+
+    if agent.agent_type == AgentType.BISHENG:
+        url = urllib.parse.unquote(url)
+        # 从 URL 中提取文件名
+        parsed_url = urllib.parse.urlparse(url)
+        filename = urllib.parse.unquote(parsed_url.path.split('/')[-1])
+        url = url.replace("http://minio:9000", settings.bisheng_base_url)
+    elif agent.agent_type == AgentType.RAGFLOW:
+        if not doc_id:
+            return Response(code=400, msg="doc_id is required")
+        url = f"{settings.ragflow_base_url}/v1/document/get/{doc_id}"
+        filename = doc_name
+    else:
+        return Response(code=400, msg="Unsupported agent type")
+
+    try:
+        # 发送GET请求获取文件内容
+        response = requests.get(url, stream=True)
+        response.raise_for_status()  # 检查请求是否成功
+
+        # 返回流式响应
+        return StreamingResponse(
+            response.iter_content(chunk_size=1024),
+            media_type="application/octet-stream",
+            headers={"Content-Disposition": f"attachment; filename*=utf-8''{urllib.parse.quote(filename)}"}
+        )
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error downloading file: {e}")

+ 2 - 1
app/api/report.py

@@ -4,7 +4,7 @@ from fastapi import WebSocket, WebSocketDisconnect, APIRouter, Depends, HTTPExce
 import asyncio
 import websockets
 from sqlalchemy.orm import Session
-from app.api import get_current_user_websocket, ResponseList, get_current_user
+from app.api import get_current_user_websocket, ResponseList, get_current_user, format_file_url, process_files
 from app.config.config import settings
 from app.models.agent_model import AgentModel, AgentType
 from app.models.base_model import get_db
@@ -68,6 +68,7 @@ async def report_chat(websocket: WebSocket,
                             t = "close"
                         else:
                             t = "stream"
+                        process_files(files, agent_id)
                         result = {"step_message": steps, "type": t, "files": files}
                         await websocket.send_json(result)
                         print(f"Forwarded to client, {chat_id}: {result}")

+ 6 - 1
app/models/token_model.py

@@ -1,6 +1,7 @@
 from datetime import datetime
+from typing import Type
 
-from sqlalchemy import Column, Integer, String, DateTime, Text
+from sqlalchemy import Column, Integer, DateTime, Text
 from sqlalchemy.orm import Session
 
 from app.models.base_model import Base
@@ -49,3 +50,7 @@ def upsert_token(db: Session, user_id: int, access_token: str, bisheng_token: st
     except Exception as e:
         # 异常处理
         db.rollback()  # 回滚事务
+
+
+def get_token(db: Session, user_id: int) -> Type[TokenModel] | None:
+    return db.query(TokenModel).filter_by(user_id=user_id).first()

+ 3 - 2
pip_install.sh

@@ -1,4 +1,5 @@
 pip install PyMySQL & pip install fastapi & pip install sqlalchemy & pip install PyJWT & pip install rsa & pip install httpx & pip install uvicorn & pip install bcrypt & pip install PyYAML & pip install pycryptodomex & pip install passlib
 pip install werkzeug
-pip install xlwings
-pip install python-multipart
+pip install openpyxl
+pip install python-multipart
+pip install requests