会话缓存本地数据库

app/api/agent.py

@@ -12,6 +12,7 @@ from app.models.agent_model import AgentType, AgentModel
 from app.models.base_model import get_db
 from app.models.user_model import UserModel
 from app.service.bisheng import BishengService
+from app.service.dialog import get_session_history
 from app.service.ragflow import RagflowService
 from app.service.service_token import get_ragflow_token, get_bisheng_token
 
@@ -41,6 +42,8 @@ async def chat_list(
         try:
             token = get_ragflow_token(db, current_user.id)
             result = await ragflow_service.get_chat_sessions(token, agent_id)
+            if not result:
+                result = await get_session_history(db, current_user.id, agent_id)
         except Exception as e:
             raise HTTPException(status_code=500, detail=str(e))
         return ResponseList(code=200, msg="", data=result)

app/api/chat.py

@@ -10,6 +10,7 @@ from app.config.config import settings
 from app.models.agent_model import AgentModel, AgentType
 from app.models.base_model import get_db
 from app.models.user_model import UserModel
+from app.service.dialog import update_session_history
 from app.service.ragflow import RagflowService
 from app.service.service_token import get_bisheng_token, get_ragflow_token
 
@@ -44,6 +45,7 @@ async def handle_client(websocket: WebSocket,
         try:
             async def forward_to_ragflow():
                 while True:
+                    is_new = False
                     message = await websocket.receive_json()
                     print(f"Received from client {chat_id}: {message}")
                     chat_history = message.get('chatHistory', [])
@@ -51,9 +53,10 @@ async def handle_client(websocket: WebSocket,
                     if len(chat_history) == 0:
                         chat_history = await ragflow_service.get_session_history(token, chat_id)
                         if len(chat_history) == 0:
+                            is_new = True
                             chat_history = await ragflow_service.set_session(token, agent_id,
                                                                              message, chat_id, True)
-                            print("chat_history------------------------", chat_history)
+                            # print("chat_history------------------------", chat_history)
                             if len(chat_history) == 0:
                                 result = {"message": "内部错误：创建会话失败", "type": "close"}
                                 await websocket.send_json(result)
@@ -91,11 +94,14 @@ async def handle_client(websocket: WebSocket,
                                 complete_response = ""
                             except json.JSONDecodeError as e:
                                 print(f"Error decoding JSON: {e}")
-                                print(f"Response text: {text}")
+                                # print(f"Response text: {text}")
                         except Exception as e2:
                             result = {"message": f"内部错误： {e2}", "type": "close"}
                             await websocket.send_json(result)
                             print(f"Error process message of ragflow: {e2}")
+                    dialog_chat_history = await ragflow_service.get_session_history(token, chat_id, 1)
+                    await update_session_history(db, dialog_chat_history, current_user.id, is_new)
+
             # 启动任务处理客户端消息
             tasks = [
                 asyncio.create_task(forward_to_ragflow())

app/models/dialog_model.py

@@ -1,6 +1,6 @@
 from datetime import datetime
 
-from sqlalchemy import Column, Integer, String, Table, ForeignKey, DateTime, BigInteger, Text, Float, Boolean
+from sqlalchemy import Column, Integer, String, BigInteger, ForeignKey, DateTime, Text, JSON
 from sqlalchemy.orm import relationship, backref
 
 from app.models.base_model import Base
@@ -33,4 +33,29 @@ class DialogModel(Base):
             'description': self.description,
             'icon': self.icon,
             'status': self.status
+        }
+
+
+class ConversationModel(Base):
+    __tablename__ = 'conversation'
+    id = Column(String(32), primary_key=True)  #  id
+    create_date = Column(DateTime)             # 创建时间
+    create_time = Column(BigInteger)
+    update_date = Column(DateTime)             # 更新时间
+    update_time = Column(BigInteger)
+    tenant_id = Column(Integer)              # 创建人
+    dialog_id = Column(String(32))
+    name = Column(String(255))                 # 名称
+    message = Column(JSON)                 # 说明
+    reference = Column(JSON)                         # 图标
+
+    def get_id(self):
+        return str(self.id)
+
+
+    def to_json(self):
+        return {
+            'id': self.id,
+            'updated_time': self.update_time,
+            'name': self.name,
         }

app/service/dialog.py

@@ -1,4 +1,5 @@
-from app.models import KnowledgeModel, GroupModel, DialogModel
+from app.api.user import user_list
+from app.models import KnowledgeModel, GroupModel, DialogModel, ConversationModel
 from app.models.user_model import UserModel
 from Log import logger
 
@@ -19,4 +20,35 @@ async def get_dialog_list(db, role, user_id):
                     dialog_list.append(k)
                     kld_set.add(k.id)
 
-    return {"rows": [kld.to_json() for kld in dialog_list]}
+    return {"rows": [kld.to_json() for kld in dialog_list]}
+
+
+async def update_session_history(db, data: dict, user_id, is_new):
+    session_id = data.get("id")
+    if not session_id:
+        logger.error("更新回话记录失败！{}".format(data))
+    if is_new:
+        try:
+            data["tenant_id"] = user_id
+            conversation_model = ConversationModel(**data)
+            db.add(conversation_model)
+            db.commit()
+        except Exception as e:
+            logger.error(e)
+            db.rollback()
+    else:
+        try:
+            data["tenant_id"] = user_id
+            del data["id"]
+            db.query(ConversationModel).filter(ConversationModel.id == session_id).update(data)
+            db.commit()
+        except Exception as e:
+            logger.error(e)
+            db.rollback()
+
+
+async def get_session_history(db, user_id, dialog_id):
+    session_list = db.query(ConversationModel).filter(ConversationModel.tenant_id.__eq__(user_id),
+                                                      ConversationModel.dialog_id.__eq__(dialog_id)).order_by(
+        ConversationModel.update_time.desc()).all()
+    return [i.to_json() for i in session_list]

app/service/group.py

@@ -95,6 +95,7 @@ async def save_user_to_group(db, user_id, group_id, user_list):
             for user1 in new_users:
                 for user2 in new_users:
                     if user1 != user2:
+                        print(user1, user2)
                         await ragflow_service.add_user_tenant(token, user_dict[user1]["rg_id"],
                                                               user_dict[user2]["email"],
                                                               user_dict[user2]["rg_id"])

app/service/ragflow.py

@@ -146,12 +146,15 @@ class RagflowService:
                 }
             ] if data else []
 
-    async def get_session_history(self, token: str, chat_id: str) -> list:
+    async def get_session_history(self, token: str, chat_id: str, is_all: int=0):
         url = f"{self.base_url}/v1/conversation/get?conversation_id={chat_id}"
         headers = {"Authorization": token}
         async with httpx.AsyncClient() as client:
             response = await client.get(url, headers=headers)
             data = self._handle_response(response)
+            # print("----------------data----------------------:", data)
+            if is_all:
+                return data
             return data.get("message", [])
 
     async def upload_and_parse(self, token: str, chat_id: str, filename: str, file: bytes) -> str:
@@ -172,6 +175,7 @@ class RagflowService:
         data = {"email": email, "user_id": user_id}
         async with httpx.AsyncClient(timeout=60) as client:
             response = await client.post(url, headers=headers, json=data)
+            print(response)
             if response.status_code != 200:
                 raise Exception(f"Ragflow add user to tenant failed: {response.text}")