ouyangyouzhang/QueryRewrite
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

更新riper-5.mdc配置以启用自动应用，优化DifyApi类，新增文档子分段管理功能，包括获取、添加、更新和删除子分段的方法，同时改进OpenAiLLM类的初始化逻辑以支持环境变量配置。

Browse Source
This commit is contained in:
ouyangyouzhang
2025-07-03 11:34:31 +08:00
parent b8ccdcf287
commit 68e3677c34
3 changed files with 190 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rag2_0/dify/dify_client/dify_api.py
+172 -6
View File
@@ -112,19 +112,23 @@ class DifyApi:
return response.json().get("document", {}).get("id", "")
def get_or_create_dataset_by_name(self, dataset_name: str) -> str:
def get_or_create_dataset_by_name(self, dataset_name: str, create_if_not_exist: bool=True) -> str:
"""
通过名称获取或创建数据集。
:param dataset_name: 数据集名称。
:param create_if_not_exist: 如果数据集不存在是否创建。
:return: 数据集ID。
"""
list_dataset = self.get_all_dataset_list()
for dataset in list_dataset:
if dataset["name"] == dataset_name:
return dataset["id"]
logging.info(f"数据集不存在，创建数据集: {dataset_name}")
return self.create_dataset(dataset_name)
if create_if_not_exist:
logging.info(f"数据集不存在，创建数据集: {dataset_name}")
return self.create_dataset(dataset_name)
else:
raise Exception(f"数据集不存在: {dataset_name}")
def get_documents(self, dataset_id: str, keyword: str = None) -> Dict[str, dict]:
"""
@@ -364,10 +368,28 @@ class DifyApi:
'Content-Type': 'application/json'
}
limit = 100
page = 0
all_segments = []
try:
response = requests.get(url, headers=headers, verify=False)
response.raise_for_status() # 如果响应状态码不是200，抛出异常
return response.json().get("data", []) # 返回分段信息列表
while True:
page += 1
params = {
'page': page,
'limit': limit
}
response = requests.get(url, headers=headers, params=params, verify=False)
response.raise_for_status() # 如果响应状态码不是200，抛出异常
data = response.json()
all_segments.extend(data.get("data", []))
if not data.get("has_more", False):
break
return all_segments
except Exception as e:
logging.error(f"获取文档分段失败: {e}")
raise
@@ -464,6 +486,150 @@ class DifyApi:
logging.warning(f"第{attempt + 1}次尝试失败: {e}")
time.sleep(1) # 重试前等待1秒
def add_document_child_chunk(
self,
dataset_id: str,
document_id: str,
segment_id: str,
content: str
) -> Dict:
"""
新增文档子分段。
:param dataset_id: 数据集ID。
:param document_id: 文档ID。
:param segment_id: 分段ID。
:param content: 子分段内容。
:return: 新增的子分段信息。
:raises: Exception 如果请求失败。
"""
url = f"{self.dify_url}/datasets/{dataset_id}/documents/{document_id}/segments/{segment_id}/child_chunks"
headers = {
'Authorization': f'Bearer {self.dify_dataset_api_key}',
'Content-Type': 'application/json'
}
# 构造请求数据
data = {
"content": content
}
try:
response = requests.post(url, headers=headers, data=json.dumps(data), verify=False)
response.raise_for_status() # 如果响应状态码不是200，抛出异常
return response.json().get("data", {}) # 返回新增的子分段信息
except Exception as e:
logging.error(f"新增文档子分段失败: {e}")
raise
def get_document_child_chunks(
self,
dataset_id: str,
document_id: str,
segment_id: str,
page: int = 1,
limit: int = 100
) -> Dict:
"""
获取文档子分段列表。
:param dataset_id: 数据集ID。
:param document_id: 文档ID。
:param segment_id: 分段ID。
:param page: 页码，默认为1。
:param limit: 每页数量，默认为20。
:return: 子分段列表信息，包含分页信息。
:raises: Exception 如果请求失败。
"""
url = f"{self.dify_url}/datasets/{dataset_id}/documents/{document_id}/segments/{segment_id}/child_chunks"
headers = {
'Authorization': f'Bearer {self.dify_dataset_api_key}'
}
params = {
'page': page,
'limit': limit
}
try:
response = requests.get(url, headers=headers, params=params, verify=False)
response.raise_for_status() # 如果响应状态码不是200，抛出异常
return response.json() # 返回子分段列表信息，包含分页信息
except Exception as e:
logging.error(f"获取文档子分段列表失败: {e}")
raise
def del_document_child_chunk(
self,
dataset_id: str,
document_id: str,
segment_id: str,
child_chunk_id: str
) -> bool:
"""
删除文档子分段。
:param dataset_id: 数据集ID。
:param document_id: 文档ID。
:param segment_id: 分段ID。
:param child_chunk_id: 子分段ID。
:return: 如果删除成功返回True，否则返回False。
"""
url = f"{self.dify_url}/datasets/{dataset_id}/documents/{document_id}/segments/{segment_id}/child_chunks/{child_chunk_id}"
headers = {
'Authorization': f'Bearer {self.dify_dataset_api_key}'
}
try:
response = requests.delete(url, headers=headers, verify=False)
if response.status_code == 200:
logging.info(f"删除子分段成功: {child_chunk_id}")
return True
else:
logging.error(f"删除子分段失败，状态码: {response.status_code}, 响应: {response.text}")
return False
except Exception as e:
logging.error(f"删除子分段失败: {e}")
return False
def update_document_child_chunk(
self,
dataset_id: str,
document_id: str,
segment_id: str,
child_chunk_id: str,
content: str
) -> Dict:
"""
更新文档子分段内容。
:param dataset_id: 数据集ID。
:param document_id: 文档ID。
:param segment_id: 分段ID。
:param child_chunk_id: 子分段ID。
:param content: 更新的子分段内容。
:return: 更新后的子分段信息。
:raises: Exception 如果请求失败。
"""
url = f"{self.dify_url}/datasets/{dataset_id}/documents/{document_id}/segments/{segment_id}/child_chunks/{child_chunk_id}"
headers = {
'Authorization': f'Bearer {self.dify_dataset_api_key}',
'Content-Type': 'application/json'
}
# 构造请求数据
data = {
"content": content
}
try:
response = requests.patch(url, headers=headers, data=json.dumps(data), verify=False)
response.raise_for_status() # 如果响应状态码不是200，抛出异常
return response.json().get("data", {}) # 返回更新后的子分段信息
except Exception as e:
logging.error(f"更新文档子分段失败: {e}")
raise
if __name__ == '__main__':
from dotenv import load_dotenv
rag2_0/tool/ModelTool.py
+14 -8
View File
@@ -127,16 +127,22 @@ class XinferenceReRankerModel:
class OpenAiLLM:
def __init__(self, **kwargs):
if kwargs.get("api_key") == None or kwargs.get("base_url") == None or kwargs.get("model") == None:
raise ValueError("api_key, base_url, model 不能为空")
if "api_key" in kwargs:
self._api_key = kwargs.get("api_key")
kwargs.pop("api_key")
self._api_key = kwargs.get("api_key")
self._url = kwargs.get("base_url")
self._model = kwargs.get("model")
if "base_url" in kwargs:
self._url = kwargs.get("base_url")
kwargs.pop("base_url")
else:
self._url = os.getenv("OPENAI_API_BASE")
if "model" in kwargs:
self._model = kwargs.get("model")
kwargs.pop("model")
else:
self._model = os.getenv("LLM_MODEL_NAME")
kwargs.pop("api_key")
kwargs.pop("base_url")
kwargs.pop("model")
self._kwargs = kwargs
def invoke(self, user_prompt="你是谁？", need_retry=True):