dify报告合并

zhaoqingang

2024-12-03 319e3e93bbb4807da58db2ac72bb8f3029738dd7

 app/api/auth.py

@@ -1,25 +1,23 @@
from typing import Dict
import json

from fastapi import APIRouter, Depends, HTTPException
from fastapi.security import OAuth2PasswordBearer
from passlib.context import CryptContext
from fastapi import APIRouter, Depends
from sqlalchemy.orm import Session

from app.api import Response
from sqlalchemy.ext.asyncio import AsyncSession
from app.api import Response, pwd_context, get_current_user
from app.config.config import settings
from app.models.app_token_model import AppToken
from app.models.base_model import get_db
from app.models.token_model import upsert_token
from app.models.user import User, UserCreate, LoginData
from app.models.postgresql_base_model import get_pdb
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
from app.service.bisheng import BishengService
from app.service.ragflow import RagflowService
from sqlalchemy.future import select


router = APIRouter()

pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")


@router.post("/register", response_model=Response)
@@ -28,24 +26,25 @@
    if db_user:
        return Response(code=200, msg="Username already registered")

    bisheng_service = BishengService(settings.bisheng_base_url)
    ragflow_service = RagflowService(settings.ragflow_base_url)
    bisheng_service = BishengService(settings.sgb_base_url)
    ragflow_service = RagflowService(settings.fwr_base_url)

    # 注册到毕昇
    try:
        await bisheng_service.register(user.username, user.password)
        bisheng_info = await bisheng_service.register(user.username, user.password)
    except Exception as e:
        return Response(code=500, msg=f"Failed to register with Bisheng: {str(e)}")

    # 注册到ragflow
    try:
        await ragflow_service.register(user.username, user.password)
        ragflow_info = await ragflow_service.register(user.username, user.password)
    except Exception as e:
        return Response(code=500, msg=f"Failed to register with Ragflow: {str(e)}")

    # 存储用户信息
    hashed_password = pwd_context.hash(user.password)
    db_user = UserModel(username=user.username, hashed_password=hashed_password)
    db_user = UserModel(username=user.username, hashed_password=hashed_password, email=ragflow_info.get("email",  f"{user.username}@example.com"),ragflow_id=ragflow_info.get("id"),bisheng_id=bisheng_info.get("user_id"))
    db_user.password = db_user.encrypted_password(user.password)
    db.add(db_user)
    db.commit()
    db.refresh(db_user)
@@ -58,8 +57,8 @@
    if not user:
        return Response(code=400, msg="Incorrect username or password")

    bisheng_service = BishengService(settings.bisheng_base_url)
    ragflow_service = RagflowService(settings.ragflow_base_url)
    bisheng_service = BishengService(settings.sgb_base_url)
    ragflow_service = RagflowService(settings.fwr_base_url)

    # 登录到毕昇
    try:
@@ -74,11 +73,73 @@
        return Response(code=500, msg=f"Failed to login with Ragflow: {str(e)}")

    # 创建本地token
    access_token = create_access_token(data={"sub": user.username})
    access_token = create_access_token(data={"sub": user.username, "user_id": user.id})

    upsert_token(db, user.id, access_token, bisheng_token, ragflow_token)

    return Response(code=200, msg="Login successful", data={
        "access_token": access_token,
        "token_type": "bearer"
        "token_type": "bearer",
        "username": user.username,
        "nickname": "",
        "user": user.to_login_json()
    })


@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,
    })


@router.post("/login_test", response_model=Response)
async def login_test(login_data: LoginData, db: Session = Depends(get_db), pdb: AsyncSession = Depends(get_pdb)):
    user = authenticate_user(db, login_data.username, login_data.password)
    if not user:
        return Response(code=400, msg="Incorrect username or password")

    bisheng_service = BishengService(settings.sgb_base_url)
    ragflow_service = RagflowService(settings.fwr_base_url)

    # 登录到毕昇
    try:
        bisheng_token = await bisheng_service.login(login_data.username, login_data.password)
    except Exception as e:
        return Response(code=500, msg=f"Failed to login with Bisheng: {str(e)}")

    # 登录到ragflow
    try:
        ragflow_token = await ragflow_service.login(login_data.username, login_data.password)
    except Exception as e:
        return Response(code=500, msg=f"Failed to login with Ragflow: {str(e)}")

    # 创建本地token
    access_token = create_access_token(data={"sub": user.username, "user_id": user.id})

    upsert_token(db, user.id, access_token, bisheng_token, ragflow_token)
    result = await pdb.execute(select(AppToken).where(AppToken.id == user.id))
    db_app_token = result.scalars().first()
    if not db_app_token:
        app_token_str = json.dumps({"rag_token": ragflow_token, "bs_token":bisheng_token})
        # print(app_token_str)
        app_token = AppToken(id=user.id, token=access_token.decode(), app_token=app_token_str)
        pdb.add(app_token)
        await pdb.commit()
        await pdb.refresh(app_token)
    else:
        db_app_token.token = access_token.decode()
        db_app_token.app_token = json.dumps({"rag_token": ragflow_token, "bs_token":bisheng_token})
        await pdb.commit()
        await pdb.refresh(db_app_token)
    return Response(code=200, msg="Login successful", data={
        "access_token": access_token,
        "token_type": "bearer",
        "username": user.username,
        "nickname": "",
        # "user": user.to_login_json()
    })