zhaoqingang
2025-04-08 0650b889a36d9b9fd42415b9b9819676f839ae9b
app/models/v2/chat.py
@@ -1,5 +1,14 @@
from pydantic import BaseModel
import json
from datetime import datetime
from typing import List, Optional
from pydantic import BaseModel
from sqlalchemy import Column, Integer, String, BigInteger, ForeignKey, DateTime, Text, TEXT
from sqlalchemy.orm import Session
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
class RetrievalSetting(BaseModel):
@@ -10,4 +19,261 @@
class RetrievalRequest(BaseModel):
    knowledge_id: str
    query: str
    retrieval_setting: RetrievalSetting
    retrieval_setting: RetrievalSetting
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] = []
    def to_dict(self):
        return {
            "sessionId": self.sessionId,
            "query": self.query,
            "chatMode": self.chatMode,
            "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
class ComplexChatModel(Base):
    __tablename__ = 'complex_chat'
    __mapper_args__ = {
        # "order_by": 'SEQ'
    }
    id = Column(String(36), primary_key=True)  #  id
    create_date = Column(DateTime, default=datetime.now())             # 创建时间
    update_date = Column(DateTime, default=datetime.now(), onupdate=datetime.now())             # 更新时间
    tenant_id = Column(String(36))              # 创建人
    name = Column(String(255))                 # 名称
    description = Column(Text)                 # 说明
    icon = Column(Text, default="intelligentFrame1")                         # 图标
    status = Column(String(1), default="1")                 # 状态
    dialog_type = Column(String(1))            #  平台
    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 {
            'id': self.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'),
            'user_id': self.tenant_id,
            'name': self.name,
            'description': self.description,
            'icon': self.icon,
            'status': self.status,
            'agentType': self.dialog_type,
            'mode': self.mode,
        }
class ComplexChatDao:
    def __init__(self, db: Session):
        self.db = db
    async def create_complex_chat(self, chat_id: str, **kwargs) -> ComplexChatModel:
        new_session = ComplexChatModel(
            id=chat_id,
            create_date=current_time(),
            update_date=current_time(),
            **kwargs
        )
        self.db.add(new_session)
        self.db.commit()
        self.db.refresh(new_session)
        return new_session
    async def get_complex_chat_by_id(self, chat_id: str) -> ComplexChatModel | None:
        session = self.db.query(ComplexChatModel).filter_by(id=chat_id).first()
        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, kwargs)
        existing_session = await self.create_complex_chat(chat_id, **kwargs)
        return existing_session
    async def delete_complex_chat(self, chat_id: str) -> None:
        session = await self.get_complex_chat_by_id(chat_id)
        if session:
            self.db.delete(session)
            self.db.commit()
    async def aget_complex_chat_ids(self) -> List:
        session_list = self.db.query(ComplexChatModel).filter(ComplexChatModel.status!=Dialog_STATSU_DELETE).all()
        return [i.id for i in session_list]
    def get_complex_chat_ids(self) -> List:
        session_list = self.db.query(ComplexChatModel).filter(ComplexChatModel.status!=Dialog_STATSU_DELETE).all()
        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_ON).first()
        return session
    async def aget_complex_chat(self) -> List:
        return self.db.query(ComplexChatModel).filter(ComplexChatModel.status!=Dialog_STATSU_DELETE).all()
class ComplexChatSessionModel(Base):
    __tablename__ = "complex_chat_sessions"
    id = Column(String(36), primary_key=True)
    chat_id = Column(String(36))
    session_id = Column(String(36), index=True)
    create_date = Column(DateTime, default=current_time, index=True)  # 创建时间,默认值为当前时区时间
    update_date = Column(DateTime, default=current_time, onupdate=current_time)  # 更新时间,默认值为当前时区时间,更新时自动更新
    tenant_id = Column(Integer, index=True)  # 创建人
    agent_type = Column(Integer) # 1=rg, 3=basic,4=df
    message_type = Column(Integer)  # 1=用户,2=机器人,3=系统
    content = Column(TEXT)
    mindmap = Column(TEXT)
    query = Column(TEXT)
    node_data = Column(TEXT)
    event_type = Column(String(16))
    conversation_id = Column(String(36))
    chat_mode = Column(Integer) # 1= 普通对话,2=联网,3=知识库,4=深度
    # to_dict 方法
    def to_dict(self):
        return {
            'session_id': self.id,
            'name': self.name,
            'agent_type': self.agent_type,
            'chat_id': self.agent_id,
            'event_type': self.event_type,
            'session_type': self.session_type if self.session_type else 0,
            'create_date': self.create_date.strftime("%Y-%m-%d %H:%M:%S"),
            'update_date': self.update_date.strftime("%Y-%m-%d %H:%M:%S"),
        }
    def log_to_json(self):
        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:
    def __init__(self, db: Session):
        self.db = db
    async def get_session_by_session_id(self, session_id: str, chat_id:str) -> ComplexChatSessionModel | None:
        session = self.db.query(ComplexChatSessionModel).filter_by(chat_id=chat_id, session_id=session_id, message_type=2).first()
        return session
    async def create_session(self, message_id: str, **kwargs) -> ComplexChatSessionModel:
        new_session = ComplexChatSessionModel(
            id=message_id,
            create_date=current_time(),
            update_date=current_time(),
            **kwargs
        )
        self.db.add(new_session)
        self.db.commit()
        self.db.refresh(new_session)
        return new_session
    async def get_session_by_id(self, message_id: str) -> ComplexChatSessionModel | None:
        session = self.db.query(ComplexChatSessionModel).filter_by(id=message_id).first()
        return session
    async def update_mindmap_by_id(self, message_id: str, mindmap:str) -> ComplexChatSessionModel | None:
        # print(message)
        session = await self.get_session_by_id(message_id)
        if session:
            try:
                session.mindmap = mindmap
                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_or_insert_by_id(self, session_id: str, **kwargs) -> ComplexChatSessionModel:
        existing_session = await self.get_session_by_id(session_id)
        if existing_session:
            return await self.update_session_by_id(session_id, existing_session, kwargs.get("message"))
        existing_session = await self.create_session(session_id, **kwargs)
        return existing_session
    async def delete_session(self, session_id: str) -> None:
        session = await self.get_session_by_id(session_id)
        if session:
            self.db.delete(session)
            self.db.commit()
    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.content.like('%{}%'.format(keyword)))
        total = query.count()
        session_list = query.order_by(ComplexChatSessionModel.create_date.desc()).offset((page-1)*page_size).limit(page_size).all()
        return total, session_list