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 class RagflowService: 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", 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", 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 = 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"] } 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())