From 93f9d6fc8738c579cb815084654f6663debd80ae Mon Sep 17 00:00:00 2001 From: zhangqian <zhangqian@123.com> Date: 星期五, 15 十一月 2024 15:51:59 +0800 Subject: [PATCH] 对话列表过滤name空数据 --- app/service/ragflow.py | 137 ++++++++++++++++++++++++++++++++++----------- 1 files changed, 102 insertions(+), 35 deletions(-) diff --git a/app/service/ragflow.py b/app/service/ragflow.py index e4d5657..3c45988 100644 --- a/app/service/ragflow.py +++ b/app/service/ragflow.py @@ -1,4 +1,7 @@ import httpx +from typing import Union, Dict, List +from fastapi import HTTPException +from starlette import status from app.config.config import settings from app.utils.rsa_crypto import RagflowCrypto @@ -8,28 +11,50 @@ def __init__(self, base_url: str): self.base_url = base_url + def _handle_response(self, response: httpx.Response) -> Union[Dict, List]: + if response.status_code != 200: + return {} + + data = response.json() + ret_code = data.get("retcode") + if ret_code == 401: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="鐧诲綍杩囨湡", + ) + if ret_code != 0: + return {} + + # 妫�鏌ヨ繑鍥炵殑鏁版嵁绫诲瀷 + if isinstance(data.get("data"), dict): + return data.get("data", {}) + elif isinstance(data.get("data"), list): + return data.get("data", []) + else: + return {} + async def register(self, username: str, password: str): password = RagflowCrypto(settings.PUBLIC_KEY, settings.PRIVATE_KEY).encrypt(password) async with httpx.AsyncClient() as client: response = await client.post( f"{self.base_url}/v1/user/register", - json={"nickname": username, "email": f"{username}@example.com", "password": password}, - headers={'Content-Type': 'application/json'} + headers={'Content-Type': 'application/json'}, + json={"nickname": username, "email": f"{username}@example.com", "password": password} ) if response.status_code != 200: raise Exception(f"Ragflow registration failed: {response.text}") + return self._handle_response(response) async def login(self, username: str, password: str) -> str: password = RagflowCrypto(settings.PUBLIC_KEY, settings.PRIVATE_KEY).encrypt(password) async with httpx.AsyncClient() as client: response = await client.post( f"{self.base_url}/v1/user/login", - json={"email": f"{username}@example.com", "password": password}, - headers={'Content-Type': 'application/json'} + headers={'Content-Type': 'application/json'}, + json={"email": f"{username}@example.com", "password": password} ) if response.status_code != 200: raise Exception(f"Ragflow login failed: {response.text}") - # 浠庡搷搴斿ご涓彁鍙� Authorization 瀛楁 authorization = response.headers.get('Authorization') if not authorization: raise Exception("Authorization header not found in response") @@ -40,19 +65,20 @@ "conversation_id": chat_id, "messages": chat_history } + + print(f"send to ragflow chat: {data}") + target_url = f"{self.base_url}/v1/conversation/completion" - async with httpx.AsyncClient() as client: + async with httpx.AsyncClient(timeout=300.0) as client: headers = { 'Content-Type': 'application/json', 'Authorization': token } - # 鍒涘缓娴佸紡璇锋眰 async with client.stream("POST", target_url, json=data, headers=headers) as response: - # 妫�鏌ュ搷搴旂姸鎬佺爜 if response.status_code == 200: - # 娴佸紡璇诲彇鍝嶅簲 try: async for answer in response.aiter_text(): + print(f"response of ragflow chat: {answer}") yield answer except GeneratorExit as e: print(e) @@ -62,15 +88,10 @@ async def get_chat_sessions(self, token: str, dialog_id: str) -> list: url = f"{self.base_url}/v1/conversation/list?dialog_id={dialog_id}" - headers = { - "Authorization": token - } + headers = {"Authorization": token} async with httpx.AsyncClient() as client: response = await client.get(url, headers=headers) - if response.status_code != 200: - raise Exception(f"Failed to fetch data from Ragflow API: {response.text}") - - data = response.json().get("data", []) + data = self._handle_response(response) result = [ { "id": item["id"], @@ -78,31 +99,77 @@ "updated_time": item["update_time"] } for item in data + if "id" in item and "name" in item and item["name"] ] return result - async def set_session(self, token: str, dialog_id: str, name: str, chat_id: str, is_new: bool) -> list: + async def get_session_log(self, token: str, conversation_id: str) -> dict: + url = f"{self.base_url}/v1/conversation/get?conversation_id={conversation_id}" + headers = {"Authorization": token} + async with httpx.AsyncClient() as client: + response = await client.get(url, headers=headers) + data = self._handle_response(response) + session_log = { + "session_log": [ + { + "message": message.get("content"), + "role": message.get("role"), + } + for message in data.get("message", []) + ], + "reference": data.get("reference"), + } + return session_log + + async def set_session(self, token: str, dialog_id: str, message: dict, chat_id: str, is_new: bool) -> list: url = f"{self.base_url}/v1/conversation/set?dialog_id={dialog_id}" - headers = { - "Authorization": token + headers = {"Authorization": token} + data = { + "dialog_id": dialog_id, + "name": message["message"], + "is_new": is_new, + "conversation_id": chat_id, } - - data = {"dialog_id": dialog_id, - "name": name, - "is_new": is_new, - "conversation_id": chat_id, - } - async with httpx.AsyncClient() as client: response = await client.post(url, headers=headers, json=data) - if response.status_code != 200: - return [] - return [{ - "content": "浣犲ソ锛� 鎴戞槸浣犵殑鍔╃悊锛屾湁浠�涔堝彲浠ュ府鍒颁綘鐨勫悧锛�", - "role": "assistant" - }, + data = self._handle_response(response) + return [ { - "content": name, - "doc_ids": [], + "content": "浣犲ソ锛� 鎴戞槸浣犵殑鍔╃悊锛屾湁浠�涔堝彲浠ュ府鍒颁綘鐨勫悧锛�", + "role": "assistant" + }, + { + "content": message["message"], + "doc_ids":message.get("doc_ids", []), "role": "user" - }] + } + ] if data else [] + + async def get_session_history(self, token: str, chat_id: str) -> list: + 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) + return data.get("message", []) + + async def upload_and_parse(self, token: str, chat_id: str, filename: str, file: bytes) -> str: + url = f"{self.base_url}/v1/document/upload_and_parse" + headers = {"Authorization": token} + data = {"conversation_id": chat_id} + + # 鍒涘缓琛ㄥ崟鏁版嵁锛屽寘鍚枃浠� + files = {"file": (filename, file)} + async with httpx.AsyncClient(timeout=60) as client: + response = await client.post(url, headers=headers, files=files, data=data) + data = self._handle_response(response) + return data + + async def add_user_tenant(self, token: str, tenant_id: str, email: str, user_id: str) -> str: + url = f"{self.base_url}/v1/tenant/{tenant_id}/user" + headers = {"Authorization": token} + data = {"email": email, "user_id": user_id} + async with httpx.AsyncClient(timeout=60) as client: + response = await client.post(url, headers=headers, json=data) + if response.status_code != 200: + raise Exception(f"Ragflow add user to tenant failed: {response.text}") -- Gitblit v1.8.0