From 081a109d369966182e5d6e4ff3ab38bb8512b498 Mon Sep 17 00:00:00 2001 From: zhaoqingang <zhaoqg0118@163.com> Date: 星期三, 20 十一月 2024 15:59:01 +0800 Subject: [PATCH] 优化 --- app/service/ragflow.py | 82 +++++++++++++++++++++++++++++++++++++++- 1 files changed, 79 insertions(+), 3 deletions(-) diff --git a/app/service/ragflow.py b/app/service/ragflow.py index 7ce287d..34e8307 100644 --- a/app/service/ragflow.py +++ b/app/service/ragflow.py @@ -1,9 +1,9 @@ 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 @@ -40,10 +40,11 @@ response = await client.post( f"{self.base_url}/v1/user/register", headers={'Content-Type': 'application/json'}, - json={"nickname": username, "email": f"{username}@example.com", "password": password} + 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) @@ -99,8 +100,27 @@ "updated_time": item["update_time"] } 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}" @@ -126,12 +146,15 @@ } ] if data else [] - async def get_session_history(self, token: str, chat_id: str) -> list: + 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: @@ -145,3 +168,56 @@ 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) + 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 + + +if __name__ == "__main__": + async def a(): + a = RagflowService("http://192.168.20.119:11080") + b = await a.get_knowledge_list("ImY3ZTZlZWQwYTY2NTExZWY5ZmFiMDI0MmFjMTMwMDA2Ig.Zzxwmw.uI_HAWzOkipQuga1aeQtoeIc3IM", 1, + 10) + print(b) + + import asyncio + + asyncio.run(a()) + -- Gitblit v1.8.0