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