From 9c275b214f9619a64cd2998596ce696610185eb4 Mon Sep 17 00:00:00 2001 From: zhaoqingang <zhaoqg0118@163.com> Date: 星期二, 19 十一月 2024 10:05:56 +0800 Subject: [PATCH] 增加删除 --- app/service/ragflow.py | 171 ++++++++++++++++++++++++++++++++++++++++++++++++++++++--- 1 files changed, 162 insertions(+), 9 deletions(-) diff --git a/app/service/ragflow.py b/app/service/ragflow.py index df131f1..f18f196 100644 --- a/app/service/ragflow.py +++ b/app/service/ragflow.py @@ -1,5 +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 @@ -8,28 +12,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,22 +66,149 @@ "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) 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"] + } + for item in data + ] + 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) -> 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}") + + 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} + async with httpx.AsyncClient(timeout=60) as client: + response = await client.get(url, headers=headers, params=params) + res = self._handle_response(response) + print(res) + return res + + + async def set_user_password(self, token: str, password: str, new_password: str) -> str: + # print("password:", password) + # print("new_password:", new_password) + 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, data=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("625e1f24a26811ef8f940242ac130006", 1, + 10) + print(b) + + import asyncio + # 鍦ㄥ悓姝ヤ唬鐮佷腑鍚姩寮傛涓诲嚱鏁� + asyncio.run(a()) + -- Gitblit v1.8.0