From b2e47e75a231baf5a7beca476256ab3d94e76c46 Mon Sep 17 00:00:00 2001
From: zhaoqingang <zhaoqg0118@163.com>
Date: 星期二, 01 四月 2025 10:08:15 +0800
Subject: [PATCH] code加密

---
 app/service/v2/chat.py |   46 +++++++++++++++++++++++++++++++++++++++-------
 1 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/app/service/v2/chat.py b/app/service/v2/chat.py
index 38683a8..9f83c63 100644
--- a/app/service/v2/chat.py
+++ b/app/service/v2/chat.py
@@ -6,6 +6,7 @@
 
 import fitz
 from fastapi import HTTPException
+from sqlalchemy import or_
 
 from Log import logger
 from app.config.agent_base_url import RG_CHAT_DIALOG, DF_CHAT_AGENT, DF_CHAT_PARAMETERS, RG_CHAT_SESSIONS, \
@@ -13,7 +14,7 @@
 from app.config.config import settings
 from app.config.const import *
 from app.models import DialogModel, ApiTokenModel, UserTokenModel, ComplexChatSessionDao, ChatDataRequest, \
-    ComplexChatDao
+    ComplexChatDao, KnowledgeModel, UserModel, KnowledgeUserModel
 from app.models.v2.session_model import ChatSessionDao, ChatData
 from app.service.v2.app_driver.chat_agent import ChatAgent
 from app.service.v2.app_driver.chat_data import ChatBaseApply
@@ -33,7 +34,7 @@
 
 
 async def add_session_log(db, session_id: str, question: str, chat_id: str, user_id, event_type: str,
-                          conversation_id: str, agent_type):
+                          conversation_id: str, agent_type, query: dict=None):
     try:
         session = await ChatSessionDao(db).update_or_insert_by_id(
             session_id=session_id,
@@ -41,7 +42,7 @@
             agent_id=chat_id,
             agent_type=agent_type,
             tenant_id=user_id,
-            message={"role": "user", "content": question},
+            message={"role": "user", "content": question, "query": query},
             conversation_id=conversation_id,
             event_type=event_type
         )
@@ -87,17 +88,48 @@
         return ChatAgent(), url
 
 
-async def service_chat_dialog(db, chat_id: str, question: str, session_id: str, user_id, mode: str):
+
+async def get_user_kb(db, user_id: int, kb_ids: list) -> list:
+    res = []
+    user = db.query(UserModel).filter(UserModel.id == user_id).first()
+    if user is None:
+        return res
+    query = db.query(KnowledgeModel)
+    if user.permission != "admin":
+        klg_list = [j.id for i in user.groups for j in i.knowledges]
+        for i in db.query(KnowledgeUserModel).filter(KnowledgeUserModel.user_id == user_id, KnowledgeUserModel.status == 1).all():
+            if i.kb_id not in klg_list:
+                klg_list.append(i.kb_id)
+        query = query.filter(or_(KnowledgeModel.id.in_(klg_list), KnowledgeModel.tenant_id == str(user_id)))
+        kb_list= query.all()
+        for kb in kb_list:
+            if kb.id in kb_ids:
+                if kb.permission == "team":
+                    res.append(kb.id)
+                elif kb.tenant_id == str(user_id):
+                    res.append(kb.id)
+        return res
+    else:
+        return kb_ids
+
+
+async def service_chat_dialog(db, chat_id: str, question: str, session_id: str, user_id: int, mode: str, kb_ids: list):
     conversation_id = ""
     token = await get_chat_token(db, rg_api_token)
     url = settings.fwr_base_url + RG_CHAT_DIALOG.format(chat_id)
+    kb_id = await get_user_kb(db, user_id, kb_ids)
+    if not kb_id:
+        yield "data: " + json.dumps({"message": smart_message_error,
+                                     "error": "\n**ERROR**: The agent has no knowledge base to work with!", "status": http_400},
+                                    ensure_ascii=False) + "\n\n"
+        return
     chat = ChatDialog()
     session = await add_session_log(db, session_id, question, chat_id, user_id, mode, session_id, RG_TYPE)
     if session:
         conversation_id = session.conversation_id
     message = {"role": "assistant", "answer": "", "reference": {}}
     try:
-        async for ans in chat.chat_completions(url, await chat.request_data(question, conversation_id),
+        async for ans in chat.chat_completions(url, await chat.complex_request_data(question, kb_id, conversation_id),
                                                await chat.get_headers(token)):
             data = {}
             error = ""
@@ -175,7 +207,7 @@
     else:
         query = "start new conversation"
     session = await add_session_log(db, session_id, query if query else "start new conversation", chat_id, user_id,
-                                    mode, conversation_id, DF_TYPE)
+                                    mode, conversation_id, DF_TYPE, chat_data.to_dict())
     if session:
         conversation_id = session.conversation_id
     try:
@@ -518,7 +550,7 @@
 
     query_data["parentId"] = new_message_id
     try:
-        if chat_request.chatMode == complex_knowledge_chat:
+        if chat_request.chatMode == complex_knowledge_chat or chat_request.chatMode == complex_knowledge_chat_deep:
             if not conversation_id:
                 session = await service_chat_sessions(db, chat_id, chat_request.query)
                 # print(session)

--
Gitblit v1.8.0