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
|
from app.models.base_model import Base
|
from app.utils.common import current_time
|
|
|
class RetrievalSetting(BaseModel):
|
top_k: int
|
score_threshold: float
|
|
|
class RetrievalRequest(BaseModel):
|
knowledge_id: str
|
query: str
|
retrieval_setting: RetrievalSetting
|
|
class ChatDataRequest(BaseModel):
|
sessionId: 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,
|
}
|
|
|
|
|
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=深度
|
|
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, 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_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"))
|
|
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_DELETE).first()
|
return session
|
|
|
|
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):
|
if self.message_type == 1:
|
return {
|
'id': self.id,
|
'role': "user",
|
'content': self.content,
|
}
|
else:
|
return {
|
'id': self.id,
|
'role': "assistant",
|
'answer': self.content,
|
'node_list': json.loads(self.node_data) if self.node_data else [],
|
}
|
|
|
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
|
