zhaoqingang hace 1 año
padre
commit
51433cba2f
Se han modificado 3 ficheros con 98 adiciones y 14 borrados
  1. 90 9
      app/models/v2/session_model.py
  2. 1 1
      app/service/v2/chat.py
  3. 7 4
      main.py

+ 90 - 9
app/models/v2/session_model.py

@@ -1,26 +1,26 @@
 import json
 import json
-from datetime import datetime
-from enum import IntEnum
-from typing import Optional
-
 import pytz
 import pytz
+
+from datetime import datetime
+from sqlalchemy.orm import Session
+from typing import Optional, Type
 from pydantic import BaseModel
 from pydantic import BaseModel
 from sqlalchemy import Column, String, Integer, DateTime, JSON, TEXT, Index
 from sqlalchemy import Column, String, Integer, DateTime, JSON, TEXT, Index
 
 
+from Log import logger
 from app.models.agent_model import AgentType
 from app.models.agent_model import AgentType
-# from app.models import current_time
 from app.models.base_model import Base
 from app.models.base_model import Base
 
 
 def current_time():
 def current_time():
     tz = pytz.timezone('Asia/Shanghai')
     tz = pytz.timezone('Asia/Shanghai')
     return datetime.now(tz)
     return datetime.now(tz)
 
 
-class SessionModel(Base):
+class ChatSessionModel(Base):
     __tablename__ = "chat_sessions"
     __tablename__ = "chat_sessions"
 
 
-    __table_args__ = (
-        Index('idx_username', 'username'),
-    )
+    # __table_args__ = (
+    #     Index('idx_username', 'username'),
+    # )
 
 
     id = Column(Integer, primary_key=True)
     id = Column(Integer, primary_key=True)
     name = Column(String(255))
     name = Column(String(255))
@@ -72,3 +72,84 @@ class ChatDialogData(BaseModel):
     sessionId: Optional[str] = ""
     sessionId: Optional[str] = ""
     question: str
     question: str
     chatId: 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()

+ 1 - 1
app/service/v2/chat.py

@@ -20,7 +20,7 @@ async def service_chat_dialog(chat_id:str, question: str, session_id: str):
         for ans in chat.chat_completions(url, data, headers):
         for ans in chat.chat_completions(url, data, headers):
 
 
             yield "data:" + json.dumps({"code": 0, "message": "", "data": ans}, ensure_ascii=False) + "\n\n"
             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:
     except Exception as e:
         yield "data:" + json.dumps({"code": 500, "message": str(e),
         yield "data:" + json.dumps({"code": 500, "message": str(e),
                                     "data": {"answer": "**ERROR**: " + str(e), "reference": []}},
                                     "data": {"answer": "**ERROR**: " + str(e), "reference": []}},

+ 7 - 4
main.py

@@ -39,11 +39,14 @@ async def lifespan(app: FastAPI):
     # initialize_agents()
     # initialize_agents()
     # # 在应用启动时同步代理
     # # 在应用启动时同步代理
     # sync_agents()
     # 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
     yield
     # 在应用关闭时执行清理操作(如果需要)
     # 在应用关闭时执行清理操作(如果需要)
     pass
     pass