Compare commits
14 Commits
defeng-pat
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
| 79ae2afd46 | |||
| 41d01d302d | |||
| 6ef7705af3 | |||
| f049dd3502 | |||
| 92716fddaa | |||
| fd8b4b6614 | |||
| 715829d763 | |||
| f8c99dcfdc | |||
| 325fed5b2d | |||
| 129846dc17 | |||
| 626139226d | |||
| 1b3cd3bfc3 | |||
| 61ee346a43 | |||
| 47b3a27812 |
204
app.py
204
app.py
@ -40,7 +40,7 @@ from pydantic import BaseModel
|
|||||||
from dataprocess import _split_markdown
|
from dataprocess import _split_markdown
|
||||||
import uuid
|
import uuid
|
||||||
import json
|
import json
|
||||||
from fileparse_util import process_document
|
from fileparse_util import process_document, extract_image_caption
|
||||||
# import hanlp
|
# import hanlp
|
||||||
from itertools import groupby
|
from itertools import groupby
|
||||||
import requests
|
import requests
|
||||||
@ -57,7 +57,7 @@ from concurrent.futures import ThreadPoolExecutor
|
|||||||
import base64
|
import base64
|
||||||
from extract_excel_node_relation import get_excel_node_relation
|
from extract_excel_node_relation import get_excel_node_relation
|
||||||
from doc2pdf import Doc2PDF
|
from doc2pdf import Doc2PDF
|
||||||
from chunk_text import get_chunk_bbox, split_text_preserve_sentences, data_replace, chunk_check,merge_short_slices,find_and_read_content_list,process_pdf_file,process_other_file
|
from chunk_text import get_chunk_bbox, split_text_preserve_sentences, data_replace, chunk_check,merge_short_slices,find_and_read_content_list,process_pdf_file,process_other_file,safe_original_filename,image_base64_to_data_url,extract_image_text,prepare_split_image
|
||||||
from filename_proceess_and_kgquery import get_entity
|
from filename_proceess_and_kgquery import get_entity
|
||||||
from kg_build.extract_filtertext import get_filtertext_node_relation
|
from kg_build.extract_filtertext import get_filtertext_node_relation
|
||||||
import os
|
import os
|
||||||
@ -92,7 +92,7 @@ from default_ontology_config import (
|
|||||||
get_relationships,
|
get_relationships,
|
||||||
get_relationship_types,
|
get_relationship_types,
|
||||||
)
|
)
|
||||||
from config import NEO4J_CONFIG,URL,API_URLS,API_OTHER_URLS
|
from config import NEO4J_CONFIG,URL,API_URLS,API_OTHER_URLS,VLM_CONCURRENCY,SEARCH_DIR
|
||||||
converter = Doc2PDF()
|
converter = Doc2PDF()
|
||||||
DATA_DIR = "/app/files"
|
DATA_DIR = "/app/files"
|
||||||
# ================== 全局资源容器 ==================
|
# ================== 全局资源容器 ==================
|
||||||
@ -102,7 +102,6 @@ _resources: Dict[str, Any] = {
|
|||||||
"async_driver": None,
|
"async_driver": None,
|
||||||
"embedder": None,
|
"embedder": None,
|
||||||
}
|
}
|
||||||
|
|
||||||
# 图谱检索和索引生成
|
# 图谱检索和索引生成
|
||||||
try:
|
try:
|
||||||
from graph_search.graph_service import GraphService
|
from graph_search.graph_service import GraphService
|
||||||
@ -167,6 +166,7 @@ async def lifespan(app):
|
|||||||
|
|
||||||
|
|
||||||
app = FastAPI(max_request_size=1024 * 1024 * 10, lifespan=lifespan)
|
app = FastAPI(max_request_size=1024 * 1024 * 10, lifespan=lifespan)
|
||||||
|
VLM_SEMAPHORE = asyncio.Semaphore(max(1, VLM_CONCURRENCY))
|
||||||
|
|
||||||
setup_logging()
|
setup_logging()
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@ -1739,7 +1739,7 @@ async def kg_build(request: KGRequest):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"保存 kg_build 配置缓存失败: {e}")
|
logger.warning(f"保存 kg_build 配置缓存失败: {e}")
|
||||||
|
|
||||||
# 后续进程启动逻辑保持不变 ...
|
# 后续进程启动逻辑保持不变 ...https://gitea.zkzdht.com/Ascend/htknow
|
||||||
cancel_flag_path = f"/tmp/kg_cancel_{task_id}"
|
cancel_flag_path = f"/tmp/kg_cancel_{task_id}"
|
||||||
if os.path.exists(cancel_flag_path):
|
if os.path.exists(cancel_flag_path):
|
||||||
os.remove(cancel_flag_path)
|
os.remove(cancel_flag_path)
|
||||||
@ -2028,6 +2028,105 @@ class RequestWrapper(BaseModel):
|
|||||||
pdf_contents: List[PdfContentItem]
|
pdf_contents: List[PdfContentItem]
|
||||||
|
|
||||||
|
|
||||||
|
class Base64ImageRequest(BaseModel):
|
||||||
|
content: Optional[str] = Field(
|
||||||
|
None,
|
||||||
|
description="Optional text around the image, for example: 图3.1 船舶主发动机组成图(images/test.jpg)",
|
||||||
|
)
|
||||||
|
filename: str = Field(..., description="Original image filename, for example test.png")
|
||||||
|
image_base64: str = Field(..., description="Base64 image content, with or without data:image/... prefix")
|
||||||
|
original_filename:Optional[str] = Field(
|
||||||
|
None,
|
||||||
|
description="图片 所在pdf文件",
|
||||||
|
)
|
||||||
|
|
||||||
|
async def analyze_image_with_vlm(image_bytes: bytes, suffix: str) -> str:
|
||||||
|
image_url = image_base64_to_data_url(image_bytes, suffix)
|
||||||
|
async with VLM_SEMAPHORE:
|
||||||
|
return await OpenaiAPI.open_api_vl_without_thinking(image_url)
|
||||||
|
|
||||||
|
|
||||||
|
@app.post("/split_image")
|
||||||
|
async def image_base64(request_data: Base64ImageRequest):
|
||||||
|
logger.info(f"Received image: {request_data.filename}")
|
||||||
|
logger.info(f"内容: {request_data.content}")
|
||||||
|
logger.info(f"pdf文件名: {request_data.original_filename}")
|
||||||
|
filename = safe_original_filename(request_data.filename)
|
||||||
|
logger.info(f"filename: {filename}")
|
||||||
|
suffix = PathLib(filename).suffix.lower()
|
||||||
|
if suffix not in {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif", ".tif", ".tiff"}:
|
||||||
|
raise HTTPException(status_code=400, detail="Unsupported image type")
|
||||||
|
|
||||||
|
image_bytes, processed_suffix = prepare_split_image(request_data.image_base64, suffix)
|
||||||
|
if not image_bytes:
|
||||||
|
raise HTTPException(status_code=400, detail="decoded image is empty")
|
||||||
|
|
||||||
|
temp_path = None
|
||||||
|
image_caption = []
|
||||||
|
vlm_text = ""
|
||||||
|
try:
|
||||||
|
with tempfile.NamedTemporaryFile(delete=False, suffix=processed_suffix) as tmp:
|
||||||
|
tmp.write(image_bytes)
|
||||||
|
temp_path = PathLib(tmp.name)
|
||||||
|
|
||||||
|
ocr_data = await extract_image_text(temp_path)
|
||||||
|
try:
|
||||||
|
temp_path.unlink()
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Failed to delete temp image {temp_path}: {e}")
|
||||||
|
finally:
|
||||||
|
temp_path = None
|
||||||
|
if request_data.original_filename:
|
||||||
|
image_caption = await extract_image_caption(
|
||||||
|
SEARCH_DIR,
|
||||||
|
request_data.original_filename,
|
||||||
|
request_data.filename,
|
||||||
|
)
|
||||||
|
|
||||||
|
caption_text = "\n".join(image_caption).strip()
|
||||||
|
|
||||||
|
# image_caption 不为空时使用 caption,
|
||||||
|
# 否则使用原始图片文件名
|
||||||
|
if caption_text:
|
||||||
|
caption_text =caption_text+".jpg"
|
||||||
|
output_filename = caption_text or filename
|
||||||
|
logger.info(f"image_caption: {caption_text}")
|
||||||
|
content = (request_data.content or "").strip()
|
||||||
|
ocr_full_text = ((ocr_data or {}).get("full_text") or "").strip()
|
||||||
|
|
||||||
|
# content 和 OCR 都为空时调用 VLM
|
||||||
|
if not content and not ocr_full_text:
|
||||||
|
vlm_text = await analyze_image_with_vlm(
|
||||||
|
image_bytes,
|
||||||
|
processed_suffix,
|
||||||
|
)
|
||||||
|
|
||||||
|
image_related_content = ocr_full_text+ "\n" +content+ "\n"+vlm_text
|
||||||
|
|
||||||
|
# image_content = (
|
||||||
|
# f"{output_filename}\n"
|
||||||
|
# f"{image_related_content}"
|
||||||
|
# )
|
||||||
|
logger.info(f"output_filename: {output_filename}")
|
||||||
|
logger.info(f"{image_related_content}")
|
||||||
|
return JSONResponse(
|
||||||
|
{
|
||||||
|
"code": 200,
|
||||||
|
"message": "ok",
|
||||||
|
"data": {
|
||||||
|
"image_content": image_related_content,
|
||||||
|
"filename": filename,
|
||||||
|
"jpg_filename": output_filename
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
if temp_path and temp_path.exists():
|
||||||
|
try:
|
||||||
|
temp_path.unlink()
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"Failed to delete temp image {temp_path}: {e}")
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
@app.post("/split_content_list")
|
@app.post("/split_content_list")
|
||||||
@ -2133,7 +2232,7 @@ async def split_result(
|
|||||||
slices = []
|
slices = []
|
||||||
for ins in contents:
|
for ins in contents:
|
||||||
slices.append({"content": ins, "positions": []})
|
slices.append({"content": ins, "positions": []})
|
||||||
result_data = {"slices": slices, "images": {}}
|
result_data = {"slices": slices, "images": {},"summary":''}
|
||||||
elif suffix in [".xlsx"]:
|
elif suffix in [".xlsx"]:
|
||||||
# converter.convert(str(temp_input_path), output_dir=str(DATA_DIR))
|
# converter.convert(str(temp_input_path), output_dir=str(DATA_DIR))
|
||||||
# temp_base_name = temp_input_path.stem # 如 tmp2bb0l4cn
|
# temp_base_name = temp_input_path.stem # 如 tmp2bb0l4cn
|
||||||
@ -2221,96 +2320,6 @@ class GraphSearchRequest(BaseModel):
|
|||||||
entry_nodes: Optional[Dict[str, Any]] = None
|
entry_nodes: Optional[Dict[str, Any]] = None
|
||||||
|
|
||||||
|
|
||||||
@app.post("/search_graph——1")
|
|
||||||
async def search_graph(request: GraphSearchRequest, background_tasks: BackgroundTasks):
|
|
||||||
"""
|
|
||||||
========== ========== ========== ========== ========== ==========
|
|
||||||
纯图谱检索
|
|
||||||
|
|
||||||
请求参数:
|
|
||||||
- query: 用户查询文本(必填)
|
|
||||||
- top_k: 返回用于展示的路径数量(默认10)
|
|
||||||
- graph_timeout: 图谱检索超时时间,秒(默认2.0,暂未使用)
|
|
||||||
- entry_nodes: 入口节点信息(可选,不提供则从查询中提取)
|
|
||||||
|
|
||||||
========== ========== ========== ========== ========== ==========
|
|
||||||
"""
|
|
||||||
|
|
||||||
start_time = time.time()
|
|
||||||
query = request.query
|
|
||||||
|
|
||||||
# ========== 参数验证 ==========
|
|
||||||
if not query or not query.strip():
|
|
||||||
raise HTTPException(status_code=400, detail="查询文本不能为空")
|
|
||||||
|
|
||||||
query = query.strip()
|
|
||||||
|
|
||||||
logger.info(f"[search_graph] 图谱检索请求: query='{query}', top_k={request.top_k}")
|
|
||||||
|
|
||||||
try:
|
|
||||||
# ========== 检查依赖是否可用 ==========
|
|
||||||
if not GRAPH_AVAILABLE:
|
|
||||||
raise HTTPException(
|
|
||||||
status_code=503,
|
|
||||||
detail="图谱检索功能依赖未安装,请安装: pip install neo4j langchain langchain-community neo4j-graphrag jieba",
|
|
||||||
)
|
|
||||||
|
|
||||||
# ========== 初始化图谱服务(复用app.py的neo4j driver) ==========
|
|
||||||
# 注意:传入app.py中创建的driver,实现连接复用
|
|
||||||
graph_service = GraphService(driver=driver)
|
|
||||||
|
|
||||||
# ========== 执行图谱检索(简单版本:只判断统计/非统计) ==========
|
|
||||||
# 说明:
|
|
||||||
# - 直接调用 GraphService.search,自动判断统计/非统计查询类型
|
|
||||||
# - query_type="auto" 会自动识别统计类查询(如"多少"、"比例"等)和明细类查询
|
|
||||||
result = await graph_service.search(
|
|
||||||
query=query,
|
|
||||||
entry_nodes=request.entry_nodes,
|
|
||||||
top_k=request.top_k,
|
|
||||||
max_attempts=3,
|
|
||||||
query_type="auto", # 自动判断统计/非统计
|
|
||||||
)
|
|
||||||
|
|
||||||
# ========== 确保返回格式一致 ==========
|
|
||||||
if not isinstance(result, dict):
|
|
||||||
result = {
|
|
||||||
"code": 200,
|
|
||||||
"message": "success",
|
|
||||||
"data": result if isinstance(result, (list, dict)) else ([] if isinstance(result, list) else {}),
|
|
||||||
"meta": {"elapsed_ms": round((time.time() - start_time) * 1000, 2)},
|
|
||||||
}
|
|
||||||
|
|
||||||
# 确保data存在,可以是列表或字典(新格式:包含 nodes、links、results 的对象)
|
|
||||||
if "data" not in result:
|
|
||||||
result["data"] = {}
|
|
||||||
# 注意:不再强制转换为列表,支持新的对象格式
|
|
||||||
|
|
||||||
# 统计结果数量(兼容新旧格式)
|
|
||||||
if isinstance(result.get("data"), list):
|
|
||||||
result_count = len(result["data"])
|
|
||||||
elif isinstance(result.get("data"), dict):
|
|
||||||
# 新格式:统计 results 数组的长度
|
|
||||||
result_count = len(result["data"].get("results", []))
|
|
||||||
else:
|
|
||||||
result_count = 0
|
|
||||||
|
|
||||||
logger.info(f"[search_graph] 图谱检索完成: 返回 {result_count} 条结果")
|
|
||||||
return result
|
|
||||||
|
|
||||||
except HTTPException:
|
|
||||||
raise
|
|
||||||
except Exception as e:
|
|
||||||
elapsed = time.time() - start_time
|
|
||||||
logger.error(f"[search_graph] 图谱检索失败 (耗时: {elapsed * 1000:.2f}ms): {e}", exc_info=True)
|
|
||||||
return {
|
|
||||||
"code": 500,
|
|
||||||
"message": "图谱检索失败",
|
|
||||||
"data": [],
|
|
||||||
"meta": {"elapsed_ms": round(elapsed * 1000, 2), "error": str(e), "query_type": "error"},
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
# ========== ========== ========== ========== ========== ==========
|
# ========== ========== ========== ========== ========== ==========
|
||||||
# Neo4j 索引创建接口
|
# Neo4j 索引创建接口
|
||||||
# ========== ========== ========== ========== ========== ==========
|
# ========== ========== ========== ========== ========== ==========
|
||||||
@ -2920,6 +2929,3 @@ if __name__ == "__main__":
|
|||||||
import uvicorn
|
import uvicorn
|
||||||
|
|
||||||
uvicorn.run(app, host="0.0.0.0", port=9085)
|
uvicorn.run(app, host="0.0.0.0", port=9085)
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
934
chunk_text.py
934
chunk_text.py
File diff suppressed because it is too large
Load Diff
97
config.py
97
config.py
@ -17,7 +17,6 @@ EMBEDDING_CONFIG = {
|
|||||||
"model": "bge-m3",
|
"model": "bge-m3",
|
||||||
"base_url": "http://192.168.0.46:59700/v1/embeddings",
|
"base_url": "http://192.168.0.46:59700/v1/embeddings",
|
||||||
"api_key": "gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5",
|
"api_key": "gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5",
|
||||||
"base_url_v2" : "http://192.168.0.46:59700/v1",
|
|
||||||
}
|
}
|
||||||
|
|
||||||
# ==================== Neo4j 图数据库配置 ====================
|
# ==================== Neo4j 图数据库配置 ====================
|
||||||
@ -75,88 +74,18 @@ API_OTHER_URLS = [
|
|||||||
# "http://192.168.0.111:9975/analyze-otherfile",
|
# "http://192.168.0.111:9975/analyze-otherfile",
|
||||||
]
|
]
|
||||||
|
|
||||||
SHIP_MODEL_NAME = "E:\ZKYNLP\Hjunproject\project0506\kgrag\ship_xinghao.json"
|
# ==================== 切片图片处理配置 ====================
|
||||||
TREE_JSON_PATH = r"E:\ZKYNLP\Hjunproject\project0506\kgrag\tree_data.json"
|
MAX_SPLIT_IMAGE_INPUT_BYTES = 50 * 1024 * 1024
|
||||||
root@70964b104c5b:/app# pwd
|
MAX_SPLIT_IMAGE_BYTES = 10 * 1024 * 1024
|
||||||
/app
|
MAX_SPLIT_IMAGE_BASE64_CHARS = (MAX_SPLIT_IMAGE_INPUT_BYTES * 4 // 3) + 4096
|
||||||
root@70964b104c5b:/app# vim config.py
|
SPLIT_IMAGE_COMPRESS_THRESHOLD_BYTES = MAX_SPLIT_IMAGE_BYTES
|
||||||
root@70964b104c5b:/app# cat config.py
|
SPLIT_IMAGE_MAX_DIMENSION = 4096
|
||||||
# app/model_config.py
|
SPLIT_IMAGE_JPEG_QUALITY = 88
|
||||||
|
OCR_CONCURRENCY = 1
|
||||||
import os
|
VLM_CONCURRENCY = 2
|
||||||
|
|
||||||
# ==================== 大模型配置 ====================
|
|
||||||
LLM_CONFIG = {
|
|
||||||
"model": "46-qwen3.5-35B",
|
|
||||||
"base_url": "http://192.168.0.46:59800/v1",
|
|
||||||
"api_key": "gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5",
|
|
||||||
"temperature": 0.1,
|
|
||||||
"max_tokens": 55096,
|
|
||||||
"timeout": 30
|
|
||||||
}
|
|
||||||
|
|
||||||
# ==================== 嵌入模型配置 ====================
|
|
||||||
EMBEDDING_CONFIG = {
|
|
||||||
"model": "bge-m3",
|
|
||||||
"base_url": "http://192.168.0.46:59700/v1/embeddings",
|
|
||||||
"api_key": "gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5",
|
|
||||||
"base_url_v2" : "http://192.168.0.46:59700/v1",
|
|
||||||
}
|
|
||||||
|
|
||||||
# ==================== Neo4j 图数据库配置 ====================
|
|
||||||
NEO4J_CONFIG = {
|
|
||||||
# 优先读取环境变量,如果没有则使用默认值
|
|
||||||
# 注意:本地开发通常用 localhost,部署时用服务名 neo4j
|
|
||||||
"uri": "neo4j://192.168.0.46:57687",
|
|
||||||
"username": "neo4j",
|
|
||||||
"password": "zdht123@",
|
|
||||||
"database": "neo4j", # 新增 database 配置
|
|
||||||
}
|
|
||||||
|
|
||||||
# ==================== 检索与索引配置 ====================
|
|
||||||
SEARCH_CONFIG = {
|
|
||||||
"name_property": "名称",
|
|
||||||
"fulltext_property": "fulltext",
|
|
||||||
"embedding_property": "embedding",
|
|
||||||
"search_label": "Searchable",
|
|
||||||
"vector_index_name": "global_searchable_embedding",
|
|
||||||
"fulltext_index_name": "global_searchable_content_search",
|
|
||||||
"fulltext_field_index_name": "global_searchable_fulltext_search",
|
|
||||||
"global_entity_embedding" : "global_entity_embedding",
|
|
||||||
"global_entity_content_search": "global_entity_content_search",
|
|
||||||
"excluded_business_labels": [
|
|
||||||
"Entity", "Chunk", "Document", "_Bloom_Perspective_", "Searchable"
|
|
||||||
],
|
|
||||||
"vector_dimension": 1024,
|
|
||||||
"jieba_pos_whitelist": ["n", "nr", "ns", "nt", "nz", "vn", "eng"],
|
|
||||||
"embed_batch_size" : 64,
|
|
||||||
"embed_thread_num" : 4,
|
|
||||||
"node_fetch_page" : 10000,
|
|
||||||
}
|
|
||||||
|
|
||||||
INDEXING_ULR = "http://192.168.0.46:59085/neo4j_indexing"
|
|
||||||
|
|
||||||
URL = {
|
|
||||||
"indexing_url" : "http://192.168.0.46:59085/neo4j_indexing",
|
|
||||||
"prefix_url" : "http://192.168.0.46:59085",
|
|
||||||
}
|
|
||||||
|
|
||||||
API_URLS = [
|
|
||||||
"http://192.168.0.64:59988/analyze-pdf",
|
|
||||||
# "http://192.168.0.111:9977/analyze-pdf",
|
|
||||||
# "http://192.168.0.111:9978/analyze-pdf",
|
|
||||||
#"http://192.168.0.111:9979/analyze-pdf",
|
|
||||||
#"http://192.168.0.111:9980/analyze-pdf",
|
|
||||||
# "http://192.168.0.111:9975/analyze-pdf",
|
|
||||||
]
|
|
||||||
API_OTHER_URLS = [
|
|
||||||
"http://192.168.0.46:59988/analyze-otherfile",
|
|
||||||
# "http://192.168.0.111:9977/analyze-otherfile",
|
|
||||||
# "http://192.168.0.111:9978/analyze-otherfile",
|
|
||||||
#"http://192.168.0.111:9979/analyze-otherfile",
|
|
||||||
#"http://192.168.0.111:9980/analyze-otherfile",
|
|
||||||
# "http://192.168.0.111:9975/analyze-otherfile",
|
|
||||||
]
|
|
||||||
|
|
||||||
SHIP_MODEL_NAME = "/app/ship_xinghao.json"
|
SHIP_MODEL_NAME = "/app/ship_xinghao.json"
|
||||||
TREE_JSON_PATH = "/app/tree_data.json"
|
TREE_JSON_PATH = "/app/tree_data.json"
|
||||||
|
# ==================== mineru和kgrag目录映射地址 ====================
|
||||||
|
SEARCH_DIR ="/app/mineru_output"
|
||||||
|
IMAGE_DIR = "/app/mineru_output/images"
|
||||||
@ -13,7 +13,6 @@ import json
|
|||||||
from modelsAPI.model_api import OpenaiAPI
|
from modelsAPI.model_api import OpenaiAPI
|
||||||
from openai import OpenAI
|
from openai import OpenAI
|
||||||
from typing import Dict, List, Any
|
from typing import Dict, List, Any
|
||||||
from config import TREE_JSON_PATH
|
|
||||||
|
|
||||||
app = FastAPI(title="PDF Upload Service")
|
app = FastAPI(title="PDF Upload Service")
|
||||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s")
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s")
|
||||||
@ -541,8 +540,8 @@ def fetch_graph_sample1(driver):
|
|||||||
# 文件名处理方法
|
# 文件名处理方法
|
||||||
# =============================
|
# =============================
|
||||||
|
|
||||||
# BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||||
# TREE_JSON_PATH = os.path.join(BASE_DIR, "tree_data.json")
|
TREE_JSON_PATH = os.path.join(BASE_DIR, "tree_data.json")
|
||||||
|
|
||||||
|
|
||||||
def _load_tree_data() -> Dict:
|
def _load_tree_data() -> Dict:
|
||||||
|
|||||||
@ -3,9 +3,9 @@ import json
|
|||||||
import base64
|
import base64
|
||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
from typing import Optional, Tuple, Dict
|
from typing import Any, Optional, Tuple, Dict
|
||||||
import aiofiles # pip install aiofiles
|
import aiofiles # pip install aiofiles
|
||||||
|
from config import SEARCH_DIR,IMAGE_DIR
|
||||||
logging.basicConfig(level=logging.INFO)
|
logging.basicConfig(level=logging.INFO)
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@ -16,7 +16,7 @@ async def find_and_read_content_list(
|
|||||||
encoding: str = 'utf-8'
|
encoding: str = 'utf-8'
|
||||||
) -> Tuple[Optional[list], Optional[str]]:
|
) -> Tuple[Optional[list], Optional[str]]:
|
||||||
"""根据原始文件名查找并读取对应的 _content_list.json 文件。"""
|
"""根据原始文件名查找并读取对应的 _content_list.json 文件。"""
|
||||||
file_name, _ = os.path.splitext(original_filename)
|
file_name, _ = os.path.splitext(os.path.basename(original_filename))
|
||||||
target_filename = f"{file_name}_content_list.json"
|
target_filename = f"{file_name}_content_list.json"
|
||||||
logger.info(f"正在查找文件: {target_filename}")
|
logger.info(f"正在查找文件: {target_filename}")
|
||||||
|
|
||||||
@ -44,6 +44,93 @@ async def find_and_read_content_list(
|
|||||||
return None, None
|
return None, None
|
||||||
|
|
||||||
|
|
||||||
|
def find_image_record(content_list: Any, image_filename: str) -> Optional[dict]:
|
||||||
|
"""Find an image record by the basename stored in ``img_path``."""
|
||||||
|
if not isinstance(content_list, list) or not image_filename:
|
||||||
|
return None
|
||||||
|
|
||||||
|
target_name = os.path.basename(str(image_filename).strip().replace('\\', '/'))
|
||||||
|
if not target_name:
|
||||||
|
return None
|
||||||
|
|
||||||
|
for item in content_list:
|
||||||
|
if not isinstance(item, dict) or item.get('type') != 'image':
|
||||||
|
continue
|
||||||
|
|
||||||
|
img_path = item.get('img_path')
|
||||||
|
if isinstance(img_path, str):
|
||||||
|
path_name = os.path.basename(img_path.strip().replace('\\', '/'))
|
||||||
|
if path_name == target_name:
|
||||||
|
return item
|
||||||
|
|
||||||
|
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
async def find_image_record_by_pdf(
|
||||||
|
directory: str,
|
||||||
|
pdf_filename: str,
|
||||||
|
image_filename: str,
|
||||||
|
encoding: str = 'utf-8',
|
||||||
|
) -> Optional[dict]:
|
||||||
|
"""Find an image record from the content list matching a PDF filename."""
|
||||||
|
content_list, _ = await find_and_read_content_list(
|
||||||
|
directory, pdf_filename, encoding=encoding
|
||||||
|
)
|
||||||
|
return find_image_record(content_list, image_filename)
|
||||||
|
|
||||||
|
|
||||||
|
async def extract_image_caption(
|
||||||
|
search_directory: str,
|
||||||
|
pdf_filename: str,
|
||||||
|
image_filename: str,
|
||||||
|
) -> list[str]:
|
||||||
|
"""?? PDF ??????????? image_caption?"""
|
||||||
|
image_record = await find_image_record_by_pdf(
|
||||||
|
directory=search_directory,
|
||||||
|
pdf_filename=pdf_filename,
|
||||||
|
image_filename=image_filename,
|
||||||
|
)
|
||||||
|
|
||||||
|
if not image_record:
|
||||||
|
return []
|
||||||
|
|
||||||
|
image_caption = image_record.get("image_caption", [])
|
||||||
|
if isinstance(image_caption, str):
|
||||||
|
return [image_caption] if image_caption.strip() else []
|
||||||
|
if isinstance(image_caption, list):
|
||||||
|
return [str(item).strip() for item in image_caption if str(item).strip()]
|
||||||
|
return []
|
||||||
|
|
||||||
|
|
||||||
|
async def find_and_read_md_file(
|
||||||
|
directory: str,
|
||||||
|
original_filename: str,
|
||||||
|
encoding: str = 'utf-8'
|
||||||
|
) -> Tuple[Optional[str], Optional[str]]:
|
||||||
|
"""Find and read the Markdown file matching the uploaded document basename."""
|
||||||
|
file_name, _ = os.path.splitext(original_filename)
|
||||||
|
target_filename = f"{file_name}.md"
|
||||||
|
logger.info(f"正在查找 Markdown 文件: {target_filename}")
|
||||||
|
|
||||||
|
def _walk_for_file():
|
||||||
|
for root, _, files in os.walk(directory):
|
||||||
|
if target_filename in files:
|
||||||
|
return os.path.join(root, target_filename)
|
||||||
|
return None
|
||||||
|
|
||||||
|
file_path = await asyncio.to_thread(_walk_for_file)
|
||||||
|
if file_path is None:
|
||||||
|
return None, None
|
||||||
|
|
||||||
|
try:
|
||||||
|
async with aiofiles.open(file_path, 'r', encoding=encoding, errors='ignore') as f:
|
||||||
|
return await f.read(), file_path
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"读取 Markdown 文件 {file_path} 时发生错误: {e}")
|
||||||
|
return None, None
|
||||||
|
|
||||||
|
|
||||||
def extract_all_jpg_filenames(doc_data: list) -> set:
|
def extract_all_jpg_filenames(doc_data: list) -> set:
|
||||||
"""从文档数据中提取所有 .jpg 图片的纯文件名。纯 CPU 操作,不需要 async。"""
|
"""从文档数据中提取所有 .jpg 图片的纯文件名。纯 CPU 操作,不需要 async。"""
|
||||||
result_set = set()
|
result_set = set()
|
||||||
@ -103,10 +190,9 @@ async def encode_images_to_base64(
|
|||||||
|
|
||||||
async def process_document(
|
async def process_document(
|
||||||
input_file_name: str,
|
input_file_name: str,
|
||||||
search_dir: Optional[str] = "/app/mineru_output",
|
|
||||||
local_image_dir: Optional[str] = None,
|
local_image_dir: Optional[str] = None,
|
||||||
concurrency: int = 32,
|
concurrency: int = 32,
|
||||||
) -> Optional[Tuple[list, Dict[str, str]]]:
|
) -> Optional[Tuple[list, Dict[str, str], Optional[str]]]:
|
||||||
"""
|
"""
|
||||||
异步处理文档:查找 content_list.json,提取图片引用,并发编码为 base64。
|
异步处理文档:查找 content_list.json,提取图片引用,并发编码为 base64。
|
||||||
|
|
||||||
@ -117,9 +203,10 @@ async def process_document(
|
|||||||
concurrency: 图片编码的并发数上限,默认 32
|
concurrency: 图片编码的并发数上限,默认 32
|
||||||
|
|
||||||
Returns:
|
Returns:
|
||||||
成功: (content_list, images_dict)
|
成功: (content_list, images_dict, md_content)
|
||||||
失败: None
|
失败: None
|
||||||
"""
|
"""
|
||||||
|
search_dir = SEARCH_DIR
|
||||||
content, content_path = await find_and_read_content_list(search_dir, input_file_name)
|
content, content_path = await find_and_read_content_list(search_dir, input_file_name)
|
||||||
if content is None:
|
if content is None:
|
||||||
logger.error("❌ 未找到指定的 _content_list.json 文件。")
|
logger.error("❌ 未找到指定的 _content_list.json 文件。")
|
||||||
@ -128,6 +215,13 @@ async def process_document(
|
|||||||
logger.info(f"✅ 找到文件: {content_path}")
|
logger.info(f"✅ 找到文件: {content_path}")
|
||||||
logger.info(f"内容类型: {type(content).__name__}, 长度: {len(str(content))}")
|
logger.info(f"内容类型: {type(content).__name__}, 长度: {len(str(content))}")
|
||||||
|
|
||||||
|
md_content, md_path = await find_and_read_md_file(search_dir, input_file_name)
|
||||||
|
if md_content is None:
|
||||||
|
logger.info("未找到对应的 Markdown 文件。")
|
||||||
|
else:
|
||||||
|
logger.info(f"找到 Markdown 文件: {md_path}")
|
||||||
|
logger.info(f"Markdown 内容长度: {len(md_content)}")
|
||||||
|
|
||||||
if local_image_dir is None:
|
if local_image_dir is None:
|
||||||
local_image_dir = os.path.join(os.path.dirname(content_path), "images")
|
local_image_dir = os.path.join(os.path.dirname(content_path), "images")
|
||||||
logger.info(f"图片目录: {local_image_dir}")
|
logger.info(f"图片目录: {local_image_dir}")
|
||||||
@ -140,20 +234,19 @@ async def process_document(
|
|||||||
)
|
)
|
||||||
logger.info(f"成功编码 {len(images_dict)} 张图片")
|
logger.info(f"成功编码 {len(images_dict)} 张图片")
|
||||||
|
|
||||||
return content, images_dict
|
return content, images_dict, md_content
|
||||||
|
|
||||||
|
|
||||||
async def main():
|
async def main():
|
||||||
result = await process_document(
|
result = await process_document(
|
||||||
input_file_name="163-06A0014-B01001_发动机-维修手册.pdf",
|
input_file_name="163-06A0014-B01001_发动机-维修手册.pdf",
|
||||||
search_dir="/app/mineru_output",
|
|
||||||
concurrency=32,
|
concurrency=32,
|
||||||
)
|
)
|
||||||
|
|
||||||
if result is None:
|
if result is None:
|
||||||
return 1
|
return 1
|
||||||
|
|
||||||
content, images_dict = result
|
content, images_dict, md_content = result
|
||||||
print(f"\n=== 处理完成 ===")
|
print(f"\n=== 处理完成 ===")
|
||||||
print(f"文档段落数: {len(content)}")
|
print(f"文档段落数: {len(content)}")
|
||||||
print(f"图片数量: {len(images_dict)}")
|
print(f"图片数量: {len(images_dict)}")
|
||||||
@ -164,5 +257,11 @@ async def main():
|
|||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
exit_code = asyncio.run(main())
|
result = asyncio.run(
|
||||||
exit(exit_code)
|
extract_image_caption(
|
||||||
|
r"E:\ZKYNLP\Hjunproject\project0506\kgrag",
|
||||||
|
"163-06A0016-B01003_雷达-维修手册.pdf",
|
||||||
|
"",
|
||||||
|
)
|
||||||
|
)
|
||||||
|
print(result)
|
||||||
|
|||||||
142
generate_summary.py
Normal file
142
generate_summary.py
Normal file
@ -0,0 +1,142 @@
|
|||||||
|
"""
|
||||||
|
文件概述生成工具
|
||||||
|
把 wiki_engine 的文档概述能力独立出来
|
||||||
|
输入:markdown 格式的文件内容
|
||||||
|
输出:文件概述(SUMMARY 行 + markdown 正文)
|
||||||
|
支持并发
|
||||||
|
"""
|
||||||
|
import asyncio
|
||||||
|
from typing import List, Optional
|
||||||
|
from config import LLM_CONFIG
|
||||||
|
from openai import AsyncOpenAI
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== 配置 ====================
|
||||||
|
_MAX_CONCURRENCY = 8 # 最大并发数
|
||||||
|
|
||||||
|
_client: Optional[AsyncOpenAI] = None
|
||||||
|
_semaphore: Optional[asyncio.Semaphore] = None
|
||||||
|
|
||||||
|
|
||||||
|
def _get_client() -> AsyncOpenAI:
|
||||||
|
"""获取或创建 AsyncOpenAI 单例客户端"""
|
||||||
|
global _client
|
||||||
|
if _client is None:
|
||||||
|
_client = AsyncOpenAI(api_key=LLM_CONFIG['api_key'], base_url=LLM_CONFIG['base_url'])
|
||||||
|
return _client
|
||||||
|
|
||||||
|
|
||||||
|
def _get_semaphore() -> asyncio.Semaphore:
|
||||||
|
"""获取或创建并发信号量"""
|
||||||
|
global _semaphore
|
||||||
|
if _semaphore is None:
|
||||||
|
_semaphore = asyncio.Semaphore(_MAX_CONCURRENCY)
|
||||||
|
return _semaphore
|
||||||
|
|
||||||
|
|
||||||
|
# ==================== Prompt(与 wiki_builder.py 的 WIKI_SUMMARY_PROMPT 一致,仅去掉 extracted_slugs 相关部分) ====================
|
||||||
|
SUMMARY_PROMPT = """You are a wiki editor. Given the following document content, create a structured wiki summary page in Markdown format.
|
||||||
|
|
||||||
|
<document>
|
||||||
|
<content>
|
||||||
|
{content}
|
||||||
|
</content>
|
||||||
|
</document>
|
||||||
|
|
||||||
|
<instructions>
|
||||||
|
1. The FIRST line of your output MUST be: SUMMARY: {{one sentence, 15-40 words, describing what this document is about for wiki index listing}}
|
||||||
|
2. Create a concise but useful document-level summary.
|
||||||
|
3. Include the document's main subject, scope, important procedures, standards, systems, equipment, tables, fields.
|
||||||
|
4. Do NOT invent facts. Stay grounded in the document content.
|
||||||
|
5. Write in Chinese.
|
||||||
|
6. If the content is empty or has no substantive information, output exactly: "SUMMARY: No textual content was extractable from this document." followed by a brief note.
|
||||||
|
</instructions>
|
||||||
|
|
||||||
|
Output the SUMMARY line first, then the Markdown content. Do not include any other preamble."""
|
||||||
|
|
||||||
|
|
||||||
|
async def generate_summary(content: str) -> str:
|
||||||
|
"""
|
||||||
|
从 markdown 内容生成文件概述
|
||||||
|
|
||||||
|
Args:
|
||||||
|
content: markdown 格式的文件内容
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
文件概述(SUMMARY 行 + markdown 正文)
|
||||||
|
"""
|
||||||
|
if not content or not content.strip():
|
||||||
|
return ""
|
||||||
|
|
||||||
|
prompt = SUMMARY_PROMPT.format(content=content)
|
||||||
|
client = _get_client()
|
||||||
|
|
||||||
|
async with _get_semaphore():
|
||||||
|
try:
|
||||||
|
response = await client.chat.completions.create(
|
||||||
|
model=LLM_CONFIG['model'],
|
||||||
|
messages=[
|
||||||
|
{"role": "system", "content": "You are a grounded wiki editor. Do not invent facts."},
|
||||||
|
{"role": "user", "content": prompt},
|
||||||
|
],
|
||||||
|
temperature=0.1,
|
||||||
|
max_tokens=LLM_CONFIG['max_tokens'],
|
||||||
|
stream=False,
|
||||||
|
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
|
||||||
|
)
|
||||||
|
return (response.choices[0].message.content or "").strip()
|
||||||
|
except Exception as exc:
|
||||||
|
print(f"[generate_summary] 概述生成失败: {exc}")
|
||||||
|
return ""
|
||||||
|
|
||||||
|
|
||||||
|
async def generate_summaries(contents: List[str]) -> List[str]:
|
||||||
|
"""
|
||||||
|
并发生成多个文件的概述
|
||||||
|
|
||||||
|
Args:
|
||||||
|
contents: markdown 格式的文件内容列表
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
文件概述列表,顺序与输入一致
|
||||||
|
"""
|
||||||
|
return await asyncio.gather(*(generate_summary(c) for c in contents))
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
# import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# if len(sys.argv) > 1:
|
||||||
|
# with open(sys.argv[1], "r", encoding="utf-8") as f:
|
||||||
|
# test_content = f.read()
|
||||||
|
# else:
|
||||||
|
# test_content = """
|
||||||
|
# # 船舶动力系统维护规程
|
||||||
|
|
||||||
|
# ## 概述
|
||||||
|
# 本文档详细介绍了船舶动力系统的日常维护和故障处理流程。
|
||||||
|
|
||||||
|
# ## 发动机日常检查
|
||||||
|
# - 润滑油位检查:每日检查发动机润滑油位,保持在标尺正常范围
|
||||||
|
# - 冷却液位检查:确保冷却系统液位正常,无泄漏
|
||||||
|
# - 皮带张紧度:检查传动皮带张紧度,过松或过紧均需调整
|
||||||
|
|
||||||
|
# ## 冷却系统维护
|
||||||
|
# 定期清洗热交换器,检查水泵密封性,更换老化管路。
|
||||||
|
|
||||||
|
# ## 故障诊断流程
|
||||||
|
# 1. 现象观察:记录故障现象和发生条件
|
||||||
|
# 2. 数据采集:收集运行参数和报警信息
|
||||||
|
# 3. 原因分析:对照标准参数分析故障原因
|
||||||
|
# 4. 处理方案:制定维修方案并执行
|
||||||
|
|
||||||
|
# ## 安全注意事项
|
||||||
|
# 所有维护操作必须在停机状态下进行,操作人员需佩戴防护装备。
|
||||||
|
# """
|
||||||
|
md_path = Path(r"E:\ZKYNLP\Hjunproject\project0506\kgrag\船舶主机燃油泵自动控制系统故障树分析.md")
|
||||||
|
with open(md_path, "r", encoding="utf-8") as f:
|
||||||
|
test_content = f.read()
|
||||||
|
|
||||||
|
result = asyncio.run(generate_summary(test_content))
|
||||||
|
print(result)
|
||||||
@ -5,9 +5,13 @@
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
from typing import Dict, List, Any, Optional
|
import re
|
||||||
|
from typing import Dict, List, Any, Optional, Set
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
EXCLUDED_RESULT_NODE_LABELS = {"维修工作", "操作程序", "操作使用"}
|
||||||
|
_FILTERED_OUT = object()
|
||||||
|
|
||||||
|
|
||||||
class SimpleNode:
|
class SimpleNode:
|
||||||
"""
|
"""
|
||||||
@ -537,6 +541,110 @@ def format_entry_nodes_as_results(entry_nodes: dict) -> List[Dict[str, Any]]:
|
|||||||
return results
|
return results
|
||||||
|
|
||||||
|
|
||||||
|
def _normalize_label_values(labels: Any) -> Set[str]:
|
||||||
|
"""将 labels/标签/type 等字段统一成标签集合。"""
|
||||||
|
if labels is None:
|
||||||
|
return set()
|
||||||
|
|
||||||
|
if isinstance(labels, str):
|
||||||
|
parts = re.split(r"[,,、/\s]+", labels)
|
||||||
|
return {part.strip("[]'\" ") for part in parts if part.strip("[]'\" ")}
|
||||||
|
|
||||||
|
if isinstance(labels, (list, tuple, set, frozenset)):
|
||||||
|
return {str(label).strip() for label in labels if str(label).strip()}
|
||||||
|
|
||||||
|
return {str(labels).strip()} if str(labels).strip() else set()
|
||||||
|
|
||||||
|
|
||||||
|
def _dict_label_values(data: Dict[str, Any]) -> Set[str]:
|
||||||
|
label_keys = ("labels", "标签", "label", "类型", "type", "节点类型", "node_type")
|
||||||
|
labels = set()
|
||||||
|
for key in label_keys:
|
||||||
|
if key in data:
|
||||||
|
labels.update(_normalize_label_values(data.get(key)))
|
||||||
|
return labels
|
||||||
|
|
||||||
|
|
||||||
|
def _has_excluded_result_label(value: Any) -> bool:
|
||||||
|
"""判断一个结果对象是否明确属于需要从 results 隐藏的节点类型。"""
|
||||||
|
labels = set()
|
||||||
|
|
||||||
|
if isinstance(value, dict):
|
||||||
|
labels.update(_dict_label_values(value))
|
||||||
|
elif hasattr(value, "labels") and (hasattr(value, "element_id") or hasattr(value, "id")):
|
||||||
|
try:
|
||||||
|
labels.update(_normalize_label_values(value.labels))
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
|
||||||
|
return bool(labels & EXCLUDED_RESULT_NODE_LABELS)
|
||||||
|
|
||||||
|
|
||||||
|
def _mentions_excluded_result_label(text: Any) -> bool:
|
||||||
|
if text is None:
|
||||||
|
return False
|
||||||
|
return any(label in str(text) for label in EXCLUDED_RESULT_NODE_LABELS)
|
||||||
|
|
||||||
|
|
||||||
|
def _is_empty_result_content(value: Any) -> bool:
|
||||||
|
return value in (None, {}, [])
|
||||||
|
|
||||||
|
|
||||||
|
def _filter_excluded_result_content(value: Any, parent_key: Any = None) -> Any:
|
||||||
|
"""
|
||||||
|
只用于 results 字段:过滤掉维修工作、操作程序、操作使用这类节点内容。
|
||||||
|
nodes/links 的原始返回不会经过这个函数。
|
||||||
|
"""
|
||||||
|
if parent_key is not None and _mentions_excluded_result_label(parent_key):
|
||||||
|
return _FILTERED_OUT
|
||||||
|
|
||||||
|
if _has_excluded_result_label(value):
|
||||||
|
return _FILTERED_OUT
|
||||||
|
|
||||||
|
if isinstance(value, dict):
|
||||||
|
name_value = value.get("名称", value.get("name"))
|
||||||
|
if _mentions_excluded_result_label(name_value):
|
||||||
|
return _FILTERED_OUT
|
||||||
|
|
||||||
|
filtered = {}
|
||||||
|
for key, child in value.items():
|
||||||
|
if _mentions_excluded_result_label(key):
|
||||||
|
continue
|
||||||
|
filtered_child = _filter_excluded_result_content(child, key)
|
||||||
|
if filtered_child is not _FILTERED_OUT:
|
||||||
|
filtered[key] = filtered_child
|
||||||
|
|
||||||
|
return _FILTERED_OUT if _is_empty_result_content(filtered) else filtered
|
||||||
|
|
||||||
|
if isinstance(value, list):
|
||||||
|
filtered_items = []
|
||||||
|
for item in value:
|
||||||
|
filtered_item = _filter_excluded_result_content(item)
|
||||||
|
if filtered_item is not _FILTERED_OUT and not _is_empty_result_content(filtered_item):
|
||||||
|
filtered_items.append(filtered_item)
|
||||||
|
return filtered_items
|
||||||
|
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
|
def _filter_results_for_output(results: Any) -> List[Any]:
|
||||||
|
filtered = _filter_excluded_result_content(results)
|
||||||
|
if filtered is _FILTERED_OUT or _is_empty_result_content(filtered):
|
||||||
|
return []
|
||||||
|
if isinstance(filtered, list):
|
||||||
|
return filtered
|
||||||
|
return [filtered]
|
||||||
|
|
||||||
|
|
||||||
|
def _filter_result_nodes_for_output(nodes: Any) -> List[Dict[str, Any]]:
|
||||||
|
if not isinstance(nodes, list):
|
||||||
|
return []
|
||||||
|
return [
|
||||||
|
node for node in nodes
|
||||||
|
if isinstance(node, dict) and not _has_excluded_result_label(node)
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
def _format_value(value: Any) -> str:
|
def _format_value(value: Any) -> str:
|
||||||
"""
|
"""
|
||||||
@ -649,8 +757,9 @@ def serialize_graph_for_llm(graph_response: dict) -> str:
|
|||||||
- results(任意查询返回值:聚合/记录/字符串等)
|
- results(任意查询返回值:聚合/记录/字符串等)
|
||||||
"""
|
"""
|
||||||
data = graph_response.get("data", {})
|
data = graph_response.get("data", {})
|
||||||
nodes = data.get("nodes", [])
|
# 这里的过滤只影响 results 字段里的文本展示,不改动接口返回的 data.nodes/data.links。
|
||||||
results = data.get("results", [])
|
nodes = _filter_result_nodes_for_output(data.get("nodes", []))
|
||||||
|
results = _filter_results_for_output(data.get("results", []))
|
||||||
|
|
||||||
output_lines = []
|
output_lines = []
|
||||||
|
|
||||||
@ -695,3 +804,4 @@ def serialize_graph_for_llm(graph_response: dict) -> str:
|
|||||||
return "图谱查询完成,但未返回有效数据。"
|
return "图谱查询完成,但未返回有效数据。"
|
||||||
|
|
||||||
return "\n".join(output_lines).rstrip()
|
return "\n".join(output_lines).rstrip()
|
||||||
|
|
||||||
|
|||||||
@ -15,7 +15,19 @@ from neo4j_graphrag.embeddings.base import Embedder
|
|||||||
|
|
||||||
from config import LLM_CONFIG, EMBEDDING_CONFIG
|
from config import LLM_CONFIG, EMBEDDING_CONFIG
|
||||||
|
|
||||||
|
# 模块级单例客户端,避免每次调用都创建新连接
|
||||||
|
_async_client: AsyncOpenAI = None
|
||||||
|
|
||||||
|
|
||||||
|
def _get_async_client() -> AsyncOpenAI:
|
||||||
|
"""获取或创建 AsyncOpenAI 单例客户端"""
|
||||||
|
global _async_client
|
||||||
|
if _async_client is None:
|
||||||
|
_async_client = AsyncOpenAI(
|
||||||
|
api_key=LLM_CONFIG["api_key"],
|
||||||
|
base_url=LLM_CONFIG["base_url"],
|
||||||
|
)
|
||||||
|
return _async_client
|
||||||
class OpenaiAPI:
|
class OpenaiAPI:
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
@ -148,7 +160,79 @@ class OpenaiAPI:
|
|||||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}}
|
extra_body={"chat_template_kwargs": {"enable_thinking": False}}
|
||||||
)
|
)
|
||||||
return response.choices[0].message.content
|
return response.choices[0].message.content
|
||||||
|
@staticmethod
|
||||||
|
async def open_api_vl_without_thinking(
|
||||||
|
image_url: str
|
||||||
|
) -> str:
|
||||||
|
"""
|
||||||
|
专门用于从图片中提取【设备名称】和【故障现象】。
|
||||||
|
自动过滤思考过程,仅返回核心结果。
|
||||||
|
|
||||||
|
参数:
|
||||||
|
image_url: Base64 格式的图片数据 (data:image/...;base64,...)
|
||||||
|
model: 模型名称,默认为配置中的模型
|
||||||
|
custom_instruction: 额外的特定指令 (可选)
|
||||||
|
"""
|
||||||
|
|
||||||
|
# 1. 确定模型名称
|
||||||
|
model = LLM_CONFIG.get("model")
|
||||||
|
|
||||||
|
# 2. 构建强约束的 System Prompt
|
||||||
|
# 核心目标:禁止思考标签,禁止废话,只给结果
|
||||||
|
system_prompt = (
|
||||||
|
"你是一个视觉分析专家。你的任务是从图片中识别内容并进行描述。\n"
|
||||||
|
"【严格约束】\n"
|
||||||
|
"1. 绝对禁止输出 <think>, <think>, <reasoning> 等任何思考过程标签。\n"
|
||||||
|
"2. 绝对禁止输出'好的'、'根据图片'、'分析如下'等开场白或结束语。\n"
|
||||||
|
"3. 直接输出最终结论,格式必须严格遵守下面的模板。"
|
||||||
|
)
|
||||||
|
|
||||||
|
# 3. 构建针对性的 User Prompt
|
||||||
|
default_task = (
|
||||||
|
"请分析这张图片,描述图片内容,重点关注以下信息\n"
|
||||||
|
"1. 图片整体阐述了什么主题,描述了哪些内容\n"
|
||||||
|
"请用专业、客观、简短的语言描述。"
|
||||||
|
)
|
||||||
|
|
||||||
|
final_user_prompt = default_task
|
||||||
|
|
||||||
|
# 4. 初始化客户端
|
||||||
|
client = _get_async_client()
|
||||||
|
|
||||||
|
# 5. 构建请求参数
|
||||||
|
kwargs = {
|
||||||
|
"model": model,
|
||||||
|
"messages": [
|
||||||
|
{"role": "system", "content": system_prompt},
|
||||||
|
{
|
||||||
|
"role": "user",
|
||||||
|
"content": [
|
||||||
|
{"type": "text", "text": final_user_prompt},
|
||||||
|
{"type": "image_url", "image_url": {"url": image_url}}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"temperature": 0.1, # 低温度以保证事实准确性
|
||||||
|
"stream": False,
|
||||||
|
"max_tokens": LLM_CONFIG.get("max_tokens"),
|
||||||
|
}
|
||||||
|
|
||||||
|
# 尝试通过参数关闭思考 (取决于后端支持情况)
|
||||||
|
kwargs["extra_body"] = {"chat_template_kwargs": {"enable_thinking": False}}
|
||||||
|
|
||||||
|
try:
|
||||||
|
print(f"🚀 正在调用 {model} 进行图片内容分析...")
|
||||||
|
response = await client.chat.completions.create(**kwargs)
|
||||||
|
raw_content = response.choices[0].message.content or ""
|
||||||
|
|
||||||
|
print(f"✅ 分析完成:\n{raw_content}")
|
||||||
|
return raw_content
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
import traceback
|
||||||
|
error_msg = f"❌ 图片内容分析失败:{str(e)}\n{traceback.format_exc()}"
|
||||||
|
print(error_msg)
|
||||||
|
return f"Error: {str(e)}"
|
||||||
@staticmethod
|
@staticmethod
|
||||||
async def open_api_chat_async_json(query: str, model: str = None, temperature: float = None) -> str:
|
async def open_api_chat_async_json(query: str, model: str = None, temperature: float = None) -> str:
|
||||||
if model is None:
|
if model is None:
|
||||||
|
|||||||
@ -35,6 +35,8 @@ EMBED_MAX_WORKERS = SEARCH_CONFIG["embed_thread_num"] # embedding 批次并行
|
|||||||
NODE_FETCH_PAGE = SEARCH_CONFIG["node_fetch_page"] # 单次拉取待生成 embedding 节点的分页大小
|
NODE_FETCH_PAGE = SEARCH_CONFIG["node_fetch_page"] # 单次拉取待生成 embedding 节点的分页大小
|
||||||
|
|
||||||
# 全局索引名称(覆盖所有节点标签)
|
# 全局索引名称(覆盖所有节点标签)
|
||||||
|
GLOBAL_VECTOR_INDEX_NAME = SEARCH_CONFIG["global_entity_embedding"]
|
||||||
|
GLOBAL_FULLTEXT_INDEX_NAME = SEARCH_CONFIG["global_entity_content_search"]
|
||||||
DEFAULT_URI = NEO4J_CONFIG["uri"]
|
DEFAULT_URI = NEO4J_CONFIG["uri"]
|
||||||
DEFAULT_AUTH = (
|
DEFAULT_AUTH = (
|
||||||
NEO4J_CONFIG["username"],
|
NEO4J_CONFIG["username"],
|
||||||
@ -55,6 +57,8 @@ OLD_INDEX_NAMES: List[str] = [
|
|||||||
VECTOR_INDEX_NAME,
|
VECTOR_INDEX_NAME,
|
||||||
FULLTEXT_INDEX_NAME,
|
FULLTEXT_INDEX_NAME,
|
||||||
FULLTEXT_FIELD_INDEX_NAME,
|
FULLTEXT_FIELD_INDEX_NAME,
|
||||||
|
GLOBAL_VECTOR_INDEX_NAME,
|
||||||
|
GLOBAL_FULLTEXT_INDEX_NAME,
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
@ -188,6 +192,390 @@ def drop_index_without_constraint(driver):
|
|||||||
driver.execute_query(f"drop index {record['name']} if exists")
|
driver.execute_query(f"drop index {record['name']} if exists")
|
||||||
|
|
||||||
|
|
||||||
|
# --------- 创建全局混合索引(覆盖所有节点标签) ---------
|
||||||
|
def create_global_indexes(driver, force_refresh=False):
|
||||||
|
"""
|
||||||
|
创建覆盖所有节点标签的全局向量索引和全文索引。
|
||||||
|
与按标签创建的索引互补,支持跨标签统一检索。
|
||||||
|
|
||||||
|
Args:
|
||||||
|
driver: Neo4j driver
|
||||||
|
force_refresh: True=先删除旧索引再重建,False=IF NOT EXISTS 跳过已有索引
|
||||||
|
"""
|
||||||
|
if force_refresh:
|
||||||
|
for index_name in (GLOBAL_VECTOR_INDEX_NAME, GLOBAL_FULLTEXT_INDEX_NAME):
|
||||||
|
try:
|
||||||
|
driver.execute_query(f"DROP INDEX {index_name} IF EXISTS")
|
||||||
|
logger.info(f"强制刷新:已删除全局索引 {index_name}")
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"删除全局索引 {index_name} 时出错(可能不存在): {e}")
|
||||||
|
|
||||||
|
# 全局向量索引(作用于所有节点的 embedding 属性)
|
||||||
|
try:
|
||||||
|
driver.execute_query(f"""
|
||||||
|
CREATE VECTOR INDEX {GLOBAL_VECTOR_INDEX_NAME} IF NOT EXISTS
|
||||||
|
FOR (n) ON n.embedding
|
||||||
|
OPTIONS {{indexConfig: {{
|
||||||
|
`vector.dimensions`: {vector_dim},
|
||||||
|
`vector.similarity_function`: 'cosine'
|
||||||
|
}}}}
|
||||||
|
""")
|
||||||
|
logger.info(f"全局向量索引 '{GLOBAL_VECTOR_INDEX_NAME}' 已就绪")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"创建全局向量索引失败: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
# 全局全文索引(作用于所有节点的 名称 和 内容 属性)
|
||||||
|
try:
|
||||||
|
driver.execute_query(f"""
|
||||||
|
CREATE FULLTEXT INDEX {GLOBAL_FULLTEXT_INDEX_NAME} IF NOT EXISTS
|
||||||
|
FOR (n) ON EACH [n.名称, n.内容]
|
||||||
|
""")
|
||||||
|
logger.info(f"全局全文索引 '{GLOBAL_FULLTEXT_INDEX_NAME}' 已就绪")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"创建全局全文索引失败: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
|
||||||
|
# --------- 创建向量索引 ---------
|
||||||
|
def vector_indexing(driver, label, force_refresh=False):
|
||||||
|
"""
|
||||||
|
创建向量索引(如果节点没有 embedding,使用"名称"属性生成 embedding)
|
||||||
|
|
||||||
|
说明:
|
||||||
|
- 待生成 embedding 的节点通过分页循环全部处理,不再受单次 LIMIT 限制,
|
||||||
|
避免节点数超过一万时被静默丢弃。
|
||||||
|
- embedding 计算采用并行批次以缩短耗时。
|
||||||
|
|
||||||
|
Args:
|
||||||
|
driver: Neo4j driver
|
||||||
|
label: 节点标签
|
||||||
|
force_refresh: 是否强制刷新(如果索引已存在,force_refresh=False 时会跳过,True 时会先删除再创建)
|
||||||
|
"""
|
||||||
|
|
||||||
|
# === 分页处理所有缺失 embedding 的节点 ===
|
||||||
|
# 每轮取一页未生成 embedding 的节点并写回,直到没有剩余节点为止。
|
||||||
|
# 这样即使待处理节点远超一万也能全部覆盖。
|
||||||
|
total_generated = 0
|
||||||
|
while True:
|
||||||
|
query_no_embedding = f"""
|
||||||
|
MATCH (n:{label})
|
||||||
|
WHERE n.embedding IS NULL
|
||||||
|
AND n.名称 IS NOT NULL
|
||||||
|
RETURN elementId(n) AS id, n.名称 AS 名称
|
||||||
|
LIMIT {NODE_FETCH_PAGE}
|
||||||
|
"""
|
||||||
|
result_no_embedding = driver.execute_query(query_no_embedding)
|
||||||
|
nodes_to_process = [
|
||||||
|
(record["id"], record["名称"]) for record in result_no_embedding.records
|
||||||
|
]
|
||||||
|
|
||||||
|
if not nodes_to_process:
|
||||||
|
break
|
||||||
|
|
||||||
|
logger.info(f"标签 {label} 本轮发现 {len(nodes_to_process)} 个节点需要生成 embedding")
|
||||||
|
|
||||||
|
ids, names = zip(*nodes_to_process)
|
||||||
|
ids = list(ids)
|
||||||
|
names = list(names)
|
||||||
|
|
||||||
|
# 并行计算 embedding
|
||||||
|
logger.info(f"开始为 {len(names)} 个节点计算 embedding(并行批次)...")
|
||||||
|
try:
|
||||||
|
all_embeddings = _compute_embeddings_parallel(names)
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"批量嵌入向量计算失败: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
# 转换为 numpy 数组并做 L2 归一化
|
||||||
|
embeddings = np.array(all_embeddings)
|
||||||
|
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
|
||||||
|
norms = np.where(norms == 0, 1, norms) # 避免除零
|
||||||
|
embeddings = embeddings / norms
|
||||||
|
|
||||||
|
# 保存 embedding 到节点
|
||||||
|
logger.info(f"写入 {len(ids)} 个节点的 embedding...")
|
||||||
|
upsert_vectors(
|
||||||
|
driver,
|
||||||
|
ids=ids,
|
||||||
|
embedding_property="embedding",
|
||||||
|
embeddings=embeddings,
|
||||||
|
)
|
||||||
|
total_generated += len(ids)
|
||||||
|
logger.info(f"本轮成功为 {len(ids)} 个节点生成并保存 embedding")
|
||||||
|
|
||||||
|
# 不足一页说明已处理完,提前结束
|
||||||
|
if len(nodes_to_process) < NODE_FETCH_PAGE:
|
||||||
|
break
|
||||||
|
|
||||||
|
if total_generated:
|
||||||
|
logger.info(f"标签 {label} 共生成 {total_generated} 个节点的 embedding")
|
||||||
|
|
||||||
|
# 检查是否有节点存在 embedding(包括刚生成的)
|
||||||
|
check_query = f"MATCH (n:{label}) WHERE n.embedding IS NOT NULL RETURN count(n) AS count"
|
||||||
|
result = driver.execute_query(check_query)
|
||||||
|
node_count = result.records[0]["count"] if result.records else 0
|
||||||
|
|
||||||
|
if node_count == 0:
|
||||||
|
logger.warning(f"标签 {label} 没有节点包含 embedding 属性,跳过向量索引创建")
|
||||||
|
return
|
||||||
|
|
||||||
|
# 创建向量索引
|
||||||
|
index_name = f"{label.lower()}_vector"
|
||||||
|
|
||||||
|
# 如果强制刷新且索引已存在,先删除它
|
||||||
|
if force_refresh:
|
||||||
|
try:
|
||||||
|
driver.execute_query(f"DROP INDEX {index_name} IF EXISTS")
|
||||||
|
logger.info(f"强制刷新模式:已删除旧索引 {index_name}")
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"删除索引 {index_name} 时出错(可能不存在): {e}")
|
||||||
|
|
||||||
|
try:
|
||||||
|
create_vector_index(
|
||||||
|
driver,
|
||||||
|
name=index_name, # 索引的唯一名称
|
||||||
|
label=label, # 要索引的节点标签
|
||||||
|
embedding_property="embedding", # 包含嵌入向量值的节点属性键
|
||||||
|
dimensions=vector_dim, # 向量嵌入维度,1024与使用的 bge-m3 嵌入模型一致
|
||||||
|
similarity_fn="cosine", # 向量相似度函数,可选 "euclidean" 或 "cosine"
|
||||||
|
)
|
||||||
|
logger.info(f"成功为标签 {label} 创建向量索引(共 {node_count} 个节点)")
|
||||||
|
except Exception as e:
|
||||||
|
# 如果是增量更新模式且索引已存在,这是正常的,只记录警告
|
||||||
|
if not force_refresh and "already exists" in str(e).lower():
|
||||||
|
logger.info(f"标签 {label} 的向量索引已存在,跳过创建")
|
||||||
|
else:
|
||||||
|
logger.error(f"为标签 {label} 创建向量索引失败: {e}")
|
||||||
|
raise
|
||||||
|
|
||||||
|
|
||||||
|
# --------- 创建全文索引 ---------
|
||||||
|
def fulltext_indexing(driver, label, force_refresh=False):
|
||||||
|
"""
|
||||||
|
创建全文索引,并添加节点属性(自动获取所有属性)
|
||||||
|
|
||||||
|
全文索引:词 ---》 id
|
||||||
|
后续检索,分词后,根据词反向找到节点
|
||||||
|
|
||||||
|
Args:
|
||||||
|
driver: Neo4j driver
|
||||||
|
label: 节点标签
|
||||||
|
force_refresh: 是否强制刷新所有节点的全文索引(True=清空后重新生成,False=只处理缺失的节点)
|
||||||
|
"""
|
||||||
|
# 排除内部属性
|
||||||
|
excluded_props = {"embedding", "fulltext", "elementId", "id"}
|
||||||
|
|
||||||
|
# 根据 force_refresh 决定查询条件
|
||||||
|
if force_refresh:
|
||||||
|
# 强制刷新:处理所有节点(包括已有 fulltext 的节点)
|
||||||
|
query = f"""MATCH (n:{label})
|
||||||
|
WITH n, elementId(n) AS id, properties(n) AS props
|
||||||
|
RETURN id, props"""
|
||||||
|
else:
|
||||||
|
# 增量更新:只处理没有 fulltext 的节点
|
||||||
|
query = f"""MATCH (n:{label})
|
||||||
|
WHERE n.fulltext IS NULL
|
||||||
|
WITH n, elementId(n) AS id, properties(n) AS props
|
||||||
|
RETURN id, props"""
|
||||||
|
|
||||||
|
# 执行查询
|
||||||
|
records = driver.execute_query(query).records
|
||||||
|
|
||||||
|
if not records:
|
||||||
|
if force_refresh:
|
||||||
|
logger.info(f"{label} 没有节点需要处理")
|
||||||
|
else:
|
||||||
|
logger.info(f"{label} 所有节点皆存在全文索引属性")
|
||||||
|
return
|
||||||
|
|
||||||
|
# 构建文本描述(使用所有属性)
|
||||||
|
record_list = []
|
||||||
|
for r in records:
|
||||||
|
node_id = r["id"]
|
||||||
|
props = r["props"]
|
||||||
|
# 过滤掉内部属性
|
||||||
|
filtered_props = {k: v for k, v in props.items() if k not in excluded_props}
|
||||||
|
# 使用所有属性构建文本
|
||||||
|
text = build_text_from_all_properties(filtered_props)
|
||||||
|
if text:
|
||||||
|
record_list.append({"id": node_id, "text": text})
|
||||||
|
|
||||||
|
if not record_list:
|
||||||
|
logger.info(f"{label} 没有可索引的文本内容")
|
||||||
|
return
|
||||||
|
|
||||||
|
# 创建全文索引(在写入数据之前创建索引)
|
||||||
|
try:
|
||||||
|
create_fulltext_index(
|
||||||
|
driver,
|
||||||
|
name=f"{label.lower()}_fulltext", # 索引的唯一名称
|
||||||
|
label=label, # 要创建索引的节点标签
|
||||||
|
node_properties=["fulltext"], # 要创建全文索引的节点属性列表
|
||||||
|
)
|
||||||
|
except Exception as e:
|
||||||
|
logger.warning(f"创建全文索引可能已存在(继续处理): {e}")
|
||||||
|
|
||||||
|
# 文本分词,作为全文索引属性
|
||||||
|
logger.info(f"计算 {label} ({len(record_list)}) 的全文索引")
|
||||||
|
pattern = re.compile(r"[a-zA-Z0-9\u4e00-\u9fa5]+") # 匹配英文字母、数字和中文字符
|
||||||
|
fulltext_tuple_list = []
|
||||||
|
for id_, text in [(r["id"], r["text"]) for r in record_list]:
|
||||||
|
# 使用 jieba 分词
|
||||||
|
words = jieba.lcut(text)
|
||||||
|
# 过滤并连接
|
||||||
|
filtered_words = [
|
||||||
|
word.strip()
|
||||||
|
for word in words
|
||||||
|
if word.strip() and pattern.fullmatch(word.strip())
|
||||||
|
]
|
||||||
|
fulltext_value = " ".join(filtered_words)
|
||||||
|
if fulltext_value:
|
||||||
|
fulltext_tuple_list.append((id_, fulltext_value))
|
||||||
|
|
||||||
|
if not fulltext_tuple_list:
|
||||||
|
logger.warning(f"{label} 分词后没有有效内容")
|
||||||
|
return
|
||||||
|
|
||||||
|
ids, fulltexts = zip(*fulltext_tuple_list)
|
||||||
|
ids = list(ids)
|
||||||
|
fulltexts = list(fulltexts)
|
||||||
|
|
||||||
|
# 按 elementId 添加全文索引属性
|
||||||
|
logger.info(f"写入 {label} ({len(fulltexts)}) 的全文索引")
|
||||||
|
insert_batch_size = 1000
|
||||||
|
for i in range(0, len(fulltext_tuple_list), insert_batch_size):
|
||||||
|
batch_rows = [
|
||||||
|
{"id": id_, "fulltext": ft}
|
||||||
|
for id_, ft in zip(
|
||||||
|
ids[i: i + insert_batch_size],
|
||||||
|
fulltexts[i: i + insert_batch_size],
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
# UNWIND:将列表数据展开为多行记录
|
||||||
|
driver.execute_query(
|
||||||
|
"UNWIND $rows AS row " # 将传入的 rows 列表展开,每项作为一行数据,命名为 row
|
||||||
|
"MATCH (n) "
|
||||||
|
"WHERE elementId(n) = row.id "
|
||||||
|
"SET n.fulltext = row.fulltext ",
|
||||||
|
{"rows": batch_rows},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def create_all_indexes(neo4j_uri=None, neo4j_username=None, neo4j_password=None, force_refresh=False, driver=None):
|
||||||
|
"""
|
||||||
|
创建所有索引的主函数(可被外部调用)
|
||||||
|
|
||||||
|
注意:Neo4j 连接配置优先从 .env 文件中读取(NEO4J_URI, NEO4J_USERNAME, NEO4J_PASSWORD)
|
||||||
|
如果 .env 文件中没有配置,则使用参数中的值(如果提供)或默认值
|
||||||
|
|
||||||
|
Args:
|
||||||
|
neo4j_uri: Neo4j URI(可选,优先级低于 .env,提供 driver 时忽略)
|
||||||
|
neo4j_username: Neo4j 用户名(可选,优先级低于 .env,提供 driver 时忽略)
|
||||||
|
neo4j_password: Neo4j 密码(可选,优先级低于 .env,提供 driver 时忽略)
|
||||||
|
force_refresh: 是否强制刷新所有索引(True=清空重建,False=增量更新)
|
||||||
|
driver: 可选的 Neo4j driver 实例,提供则复用外部 driver
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
dict: 包含执行结果的字典
|
||||||
|
"""
|
||||||
|
use_external_driver = driver is not None
|
||||||
|
|
||||||
|
if not use_external_driver:
|
||||||
|
# 优先从 .env 文件读取配置,如果 .env 中没有则使用参数或默认值
|
||||||
|
neo4j_uri = os.getenv("NEO4J_URI") or neo4j_uri or "bolt://192.168.0.46:57687"
|
||||||
|
neo4j_username = os.getenv("NEO4J_USERNAME") or neo4j_username or "neo4j"
|
||||||
|
neo4j_password = os.getenv("NEO4J_PASSWORD") or neo4j_password or "zdht123@"
|
||||||
|
logger.info(f"连接到 Neo4j: {neo4j_uri}")
|
||||||
|
else:
|
||||||
|
logger.info("使用外部提供的 Neo4j driver")
|
||||||
|
|
||||||
|
# 定义内部函数来执行索引创建逻辑(避免代码重复)
|
||||||
|
def _execute_indexing(driver):
|
||||||
|
# 1、根据 force_refresh 决定是否清空索引和约束
|
||||||
|
if force_refresh:
|
||||||
|
logger.info("强制刷新模式:清空所有约束...")
|
||||||
|
drop_constraint(driver)
|
||||||
|
logger.info("强制刷新模式:清空所有索引...")
|
||||||
|
drop_index_without_constraint(driver)
|
||||||
|
logger.info("强制刷新模式:清空所有节点的 fulltext 属性...")
|
||||||
|
driver.execute_query("MATCH (n) WHERE n.fulltext IS NOT NULL REMOVE n.fulltext")
|
||||||
|
else:
|
||||||
|
logger.info("增量更新模式:保留现有索引和约束")
|
||||||
|
|
||||||
|
# 2、自动获取所有节点标签并创建向量索引
|
||||||
|
logger.info("=" * 50)
|
||||||
|
logger.info("开始创建向量索引...")
|
||||||
|
logger.info("=" * 50)
|
||||||
|
|
||||||
|
all_labels = get_all_labels(driver)
|
||||||
|
logger.info(f"发现 {len(all_labels)} 个节点标签: {all_labels}")
|
||||||
|
|
||||||
|
vector_index_errors = []
|
||||||
|
for label in all_labels:
|
||||||
|
logger.info(f"处理标签: {label}")
|
||||||
|
try:
|
||||||
|
vector_indexing(driver, label, force_refresh=force_refresh)
|
||||||
|
except Exception as e:
|
||||||
|
error_msg = f"为标签 {label} 创建向量索引时出错: {e}"
|
||||||
|
logger.error(error_msg)
|
||||||
|
vector_index_errors.append(error_msg)
|
||||||
|
continue
|
||||||
|
|
||||||
|
# 3、创建全文索引(自动获取所有属性)
|
||||||
|
logger.info("=" * 50)
|
||||||
|
logger.info("开始创建全文索引...")
|
||||||
|
logger.info("=" * 50)
|
||||||
|
|
||||||
|
fulltext_index_errors = []
|
||||||
|
for label in all_labels:
|
||||||
|
logger.info(f"处理标签: {label}")
|
||||||
|
try:
|
||||||
|
fulltext_indexing(driver, label, force_refresh=force_refresh)
|
||||||
|
except Exception as e:
|
||||||
|
error_msg = f"为标签 {label} 创建全文索引时出错: {e}"
|
||||||
|
logger.error(error_msg)
|
||||||
|
fulltext_index_errors.append(error_msg)
|
||||||
|
continue
|
||||||
|
|
||||||
|
logger.info("=" * 50)
|
||||||
|
logger.info("所有索引创建完成!")
|
||||||
|
logger.info("=" * 50)
|
||||||
|
|
||||||
|
has_errors = len(vector_index_errors) > 0 or len(fulltext_index_errors) > 0
|
||||||
|
mode_text = "强制刷新" if force_refresh else "增量更新"
|
||||||
|
return {
|
||||||
|
"success": not has_errors,
|
||||||
|
"message": f"{mode_text}模式:索引创建完成" if not has_errors
|
||||||
|
else f"{mode_text}模式:索引创建完成,但部分标签出现错误",
|
||||||
|
"mode": "force_refresh" if force_refresh else "incremental",
|
||||||
|
"labels": all_labels,
|
||||||
|
"vector_index_errors": vector_index_errors,
|
||||||
|
"fulltext_index_errors": fulltext_index_errors,
|
||||||
|
}
|
||||||
|
|
||||||
|
try:
|
||||||
|
if use_external_driver:
|
||||||
|
# 使用外部 driver,直接执行,不关闭 driver
|
||||||
|
return _execute_indexing(driver)
|
||||||
|
else:
|
||||||
|
# 使用内部 driver,使用上下文管理器自动关闭
|
||||||
|
with GraphDatabase.driver(neo4j_uri, auth=(neo4j_username, neo4j_password)) as driver:
|
||||||
|
return _execute_indexing(driver)
|
||||||
|
except Exception as e:
|
||||||
|
error_msg = f"创建索引过程中发生错误: {e}"
|
||||||
|
logger.error(error_msg, exc_info=True)
|
||||||
|
return {
|
||||||
|
"success": False,
|
||||||
|
"message": error_msg,
|
||||||
|
"mode": "error",
|
||||||
|
"labels": [],
|
||||||
|
"vector_index_errors": [],
|
||||||
|
"fulltext_index_errors": [],
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
def print_neo4j_version(driver) -> None:
|
def print_neo4j_version(driver) -> None:
|
||||||
"""打印当前 Neo4j 实例的版本和版本类型。"""
|
"""打印当前 Neo4j 实例的版本和版本类型。"""
|
||||||
with driver.session() as session:
|
with driver.session() as session:
|
||||||
@ -407,12 +795,13 @@ build_hyrid_indexes = build_hybrid_indexes
|
|||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
# 可以通过命令行参数或环境变量控制是否强制刷新
|
||||||
|
import sys
|
||||||
|
|
||||||
result = build_hybrid_indexes(
|
force_refresh = "--force" in sys.argv or os.getenv("FORCE_REFRESH", "false").lower() == "true"
|
||||||
drop_old=True
|
|
||||||
)
|
|
||||||
|
|
||||||
|
result = create_all_indexes(force_refresh=force_refresh)
|
||||||
if result["success"]:
|
if result["success"]:
|
||||||
logger.info("混合索引创建成功")
|
logger.info(f"索引创建成功完成(模式: {result['mode']})")
|
||||||
else:
|
else:
|
||||||
logger.error(result["message"])
|
logger.error(f"索引创建完成但有错误: {result['message']}")
|
||||||
|
|||||||
184
resetlevel.py
Normal file
184
resetlevel.py
Normal file
@ -0,0 +1,184 @@
|
|||||||
|
import copy
|
||||||
|
import json
|
||||||
|
import re
|
||||||
|
import sys
|
||||||
|
from typing import Any
|
||||||
|
|
||||||
|
|
||||||
|
CN_NUM = "\u4e00\u4e8c\u4e09\u56db\u4e94\u516d\u4e03\u516b\u4e5d\u5341\u767e\u5343\u4e07\u3007\u96f6\u4e24"
|
||||||
|
|
||||||
|
CHAPTER_PATTERN = re.compile(rf"^\s*\u7b2c\s*([{CN_NUM}\d]+)\s*\u7ae0")
|
||||||
|
CN_LEVEL1_PATTERN = re.compile(rf"^\s*([{CN_NUM}]+)(?=\s|[\u3001\u3002\uff0e.])")
|
||||||
|
|
||||||
|
CASE_LIKE_PATTERN = re.compile(r"^\s*\d+(?:\.|\uff0e)(?!\d|\s|$)")
|
||||||
|
ARABIC_PATTERN = re.compile(
|
||||||
|
r"^\s*(\d+(?:\.\d+)*)(?=\s|[\u4e00-\u9fff\u3001\u3002\uff0e]|$|\.(?:\s|$))"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def has_valid_text_level(item: dict) -> bool:
|
||||||
|
return "text_level" in item and item.get("text_level") not in (None, "")
|
||||||
|
|
||||||
|
|
||||||
|
def parse_heading(text: str):
|
||||||
|
text = text.strip()
|
||||||
|
if not text:
|
||||||
|
return None, None
|
||||||
|
|
||||||
|
if CASE_LIKE_PATTERN.match(text):
|
||||||
|
return "case_like", None
|
||||||
|
|
||||||
|
if CHAPTER_PATTERN.match(text):
|
||||||
|
return "level1", None
|
||||||
|
|
||||||
|
if CN_LEVEL1_PATTERN.match(text):
|
||||||
|
return "level1", None
|
||||||
|
|
||||||
|
match = ARABIC_PATTERN.match(text)
|
||||||
|
if match:
|
||||||
|
return "arabic", match.group(1)
|
||||||
|
|
||||||
|
return None, None
|
||||||
|
|
||||||
|
|
||||||
|
def arabic_depth(number: str) -> int:
|
||||||
|
return number.count(".") + 1
|
||||||
|
|
||||||
|
|
||||||
|
def parent_number(number: str) -> str | None:
|
||||||
|
parts = number.split(".")
|
||||||
|
if len(parts) <= 1:
|
||||||
|
return None
|
||||||
|
return ".".join(parts[:-1])
|
||||||
|
|
||||||
|
|
||||||
|
def collect_headings(data: list) -> list:
|
||||||
|
headings = []
|
||||||
|
|
||||||
|
for item in data:
|
||||||
|
if not isinstance(item, dict):
|
||||||
|
continue
|
||||||
|
if item.get("type") != "text":
|
||||||
|
continue
|
||||||
|
if not has_valid_text_level(item):
|
||||||
|
continue
|
||||||
|
|
||||||
|
kind, number = parse_heading(item.get("text", ""))
|
||||||
|
if kind in {"level1", "arabic", "case_like"}:
|
||||||
|
headings.append({"kind": kind, "number": number})
|
||||||
|
|
||||||
|
return headings
|
||||||
|
|
||||||
|
|
||||||
|
def is_valid_arabic_hierarchy(numbers: list[str], allow_missing_depth1_parent: bool) -> bool:
|
||||||
|
seen = set()
|
||||||
|
has_depth2 = False
|
||||||
|
has_depth3 = False
|
||||||
|
|
||||||
|
for number in numbers:
|
||||||
|
depth = arabic_depth(number)
|
||||||
|
if depth >= 2:
|
||||||
|
has_depth2 = True
|
||||||
|
if depth >= 3:
|
||||||
|
has_depth3 = True
|
||||||
|
|
||||||
|
parent = parent_number(number)
|
||||||
|
if parent is not None:
|
||||||
|
parent_depth = arabic_depth(parent)
|
||||||
|
parent_is_missing_depth1 = allow_missing_depth1_parent and parent_depth == 1
|
||||||
|
if parent not in seen and not parent_is_missing_depth1:
|
||||||
|
return False
|
||||||
|
|
||||||
|
seen.add(number)
|
||||||
|
|
||||||
|
return has_depth2 and has_depth3
|
||||||
|
|
||||||
|
|
||||||
|
def detect_valid_mode(headings: list):
|
||||||
|
if not headings:
|
||||||
|
return None
|
||||||
|
|
||||||
|
if any(heading["kind"] == "case_like" for heading in headings):
|
||||||
|
return None
|
||||||
|
|
||||||
|
first = headings[0]
|
||||||
|
|
||||||
|
if first["kind"] == "arabic":
|
||||||
|
numbers = [heading["number"] for heading in headings if heading["kind"] == "arabic"]
|
||||||
|
if numbers and arabic_depth(numbers[0]) == 1:
|
||||||
|
if is_valid_arabic_hierarchy(numbers, allow_missing_depth1_parent=False):
|
||||||
|
return "A"
|
||||||
|
return None
|
||||||
|
|
||||||
|
if first["kind"] == "level1":
|
||||||
|
numbers = [heading["number"] for heading in headings[1:] if heading["kind"] == "arabic"]
|
||||||
|
if not numbers:
|
||||||
|
return None
|
||||||
|
|
||||||
|
first_depth = arabic_depth(numbers[0])
|
||||||
|
if first_depth == 2 and is_valid_arabic_hierarchy(numbers, allow_missing_depth1_parent=True):
|
||||||
|
return "B"
|
||||||
|
|
||||||
|
if first_depth == 1 and is_valid_arabic_hierarchy(numbers, allow_missing_depth1_parent=False):
|
||||||
|
return "C"
|
||||||
|
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
JsonContent = Any
|
||||||
|
|
||||||
|
|
||||||
|
def reset_textlevel(json_content: JsonContent) -> JsonContent:
|
||||||
|
"""
|
||||||
|
输入 JSON 内容,返回更新 text_level 后的新 JSON 内容。
|
||||||
|
|
||||||
|
json_content 可以是已经 json.load/json.loads 后的 Python 对象,
|
||||||
|
也可以是 JSON 字符串。函数不会修改原始入参。
|
||||||
|
"""
|
||||||
|
if isinstance(json_content, str):
|
||||||
|
data = json.loads(json_content)
|
||||||
|
else:
|
||||||
|
data = copy.deepcopy(json_content)
|
||||||
|
|
||||||
|
if not isinstance(data, list):
|
||||||
|
return data
|
||||||
|
|
||||||
|
headings = collect_headings(data)
|
||||||
|
mode = detect_valid_mode(headings)
|
||||||
|
|
||||||
|
if mode is None:
|
||||||
|
return data
|
||||||
|
|
||||||
|
for item in data:
|
||||||
|
if not isinstance(item, dict):
|
||||||
|
continue
|
||||||
|
if item.get("type") != "text":
|
||||||
|
continue
|
||||||
|
if not has_valid_text_level(item):
|
||||||
|
continue
|
||||||
|
|
||||||
|
kind, number = parse_heading(item.get("text", ""))
|
||||||
|
if kind == "level1":
|
||||||
|
item["text_level"] = 1
|
||||||
|
continue
|
||||||
|
if kind != "arabic":
|
||||||
|
continue
|
||||||
|
|
||||||
|
depth = arabic_depth(number)
|
||||||
|
item["text_level"] = depth if mode in {"A", "B"} else depth + 1
|
||||||
|
|
||||||
|
return data
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
json_file_path = r'E:\ZKYNLP\Hjunproject\project0506\kgrag\163-06A0014-B01001_发动机-维修手册_content_list.json'
|
||||||
|
with open(json_file_path, "r", encoding="utf-8") as f:
|
||||||
|
json_content = json.load(f)
|
||||||
|
result= reset_textlevel(json_content)
|
||||||
|
|
||||||
|
for ins in result[:30]:
|
||||||
|
print(ins)
|
||||||
|
# print(json.dumps(result, ensure_ascii=False, indent=2))
|
||||||
3514
testfiles/output/122-06A0014-B01001_发动机-维修手册测试_content_list.json
Normal file
3514
testfiles/output/122-06A0014-B01001_发动机-维修手册测试_content_list.json
Normal file
File diff suppressed because one or more lines are too long
3802
testfiles/output/163-06A0014-B01002_螺旋桨-操作使用手册_content_list.json
Normal file
3802
testfiles/output/163-06A0014-B01002_螺旋桨-操作使用手册_content_list.json
Normal file
File diff suppressed because one or more lines are too long
4434
testfiles/output/163-06A0015-B01001_发动机-操作使用手册_content_list.json
Normal file
4434
testfiles/output/163-06A0015-B01001_发动机-操作使用手册_content_list.json
Normal file
File diff suppressed because one or more lines are too long
3591
testfiles/output/163-06A0015-B01001_发动机-维修手册_content_list.json
Normal file
3591
testfiles/output/163-06A0015-B01001_发动机-维修手册_content_list.json
Normal file
File diff suppressed because one or more lines are too long
419
testfiles/output/163-06A0016-B01003_雷达-图解目录_content_list.json
Normal file
419
testfiles/output/163-06A0016-B01003_雷达-图解目录_content_list.json
Normal file
File diff suppressed because one or more lines are too long
Loading…
x
Reference in New Issue
Block a user