From 48e33c7fbd6969d05b2affda0321595d2d6b370a Mon Sep 17 00:00:00 2001 From: zhaoqingang <zhaoqg0118@163.com> Date: 星期一, 10 二月 2025 14:51:06 +0800 Subject: [PATCH] 增加ping機制 --- app/service/ragflow.py | 219 +++++++++++++++++++++++++++++++++++++++++++++++++++++- 1 files changed, 212 insertions(+), 7 deletions(-) diff --git a/app/service/ragflow.py b/app/service/ragflow.py index 699ffbe..f10902f 100644 --- a/app/service/ragflow.py +++ b/app/service/ragflow.py @@ -1,5 +1,11 @@ -import httpx +from datetime import datetime +import httpx +from typing import Union, Dict, List +from fastapi import HTTPException +from starlette import status + +# from Log import logger from app.config.config import settings from app.utils.rsa_crypto import RagflowCrypto @@ -8,25 +14,224 @@ def __init__(self, base_url: str): self.base_url = base_url - async def register(self, username: str, password: str): + 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, token: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}") - return response.json().get('data', {}).get('access_token') + authorization = response.headers.get('Authorization') + if not authorization: + raise Exception("Authorization header not found in response") + return authorization + + async def chat(self, token: str, chat_id: str, chat_history: list): + data = { + "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(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) + return + else: + yield f"Error: {response.status_code}" + + 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} + async with httpx.AsyncClient() as client: + response = await client.get(url, headers=headers) + data = self._handle_response(response) + result = [ + { + "id": item["id"], + "name": item["name"], + "updated_time": item["update_time"], + "update_date": datetime.fromtimestamp(item["update_time"] / 1000.0).strftime('%Y-%m-%d %H:%M:%S') + } + for item in data + if "id" in item and "name" in item and item["name"] + ] + return result + + 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} + data = { + "dialog_id": dialog_id, + "name": message["message"], + "is_new": is_new, + "conversation_id": chat_id, + } + async with httpx.AsyncClient() as client: + response = await client.post(url, headers=headers, json=data) + data = self._handle_response(response) + return [ + { + "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, 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: + 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) + print(response.text) + if response.status_code != 200: + raise Exception(f"Ragflow add user to tenant failed: {response.text}") + + async def get_knowledge_list(self, token: str, page_index: int, page_size: int) -> str: + url = f"{self.base_url}/v1/kb/list" + headers = {"Authorization": token} + # params = {"page": page_index, "page_size": page_size} + params = {} + async with httpx.AsyncClient(timeout=60) as client: + response = await client.get(url, headers=headers, params=params) + res = self._handle_response(response) + return res + + + async def set_user_password(self, token: str, password: str, new_password: str) -> str: + # print("password:", password) + # print("new_password:", new_password) + headers = { + 'Content-Type': 'application/json', + 'Authorization': token + } + password = RagflowCrypto(settings.PUBLIC_KEY, settings.PRIVATE_KEY).encrypt(password) + new_password = RagflowCrypto(settings.PUBLIC_KEY, settings.PRIVATE_KEY).encrypt(new_password) + url = f"{self.base_url}/v1/user/setting" + # headers = {"Authorization": token} + data = {"password": password, "new_password":new_password} + + async with httpx.AsyncClient(timeout=60) as client: + response = await client.post(url, headers=headers, json=data) + data = self._handle_response(response) + # logger.info("set_user_password:{}".format(data)) + return data + + async def agree_tenant_role(self, token: str, tenant_id: str) -> str: + url = f"{self.base_url}/v1/tenant/agree/{tenant_id}" + headers = {"Authorization": token} + data = {} + async with httpx.AsyncClient(timeout=60) as client: + response = await client.put(url, headers=headers, json=data) + print(response.text) + if response.status_code != 200: + raise Exception(f"Ragflow add user to tenant failed: {response.text}") + + +if __name__ == "__main__": + async def a(): + a = RagflowService("http://192.168.20.119:11080") + b = await a.agree_tenant_role("ImQ3MGY1NGFhYzFkZjExZWY5MDFmMDI0MmFjMTMwMDA2Ig.Z2qJfA.hJX4rPnnmlkUzpjIl0YGcM-yp1M", "fe24dd2c9be611ef92880242ac160006") + print(b) + print(b) + + import asyncio + + asyncio.run(a()) + # password = RagflowCrypto(settings.PUBLIC_KEY, settings.PRIVATE_KEY).encrypt("123456") + # print(password) -- Gitblit v1.8.0