2926 lines
115 KiB
Python
2926 lines
115 KiB
Python
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 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知识图谱混合索引创建接口
|
||
|
||
当前版本：
|
||
只创建 Searchable 混合索引体系：
|
||
|
||
1. searchable_vector
|
||
- embedding语义检索
|
||
|
||
2. searchable_name_fulltext
|
||
- 名称关键词检索
|
||
|
||
3. searchable_fulltext
|
||
- 内容关键词检索
|
||
"""
|
||
|
||
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":"已有Neo4j索引创建任务正在执行",
|
||
"data":None
|
||
}
|
||
)
|
||
async with _neo4j_indexing_lock:
|
||
try:
|
||
if not INDEXING_AVAILABLE:
|
||
raise HTTPException(status_code=503,detail="Neo4j索引模块不可用")
|
||
kwargs = {
|
||
"drop_old": True,
|
||
}
|
||
# 如果接口传入连接参数，则覆盖
|
||
if request.neo4j_uri:
|
||
kwargs["uri"] = request.neo4j_uri
|
||
if request.neo4j_username:
|
||
kwargs["user"] = request.neo4j_username
|
||
if request.neo4j_password:
|
||
kwargs["password"] = request.neo4j_password
|
||
|
||
result = await asyncio.to_thread(build_hybrid_indexes,**kwargs)
|
||
elapsed = time.time()-start_time
|
||
if result.get("success"):
|
||
logger.info(
|
||
f"[neo4j_indexing] 创建成功 "
|
||
f"耗时:{elapsed:.2f}s")
|
||
return {
|
||
"code":200,
|
||
"message":"success",
|
||
"data":{
|
||
"success":True,
|
||
"message":
|
||
result.get("message","混合索引创建完成"),
|
||
"searchable_added":
|
||
result.get("searchable_added",0),
|
||
"elapsed_seconds":
|
||
round(elapsed,2)
|
||
}
|
||
}
|
||
else:
|
||
logger.error(
|
||
f"[neo4j_indexing] 创建失败:"
|
||
f"{result.get('message')}"
|
||
)
|
||
return {
|
||
"code":207,
|
||
"message":"partial success",
|
||
"data":{
|
||
"success":False,
|
||
"message":
|
||
result.get("message","未知错误"),
|
||
"searchable_added":0,
|
||
"elapsed_seconds":round(elapsed,2)
|
||
}
|
||
}
|
||
|
||
except HTTPException:
|
||
raise
|
||
|
||
except Exception as e:
|
||
elapsed=time.time()-start_time
|
||
error_msg=f"索引创建失败:{str(e)}"
|
||
logger.error(error_msg,exc_info=True)
|
||
return {
|
||
"code":500,
|
||
"message":"索引创建失败",
|
||
"data":{
|
||
"success":False,
|
||
"message":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)
|
||
|
||
|
||
|