From 4e301c1fcf4a0b6e2e797acaa3169d5ff4fe92f5 Mon Sep 17 00:00:00 2001 From: zhaoqingang <zhaoqg0118@163.com> Date: 星期三, 08 一月 2025 13:50:52 +0800 Subject: [PATCH] sse --- app/models/v2/session_model.py | 82 +++++++++++++--------------------------- 1 files changed, 27 insertions(+), 55 deletions(-) diff --git a/app/models/v2/session_model.py b/app/models/v2/session_model.py index 94cf32b..aa85f19 100644 --- a/app/models/v2/session_model.py +++ b/app/models/v2/session_model.py @@ -11,9 +11,11 @@ from app.models.agent_model import AgentType from app.models.base_model import Base + def current_time(): tz = pytz.timezone('Asia/Shanghai') return datetime.now(tz) + class ChatSessionModel(Base): __tablename__ = "chat_sessions" @@ -22,18 +24,17 @@ # Index('idx_username', 'username'), # ) - id = Column(Integer, primary_key=True) + id = Column(String(36), primary_key=True) name = Column(String(255)) agent_id = Column(String(255)) agent_type = Column(Integer) # 鐩墠鍙瓨basic鐨勶紝ragflow鍜宐isheng鐨勮皟鎺ュ彛鑾峰彇 create_date = Column(DateTime, default=current_time) # 鍒涘缓鏃堕棿锛岄粯璁ゅ�间负褰撳墠鏃跺尯鏃堕棿 update_date = Column(DateTime, default=current_time, onupdate=current_time, index=True) # 鏇存柊鏃堕棿锛岄粯璁ゅ�间负褰撳墠鏃跺尯鏃堕棿锛屾洿鏂版椂鑷姩鏇存柊 - tenant_id = Column(Integer) # 鍒涘缓浜� - message = Column(TEXT) # 璇存槑 - reference = Column(TEXT) # 璇存槑 - conversation_id = Column(String(64)) - session_id = Column(String(36), index=True) - chat_mode = Column(Integer) + tenant_id = Column(Integer, index=True) # 鍒涘缓浜� + message = Column(TEXT) + reference = Column(TEXT) + conversation_id = Column(String(36), index=True) + event_type = Column(String(16)) # to_dict 鏂规硶 def to_dict(self): @@ -64,8 +65,10 @@ msg = json.loads(self.message) msg.append(message) except Exception as e: + print(e) return self.message = json.dumps(msg) + class ChatDialogData(BaseModel): @@ -74,45 +77,34 @@ chatId: str - class ChatSessionDao: def __init__(self, db: Session): self.db = db - def create_session(self, session_id: str, name: str, agent_id: str, agent_type: int, user_id: int, message: str,reference:str) -> ChatSessionModel: + async def create_session(self, session_id: str, **kwargs) -> ChatSessionModel: new_session = ChatSessionModel( id=session_id, - name=name[0:255], - agent_id=agent_id, - agent_type=agent_type, create_date=current_time(), update_date=current_time(), - tenant_id=user_id, - message=message, - reference=reference, + **kwargs ) + new_session.message = json.dumps([new_session.message]) self.db.add(new_session) self.db.commit() self.db.refresh(new_session) return new_session - def get_session_by_id(self, session_id: str) -> Type[ChatSessionModel] | None: + async def get_session_by_id(self, session_id: str) -> ChatSessionModel | None: session = self.db.query(ChatSessionModel).filter_by(id=session_id).first() - if session and session.message is None: - session.message = '[]' return session - def update_session_by_id(self, session_id: str, **kwargs) -> Type[ChatSessionModel] | None: - session = self.get_session_by_id(session_id) + async def update_session_by_id(self, session_id: str, session, message: dict) -> ChatSessionModel | None: + if not session: + session = await self.get_session_by_id(session_id) if session: - if "message" in kwargs: - session.add_message(kwargs["message"]) - # 鏇挎崲鍏朵粬瀛楁 - for key, value in kwargs.items(): - if key != "message": - setattr(session, key, value) - session.update_date = current_time() try: + session.add_message(message) + session.update_date = current_time() self.db.commit() self.db.refresh(session) except Exception as e: @@ -120,36 +112,16 @@ self.db.rollback() return session - def create_session(self, session_id: str, name: str, agent_id: str, agent_type: AgentType, user_id: int) -> ChatSessionModel: - existing_session = self.get_session_by_id(session_id) + async def update_or_insert_by_id(self, session_id: str, **kwargs) -> ChatSessionModel: + existing_session = await self.get_session_by_id(session_id) if existing_session: - existing_session.add_message({"role": "user", "content": name}) - existing_session.update_date = current_time() - self.db.commit() - self.db.refresh(existing_session) - return existing_session + return await self.update_session_by_id(session_id, existing_session, kwargs.get("message")) - new_session = ChatSessionModel( - id=session_id, - name=name[0:50], - agent_id=agent_id, - agent_type=agent_type, - tenant_id=user_id, - message=json.dumps([{"role": "user", "content": name}]) - ) - self.db.add(new_session) - self.db.commit() - self.db.refresh(new_session) - return new_session + existing_session = await self.create_session(session_id, **kwargs) + return existing_session - def delete_session(self, session_id: str) -> None: - """ - 鍒犻櫎浼氳瘽璁板綍銆� - - 鍙傛暟: - session_id (str): 浼氳瘽ID銆� - """ - session = self.get_session_by_id(session_id) + 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() \ No newline at end of file + self.db.commit() -- Gitblit v1.8.0