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