zhaoqingang
2025-04-08 0650b889a36d9b9fd42415b9b9819676f839ae9b
app/models/v2/chat.py
@@ -6,7 +6,7 @@
from sqlalchemy import Column, Integer, String, BigInteger, ForeignKey, DateTime, Text, TEXT
from sqlalchemy.orm import Session
from app.config.const import Dialog_STATSU_DELETE
from app.config.const import Dialog_STATSU_DELETE, Dialog_STATSU_ON, complex_knowledge_chat, complex_knowledge_chat_deep
from app.models.base_model import Base
from app.utils.common import current_time
@@ -23,9 +23,11 @@
class ChatDataRequest(BaseModel):
    sessionId: str
    parentId: Optional[str] = ""
    query: str
    chatMode: Optional[int] = 1  # 1= 普通对话,2=联网,3=知识库,4=深度
    isDeep: Optional[int] = 1  # 1= 普通, 2=深度
    optimizeType: Optional[str] = ""  # 优化类型:润色,扩写,缩写,调整语气,自定义
    knowledgeId: Optional[list] = []
    files: Optional[list] = []
@@ -38,8 +40,17 @@
            "knowledgeId": self.knowledgeId,
            "files": self.files,
            "isDeep": self.isDeep,
            "optimizeType": self.optimizeType,
            "parentId": self.parentId,
        }
class SetModelRequest(BaseModel):
    chatType: int
    modelType: int
    modelName: str
    modelProvider: str
@@ -60,6 +71,9 @@
    mode = Column(String(36))
    parameters = Column(Text)
    chat_mode = Column(Integer) #1= 普通对话,2=联网,3=知识库,4=深度
    chat_model = Column(String(255)) # 模型
    chat_model_ds = Column(String(255)) # 模型
    chat_provider = Column(String(255)) # 模型提供商
    def to_json(self):
        return {
@@ -95,24 +109,20 @@
        session = self.db.query(ComplexChatModel).filter_by(id=chat_id).first()
        return session
    async def update_complex_chat_by_id(self, chat_id: str, session, message: dict, conversation_id=None) -> ComplexChatModel | None:
        if not session:
            session = await self.get_complex_chat_by_id(chat_id)
        if session:
            try:
                # TODO
                session.update_date = current_time()
                self.db.commit()
                self.db.refresh(session)
            except Exception as e:
                # logger.error(e)
                self.db.rollback()
        return session
    async def update_complex_chat_by_id(self, chat_id: str, kwargs:dict) -> None:
        try:
            self.db.query(ComplexChatModel).filter_by(id=chat_id).update(kwargs)
            self.db.commit()
        except Exception as e:
            # logger.error(e)
            self.db.rollback()
    async def update_or_insert_by_id(self, chat_id: str, **kwargs) -> ComplexChatModel:
        existing_session = await self.get_complex_chat_by_id(chat_id)
        if existing_session:
            return await self.update_complex_chat_by_id(chat_id, existing_session, kwargs.get("message"))
            return await self.update_complex_chat_by_id(chat_id, kwargs)
        existing_session = await self.create_complex_chat(chat_id, **kwargs)
        return existing_session
@@ -134,8 +144,12 @@
        return [i.id for i in session_list]
    async def get_complex_chat_by_mode(self, chat_mode: int) -> ComplexChatModel | None:
        session = self.db.query(ComplexChatModel).filter(ComplexChatModel.chat_mode==chat_mode, ComplexChatModel.status!=Dialog_STATSU_DELETE).first()
        session = self.db.query(ComplexChatModel).filter(ComplexChatModel.chat_mode==chat_mode, ComplexChatModel.status==Dialog_STATSU_ON).first()
        return session
    async def aget_complex_chat(self) -> List:
        return self.db.query(ComplexChatModel).filter(ComplexChatModel.status!=Dialog_STATSU_DELETE).all()
@@ -173,15 +187,32 @@
        }
    def log_to_json(self):
        return {
            'id': self.id,
            'name': self.name,
            'agent_type': self.agent_type,
            'chat_id': self.agent_id,
            'create_date': self.create_date.strftime("%Y-%m-%d %H:%M:%S"),
            'update_date': self.update_date.strftime("%Y-%m-%d %H:%M:%S"),
            'message': json.loads(self.message)
        }
        query = {}
        if self.query:
            query = json.loads(self.query)
        if self.message_type == 1:
            return {
                'id': self.id,
                'role': "user",
                'content': self.content,
                'files': query.get("files", []),
            }
        else:
            res = {
                'id': self.id,
                'role': "assistant",
                'answer': self.content,
                'chat_mode': self.chat_mode,
                "parentId": query.get("parentId"),
                "isDeep": query.get("isDeep", 1),
                "mindmap": True if self.mindmap else False,
            }
            if self.chat_mode == complex_knowledge_chat or self.chat_mode == complex_knowledge_chat_deep:
                res['reference'] = json.loads(self.node_data) if self.node_data else {}
            else:
                res['node_list'] = json.loads(self.node_data) if self.node_data else []
            return res
class ComplexChatSessionDao:
@@ -238,12 +269,11 @@
            self.db.delete(session)
            self.db.commit()
    async def get_session_list(self, user_id: int, agent_id: str, keyword:str, page: int, page_size: int) -> any:
        query = self.db.query(ComplexChatSessionModel).filter(ComplexChatSessionModel.tenant_id==user_id)
        if agent_id:
            query = query.filter(ComplexChatSessionModel.agent_id==agent_id)
    async def get_session_list(self, session_id: int, keyword:str="", page: int=1, page_size: int=100) -> any:
        query = self.db.query(ComplexChatSessionModel).filter(ComplexChatSessionModel.session_id==session_id)
        if keyword:
            query = query.filter(ComplexChatSessionModel.name.like('%{}%'.format(keyword)))
            query = query.filter(ComplexChatSessionModel.content.like('%{}%'.format(keyword)))
        total = query.count()
        session_list = query.order_by(ComplexChatSessionModel.update_date.desc()).offset((page-1)*page_size).limit(page_size).all()
        session_list = query.order_by(ComplexChatSessionModel.create_date.desc()).offset((page-1)*page_size).limit(page_size).all()
        return total, session_list