zhaoqingang 1 éve
szülő
commit
51433cba2f
3 módosított fájl, 98 hozzáadás és 14 törlés
  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
-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 @@ class ChatDialogData(BaseModel):
     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()

+ 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):
 
             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": []}},

+ 7 - 4
main.py

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