| | |
| | | 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): |
| | |
| | | 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 |