257 lines
7.7 KiB
Python
257 lines
7.7 KiB
Python
|
from elasticsearch import Elasticsearch
|
|||
|
|
#from src.add_chunk_cli_pdf_img import update_positon_img_id_in_elasticsearch
|
|||
|
|
# 初始化 Elasticsearch 用户名elastic，密码infini_rag_flow
|
|||
|
|
|
|||
|
|
from dotenv import load_dotenv # 新增
|
|||
|
|
import os
|
|||
|
|
import json
|
|||
|
|
# 加载 .env 文件中的环境变量
|
|||
|
|
load_dotenv()
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 初始化 Elasticsearch
|
|||
|
|
es = Elasticsearch(
|
|||
|
|
[{
|
|||
|
|
'host': os.getenv("ELASTIC_HOST"),
|
|||
|
|
'port': int(os.getenv("ELASTIC_PORT")),
|
|||
|
|
'scheme': 'http'
|
|||
|
|
}],
|
|||
|
|
basic_auth=(
|
|||
|
|
os.getenv("ELASTIC_USERNAME"),
|
|||
|
|
os.getenv("ELASTIC_PASSWORD")
|
|||
|
|
)
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_index_mapping(tenant_id):
|
|||
|
|
"""
|
|||
|
|
获取指定索引的 mapping 信息
|
|||
|
|
|
|||
|
|
:param tenant_id: 租户 ID
|
|||
|
|
:return: mapping 信息
|
|||
|
|
"""
|
|||
|
|
index_name = f"ragflow_{tenant_id}"
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
mapping = es.indices.get_mapping(index=index_name)
|
|||
|
|
# 将 ObjectApiResponse 转换为普通字典
|
|||
|
|
mapping_dict = dict(mapping)
|
|||
|
|
return {"code": 0, "message": "", "data": mapping_dict}
|
|||
|
|
except Exception as e:
|
|||
|
|
return {"code": 500, "message": str(e), "data": {}}
|
|||
|
|
|
|||
|
|
def update_positon_in_elasticsearch(tenant_id, doc_id, chunk_id, positions):
|
|||
|
|
"""
|
|||
|
|
在 Elasticsearch 中更新指定文档块的position and img_id。
|
|||
|
|
|
|||
|
|
:param tenant_id: 租户 ID
|
|||
|
|
:param doc_id: 文档 ID
|
|||
|
|
:param chunk_id: 文档块 ID
|
|||
|
|
:param new_img_id: 新的 img_id
|
|||
|
|
:param position: 位置信息
|
|||
|
|
:return: 更新结果
|
|||
|
|
"""
|
|||
|
|
if not positions:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
position_int = []
|
|||
|
|
|
|||
|
|
for pos in positions:
|
|||
|
|
if len(pos) != 5:
|
|||
|
|
continue # Skip invalid positions
|
|||
|
|
|
|||
|
|
pn, left, right, top, bottom = pos
|
|||
|
|
# 使用元组格式，与原始RAGFlow保持一致
|
|||
|
|
position_int.append((int(pn + 1), int(left), int(right), int(top), int(bottom)))
|
|||
|
|
|
|||
|
|
if position_int: # Only add if we have valid positions
|
|||
|
|
# 仅添加精确位置信息，不修改排序字段
|
|||
|
|
|
|||
|
|
# 构建索引名称
|
|||
|
|
index_name = f"ragflow_{tenant_id}"
|
|||
|
|
|
|||
|
|
# 构建查询条件
|
|||
|
|
query = {
|
|||
|
|
"bool": {
|
|||
|
|
"must": [
|
|||
|
|
{"term": {"doc_id": doc_id}},
|
|||
|
|
{"term": {"_id": chunk_id}}
|
|||
|
|
]
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 搜索目标文档
|
|||
|
|
result = es.search(index=index_name, body={"query": query})
|
|||
|
|
|
|||
|
|
# 检查是否找到目标文档
|
|||
|
|
if result['hits']['total']['value'] == 0:
|
|||
|
|
print(f"在 Elasticsearch 中未找到文档: index={index_name}, doc_id={doc_id}, chunk_id={chunk_id}")
|
|||
|
|
return {"code": 102, "message": f"Can't find this chunk {chunk_id}"}
|
|||
|
|
|
|||
|
|
# 获取目标文档的 ID
|
|||
|
|
hit = result['hits']['hits'][0]
|
|||
|
|
doc_id_in_es = hit['_id']
|
|||
|
|
|
|||
|
|
# 构建更新请求 - 只更新存在的字段
|
|||
|
|
update_body = {"doc": {}}
|
|||
|
|
update_body["doc"]["position_int"] = position_int
|
|||
|
|
update_body["doc"]["page_num_int"] = [position_int[0][0]]
|
|||
|
|
update_body["doc"]["top_int"] = [position_int[0][3]]
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 更新文档
|
|||
|
|
update_result = es.update(
|
|||
|
|
index=index_name,
|
|||
|
|
id=doc_id_in_es,
|
|||
|
|
body=update_body,
|
|||
|
|
refresh=True # 确保更新立即可见
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
print(f"Elasticsearch 更新结果: index={index_name}, id={doc_id_in_es}, result={update_result}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
def update_positon_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id, position, new_img_id):
|
|||
|
|
"""
|
|||
|
|
在 Elasticsearch 中更新指定文档块的position and img_id。
|
|||
|
|
|
|||
|
|
:param tenant_id: 租户 ID
|
|||
|
|
:param doc_id: 文档 ID
|
|||
|
|
:param chunk_id: 文档块 ID
|
|||
|
|
:param new_img_id: 新的 img_id
|
|||
|
|
:param position: 位置信息
|
|||
|
|
:return: 更新结果
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
|
|||
|
|
# 构建索引名称
|
|||
|
|
index_name = f"ragflow_{tenant_id}"
|
|||
|
|
|
|||
|
|
# 构建查询条件
|
|||
|
|
query = {
|
|||
|
|
"bool": {
|
|||
|
|
"must": [
|
|||
|
|
{"term": {"doc_id": doc_id}},
|
|||
|
|
{"term": {"_id": chunk_id}}
|
|||
|
|
]
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# 搜索目标文档
|
|||
|
|
result = es.search(index=index_name, body={"query": query})
|
|||
|
|
|
|||
|
|
# 检查是否找到目标文档
|
|||
|
|
if result['hits']['total']['value'] == 0:
|
|||
|
|
print(f"在 Elasticsearch 中未找到文档: index={index_name}, doc_id={doc_id}, chunk_id={chunk_id}")
|
|||
|
|
return {"code": 102, "message": f"Can't find this chunk {chunk_id}"}
|
|||
|
|
|
|||
|
|
# 获取目标文档的 ID
|
|||
|
|
hit = result['hits']['hits'][0]
|
|||
|
|
doc_id_in_es = hit['_id']
|
|||
|
|
|
|||
|
|
# 构建更新请求 - 只更新存在的字段
|
|||
|
|
update_body = {"doc": {}}
|
|||
|
|
|
|||
|
|
#只有当 new_img_id 存在时才更新 img_id
|
|||
|
|
if new_img_id is not None:
|
|||
|
|
update_body["doc"]["img_id"] = new_img_id
|
|||
|
|
|
|||
|
|
# 只有当 position 存在时才更新 positions
|
|||
|
|
if position is not None:
|
|||
|
|
|
|||
|
|
update_body["doc"]["positions"] = position
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 如果没有需要更新的字段，直接返回成功
|
|||
|
|
if not update_body["doc"]:
|
|||
|
|
print("没有需要更新的字段")
|
|||
|
|
return {"code": 0, "message": "No fields to update"}
|
|||
|
|
|
|||
|
|
# 更新文档
|
|||
|
|
update_result = es.update(
|
|||
|
|
index=index_name,
|
|||
|
|
id=doc_id_in_es,
|
|||
|
|
body=update_body,
|
|||
|
|
refresh=True # 确保更新立即可见
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
print(f"Elasticsearch 更新结果: index={index_name}, id={doc_id_in_es}, result={update_result}")
|
|||
|
|
|
|||
|
|
# 验证更新
|
|||
|
|
verify_doc = es.get(index=index_name, id=doc_id_in_es)
|
|||
|
|
|
|||
|
|
# 检查 img_id 是否已更新（如果提供了 new_img_id）
|
|||
|
|
img_id_updated = True
|
|||
|
|
if new_img_id is not None:
|
|||
|
|
img_id_updated = verify_doc['_source'].get('img_id') == new_img_id
|
|||
|
|
if img_id_updated:
|
|||
|
|
print(f"成功更新 img_id 为: {new_img_id}")
|
|||
|
|
else:
|
|||
|
|
print(f"更新验证失败，当前 img_id: {verify_doc['_source'].get('img_id')}")
|
|||
|
|
|
|||
|
|
# 检查 position 是否已更新（如果提供了 position）
|
|||
|
|
position_updated = True
|
|||
|
|
if position is not None:
|
|||
|
|
position_updated = verify_doc['_source'].get('positions') == position
|
|||
|
|
if position_updated:
|
|||
|
|
print(f"成功更新 position 为: {position}")
|
|||
|
|
else:
|
|||
|
|
print(f"更新验证失败，当前 position: {verify_doc['_source'].get('positions')}")
|
|||
|
|
|
|||
|
|
# 统一返回结果
|
|||
|
|
if img_id_updated and position_updated:
|
|||
|
|
return {"code": 0, "message": ""}
|
|||
|
|
else:
|
|||
|
|
return {"code": 100, "message": "Failed to verify update"}
|
|||
|
|
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"更新 Elasticsearch 时发生错误: {str(e)}")
|
|||
|
|
return {"code": 101, "message": f"Error updating img_id: {str(e)}"}
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 示例调用 - 列出特定文档的所有 chunks
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
try:
|
|||
|
|
print(es.info())
|
|||
|
|
except Exception as e:
|
|||
|
|
print("连接失败：", e)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 单位电脑
|
|||
|
|
tenant_id = "d669205e57a211f0b9e7324e7f243034"
|
|||
|
|
new_img_id ="10345832587311f0919f3a2728512a4b-bd04866cd05337281"
|
|||
|
|
doc_id="ea8d75966df811f0925ac6e8db75f472"
|
|||
|
|
chunk_id="4a4927560a7e6d80"
|
|||
|
|
# 添加以下代码来检查索引映射
|
|||
|
|
# mapping_result = get_index_mapping(tenant_id)
|
|||
|
|
# print("Positions field mapping:", mapping_result["data"][f"ragflow_{tenant_id}"]["mappings"]["properties"]["positions"])
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 左，右 -->
|
|||
|
|
#上， 下| 上面最小，下面最大
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
pos = [[4, 0, 100, 200, 510]]
|
|||
|
|
#pos_string = json.dumps(pos) # 转换为 JSON 字符串
|
|||
|
|
update_positon_in_elasticsearch(tenant_id, doc_id, chunk_id, pos)
|
|||
|
|
|
|||
|
|
|
|||
|
|
#update_positon_img_id_in_elasticsearch(tenant_id, doc_id, chunk_id, pos, "")
|