rag-gateway.git

			@@ -1,4 +1,5 @@
			import asyncio
			import datetime
			import io
			import json
			import time
			@@ -6,14 +7,15 @@

			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, \
			DF_CHAT_WORKFLOW, DF_UPLOAD_FILE, RG_ORIGINAL_URL
			DF_CHAT_WORKFLOW, DF_UPLOAD_FILE, RG_ORIGINAL_URL, RG_CHAT_UPDATE_URL, DF_WORKFLOW_DRAFT, DF_WORKFLOW_PUBLISH
			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 +35,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 +43,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
			)
			@@ -63,6 +65,10 @@
			if app_token:
			return app_token.token
			return ""

			async def get_workflow_token(db):
			user_token = db.query(UserTokenModel).filter(UserTokenModel.id == workflow_server).first()
			return user_token.access_token if user_token else ""


			async def add_chat_token(db, data):
			@@ -87,17 +93,49 @@
			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": "\nERROR: 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 +213,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:
			@@ -253,15 +291,11 @@
			"error": error}, conversation_id)




			async def service_chat_basic(db, chat_id: str, chat_data: ChatData, session_id: str, user_id, mode: str):

			if chat_id == basic_report_talk:
			complex_chat = await ComplexChatDao(db).get_complex_chat_by_mode(chat_data.report_mode)
			if complex_chat:
			...



			async def service_chat_parameters(db, chat_id, user_id):
			@@ -296,9 +330,8 @@
			session_log = await ChatSessionDao(db).get_session_by_id(session_id)
			if not session_log:
			return {}
			log_info =session_log.log_to_json()
			if session_log.event_type == complex_chat:

			log_info = session_log.log_to_json()
			if session_log.event_type == complex_chat:
			total, message_list = await ComplexChatSessionDao(db).get_session_list(session_id)
			log_info["message"] = [message.log_to_json() for message in message_list[::-1]]

			@@ -446,7 +479,8 @@
			return records


			async def add_complex_log(db, message_id, chat_id, session_id, chat_mode, query, user_id, mode, agent_type, message_type, conversation_id="", node_data=None, query_data=None):
			async def add_complex_log(db, message_id, chat_id, session_id, chat_mode, query, user_id, mode, agent_type,
			message_type, conversation_id="", node_data=None, query_data=None):
			if not node_data:
			node_data = []
			if not query_data:
			@@ -461,22 +495,23 @@
			if session:
			conversation_id = session.conversation_id
			await complex_log.create_session(message_id,
			chat_id=chat_id,
			session_id=session_id,
			chat_mode=chat_mode,
			message_type=message_type,
			content=query,
			event_type=mode,
			tenant_id=user_id,
			conversation_id=conversation_id,
			node_data=json.dumps(node_data),
			query=json.dumps(query_data),
			agent_type=agent_type)
			chat_id=chat_id,
			session_id=session_id,
			chat_mode=chat_mode,
			message_type=message_type,
			content=query,
			event_type=mode,
			tenant_id=user_id,
			conversation_id=conversation_id,
			node_data=json.dumps(node_data),
			query=json.dumps(query_data),
			agent_type=agent_type)
			return conversation_id, True

			except Exception as e:
			logger.error(e)
			return conversation_id, False


			async def add_query_files(db, message_id):
			query = {}
			@@ -484,6 +519,7 @@
			if complex_log:
			query = json.loads(complex_log.query)
			return query.get("files", [])


			async def service_complex_chat(db, chat_id, mode, user_id, chat_request: ChatDataRequest):
			answer_event = ""
			@@ -495,7 +531,7 @@
			task_id = ""
			error = ""
			node_list = []
			reference= {}
			reference = {}
			conversation_id = ""
			query_data = chat_request.to_dict()
			new_message_id = str(uuid.uuid4())
			@@ -507,9 +543,12 @@
			if not files and chat_request.parentId:
			files = await add_query_files(db, chat_request.parentId)
			if chat_request.chatMode != complex_content_optimization_chat:
			await add_session_log(db, chat_request.sessionId, chat_request.query if chat_request.query else "未命名会话", chat_id, user_id,
			mode, "", DF_TYPE)
			conversation_id, message = await add_complex_log(db, new_message_id, chat_id, chat_request.sessionId, chat_request.chatMode, chat_request.query, user_id, mode, DF_TYPE, 1, query_data=query_data)
			await add_session_log(db, chat_request.sessionId, chat_request.query if chat_request.query else "未命名会话",
			chat_id, user_id,
			mode, "", DF_TYPE)
			conversation_id, message = await add_complex_log(db, new_message_id, chat_id, chat_request.sessionId,
			chat_request.chatMode, chat_request.query, user_id, mode,
			DF_TYPE, 1, query_data=query_data)
			if not message:
			yield "data: " + json.dumps({"message": smart_message_error,
			"error": "\nERROR: 创建会话失败！", "status": http_500},
			@@ -518,12 +557,12 @@

			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)
			if not session or session.get("code") != 0:
			yield "data: " + json.dumps(
			yield "data: " + json.dumps(
			{"message": smart_message_error, "error": "\nERROR: chat agent error", "status": http_500})
			return
			conversation_id = session.get("data", {}).get("id")
			@@ -531,7 +570,9 @@
			url = settings.fwr_base_url + RG_CHAT_DIALOG.format(chat_id)
			chat = ChatDialog()
			try:
			async for ans in chat.chat_completions(url, await chat.complex_request_data(chat_request.query, chat_request.knowledgeId, conversation_id),
			async for ans in chat.chat_completions(url, await chat.complex_request_data(chat_request.query,
			chat_request.knowledgeId,
			conversation_id),
			await chat.get_headers(token)):
			data = {}
			error = ""
			@@ -555,9 +596,9 @@
			reference = data.get("reference", {})
			event = smart_message_cover
			message_str = "data: " + json.dumps(
			{"event": event, "data": data, "error": error, "status": status, "message_id":message_id,
			"parent_id": new_message_id,
			"session_id": chat_request.sessionId},
			{"event": event, "data": data, "error": error, "status": status, "message_id": message_id,
			"parent_id": new_message_id,
			"session_id": chat_request.sessionId},
			ensure_ascii=False) + "\n\n"
			for i in range(0, len(message_str), max_chunk_size):
			chunk = message_str[i:i + max_chunk_size]
			@@ -576,7 +617,9 @@
			token = await get_chat_token(db, chat_id)
			chat, url = await get_chat_object(mode)
			async for ans in chat.chat_completions(url,
			await chat.complex_request_data(chat_request.query, conversation_id, str(user_id), files=files, inputs=inputs),
			await chat.complex_request_data(chat_request.query, conversation_id,
			str(user_id), files=files,
			inputs=inputs),
			await chat.get_headers(token)):
			# print(ans)
			data = {}
			@@ -630,7 +673,8 @@
			continue

			yield "data: " + json.dumps(
			{"event": event, "data": data, "error": error, "status": status, "task_id": task_id, "message_id":message_id,
			{"event": event, "data": data, "error": error, "status": status, "task_id": task_id,
			"message_id": message_id,
			"parent_id": new_message_id,
			"session_id": chat_request.sessionId},
			ensure_ascii=False) + "\n\n"
			@@ -650,7 +694,11 @@
			# "node_list": node_list, "task_id": task_id, "id": message_id,
			# "error": error}, conversation_id)
			if message_id:
			await add_complex_log(db, message_id, chat_id, chat_request.sessionId, chat_request.chatMode, answer_event or answer_agent or answer_workflow or answer_dialog or error, user_id, mode, DF_TYPE, 2, conversation_id, node_data=node_list or reference, query_data=query_data)
			await add_complex_log(db, message_id, chat_id, chat_request.sessionId, chat_request.chatMode,
			answer_event or answer_agent or answer_workflow or answer_dialog or error, user_id,
			mode, DF_TYPE, 2, conversation_id, node_data=node_list or reference,
			query_data=query_data)


			async def service_complex_upload(db, chat_id, file, user_id):
			files = []
			@@ -674,6 +722,92 @@
			logger.error(e)
			return json.dumps(files) if files else ""


			async def service_complex_model(db, chat_type, model_type, model_name, model_provider):
			if chat_type == 1 and model_type == 1:
			return await set_dialog_model(db, complex_knowledge_chat, model_name)
			elif chat_type == 1 and model_type == 2:
			return await set_dialog_model(db, complex_knowledge_chat_deep, model_name, model_provider)
			else:
			if model_type == 1:
			chats = [complex_dialog_chat, complex_network_chat, complex_mindmap_chat, complex_content_optimization_chat]
			else:
			chats = [complex_dialog_chat, complex_network_chat]
			return await set_workflow_model(db,chats
			, # , complex_network_chat, complex_mindmap_chat, complex_content_optimization_chat
			model_type, model_name, model_provider)


			async def set_dialog_model(db, chat_mode, model_name, model_provider):
			chat = await ComplexChatDao(db).get_complex_chat_by_mode(chat_mode)
			if chat:
			access_token = await get_chat_token(db, rg_api_token)
			url = settings.fwr_base_url + RG_CHAT_UPDATE_URL.format(chat.id)
			chat_base = ChatBaseApply()
			payload = {
			"name": chat.name,
			"llm": {
			"model_name": model_name
			}
			}
			response = await chat_base.chat_put(url, payload, await chat_base.get_headers(access_token))
			# print(response)
			if not response:
			return "服务异常,修改失败！"
			await ComplexChatDao(db).update_complex_chat_by_id(chat.id, {"chat_model": model_name, "chat_model_ds": model_name, "chat_provider": model_provider, "update_date": datetime.datetime.now()})
			return ""

			async def set_workflow_model(db, chat_modes, model_type, model_name, model_provider):
			chat_base = ChatBaseApply()
			token = await get_workflow_token(db)
			for chat_mode in chat_modes:
			chat = await ComplexChatDao(db).get_complex_chat_by_mode(chat_mode)
			if chat:

			get_draft_url = settings.dify_base_url + DF_WORKFLOW_DRAFT.format(chat.id)
			draft_data = await chat_base.chat_get(get_draft_url, {}, await chat_base.get_headers(token))
			if draft_data:
			graph = draft_data.get("graph")
			for node in graph.get("nodes"):
			if node.get("data", {}).get("type") == "llm":
			if model_type == 1 and "深度搜索" not in node.get("data", {}).get("title"):
			node["data"]["model"]["name"] = model_name
			node["data"]["model"]["provider"] = model_provider
			elif model_type == 2 and "深度搜索" in node.get("data", {}).get("title"):
			node["data"]["model"]["name"] = model_name
			node["data"]["model"]["provider"] = model_provider

			draft_data_query = {"conversation_variables": draft_data.get("conversation_variables"),
			"environment_variables": draft_data.get("environment_variables"),
			"hash": draft_data.get("hash"),
			"features": draft_data.get("features"),
			"graph": graph}
			set_draft_data = await chat_base.chat_post(get_draft_url, draft_data_query, await chat_base.get_headers(token))
			if set_draft_data and set_draft_data.get("result") == "success":
			publish_url = settings.dify_base_url + DF_WORKFLOW_PUBLISH.format(chat.id)
			publish_data = await chat_base.chat_post(publish_url, {}, await chat_base.get_headers(token))
			if publish_data and publish_data.get("result") == "success":
			update_kwargs = {"chat_provider": model_provider, "update_date": datetime.datetime.now()}
			if model_type == 1:
			update_kwargs["chat_model"] = model_name
			else:
			update_kwargs["chat_model_ds"] = model_name
			await ComplexChatDao(db).update_complex_chat_by_id(chat.id, update_kwargs)


			async def service_get_complex_model(db):
			res = {}
			for complexs in await ComplexChatDao(db).aget_complex_chat():
			if complexs.chat_mode == complex_knowledge_chat:
			res["dialog"] = {"modelName": complexs.chat_model, "modelProvider": complexs.chat_provider}
			elif complexs.chat_mode == complex_knowledge_chat_deep:
			res["dialog_ds"] = {"modelName": complexs.chat_model_ds, "modelProvider": complexs.chat_provider}
			else:
			res["workflow"] = {"modelName": complexs.chat_model, "modelProvider": complexs.chat_provider}
			res["workflow_ds"] = {"modelName": complexs.chat_model_ds, "modelProvider": complexs.chat_provider}

			return json.dumps(res)

			if __name__ == "__main__":
			q = json.dumps({"query": "设备", "dataset_ids": ["fc68db52f43111efb94a0242ac120004"]})
			top_k = 2