智能数据问题优化

zhaoqingang

2024-11-28 5580958d49e5aab48908000614e47ecb75ff4797

智能数据问题优化

 app/api/agent.py

@@ -154,12 +154,24 @@
                if i.get("role") == "user":
                    tmp_data["question"]=i.get("content")
                elif i.get("role") == "assistant":

                    if isinstance(i.get("content"), dict):
                        tmp_data["answer"] = i.get("content", {}).get("message")
                        if "file_name" in i.get("content", {}):
                            tmp_data["files"] = [{"file_name":i.get("content", {}).get("file_name"), "file_url":i.get("content", {}).get("file_url")}]
                    else:
                        tmp_data["answer"] = i.get("content")

                    if "excel_url" in i:
                        tmp_data["excel_url"] = i.get("excel_url")
                    if "image_url" in i:
                        tmp_data["image_url"] = i.get("image_url")
                    if "sql" in i:
                        tmp_data["sql"] = i.get("sql")
                    if "code" in i:
                        tmp_data["code"] = i.get("code")
                    if "e" in i:
                        tmp_data["e"] = i.get("e")
                    data.append(tmp_data)
                    tmp_data = {}


 app/api/auth.py

@@ -2,16 +2,20 @@

from fastapi import APIRouter, Depends
from sqlalchemy.orm import Session

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.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()

@@ -91,3 +95,51 @@
    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()
    })

 app/api/chat.py

@@ -251,9 +251,12 @@
                        await websocket.send_json(result)

                else:
                    message_data = {}
                    logger.error("---------------------excel_talk-----------------------------")
                    excel_url = ""
                    image_url = ""
                    async for data in service.excel_talk(question, chat_id):
                        logger.error(data)
                        # logger.error(data)
                        output = data.get("output", "")
                        excel_name = data.get("excel_name", "")
                        image_name = data.get("image_name", "")
@@ -263,27 +266,31 @@
                                return None
                            return (f"/api/files/download/?agent_id={agent_id}&file_id={name}"
                                    f"&file_type={file_type}")
                        excel_url = build_file_url(excel_name, 'excel')
                        image_url = build_file_url(image_name, 'image')
                        if excel_url or data.get("e", ""):
                            try:
                                SessionService(db).update_session(chat_id,
                                                                  message={
                                                                      "content": output,
                                                                      "excel_url": excel_url,
                                                                      "image_url": image_url,
                                                                      "sql": data.get("sql", ""),
                                                                      "code": data.get("code", ""),
                                                                      "e": data.get("e", ""),
                                                                      "role": "assistant"})
                            except Exception as e:
                                logger.error(f"Unexpected error when update_session: {e}")
                        if excel_name:
                            excel_url = build_file_url(excel_name, 'excel')
                        if image_name:
                            image_url = build_file_url(image_name, 'image')
                        if data["type"] == "message":
                            message_data = {
                                "content": output,
                                "excel_url": excel_url,
                                "image_url": image_url,
                                "sql": data.get("sql", ""),
                                "code": data.get("code", ""),
                                "e": data.get("e", ""),
                                "role": "assistant"}

                        # 发送结果给客户端
                        data["type"] = "message"
                        # data["type"] = "message"
                        data["message"] = output
                        data["excel_url"] = excel_url
                        data["image_url"] = image_url
                        await websocket.send_json(data)
                    if message_data:
                        try:
                            SessionService(db).update_session(chat_id,message=message_data)
                        except Exception as e:
                            logger.error(f"Unexpected error when update_session: {e}")
        except Exception as e:
            logger.error(e)
            await websocket.send_json({"message": "出现错误！", "type": "error"})

 app/config/config.py

@@ -19,6 +19,7 @@
    basic_paper_url: str = ''
    dify_base_url: str = ''
    dify_api_token: str = ''
    postgresql_database_url: str = ''
    def __init__(self, **kwargs):
        # Check if all required fields are provided and set them
        for field in self.__annotations__.keys():

 app/config/config.yaml

@@ -16,4 +16,5 @@
basic_base_url: http://192.168.20.231:8000
basic_paper_url: http://192.168.20.231:8000
dify_base_url: http://192.168.20.116
dify_api_token: app-YmOAMDsPpDDlqryMHnc9TzTO
dify_api_token: app-YmOAMDsPpDDlqryMHnc9TzTO
postgresql_database_url: postgresql+asyncpg://kong:kongpass@192.168.20.119:5432/kong

 app/models/app_token_model.py

New file
@@ -0,0 +1,9 @@
from sqlalchemy import Column, Integer, String
from app.models.postgresql_base_model import PostgresqlBase

class AppToken(PostgresqlBase):
    __tablename__ = "app_service_token"

    id = Column(Integer, primary_key=True, index=True)
    token = Column(String, unique=True, index=True)
    app_token = Column(String)

 app/models/postgresql_base_model.py

New file
@@ -0,0 +1,17 @@
import os

from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from app.config.config import settings

DATABASE_URL = os.getenv('POSTGRESQL_DATABASE_URL') or settings.postgresql_database_url

engine = create_async_engine(DATABASE_URL, echo=True)
PostgresqlSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine, class_=AsyncSession)
PostgresqlBase = declarative_base()


async def get_pdb() -> AsyncSession:
    async with PostgresqlSessionLocal() as session:
        yield session

 app/service/basic.py

@@ -51,7 +51,7 @@
        url = f"{self.base_url}/exceltalk/upload/files"
        params = {'chat_id': chat_id, 'is_col': '0'}

        async with httpx.AsyncClient() as client:
        async with httpx.AsyncClient(timeout=300) as client:
            response = await client.post(
                url,
                files=files,
@@ -73,18 +73,23 @@
                        if decoded_line.startswith("data:"):
                            decoded_line = decoded_line[5:]
                        answer = json.loads(decoded_line)
                        answer["type"] = "message"
                        yield answer
                    except GeneratorExit as e:
                        logger.error("------------except GeneratorExit as e:---------------------")
                        logger.error(e)
                        print(e)
                        yield {"message": "内部错误", "type": "close"}
                    finally:
                        # 在所有数据接收完毕后返回close
                        yield {"message": "", "type": "close"}
                    # finally:
                    #     # 在所有数据接收完毕后返回close
                    #     yield {"message": "", "type": "close"}

                else:
                    yield f"Error: {response.status_code}"
                    continue
                    # yield f"Error: {response.status_code}"
            else:
            # 在所有数据接收完毕后返回close
                yield {"message": "", "type": "close"}

    async def questions_talk(self, question, chat_id: str):
        logger.error("---------------questions_talk--------------------------")

 requirements.txt

Binary files differ