1044 lines
38 KiB
Python
1044 lines
38 KiB
Python
import aiohttp
|
||
import asyncio
|
||
import json
|
||
from typing import Optional, List, Dict, Any, Union, AsyncGenerator
|
||
import os
|
||
from pathlib import Path
|
||
from urllib.parse import unquote
|
||
|
||
|
||
class RAGFlowError(Exception):
|
||
"""RAGFlow API错误异常"""
|
||
def __init__(self, code: int, message: str):
|
||
self.code = code
|
||
self.message = message
|
||
super().__init__(f"Error {code}: {message}")
|
||
|
||
|
||
class AsyncRAGFlowClient:
|
||
"""异步RAGFlow API客户端"""
|
||
|
||
def __init__(self, base_url: str, api_key: str, timeout: int = 30):
|
||
"""
|
||
初始化RAGFlow客户端
|
||
|
||
Args:
|
||
base_url: RAGFlow服务器地址
|
||
api_key: API密钥
|
||
timeout: 请求超时时间（秒）
|
||
"""
|
||
try:
|
||
self.base_url = base_url.rstrip('/')
|
||
except:
|
||
self.base_url = base_url
|
||
self.api_key = api_key
|
||
self.timeout = timeout
|
||
self.headers = {
|
||
'Authorization': f'Bearer {api_key}',
|
||
'Content-Type': 'application/json'
|
||
}
|
||
self._session = None
|
||
|
||
async def __aenter__(self):
|
||
"""异步上下文管理器入口"""
|
||
await self.create_session()
|
||
return self
|
||
|
||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||
"""异步上下文管理器出口"""
|
||
await self.close_session()
|
||
|
||
async def create_session(self):
|
||
"""创建aiohttp会话"""
|
||
if self._session is None:
|
||
timeout = aiohttp.ClientTimeout(total=self.timeout)
|
||
self._session = aiohttp.ClientSession(timeout=timeout)
|
||
|
||
async def close_session(self):
|
||
"""关闭aiohttp会话"""
|
||
if self._session:
|
||
await self._session.close()
|
||
self._session = None
|
||
|
||
async def _request(self, method: str, endpoint: str, **kwargs) -> Dict[str, Any]:
|
||
"""发送HTTP请求"""
|
||
if not self._session:
|
||
await self.create_session()
|
||
|
||
url = f"{self.base_url}{endpoint}"
|
||
|
||
# 处理headers
|
||
headers = kwargs.pop('headers', self.headers.copy())
|
||
|
||
async with self._session.request(method, url, headers=headers, **kwargs) as response:
|
||
try:
|
||
result = await response.json()
|
||
except (aiohttp.ContentTypeError, json.JSONDecodeError):
|
||
if response.status == 200:
|
||
content = await response.read()
|
||
return {'code': 0, 'data': content}
|
||
else:
|
||
text = await response.text()
|
||
raise RAGFlowError(response.status, text)
|
||
|
||
if result.get('code', 0) != 0:
|
||
raise RAGFlowError(result.get('code'), result.get('message', 'Unknown error'))
|
||
|
||
return result
|
||
|
||
async def _stream_request(self, method: str, endpoint: str, **kwargs) -> AsyncGenerator[Dict[str, Any], None]:
|
||
"""发送流式HTTP请求"""
|
||
if not self._session:
|
||
await self.create_session()
|
||
|
||
url = f"{self.base_url}{endpoint}"
|
||
headers = kwargs.pop('headers', self.headers.copy())
|
||
|
||
async with self._session.request(method, url, headers=headers, **kwargs) as response:
|
||
async for line in response.content:
|
||
if line:
|
||
line_str = line.decode('utf-8').strip()
|
||
if line_str.startswith('data:'):
|
||
try:
|
||
data = json.loads(line_str[5:].strip())
|
||
yield data
|
||
except json.JSONDecodeError:
|
||
continue
|
||
|
||
# ====================
|
||
# OpenAI兼容API
|
||
# ====================
|
||
|
||
async def create_chat_completion(self, chat_id: str, model: str, messages: List[Dict[str, str]],
|
||
stream: bool = False) -> Union[Dict[str, Any], AsyncGenerator[Dict[str, Any], None]]:
|
||
"""
|
||
创建聊天完成
|
||
|
||
Args:
|
||
chat_id: 聊天ID
|
||
model: 模型名称
|
||
messages: 消息列表
|
||
stream: 是否流式返回
|
||
"""
|
||
endpoint = f"/api/v1/chats_openai/{chat_id}/chat/completions"
|
||
data = {
|
||
"model": model,
|
||
"messages": messages,
|
||
"stream": stream
|
||
}
|
||
|
||
if stream:
|
||
return self._stream_request('POST', endpoint, json=data)
|
||
else:
|
||
return await self._request('POST', endpoint, json=data)
|
||
|
||
async def create_agent_completion(self, agent_id: str, model: str, messages: List[Dict[str, str]],
|
||
stream: bool = False) -> Union[Dict[str, Any], AsyncGenerator[Dict[str, Any], None]]:
|
||
"""
|
||
创建代理完成
|
||
|
||
Args:
|
||
agent_id: 代理ID
|
||
model: 模型名称
|
||
messages: 消息列表
|
||
stream: 是否流式返回
|
||
"""
|
||
endpoint = f"/api/v1/agents_openai/{agent_id}/chat/completions"
|
||
data = {
|
||
"model": model,
|
||
"messages": messages,
|
||
"stream": stream
|
||
}
|
||
|
||
if stream:
|
||
return self._stream_request('POST', endpoint, json=data)
|
||
else:
|
||
return await self._request('POST', endpoint, json=data)
|
||
|
||
# ====================
|
||
# 数据集管理
|
||
# ====================
|
||
|
||
async def create_dataset(self, name: str, avatar: Optional[str] = None, description: Optional[str] = None,
|
||
embedding_model: Optional[str] = None, permission: str = "me",
|
||
chunk_method: str = "naive",
|
||
parser_config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
"""
|
||
创建数据集
|
||
|
||
Args:
|
||
name: 数据集名称
|
||
avatar: Base64编码的头像
|
||
description: 描述
|
||
embedding_model: 嵌入模型
|
||
permission: 权限设置 ("me" 或 "team")
|
||
chunk_method: 分块方法
|
||
parser_config: 解析器配置
|
||
"""
|
||
endpoint = "/api/v1/datasets"
|
||
data = {
|
||
"name": name,
|
||
"permission": permission,
|
||
"chunk_method": chunk_method,
|
||
}
|
||
|
||
if avatar:
|
||
data["avatar"] = avatar
|
||
if description:
|
||
data["description"] = description
|
||
if embedding_model:
|
||
data["embedding_model"] = embedding_model
|
||
if parser_config:
|
||
data["parser_config"] = parser_config
|
||
|
||
return await self._request('POST', endpoint, json=data)
|
||
|
||
async def delete_datasets(self, ids: Optional[List[str]] = None) -> Dict[str, Any]:
|
||
"""
|
||
删除数据集
|
||
|
||
Args:
|
||
ids: 要删除的数据集ID列表，None表示删除所有
|
||
"""
|
||
endpoint = "/api/v1/datasets"
|
||
data = {"ids": ids}
|
||
return await self._request('DELETE', endpoint, json=data)
|
||
|
||
async def update_dataset(self, dataset_id: str, name: Optional[str] = None,
|
||
avatar: Optional[str] = None, description: Optional[str] = None,
|
||
embedding_model: Optional[str] = None, permission: Optional[str] = None,
|
||
chunk_method: Optional[str] = None, pagerank: Optional[int] = None,
|
||
parser_config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
"""
|
||
更新数据集
|
||
"""
|
||
endpoint = f"/api/v1/datasets/{dataset_id}"
|
||
data = {}
|
||
|
||
if name is not None:
|
||
data["name"] = name
|
||
if avatar is not None:
|
||
data["avatar"] = avatar
|
||
if description is not None:
|
||
data["description"] = description
|
||
if embedding_model is not None:
|
||
data["embedding_model"] = embedding_model
|
||
if permission is not None:
|
||
data["permission"] = permission
|
||
if chunk_method is not None:
|
||
data["chunk_method"] = chunk_method
|
||
if pagerank is not None:
|
||
data["pagerank"] = pagerank
|
||
if parser_config is not None:
|
||
data["parser_config"] = parser_config
|
||
|
||
return await self._request('PUT', endpoint, json=data)
|
||
|
||
async def list_datasets(self, page: int = 1, page_size: int = 30, orderby: str = "create_time",
|
||
desc: str = "true", name: Optional[str] = None,
|
||
dataset_id: Optional[str] = None) -> Dict[str, Any]:
|
||
"""
|
||
列出数据集
|
||
"""
|
||
endpoint = "/api/v1/datasets"
|
||
params = {
|
||
"page": page,
|
||
"page_size": page_size,
|
||
"orderby": orderby,
|
||
"desc": desc
|
||
}
|
||
|
||
if name:
|
||
params["name"] = name
|
||
if dataset_id:
|
||
params["id"] = dataset_id
|
||
|
||
return await self._request('GET', endpoint, params=params)
|
||
|
||
# ====================
|
||
# 文档管理
|
||
# ====================
|
||
|
||
async def upload_documents(self, dataset_id: str, file_paths: List[str]) -> Dict[str, Any]:
|
||
"""
|
||
上传文档到数据集
|
||
|
||
Args:
|
||
dataset_id: 数据集ID
|
||
file_paths: 文件路径列表
|
||
"""
|
||
if not self._session:
|
||
await self.create_session()
|
||
|
||
endpoint = f"/api/v1/datasets/{dataset_id}/documents"
|
||
url = f"{self.base_url}{endpoint}"
|
||
|
||
# 准备multipart数据
|
||
data = aiohttp.FormData()
|
||
for file_path in file_paths:
|
||
if os.path.exists(file_path):
|
||
file_name = Path(file_path).name
|
||
with open(file_path, 'rb') as f:
|
||
data.add_field('file', f.read(), filename=file_name)
|
||
|
||
headers = {'Authorization': f'Bearer {self.api_key}'}
|
||
|
||
async with self._session.post(url, headers=headers, data=data) as response:
|
||
try:
|
||
result = await response.json()
|
||
if result.get('code', 0) != 0:
|
||
raise RAGFlowError(result.get('code'), result.get('message'))
|
||
return result
|
||
except (aiohttp.ContentTypeError, json.JSONDecodeError):
|
||
text = await response.text()
|
||
raise RAGFlowError(response.status, text)
|
||
|
||
async def upload_documents_bytes(self, dataset_id: str, file_bytes: List) -> Dict[str, Any]:
|
||
"""
|
||
上传文档到数据集
|
||
|
||
Args:
|
||
dataset_id: 数据集ID
|
||
file_name: 文件名
|
||
file_bytes: 文件二进制列表
|
||
"""
|
||
if not self._session:
|
||
await self.create_session()
|
||
|
||
endpoint = f"/api/v1/datasets/{dataset_id}/documents"
|
||
url = f"{self.base_url}{endpoint}"
|
||
|
||
# 准备multipart数据
|
||
data = aiohttp.FormData()
|
||
for file in file_bytes:
|
||
|
||
data.add_field('file', file.file.read(),filename=file.filename)
|
||
|
||
headers = {'Authorization': f'Bearer {self.api_key}'}
|
||
|
||
async with self._session.post(url, headers=headers, data=data) as response:
|
||
try:
|
||
result = await response.json()
|
||
if result.get('code', 0) != 0:
|
||
raise RAGFlowError(result.get('code'), result.get('message'))
|
||
return result
|
||
except (aiohttp.ContentTypeError, json.JSONDecodeError):
|
||
text = await response.text()
|
||
raise RAGFlowError(response.status, text)
|
||
|
||
async def update_document(self, dataset_id: str, document_id: str, name: Optional[str] = None,
|
||
meta_fields: Optional[Dict[str, Any]] = None,
|
||
chunk_method: Optional[str] = None,
|
||
parser_config: Optional[Dict[str, Any]] = None,
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
更新文档配置
|
||
"""
|
||
endpoint = f"/api/v1/datasets/{dataset_id}/documents/{document_id}"
|
||
data = {}
|
||
|
||
if name is not None:
|
||
data["name"] = name
|
||
if meta_fields is not None:
|
||
data["meta_fields"] = meta_fields
|
||
if chunk_method is not None:
|
||
data["chunk_method"] = chunk_method
|
||
if parser_config is not None:
|
||
data["parser_config"] = parser_config
|
||
|
||
|
||
return await self._request('PUT', endpoint, json=data)
|
||
|
||
async def download_document(self, dataset_id: str, document_id: str, save_path: str) -> None:
|
||
"""
|
||
下载文档
|
||
"""
|
||
if not self._session:
|
||
await self.create_session()
|
||
|
||
endpoint = f"/api/v1/datasets/{dataset_id}/documents/{document_id}"
|
||
url = f"{self.base_url}{endpoint}"
|
||
headers = {'Authorization': f'Bearer {self.api_key}'}
|
||
|
||
async with self._session.get(url, headers=headers) as response:
|
||
if response.status == 200:
|
||
content = await response.read()
|
||
with open(save_path, 'wb') as f:
|
||
f.write(content)
|
||
else:
|
||
try:
|
||
error = await response.json()
|
||
raise RAGFlowError(error.get('code'), error.get('message'))
|
||
except (aiohttp.ContentTypeError, json.JSONDecodeError):
|
||
text = await response.text()
|
||
raise RAGFlowError(response.status, text)
|
||
|
||
async def list_documents(self, dataset_id: str, page: int = 1, page_size: int = 30,
|
||
orderby: str = "create_time", desc: str = "true",
|
||
keywords: Optional[str] = None, document_id: Optional[str] = None,
|
||
document_name: Optional[str] = None) -> Dict[str, Any]:
|
||
"""
|
||
列出数据集中的文档
|
||
"""
|
||
endpoint = f"/api/v1/datasets/{dataset_id}/documents"
|
||
params = {
|
||
"page": page,
|
||
"page_size": page_size,
|
||
"orderby": orderby,
|
||
"desc": desc
|
||
}
|
||
|
||
if keywords:
|
||
params["keywords"] = keywords
|
||
if document_id:
|
||
params["id"] = document_id
|
||
if document_name:
|
||
params["name"] = document_name
|
||
|
||
return await self._request('GET', endpoint, params=params)
|
||
|
||
async def delete_documents(self, dataset_id: str, ids: Optional[List[str]] = None) -> Dict[str, Any]:
|
||
"""
|
||
删除文档
|
||
"""
|
||
endpoint = f"/api/v1/datasets/{dataset_id}/documents"
|
||
data = {"ids": ids} if ids else {}
|
||
return await self._request('DELETE', endpoint, json=data)
|
||
|
||
async def parse_documents(self, dataset_id: str, document_ids: List[str]) -> Dict[str, Any]:
|
||
"""
|
||
解析文档
|
||
"""
|
||
endpoint = f"/api/v1/datasets/{dataset_id}/chunks"
|
||
data = {"document_ids": document_ids}
|
||
return await self._request('POST', endpoint, json=data)
|
||
|
||
async def stop_parsing_documents(self, dataset_id: str, document_ids: List[str]) -> Dict[str, Any]:
|
||
"""
|
||
停止解析文档
|
||
"""
|
||
endpoint = f"/api/v1/datasets/{dataset_id}/chunks"
|
||
data = {"document_ids": document_ids}
|
||
return await self._request('DELETE', endpoint, json=data)
|
||
|
||
# ====================
|
||
# 分块管理
|
||
# ====================
|
||
|
||
async def add_chunk(self, dataset_id: str, document_id: str, content: str,
|
||
important_keywords: Optional[List[str]] = None,
|
||
questions: Optional[List[str]] = None) -> Dict[str, Any]:
|
||
"""
|
||
添加分块
|
||
"""
|
||
endpoint = f"/api/v1/datasets/{dataset_id}/documents/{document_id}/chunks"
|
||
data = {"content": content}
|
||
|
||
if important_keywords:
|
||
data["important_keywords"] = important_keywords
|
||
if questions:
|
||
data["questions"] = questions
|
||
|
||
return await self._request('POST', endpoint, json=data)
|
||
|
||
async def list_chunks(self, dataset_id: str, document_id: str, keywords: Optional[str] = None,
|
||
page: int = 1, page_size: int = 1024,
|
||
chunk_id: Optional[str] = None) -> Dict[str, Any]:
|
||
"""
|
||
列出分块
|
||
"""
|
||
endpoint = f"/api/v1/datasets/{dataset_id}/documents/{document_id}/chunks"
|
||
params = {"page": page, "page_size": page_size}
|
||
|
||
if keywords:
|
||
params["keywords"] = keywords
|
||
if chunk_id:
|
||
params["id"] = chunk_id
|
||
|
||
return await self._request('GET', endpoint, params=params)
|
||
|
||
async def delete_chunks(self, dataset_id: str, document_id: str,
|
||
chunk_ids: Optional[List[str]] = None) -> Dict[str, Any]:
|
||
"""
|
||
删除分块
|
||
"""
|
||
endpoint = f"/api/v1/datasets/{dataset_id}/documents/{document_id}/chunks"
|
||
data = {"chunk_ids": chunk_ids} if chunk_ids else {}
|
||
return await self._request('DELETE', endpoint, json=data)
|
||
|
||
async def update_chunk(self, dataset_id: str, document_id: str, chunk_id: str,
|
||
content: Optional[str] = None,
|
||
important_keywords: Optional[List[str]] = None,
|
||
available: Optional[bool] = None) -> Dict[str, Any]:
|
||
"""
|
||
更新分块
|
||
"""
|
||
endpoint = f"/api/v1/datasets/{dataset_id}/documents/{document_id}/chunks/{chunk_id}"
|
||
data = {}
|
||
|
||
if content is not None:
|
||
data["content"] = content
|
||
if important_keywords is not None:
|
||
data["important_keywords"] = important_keywords
|
||
if available is not None:
|
||
data["available"] = available
|
||
|
||
return await self._request('PUT', endpoint, json=data)
|
||
|
||
async def retrieve_chunks(self, question: str, dataset_ids: Optional[List[str]] = None,
|
||
document_ids: Optional[List[str]] = None, page: int = 1,
|
||
page_size: int = 30, similarity_threshold: float = 0.2,
|
||
vector_similarity_weight: float = 0.3, top_k: int = 1024,
|
||
rerank_id: Optional[str] = None, keyword: bool = False,
|
||
highlight: bool = False) -> Dict[str, Any]:
|
||
"""
|
||
检索分块
|
||
"""
|
||
endpoint = "/api/v1/retrieval"
|
||
data = {
|
||
"question": question,
|
||
"page": page,
|
||
"page_size": page_size,
|
||
"similarity_threshold": similarity_threshold,
|
||
"vector_similarity_weight": vector_similarity_weight,
|
||
"top_k": top_k,
|
||
"keyword": keyword,
|
||
"highlight": highlight
|
||
}
|
||
|
||
if dataset_ids:
|
||
data["dataset_ids"] = dataset_ids
|
||
if document_ids:
|
||
data["document_ids"] = document_ids
|
||
if rerank_id:
|
||
data["rerank_id"] = rerank_id
|
||
|
||
return await self._request('POST', endpoint, json=data)
|
||
|
||
# ====================
|
||
# 聊天助手管理
|
||
# ====================
|
||
|
||
async def create_chat_assistant(self, name: str, avatar: Optional[str] = None,
|
||
dataset_ids: Optional[List[str]] = None,
|
||
llm: Optional[Dict[str, Any]] = None,
|
||
prompt: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
"""
|
||
创建聊天助手
|
||
"""
|
||
endpoint = "/api/v1/chats"
|
||
data = {"name": name}
|
||
|
||
if avatar:
|
||
data["avatar"] = avatar
|
||
if dataset_ids:
|
||
data["dataset_ids"] = dataset_ids
|
||
if llm:
|
||
data["llm"] = llm
|
||
if prompt:
|
||
data["prompt"] = prompt
|
||
|
||
return await self._request('POST', endpoint, json=data)
|
||
|
||
async def update_chat_assistant(self, chat_id: str, name: Optional[str] = None,
|
||
avatar: Optional[str] = None,
|
||
dataset_ids: Optional[List[str]] = None,
|
||
llm: Optional[Dict[str, Any]] = None,
|
||
prompt: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
"""
|
||
更新聊天助手
|
||
"""
|
||
endpoint = f"/api/v1/chats/{chat_id}"
|
||
data = {}
|
||
|
||
if name is not None:
|
||
data["name"] = name
|
||
if avatar is not None:
|
||
data["avatar"] = avatar
|
||
if dataset_ids is not None:
|
||
data["dataset_ids"] = dataset_ids
|
||
if llm is not None:
|
||
data["llm"] = llm
|
||
if prompt is not None:
|
||
data["prompt"] = prompt
|
||
|
||
return await self._request('PUT', endpoint, json=data)
|
||
|
||
async def delete_chat_assistants(self, ids: Optional[List[str]] = None) -> Dict[str, Any]:
|
||
"""
|
||
删除聊天助手
|
||
"""
|
||
endpoint = "/api/v1/chats"
|
||
data = {"ids": ids} if ids else {}
|
||
return await self._request('DELETE', endpoint, json=data)
|
||
|
||
async def list_chat_assistants(self, page: int = 1, page_size: int = 30,
|
||
orderby: str = "create_time", desc: str = "true",
|
||
name: Optional[str] = None,
|
||
chat_id: Optional[str] = None) -> Dict[str, Any]:
|
||
"""
|
||
列出聊天助手
|
||
"""
|
||
endpoint = "/api/v1/chats"
|
||
params = {
|
||
"page": page,
|
||
"page_size": page_size,
|
||
"orderby": orderby,
|
||
"desc": desc
|
||
}
|
||
|
||
if name:
|
||
params["name"] = name
|
||
if chat_id:
|
||
params["id"] = chat_id
|
||
|
||
return await self._request('GET', endpoint, params=params)
|
||
|
||
# ====================
|
||
# 会话管理
|
||
# ====================
|
||
|
||
async def create_session_with_chat(self, chat_id: str, name: str,
|
||
user_id: Optional[str] = None) -> Dict[str, Any]:
|
||
"""
|
||
创建与聊天助手的会话
|
||
"""
|
||
endpoint = f"/api/v1/chats/{chat_id}/sessions"
|
||
data = {"name": name}
|
||
|
||
if user_id:
|
||
data["user_id"] = user_id
|
||
|
||
return await self._request('POST', endpoint, json=data)
|
||
|
||
async def update_chat_session(self, chat_id: str, session_id: str, name: Optional[str] = None,
|
||
user_id: Optional[str] = None) -> Dict[str, Any]:
|
||
"""
|
||
更新聊天会话
|
||
"""
|
||
endpoint = f"/api/v1/chats/{chat_id}/sessions/{session_id}"
|
||
data = {}
|
||
|
||
if name is not None:
|
||
data["name"] = name
|
||
if user_id is not None:
|
||
data["user_id"] = user_id
|
||
|
||
return await self._request('PUT', endpoint, json=data)
|
||
|
||
async def list_chat_sessions(self, chat_id: str, page: int = 1, page_size: int = 30,
|
||
orderby: str = "create_time", desc: str = "true",
|
||
name: Optional[str] = None, session_id: Optional[str] = None,
|
||
user_id: Optional[str] = None) -> Dict[str, Any]:
|
||
"""
|
||
列出与指定聊天助手相关的聊天会话
|
||
"""
|
||
endpoint = f"/api/v1/chats/{chat_id}/sessions"
|
||
params = {
|
||
"page": page,
|
||
"page_size": page_size,
|
||
"orderby": orderby,
|
||
"desc": desc
|
||
}
|
||
|
||
if name:
|
||
params["name"] = name
|
||
if session_id:
|
||
params["id"] = session_id
|
||
if user_id:
|
||
params["user_id"] = user_id
|
||
|
||
return await self._request('GET', endpoint, params=params)
|
||
|
||
async def delete_chat_sessions(self, chat_id: str, ids: Optional[List[str]] = None) -> Dict[str, Any]:
|
||
"""
|
||
删除聊天会话
|
||
"""
|
||
endpoint = f"/api/v1/chats/{chat_id}/sessions"
|
||
data = {"ids": ids} if ids else {}
|
||
return await self._request('DELETE', endpoint, json=data)
|
||
|
||
async def converse_with_chat_assistant(self, chat_id: str, question: str, stream: bool = True,
|
||
session_id: Optional[str] = None,
|
||
user_id: Optional[str] = None) -> Union[Dict[str, Any], AsyncGenerator[Dict[str, Any], None]]:
|
||
"""
|
||
与聊天助手对话
|
||
"""
|
||
endpoint = f"/api/v1/chats/{chat_id}/completions"
|
||
data = {"question": question, "stream": stream}
|
||
|
||
if session_id:
|
||
data["session_id"] = session_id
|
||
|
||
print(f"开始对话: {question} {session_id}")
|
||
if user_id:
|
||
data["user_id"] = user_id
|
||
|
||
if stream:
|
||
return self._stream_request('POST', endpoint, json=data)
|
||
else:
|
||
return await self._request('POST', endpoint, json=data)
|
||
|
||
# ====================
|
||
# 代理管理
|
||
# ====================
|
||
|
||
async def create_session_with_agent(self, agent_id: str, user_id: Optional[str] = None,
|
||
file_data: Optional[Dict[str, Any]] = None,
|
||
**kwargs) -> Dict[str, Any]:
|
||
"""
|
||
创建与代理的会话
|
||
"""
|
||
if not self._session:
|
||
await self.create_session()
|
||
|
||
endpoint = f"/api/v1/agents/{agent_id}/sessions"
|
||
url = f"{self.base_url}{endpoint}"
|
||
params = {}
|
||
|
||
if user_id:
|
||
params["user_id"] = user_id
|
||
|
||
if file_data:
|
||
# 处理文件上传
|
||
data = aiohttp.FormData()
|
||
|
||
for key, file_path in file_data.items():
|
||
if os.path.exists(file_path):
|
||
file_name = Path(file_path).name
|
||
with open(file_path, 'rb') as f:
|
||
data.add_field(key, f.read(), filename=file_name)
|
||
|
||
headers = {'Authorization': f'Bearer {self.api_key}'}
|
||
|
||
async with self._session.post(url, headers=headers, data=data, params=params) as response:
|
||
try:
|
||
result = await response.json()
|
||
if result.get('code', 0) != 0:
|
||
raise RAGFlowError(result.get('code'), result.get('message'))
|
||
return result
|
||
except (aiohttp.ContentTypeError, json.JSONDecodeError):
|
||
text = await response.text()
|
||
raise RAGFlowError(response.status, text)
|
||
else:
|
||
# 普通JSON请求
|
||
data = kwargs
|
||
return await self._request('POST', endpoint, json=data, params=params)
|
||
|
||
async def converse_with_agent(self, agent_id: str, question: str, stream: bool = True,
|
||
session_id: Optional[str] = None, user_id: Optional[str] = None,
|
||
sync_dsl: bool = False, **kwargs) -> Union[Dict[str, Any], AsyncGenerator[Dict[str, Any], None]]:
|
||
"""
|
||
与代理对话
|
||
"""
|
||
endpoint = f"/api/v1/agents/{agent_id}/completions"
|
||
data = {"question": question, "stream": stream, "sync_dsl": sync_dsl}
|
||
|
||
if session_id:
|
||
data["session_id"] = session_id
|
||
if user_id:
|
||
data["user_id"] = user_id
|
||
|
||
# 添加其他Begin组件参数
|
||
data.update(kwargs)
|
||
|
||
if stream:
|
||
return self._stream_request('POST', endpoint, json=data)
|
||
else:
|
||
return await self._request('POST', endpoint, json=data)
|
||
|
||
async def list_agent_sessions(self, agent_id: str, page: int = 1, page_size: int = 30,
|
||
orderby: str = "create_time", desc: str = "true",
|
||
session_id: Optional[str] = None, user_id: Optional[str] = None,
|
||
dsl: bool = True) -> Dict[str, Any]:
|
||
"""
|
||
列出代理会话
|
||
"""
|
||
endpoint = f"/api/v1/agents/{agent_id}/sessions"
|
||
params = {
|
||
"page": page,
|
||
"page_size": page_size,
|
||
"orderby": orderby,
|
||
"desc": desc,
|
||
"dsl": dsl
|
||
}
|
||
|
||
if session_id:
|
||
params["id"] = session_id
|
||
if user_id:
|
||
params["user_id"] = user_id
|
||
|
||
return await self._request('GET', endpoint, params=params)
|
||
|
||
async def delete_agent_sessions(self, agent_id: str, ids: Optional[List[str]] = None) -> Dict[str, Any]:
|
||
"""
|
||
删除代理会话
|
||
"""
|
||
endpoint = f"/api/v1/agents/{agent_id}/sessions"
|
||
data = {"ids": ids} if ids else {}
|
||
return await self._request('DELETE', endpoint, json=data)
|
||
|
||
async def get_related_questions(self, question: str, login_token: str) -> Dict[str, Any]:
|
||
"""
|
||
生成相关问题
|
||
"""
|
||
endpoint = "/v1/sessions/related_questions"
|
||
headers = {
|
||
'Authorization': f'Bearer {login_token}',
|
||
'Content-Type': 'application/json'
|
||
}
|
||
data = {"question": question}
|
||
|
||
return await self._request('POST', endpoint, headers=headers, json=data)
|
||
|
||
async def list_agents(self, page: int = 1, page_size: int = 30, orderby: str = "create_time",
|
||
desc: str = "true", name: Optional[str] = None,
|
||
agent_id: Optional[str] = None) -> Dict[str, Any]:
|
||
"""
|
||
列出代理
|
||
"""
|
||
endpoint = "/api/v1/agents"
|
||
params = {
|
||
"page": page,
|
||
"page_size": page_size,
|
||
"orderby": orderby,
|
||
"desc": desc
|
||
}
|
||
|
||
if name:
|
||
params["name"] = name
|
||
if agent_id:
|
||
params["id"] = agent_id
|
||
|
||
return await self._request('GET', endpoint, params=params)
|
||
|
||
async def create_agent(self, title: str, description: Optional[str] = None,
|
||
dsl: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
"""
|
||
创建代理
|
||
"""
|
||
endpoint = "/api/v1/agents"
|
||
data = {"title": title}
|
||
|
||
if description is not None:
|
||
data["description"] = description
|
||
if dsl is not None:
|
||
data["dsl"] = dsl
|
||
|
||
return await self._request('POST', endpoint, json=data)
|
||
|
||
async def update_agent(self, agent_id: str, title: Optional[str] = None,
|
||
description: Optional[str] = None,
|
||
dsl: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||
"""
|
||
更新代理
|
||
|
||
Args:
|
||
agent_id: 代理ID
|
||
title: 新标题
|
||
description: 新描述
|
||
dsl: 新DSL配置
|
||
"""
|
||
endpoint = f"/api/v1/agents/{agent_id}"
|
||
data = {}
|
||
|
||
if title is not None:
|
||
data["title"] = title
|
||
if description is not None:
|
||
data["description"] = description
|
||
if dsl is not None:
|
||
data["dsl"] = dsl
|
||
|
||
return await self._request('PUT', endpoint, json=data)
|
||
|
||
async def delete_agent(self, agent_id: str) -> Dict[str, Any]:
|
||
"""
|
||
删除代理
|
||
|
||
Args:
|
||
agent_id: 代理ID
|
||
"""
|
||
endpoint = f"/api/v1/agents/{agent_id}"
|
||
return await self._request('DELETE', endpoint)
|
||
|
||
|
||
# ====================
|
||
# 使用示例
|
||
# ====================
|
||
|
||
async def example_usage():
|
||
"""
|
||
异步RAGFlow SDK使用示例
|
||
"""
|
||
# 使用异步上下文管理器
|
||
async with AsyncRAGFlowClient(
|
||
base_url="http://10.0.0.202:82",
|
||
api_key="ragflow-hlMjRmNzE2ODNiNTExZjA4ZTNlMDI0Mm"
|
||
) as client:
|
||
|
||
try:
|
||
# 删除数据集
|
||
await client.delete_datasets(ids=["afe1387883bb11f0a0fd0242ac170006"])
|
||
print("删除数据集成功")
|
||
|
||
# 1. 创建数据集
|
||
dataset = await client.create_dataset(
|
||
name="我的数据集",
|
||
description="这是一个测试数据集",
|
||
chunk_method="naive"
|
||
)
|
||
dataset_id = dataset['data']['id']
|
||
print(f"创建数据集成功: {dataset_id}")
|
||
|
||
# 2. 上传文档
|
||
documents = await client.upload_documents(
|
||
dataset_id=dataset_id,
|
||
file_paths=[
|
||
"/home/admin-root/haotian/康达瑞贝斯机器人后台/ruoyi-fastapi-backend/requirements.txt",
|
||
"/home/admin-root/haotian/康达瑞贝斯机器人后台/ruoyi-fastapi-backend/requirements-pg.txt"
|
||
]
|
||
)
|
||
print("文档上传成功")
|
||
|
||
# 3. 解析文档
|
||
document_ids = [doc['id'] for doc in documents['data']]
|
||
await client.parse_documents(dataset_id, document_ids)
|
||
print("开始解析文档")
|
||
|
||
# 等待解析完成
|
||
await asyncio.sleep(5)
|
||
|
||
# 4. 创建聊天助手
|
||
chat_assistant = await client.create_chat_assistant(
|
||
name="我的AI助手",
|
||
dataset_ids=[dataset_id]
|
||
)
|
||
chat_id = chat_assistant['data']['id']
|
||
print(f"创建聊天助手成功: {chat_id}")
|
||
|
||
# 5. 创建会话
|
||
session = await client.create_session_with_chat(
|
||
chat_id=chat_id,
|
||
name="测试会话"
|
||
)
|
||
session_id = session['data']['id']
|
||
print(f"创建会话成功: {session_id}")
|
||
|
||
# 6. 开始对话（流式）
|
||
responses = client.converse_with_chat_assistant(
|
||
chat_id=chat_id,
|
||
question="你好，请介绍一下自己",
|
||
stream=True,
|
||
session_id=session_id
|
||
)
|
||
|
||
print("AI回复:")
|
||
async for response in responses:
|
||
if response.get('data') and isinstance(response['data'], dict):
|
||
answer = response['data'].get('answer', '')
|
||
if answer:
|
||
print(answer, end='', flush=True)
|
||
print()
|
||
|
||
# 7. 检索相关文档块
|
||
chunks = await client.retrieve_chunks(
|
||
question="RAGFlow的优势是什么？",
|
||
dataset_ids=[dataset_id],
|
||
top_k=5,
|
||
highlight=True
|
||
)
|
||
print(f"检索到 {chunks['data']['total']} 个相关文档块")
|
||
|
||
# 8. 列出数据集
|
||
datasets = await client.list_datasets(page=1, page_size=10)
|
||
print(f"当前有 {len(datasets['data'])} 个数据集")
|
||
|
||
except RAGFlowError as e:
|
||
print(f"RAGFlow API错误: {e}")
|
||
except Exception as e:
|
||
print(f"其他错误: {e}")
|
||
|
||
|
||
async def example_usage_1():
|
||
"""测试获取列表方法的异步版本"""
|
||
|
||
# 使用异步上下文管理器
|
||
async with AsyncRAGFlowClient(
|
||
base_url="http://10.0.0.202:82",
|
||
api_key="ragflow-hlMjRmNzE2ODNiNTExZjA4ZTNlMDI0Mm"
|
||
) as client:
|
||
|
||
try:
|
||
# 1. 获取数据集列表
|
||
results_dataset = await client.list_datasets()
|
||
print(f"获取数据集列表成功，共有 {len(results_dataset['data'])} 个数据集")
|
||
print("数据集ID:\n", [result["id"] for result in results_dataset['data']])
|
||
|
||
# 2. 获取数据集中文档列表
|
||
for result in results_dataset['data']:
|
||
print(f"数据集 {result['id']} 的文档列表为：")
|
||
results_doc = await client.list_documents(dataset_id=result['id'])
|
||
# 文档名称
|
||
print([t["name"] for t in results_doc["data"]["docs"]])
|
||
|
||
except RAGFlowError as e:
|
||
print(f"RAGFlow API错误: {e}")
|
||
except Exception as e:
|
||
print(f"其他错误: {e}")
|
||
|
||
|
||
async def batch_operations_example():
|
||
"""
|
||
批量操作示例 - 演示异步并发处理
|
||
"""
|
||
async with AsyncRAGFlowClient(
|
||
base_url="http://10.0.0.202:82",
|
||
api_key="ragflow-hlMjRmNzE2ODNiNTExZjA4ZTNlMDI0Mm"
|
||
) as client:
|
||
|
||
try:
|
||
# 并发获取数据集和代理列表
|
||
datasets_task = client.list_datasets()
|
||
agents_task = client.list_agents()
|
||
chats_task = client.list_chat_assistants()
|
||
|
||
# 等待所有任务完成
|
||
datasets, agents, chats = await asyncio.gather(
|
||
datasets_task, agents_task, chats_task
|
||
)
|
||
|
||
print(f"数据集数量: {len(datasets['data'])}")
|
||
print(f"代理数量: {len(agents['data'])}")
|
||
print(f"聊天助手数量: {len(chats['data'])}")
|
||
|
||
# 如果有数据集，并发获取每个数据集的文档
|
||
if datasets['data']:
|
||
doc_tasks = [
|
||
client.list_documents(dataset_id=dataset['id'])
|
||
for dataset in datasets['data'][:3] # 限制前3个
|
||
]
|
||
|
||
doc_results = await asyncio.gather(*doc_tasks, return_exceptions=True)
|
||
|
||
for i, result in enumerate(doc_results):
|
||
if isinstance(result, Exception):
|
||
print(f"获取数据集 {datasets['data'][i]['id']} 的文档时出错: {result}")
|
||
else:
|
||
print(f"数据集 {datasets['data'][i]['id']} 有 {len(result['data']['docs'])} 个文档")
|
||
|
||
except Exception as e:
|
||
print(f"批量操作出错: {e}")
|
||
|
||
|
||
# 运行示例的辅助函数
|
||
def run_example():
|
||
"""运行异步示例"""
|
||
# 可以选择运行不同的示例
|
||
asyncio.run(example_usage_1())
|
||
# asyncio.run(example_usage())
|
||
# asyncio.run(batch_operations_example())
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 运行示例
|
||
run_example() |