From 51433cba2f35b9a2571023236006ebc69d1d4d2d Mon Sep 17 00:00:00 2001
From: zhaoqingang <zhaoqg0118@163.com>
Date: 星期二, 07 一月 2025 16:04:31 +0800
Subject: [PATCH] tnp

---
 app/service/v2/chat.py         |    2 
 app/models/v2/session_model.py |   99 +++++++++++++++++++++++++++++++++++++++++++++----
 main.py                        |   11 +++--
 3 files changed, 98 insertions(+), 14 deletions(-)

diff --git a/app/models/v2/session_model.py b/app/models/v2/session_model.py
index 1bebe2f..94cf32b 100644
--- a/app/models/v2/session_model.py
+++ b/app/models/v2/session_model.py
@@ -1,26 +1,26 @@
 import json
-from datetime import datetime
-from enum import IntEnum
-from typing import Optional
-
 import pytz
+
+from datetime import datetime
+from sqlalchemy.orm import Session
+from typing import Optional, Type
 from pydantic import BaseModel
 from sqlalchemy import Column, String, Integer, DateTime, JSON, TEXT, Index
 
+from Log import logger
 from app.models.agent_model import AgentType
-# from app.models import current_time
 from app.models.base_model import Base
 
 def current_time():
     tz = pytz.timezone('Asia/Shanghai')
     return datetime.now(tz)
 
-class SessionModel(Base):
+class ChatSessionModel(Base):
     __tablename__ = "chat_sessions"
 
-    __table_args__ = (
-        Index('idx_username', 'username'),
-    )
+    # __table_args__ = (
+    #     Index('idx_username', 'username'),
+    # )
 
     id = Column(Integer, primary_key=True)
     name = Column(String(255))
@@ -72,3 +72,84 @@
     sessionId: Optional[str] = ""
     question: str
     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:
+        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,
+        )
+        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:
+        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)
+        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:
+                self.db.commit()
+                self.db.refresh(session)
+            except Exception as e:
+                logger.error(e)
+                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)
+        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
+
+        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
+
+    def delete_session(self, session_id: str) -> None:
+        """
+        鍒犻櫎浼氳瘽璁板綍銆�
+
+        鍙傛暟:
+            session_id (str): 浼氳瘽ID銆�
+        """
+        session = self.get_session_by_id(session_id)
+        if session:
+            self.db.delete(session)
+            self.db.commit()
\ No newline at end of file
diff --git a/app/service/v2/chat.py b/app/service/v2/chat.py
index f35604a..136a113 100644
--- a/app/service/v2/chat.py
+++ b/app/service/v2/chat.py
@@ -20,7 +20,7 @@
         for ans in chat.chat_completions(url, data, headers):
 
             yield "data:" + json.dumps({"code": 0, "message": "", "data": ans}, ensure_ascii=False) + "\n\n"
-        ConversationService.update_by_id(conv.id, conv.to_dict())
+        ChatSessionModel.update_by_id(conv.id, conv.to_dict())
     except Exception as e:
         yield "data:" + json.dumps({"code": 500, "message": str(e),
                                     "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
diff --git a/main.py b/main.py
index 16f7a3b..881517c 100644
--- a/main.py
+++ b/main.py
@@ -39,11 +39,14 @@
     # initialize_agents()
     # # 鍦ㄥ簲鐢ㄥ惎鍔ㄦ椂鍚屾浠g悊
     # sync_agents()
-    await sync_default_data()
 
-    sync_agents_v2()
-    sync_knowledge()
-    sync_resources_from_json()
+
+
+    # await sync_default_data()
+    #
+    # sync_agents_v2()
+    # sync_knowledge()
+    # sync_resources_from_json()
     yield
     # 鍦ㄥ簲鐢ㄥ叧闂椂鎵ц娓呯悊鎿嶄綔锛堝鏋滈渶瑕侊級
     pass

--
Gitblit v1.8.0