from filename_proceess_and_kgquery import build_graph_from_paths, fetch_graph_sample, node_to_json, rel_to_json from graph_search.operation_graph_search import operation_graph_search from graph_search.fault_graph_search import fault_graph_reason_and_projects from contextlib import asynccontextmanager from langchain_community.graphs import Neo4jGraph from neo4j import GraphDatabase from neo4j.exceptions import Neo4jError from fastapi.concurrency import run_in_threadpool from io import BytesIO import aiofiles import hashlib import logging import random import os import re import tempfile import mimetypes import time import asyncio from fastapi import FastAPI, File, Path, UploadFile, HTTPException, Form, Request, Header, Body, BackgroundTasks from fastapi.responses import JSONResponse, StreamingResponse from modelsAPI.model_api import OpenaiAPI from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field from dotenv import load_dotenv from pathlib import PurePath, Path as PathLib import httpx import pdfplumber from typing import List, Optional, Dict from collections import defaultdict from fuzzywuzzy import fuzz, process from collections import defaultdict, Counter from kg_build.filter_node_relation import filter_node_relationship from kg_build.kg_save import MilitaryKnowledgeGraph from kgbuild.table_html_extract import extract_operation_tables, extract_repair_tables_keep_first, extract_tables from kgbuild.entity_rela_reset import create_extraction_prompt from filename_proceess_and_kgquery import get_entity,build_tree from typing import List, Dict, Optional, Union, Any, Tuple from pydantic import BaseModel from dataprocess import _split_markdown import uuid import json from fileparse_util import process_document # import hanlp from itertools import groupby import requests from datetime import datetime from excel_process import transform_excel_json, read_excel_with_merged_cells, dataframe_to_row_json, json_to_markdown import traceback import signal import threading from multiprocessing import Process, Queue import asyncio from fastapi.staticfiles import StaticFiles from extract_text_node_relation import get_md_node_relation, check_filename_ocr, process_file_content from concurrent.futures import ThreadPoolExecutor import base64 from extract_excel_node_relation import get_excel_node_relation 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 filename_proceess_and_kgquery import get_entity from kg_build.extract_filtertext import get_filtertext_node_relation import os import logging import sys import asyncio from time import sleep from graph_search.Hybrid_graphsearch import ( retrieve_entities, query_with_cypher, fetch_relations_among_results, format_relations, LocalBgeM3Embeddings, ) from neo4j import AsyncGraphDatabase, GraphDatabase from graph_search.cypher_excute import excute_cypher if sys.platform == "win32": asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) from contextlib import asynccontextmanager import mimetypes from default_ontology_config import ( ONTOLOGY as ontology, ONTOLOGY_DESCRIPTION as Ontology_DESCRIPTION, NODE_TYPES as node_types, RELATIONSHIPS as relationships, RELATIONSHIP_TYPES as relationship_types, # 也可以导入便捷函数 get_ontology_config, get_ontology_list, get_ontology_description, get_node_types, get_relationships, get_relationship_types, ) from config import NEO4J_CONFIG,URL,API_URLS,API_OTHER_URLS converter = Doc2PDF() DATA_DIR = "/app/files" # ================== 全局资源容器 ================== # 用一个字典持有重资源，避免每次请求都重建 _resources: Dict[str, Any] = { "sync_driver": None, "async_driver": None, "embedder": None, } # 图谱检索和索引生成 try: from graph_search.graph_service import GraphService from neo4j_indexing import create_all_indexes, create_global_indexes, build_hybrid_indexes INDEXING_AVAILABLE = True GRAPH_AVAILABLE = True except ImportError as e: INDEXING_AVAILABLE = False GRAPH_AVAILABLE = False kg_running_tasks: Dict[str, dict] = {} _neo4j_indexing_lock = asyncio.Lock() def setup_logging(): fmt = "%(asctime)s - %(levelname)s - %(name)s - %(message)s" handlers = [ logging.StreamHandler(), logging.FileHandler("app.log", encoding="utf-8"), ] logging.basicConfig(level=logging.INFO, format=fmt, handlers=handlers, force=True) logging.disable(logging.NOTSET) for name in list(logging.root.manager.loggerDict.keys()): lg = logging.getLogger(name) lg.disabled = False lg.setLevel(logging.NOTSET) # ← 加这行，强制清除所有子 logger 的独立 level for name in ("uvicorn", "uvicorn.error", "uvicorn.access"): uv_logger = logging.getLogger(name) uv_logger.handlers = list(handlers) uv_logger.propagate = False @asynccontextmanager async def lifespan(app): setup_logging() """应用启动时初始化驱动和 embedder，关闭时释放。""" logger.info("[lifespan] 初始化 Neo4j 驱动与 embedder ...") # 同步驱动（HybridCypherRetriever 用） _resources["sync_driver"] = GraphDatabase.driver(NEO4J_URI, auth=("neo4j", "zdht123@")) # 异步驱动（fulltext 检索 + Cypher 执行用） _resources["async_driver"] = AsyncGraphDatabase.driver(NEO4J_URI, auth=("neo4j", "zdht123@")) # embedder _resources["embedder"] = LocalBgeM3Embeddings( base_url=os.getenv("LOCAL_EMBEDDING_BASE"), api_key=os.getenv("LOCAL_API_KEY"), ) logger.info("[lifespan] 初始化完成") try: yield finally: logger.info("[lifespan] 释放资源 ...") if _resources["sync_driver"] is not None: _resources["sync_driver"].close() if _resources["async_driver"] is not None: await _resources["async_driver"].close() logger.info("[lifespan] 资源已释放") yield app = FastAPI(max_request_size=1024 * 1024 * 10, lifespan=lifespan) setup_logging() logger = logging.getLogger(__name__) app.add_middleware( CORSMiddleware, allow_origins=["*"], # Allow all origins allow_credentials=True, allow_methods=["*"], # Allow all HTTP methods allow_headers=["*"], # Allow all headers ) app.mount("/upload", StaticFiles(directory="/app/files/uploads_v1"), name="upload") app.mount("/files", StaticFiles(directory="/app/files"), name="files") """ 图查询接口 """ load_dotenv() # ===== 本地 Neo4j 强制配置 ===== # NEO4J_URI = os.getenv("NEO4J_URI", "bolt://localhost:7687") # NEO4J_URI = os.getenv("NEO4J_URI", "bolt://localhost:7687") NEO4J_URI = NEO4J_CONFIG["uri"] NEO4J_USER = NEO4J_CONFIG["username"] NEO4J_PASSWORD = NEO4J_CONFIG["password"] # 不设默认值，强制要求提供 NEO4J_DATABASE =NEO4J_CONFIG["database"] indexing_url = URL['indexing_url'] # 根据实际部署调整 PREFIX_URL = URL['prefix_url'] API_URLS = API_URLS API_other_URLS = API_OTHER_URLS _api_url_inflight = [0] * len(API_URLS) _api_url_lock = threading.Lock() def get_api_url(task_id: str = "") -> tuple: """返回 (analyze_url, contentlist_url)，选当前 inflight 最少的实例。""" with _api_url_lock: idx = _api_url_inflight.index(min(_api_url_inflight)) _api_url_inflight[idx] += 1 base = API_URLS[idx] return base, base.replace("/analyze-pdf", "/analyze-otherfile") def release_api_url(url: str): with _api_url_lock: if url in API_URLS: idx = API_URLS.index(url) _api_url_inflight[idx] = max(0, _api_url_inflight[idx] - 1) # 可选：如果密码未设置，抛出明确错误 if NEO4J_PASSWORD is None: raise ValueError("Environment variable 'NEO4J_PASSWORD' is required.") driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD), database=NEO4J_DATABASE) DEFAULT_ONTOLOGY: list = ontology DEFAULT_ONTOLOGY_DESCRIPTION: dict = Ontology_DESCRIPTION DEFAULT_NODE_TYPES: dict = node_types DEFAULT_RELATIONSHIPS: list = relationships DEFAULT_RELATIONSHIP_TYPES: dict = relationship_types # ============================================================ # 2. 合并工具函数 # 规则： # - dict → 以默认值为底，用户值 key 已存在则合并 value list； # key 不存在则直接新增 # ============================================================ def _merge_list(default: list, user: list | None) -> list: """去重合并两个列表，保持默认顺序，用户新增项追加到末尾。""" if not user: return list(default) merged = list(default) for item in user: if item not in merged: merged.append(item) return merged def _merge_dict_of_lists(default: dict, user: dict | None) -> dict: """ 合并两个 {key: list} 结构的字典。 - key 存在于 default：合并 list（去重），用户值追加到末尾 - key 只存在于 user：直接新增 """ if not user: return {k: list(v) for k, v in default.items()} merged = {k: list(v) for k, v in default.items()} for key, val in user.items(): if key in merged: for item in val: if item not in merged[key]: merged[key].append(item) else: merged[key] = list(val) return merged def _load_registry_config(): """每次调用都从 entity_registry.json 读取最新注册的实体类型，无需重启生效。""" try: import importlib import entity_registry as _reg importlib.reload(_reg) # 强制重新读文件 return _reg.get_ontology_list(), _reg.get_ontology_description(), _reg.get_node_types() except Exception: return [], {}, {} def build_kg_config( user_ontology: list | None, user_ontology_desc: dict | None, user_relationships: list | None, user_node_types: dict | None, user_relationship_types: dict | None, ) -> dict: """ 将用户传入的参数与全局默认值合并，返回最终执行配置。 任何字段未传（None）时直接使用默认值；传了则合并。 registry 中的新类型每次请求动态合并，无需重启。 """ reg_list, reg_desc, reg_nodes = _load_registry_config() ontology = _merge_list(DEFAULT_ONTOLOGY, reg_list) ontology_desc = {**reg_desc, **DEFAULT_ONTOLOGY_DESCRIPTION} # 手写配置优先 node_types = {**reg_nodes, **DEFAULT_NODE_TYPES} # 手写配置优先 return { "ontology": _merge_list(ontology, user_ontology), "ontology_description": _merge_dict_of_lists(ontology_desc, user_ontology_desc), "relationships": _merge_list(DEFAULT_RELATIONSHIPS, user_relationships), "node_types": _merge_dict_of_lists(node_types, user_node_types), "relationship_types": _merge_dict_of_lists(DEFAULT_RELATIONSHIP_TYPES, user_relationship_types), } # ============================================================ # 3. 修改后的 Pydantic 模型 # 所有字段改为 Optional（默认 None），表示"用户未传" # 真正的默认值由 build_kg_config() 在运行时合并产生 # ============================================================ class FileContent(BaseModel): """单个文件""" file_name: str # 修改：改为可选，默认为空字符串 markdown_content: str = Field(default="", description="Markdown格式的内容") # 修改：改为可选，默认为空列表 # 注意：List[Dict[str, Any]] 不能直接赋值 default=[]，需要用 Field 或 default_factory 防止引用共享 content_middle: List[Dict[str, Any]] = Field(default_factory=list, description="中间处理结果") file_base64: str = Field(default="", description="文件内容的Base64编码字符串", validate_default=True) class KGRequest(BaseModel): task_id: str files: Union[FileContent, List[FileContent]] # 全部改为 Optional[...] = None，不再在模型层携带默认值 ontology: Optional[List[str]] = None Ontology_DESCRIPTION: Optional[Dict[str, List[str]]] = None relationships: Optional[List[str]] = None node_types: Optional[Dict[str, List[str]]] = None relationship_types: Optional[Dict[str, List[str]]] = None class FileContentReset(BaseModel): """单个文件""" file_name: str # 保留文件名（可用于日志或调试） markdown_content: str class resetRequest(BaseModel): files: Union[FileContentReset, List[FileContentReset]] ontology: Optional[List[str]] = None Ontology_DESCRIPTION: Optional[Dict[str, List[str]]] = None relationships: Optional[List[str]] = None node_types: Optional[Dict[str, List[str]]] = None relationship_types: Optional[Dict[str, List[str]]] = None def clear_database(): graph = Neo4jGraph( refresh_schema=False, url=NEO4J_URI, username=NEO4J_USER, password=NEO4J_PASSWORD, database=NEO4J_DATABASE ) # 清空数据库内容 graph.query("MATCH (n) DETACH DELETE n") # ========================= # 接口：head / tail / 查询，并返回最短路径。node+depth查询 # ======================== def add_node_style(node): node_type = node.get("labels", []) if node_type and len(node_type) > 0: node_type = node_type[0] else: node_type = "" if node_type == "舰艇": node["color"] = "#F89A64" node["size"] = 86 elif node_type == "系统": node["color"] = "#C25C21" node["size"] = 86 elif node_type == "子系统": node["color"] = "#7D2D00" node["size"] = 86 elif node_type == "设备": node["color"] = "#540702C1" node["size"] = 86 elif node_type in [ "功能", "环境条件", "接口", "安全警告", "调试", "器材保障", "图册", "设计单位", "维修项目", "操作项目", ]: node["color"] = "#2261DF" node["size"] = 52 else: node["color"] = "#68C3FF" node["size"] = 53 return node def add_link_style(link, nodes_dict): source_id = link.get("source") target_id = link.get("target") source_type = "" target_type = "" if source_id and source_id in nodes_dict: source_labels = nodes_dict[source_id].get("labels", []) if source_labels and len(source_labels) > 0: source_type = source_labels[0] if target_id and target_id in nodes_dict: target_labels = nodes_dict[target_id].get("labels", []) if target_labels and len(target_labels) > 0: target_type = target_labels[0] hierarchy_types = ["舰艇", "系统", "子系统", "设备"] if source_type in hierarchy_types and target_type in hierarchy_types: link["color"] = "#F59E0B" else: link["color"] = "#FFFFFF" return link def apply_styles(graph_data): nodes = graph_data.get("nodes", []) links = graph_data.get("links", []) nodes_dict = {n.get("id"): n for n in nodes if n.get("id")} for node in nodes: add_node_style(node) for link in links: add_link_style(link, nodes_dict) return graph_data class GraphQueryRequest(BaseModel): filename: Optional[str] = None queryview: Optional[str] = None querylink: Optional[str] = None head: Optional[str] = None tail: Optional[str] = None node: Optional[str] = None source_node: Optional[str] = None # 在这里设置默认值为 1 depth: Optional[int] = 1 simple_graphData_copy: Optional[Dict[str, Any]] = None @app.post("/api/v2/neo4j/query") async def graph_query(request: GraphQueryRequest): """ 返回结构： { "message": str, "simple_graphData": {...}, # 200节点采样（带 level 属性） "graphData": {...}, # 查询结果图（头尾路径 / 节点展开等） "short_routes": {...} # 最短路径（仅头尾模式有效） } """ filename = request.filename queryview = request.queryview querylink = request.querylink head = request.head tail = request.tail node = request.node depth = request.depth source_node = request.source_node simple_graphData_copy = request.simple_graphData_copy simple_graphData = {"nodes": [], "links": []} logger.info("Fetching graph sample") try: # 只有当没有特定查询意图（头、尾、节点、文件名、视图）时，才加载全量采样 if not any([filename, queryview, head, tail, node]): simple_graphData = await run_in_threadpool(fetch_graph_sample, driver) # # 处理外部传入的 copy 数据 (如果存在) if simple_graphData_copy: # 注意：这里假设传入的结构是 {"simple_graphData_copy": {...}} # 如果传入的就是图数据本身，请调整 key 的访问方式 data_source = simple_graphData_copy.get("simple_graphData_copy", simple_graphData_copy) # 删除 labels 字段（如果存在） data_source.pop("labels", None) # 删除 nodes 中每个 node 的 color 字段 for n in data_source.get("nodes", []): if isinstance(n, dict): n.pop("color", None) # 更新 simple_graphData 为处理后的数据 simple_graphData = data_source except Exception: logger.exception("Neo4j sample query error") # 可以选择继续执行（使用空图），或者抛出错误 # 这里选择继续执行，simple_graphData 保持为空或默认值 simple_graphData = {"nodes": [], "links": []} def db_task(): try: with driver.session() as session: # 默认返回值 res = { "graphData": {"nodes": [], "links": []}, "short_routes": {"nodes": [], "links": []}, "message": "", } # =============================== # 情况 1: 头尾节点都有 # =============================== if head and tail: logger.info(f"Mode: HEAD + TAIL | {head} -> {tail}") h_exists = session.run("MATCH (h {名称: $head}) RETURN h", head=head).single() t_exists = session.run("MATCH (t {名称: $tail}) RETURN t", tail=tail).single() if h_exists and t_exists: res["message"] = "头节点和尾节点展开成功" if not h_exists and not t_exists: res["message"] = "头节点和尾节点均不存在" res["graphData"] = {"nodes": [], "links": []} return res elif not h_exists: res["message"] = "头节点不存在" # 返回存在的尾节点（如果存在） nodes_list = [] if t_exists: nodes_list.append(node_to_json(t_exists["t"])) res["graphData"] = {"nodes": nodes_list, "links": []} return res elif not t_exists: res["message"] = "尾节点不存在" nodes_list = [node_to_json(h_exists["h"])] res["graphData"] = {"nodes": nodes_list, "links": []} return res # 两者都存在，查路径 sp_paths = list( session.run( """ MATCH path = allShortestPaths((h {名称: $head})-[*1..10]->(t {名称: $tail})) RETURN path AS p """, head=head, tail=tail, ) ) all_paths = list( session.run( """ MATCH path = (h {名称: $head})-[*1..10]->(t {名称: $tail}) UNWIND nodes(path) AS x WITH path, COUNT(DISTINCT x) AS nodeCount WHERE nodeCount = LENGTH(path) + 1 RETURN path AS p """, head=head, tail=tail, ) ) if not sp_paths and not all_paths: # 无连接，返回两个孤立节点 nodes_list = [node_to_json(h_exists["h"]), node_to_json(t_exists["t"])] res["graphData"] = {"nodes": nodes_list, "links": []} res["short_routes"] = {"nodes": [], "links": []} else: res["graphData"] = build_graph_from_paths(all_paths) res["short_routes"] = build_graph_from_paths(sp_paths) # === 新增：为 graphData 中每个节点计算 level（最大无向跳数，≤3）=== nodes_list = res["graphData"]["nodes"] if not nodes_list: return res # 确保每个节点都有 "名称" valid_nodes = [n for n in nodes_list if "name" in n and n["name"]] node_names = [n["name"] for n in valid_nodes] name_to_node = {n["name"]: n for n in valid_nodes} for name in node_names: try: result = session.run( """ MATCH (start {名称: $name}) OPTIONAL MATCH path = (start)-[*1]-(end) WHERE start <> end RETURN coalesce(max(length(path)), 0) AS max_level """, name=name, ).single() if result: level = min(result["max_level"], 3) name_to_node[name]["level"] = level else: name_to_node[name]["level"] = 0 except Exception as e: logger.warning(f"Failed to compute level for {name}: {e}") name_to_node[name]["level"] = 0 apply_styles(res["graphData"]) apply_styles(res["short_routes"]) return res # =============================== # 情况 2: 单节点查询 # =============================== elif node: logger.info(f"Mode: NODE (1-hop only) | {node}") # === 只查询 1 跳关系 === paths = session.run( """ MATCH path = (n {名称: $node})-[*1]-(m) UNWIND nodes(path) AS x WITH path, COUNT(DISTINCT x) AS nodeCount WHERE nodeCount = LENGTH(path) + 1 RETURN path AS p """, node=node, ) # 构建图 graph = build_graph_from_paths(paths) # ============================== # ✅ 新增：节点类型过滤（核心） # ============================== FILTER_TYPES = { "功能", "环境条件", "组成", "技术指标参数", "接口", "安全警告", "操作使用", "维修工作", "调试", "器材保障", "图册", "试验鉴定", "设计单位", "组成部件", "无人机", "无人机性能", "外部设备", "维修", "操作", "故障", "舰艇", "系统", "子系统", "设备", } nodes = graph.get("nodes", []) links = graph.get("links", []) # 👉 1. 过滤节点（注意字段名） filtered_nodes = [ n for n in nodes if n.get("labels")[0] in FILTER_TYPES # ⚠️ type 要和你的结构一致 ] # 👉 2. 保留合法节点ID valid_ids = {n["id"] for n in filtered_nodes} # 👉 3. 过滤边（避免“断边”） filtered_links = [l for l in links if l.get("source") in valid_ids and l.get("target") in valid_ids] # 👉 4. 回写 graph graph["nodes"] = filtered_nodes graph["links"] = filtered_links # ============================== center_node_result = session.run("MATCH (n {名称: $node}) RETURN n", node=node).single() center_node_json = None if center_node_result: center_node_json = node_to_json(center_node_result["n"]) short_routes = {"nodes": [center_node_json] if center_node_json else [], "links": []} res["short_routes"] = short_routes res["graphData"] = graph # === level 计算 === nodes_list = graph.get("nodes", []) if not nodes_list: res["message"] = "未查询到相关节点" return res valid_nodes = [n for n in nodes_list if n.get("name")] name_to_node = {n["name"]: n for n in valid_nodes} for name in name_to_node: try: result = session.run( """ MATCH (start {名称: $name}) OPTIONAL MATCH (start)-[*1]-(end) WHERE start <> end RETURN COUNT(end) AS cnt """, name=name, ).single() level = 1 if result and result["cnt"] > 0 else 0 name_to_node[name]["level"] = level except Exception as e: logger.warning(f"Failed to compute level for {name}: {e}") name_to_node[name]["level"] = 0 apply_styles(graph) apply_styles(short_routes) res["message"] = "节点展开成功" return res # =============================== # 情况 3: 有simple_graphData_copy和source_node # =============================== elif simple_graphData_copy and source_node: logger.info(f"Mode: SIMPLE_GRAPHDATA_COPY + SOURCE_NODE | source_node={source_node}") data_source = simple_graphData_copy.get("simple_graphData_copy", simple_graphData_copy) data_source.pop("labels", None) for n in data_source.get("nodes", []): if isinstance(n, dict): n.pop("color", None) existing_nodes = {n["id"]: n for n in data_source.get("nodes", []) if n.get("id")} existing_links = {l["id"]: l for l in data_source.get("links", []) if l.get("id")} s_exists = session.run("MATCH (s {名称: $source_node}) RETURN s", source_node=source_node).single() if not s_exists: res["message"] = "source_node 节点不存在" res["short_routes"] = { "nodes": list(existing_nodes.values()), "links": list(existing_links.values()), } res["graphData"] = { "nodes": list(existing_nodes.values()), "links": list(existing_links.values()), } return res source_node_json = node_to_json(s_exists["s"]) source_id = s_exists["s"].element_id if source_id not in existing_nodes: existing_nodes[source_id] = source_node_json paths = session.run( """ MATCH path = (s {名称: $source_node})-[*1]-(m) RETURN path """, source_node=source_node, ) for p_rec in paths: path = p_rec["path"] for n in path.nodes: nid = n.element_id if nid not in existing_nodes: existing_nodes[nid] = node_to_json(n) for rel in path.relationships: rid = rel.element_id if rid not in existing_links: existing_links[rid] = rel_to_json(rel) short_routes_data = simple_graphData_copy.get("simple_graphData_copy", simple_graphData_copy) short_routes_nodes = list(short_routes_data.get("nodes", [])) short_routes_links = list(short_routes_data.get("links", [])) graph_data = {"nodes": list(existing_nodes.values()), "links": list(existing_links.values())} FILTER_TYPES = { "功能", "环境条件", "组成", "技术指标参数", "接口", "安全警告", "操作使用", "维修工作", "调试", "器材保障", "图册", "试验鉴定", "设计单位", "组成部件", "无人机", "无人机性能", "外部设备", "维修", "操作", "故障", "舰艇", "系统", "子系统", "设备", } nodes = graph_data.get("nodes", []) links = graph_data.get("links", []) filtered_nodes = [ n for n in nodes if n.get("labels") and len(n["labels"]) > 0 and n["labels"][0] in FILTER_TYPES ] valid_ids = {n["id"] for n in filtered_nodes} filtered_links = [l for l in links if l.get("source") in valid_ids and l.get("target") in valid_ids] graph_data["nodes"] = filtered_nodes graph_data["links"] = filtered_links short_routes = {"nodes": short_routes_nodes, "links": short_routes_links} valid_nodes = [n for n in graph_data["nodes"] if n.get("name")] name_to_node = {n["name"]: n for n in valid_nodes if n.get("name")} for name in name_to_node: try: result = session.run( """ MATCH (start {名称: $name}) OPTIONAL MATCH path = (start)-[*1]-(end) WHERE start <> end RETURN coalesce(max(length(path)), 0) AS max_level """, name=name, ).single() level = min(result["max_level"], 3) if result else 0 name_to_node[name]["level"] = level except Exception as e: logger.warning(f"Failed to compute level for {name}: {e}") name_to_node[name]["level"] = 0 apply_styles(graph_data) apply_styles(short_routes) res["graphData"] = graph_data res["short_routes"] = short_routes res["message"] = f"查询节点成功！" return res # =============================== # 情况 4: 只有头节点（有向，最多3跳） # =============================== elif head: logger.info(f"Mode: HEAD ONLY | {head}") h_exists = session.run("MATCH (h {名称: $head}) RETURN h", head=head).single() if not h_exists: res["message"] = "头节点不存在" return res paths = session.run( """ MATCH path = (h {名称: $head})-[*1]->(n) UNWIND nodes(path) AS x WITH path, COUNT(DISTINCT x) AS nodeCount WHERE nodeCount = LENGTH(path) + 1 RETURN path AS p """, head=head, ) paths = list(paths) if not paths: # 3. 如果没有有向路径，返回单个节点信息 res["graphData"] = {"nodes": [node_to_json(h_exists["h"])], "links": []} else: res["graphData"] = build_graph_from_paths(paths) # === 新增：为 graphData 中每个节点计算 level（最大无向跳数，≤3）=== nodes_list = res["graphData"]["nodes"] if not nodes_list: return res # 确保每个节点都有 "名称" valid_nodes = [n for n in nodes_list if "name" in n and n["name"]] node_names = [n["name"] for n in valid_nodes] name_to_node = {n["name"]: n for n in valid_nodes} for name in node_names: try: result = session.run( """ MATCH (start {名称: $name}) OPTIONAL MATCH path = (start)-[*1]-(end) WHERE start <> end RETURN coalesce(max(length(path)), 0) AS max_level """, name=name, ).single() if result: level = min(result["max_level"], 3) name_to_node[name]["level"] = level else: name_to_node[name]["level"] = 0 except Exception as e: logger.warning(f"Failed to compute level for {name}: {e}") name_to_node[name]["level"] = 0 apply_styles(res["graphData"]) res["message"] = "头节点展开成功" return res # =============================== # 情况 5: 只有尾节点（应查入边！即 ←） # =============================== elif tail: logger.info(f"Mode: TAIL ONLY | {tail}") t_exists = session.run("MATCH (t {名称: $tail}) RETURN t", tail=tail).single() if not t_exists: res["message"] = "尾节点不存在" return res # 注意：尾节点作为终点，应查找指向它的路径（←） paths = session.run( """ MATCH path = (n)-[*1]->(t {名称: $tail}) UNWIND nodes(path) AS x WITH path, COUNT(DISTINCT x) AS nodeCount WHERE nodeCount = LENGTH(path) + 1 RETURN path AS p """, tail=tail, ) paths = list(paths) if not paths: # 3. 如果没有有向路径，返回单个节点信息 res["graphData"] = {"nodes": [node_to_json(t_exists["t"])], "links": []} else: res["graphData"] = build_graph_from_paths(paths) # === 新增：为 graphData 中每个节点计算 level（最大无向跳数，≤3）=== nodes_list = res["graphData"]["nodes"] if not nodes_list: return res # 确保每个节点都有 "名称" valid_nodes = [n for n in nodes_list if "name" in n and n["name"]] node_names = [n["name"] for n in valid_nodes] name_to_node = {n["name"]: n for n in valid_nodes} for name in node_names: try: result = session.run( """ MATCH (start {名称: $name}) OPTIONAL MATCH path = (start)-[*1]-(end) WHERE start <> end RETURN coalesce(max(length(path)), 0) AS max_level """, name=name, ).single() if result: level = min(result["max_level"], 3) name_to_node[name]["level"] = level else: name_to_node[name]["level"] = 0 except Exception as e: logger.warning(f"Failed to compute level for {name}: {e}") name_to_node[name]["level"] = 0 apply_styles(res["graphData"]) res["message"] = "尾节点展开成功" return res # =============================== # 情况 6: 查询节点展开 # =============================== elif queryview: logger.info(f"Mode: QUERYVIEW + QUERYLINK | nodes={queryview}, links={querylink}") node_ids = [x.strip() for x in queryview.split(",") if x.strip()] if querylink: link_ids = [x.strip() for x in querylink.split(",") if x.strip()] else: link_ids = [] if not node_ids and not link_ids: res["message"] = "queryview 和 querylink 均为空" return res graph_data_nodes = {} graph_data_links = {} if node_ids: nodes_result = session.run("MATCH (n) WHERE elementId(n) IN $ids RETURN n", ids=node_ids) for rec in nodes_result: n = rec["n"] nid = n.element_id if nid not in graph_data_nodes: graph_data_nodes[nid] = node_to_json(n) if link_ids: rels_result = session.run( """ MATCH ()-[r]->() WHERE elementId(r) IN $link_ids RETURN r, elementId(startNode(r)) AS start_id, elementId(endNode(r)) AS end_id """, link_ids=link_ids, ) for record in rels_result: rel = record["r"] rid = rel.element_id if rid not in graph_data_links: graph_data_links[rid] = rel_to_json(rel) start_id = record["start_id"] end_id = record["end_id"] if start_id not in graph_data_nodes: node_result = session.run( "MATCH (n) WHERE elementId(n) = $id RETURN n", id=start_id ).single() if node_result: graph_data_nodes[start_id] = node_to_json(node_result["n"]) if end_id not in graph_data_nodes: node_result = session.run( "MATCH (n) WHERE elementId(n) = $id RETURN n", id=end_id ).single() if node_result: graph_data_nodes[end_id] = node_to_json(node_result["n"]) graph_data = {"nodes": list(graph_data_nodes.values()), "links": list(graph_data_links.values())} short_routes = {"nodes": list(graph_data_nodes.values()), "links": list(graph_data_links.values())} valid_nodes = [n for n in graph_data["nodes"] if n.get("name")] name_to_node = {n["name"]: n for n in valid_nodes if n.get("name")} for name in name_to_node: try: result = session.run( """ MATCH (start {名称: $name}) OPTIONAL MATCH path = (start)-[*1]-(end) WHERE start <> end RETURN coalesce(max(length(path)), 0) AS max_level """, name=name, ).single() level = min(result["max_level"], 3) if result else 0 name_to_node[name]["level"] = level except Exception as e: logger.warning(f"Failed to compute level for {name}: {e}") name_to_node[name]["level"] = 0 apply_styles(graph_data) apply_styles(short_routes) res["graphData"] = graph_data res["short_routes"] = short_routes res["message"] = f"查询拓展加载成功！" return res # =============================== # 情况 7: 按文件名检索所有匹配节点 # =============================== elif filename: logger.info(f"Mode: FILENAME SEARCH | filename={filename}") try: search_pattern = f'"filename": "{filename}"' result = session.run( """ MATCH (n) WHERE n.knowledge_source IS NOT NULL AND n.knowledge_source CONTAINS $search_pattern RETURN n """, search_pattern=search_pattern, ) nodes_list = [] node_ids = [] # ========= 查询节点 ========= for record in result: n = record["n"] node_json = node_to_json(n) # 使用 elementId 作为前端 graph 的唯一 id eid = n.element_id node_json["id"] = eid node_ids.append(eid) identifier = node_json.get("name") or node_json.get("名称") if identifier: try: level_result = session.run( """ MATCH (start {名称: $id}) OPTIONAL MATCH path = (start)-[*1..3]-(end) WHERE start <> end RETURN coalesce(max(length(path)), 0) AS max_level """, id=identifier, ).single() level = min(level_result["max_level"], 3) if level_result else 0 node_json["level"] = level except Exception as e: logger.warning(f"Level calc failed for {identifier}: {e}") node_json["level"] = 0 else: node_json["level"] = 0 nodes_list.append(node_json) # ========= 查询节点之间的边 ========= links_list = [] if node_ids: rel_result = session.run( """ MATCH (a)-[r]->(b) WHERE elementId(a) IN $ids AND elementId(b) IN $ids RETURN r, elementId(a) AS source, elementId(b) AS target """, ids=node_ids, ) for record in rel_result: r = record["r"] link_json = rel_to_json(r) # 你已有的方法 # # 确保 source / target 正确 # link_json["source"] = record["source"] # link_json["target"] = record["target"] links_list.append(link_json) # ========= 返回结果 ========= if nodes_list: graph_data = {"nodes": nodes_list, "links": links_list} apply_styles(graph_data) res["graphData"] = graph_data res["message"] = ( f"找到 {len(nodes_list)} 个关联文件名 '{filename}' 的节点，{len(links_list)} 条关系" ) else: res["graphData"] = {"nodes": [], "links": []} res["message"] = f"未找到关联文件名 '{filename}' 的节点" return res except Exception as e: logger.exception("Filename search error") res["message"] = f"按文件名查询节点时出错: {str(e)}" return res # =============================== # 兜底：无任何输入 # =============================== else: res["message"] = "查询成功！" return res except Exception as e: logger.exception("Internal db_task error") return { "graphData": {"nodes": [], "links": []}, "short_routes": {"nodes": [], "links": []}, "message": f"数据库查询异常: {str(e)}", } # 执行数据库任务 query_result = await run_in_threadpool(db_task) # 组装最终响应 if head or tail: return { "message": query_result["message"], "simple_graphData": query_result["graphData"], "graphData": query_result["graphData"], "short_routes": query_result.get("short_routes", {"nodes": [], "links": []}), } if node and depth: return { "message": query_result["message"], "simple_graphData": query_result["graphData"], "graphData": query_result["graphData"], "short_routes": query_result.get("short_routes", {"nodes": [], "links": []}), } if simple_graphData_copy and source_node: return { "message": query_result["message"], "simple_graphData": query_result["graphData"], "graphData": query_result["graphData"], "short_routes": query_result.get("short_routes", {"nodes": [], "links": []}), } if queryview: return { "message": query_result["message"], "simple_graphData": query_result["graphData"], "graphData": query_result["graphData"], "short_routes": query_result.get("short_routes", {"nodes": [], "links": []}), } if filename: return { "message": query_result["message"], "simple_graphData": query_result["graphData"], "graphData": query_result["graphData"], "short_routes": query_result.get("short_routes", {"nodes": [], "links": []}), } else: apply_styles(simple_graphData) tree_graphData = build_tree(simple_graphData["nodes"], simple_graphData["links"]) return { "message": query_result["message"], "simple_graphData": simple_graphData, "graphData": simple_graphData, "short_routes": query_result.get("short_routes", {"nodes": [], "links": []}), "tree_graphData": tree_graphData } class FuzzyNodeOnlyRequest(BaseModel): node: str limit: Optional[int] = 20 @app.post("/api/v2/neo4j/fuzzy_query") async def fuzzy_node_only(request: FuzzyNodeOnlyRequest): """ 节点名称模糊匹配（仅返回节点，不返回关系） """ keyword = request.node.strip() limit = request.limit or 50 if not keyword: raise HTTPException(status_code=400, detail="keyword 不能为空") def db_task(): try: with driver.session() as session: logger.info(f"Mode: FUZZY NODE ONLY | keyword='{keyword}', limit={limit}") result = session.run( """ MATCH (n) WHERE n.名称 CONTAINS $kw RETURN n LIMIT $limit """, kw=keyword, limit=limit, ) nodes_list = [] for record in result: n = record["n"] node_json = node_to_json(n) # 统一前端使用 elementId 作为 id node_json["id"] = n.element_id nodes_list.append(node_json) if nodes_list: return {"message": f"模糊匹配 '{keyword}' 成功，找到 {len(nodes_list)} 个节点", "nodes": nodes_list} else: return {"message": f"未找到名称包含 '{keyword}' 的节点", "nodes": []} except Exception as e: logger.exception("Fuzzy node only query error") return {"message": f"节点模糊查询异常: {str(e)}", "nodes": []} return await run_in_threadpool(db_task) class ExpandQueryRequest(BaseModel): node: str @app.post("/api/v2/neo4j/expand") async def expand_node(request: ExpandQueryRequest): """ 无向节点拓展： 仅返回指定节点 1 跳范围内的所有无向关联及节点信息。 """ node_name = request.node if not node_name: raise HTTPException(status_code=400, detail="Node name is required") logger.info(f"Expanding node: {node_name} | depth=1 (undirected)") def db_task_expand(): try: with driver.session() as session: # 1. 执行 1 跳无向查询 # 使用 -[r]- 表示无向，匹配起点 node_name 的所有直接邻居 cypher_query = """ MATCH path = (n {名称: $name})-[r]-(m) RETURN path AS p """ result = session.run(cypher_query, name=node_name) # 将结果转为列表，以便多次处理或判断 records = list(result) if not records: logger.info(f"No connections found for node: {node_name}") return {"nodes": [], "links": []} # 使用 build_graph_from_paths 解析路径（内部已过滤 Searchable） return build_graph_from_paths(records) except Exception as e: logger.exception("Neo4j expansion query error") raise e # 执行数据库任务 try: related_data = await run_in_threadpool(db_task_expand) except Exception as e: raise HTTPException(status_code=500, detail=f"Expansion failed: {str(e)}") return {"message": "200", "related_graphData": related_data} def _sync_kg_build_task( task_id: str, files_data: List[dict], kg_config: dict, result_queue: Queue, cancel_flag_path: str ): """子进程中执行知识图谱构建（并发版：最多 4 个文件同时处理）""" logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) complete_files = [] _complete_files_lock = threading.Lock() # 保护 complete_files 和 result_queue _kg_lock = threading.Lock() # 保护 Neo4j 入库（kg.process_data / filter_node_relationship） def check_cancelled(): return os.path.exists(cancel_flag_path) # 配置加载 NEO4J_URI = kg_config["neo4j_uri"] NEO4J_USER = kg_config["neo4j_user"] NEO4J_PASSWORD = kg_config["neo4j_password"] DATA_DIR = "/app/files" Ontology = kg_config["ontology"] Ontology_DESC = kg_config["ontology_description"] relationships = kg_config["relationships"] node_types = kg_config["node_types"] relationship_types = kg_config["relationship_types"] kg = MilitaryKnowledgeGraph(uri=NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD), Ontology=Ontology) # ---- 单文件处理函数（供线程池调用） ---- def _process_single_file(idx: int, file_data: dict): """处理单个文件：OCR/提取 → 入库。线程安全。""" entities = [] relations = [] filename_entity = [] filename_relation = [] lower_entity_device = None lower_entity_type = None file_name = file_data.get("file_name") or f"unknown_{idx}" _t_file_start = time.time() try: if check_cancelled(): raise InterruptedError("任务已取消") file_name = file_data["file_name"] file_base64 = file_data["file_base64"] logger.info(f"[线程] 处理文件: {file_name}") # --- 1. 文件名相关的简单实体提取 --- try: final_result = get_entity(file_name) filename_entity = final_result.get("entities", []) print(f"文件名实体提取结果: {filename_entity}") filename_relation = final_result.get("relationships", []) print(f"文件名关系提取结果: {filename_relation}") lower_entity = final_result.get("low_level") for ins in filename_entity: if ins.get("properties", {}).get("名称") == lower_entity: lower_entity_device = ins["properties"].get("名称") lower_entity_type = ins["type"] break except Exception as e: logger.warning(f"文件名实体提取失败: {e}") filename_entity = [] filename_relation = [] # --- 2. 保存文件到磁盘 --- os.makedirs(DATA_DIR, exist_ok=True) FILE_PATH = os.path.join(DATA_DIR, file_name) file_binary = base64.b64decode(file_base64) with open(FILE_PATH, "wb") as f: f.write(file_binary) logger.info(f"? 已保存文件: {FILE_PATH}") content_middle = file_data.get("content_middle") or [] check_result = check_filename_ocr(file_name) print("#####################################") print(check_result) if content_middle: # --- 3a. 已有预解析内容，直接抽取，跳过 MinerU --- logger.info(f"文件 {file_name} 使用预解析内容，跳过 MinerU") _t_extract_start = time.time() async def run_extraction_with_content(): return await get_md_node_relation( content_list=content_middle, input_pdf_path=FILE_PATH, Ontology=Ontology, Ontology_DESCRIPTION=Ontology_DESC, Relationships=relationships, RELATIONSHIP_TYPES=relationship_types, NODE_TYPES=node_types, PREFIX_URL=PREFIX_URL, device=lower_entity_device, device_type=lower_entity_type, ) entities, relations = asyncio.run(run_extraction_with_content()) elif check_result == "ocr_true": logger.info(f"文件 {file_name} 需要 OCR 处理") # --- 3b. 调用 MinerU 进行 OCR（get_api_url 轮询不同实例） --- _t_extract_start = time.time() entities, relations = process_file_content( file_name=file_name, file_path=FILE_PATH, task_id=task_id, data_dir=DATA_DIR, ontology=Ontology, ontology_desc=Ontology_DESC, relationships=relationships, node_types=node_types, relationship_types=relationship_types, prefix_url=PREFIX_URL, api_url=file_data["api_url"], check_cancelled_callback=check_cancelled, device=lower_entity_device, device_type=lower_entity_type, ) elif check_result == "ocr_false": logger.info(f"文件 {file_name} 不需要 OCR 处理，调用 get_filtertext_node_relation") _t_extract_start = time.time() async def run_filtertext(): return await get_filtertext_node_relation( input_pdf_path=FILE_PATH, data=[], prefix_url=PREFIX_URL, device=lower_entity_device ) try: filter_result = asyncio.run(run_filtertext()) if isinstance(filter_result, tuple) and len(filter_result) >= 2: entities, relations = filter_result[0], filter_result[1] logger.info(f"get_filtertext_node_relation 返回: 实体 {len(entities)}, 关系 {len(relations)}") else: logger.error( f"get_filtertext_node_relation 返回格式异常: {type(filter_result)}, 值: {filter_result}" ) entities, relations = [], [] except Exception as e: logger.error(f"get_filtertext_node_relation 调用失败: {e}") logger.error(traceback.format_exc()) entities, relations = [], [] else: logger.warning(f"未知的 check_result: {check_result}，默认使用空值") _t_extract_start = time.time() entities, relations = [], [] _t_extract_cost = time.time() - _t_extract_start logger.info(f"[计时][{file_name}] OCR/抽取耗时: {_t_extract_cost:.1f}s") # --- 4. 后续合并与入库逻辑 --- all_entity = entities + filename_entity all_relation = relations + filename_relation all_entity_new = [ins for ins in all_entity if ins.get("properties", {}).get("名称")] # 生成 Embedding _t_embed_start = time.time() for entity in all_entity_new: name = entity["properties"].get("名称") if "未知" in name: continue if name: try: embedding_vector = OpenaiAPI.generate_embedding(name) entity["properties"]["embedding"] = embedding_vector except Exception as e: logger.warning(f"生成 embedding 失败 for {name}: {e}") _t_embed_cost = time.time() - _t_embed_start logger.info(f"[计时][{file_name}] Embedding耗时: {_t_embed_cost:.1f}s (实体数: {len(all_entity_new)})") merged_output_data = {"entities": all_entity_new, "relationships": all_relation} # 调试输出 debug_path = f"/app/output_{task_id}_{idx}.json" try: with open(debug_path, "w", encoding="utf-8") as f: json.dump(merged_output_data, f, ensure_ascii=False, indent=2) except Exception as e: logger.warning(f"保存调试文件失败: {e}") # 入库（加锁，Neo4j 写操作串行化，避免并发冲突） _t_lock_wait_start = time.time() with _kg_lock: _t_lock_wait_cost = time.time() - _t_lock_wait_start logger.info(f"[计时][{file_name}] 等锁耗时: {_t_lock_wait_cost:.1f}s") _t_filter_start = time.time() try: filter_data = filter_node_relationship( URI=NEO4J_URI, USERNAME=NEO4J_USER, PASSWORD=NEO4J_PASSWORD, merged_output_data=merged_output_data, ) except Exception as e: logger.error(f"过滤节点关系失败: {e}") raise _t_filter_cost = time.time() - _t_filter_start logger.info(f"[计时][{file_name}] filter去重耗时: {_t_filter_cost:.1f}s") _t_write_start = time.time() entity_count, rel_count = kg.process_data(filter_data) _t_write_cost = time.time() - _t_write_start logger.info( f"[计时][{file_name}] Neo4j写入耗时: {_t_write_cost:.1f}s (实体: {entity_count}, 关系: {rel_count})" ) # 建立索引 try: data = { "neo4j_uri": NEO4J_URI, "neo4j_username": NEO4J_USER, "neo4j_password": NEO4J_PASSWORD, "force_refresh": True, } requests.post(indexing_url, json=data, timeout=30) except Exception as e: logger.warning(f"建立索引失败: {e}") _t_file_total = time.time() - _t_file_start logger.info( f"[计时][{file_name}] ===总耗时: {_t_file_total:.1f}s === " f"(抽取:{_t_extract_cost:.1f}s embed:{_t_embed_cost:.1f}s " f"等锁:{_t_lock_wait_cost:.1f}s filter:{_t_filter_cost:.1f}s 写入:{_t_write_cost:.1f}s)" ) # 记录结果（加锁） with _complete_files_lock: file_result = { "file_name": file_name, "status": "success", "entity_count": entity_count, "relation_count": rel_count, } complete_files.append(file_result) result_queue.put( { "code": 200, "type": "progress", "file_index": idx, "total_files": len(files_data), "file_name": file_name, "status": "success", "entity_count": entity_count, "relation_count": rel_count, "complete_files": complete_files.copy(), } ) except InterruptedError: with _complete_files_lock: complete_files.append({"file_name": file_name, "status": "cancelled"}) raise except Exception as e: error_msg = f"[{task_id}] 文件 {file_name} 处理失败: {str(e)}" logger.error(error_msg) logger.error(traceback.format_exc()) with _complete_files_lock: file_result = {"file_name": file_name, "status": "failed", "error": str(e)} complete_files.append(file_result) result_queue.put( { "code": 200, "type": "progress", "file_index": idx, "total_files": len(files_data), "file_name": file_name, "status": "failed", "error": str(e), "complete_files": complete_files.copy(), } ) # ---- 并发执行 ---- # 预解析文件不经过 MinerU，并发数不受 API_URLS 限制 # 需要 MinerU 的文件受 API_URLS 实例数约束，取两者中较大值保证并发 _files_need_mineru = sum(1 for fd in files_data if not (fd.get("content_middle") or [])) CONCURRENT_WORKERS = ( max(len(API_URLS), min(len(files_data), 100)) if _files_need_mineru == 0 else max(len(API_URLS), 1) ) logger.info(f"[{task_id}] 并发数: {CONCURRENT_WORKERS}，需MinerU文件: {_files_need_mineru}/{len(files_data)}") cancelled = False with ThreadPoolExecutor(max_workers=CONCURRENT_WORKERS) as executor: futures = {executor.submit(_process_single_file, idx, fd): (idx, fd) for idx, fd in enumerate(files_data)} for future in futures: try: future.result() # 等待完成，异常已在内部处理 except InterruptedError: cancelled = True # 取消尚未开始的任务 for f in futures: f.cancel() break except Exception: pass # 已在 _process_single_file 内部处理并汇报 if cancelled: result_queue.put( {"code": 200, "status": "cancelled", "message": "任务已取消", "complete_files": complete_files} ) return # 所有文件处理完成 result_queue.put( {"code": 200, "type": "completed", "status": "success", "message": "构建完成", "complete_files": complete_files} ) @app.post("/kg_build") async def kg_build(request: KGRequest): task_id = request.task_id if task_id in kg_running_tasks: existing = kg_running_tasks[task_id] if existing.get("process") and existing["process"].is_alive(): return {"code": 400, "message": f"任务 {task_id} 已在运行中"} files = request.files if isinstance(request.files, list) else [request.files] files_data = [] for i, f in enumerate(files): _api, _cache = get_api_url(f"{task_id}_{i}") # 哈希取模，多 worker 进程安全 logger.info(f"[{task_id}] 文件 {f.file_name} → MinerU {_api}") files_data.append( { "file_name": f.file_name, "markdown_content": f.markdown_content, "file_base64": f.file_base64, "content_middle": f.content_middle, "api_url": _api, "get_contentlist_url": _cache, } ) logger.info(f"[{task_id}] 开始构建知识图谱，文件数量: {len(files_data)}") # ★ 关键改动：合并默认值与用户传入值 merged = build_kg_config( user_ontology=request.ontology, user_ontology_desc=request.Ontology_DESCRIPTION, user_relationships=request.relationships, user_node_types=request.node_types, user_relationship_types=request.relationship_types, ) kg_config = { "neo4j_uri": NEO4J_URI, "neo4j_user": NEO4J_USER, "neo4j_password": NEO4J_PASSWORD, "data_dir": DATA_DIR, **merged, # ontology / ontology_description / relationships / node_types / relationship_types } # 持久化配置缓存（使用合并后的值） try: cache_dir = DATA_DIR if os.path.isdir(DATA_DIR) else os.path.dirname(__file__) os.makedirs(cache_dir, exist_ok=True) cache_path = os.path.join(cache_dir, "kg_config_cache.json") async with aiofiles.open(cache_path, "w", encoding="utf-8") as f: await f.write(json.dumps( { "ontology": kg_config["ontology"], "ontology_description": kg_config["ontology_description"], "relationships": kg_config["relationships"], "node_types": kg_config["node_types"], "relationship_types": kg_config["relationship_types"], "updated_at": datetime.utcnow().isoformat() + "Z", }, ensure_ascii=False, indent=2, )) logger.info(f"已保存 kg_build 配置缓存: {cache_path}") except Exception as e: logger.warning(f"保存 kg_build 配置缓存失败: {e}") # 后续进程启动逻辑保持不变 ... cancel_flag_path = f"/tmp/kg_cancel_{task_id}" if os.path.exists(cancel_flag_path): os.remove(cancel_flag_path) result_queue = Queue() proc = Process( target=_sync_kg_build_task, args=(task_id, files_data, kg_config, result_queue, cancel_flag_path), daemon=True, ) proc.start() # 记录任务（用于取消） task_data = { "process": proc, "result_queue": result_queue, "cancel_flag_path": cancel_flag_path, "status": "processing", } kg_running_tasks[task_id] = task_data # 流式返回生成器 async def event_generator(): """生成器：实时返回每个文件的处理进度""" timeout = 3600 # 2小时超时 start = time.time() try: while True: # 超时检查 if time.time() - start > timeout: proc.terminate() proc.join(timeout=5) if proc.is_alive(): proc.kill() kg_running_tasks.pop(task_id, None) yield ( json.dumps( {"code": 500, "type": "error", "message": "任务执行超时", "complete_files": []}, ensure_ascii=False, ) + "\n" ) return # 检查是否有结果 try: # 非阻塞获取队列消息 while not result_queue.empty(): result = result_queue.get_nowait() yield json.dumps(result, ensure_ascii=False) + "\n" # 如果是完成或错误消息，结束生成器 if result.get("type") in ["completed", "error", "cancelled"]: kg_running_tasks.pop(task_id, None) return except Exception as e: logger.warning(f"[{task_id}] Queue读取异常: {e}") # 检查进程是否还活着 if not proc.is_alive(): # 进程结束，读取剩余消息 while not result_queue.empty(): result = result_queue.get_nowait() yield json.dumps(result, ensure_ascii=False) + "\n" # 如果队列为空但进程异常退出 if result_queue.empty(): yield ( json.dumps( { "code": 500, "type": "error", "message": f"进程异常退出: {proc.exitcode}", "complete_files": [], }, ensure_ascii=False, ) + "\n" ) kg_running_tasks.pop(task_id, None) return # 短暂休眠，避免CPU空转 await asyncio.sleep(0.3) except Exception as e: kg_running_tasks.pop(task_id, None) logger.error(f"[{task_id}] 流式返回异常: {str(e)}") yield ( json.dumps({"code": 500, "type": "error", "message": str(e), "complete_files": []}, ensure_ascii=False) + "\n" ) return StreamingResponse(event_generator(), media_type="application/x-ndjson") # ==================== kg_task_cancel 接口 ==================== @app.post("/kg_task_cancel/{task_id}") async def kg_task_cancel(task_id: str): """取消知识图谱构建任务""" task = kg_running_tasks.get(task_id) if not task: raise HTTPException(status_code=404, detail=f"任务 {task_id} 不存在") proc = task.get("process") if not proc or not proc.is_alive(): return {"code": 400, "message": "任务已结束", "status": task.get("status")} cancel_flag_path = task.get("cancel_flag_path") try: # 1. 设置取消标志 async with aiofiles.open(cancel_flag_path, "w") as f: await f.write(str(time.time())) await asyncio.sleep(3) # 2. 强制终止 if proc.is_alive(): os.kill(proc.pid, signal.SIGTERM) proc.join(timeout=5) if proc.is_alive(): os.kill(proc.pid, signal.SIGKILL) proc.join() task["status"] = "cancelled" if os.path.exists(cancel_flag_path): os.remove(cancel_flag_path) kg_running_tasks.pop(task_id, None) logger.info(f"[{task_id}] 任务已取消") return {"code": 200, "message": "任务已取消", "task_id": task_id} except Exception as e: return {"code": 500, "message": f"取消失败: {str(e)}"} @app.post("/ontology_reset") async def ontology_reset(request: resetRequest): # ★ 关键改动：合并默认值与用户传入值 merged = build_kg_config( user_ontology=request.ontology, user_ontology_desc=request.Ontology_DESCRIPTION, user_relationships=request.relationships, user_node_types=request.node_types, user_relationship_types=request.relationship_types, ) Ontology = merged["ontology"] Ontology_DESC = merged["ontology_description"] relationships = merged["relationships"] node_types = merged["node_types"] relationship_types = merged["relationship_types"] files = request.files if isinstance(request.files, list) else [request.files] logger.info(f"开始本体衍生，文件数量: {len(files)}") async def event_generator(): total_files = len(files) for file_idx, file in enumerate(files): markdown_content = file.markdown_content file_name = file.file_name new_filename = re.sub(r"^[0-9a-f]{8}-([0-9a-f]{4}-){3}[0-9a-f]{12}_", "", file_name) wx_tables, cleaned_wx_content = extract_repair_tables_keep_first(content=markdown_content) opration_tables, cleaned_opration_content = extract_operation_tables(cleaned_wx_content) nromaltabel, contexts, cleaned_normal_content = extract_tables(cleaned_opration_content) content, headers = _split_markdown(cleaned_normal_content, chunk_size=15000, chunk_overlap=500) total_chunks = len(content) for chunk_idx, ins in enumerate(content): final_prompt = create_extraction_prompt( text=ins, ontology=Ontology, Ontology_DESCRIPTION=Ontology_DESC, relationships=relationships, relationship_types=relationship_types, node_types=node_types, ) final_prompt = create_extraction_prompt( text=ins, ontology=Ontology, Ontology_DESCRIPTION=Ontology_DESCRIPTION, relationships=relationships, relationship_types=relationship_types, node_types=node_types, ) res_str = await OpenaiAPI.openai_chat_aysnc(final_prompt) # 处理 API 调用失败 if res_str is None: error_msg = { "code": 500, "error": "API调用失败", "file_index": file_idx, "total_files": total_files, "chunk_index": chunk_idx, "total_chunks": total_chunks, "file_name": new_filename, } yield json.dumps(error_msg, ensure_ascii=False) + "\n" continue # 尝试解析 JSON try: res = json.loads(res_str) except json.JSONDecodeError as e: error_msg = { "code": 500, "error": f"JSON解析失败: {str(e)}", "raw_response": repr(res_str), "file_index": file_idx, "total_files": total_files, "chunk_index": chunk_idx, "total_chunks": total_chunks, "file_name": new_filename, } yield json.dumps(error_msg, ensure_ascii=False) + "\n" continue new_entities = res.get("new_entities", []) logger.info(f"new_entities: {new_entities}") new_relationships = res.get("new_relationships", []) logger.info(f"new_relationships: {new_relationships}") # 简化逻辑：file_index 就是当前文件索引（从0开始） result = { "code": 200, "file_index": file_idx, "total_files": total_files, "chunk_index": chunk_idx, "total_chunks": total_chunks, "new_entities": new_entities, "new_relationships": new_relationships, "file_name": new_filename, } logger.info(f"衍生本体如下：{result}") # 流式返回（JSON Lines） yield json.dumps(result, ensure_ascii=False) + "\n" # 当前文件所有chunk处理完毕后，发送文件完成标记 yield ( json.dumps( { "code": 200, "status": "file_completed", "file_index": file_idx, "total_files": total_files, "file_name": new_filename, "message": f"文件 {new_filename} 处理完成", }, ensure_ascii=False, ) + "\n" ) # 最后发送全部完成标记 yield ( json.dumps({"code": 200, "status": "completed", "message": "所有文件本体衍生完毕"}, ensure_ascii=False) + "\n" ) logger.info("所有文件本体衍生完毕") return StreamingResponse(event_generator(), media_type="application/x-ndjson") class PdfContentItem(BaseModel): page_idx: int bbox: Optional[str] = None text: Optional[str] = None text_level: Optional[int] = None img_path: Optional[str] = None table_body: Optional[str] = None class RequestWrapper(BaseModel): pdf_contents: List[PdfContentItem] @app.post("/split_content_list") async def split_content_list(request_data: RequestWrapper): pdf_contents = request_data.pdf_contents processed_list = [] for item in pdf_contents: # 将 Pydantic 模型转换为字典，方便修改 item_dict = item.model_dump() # --- 核心逻辑：判断并添加 type 字段 --- item_type = None # 1. 如果 text 非空 -> type: "text" if item.text: item_type = "text" # 2. 如果 table_body 非空 -> type: "table" # 注意：如果 text 和 table_body 都有，根据代码顺序，table 会覆盖 text。 # 如果你的业务逻辑是互斥的或需要优先级，请调整此处顺序。 elif item.table_body: item_type = "table" # 3. 如果 img_path 非空 且 table_body 为空 -> type: "image" # 注意：上面用了 elif 判断 table_body，所以到这里 table_body 肯定为空， # 但为了逻辑清晰，显式写出条件。 elif item.img_path and not item.table_body: item_type = "image" # 如果匹配到了类型，写入字典 if item_type: item_dict["type"] = item_type processed_list.append(item_dict) slices = get_chunk_bbox(processed_list) slices_check = chunk_check(slices, 8000) return {"code": 200, "message": "ok", "data": {"slices": slices_check}} @app.post("/split_result") async def split_result( file: UploadFile = File(...), image_prefix: str = Form("/api/v1/knowledge/files/images/"), ): if not file.filename: raise HTTPException(status_code=400, detail="File name is missing") chunk_size = 1024 chunk_overlap = 100 filename = file.filename suffix = PathLib(filename).suffix.lower() # 获取文件后缀 temp_input_path: Optional[PathLib] = None converted_pdf_path: Optional[PathLib] = None logger.info("调用知识切分接口") try: # ? 创建临时文件时指定后缀名 with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_in: content = await file.read() logger.info(f"文件大小:{len(content)}") logger.info(f"前10字节:{content[:10]}") # if not content.startwith(b"%PDF"): # logger.error("不是合法PDF") # raise HTTPException(status_code=400,detail="Uploaded file is not a valid PDF") tmp_in.write(content) temp_input_path = PathLib(tmp_in.name) # 2. 根据文件类型处理 if suffix == ".pdf": logger.info("调用pdf切分方式") result_data = await process_pdf_file(temp_input_path, image_prefix,filename) if result_data is None: raise HTTPException(status_code=502, detail="PDF processing failed") elif suffix == ".docx": logger.info("调用docx切分方式") logger.info(f"转换 DOCX -> PDF: {filename}") # 执行转换 converter.convert(str(temp_input_path), output_dir=str(DATA_DIR)) # ? 使用临时文件的基础名查找 PDF temp_base_name = temp_input_path.stem # 如 tmp2bb0l4cn converted_pdf_path = PathLib(DATA_DIR) / f"{temp_base_name}.pdf" if not converted_pdf_path.exists(): raise HTTPException(status_code=500, detail=f"PDF 转换失败: {converted_pdf_path}") result_data = await process_pdf_file(converted_pdf_path, image_prefix) if result_data is None: raise HTTPException(status_code=502, detail="DOCX to PDF processing failed") elif suffix in (".txt", ".md"): logger.info("调用txt切分方式") # 直接读取文本内容 async with aiofiles.open(temp_input_path, "r", encoding="utf-8", errors="ignore") as f: text_content = await f.read() contents = split_text_preserve_sentences(text_content) # 构造统一结构（无 images） slices = [] for ins in contents: slices.append({"content": ins, "positions": []}) result_data = {"slices": slices, "images": {}} elif suffix in [".xlsx"]: # converter.convert(str(temp_input_path), output_dir=str(DATA_DIR)) # temp_base_name = temp_input_path.stem # 如 tmp2bb0l4cn # converted_pdf_path = PathLib(DATA_DIR) / f"{temp_base_name}.pdf" # if not converted_pdf_path.exists(): # raise HTTPException(status_code=500, detail=f"PDF 转换失败: {converted_pdf_path}") logger.info("调用xlsx切分方式") result_data = await process_other_file(temp_input_path, image_prefix,filename) if result_data is None: raise HTTPException(status_code=502, detail="DOCX parse failed") elif suffix in [".ppt", ".pptx"]: logger.info("调用 PPT/PPTX 切分方式") # 1. 判断是否为旧版 .ppt 格式，如果是，则先进行转换 if suffix == ".ppt": logger.info(f"检测到旧版 PPT 文件，正在转换为 PPTX: {temp_input_path}") try: # 定义 LibreOffice 转换命令 command = [ 'libreoffice', '--headless', # 无头模式，不弹出图形界面 '--convert-to', 'pptx', '--outdir', os.path.dirname(temp_input_path), temp_input_path ] # 2. 使用 run_in_executor 将阻塞的同步转换操作放入线程池，避免阻塞异步事件循环 loop = asyncio.get_event_loop() result = await loop.run_in_executor( None, lambda: subprocess.run(command, capture_output=True, text=True) ) # 3. 检查转换是否成功 if result.returncode != 0: logger.error(f"PPT 转换失败: {result.stderr}") raise HTTPException(status_code=500, detail="PPT to PPTX conversion failed") # 4. 转换成功后，更新输入路径为新生成的 .pptx 文件路径 temp_input_path = os.path.splitext(temp_input_path)[0] + ".pptx" logger.info(f"PPT 转换成功，新文件路径: {temp_input_path}") except Exception as e: logger.exception(f"PPT 转换过程发生异常: {e}") raise HTTPException(status_code=500, detail=f"PPT conversion error: {str(e)}") # 5. 无论是原本就是 .pptx，还是刚刚转换完成的 .pptx，都统一走后续处理逻辑 result_data = await process_other_file(temp_input_path, image_prefix, filename) if result_data is None: raise HTTPException(status_code=502, detail="PPTX parse failed") else: raise HTTPException( status_code=400, detail="Unsupported file type. Only PDF, DOCX, TXT, and MD are supported." ) return JSONResponse({"code": 200, "message": "ok", "data": result_data}) except requests.RequestException as e: logger.error(f"Split service request failed: {e}") raise HTTPException(status_code=502, detail=f"Split service error: {str(e)}") except Exception as e: logger.error(f"Internal error: {e}") raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") finally: # 清理临时文件 for path in [temp_input_path, converted_pdf_path]: if path and path.exists(): try: path.unlink() except Exception as e: logger.warning(f"Failed to delete temp file {path}: {e}") # ========== ========== ========== ========== ========== ========== # 图谱检索接口 # ========== ========== ========== ========== ========== ========== class GraphSearchRequest(BaseModel): """图谱检索请求模型""" query: Optional[str] = None top_k: Optional[int] = 10 # 图谱检索返回结果数量 graph_timeout: Optional[float] = 2.0 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 索引创建接口 # ========== ========== ========== ========== ========== ========== class Neo4jIndexingRequest(BaseModel): """Neo4j索引创建请求模型""" neo4j_uri: Optional[str] = None # 可选，如果不提供则从环境变量获取 neo4j_username: Optional[str] = None # 可选，如果不提供则从环境变量获取 neo4j_password: Optional[str] = None # 可选，如果不提供则从环境变量获取 force_refresh: Optional[bool] = ( False # 是否强制刷新所有索引（True=清空重建，False=增量更新，只处理新节点或缺失的索引） ) @app.post("/neo4j_indexing") async def neo4j_indexing(request: Neo4jIndexingRequest): """ ========== ========== ========== ========== ========== ========== Neo4j 知识图谱索引创建接口 功能：为 Neo4j 知识图谱中的所有节点标签创建向量索引和全文索引， 并构建跨标签的全局混合(Searchable)索引。 请求参数（均为可选）： - neo4j_uri: Neo4j 连接 URI（可选，优先使用 .env 中的 NEO4J_URI） - neo4j_username: Neo4j 用户名（可选，优先使用 .env 中的 NEO4J_USERNAME） - neo4j_password: Neo4j 密码（可选，优先使用 .env 中的 NEO4J_PASSWORD） - force_refresh: 是否强制刷新所有索引（默认 False） * False（增量更新）：只处理新节点或缺失索引的节点，保留现有索引 * True（强制刷新）：清空所有索引和约束后重新创建 返回格式： { "code": 200 / 207 / 409 / 500, "message": "...", "data": { "success": bool, "message": "...", "mode": "force_refresh" / "incremental", "labels": [...], "labels_count": int, "vector_index_errors": [...], "fulltext_index_errors": [...], "hybrid_index_status": "success" / "failed", "hybrid_index_error": str / null, "elapsed_seconds": float } } 并发说明： - 索引创建为重操作，接口通过全局锁串行化； 已有任务在跑时，新请求立即返回 409，不会排队等待。 ========== ========== ========== ========== ========== ========== """ start_time = time.time() logger.info("[neo4j_indexing] 收到索引创建请求") # ========== 并发保护：已有索引任务在跑时直接拒绝 ========== if _neo4j_indexing_lock.locked(): logger.warning("[neo4j_indexing] 已有索引任务在执行，拒绝并发请求") return JSONResponse( status_code=409, content={ "code": 409, "message": "已有索引创建任务正在执行，请稍后再试", "data": None, }, ) async with _neo4j_indexing_lock: try: # ========== 检查依赖是否可用 ========== if not INDEXING_AVAILABLE: raise HTTPException( status_code=503, detail="Neo4j索引创建功能不可用，请确保已正确安装相关依赖", ) # force_refresh 在 Pydantic 模型中已有默认值 False，不会是 None force_refresh = request.force_refresh # 构建函数调用参数，只传递非 None 的参数 kwargs = {"force_refresh": force_refresh} if request.neo4j_uri: kwargs["neo4j_uri"] = request.neo4j_uri if request.neo4j_username: kwargs["neo4j_username"] = request.neo4j_username if request.neo4j_password: kwargs["neo4j_password"] = request.neo4j_password # ========== 第一步：按标签创建向量/全文索引 ========== # 同步操作放到线程池执行，避免阻塞事件循环 result = await asyncio.to_thread(create_all_indexes, **kwargs) # ========== 第二步：构建全局混合(Searchable)索引 ========== # 单独捕获异常：即使这一步失败，也不能丢掉 create_all_indexes 的成功结果 hybrid_status = "success" hybrid_error = None try: hybrid_result = await asyncio.to_thread(build_hybrid_indexes, driver) if not hybrid_result.get("success", False): hybrid_status = "failed" hybrid_error = hybrid_result.get("message", "未知错误") logger.error(f"[neo4j_indexing] build_hybrid_indexes 失败: {hybrid_error}") except Exception as e: hybrid_status = "failed" hybrid_error = str(e) logger.error(f"[neo4j_indexing] build_hybrid_indexes 异常: {e}", exc_info=True) elapsed = time.time() - start_time # ========== 构建响应 ========== overall_success = result["success"] and hybrid_status == "success" response_data = { "success": overall_success, "message": result["message"], "mode": result.get("mode", "unknown"), "labels": result["labels"], "labels_count": len(result["labels"]), "vector_index_errors": result["vector_index_errors"], "fulltext_index_errors": result["fulltext_index_errors"], "hybrid_index_status": hybrid_status, "hybrid_index_error": hybrid_error, "elapsed_seconds": round(elapsed, 2), } logger.info(f"[neo4j_indexing] 索引创建完成，耗时: {elapsed:.2f}秒") return { "code": 200 if overall_success else 207, # 207 表示部分成功 "message": "success", "data": response_data, } except HTTPException: raise except Exception as e: elapsed = time.time() - start_time error_msg = f"索引创建失败: {str(e)}" logger.error(f"[neo4j_indexing] {error_msg} (耗时: {elapsed:.2f}秒)", exc_info=True) return { "code": 500, "message": "索引创建失败", "data": { "success": False, "message": error_msg, "mode": "error", "labels": [], "labels_count": 0, "vector_index_errors": [], "fulltext_index_errors": [], "hybrid_index_status": "error", "hybrid_index_error": error_msg, "elapsed_seconds": round(elapsed, 2), }, } # ========== ========== ========== ========== ========== ========== # Neo4j 索引创建接口结束 # ========== ========== ========== ========== ========== ========== # ========== ========== ========== ========== ========== ========== # fault 图谱检索接口 # ========== ========== ========== ========== ========== ========== class FaultGraphSearchRequest(BaseModel): """Fault 图谱检索请求模型""" ship_number: Optional[str] = None device_name: Optional[str] = None fault_symptom: Optional[str] = None @app.post("/fault_graph_search") async def fault_graph_search(request: FaultGraphSearchRequest): result = await fault_graph_reason_and_projects(request.ship_number, request.device_name, request.fault_symptom) return result # ========== ========== ========== ========== ========== ========== # fault 图谱检索接口结束 # ========== ========== ========== ========== ========== ========== # ========== ========== ========== ========== ========== ========== # operation 图谱检索接口 # ========== ========== ========== ========== ========== ========== class OperationGraphSearchRequest(BaseModel): """Operation 图谱检索请求模型""" ship_number: Optional[str] = None device_name: Optional[str] = None operation_item: Optional[str] = None @app.post("/operation_graph_search") async def operation_graph_search_endpoint(request: OperationGraphSearchRequest): result = await operation_graph_search(request.ship_number, request.device_name, request.operation_item) return result # ========== ========== ========== ========== ========== ========== # operation 图谱检索接口结束 # ========== ========== ========== ========== ========== ========== # ========== ========== ========== ========== ========== ========== # 设备到系统反向检索接口 # ========== ========== ========== ========== ========== ========== from graph_search.device_to_system import find_system_by_device class DeviceToSystemRequest(BaseModel): """设备到系统反向检索请求模型""" device_name: Optional[str] = None min_hops: Optional[int] = 2 max_hops: Optional[int] = 8 top_k: Optional[int] = 10 @app.post("/device_to_system") async def device_to_system_endpoint(request: DeviceToSystemRequest): """ ========== ========== ========== ========== ========== ========== 设备到系统反向检索接口 功能：根据设备名称，使用混合检索找到设备节点，然后反向遍历找到对应的系统节点 请求参数： - device_name: 设备名称（必填） - min_hops: 最小跳数（默认2级） - max_hops: 最大跳数（默认8级） - top_k: 混合检索返回的候选节点数量（默认10） 返回格式： { "success": true/false, "device": {设备节点信息}, "systems": [系统节点列表], "graph": {节点和边的图谱数据}, "meta": {元数据} } ========== ========== ========== ========== ========== ========== """ result = await find_system_by_device( request.device_name, min_hops=request.min_hops or 2, max_hops=request.max_hops or 8, top_k=request.top_k or 10, driver=driver ) return result # ========== ========== ========== ========== ========== ========== # 设备到系统反向检索接口结束 # ========== ========== ========== ========== ========== ========== # ========== ========== ========== ========== ========== ========== # 图册检索接口 # ========== ========== ========== ========== ========== ========== from graph_search.atlas_retrieval import find_atlas_by_nodes class AtlasRetrievalRequest(BaseModel): """图册检索请求模型""" node_names: Optional[List[str]] = None top_k: Optional[int] = 10 @app.post("/atlas_retrieval") async def atlas_retrieval_endpoint(request: AtlasRetrievalRequest): """ ========== ========== ========== ========== ========== ========== 图册检索接口 功能：根据节点名称（可以是系统、子系统或设备），使用混合检索匹配节点 然后找到一跳关系的图册节点，提取图册名称和PDF URL 请求参数： - node_names: 节点名称列表（可以是系统、子系统或设备，必填） - top_k: 混合检索返回的候选节点数量（默认10） 返回格式（层级结构）： { "success": true/false, "data": { "节点名称1": { "图册名称1": ["PDF_URL1", "PDF_URL2"], "图册名称2": ["PDF_URL3"] }, "节点名称2": { ... } } } ========== ========== ========== ========== ========== ========== """ result = await find_atlas_by_nodes( node_names=request.node_names, top_k=request.top_k or 10, driver=driver ) return result # ========== ========== ========== ========== ========== ========== # 图册检索接口结束 # ========== ========== ========== ========== ========== ========== # ====================================================== # 处理excel解析请求 # ====================================================== @app.post("/api/v2/parse/excel") async def parse_excel(file: UploadFile = File(...), sheet_name: Optional[str] = Form(None)): """ 解析Excel文件的API端点 Args: file (UploadFile): 上传的Excel文件 sheet_name (Optional[str]): 工作表名称 Returns: JSONResponse: 解析结果 """ start_time = datetime.now() try: logger.info(f"Received request to parse Excel file: {file.filename}, sheet: {sheet_name}") # if not file.filename.endswith(".xlsx"): # logger.warning(f"Invalid file type: {file.filename}") # raise HTTPException(status_code=400, detail="仅支持 .xlsx 文件") # with tempfile.NamedTemporaryFile(delete=False, suffix=".xlsx") as tmp: # tmp.write(await file.read()) # tmp_path = tmp.name # logger.info(f"Saved uploaded file to temporary path: {tmp_path}") allowed_extensions = {".xlsx", ".xls"} # 获取文件扩展名，并转为小写 file_ext = os.path.splitext(file.filename)[1].lower() if file_ext not in allowed_extensions: logger.warning(f"Invalid file type: {file.filename}") raise HTTPException(status_code=400, detail="仅支持 .xls 和 .xlsx 文件") with tempfile.NamedTemporaryFile(delete=False, suffix=file_ext) as tmp: tmp.write(await file.read()) tmp_path = tmp.name logger.info(f"Saved uploaded file to temporary path: {tmp_path}") try: df = read_excel_with_merged_cells(tmp_path, sheet_name) rows = dataframe_to_row_json(df) result = transform_excel_json(rows) with open("output.json", "w", encoding="utf-8") as f: json.dump(result, f, ensure_ascii=False, indent=2) md_text = json_to_markdown(result) with open("output.md", "w", encoding="utf-8") as f: f.write(md_text) processing_time = (datetime.now() - start_time).total_seconds() logger.info(f"Successfully processed file in {processing_time:.2f}s, result length: {len(result)}") return JSONResponse( { "message": "success", "data": result, "markdown": md_text, "processing_time": processing_time, "row_count": len(result), } ) except Exception as e: logger.error(f"Error processing Excel file: {e}") logger.error(traceback.format_exc()) raise HTTPException(status_code=500, detail=f"处理Excel文件时发生错误: {str(e)}") finally: if os.path.exists(tmp_path): os.remove(tmp_path) logger.info(f"Removed temporary file: {tmp_path}") except HTTPException: # 重新抛出HTTP异常，不记录为错误 logger.warning(f"HTTP Exception: {traceback.format_exc()}") raise except Exception as e: logger.error(f"Unexpected error in parse_excel: {e}") logger.error(traceback.format_exc()) raise HTTPException(status_code=500, detail=f"服务器内部错误: {str(e)}") # ================== 请求模型 ================== class CypherQueryRequest(BaseModel): """Cypher 图谱问答请求模型""" query: Optional[str] = None use_rerank: Optional[bool] = False top_k: Optional[int] = 70 fulltext_top_k: Optional[int] = 15 rerank_top_k: Optional[int] = 15 hops: Optional[int] = 3 verbose: Optional[bool] = False # 接口默认关闭冗长日志 @app.post("/search_graph") async def search_graph(request: CypherQueryRequest): """ Cypher 图谱问答接口 """ start_time = time.time() query = request.query # 1. 基础校验 if not query or not query.strip(): raise HTTPException(status_code=400, detail="查询文本不能为空") query = query.strip() logger.info(f"[query_cypher] 请求: query='{query}', use_rerank={getattr(request, 'use_rerank', False)}") # ========== 检查资源是否就绪 ========== sync_driver = _resources.get("sync_driver") async_driver = _resources.get("async_driver") embedder = _resources.get("embedder") if sync_driver is None or async_driver is None or embedder is None: raise HTTPException(status_code=503, detail="图谱服务未初始化，请稍后重试") # 初始化变量，防止 except 中引用未定义变量 final_cypher = "" elapsed_ms = 0.0 nodes = [] links = [] text_results = "" try: # ========== ① 实体检索 ========== retrieval = await retrieve_entities( query=query, sync_driver=sync_driver, async_driver=async_driver, embedder=embedder, hybrid_top_k=getattr(request, 'top_k', 40), fulltext_top_k=getattr(request, 'fulltext_top_k', 15), rerank_top_k=getattr(request, 'rerank_top_k', 15), use_rerank=getattr(request, 'use_rerank', False), hops=getattr(request, 'hops', 3), verbose=getattr(request, 'verbose', False), ) # ========== ② 增加已提取实体的一跳关系提取 ========== relations = await fetch_relations_among_results( search_results=retrieval["search_results"], async_driver=async_driver, verbose=True, ) formatted_relations = format_relations(relations) retrieval["formatted_relations"] = formatted_relations # ========== ② Cypher 生成 + 修复 + 查询 ========== result = await query_with_cypher( query=query, retrieval_result=retrieval, driver=async_driver, verbose=getattr(request, 'verbose', False), ) # ========== ③ 执行 Cypher 获取原始数据 ========== final_cypher = result.get("final_cypher", result.get("original_cypher", "")) # 如果 Cypher 生成失败或为空 if not final_cypher: elapsed_ms = round((time.time() - start_time) * 1000, 2) return { "code": 200, "message": "success", "data": { "nodes": [], "links": [], "results": "未能生成有效的 Cypher 查询语句" }, "meta": { "elapsed_ms": elapsed_ms, "result_count": 0, "cypher": "", "query_classification": result.get("query_classification") } } logger.info(f"[query_cypher] 准备执行 Cypher: {final_cypher}") # 执行 Cypher excute_result = None try: # 假设 excute_cypher 是异步的，传入 async_driver excute_result = await excute_cypher(async_driver, final_cypher) except Exception as e: logger.error(f"[query_cypher] excute_cypher 调用失败: {e}，尝试降级为直接执行") # 降级方案：直接使用 async_driver 执行原始 Cypher try: records, _, _ = await async_driver.execute_query(final_cypher) # 将 Record 转换为 dict excute_result = [r.data() for r in records] except Exception as e2: logger.error(f"[query_cypher] 降级执行也失败: {e2}") excute_result = [] logger.info(f"[query_cypher] 执行结果类型: {type(excute_result)}") # ========== ④ 数据处理与格式化 (健壮性核心) ========== # A. 尝试从 result (LLM/GraphService 输出) 中提取结构化图数据 # 优先使用 query_with_cypher 返回的可能存在的图结构 if isinstance(result.get("data"), dict): graph_data = result["data"] nodes = graph_data.get("nodes", []) links = graph_data.get("links", []) text_results = graph_data.get("results", "") # B. 如果 result 中没有图数据，则处理 excute_result # 注意：这里需要兼容 excute_result 的不同返回格式 if not nodes and not links: if isinstance(excute_result, dict): # 情况1: excute_result 已经是格式化好的 {'data': {'nodes':..., 'links':..., 'results':...}} if "data" in excute_result and isinstance(excute_result["data"], dict): data_part = excute_result["data"] nodes = data_part.get("nodes", []) links = data_part.get("links", []) # results 可能是字符串或列表 res_val = data_part.get("results", "") if isinstance(res_val, list): text_results = json.dumps(res_val, ensure_ascii=False, indent=2) else: text_results = str(res_val) # 情况2: excute_result 是 {'data': [rows]} 这种普通数据库返回 elif "data" in excute_result and isinstance(excute_result["data"], list): rows_data = excute_result["data"] if rows_data: text_results = json.dumps(rows_data[:50], ensure_ascii=False, indent=2) else: text_results = "查询无结果" else: # 其他字典结构，直接序列化 text_results = json.dumps(excute_result, ensure_ascii=False, indent=2)[:2000] elif isinstance(excute_result, list): # 情况3: 直接返回列表 if excute_result: text_results = json.dumps(excute_result[:50], ensure_ascii=False, indent=2) else: text_results = "查询无结果" else: text_results = str(excute_result) if excute_result else "查询无结果" # C. 计算耗时 elapsed_ms = round((time.time() - start_time) * 1000, 2) # D. 构建 Meta # 计算 result_count count_val = 0 if nodes: count_val = len(nodes) elif links: count_val = len(links) elif isinstance(excute_result, list): count_val = len(excute_result) elif isinstance(excute_result, dict) and "data" in excute_result and isinstance(excute_result["data"], list): count_val = len(excute_result["data"]) meta = { "elapsed_ms": elapsed_ms, "result_count": count_val, "cypher": final_cypher, "query_type": "graph_search" if (nodes or links) else "cypher_query" } if result.get("query_classification"): meta["query_classification"] = result.get("query_classification") # ========== ⑤ 最终返回 ========== return { "code": 200, "message": "success", "data": { "nodes": nodes, "links": links, "results": text_results, }, "meta": meta, } except HTTPException: raise except Exception as e: # 确保 elapsed_ms 有值 if elapsed_ms == 0.0: elapsed_ms = round((time.time() - start_time) * 1000, 2) logger.exception(f"[query_cypher] Critical Error: {str(e)}") return { "code": 500, "message": f"Query failed: {str(e)}", "data": { "nodes": [], "links": [], "results": "" }, "meta": { "elapsed_ms": elapsed_ms, "result_count": 0, "cypher": final_cypher, "error": str(e) } } # 接口：读取 kg_build 配置缓存 @app.get("/kg_config_cache") async def get_kg_config_cache(): """返回最近一次通过 `/kg_build` 保存的配置缓存（kg_config_cache.json）。""" try: cache_path = os.path.join(DATA_DIR, "kg_config_cache.json") if not os.path.exists(cache_path): # 回退到代码目录中的缓存文件（如果 DATA_DIR 不可写或不存在） cache_path = os.path.join(os.path.dirname(__file__), "kg_config_cache.json") if not os.path.exists(cache_path): return JSONResponse({"code": 404, "message": "kg_config_cache.json 未找到", "data": None}, status_code=404) async with aiofiles.open(cache_path, "r", encoding="utf-8") as f: raw = await f.read() try: data = json.loads(raw) except Exception as e: return JSONResponse({"code": 500, "message": f"解析缓存文件失败: {e}", "data": None}, status_code=500) return {"code": 200, "message": "success", "data": data} except Exception as e: return JSONResponse({"code": 500, "message": str(e), "data": None}, status_code=500) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=9085)