新增 chunk_pos.py 文件，集成 Elasticsearch 功能，支持文档块位置和图像 ID 更新，优化索引映射获取逻辑

chunk_pos.py

@@ -0,0 +1,256 @@
+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, "")