Add KGRAG code and image backup
This commit is contained in:
commit
cdb8193ed5
42
.env
Normal file
42
.env
Normal file
@ -0,0 +1,42 @@
|
||||
#大模型API地址
|
||||
OPENAI_API_BASE=http://192.168.0.46:59800/v1
|
||||
#大模型API密钥
|
||||
OPENAI_API_KEY=gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5
|
||||
#Embedding 模型地址 :
|
||||
OPENAI_API_EMBEDDING_BASE=http://172.18.30.122:40044/v1
|
||||
# OPENAI_API_EMBEDDING_BASE=http://172.18.127.124:40050/v1
|
||||
#Embedding 模型名称
|
||||
OPENAI_EMBEDDING_MODEL=bge-m3
|
||||
# OPENAI_EMBEDDING_MODEL=Qwen3-Embedding
|
||||
#大模型名称
|
||||
# OPENAI_MODEL=Qwen3.5-35B-A3B
|
||||
OPENAI_MODEL=Qwen3-32B
|
||||
#重排序模型地址
|
||||
OPENAI_API_RERANKER_BASE=http://172.18.30.122:40014/v1
|
||||
# OPENAI_API_RERANKER_BASE=http://172.18.127.124:40022/v1
|
||||
#重排序模型名称
|
||||
OPENAI_RERANKER_MODEL=bge-rerank
|
||||
# OPENAI_RERANKER_MODEL=Qwen3-Reranker-8B
|
||||
#向量数据库用户名
|
||||
Milvus_USER=user
|
||||
#向量数据库密码
|
||||
Milvus_PASSWORD=password
|
||||
#向量数据库地址
|
||||
Milvus_URI=http://192.168.0.46:59530
|
||||
#文档解析
|
||||
Mineru_URL=http://172.18.30.122:9988
|
||||
#word 转pdf
|
||||
CONVERT_DOC_URL=http://172.18.30.122:9000
|
||||
#3D_模型召回 置信度
|
||||
3D_RECALL_SCORE = 0.3
|
||||
#服务器所在地址
|
||||
HOST_IP=http://172.18.30.122:9988
|
||||
|
||||
# neo4j
|
||||
NEO4J_URI=bolt://192.168.0.46:57687
|
||||
NEO4J_USERNAME=neo4j
|
||||
#NEO4J_PASSWORD=123456hdf
|
||||
NEO4J_PASSWORD=zdht123@
|
||||
//NEO4J_URI=neo4j://192.168.0.46:7687
|
||||
//NEO4J_USER=neo4j
|
||||
//NEO4J_PASSWORD=123456hdf
|
||||
22
.env.example
Normal file
22
.env.example
Normal file
@ -0,0 +1,22 @@
|
||||
#大模型API地址
|
||||
OPENAI_API_BASE=http://172.18.127.124:40058/v1
|
||||
#大模型API密钥
|
||||
OPENAI_API_KEY=gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5
|
||||
#Embedding 模型地址 :
|
||||
OPENAI_API_EMBEDDING_BASE=http://172.18.127.124:40023/v1
|
||||
#Embedding 模型名称
|
||||
OPENAI_EMBEDDING_MODEL=bge-m3
|
||||
#大模型名称
|
||||
OPENAI_MODEL=qwen3-32b-gguf
|
||||
#重排序模型地址
|
||||
OPENAI_API_RERANKER_BASE=http://172.18.30.122:20507/v1
|
||||
#重排序模型名称
|
||||
OPENAI_RERANKER_MODEL=bge-rerank
|
||||
#向量数据库用户名
|
||||
Milvus_USER=user
|
||||
#向量数据库密码
|
||||
Milvus_PASSWORD=password
|
||||
#向量数据库地址
|
||||
Milvus_URI=http://172.18.30.122:19530
|
||||
#文档解析
|
||||
Mineru_URL=http://172.18.30.122:9988
|
||||
1
.gitattributes
vendored
Normal file
1
.gitattributes
vendored
Normal file
@ -0,0 +1 @@
|
||||
docker-images/*.tar.gz filter=lfs diff=lfs merge=lfs -text
|
||||
21
.gitignore
vendored
Normal file
21
.gitignore
vendored
Normal file
@ -0,0 +1,21 @@
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
.pytest_cache/
|
||||
.cache/
|
||||
|
||||
*.log
|
||||
*.tmp
|
||||
*.swp
|
||||
.DS_Store
|
||||
.nfs*
|
||||
|
||||
files/
|
||||
files_part1/
|
||||
mineru_output/
|
||||
kg_output/
|
||||
parsefiles/
|
||||
upload/
|
||||
|
||||
output*.json
|
||||
output*.txt
|
||||
result_tmp.json
|
||||
9636
122-06A0014-B01001_发动机-维修手册_content_list.json
Normal file
9636
122-06A0014-B01001_发动机-维修手册_content_list.json
Normal file
File diff suppressed because one or more lines are too long
BIN
122-06A0014-B01003_信号灯-操作使用手册.pdf
Normal file
BIN
122-06A0014-B01003_信号灯-操作使用手册.pdf
Normal file
Binary file not shown.
BIN
122-06A0014-B01003_雷达-使用说明书.docx
Normal file
BIN
122-06A0014-B01003_雷达-使用说明书.docx
Normal file
Binary file not shown.
BIN
122-06A0014-B01003_雷达-使用说明书.pdf
Normal file
BIN
122-06A0014-B01003_雷达-使用说明书.pdf
Normal file
Binary file not shown.
37
README.md
Normal file
37
README.md
Normal file
@ -0,0 +1,37 @@
|
||||
# KGRAG ???????
|
||||
|
||||
?????`192.168.1.108`
|
||||
|
||||
???????`/home/workspace/kgrag`
|
||||
|
||||
?????`kgrag:arm64`
|
||||
|
||||
## ????
|
||||
|
||||
- `app.py`?KGRAG ?????
|
||||
- `requirements.txt`?Python ??
|
||||
- `start.sh` / `stop.sh`????????
|
||||
- `docker-images/kgrag_arm64.docker-save.tar.gz`?`kgrag:arm64` ????
|
||||
|
||||
## ?????
|
||||
|
||||
??????????????????????? `.git` ????? `.bak` ??????????????
|
||||
|
||||
## ?????
|
||||
|
||||
????? Git LFS ?? Docker ????
|
||||
|
||||
????????
|
||||
|
||||
```bash
|
||||
git lfs install
|
||||
git lfs pull
|
||||
```
|
||||
|
||||
????????? `version https://git-lfs.github.com/spec/v1`??????? LFS ?????????? `git lfs pull` ???????
|
||||
|
||||
## ????
|
||||
|
||||
```bash
|
||||
gunzip -c docker-images/kgrag_arm64.docker-save.tar.gz | docker load
|
||||
```
|
||||
0
__init__.py
Normal file
0
__init__.py
Normal file
BIN
aiofiles-25.1.0-py3-none-any.whl
Normal file
BIN
aiofiles-25.1.0-py3-none-any.whl
Normal file
Binary file not shown.
32
app_pkg/.env
Normal file
32
app_pkg/.env
Normal file
@ -0,0 +1,32 @@
|
||||
#大模型API地址
|
||||
OPENAI_API_BASE=http://172.18.127.124:40063/v1
|
||||
#大模型API密钥
|
||||
OPENAI_API_KEY=gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5
|
||||
#Embedding 模型地址 :
|
||||
OPENAI_API_EMBEDDING_BASE=http://172.18.30.122:40044/v1
|
||||
# OPENAI_API_EMBEDDING_BASE=http://172.18.127.124:40050/v1
|
||||
#Embedding 模型名称
|
||||
OPENAI_EMBEDDING_MODEL=bge-m3
|
||||
# OPENAI_EMBEDDING_MODEL=Qwen3-Embedding
|
||||
#大模型名称
|
||||
OPENAI_MODEL=Qwen3-32B
|
||||
#重排序模型地址
|
||||
OPENAI_API_RERANKER_BASE=http://172.18.30.122:40014/v1
|
||||
# OPENAI_API_RERANKER_BASE=http://172.18.127.124:40022/v1
|
||||
#重排序模型名称
|
||||
OPENAI_RERANKER_MODEL=bge-rerank
|
||||
# OPENAI_RERANKER_MODEL=Qwen3-Reranker-8B
|
||||
#向量数据库用户名
|
||||
Milvus_USER=user
|
||||
#向量数据库密码
|
||||
Milvus_PASSWORD=password
|
||||
#向量数据库地址
|
||||
Milvus_URI=http://192.168.0.46:59530
|
||||
#文档解析
|
||||
Mineru_URL=http://172.18.30.122:9988
|
||||
#word 转pdf
|
||||
CONVERT_DOC_URL=http://172.18.30.122:9000
|
||||
#3D_模型召回 置信度
|
||||
3D_RECALL_SCORE = 0.3
|
||||
#服务器所在地址
|
||||
HOST_IP=http://172.18.30.122:9988
|
||||
22
app_pkg/.env.example
Normal file
22
app_pkg/.env.example
Normal file
@ -0,0 +1,22 @@
|
||||
#大模型API地址
|
||||
OPENAI_API_BASE=http://172.18.127.124:40058/v1
|
||||
#大模型API密钥
|
||||
OPENAI_API_KEY=gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5
|
||||
#Embedding 模型地址 :
|
||||
OPENAI_API_EMBEDDING_BASE=http://172.18.127.124:40023/v1
|
||||
#Embedding 模型名称
|
||||
OPENAI_EMBEDDING_MODEL=bge-m3
|
||||
#大模型名称
|
||||
OPENAI_MODEL=qwen3-32b-gguf
|
||||
#重排序模型地址
|
||||
OPENAI_API_RERANKER_BASE=http://172.18.30.122:20507/v1
|
||||
#重排序模型名称
|
||||
OPENAI_RERANKER_MODEL=bge-rerank
|
||||
#向量数据库用户名
|
||||
Milvus_USER=user
|
||||
#向量数据库密码
|
||||
Milvus_PASSWORD=password
|
||||
#向量数据库地址
|
||||
Milvus_URI=http://172.18.30.122:19530
|
||||
#文档解析
|
||||
Mineru_URL=http://172.18.30.122:9988
|
||||
0
app_pkg/__init__.py
Normal file
0
app_pkg/__init__.py
Normal file
1358
app_pkg/app.py
Normal file
1358
app_pkg/app.py
Normal file
File diff suppressed because it is too large
Load Diff
62
app_pkg/ceshi.py
Normal file
62
app_pkg/ceshi.py
Normal file
@ -0,0 +1,62 @@
|
||||
from phrasetree.tree import Tree
|
||||
import hanlp
|
||||
import requests
|
||||
content = """
|
||||
汽车底盘主要包括以下几个部分:
|
||||
1. **车架**:作为汽车的基础骨架,支撑着发动机、车身和其它部件。
|
||||
2. **悬架系统**:包括弹簧和减震器等,用于缓冲路面不平带来的震动。
|
||||
3. **转向系统**:使汽车能够按照驾驶员的意愿改变行驶方向。
|
||||
4. **制动系统**:用于减速或停车,包括制动踏板、制动总泵、制动片等。
|
||||
5. **传动系统**:将发动机的动力传递给车轮,包括离合器、变速器、驱动轴等。
|
||||
例如,陕汽德龙X5000 重卡的底盘采用了高强度轻量化设计的车架及悬架系统,通过低碳合金钢和 CAE 模拟优化设计,等宽直大梁车架结构和少片变截面钢板弹簧结构设计,屈服强度提高近 $50\%$,承载力、稳定性、平顺性远优于国内同级车型。底盘应用大量轻量化材料及工艺,整车轻量化水平行业领先。
|
||||
相关图示:
|
||||
"""
|
||||
print("正在加载 HanLP 模型...")
|
||||
HanLPmodel = hanlp.load(hanlp.pretrained.mtl.CLOSE_TOK_POS_NER_SRL_DEP_SDP_CON_ELECTRA_SMALL_ZH)
|
||||
|
||||
def extract_unique_nouns_from_file(content, model):
|
||||
"""
|
||||
从指定的文本文件中提取所有唯一的名词(NN),去除 '**' 并去重,保持首次出现顺序。
|
||||
|
||||
参数:
|
||||
file_path (str): 输入文本文件的路径。
|
||||
model: HanLP 预训练模型实例(可选)。如果未提供,则自动加载默认模型。
|
||||
|
||||
返回:
|
||||
list: 提取到的唯一名词字符串列表,如 ['转向节', '螺栓', ...]
|
||||
"""
|
||||
|
||||
# 1. 使用 HanLP 进行句法分析,获取成分句法树(con)
|
||||
doc = model(content)
|
||||
|
||||
# 2. 定义递归函数提取所有叶子节点为字符串的 NN 节点
|
||||
def extract_nn(tree):
|
||||
res = []
|
||||
if isinstance(tree, Tree):
|
||||
if tree.label() == 'NN' and len(tree) == 1 and isinstance(tree[0], str):
|
||||
res.append(['NN', [tree[0]]])
|
||||
else:
|
||||
for child in tree:
|
||||
res.extend(extract_nn(child))
|
||||
return res
|
||||
|
||||
# 3. 提取所有 NN 节点
|
||||
nn_list = extract_nn(doc['con'])
|
||||
print(f'共提取到 NN 节点:{len(nn_list)}')
|
||||
|
||||
# 4. 过滤掉词为 '**' 的项
|
||||
filtered = [rec for rec in nn_list if rec[1][0] != '**']
|
||||
|
||||
# 5. 去重,保持顺序
|
||||
seen = set()
|
||||
unique_nn = []
|
||||
for rec in filtered:
|
||||
key = (rec[0], rec[1][0]) # 使用 (标签, 词语) 作为唯一键
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique_nn.append(rec[1][0])
|
||||
|
||||
print(f'共提取到 NN 节点(去重并过滤 ** 后):{len(unique_nn)}')
|
||||
return unique_nn
|
||||
fenci_word = extract_unique_nouns_from_file(content,model=HanLPmodel)
|
||||
print(fenci_word)
|
||||
48
app_pkg/contentprocess/contentprocess.py
Normal file
48
app_pkg/contentprocess/contentprocess.py
Normal file
@ -0,0 +1,48 @@
|
||||
from phrasetree.tree import Tree
|
||||
import hanlp
|
||||
|
||||
def extract_unique_nouns_from_file(content, model):
|
||||
"""
|
||||
从指定的文本文件中提取所有唯一的名词(NN),去除 '**' 并去重,保持首次出现顺序。
|
||||
|
||||
参数:
|
||||
file_path (str): 输入文本文件的路径。
|
||||
model: HanLP 预训练模型实例(可选)。如果未提供,则自动加载默认模型。
|
||||
|
||||
返回:
|
||||
list: 提取到的唯一名词字符串列表,如 ['转向节', '螺栓', ...]
|
||||
"""
|
||||
|
||||
# 1. 使用 HanLP 进行句法分析,获取成分句法树(con)
|
||||
doc = model(content)
|
||||
|
||||
# 2. 定义递归函数提取所有叶子节点为字符串的 NN 节点
|
||||
def extract_nn(tree):
|
||||
res = []
|
||||
if isinstance(tree, Tree):
|
||||
if tree.label() == 'NN' and len(tree) == 1 and isinstance(tree[0], str):
|
||||
res.append(['NN', [tree[0]]])
|
||||
else:
|
||||
for child in tree:
|
||||
res.extend(extract_nn(child))
|
||||
return res
|
||||
|
||||
# 3. 提取所有 NN 节点
|
||||
nn_list = extract_nn(doc['con'])
|
||||
# print(f'共提取到 NN 节点:{len(nn_list)}')
|
||||
|
||||
# 4. 过滤掉词为 '**' 的项
|
||||
filtered = [rec for rec in nn_list if rec[1][0] != '**']
|
||||
|
||||
# 5. 去重,保持顺序
|
||||
seen = set()
|
||||
unique_nn = []
|
||||
for rec in filtered:
|
||||
key = (rec[0], rec[1][0]) # 使用 (标签, 词语) 作为唯一键
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique_nn.append(rec[1][0])
|
||||
|
||||
print(f'共提取到 (去重并过滤后)关键词:{len(unique_nn)}')
|
||||
print(f'输出提取到的关键词:{unique_nn}')
|
||||
return unique_nn
|
||||
BIN
app_pkg/converted_pdfs/122-06a0014-b01001_发动机-图解目录_0f8631ec.pdf
Normal file
BIN
app_pkg/converted_pdfs/122-06a0014-b01001_发动机-图解目录_0f8631ec.pdf
Normal file
Binary file not shown.
BIN
app_pkg/converted_pdfs/122-06a0014-b01001_发动机-图解目录_5b08bbf4.pdf
Normal file
BIN
app_pkg/converted_pdfs/122-06a0014-b01001_发动机-图解目录_5b08bbf4.pdf
Normal file
Binary file not shown.
BIN
app_pkg/converted_pdfs/122-06a0014-b01001_发动机-图解目录_7105b60e.pdf
Normal file
BIN
app_pkg/converted_pdfs/122-06a0014-b01001_发动机-图解目录_7105b60e.pdf
Normal file
Binary file not shown.
BIN
app_pkg/converted_pdfs/122-06a0014-b01001_发动机-图解目录_c70d577c.pdf
Normal file
BIN
app_pkg/converted_pdfs/122-06a0014-b01001_发动机-图解目录_c70d577c.pdf
Normal file
Binary file not shown.
BIN
app_pkg/converted_pdfs/122-06a0014-b01001_发动机-图解目录_dc1d1aa0.pdf
Normal file
BIN
app_pkg/converted_pdfs/122-06a0014-b01001_发动机-图解目录_dc1d1aa0.pdf
Normal file
Binary file not shown.
BIN
app_pkg/converted_pdfs/图查询界面接口说明_ded8f0cc.pdf
Normal file
BIN
app_pkg/converted_pdfs/图查询界面接口说明_ded8f0cc.pdf
Normal file
Binary file not shown.
0
app_pkg/database/__init__.py
Normal file
0
app_pkg/database/__init__.py
Normal file
860
app_pkg/database/milvus_api.py
Normal file
860
app_pkg/database/milvus_api.py
Normal file
@ -0,0 +1,860 @@
|
||||
import json
|
||||
import os
|
||||
from pymilvus import (
|
||||
connections,
|
||||
utility,
|
||||
FieldSchema,
|
||||
CollectionSchema,
|
||||
DataType,
|
||||
Collection,
|
||||
MilvusClient,
|
||||
db,
|
||||
Function,
|
||||
FunctionType,
|
||||
AnnSearchRequest,
|
||||
WeightedRanker
|
||||
)
|
||||
from typing import List, Dict, Union, Optional, Any
|
||||
from pypinyin import lazy_pinyin
|
||||
from typing import List, Dict, Tuple, Optional
|
||||
|
||||
|
||||
class Milvus_Database:
|
||||
def __init__(self, user=os.getenv("Milvus_USER", "root"), password=os.getenv("Milvus_PASSWORD", "Milvus"),
|
||||
uri=os.getenv("Milvus_URI")):
|
||||
self.client = MilvusClient(
|
||||
uri=uri,
|
||||
token=f"{user}:{password}"
|
||||
)
|
||||
# self.conn = connections.connect(host="172.18.127.124",port=19530,user=user,password=password)
|
||||
|
||||
# database = db.create_database(database_name)
|
||||
# self.client.list_databases()
|
||||
|
||||
def create_database(self, database_name) -> List[str]:
|
||||
database = self.client.create_database(database_name)
|
||||
return self.client.list_databases()
|
||||
|
||||
def list_databases(self) -> List[str]:
|
||||
return self.client.list_databases()
|
||||
|
||||
def using_database(self, database_name):
|
||||
self.client.use_database(database_name)
|
||||
|
||||
def create_hierarchical_collection(self, collection_name, database_name):
|
||||
self.client.use_database(database_name)
|
||||
schema = MilvusClient.create_schema(
|
||||
auto_id=True,
|
||||
enable_dynamic_field=True,
|
||||
)
|
||||
|
||||
analyzer_params_custom = {
|
||||
"tokenizer": "jieba",
|
||||
"type": "chinese",
|
||||
"filter": ["cnalphanumonly"]
|
||||
}
|
||||
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
|
||||
# 层次召回中的摘要部分
|
||||
schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=65535, enable_analyzer=True,
|
||||
analyzer_params=analyzer_params_custom)
|
||||
# 层次召回中的完整文档
|
||||
schema.add_field(field_name="origin_text", datatype=DataType.VARCHAR, max_length=65535, enable_analyzer=True,
|
||||
analyzer_params=analyzer_params_custom)
|
||||
schema.add_field(field_name="bm_25", datatype=DataType.SPARSE_FLOAT_VECTOR)
|
||||
# Qwen3-8B-Embedding 维度4096
|
||||
schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=4096, )
|
||||
schema.add_field(field_name="resource", datatype=DataType.VARCHAR, max_length=1024, default_value='EMPTY'),
|
||||
schema.add_field(field_name="img_path", datatype=DataType.VARCHAR, max_length=4096, default_value="EMPTY")
|
||||
schema.add_field(field_name="page_idx", datatype=DataType.INT64, default_value=0)
|
||||
schema.add_field(field_name="origin_id", datatype=DataType.VARCHAR, max_length=256, default_value="EMPTY")
|
||||
# schema.add_field(field_name="img_information",datatype=DataType.Json,default_value="EMPTY")
|
||||
schema.add_field(field_name="header_1", datatype=DataType.VARCHAR, max_length=1024, default_value='EMPTY')
|
||||
schema.add_field(field_name="header_2", datatype=DataType.VARCHAR, max_length=1024, default_value='EMPTY')
|
||||
schema.add_field(field_name="header_3", datatype=DataType.VARCHAR, max_length=1024, default_value='EMPTY')
|
||||
|
||||
bm25_function = Function(
|
||||
name="bm25_emb", # Function name
|
||||
input_field_names=["text"], # Name of the VARCHAR field containing raw text data
|
||||
output_field_names=["bm_25"],
|
||||
# Name of the SPARSE_FLOAT_VECTOR field reserved to store generated embeddings
|
||||
function_type=FunctionType.BM25,
|
||||
)
|
||||
schema.add_function(bm25_function)
|
||||
index_params = self.client.prepare_index_params()
|
||||
# Add indexes
|
||||
index_params.add_index(
|
||||
field_name="vector",
|
||||
index_name="vector_index",
|
||||
index_type="HNSW", # "IVF_FLAT". AUTOINDEX
|
||||
metric_type="IP", # COSINE
|
||||
# params={"nlist": 128},
|
||||
)
|
||||
|
||||
index_params.add_index(
|
||||
field_name="bm_25",
|
||||
index_name="bm_25_index",
|
||||
index_type="SPARSE_INVERTED_INDEX", # Index type for sparse vectors
|
||||
metric_type="BM25", # Set to `BM25` when using function to generate sparse vectors
|
||||
params={"inverted_index_algo": "DAAT_MAXSCORE"},
|
||||
# The ratio of small vector values to be dropped during indexing
|
||||
)
|
||||
|
||||
self.client.create_collection(
|
||||
collection_name=collection_name,
|
||||
schema=schema,
|
||||
index_params=index_params
|
||||
)
|
||||
self.release_collection(collection_name)
|
||||
return "collections created"
|
||||
|
||||
def create_hybrid_collection(self, collection_name, database_name: str, save_title: bool = True):
|
||||
"""
|
||||
中文专用混合检索 表
|
||||
"""
|
||||
self.client.use_database(database_name)
|
||||
schema = MilvusClient.create_schema(
|
||||
auto_id=True,
|
||||
enable_dynamic_field=True,
|
||||
)
|
||||
|
||||
analyzer_params_custom = {
|
||||
"tokenizer": "jieba",
|
||||
"type": "chinese",
|
||||
"filter": ["cnalphanumonly"]
|
||||
}
|
||||
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
|
||||
schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=65535, enable_analyzer=True,
|
||||
analyzer_params=analyzer_params_custom)
|
||||
schema.add_field(field_name="origin_text", datatype=DataType.VARCHAR, max_length=65535, default_value="EMPTY")
|
||||
schema.add_field(field_name="bm_25", datatype=DataType.SPARSE_FLOAT_VECTOR)
|
||||
schema.add_field(field_name="vector", datatype=DataType.FLOAT_VECTOR, dim=1024, )
|
||||
schema.add_field(field_name="resource", datatype=DataType.VARCHAR, max_length=1024, default_value='EMPTY')
|
||||
schema.add_field(field_name="img_path", datatype=DataType.VARCHAR, max_length=4096, default_value="EMPTY")
|
||||
schema.add_field(field_name="page_idx", datatype=DataType.INT64, default_value=0)
|
||||
schema.add_field(field_name="origin_id", datatype=DataType.VARCHAR, max_length=256, default_value="EMPTY")
|
||||
# schema.add_field(field_name="img_information",datatype=DataType.Json,default_value="EMPTY")
|
||||
if save_title:
|
||||
schema.add_field(field_name="header_1", datatype=DataType.VARCHAR, max_length=1024, default_value='EMPTY')
|
||||
schema.add_field(field_name="header_2", datatype=DataType.VARCHAR, max_length=1024, default_value='EMPTY')
|
||||
schema.add_field(field_name="header_3", datatype=DataType.VARCHAR, max_length=1024, default_value='EMPTY')
|
||||
|
||||
bm25_function = Function(
|
||||
name="bm25_emb", # Function name
|
||||
input_field_names=["text"], # Name of the VARCHAR field containing raw text data
|
||||
output_field_names=["bm_25"],
|
||||
# Name of the SPARSE_FLOAT_VECTOR field reserved to store generated embeddings
|
||||
function_type=FunctionType.BM25,
|
||||
)
|
||||
schema.add_function(bm25_function)
|
||||
index_params = self.client.prepare_index_params()
|
||||
# Add indexes
|
||||
index_params.add_index(
|
||||
field_name="vector",
|
||||
index_name="vector_index",
|
||||
index_type="HNSW", # "IVF_FLAT". AUTOINDEX
|
||||
metric_type="IP", # COSINE
|
||||
# params={"nlist": 128},
|
||||
)
|
||||
|
||||
index_params.add_index(
|
||||
field_name="bm_25",
|
||||
index_name="bm_25_index",
|
||||
index_type="SPARSE_INVERTED_INDEX", # Index type for sparse vectors
|
||||
metric_type="BM25", # Set to `BM25` when using function to generate sparse vectors
|
||||
params={"inverted_index_algo": "DAAT_MAXSCORE"},
|
||||
# The ratio of small vector values to be dropped during indexing
|
||||
)
|
||||
|
||||
self.client.create_collection(
|
||||
collection_name=collection_name,
|
||||
schema=schema,
|
||||
index_params=index_params
|
||||
)
|
||||
self.release_collection(collection_name)
|
||||
return "collections created"
|
||||
|
||||
def drop_collection(self, collection_name, database_name: str):
|
||||
self.client.use_database(database_name)
|
||||
self.load_collection(collection_name=collection_name)
|
||||
self.client.drop_collection(collection_name)
|
||||
|
||||
def load_collection(self, collection_name):
|
||||
self.client.load_collection(collection_name=collection_name)
|
||||
res = self.client.get_load_state(
|
||||
collection_name=collection_name
|
||||
)
|
||||
return res
|
||||
|
||||
def release_collection(self, collection_name):
|
||||
self.client.release_collection(collection_name=collection_name)
|
||||
res = self.client.get_load_state(
|
||||
collection_name=collection_name
|
||||
)
|
||||
return res
|
||||
|
||||
def delete_resource(self, collection_name, resource, database_name):
|
||||
self.client.use_database(database_name)
|
||||
self.load_collection(collection_name)
|
||||
res = self.client.delete(
|
||||
collection_name=collection_name,
|
||||
filter="resource == '{}'".format(resource),
|
||||
)
|
||||
return res
|
||||
|
||||
def has_collection(self, collection_name, database_name):
|
||||
"""
|
||||
Before using list collection, You need to using using_database method first
|
||||
"""
|
||||
self.using_database(database_name=database_name)
|
||||
return self.client.has_collection(collection_name)
|
||||
|
||||
def list_collections(self, database_name):
|
||||
"""
|
||||
Before using list collection, You need to using using_database method first
|
||||
"""
|
||||
self.client.use_database(database_name)
|
||||
return self.client.list_collections()
|
||||
|
||||
def hierarchical_insert(self, params, collection_name: str, database_name: str):
|
||||
try:
|
||||
self.client.use_database(database_name)
|
||||
self.load_collection(collection_name)
|
||||
res = self.client.insert(
|
||||
collection_name=collection_name,
|
||||
data=params
|
||||
)
|
||||
return res
|
||||
finally:
|
||||
try:
|
||||
self.client.release_collection(collection_name)
|
||||
except Exception as e:
|
||||
print(f"释放集合失败: {e}")
|
||||
|
||||
def hybrid_insert(self, params, collection_name: str, database_name: str):
|
||||
try:
|
||||
self.client.use_database(database_name)
|
||||
self.load_collection(collection_name)
|
||||
res = self.client.insert(
|
||||
collection_name=collection_name,
|
||||
data=params
|
||||
)
|
||||
return res
|
||||
finally:
|
||||
try:
|
||||
self.client.release_collection(collection_name)
|
||||
except Exception as e:
|
||||
print(f"释放集合失败: {e}")
|
||||
|
||||
# def insert(self, params,collection_name:str = "MarkdownTest",database_name:str = "JXTest"):
|
||||
# """
|
||||
# params: List[List]
|
||||
# 例子: [[id, header_1, header_2, header_3, text, vector]]
|
||||
# """
|
||||
# # self.client.insert(collection_name=collection_name, data=params)
|
||||
# db.using_database(database_name)
|
||||
# collection = Collection(collection_name)
|
||||
# mr = collection.insert(params)
|
||||
# return mr
|
||||
def get_all_chunk_by_resource(self, collection_name, database_name, resource):
|
||||
"""
|
||||
return List :
|
||||
"""
|
||||
try:
|
||||
self.client.using_database(db_name=database_name)
|
||||
if not self.client.has_collection(collection_name):
|
||||
return []
|
||||
else:
|
||||
self.client.load_collection(collection_name)
|
||||
iterator = self.client.query_iterator(
|
||||
collection_name=collection_name,
|
||||
batch_size=5,
|
||||
filter=f"resource like \"{resource}\"",
|
||||
output_fields=["id", "text", "resource", "Header_1", "Header_2", "Header_3"],
|
||||
)
|
||||
|
||||
results = []
|
||||
|
||||
while True:
|
||||
result = iterator.next()
|
||||
if not result:
|
||||
iterator.close()
|
||||
break
|
||||
results += result
|
||||
return results
|
||||
finally:
|
||||
# 3. 无论成功还是失败,最终释放集合
|
||||
try:
|
||||
self.client.release_collection(collection_name)
|
||||
except Exception as e:
|
||||
print(f"释放集合失败: {e}")
|
||||
|
||||
def get_chunk_by_resource(self, collection_name, database_name, resource):
|
||||
"""
|
||||
return List :
|
||||
"""
|
||||
try:
|
||||
self.client.using_database(db_name=database_name)
|
||||
if not self.client.has_collection(collection_name):
|
||||
return []
|
||||
else:
|
||||
self.client.load_collection(collection_name)
|
||||
iterator = self.client.query_iterator(
|
||||
collection_name=collection_name,
|
||||
batch_size=5,
|
||||
filter=f"resource like \"{resource}\"",
|
||||
output_fields=["id", "text", "resource"],
|
||||
)
|
||||
|
||||
results = []
|
||||
|
||||
while True:
|
||||
result = iterator.next()
|
||||
if not result:
|
||||
iterator.close()
|
||||
break
|
||||
results += result
|
||||
return results
|
||||
finally:
|
||||
# 3. 无论成功还是失败,最终释放集合
|
||||
try:
|
||||
self.client.release_collection(collection_name)
|
||||
except Exception as e:
|
||||
print(f"释放集合失败: {e}")
|
||||
|
||||
def get_chunk_by_id(self, collection_name, database_name, id):
|
||||
"""
|
||||
return List :
|
||||
"""
|
||||
try:
|
||||
self.client.using_database(db_name=database_name)
|
||||
if not self.client.has_collection(collection_name):
|
||||
return []
|
||||
else:
|
||||
self.client.load_collection(collection_name)
|
||||
iterator = self.client.query_iterator(
|
||||
collection_name=collection_name,
|
||||
batch_size=5,
|
||||
filter=f"id=={id + 1}",
|
||||
output_fields=["id", "text", "resource"],
|
||||
)
|
||||
|
||||
results = []
|
||||
|
||||
while True:
|
||||
result = iterator.next()
|
||||
if not result:
|
||||
iterator.close()
|
||||
break
|
||||
results += result
|
||||
return results
|
||||
finally:
|
||||
# 3. 无论成功还是失败,最终释放集合
|
||||
try:
|
||||
self.client.release_collection(collection_name)
|
||||
except Exception as e:
|
||||
print(f"释放集合失败: {e}")
|
||||
|
||||
def get_text_by_resource(self, collection_name, database_name, resource, logger):
|
||||
"""
|
||||
return List :
|
||||
"""
|
||||
|
||||
# def prepare_milvus_filter( resource):
|
||||
# """安全构造 Milvus 过滤条件"""
|
||||
# escaped_resource = json.dumps(resource)[1:-1]
|
||||
# return f'resource="{escaped_resource}"'
|
||||
def prepare_milvus_filter(resource):
|
||||
"""安全构造 Milvus 过滤条件"""
|
||||
# 替换资源名称中的双引号
|
||||
safe_resource = resource.replace('"', '\\"')
|
||||
return f'resource == "{safe_resource}"'
|
||||
|
||||
self.client.using_database(db_name=database_name)
|
||||
|
||||
# 检查集合是否存在
|
||||
if not self.client.has_collection(collection_name):
|
||||
logger.info(f"集合 {collection_name} 不存在")
|
||||
return []
|
||||
else:
|
||||
self.client.load_collection(collection_name)
|
||||
try:
|
||||
# 2. 执行查询
|
||||
filter_expr = prepare_milvus_filter(resource)
|
||||
iterator = self.client.query_iterator(
|
||||
collection_name=collection_name,
|
||||
batch_size=5,
|
||||
filter=filter_expr,
|
||||
output_fields=["id", "resource", "text"],
|
||||
)
|
||||
|
||||
results = []
|
||||
while True:
|
||||
try:
|
||||
result = iterator.next()
|
||||
if not result:
|
||||
break
|
||||
results.extend(result)
|
||||
except StopIteration:
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"查询迭代出错: {e}")
|
||||
break
|
||||
|
||||
return results
|
||||
|
||||
except Exception as e:
|
||||
print(f"查询失败: {e}")
|
||||
return []
|
||||
|
||||
finally:
|
||||
# 3. 无论成功还是失败,最终释放集合
|
||||
try:
|
||||
self.client.release_collection(collection_name)
|
||||
except Exception as e:
|
||||
print(f"释放集合失败: {e}")
|
||||
|
||||
def embedding_search(self, query_vector: List, top_k=3, collection_name: str = "MarkdownTest",
|
||||
database_name: str = "JXTest") -> List[dict | str]:
|
||||
"""
|
||||
纯向量检索
|
||||
"""
|
||||
self.client.use_database(database_name)
|
||||
self.load_collection(collection_name)
|
||||
res = self.client.search(
|
||||
collection_name=collection_name,
|
||||
anns_field="vector",
|
||||
limit=top_k,
|
||||
data=[query_vector],
|
||||
search_params={
|
||||
"metric_type": "IP", # COSINE
|
||||
"params": {"ef": 64} # "params": {"nprobe": 10}
|
||||
},
|
||||
output_fields=["id", "Header_1", "Header_2", "Header_3", "text", "resource", "img_path", "page_idx",
|
||||
"origin_id", "*"])
|
||||
search_results = []
|
||||
search_resources = []
|
||||
recall_results = []
|
||||
try:
|
||||
for i, hits in enumerate(res):
|
||||
for j, hit in enumerate(hits):
|
||||
id = hit["entity"].get("id")
|
||||
search_result_text = hit["entity"].get("text")
|
||||
distance = hit["distance"]
|
||||
page_idx = hit["entity"].get("page_idx")
|
||||
search_result_header_1 = hit["entity"].get("header_1")
|
||||
search_result_header_2 = hit["entity"].get("header_2")
|
||||
search_result_header_3 = hit["entity"].get("header_3")
|
||||
search_result_resource = hit["entity"].get("resource")
|
||||
search_result_img_path = hit["entity"].get("img_path")
|
||||
origin_id = hit["entity"].get("origin_id")
|
||||
dynamic_field_name = hit["entity"].get_properties()
|
||||
# search_score = hit["entity"].get("score")
|
||||
# search_result_img_information = hit["entity"].get("img_information")
|
||||
|
||||
if search_result_header_1 is None:
|
||||
search_result_header_2 = ""
|
||||
else:
|
||||
search_result_header_1 = f"#{search_result_header_1}"
|
||||
if search_result_header_2 is None:
|
||||
search_result_header_2 = ""
|
||||
else:
|
||||
search_result_header_2 = f"##{search_result_header_2}"
|
||||
if search_result_header_3 is None:
|
||||
search_result_header_3 = ""
|
||||
else:
|
||||
search_result_header_3 = f"###{search_result_header_3}"
|
||||
# search_result = f"#{search_result_header_1}\n{search_result_header_2}\n{search_result_header_3}\n\n{search_result_text}"
|
||||
search_result = search_result_text
|
||||
search_results.append(search_result)
|
||||
search_resources.append({"resource": search_result_resource, "img_path": search_result_img_path})
|
||||
# search_resources.append({"resource":search_result_resource,"img_path":search_result_img_path,"img_information":search_result_img_imformation})
|
||||
recall_results.append(
|
||||
{"id": origin_id, "index": i + j + 1, "text": search_result, "resource": search_result_resource,
|
||||
"img_path": search_result_img_path, "score": distance, "page_idx": page_idx,
|
||||
"dynamic_field_name": dynamic_field_name})
|
||||
# recall_results.append({"index":i+j+1,"text":search_result,"resource":search_result_resource,"img_path":search_result_img_path,"img_information":search_result_img_imformation})
|
||||
return search_results, search_resources, recall_results
|
||||
finally:
|
||||
# 3. 无论成功还是失败,最终释放集合
|
||||
try:
|
||||
self.client.release_collection(collection_name)
|
||||
except Exception as e:
|
||||
print(f"释放集合失败: {e}")
|
||||
|
||||
def hierarchical_search(self, query_vector: List, query, top_k=3, collection_name: str = "MarkdownTest",
|
||||
database_name: str = "JXTest") -> List[dict | str]:
|
||||
"""
|
||||
层级检索
|
||||
"""
|
||||
try:
|
||||
self.client.use_database(database_name)
|
||||
self.load_collection(collection_name)
|
||||
search_param_1 = {
|
||||
"data": [query_vector],
|
||||
"anns_field": "vector",
|
||||
"param": {
|
||||
"metric_type": "IP", # COSINE
|
||||
"params": {"ef": 64} # "params": {"nprobe": 10}
|
||||
},
|
||||
"limit": top_k
|
||||
}
|
||||
request_1 = AnnSearchRequest(**search_param_1)
|
||||
|
||||
search_param_2 = {
|
||||
"data": [query],
|
||||
"anns_field": "bm_25",
|
||||
"param": {
|
||||
"metric_type": "BM25",
|
||||
},
|
||||
"limit": top_k
|
||||
}
|
||||
request_2 = AnnSearchRequest(**search_param_2)
|
||||
|
||||
reqs = [request_1, request_2]
|
||||
ranker = WeightedRanker(0.3, 0.7)
|
||||
res = self.client.hybrid_search(
|
||||
collection_name=collection_name,
|
||||
reqs=reqs,
|
||||
ranker=ranker,
|
||||
limit=top_k,
|
||||
output_fields=["id", "Header_1", "Header_2", "Header_3", "text", "resource", "img_path", "page_idx",
|
||||
"origin_id"]
|
||||
# output_fields=["Header_1","Header_2","Header_3","text","resource","img_path","img_information"]
|
||||
)
|
||||
|
||||
search_results = []
|
||||
search_resources = []
|
||||
recall_results = []
|
||||
for i, hits in enumerate(res):
|
||||
for j, hit in enumerate(hits):
|
||||
id = hit["entity"].get("id")
|
||||
search_result_text = hit["entity"].get("text")
|
||||
distance = hit["distance"]
|
||||
page_idx = hit["entity"].get("page_idx")
|
||||
text = hit["entity"].get("text")
|
||||
search_result_header_1 = hit["entity"].get("header_1")
|
||||
search_result_header_2 = hit["entity"].get("header_2")
|
||||
search_result_header_3 = hit["entity"].get("header_3")
|
||||
search_result_resource = hit["entity"].get("resource")
|
||||
search_result_img_path = hit["entity"].get("img_path")
|
||||
origin_id = hit["entity"].get("origin_id")
|
||||
origin_text = hit["entity".get("origin_text")]
|
||||
# search_result_img_information = hit["entity"].get("img_information")
|
||||
|
||||
if search_result_header_1 is None:
|
||||
search_result_header_2 = ""
|
||||
else:
|
||||
search_result_header_1 = f"#{search_result_header_1}"
|
||||
if search_result_header_2 is None:
|
||||
search_result_header_2 = ""
|
||||
else:
|
||||
search_result_header_2 = f"##{search_result_header_2}"
|
||||
if search_result_header_3 is None:
|
||||
search_result_header_3 = ""
|
||||
else:
|
||||
search_result_header_3 = f"###{search_result_header_3}"
|
||||
# search_result = f"#{search_result_header_1}\n{search_result_header_2}\n{search_result_header_3}\n\n{search_result_text}"
|
||||
search_result = search_result_text
|
||||
search_results.append(search_result)
|
||||
search_resources.append({"resource": search_result_resource, "img_path": search_result_img_path})
|
||||
# search_resources.append({"resource":search_result_resource,"img_path":search_result_img_path,"img_information":search_result_img_imformation})
|
||||
# recall_results.append({"index":i+j+1,"id":origin_id,"text":origin_text,"resource":search_result_resource,"img_path":search_result_img_path,"score":distance,"page_idx":page_idx})
|
||||
recall_results.append(
|
||||
{"index": i + j + 1, "id": origin_id, "text": text, "resource": search_result_resource,
|
||||
"img_path": search_result_img_path, "score": distance, "page_idx": page_idx})
|
||||
self.release_collection(collection_name)
|
||||
return search_results, search_resources, recall_results
|
||||
finally:
|
||||
try:
|
||||
self.client.release_collection(collection_name)
|
||||
except Exception as e:
|
||||
print(f"释放集合失败: {e}")
|
||||
|
||||
def hybrid_search(self, query_vector: List, query, top_k=3, collection_name: str = "MarkdownTest",
|
||||
database_name: str = "JXTest", embedding_rate=0.3) -> List[dict | str]:
|
||||
"""
|
||||
|
||||
"""
|
||||
try:
|
||||
self.client.use_database(database_name)
|
||||
self.load_collection(collection_name)
|
||||
search_param_1 = {
|
||||
"data": [query_vector],
|
||||
"anns_field": "vector",
|
||||
"param": {
|
||||
"metric_type": "IP", # COSINE
|
||||
"params": {"ef": 64} # "params": {"nprobe": 10}
|
||||
},
|
||||
"limit": top_k
|
||||
}
|
||||
request_1 = AnnSearchRequest(**search_param_1)
|
||||
|
||||
search_param_2 = {
|
||||
"data": [query],
|
||||
"anns_field": "bm_25",
|
||||
"param": {
|
||||
"metric_type": "BM25",
|
||||
},
|
||||
"limit": top_k
|
||||
}
|
||||
request_2 = AnnSearchRequest(**search_param_2)
|
||||
|
||||
reqs = [request_1, request_2]
|
||||
ranker = WeightedRanker(float(embedding_rate), 1 - float(embedding_rate))
|
||||
res = self.client.hybrid_search(
|
||||
collection_name=collection_name,
|
||||
reqs=reqs,
|
||||
ranker=ranker,
|
||||
limit=top_k,
|
||||
output_fields=["id", "Header_1", "Header_2", "Header_3", "text", "resource", "img_path", "page_idx",
|
||||
"origin_id", "origin_text"]
|
||||
# output_fields=["Header_1","Header_2","Header_3","text","resource","img_path","img_information"]
|
||||
)
|
||||
|
||||
search_results = []
|
||||
search_resources = []
|
||||
recall_results = []
|
||||
for i, hits in enumerate(res):
|
||||
for j, hit in enumerate(hits):
|
||||
id = hit["entity"].get("id")
|
||||
search_result_text = hit["entity"].get("text")
|
||||
distance = hit["distance"]
|
||||
page_idx = hit["entity"].get("page_idx")
|
||||
search_result_header_1 = hit["entity"].get("Header_1")
|
||||
search_result_header_2 = hit["entity"].get("Header_2")
|
||||
search_result_header_3 = hit["entity"].get("Header_3")
|
||||
search_result_resource = hit["entity"].get("resource")
|
||||
search_result_img_path = hit["entity"].get("img_path")
|
||||
origin_id = hit["entity"].get("origin_id")
|
||||
header_1_num = len(search_result_header_1) if search_result_header_1 else 0
|
||||
header_2_num = len(search_result_header_2) if search_result_header_2 else 0
|
||||
header_3_num = len(search_result_header_3) if search_result_header_3 else 0
|
||||
origin_text = hit["entity"].get("origin_text")
|
||||
# search_result_img_information = hit["entity"].get("img_information")
|
||||
|
||||
if search_result_header_1 is None:
|
||||
search_result_header_2 = ""
|
||||
else:
|
||||
search_result_header_1 = f"#{search_result_header_1}"
|
||||
if search_result_header_2 is None:
|
||||
search_result_header_2 = ""
|
||||
else:
|
||||
search_result_header_2 = f"##{search_result_header_2}"
|
||||
if search_result_header_3 is None:
|
||||
search_result_header_3 = ""
|
||||
else:
|
||||
search_result_header_3 = f"###{search_result_header_3}"
|
||||
# search_result = f"#{search_result_header_1}\n{search_result_header_2}\n{search_result_header_3}\n\n{search_result_text}"
|
||||
search_result = search_result_text
|
||||
origin_text = search_result[header_1_num + header_2_num + header_3_num + 3:]
|
||||
search_results.append(search_result)
|
||||
search_resources.append({"resource": search_result_resource, "img_path": search_result_img_path})
|
||||
|
||||
# 获取片段的下文,并且拼接
|
||||
# next_chunk = self.get_chunk_by_id(collection_name=collection_name,database_name=database_name,id =id)
|
||||
|
||||
# next_text = next_chunk[0].get("text")
|
||||
|
||||
# search_result += next_text
|
||||
# origin_text += next_text
|
||||
# search_resources.append({"resource":search_result_resource,"img_path":search_result_img_path,"img_information":search_result_img_imformation})
|
||||
recall_results.append(
|
||||
{"index": i + j + 1, "id": origin_id, "text": search_result, "resource": search_result_resource,
|
||||
"img_path": search_result_img_path, "score": distance, "page_idx": page_idx,
|
||||
"origin_text": origin_text})
|
||||
# recall_results.append({"index":i+j+1,"text":search_result,"resource":search_result_resource,"img_path":search_result_img_path,"img_information":search_result_img_imformation})
|
||||
self.release_collection(collection_name)
|
||||
return search_results, search_resources, recall_results
|
||||
finally:
|
||||
try:
|
||||
self.client.release_collection(collection_name)
|
||||
except Exception as e:
|
||||
print(f"释放集合失败: {e}")
|
||||
|
||||
# return search_results
|
||||
|
||||
def hybrid_search_batch(
|
||||
self,
|
||||
query_vectors: List[List[float]],
|
||||
queries: List[str],
|
||||
top_k: int = 3,
|
||||
collection_name: str = "MarkdownTest",
|
||||
database_name: str = "JXTest",
|
||||
embedding_rate: float = 0.3
|
||||
) -> Tuple[List[List[Dict]], List[List[Dict]], List[List[Dict]]]:
|
||||
"""
|
||||
执行批量混合检索(稠密向量 + BM25 稀疏检索),对每个查询分别进行 hybrid search。
|
||||
|
||||
Args:
|
||||
query_vectors (List[List[float]]): 每个查询对应的稠密向量列表,形状为 [N, dim]
|
||||
queries (List[str]): 原始文本查询列表,长度为 N
|
||||
top_k (int): 返回每个查询的 top-k 结果
|
||||
collection_name (str): 集合名称
|
||||
database_name (str): 数据库名称
|
||||
embedding_rate (float): 稠密向量检索结果的权重(BM25 权重 = 1 - embedding_rate)
|
||||
|
||||
Returns:
|
||||
Tuple[
|
||||
List[List[str]]: 每个查询对应的结果文本列表(外层 list 是 query 数量,内层是 top-k 结果),
|
||||
List[List[dict]]: 每个查询对应的资源信息列表 ,
|
||||
List[List[dict]]: 每个查询对应的完整召回详情列表(含 score、id、origin_text 等)
|
||||
]
|
||||
"""
|
||||
if len(query_vectors) != len(queries):
|
||||
raise ValueError("query_vectors 和 queries 的数量必须一致")
|
||||
|
||||
try:
|
||||
self.client.use_database(database_name)
|
||||
self.load_collection(collection_name)
|
||||
|
||||
all_search_results = [] # [[res1_q1, res2_q1], [res1_q2, ...]]
|
||||
all_search_resources = []
|
||||
all_recall_results = []
|
||||
|
||||
for i, (vec, text) in enumerate(zip(query_vectors, queries)):
|
||||
try:
|
||||
# 构建稠密向量请求
|
||||
req1 = AnnSearchRequest(
|
||||
data=[vec],
|
||||
anns_field="vector",
|
||||
param={"metric_type": "IP", "params": {"ef": 64}},
|
||||
limit=top_k
|
||||
)
|
||||
|
||||
# 构建稀疏向量请求(BM25)
|
||||
req2 = AnnSearchRequest(
|
||||
data=[text],
|
||||
anns_field="bm_25",
|
||||
param={"metric_type": "BM25"},
|
||||
limit=top_k
|
||||
)
|
||||
|
||||
# 混合搜索
|
||||
ranker = WeightedRanker(embedding_rate, 1 - embedding_rate)
|
||||
res = self.client.hybrid_search(
|
||||
collection_name=collection_name,
|
||||
reqs=[req1, req2],
|
||||
ranker=ranker,
|
||||
limit=top_k,
|
||||
output_fields=[
|
||||
"id", "Header_1", "Header_2", "Header_3", "text", "resource",
|
||||
"img_path", "page_idx", "origin_id", "origin_text"
|
||||
],
|
||||
timeout=30
|
||||
)
|
||||
|
||||
# 解析单个查询的结果
|
||||
search_results = []
|
||||
search_resources = []
|
||||
recall_results = []
|
||||
|
||||
for j, hit in enumerate(res[0]): # res 是 List[List[Hit]], res[0] 对应第一个 request 的 hits
|
||||
entity = hit["entity"]
|
||||
id_val = entity.get("id")
|
||||
origin_id = entity.get("origin_id")
|
||||
text_val = entity.get("text") or ""
|
||||
header_1 = entity.get("Header_1") or ""
|
||||
header_2 = entity.get("Header_2") or ""
|
||||
header_3 = entity.get("Header_3") or ""
|
||||
resource = entity.get("resource")
|
||||
img_path = entity.get("img_path")
|
||||
page_idx = entity.get("page_idx")
|
||||
distance = hit["distance"]
|
||||
origin_text = entity.get("origin_text") or text_val
|
||||
|
||||
# 格式化标题
|
||||
header_1 = f"#{header_1}" if header_1 else ""
|
||||
header_2 = f"##{header_2}" if header_2 else ""
|
||||
header_3 = f"###{header_3}" if header_3 else ""
|
||||
|
||||
# 组合内容(可根据需要调整)
|
||||
combined_text = "\n".join(filter(None, [header_1, header_2, header_3, text_val])).strip()
|
||||
|
||||
search_results.append(combined_text)
|
||||
search_resources.append({"resource": resource, "img_path": img_path})
|
||||
|
||||
recall_results.append({
|
||||
"keyword": text,
|
||||
"index": j + 1,
|
||||
"id": origin_id,
|
||||
"text": combined_text,
|
||||
"resource": resource,
|
||||
"img_path": img_path,
|
||||
"score": distance,
|
||||
"page_idx": page_idx,
|
||||
"origin_text": origin_text
|
||||
})
|
||||
|
||||
all_search_results.append(search_results)
|
||||
all_search_resources.append(search_resources)
|
||||
all_recall_results.append(recall_results)
|
||||
|
||||
except Exception as e:
|
||||
print(f"第 {i + 1} 个查询 '{text}' 检索失败: {e}")
|
||||
# 失败时返回空列表占位
|
||||
all_search_results.append([])
|
||||
all_search_resources.append([])
|
||||
all_recall_results.append([])
|
||||
|
||||
return all_search_results, all_search_resources, all_recall_results
|
||||
|
||||
except Exception as e:
|
||||
print(f"批量混合检索整体失败: {e}")
|
||||
raise
|
||||
|
||||
finally:
|
||||
try:
|
||||
self.client.release_collection(collection_name)
|
||||
except Exception as e:
|
||||
print(f"释放集合失败: {e}")
|
||||
# return search_results
|
||||
|
||||
def close(self):
|
||||
self.client.close()
|
||||
|
||||
|
||||
def sanitize_collection_name(name: str) -> str:
|
||||
"""
|
||||
中文转换为拼音
|
||||
"""
|
||||
|
||||
def is_valid_milvus_collection_name(name: str) -> bool:
|
||||
"""检查名称是否符合 Milvus 集合命名规则"""
|
||||
import re
|
||||
pattern = r'^[a-zA-Z_][a-zA-Z0-9_]*$' # 首字符必须是字母或_,后续可以是字母、数字、_
|
||||
return bool(re.fullmatch(pattern, name))
|
||||
|
||||
if is_valid_milvus_collection_name(name):
|
||||
return name
|
||||
pinyin = "_".join(lazy_pinyin(name))
|
||||
# 确保首字符是字母或下划线
|
||||
if not pinyin[0].isalpha() and pinyin[0] != "_":
|
||||
pinyin = f"_{pinyin}"
|
||||
return pinyin
|
||||
|
||||
|
||||
##应用示例
|
||||
# if __name__ == '__main__':
|
||||
# #创建collection
|
||||
# database = Milvus_Database("XIAN")
|
||||
|
||||
# result = database.create_connection("MarkdownTest",drop_=False)
|
||||
# print(result)
|
||||
# print(database.list_partitions("MarkdownTest"))
|
||||
|
||||
|
||||
# 创建db
|
||||
# from pymilvus import connections, db
|
||||
# #
|
||||
# conn = Milvus_Database(user='root',
|
||||
# password='Milvus',
|
||||
# uri="http://192.168.0.46:19530")
|
||||
# print(conn.list_collections("XIAN"))
|
||||
|
||||
# conn = connections.connect(user='root',
|
||||
# password='Milvus',
|
||||
# uri="http://172.18.30.165:19530")
|
||||
|
||||
# database = db.create_database("JXTest")
|
||||
# db.using_database("JXTest")
|
||||
|
||||
0
app_pkg/examples/__init__.py
Normal file
0
app_pkg/examples/__init__.py
Normal file
177
app_pkg/examples/convert_word_pdf.py
Normal file
177
app_pkg/examples/convert_word_pdf.py
Normal file
@ -0,0 +1,177 @@
|
||||
from fastapi import UploadFile
|
||||
import os
|
||||
from typing import Optional
|
||||
import tempfile
|
||||
import requests
|
||||
import time
|
||||
from tqdm import tqdm
|
||||
from pathlib import Path
|
||||
|
||||
class ProgressFileReader:
|
||||
def __init__(self, file_path, progress_bar):
|
||||
self.file = open(file_path, 'rb')
|
||||
self.progress_bar = progress_bar
|
||||
|
||||
def read(self, size=-1):
|
||||
data = self.file.read(size)
|
||||
if data:
|
||||
self.progress_bar.update(len(data))
|
||||
return data
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.file.close()
|
||||
async def convert_doc_to_pdf(api_url: str, file: UploadFile, timeout: int = 300) -> Optional[str]:
|
||||
"""
|
||||
将DOC文件转换为pdf
|
||||
|
||||
参数:
|
||||
api_url: 转换服务的URL (如 "http://172.18.127.124:8000")
|
||||
file: 要转换的UploadFile对象
|
||||
timeout: 请求超时时间(秒)
|
||||
"""
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".doc") as tmp_doc:
|
||||
content = await file.read()
|
||||
tmp_doc.write(content)
|
||||
tmp_path = tmp_doc.name
|
||||
|
||||
file_name = file.filename
|
||||
file_size = os.path.getsize(tmp_path)
|
||||
output_name = Path(file_name).stem + ".pdf"
|
||||
|
||||
with tqdm(total=file_size, unit='B', unit_scale=True, desc="上传进度") as pbar:
|
||||
with ProgressFileReader(tmp_path, pbar) as file_reader:
|
||||
files = {'file': (file_name, file_reader)}
|
||||
params = {'src': 'doc', 'to': 'pdf'}
|
||||
|
||||
response = requests.post(
|
||||
f"{api_url}/transfer",
|
||||
files=files,
|
||||
params=params,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_docx:
|
||||
tmp_docx.write(response.content)
|
||||
output_path = tmp_docx.name
|
||||
|
||||
os.unlink(tmp_path)
|
||||
|
||||
return output_path
|
||||
else:
|
||||
raise Exception(f"转换失败: {response.text}")
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
raise Exception(f"请求超时(超过{timeout}秒)")
|
||||
except requests.exceptions.RequestException as e:
|
||||
raise Exception(f"请求错误: {str(e)}")
|
||||
except Exception as e:
|
||||
# 确保临时文件被清理
|
||||
if 'tmp_path' in locals() and os.path.exists(tmp_path):
|
||||
os.unlink(tmp_path)
|
||||
if 'output_path' in locals() and os.path.exists(output_path):
|
||||
os.unlink(output_path)
|
||||
raise Exception(f"转换过程中出错: {str(e)}")
|
||||
|
||||
async def convert_docx_to_pdf(api_url: str, file: UploadFile, timeout: int = 300) -> Optional[str]:
|
||||
"""
|
||||
将DOCX文件转换为PDF(解决中文缺失问题)
|
||||
|
||||
参数:
|
||||
api_url: 转换服务的URL (如 "http://172.18.127.124:8000")
|
||||
file: 要转换的UploadFile对象
|
||||
timeout: 请求超时时间(秒)
|
||||
"""
|
||||
try:
|
||||
# 创建临时文件保存上传的DOCX
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".docx") as tmp_docx:
|
||||
content = await file.read()
|
||||
tmp_docx.write(content)
|
||||
tmp_path = tmp_docx.name
|
||||
|
||||
# 获取文件信息
|
||||
file_name = file.filename
|
||||
file_size = os.path.getsize(tmp_path)
|
||||
output_name = Path(file_name).stem + ".pdf"
|
||||
|
||||
print(f"正在转换: {file_name} ({file_size/1024:.1f} KB) -> PDF")
|
||||
start_time = time.time()
|
||||
|
||||
# 使用带进度条的文件读取器
|
||||
with tqdm(total=file_size, unit='B', unit_scale=True, desc="上传进度") as pbar:
|
||||
with ProgressFileReader(tmp_path, pbar) as file_reader:
|
||||
files = {'file': (file_name, file_reader)}
|
||||
params = {
|
||||
'src': 'docx',
|
||||
'to': 'pdf',
|
||||
'font': 'SimSun',
|
||||
'encoding': 'utf-8'
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
f"{api_url}/transfer",
|
||||
files=files,
|
||||
params=params,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
# 处理响应
|
||||
if response.status_code == 200:
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_pdf:
|
||||
tmp_pdf.write(response.content)
|
||||
output_path = tmp_pdf.name
|
||||
|
||||
|
||||
if os.path.getsize(output_path) < 1024: # 如果PDF文件过小
|
||||
raise Exception("生成的PDF可能不包含有效内容")
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"\n转换成功! 保存为: {output_path}")
|
||||
print(f"总耗时: {elapsed:.2f}秒")
|
||||
|
||||
os.unlink(tmp_path)
|
||||
|
||||
return output_path
|
||||
else:
|
||||
raise Exception(f"转换失败: {response.text}")
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
raise Exception(f"请求超时(超过{timeout}秒)")
|
||||
except requests.exceptions.RequestException as e:
|
||||
raise Exception(f"请求错误: {str(e)}")
|
||||
except Exception as e:
|
||||
# 确保临时文件被清理
|
||||
if 'tmp_path' in locals() and os.path.exists(tmp_path):
|
||||
os.unlink(tmp_path)
|
||||
if 'output_path' in locals() and os.path.exists(output_path):
|
||||
os.unlink(output_path)
|
||||
raise Exception(f"转换过程中出错: {str(e)}")
|
||||
|
||||
|
||||
|
||||
#调用示例:
|
||||
async def split_result(file: UploadFile = File(...),
|
||||
chunk_type: str = Form("01"),
|
||||
chunk_size: int = Form(256),
|
||||
chunk_overlap: int = Form(50),
|
||||
max_level: Optional[int]=Form(3)):
|
||||
"""
|
||||
拆分结果
|
||||
"""
|
||||
filename = file.filename
|
||||
file_ext = os.path.splitext(filename)[1]
|
||||
|
||||
if file_ext == ".docx":
|
||||
# 处理DOCX文件
|
||||
converted_path = await convert_docx_to_pdf(os.getenv("CONVERT_DOC_URL","http://172.18.127.124:8000"), file)
|
||||
|
||||
# 读取转换后的DOCX文件
|
||||
with open(converted_path, 'rb') as f:
|
||||
converted_file = UploadFile(filename=f"{PATH(filename).stem}.pdf", file=f,headers=headers)
|
||||
#调用Mineru
|
||||
md_content = await process_pdf_logic(converted_file)
|
||||
28
app_pkg/examples/create_collections.py
Normal file
28
app_pkg/examples/create_collections.py
Normal file
@ -0,0 +1,28 @@
|
||||
from pymilvus import CollectionSchema, FieldSchema, DataType
|
||||
from pymilvus import Collection, db, connections,utility
|
||||
|
||||
def create_collections(collections_name,database_name : str ="JXTest"):
|
||||
db.using_database(database_name)
|
||||
|
||||
id = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True)
|
||||
text = FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=4096)
|
||||
header_1 = FieldSchema(name="header_1", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
|
||||
header_2 = FieldSchema(name="header_2", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
|
||||
header_3 = FieldSchema(name="header_3", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
|
||||
vector = FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=1024,)
|
||||
|
||||
schema = CollectionSchema(
|
||||
fields=[id, header_1, header_2, header_3, text, vector],
|
||||
auto_id=True,
|
||||
enable_dynamic_field=True
|
||||
)
|
||||
index_params = {
|
||||
"metric_type": "L2",
|
||||
"index_type": "IVF_FLAT", #FLAT
|
||||
"params": {"nlist": 128}
|
||||
}
|
||||
collection = Collection(name=collections_name, schema=schema, using='default', shards_num=2)
|
||||
collection.create_index(field_name="vector", index_params=index_params)
|
||||
|
||||
utility.index_building_progress("word_vector")
|
||||
return "collections created"
|
||||
58
app_pkg/examples/create_database.py
Normal file
58
app_pkg/examples/create_database.py
Normal file
@ -0,0 +1,58 @@
|
||||
# from pymilvus import Collection, utility, connections, db
|
||||
# from pymilvus import CollectionSchema, FieldSchema, DataType
|
||||
# from pymilvus import Collection, db, connections
|
||||
|
||||
# conn = connections.connect(host="172.18.127.124", port=19530)
|
||||
# # database = db.create_database("sample_db")
|
||||
# print(db.list_database())
|
||||
|
||||
# db.using_database("sample_db")
|
||||
|
||||
# # m_id = FieldSchema(name="m_id", dtype=DataType.INT64, is_primary=True,)
|
||||
# # embeding = FieldSchema(name="embeding", dtype=DataType.FLOAT_VECTOR, dim=768,)
|
||||
# # count = FieldSchema(name="count", dtype=DataType.INT64,)
|
||||
# # desc = FieldSchema(name="desc", dtype=DataType.VARCHAR, max_length=256,)
|
||||
# # schema = CollectionSchema(
|
||||
# # fields=[m_id, embeding, desc, count],
|
||||
# # description="Test embeding search",
|
||||
# # enable_dynamic_field=True
|
||||
# # )
|
||||
|
||||
# # collection_name = "word_vector"
|
||||
# # collection = Collection(name=collection_name, schema=schema, using='default', shards_num=2)
|
||||
|
||||
# index_params = {
|
||||
# "metric_type": "IP",
|
||||
# "index_type": "IVF_FLAT",
|
||||
# "params": {"nlist": 1024}
|
||||
# }
|
||||
|
||||
# collection = Collection("word_vector")
|
||||
# collection.create_index(
|
||||
# field_name="embeding",
|
||||
# index_params=index_params
|
||||
# )
|
||||
|
||||
# utility.index_building_progress("word_vector")
|
||||
|
||||
|
||||
from pymilvus import Collection, db, connections
|
||||
import numpy as np
|
||||
|
||||
conn = connections.connect(host="172.18.107.78", port=19530)
|
||||
database = db.create_database("XIAN")
|
||||
print(db.list_database())
|
||||
# db.using_database("sample_db")
|
||||
# coll_name = 'word_vector'
|
||||
|
||||
# mids, embedings, counts, descs = [], [], [], []
|
||||
# data_num = 100
|
||||
# for idx in range(0, data_num):
|
||||
# mids.append(idx)
|
||||
# embedings.append(np.random.normal(0, 0.1, 768).tolist())
|
||||
# descs.append(f'random num {idx}')
|
||||
# counts.append(idx)
|
||||
|
||||
# collection = Collection(coll_name)
|
||||
# mr = collection.insert([mids, embedings, descs, counts])
|
||||
# print(mr)
|
||||
118
app_pkg/examples/create_hybrid_collection.py
Normal file
118
app_pkg/examples/create_hybrid_collection.py
Normal file
@ -0,0 +1,118 @@
|
||||
from pymilvus import CollectionSchema, FieldSchema, DataType,client
|
||||
from pymilvus import Collection, db, connections,utility,Function,FunctionType
|
||||
from pymilvus import (
|
||||
MilvusClient, DataType
|
||||
)
|
||||
|
||||
def new_create_coll():
|
||||
client = MilvusClient(
|
||||
uri="http://172.18.107.78:19530",
|
||||
token="root:Milvus"
|
||||
)
|
||||
client.use_database("JXTest")
|
||||
|
||||
schema = MilvusClient.create_schema(
|
||||
auto_id=True,
|
||||
enable_dynamic_field=True,
|
||||
)
|
||||
|
||||
analyzer_params_custom = {
|
||||
"tokenizer": "jieba",
|
||||
"type": "chinese",
|
||||
"filter": ["cnalphanumonly"]
|
||||
}
|
||||
|
||||
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
|
||||
schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=1000,enable_analyzer=True,analyzer_params=analyzer_params_custom)
|
||||
schema.add_field(field_name="sparse", datatype=DataType.SPARSE_FLOAT_VECTOR)
|
||||
schema.add_field(field_name="dense", datatype=DataType.FLOAT_VECTOR, dim=5)
|
||||
|
||||
bm25_function = Function(
|
||||
name="text_bm25_emb", # Function name
|
||||
input_field_names=["text"], # Name of the VARCHAR field containing raw text data
|
||||
output_field_names=["sparse"], # Name of the SPARSE_FLOAT_VECTOR field reserved to store generated embeddings
|
||||
function_type=FunctionType.BM25,
|
||||
)
|
||||
|
||||
schema.add_function(bm25_function)
|
||||
|
||||
index_params = client.prepare_index_params()
|
||||
|
||||
# Add indexes
|
||||
index_params.add_index(
|
||||
field_name="dense",
|
||||
index_name="dense_index",
|
||||
index_type="AUTOINDEX",#"IVF_FLAT"
|
||||
metric_type="COSINE", #L2
|
||||
params={"nlist": 128},
|
||||
)
|
||||
|
||||
index_params.add_index(
|
||||
field_name="sparse",
|
||||
index_name="sparse_index",
|
||||
index_type="SPARSE_INVERTED_INDEX", # Index type for sparse vectors
|
||||
metric_type="BM25", # Set to `BM25` when using function to generate sparse vectors
|
||||
params={"inverted_index_algo": "DAAT_MAXSCORE"}, # The ratio of small vector values to be dropped during indexing
|
||||
)
|
||||
|
||||
client.create_collection(
|
||||
collection_name="hybrid_search_collection",
|
||||
schema=schema,
|
||||
index_params=index_params
|
||||
)
|
||||
return "Sucessful"
|
||||
|
||||
|
||||
|
||||
def create_collections(collections_name,database_name : str ="JXTest"):
|
||||
db.using_database(database_name)
|
||||
|
||||
id = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True)
|
||||
text = FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=4096,enable_analyzer=True)
|
||||
header_1 = FieldSchema(name="header_1", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
|
||||
header_2 = FieldSchema(name="header_2", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
|
||||
header_3 = FieldSchema(name="header_3", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY')
|
||||
vector = FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=1024,)
|
||||
sparse = FieldSchema(name="sparse_bm25", dtype=DataType.SPARSE_FLOAT_VECTOR)
|
||||
schema = CollectionSchema(
|
||||
fields=[id, header_1, header_2, header_3, text, vector,sparse],
|
||||
auto_id=True,
|
||||
enable_dynamic_field=True
|
||||
)
|
||||
index_params = {
|
||||
"field_name": "vector",
|
||||
"index_name": "vector_index",
|
||||
"metric_type": "L2",
|
||||
"index_type": "IVF_FLAT", #FLAT
|
||||
"params": {"nlist": 128},
|
||||
"field_name":"sparse_bm25",
|
||||
"index_name":"sparse_bm25_index",
|
||||
"index_type":"SPARSE_WAND",
|
||||
"metric_type":"BM25",
|
||||
"params":{
|
||||
"inverted_index_algo": "DAAT_MAXSCORE", # Algorithm for building and querying the index. Valid values: DAAT_MAXSCORE, DAAT_WAND, TAAT_NAIVE.
|
||||
"bm25_k1": 1.2,
|
||||
"bm25_b": 0.75
|
||||
}
|
||||
|
||||
}
|
||||
bm25_function = Function(
|
||||
name="bm25",
|
||||
function_type=FunctionType.BM25,
|
||||
input_field_names=["text"],
|
||||
output_field_names="sparse_bm25",
|
||||
)
|
||||
schema.add_function(bm25_function)
|
||||
|
||||
collection = Collection(name=collections_name, schema=schema, using='default', shards_num=2)
|
||||
collection.create_index(field_name="vector", index_params=index_params)
|
||||
|
||||
utility.index_building_progress("Hybrid")
|
||||
return "collections created"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# conn = connections.connect(host="172.18.107.78",port=19530,user="user",password="Milvus")
|
||||
# res = create_collections("HybirdTest","JXTest")
|
||||
# print(res)
|
||||
res = new_create_coll()
|
||||
18
app_pkg/examples/create_index.py
Normal file
18
app_pkg/examples/create_index.py
Normal file
@ -0,0 +1,18 @@
|
||||
from pymilvus import Collection, utility, connections, db
|
||||
|
||||
conn = connections.connect(host="192.168.1.156", port=19530)
|
||||
db.using_database("sample_db")
|
||||
|
||||
index_params = {
|
||||
"metric_type": "IP",
|
||||
"index_type": "IVF_FLAT",
|
||||
"params": {"nlist": 1024}
|
||||
}
|
||||
|
||||
collection = Collection("word_vector")
|
||||
collection.create_index(
|
||||
field_name="embeding",
|
||||
index_params=index_params
|
||||
)
|
||||
|
||||
utility.index_building_progress("word_vector")
|
||||
27
app_pkg/examples/embedding_score.py
Normal file
27
app_pkg/examples/embedding_score.py
Normal file
@ -0,0 +1,27 @@
|
||||
import json
|
||||
import requests
|
||||
def embedding_cosine_similarity(text_1, text_2):
|
||||
# 请求头
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': f'Bearer gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5'
|
||||
}
|
||||
|
||||
# 请求体
|
||||
data = {
|
||||
"text_1": text_1,
|
||||
"text_2": text_2,
|
||||
"model": "bge-rerank"
|
||||
}
|
||||
|
||||
# 发送 POST 请求
|
||||
response = requests.post("http://172.18.30.122:9500/v1" + "/score", headers=headers, data=json.dumps(data))
|
||||
if response.status_code == 200:
|
||||
result_json= response.json()
|
||||
# print(result_json)
|
||||
score = result_json["data"][0]["score"]
|
||||
return score
|
||||
else:
|
||||
raise Exception(f"Error: {response.status_code}\n{response.text}") # 抛出异常
|
||||
|
||||
print(embedding_cosine_similarity("你好","Hello"))
|
||||
15
app_pkg/examples/embedding_text.py
Normal file
15
app_pkg/examples/embedding_text.py
Normal file
@ -0,0 +1,15 @@
|
||||
from typing import List
|
||||
from openai import OpenAI
|
||||
import os
|
||||
def get_embeddings(embedding_texts: List[str]):
|
||||
response = OpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY","EMPTY"),
|
||||
base_url=os.getenv("OPENAI_API_EMBEDDING_BASE","http://172.18.30.122:40060/v1"),
|
||||
).embeddings.create(
|
||||
model=os.getenv("OPENAI_EMBEDDING_MODEL","bge-m3"),
|
||||
input=embedding_texts,
|
||||
)
|
||||
emebdding_list = response.data[0].embedding
|
||||
return emebdding_list
|
||||
|
||||
print(get_embeddings(["你好","世界"]))
|
||||
46
app_pkg/examples/get_3D.py
Normal file
46
app_pkg/examples/get_3D.py
Normal file
@ -0,0 +1,46 @@
|
||||
import os
|
||||
import json
|
||||
from typing import List
|
||||
import httpx
|
||||
import requests
|
||||
|
||||
|
||||
def get_embeddings(embedding_texts: List[str]):
|
||||
base_url=os.getenv("OPENAI_API_EMBEDDING_BASE","http://172.18.30.122:40060/v1")
|
||||
embedding_base_url = os.path.join(base_url, "embeddings")
|
||||
with httpx.Client(timeout=60) as client:
|
||||
response = client.post(
|
||||
embedding_base_url,
|
||||
json={"model": "bge-m3", "input": embedding_texts},
|
||||
headers={"Authorization": os.getenv("OPENAI_API_KEY","EMPTY")},
|
||||
)
|
||||
embedding_list = response.json()["data"][0]["embedding"]
|
||||
return embedding_list
|
||||
# 1. 组装请求体(与 SearchRequest 一一对应)
|
||||
|
||||
|
||||
payload = { # 若不提供 query,可改为传 query_vector
|
||||
"top_k": 3,
|
||||
"database_name": "XIAN",
|
||||
"collection_name": "b9d47d5a6d66cbe2cd44a8ea4ad021fc",
|
||||
"query_vector": get_embeddings(["和尚"]) # 如果本地已算好向量,可直接把 list[float] 写进来
|
||||
}
|
||||
|
||||
# 2. 发送 POST 请求
|
||||
url = "http://172.18.30.122:9079/search_3D"
|
||||
headers = {"Content-Type": "application/json"}
|
||||
|
||||
resp = requests.post(url, data=json.dumps(payload), headers=headers, timeout=30)
|
||||
|
||||
# 3. 处理返回
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
if data.get("code") == 200:
|
||||
print("检索成功:")
|
||||
for item in data.get("data", []):
|
||||
print("score:", item["score"])
|
||||
print("text :", item["text"][:200], "...")
|
||||
else:
|
||||
print("业务逻辑异常:", data.get("message"))
|
||||
else:
|
||||
print("HTTP 异常:", resp.status_code, resp.text)
|
||||
53
app_pkg/examples/get_chunk.py
Normal file
53
app_pkg/examples/get_chunk.py
Normal file
@ -0,0 +1,53 @@
|
||||
import os
|
||||
from pymilvus import (
|
||||
connections,
|
||||
utility,
|
||||
FieldSchema,
|
||||
CollectionSchema,
|
||||
DataType,
|
||||
Collection,
|
||||
MilvusClient,
|
||||
db,
|
||||
Function,
|
||||
FunctionType,
|
||||
AnnSearchRequest,
|
||||
WeightedRanker
|
||||
)
|
||||
|
||||
class Milvus_Database:
|
||||
def __init__(self,user=os.getenv("Milvus_USER","root"),password=os.getenv("Milvus_PASSWORD","Milvus"),uri="http://172.18.30.165:19530"):
|
||||
self.client = MilvusClient(
|
||||
uri="http://172.18.30.165:19530",
|
||||
token="root:Milvus"
|
||||
)
|
||||
def get_chunk_by_resource(self, collection_name, database_name,resource):
|
||||
"""
|
||||
return List :
|
||||
"""
|
||||
self.client.using_database(db_name=database_name)
|
||||
iterator = self.client.query_iterator(
|
||||
collection_name=collection_name,
|
||||
batch_size=5,
|
||||
filter=f"resource like \"{resource}\"",
|
||||
output_fields=["id","resource"],
|
||||
)
|
||||
|
||||
results = []
|
||||
|
||||
while True:
|
||||
result = iterator.next()
|
||||
if not result:
|
||||
iterator.close()
|
||||
break
|
||||
results += result
|
||||
return results
|
||||
|
||||
def get_chunk(collection_name,database_name,resource):
|
||||
milvus_db=Milvus_Database()
|
||||
results = milvus_db.get_chunk_by_resource(collection_name=collection_name,database_name=database_name,resource=resource)
|
||||
|
||||
return {"code": 200, "message": "success", "data": results}
|
||||
|
||||
if __name__ == "__main__":
|
||||
result = get_chunk(database_name="XIAN",collection_name="ce_shi_2_4dde68",resource="经桡动脉和股动脉穿刺行支架...治疗颅内动脉狭窄的效果比较_邱茜茜(1).pdf")
|
||||
print(result)
|
||||
85
app_pkg/examples/get_content_list.py
Normal file
85
app_pkg/examples/get_content_list.py
Normal file
@ -0,0 +1,85 @@
|
||||
import json
|
||||
import os
|
||||
import httpx
|
||||
|
||||
|
||||
async def process_pdf_logic(pdf_file):
|
||||
"""
|
||||
接收 PDF 文件并转发到解析服务
|
||||
"""
|
||||
# 验证文件类型
|
||||
if not pdf_file.content_type.startswith("application/pdf"):
|
||||
raise Exception(400, "仅支持 PDF 文件")
|
||||
|
||||
try:
|
||||
# 读取文件内容
|
||||
# file_content = await pdf_file.read()
|
||||
# 配置目标接口参数
|
||||
target_url = os.getenv("Mineru_URL","http://172.18.30.122:9988")
|
||||
target_url = target_url + "/analyze-pdf"
|
||||
# target_url = os.getenv("Mineru_URL","http://172.18.127.124:9088/pdf_parse")
|
||||
# target_url = "http://172.18.30.122:8998/pdf_parse"
|
||||
# params = {
|
||||
# "parse_method": "auto",
|
||||
# "is_json_md_dump": "true",
|
||||
# "output_dir": "output"
|
||||
# }
|
||||
|
||||
# 构建文件上传格式
|
||||
files = {
|
||||
"file": (pdf_file.filename, pdf_file.file, "application/pdf")
|
||||
# "pdf_file": (pdf_file.filename, pdf_file.file, "application/pdf")
|
||||
}
|
||||
|
||||
# 发送请求到解析服务
|
||||
async with httpx.AsyncClient(timeout=1200.0) as client:
|
||||
response = await client.post(
|
||||
target_url,
|
||||
# params=params,
|
||||
files=files,
|
||||
headers={"Accept": "application/json"}
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
full_content = result.get("data", {}).get("full_content", "")
|
||||
content_list = result.get("data", {}).get("content_list", [])
|
||||
# markdown_content =""
|
||||
# for content in content_list:
|
||||
# text = content["text"]
|
||||
|
||||
return full_content,content_list
|
||||
|
||||
except httpx.RequestError as e:
|
||||
raise Exception(503, f"无法连接到解析服务: {str(e)}")
|
||||
except Exception as e:
|
||||
raise Exception(500, f"服务器内部错误: {str(e)}")
|
||||
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
|
||||
async def test_request(path):
|
||||
pdf_path = Path(path)
|
||||
async with httpx.AsyncClient() as client:
|
||||
with open(pdf_path, "rb") as f:
|
||||
files = {"file": (pdf_path.name, f, "application/pdf")}
|
||||
response = await client.post(
|
||||
"http://172.18.30.122:9988/analyze-pdf",
|
||||
files=files,
|
||||
timeout=120
|
||||
)
|
||||
result = response.json()
|
||||
|
||||
full_content = result.get("data", {}).get("full_content", "")
|
||||
content_list = result.get("data", {}).get("content_list", [])
|
||||
combined_json = ''.join(content_list)
|
||||
content_data = json.loads(combined_json)
|
||||
for item in content_data:
|
||||
if isinstance(item, dict): # 确保是字典类型
|
||||
text = item.get("text", "")
|
||||
page_idx = item.get("page_idx", -1) # 默认值 -1 表示缺失
|
||||
print(f"Text: {text}\nPage Index: {page_idx}\n{'-'*50}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
# asyncio.run(test_request("/Users/lijinxuan/PycharmProjects/Sea/AutoRag/src/examples/邱茜茜.pdf"))
|
||||
asyncio.run(test_request("/Users/lijinxuan/PycharmProjects/Sea/AutoRag/src/examples/邱茜茜.pdf"))
|
||||
31
app_pkg/examples/has_collection.py
Normal file
31
app_pkg/examples/has_collection.py
Normal file
@ -0,0 +1,31 @@
|
||||
from pymilvus import (
|
||||
connections,
|
||||
utility,
|
||||
FieldSchema,
|
||||
CollectionSchema,
|
||||
DataType,
|
||||
Collection,
|
||||
MilvusClient,
|
||||
db,
|
||||
Function,
|
||||
FunctionType,
|
||||
AnnSearchRequest,
|
||||
WeightedRanker
|
||||
)
|
||||
import os
|
||||
class Milvus_Database:
|
||||
def __init__(self,user=os.getenv("Milvus_USER","root"),password=os.getenv("Milvus_PASSWORD","Milvus"),uri=os.getenv("Milvus_URI")):
|
||||
self.client = MilvusClient(
|
||||
uri=uri,
|
||||
token=f"{user}:{password}"
|
||||
)
|
||||
def has_collection(self,collection_name,database_name):
|
||||
"""
|
||||
Before using list collection, You need to using using_database method first
|
||||
"""
|
||||
self.client.using_database(db_name=database_name)
|
||||
return self.client.has_collection(collection_name)
|
||||
|
||||
if __name__ == "__main__":
|
||||
db=Milvus_Database(uri="http://172.18.30.165:19530")
|
||||
print(db.has_collection("test","XIAN"))
|
||||
164
app_pkg/examples/hubrid_insert_search.py
Normal file
164
app_pkg/examples/hubrid_insert_search.py
Normal file
@ -0,0 +1,164 @@
|
||||
from pymilvus import (
|
||||
MilvusClient,
|
||||
AnnSearchRequest,
|
||||
WeightedRanker
|
||||
)
|
||||
import os
|
||||
from openai import OpenAI
|
||||
from typing import List
|
||||
def insert():
|
||||
|
||||
docs = [
|
||||
"Artificial intelligence was founded as an academic discipline in 1956.",
|
||||
"Alan Turing was the first person to conduct substantial research in AI.",
|
||||
"Born in Maida Vale, London, Turing was raised in southern England.",
|
||||
]
|
||||
|
||||
data = [
|
||||
{"text": docs[0], "dense": [2.7242085933685303, 6.021071434020996, 0.4754035174846649, 9.358858108520508, 5.173221111297607]},
|
||||
{"text": docs[1], "dense": [8.584294319152832, 2.7640628814697266, 9.558855056762695, 2.584272861480713, 4.705013275146484]},
|
||||
{"text": docs[2], "dense": [2.5525057315826416, 3.8815805912017822, 9.343480110168457, 7.888997554779053, 4.500918388366699]},
|
||||
]
|
||||
|
||||
client = MilvusClient(
|
||||
uri="http://172.18.107.78:19530",
|
||||
token="root:Milvus"
|
||||
)
|
||||
client.use_database("JXTest")
|
||||
res = client.insert(
|
||||
collection_name="hybrid_search_collection",
|
||||
data=data
|
||||
)
|
||||
return res
|
||||
|
||||
def search():
|
||||
from pymilvus import AnnSearchRequest
|
||||
client = MilvusClient(
|
||||
uri="http://172.18.127.124:19530",
|
||||
token="root:Milvus"
|
||||
)
|
||||
search_param_1 = {
|
||||
"data": [[0.7425515055656433, 7.774101734161377, 0.7397570610046387, 2.429982900619507, 3.8253049850463867]],
|
||||
"anns_field": "dense",
|
||||
"param": {
|
||||
"metric_type": "L2",
|
||||
"params": {"nprobe": 10}
|
||||
},
|
||||
"limit": 2
|
||||
}
|
||||
request_1 = AnnSearchRequest(**search_param_1)
|
||||
|
||||
search_param_2 = {
|
||||
"data": ['Who started AI research'],
|
||||
"anns_field": "sparse",
|
||||
"param": {
|
||||
"metric_type": "BM25",
|
||||
},
|
||||
"limit": 2
|
||||
}
|
||||
request_2 = AnnSearchRequest(**search_param_2)
|
||||
|
||||
reqs = [request_1, request_2]
|
||||
return reqs
|
||||
|
||||
openai_embedding_api_base = os.getenv("OPENAI_API_EMBEDDING_BASE","http://172.18.127.124:40697/v1")
|
||||
openai_api_key = os.getenv("OPENAI_API_KEY", "gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5")
|
||||
|
||||
embedding_client = OpenAI(
|
||||
api_key=openai_api_key,
|
||||
base_url=openai_embedding_api_base,
|
||||
)
|
||||
def get_embeddings(embedding_texts: List[str]):
|
||||
response = embedding_client.embeddings.create(
|
||||
model="bge-m3",
|
||||
input=embedding_texts,
|
||||
)
|
||||
emebdding_list = response.data[0].embedding
|
||||
return emebdding_list
|
||||
|
||||
def test():
|
||||
from pymilvus import AnnSearchRequest
|
||||
client = MilvusClient(
|
||||
uri="http://172.18.127.124:19530",
|
||||
token="root:Milvus"
|
||||
)
|
||||
client.use_database("JXTest")
|
||||
query_vector = get_embeddings(["焦煤谜团中央医院收治了多少名颅内动脉狭窄患者?"])
|
||||
search_param_1 = {
|
||||
"data": [query_vector],
|
||||
"anns_field": "vector",
|
||||
"param": {
|
||||
"metric_type": "COSINE",
|
||||
"params": {"nprobe": 10}
|
||||
},
|
||||
"limit": 3
|
||||
}
|
||||
request_1 = AnnSearchRequest(**search_param_1)
|
||||
|
||||
search_param_2 = {
|
||||
"data": ["焦煤谜团中央医院收治了多少名颅内动脉狭窄患者?"],
|
||||
"anns_field": "bm_25",
|
||||
"param": {
|
||||
"metric_type": "BM25",
|
||||
},
|
||||
"limit": 3
|
||||
}
|
||||
request_2 = AnnSearchRequest(**search_param_2)
|
||||
|
||||
reqs = [request_1, request_2]
|
||||
ranker = WeightedRanker(0.5, 0.5)
|
||||
res = client.hybrid_search(
|
||||
collection_name="XianTest",
|
||||
reqs=reqs,
|
||||
ranker=ranker,
|
||||
limit=3,
|
||||
output_fields=["Header_1","Header_2","Header_3","text","resource","img_path"]
|
||||
)
|
||||
search_results = []
|
||||
search_resources = []
|
||||
for hits in res:
|
||||
for hit in hits:
|
||||
search_result_text = hit["entity"].get("text")
|
||||
|
||||
search_result_resource = hit["entity"].get("resource")
|
||||
search_result_img_path = hit["entity"].get("img_path")
|
||||
search_results.append(search_result_text)
|
||||
search_resources.append({"resource":search_result_resource,"img_path":search_result_img_path})
|
||||
client.close()
|
||||
return search_results,search_resources
|
||||
|
||||
if __name__ == '__main__':
|
||||
search_results,search_resources=test()
|
||||
print(search_results)
|
||||
print(search_resources)
|
||||
|
||||
|
||||
|
||||
# insert_res = insert()
|
||||
# print(insert_res)
|
||||
# reqs = search()
|
||||
# client = MilvusClient(
|
||||
# uri = "http://172.18.127.124:19530",
|
||||
# # uri="http://172.18.107.78:19530",
|
||||
# token="root:Milvus"
|
||||
# )
|
||||
# client.use_database("JXTest")
|
||||
# print(client.list_collections())
|
||||
|
||||
# from pymilvus import WeightedRanker
|
||||
|
||||
# ranker = WeightedRanker(0.5, 0.5)
|
||||
# res = client.hybrid_search(
|
||||
# collection_name="hybrid_search_collection",
|
||||
# reqs=reqs,
|
||||
# ranker=ranker,
|
||||
# limit=2,
|
||||
# output_fields=["text"]
|
||||
# )
|
||||
# for hits in res:
|
||||
# print("TopK results:")
|
||||
# for hit in hits:
|
||||
# print(hit)
|
||||
# Ouput :
|
||||
#{'id': 457020197413365373, 'distance': 0.3451843857765198, 'entity': {'text': 'Alan Turing was the first person to conduct substantial research in AI.'}}
|
||||
#{'id': 457020197413365372, 'distance': 0.005594015121459961, 'entity': {'text': 'Artificial intelligence was founded as an academic discipline in 1956.'}}
|
||||
18
app_pkg/examples/insert_data.py
Normal file
18
app_pkg/examples/insert_data.py
Normal file
@ -0,0 +1,18 @@
|
||||
from pymilvus import Collection, db, connections
|
||||
import numpy as np
|
||||
|
||||
conn = connections.connect(host="192.168.1.156", port=19530)
|
||||
db.using_database("sample_db")
|
||||
coll_name = 'word_vector'
|
||||
|
||||
mids, embedings, counts, descs = [], [], [], []
|
||||
data_num = 100
|
||||
for idx in range(0, data_num):
|
||||
mids.append(idx)
|
||||
embedings.append(np.random.normal(0, 0.1, 768).tolist())
|
||||
descs.append(f'random num {idx}')
|
||||
counts.append(idx)
|
||||
|
||||
collection = Collection(coll_name)
|
||||
mr = collection.insert([mids, embedings, descs, counts])
|
||||
print(mr)
|
||||
30
app_pkg/examples/list_connection.py
Normal file
30
app_pkg/examples/list_connection.py
Normal file
@ -0,0 +1,30 @@
|
||||
# from pymilvus import Collection, db, connections
|
||||
# from pymilvus import utility
|
||||
# conn = connections.connect(host="172.18.30.165", port=19530)
|
||||
# db.using_database("XIAN")
|
||||
# # result = utility.has_collection("RAGTest")
|
||||
# # print(result)
|
||||
# print(utility.list_collections())
|
||||
|
||||
|
||||
from pymilvus import MilvusClient, DataType
|
||||
|
||||
client = MilvusClient(
|
||||
uri="http://172.18.30.165:19530",
|
||||
token="root:Milvus"
|
||||
)
|
||||
client.using_database("XIAN")
|
||||
|
||||
res = client.list_collections()
|
||||
search_params = {
|
||||
"metric_type": "IP", # 距离度量类型,如 L2/IP等
|
||||
"params": {"nprobe": 10}
|
||||
}
|
||||
res = client.describe_collection(
|
||||
collection_name="ce_shi_zhi_shi_ku_1_449853"
|
||||
)
|
||||
|
||||
for hits in res:
|
||||
print(hits)
|
||||
# print(f"ID: {hit.id}, Distance: {hit.distance}, Vector: {hit.entity.get('vector_field_name')}")
|
||||
|
||||
177
app_pkg/examples/markdown_splits.py
Normal file
177
app_pkg/examples/markdown_splits.py
Normal file
@ -0,0 +1,177 @@
|
||||
from langchain.text_splitter import MarkdownTextSplitter,MarkdownHeaderTextSplitter,RecursiveCharacterTextSplitter
|
||||
from md_header_convert import convert_numbered_headings
|
||||
import re
|
||||
import json
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=256,
|
||||
chunk_overlap=50,
|
||||
length_function=len,
|
||||
separators=[" "]
|
||||
)
|
||||
headers_to_split_on = [("#", "Header_1"), ("##", "Header_2"), ("###", "Header_3")]
|
||||
|
||||
markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
|
||||
|
||||
|
||||
def parse_markdown(md_text):
|
||||
md_text = convert_numbered_headings(md_text)
|
||||
lines = md_text.strip().split("\n")
|
||||
results = []
|
||||
current_h1, current_h2, current_h3 = None, None, None
|
||||
current_text = []
|
||||
|
||||
def add_entry():
|
||||
if current_h1 or current_h2 or current_h3:
|
||||
results.append({
|
||||
"Header_1": current_h1,
|
||||
"Header_2": current_h2,
|
||||
"Header_3": current_h3,
|
||||
"text": "\n".join(current_text).strip()
|
||||
})
|
||||
current_text.clear()
|
||||
|
||||
for line in lines:
|
||||
line = line.strip()
|
||||
header_match = re.match(r'^(#{1,3})\s*(.*)', line)
|
||||
|
||||
if header_match:
|
||||
add_entry()
|
||||
level = len(header_match.group(1))
|
||||
title = header_match.group(2).strip()
|
||||
|
||||
if level == 1:
|
||||
current_h1, current_h2, current_h3 = title, None, None
|
||||
elif level == 2:
|
||||
current_h2, current_h3 = title, None
|
||||
elif level == 3:
|
||||
current_h3 = title
|
||||
else:
|
||||
current_text.append(line)
|
||||
|
||||
add_entry()
|
||||
return results
|
||||
# return json.dumps(results, indent=4, ensure_ascii=False)
|
||||
|
||||
def header_postprocess(md_header_json):
|
||||
for md in md_header_json:
|
||||
# print(md)
|
||||
if not md["Header_2"] and not md["Header_3"]:
|
||||
if md["text"]:
|
||||
continue
|
||||
else:
|
||||
md["text"] = md["Header_1"]
|
||||
if md["Header_2"] and not md["Header_3"]:
|
||||
if md["text"]:
|
||||
md["text"] = md["Header_2"] + md["text"]
|
||||
md["Header_2"] = None
|
||||
else:
|
||||
md["text"] = md["Header_2"]
|
||||
md["Header_2"] = None
|
||||
if md["Header_2"] and md["Header_3"]:
|
||||
if not md["text"]:
|
||||
md["text"] = md["Header_3"]
|
||||
md["Header_3"] = None
|
||||
else:
|
||||
md["text"] = md["Header_3"] + md["text"]
|
||||
md["Header_3"] = None
|
||||
return md_header_json
|
||||
|
||||
# text_splitter = MarkdownTextSplitter(chunk_size=512,chunk_overlap=128)
|
||||
|
||||
|
||||
# print(text)
|
||||
# print("####################")
|
||||
|
||||
def langchain_split(text):
|
||||
text = convert_numbered_headings(text)
|
||||
# print(text)
|
||||
md_header_splits = markdown_splitter.split_text(text)
|
||||
results = []
|
||||
for i, doc in enumerate(md_header_splits):
|
||||
print("-------------------------------------------------------")
|
||||
print(f"Document {i+1}:")
|
||||
print("text")
|
||||
print(doc.page_content)
|
||||
|
||||
if doc.metadata:
|
||||
header_1 = doc.metadata["Header_1"] if "Header_1" in doc.metadata.keys() else None
|
||||
header_2 = doc.metadata["Header_2"] if "Header_2" in doc.metadata.keys() else None
|
||||
header_3 = doc.metadata["Header_3"] if "Header_3" in doc.metadata.keys() else None
|
||||
text_ = doc.page_content
|
||||
# print("Header_1:",header_1)
|
||||
# print("Header_2:",header_2)
|
||||
# print("Header_3:",header_3)
|
||||
results.append({"Header_1":header_1,"Header_2":header_2,"Header_3":header_3,"text":text_})
|
||||
# print("\n")
|
||||
return results
|
||||
|
||||
# md_header_json = header_postprocess(results)
|
||||
# for md in md_header_json:
|
||||
# print(md)
|
||||
# print("======================================")
|
||||
|
||||
def split_regular(text):
|
||||
#自定义分解
|
||||
md_header_json = parse_markdown(text)
|
||||
# print(md_header_json)
|
||||
#合并处理
|
||||
md_header_json = header_postprocess(md_header_json)
|
||||
# print(md_header_json)
|
||||
|
||||
results = []
|
||||
for i ,md in enumerate(md_header_json):
|
||||
if md["text"]:
|
||||
header_1 = md["Header_1"]
|
||||
header_2 = md["Header_2"]
|
||||
header_3 = md["Header_3"]
|
||||
text = md["text"]
|
||||
# print("Header_1:",header_1)
|
||||
# print("Header_2:",header_2)
|
||||
# print("Header_3:",header_3)
|
||||
# print("text:",text_)
|
||||
# print("======================================")
|
||||
texts = text_splitter.split_text(text)
|
||||
for j,txt in enumerate(texts):
|
||||
# print(j)
|
||||
# print(texts)
|
||||
# print("======================================")
|
||||
results.append({"id":j+i+1,"Header_1":header_1,"Header_2":header_2,"Header_3":header_3,"text":txt})
|
||||
return results
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
# results = []
|
||||
# for _, i in enumerate(md_header_splits):
|
||||
# if i.metadata:
|
||||
# header_1 = i.metadata["Header_1"] if "Header_1" in i.metadata.keys() else None
|
||||
# header_2 = i.metadata["Header_2"] if "Header_2" in i.metadata.keys() else None
|
||||
# header_3 = i.metadata["Header_3"] if "Header_3" in i.metadata.keys() else None
|
||||
# text_ = i.page_content
|
||||
# print("Header_1:",header_1)
|
||||
# print("Header_2:",header_2)
|
||||
# print("Header_3:",header_3)
|
||||
# # texts = text_splitter.split_text(text)
|
||||
# # for j in texts:
|
||||
# # print(j)
|
||||
# # print(texts)
|
||||
# print("======================================")
|
||||
# # results.append({"id":i,"content":texts})
|
||||
|
||||
# for res in results:
|
||||
# print(res["id"])
|
||||
# print(res["content"])
|
||||
# print("======================================")
|
||||
|
||||
if __name__ == "__main__":
|
||||
with open("邱茜茜.md","r",encoding="utf-8") as f:
|
||||
text = f.read()
|
||||
results = split_regular(text)
|
||||
for res in results:
|
||||
print("======================================")
|
||||
for key, value in res.items():
|
||||
print(f"{key}: {value}")
|
||||
print("\n")
|
||||
|
||||
48
app_pkg/examples/md_header_convert.py
Normal file
48
app_pkg/examples/md_header_convert.py
Normal file
@ -0,0 +1,48 @@
|
||||
import re
|
||||
|
||||
def convert_numbered_headings(text):
|
||||
"""
|
||||
将编号转换为对应的Markdown标题:
|
||||
- `1.2` → `## 1.2`
|
||||
- `1.2.1` → `### 1.2.1`
|
||||
- `1. 3. 1` → `### 1.3.1`(去掉空格)
|
||||
"""
|
||||
# 替换全角 `.` 为半角 `.`
|
||||
text = text.replace('.', '.')
|
||||
|
||||
# 遍历每一行,处理标题
|
||||
lines = text.split('\n')
|
||||
|
||||
for i in range(len(lines)):
|
||||
line = lines[i].strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
# 只匹配 `1.2` 或 `1.2.1` 形式的编号,允许 `1. 3. 1` 这种格式
|
||||
match = re.match(r'^(#*\s*)(\d+(?:\s*\.\s*\d+)+)', line)
|
||||
if match:
|
||||
prefix, num = match.groups()
|
||||
level = num.count('.') # `.` 的个数决定层级
|
||||
new_prefix = '#' * (level + 1) # 1.2 → ##, 1.2.1 → ###
|
||||
|
||||
# 去掉数字和 `.` 之间的空格
|
||||
clean_num = re.sub(r'\s*\.\s*', '.', num)
|
||||
|
||||
# 替换标题
|
||||
lines[i] = f"{new_prefix} {clean_num}{line[match.end():]}"
|
||||
|
||||
return '\n'.join(lines)
|
||||
|
||||
|
||||
# # 示例文本
|
||||
# text = """
|
||||
# # 1资料与方法
|
||||
|
||||
# 1.1一般资料选取2019年12月至2021年1月焦煤集团中央医院收治的60例...
|
||||
|
||||
# # 1.2选取标准
|
||||
|
||||
# 1.2.1纳入标准(1)经CT、MRI检查确诊为颅内动脉狭窄...
|
||||
# """
|
||||
|
||||
# print(convert_numbered_headings(text))
|
||||
62
app_pkg/examples/pdf_md.py
Normal file
62
app_pkg/examples/pdf_md.py
Normal file
@ -0,0 +1,62 @@
|
||||
import httpx
|
||||
import os
|
||||
|
||||
|
||||
async def process_pdf_logic(pdf_file):
|
||||
"""
|
||||
接收 PDF 文件并转发到解析服务
|
||||
"""
|
||||
# 验证文件类型
|
||||
if not pdf_file.content_type.startswith("application/pdf"):
|
||||
raise Exception(400, "仅支持 PDF 文件")
|
||||
|
||||
try:
|
||||
# 读取文件内容
|
||||
# file_content = await pdf_file.read()
|
||||
|
||||
# 配置目标接口参数
|
||||
target_url = os.getenv("Mineru_URL","http://172.18.30.122:9988")
|
||||
target_url = target_url + "/analyze-pdf"
|
||||
# target_url = os.getenv("Mineru_URL","http://172.18.127.124:9088/pdf_parse")
|
||||
# target_url = "http://172.18.30.122:8998/pdf_parse"
|
||||
# params = {
|
||||
# "parse_method": "auto",
|
||||
# "is_json_md_dump": "true",
|
||||
# "output_dir": "output"
|
||||
# }
|
||||
|
||||
# 构建文件上传格式
|
||||
files = {
|
||||
"file": (pdf_file.filename, pdf_file.file, "application/pdf")
|
||||
# "pdf_file": (pdf_file.filename, pdf_file.file, "application/pdf")
|
||||
}
|
||||
|
||||
# 发送请求到解析服务
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
response = await client.post(
|
||||
target_url,
|
||||
# params=params,
|
||||
files=files,
|
||||
headers={"Accept": "application/json"}
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
content_list = result["data"]["full_content"]
|
||||
# markdown_content =""
|
||||
# for content in content_list:
|
||||
# text = content["text"]
|
||||
|
||||
return content_list
|
||||
|
||||
except httpx.RequestError as e:
|
||||
raise Exception(503, f"无法连接到解析服务: {str(e)}")
|
||||
except Exception as e:
|
||||
raise Exception(500, f"服务器内部错误: {str(e)}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 测试:传入 PDF 文件路径
|
||||
result = process_pdf_logic("邱茜茜.pdf")
|
||||
print(result)
|
||||
# with open("邱茜茜.md", "w", encoding="utf-8") as f:
|
||||
# f.write(result)
|
||||
158
app_pkg/examples/search.py
Normal file
158
app_pkg/examples/search.py
Normal file
@ -0,0 +1,158 @@
|
||||
from pymilvus import Collection, db, connections,
|
||||
import numpy as np
|
||||
from typing import List
|
||||
from openai import OpenAI
|
||||
import os
|
||||
|
||||
openai_embedding_api_base = os.getenv("OPENAI_API_EMBEDDING_BASE","http://172.18.127.124:40483/v1")
|
||||
openai_chat_api_base = os.getenv("OPENAI_API_BASE","http://172.18.127.124:40159/v1")
|
||||
openai_api_key = os.getenv("OPENAI_API_KEY", "gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5")
|
||||
|
||||
embedding_client = OpenAI(
|
||||
api_key=openai_api_key,
|
||||
base_url=openai_embedding_api_base,
|
||||
)
|
||||
def get_embeddings(embedding_texts: List[str]):
|
||||
response = embedding_client.embeddings.create(
|
||||
model="bge-m3",
|
||||
input=embedding_texts,
|
||||
)
|
||||
emebdding_list = response.data[0].embedding
|
||||
return emebdding_list
|
||||
|
||||
conn = connections.connect(host="172.18.127.124", port=19530)
|
||||
db.using_database("JXTest")
|
||||
coll_name = 'RAGTest'
|
||||
|
||||
search_params = {
|
||||
"metric_type": 'L2',
|
||||
"offset": 0,
|
||||
"ignore_growing": False,
|
||||
"params": {"nprobe": 16}
|
||||
}
|
||||
|
||||
collection = Collection(coll_name)
|
||||
collection.load()
|
||||
query = "手术成功率是多少?"
|
||||
embeddings = get_embeddings([query])
|
||||
results = collection.search(
|
||||
data=[embeddings],
|
||||
anns_field="vector",
|
||||
param=search_params,
|
||||
limit=16,
|
||||
expr=None,
|
||||
# output_fields=['m_id', 'embeding', 'desc', 'count'],
|
||||
output_fields=["text","header_1"],
|
||||
consistency_level="Strong"
|
||||
)
|
||||
collection.release()
|
||||
# print(results[0].ids)
|
||||
# print(results[0].distances)
|
||||
# hit = results[0][0]
|
||||
# print(hit.entity.get('text'))
|
||||
print(results)
|
||||
|
||||
|
||||
|
||||
class Milvus_Database:
|
||||
def __init__(self,user=os.getenv("Milvus_USER","root"),password=os.getenv("Milvus_PASSWORD","Milvus"),uri=os.getenv("Milvus_URI")):
|
||||
self.client = MilvusClient(
|
||||
uri=uri,
|
||||
token=f"{user}:{password}"
|
||||
)
|
||||
|
||||
def hybrid_search(self,query_vector:List,query,top_k=3,collection_name:str = "MarkdownTest",database_name:str = "JXTest",embedding_rate=0.1)-> List[dict|str]:
|
||||
"""
|
||||
|
||||
"""
|
||||
try:
|
||||
self.client.use_database(database_name)
|
||||
self.load_collection(collection_name)
|
||||
search_param_1 = {
|
||||
"data": [query_vector],
|
||||
"anns_field": "vector",
|
||||
"param": {
|
||||
"metric_type": "IP", #COSINE
|
||||
"params":{"ef":64}#"params": {"nprobe": 10}
|
||||
},
|
||||
"limit": top_k
|
||||
}
|
||||
request_1 = AnnSearchRequest(**search_param_1)
|
||||
|
||||
search_param_2 = {
|
||||
"data": [query],
|
||||
"anns_field": "bm_25",
|
||||
"param": {
|
||||
"metric_type": "BM25",
|
||||
},
|
||||
"limit": top_k
|
||||
}
|
||||
request_2 = AnnSearchRequest(**search_param_2)
|
||||
|
||||
reqs = [request_1, request_2]
|
||||
ranker = WeightedRanker(float(embedding_rate), 1-float(embedding_rate))
|
||||
res = self.client.hybrid_search(
|
||||
collection_name=collection_name,
|
||||
reqs=reqs,
|
||||
ranker=ranker,
|
||||
limit=top_k,
|
||||
output_fields=["id","Header_1","Header_2","Header_3","text","resource","img_path","page_idx","origin_id","*"]
|
||||
#output_fields=["Header_1","Header_2","Header_3","text","resource","img_path","img_information"]
|
||||
)
|
||||
|
||||
search_results = []
|
||||
search_resources = []
|
||||
recall_results = []
|
||||
for i,hits in enumerate(res):
|
||||
for j,hit in enumerate(hits):
|
||||
id = hit["entity"].get("id")
|
||||
search_result_text = hit["entity"].get("text")
|
||||
distance = hit["distance"]
|
||||
page_idx = hit["entity"].get("page_idx")
|
||||
search_result_header_1 = hit["entity"].get("Header_1")
|
||||
search_result_header_2 = hit["entity"].get("Header_2")
|
||||
search_result_header_3 = hit["entity"].get("Header_3")
|
||||
search_result_resource = hit["entity"].get("resource")
|
||||
search_result_img_path = hit["entity"].get("img_path")
|
||||
origin_id = hit["entity"].get("origin_id")
|
||||
header_1_num = len(search_result_header_1) if search_result_header_1 else 0
|
||||
header_2_num = len(search_result_header_2) if search_result_header_2 else 0
|
||||
header_3_num = len(search_result_header_3) if search_result_header_3 else 0
|
||||
|
||||
#search_result_img_information = hit["entity"].get("img_information")
|
||||
|
||||
if search_result_header_1 is None:
|
||||
search_result_header_2 = ""
|
||||
else:
|
||||
search_result_header_1 = f"#{search_result_header_1}"
|
||||
if search_result_header_2 is None:
|
||||
search_result_header_2 = ""
|
||||
else:
|
||||
search_result_header_2 = f"##{search_result_header_2}"
|
||||
if search_result_header_3 is None:
|
||||
search_result_header_3 = ""
|
||||
else:
|
||||
search_result_header_3 = f"###{search_result_header_3}"
|
||||
# search_result = f"#{search_result_header_1}\n{search_result_header_2}\n{search_result_header_3}\n\n{search_result_text}"
|
||||
search_result = search_result_text
|
||||
origin_text = search_result[header_1_num+header_2_num+header_3_num+3:]
|
||||
search_results.append(search_result)
|
||||
search_resources.append({"resource":search_result_resource,"img_path":search_result_img_path})
|
||||
|
||||
#获取片段的下文,并且拼接
|
||||
next_chunk = self.get_chunk_by_id(collection_name=collection_name,database_name=database_name,id =id)
|
||||
|
||||
next_text = next_chunk[0].get("text")
|
||||
|
||||
search_result += next_text
|
||||
origin_text += next_text
|
||||
# search_resources.append({"resource":search_result_resource,"img_path":search_result_img_path,"img_information":search_result_img_imformation})
|
||||
recall_results.append({"index":i+j+1,"id":origin_id,"text":search_result,"resource":search_result_resource,"img_path":search_result_img_path,"score":distance,"page_idx":page_idx,"origin_text":origin_text})
|
||||
#recall_results.append({"index":i+j+1,"text":search_result,"resource":search_result_resource,"img_path":search_result_img_path,"img_information":search_result_img_imformation})
|
||||
self.release_collection(collection_name)
|
||||
return search_results, search_resources,recall_results
|
||||
finally:
|
||||
try:
|
||||
self.client.release_collection(collection_name)
|
||||
except Exception as e:
|
||||
print(f"释放集合失败: {e}")
|
||||
87
app_pkg/examples/search_text.py
Normal file
87
app_pkg/examples/search_text.py
Normal file
@ -0,0 +1,87 @@
|
||||
import os
|
||||
from pymilvus import (
|
||||
connections,
|
||||
utility,
|
||||
FieldSchema,
|
||||
CollectionSchema,
|
||||
DataType,
|
||||
Collection,
|
||||
MilvusClient,
|
||||
db,
|
||||
Function,
|
||||
FunctionType,
|
||||
AnnSearchRequest,
|
||||
WeightedRanker
|
||||
)
|
||||
|
||||
class Milvus_Database:
|
||||
def __init__(self,user=os.getenv("Milvus_USER","root"),password=os.getenv("Milvus_PASSWORD","Milvus"),uri="http://172.18.30.165:19530"):
|
||||
self.client = MilvusClient(
|
||||
uri="http://172.18.30.165:19530",
|
||||
token="root:Milvus"
|
||||
)
|
||||
def get_text_by_resource(self, collection_name, database_name,resource,text):
|
||||
"""
|
||||
return List :
|
||||
"""
|
||||
self.client.using_database(db_name=database_name)
|
||||
iterator = self.client.query_iterator(
|
||||
collection_name=collection_name,
|
||||
batch_size=5,
|
||||
filter=f"text like \"{text}\" and resource like \"{resource}\"",
|
||||
output_fields=["id","resource","text"],
|
||||
)
|
||||
|
||||
results = []
|
||||
|
||||
while True:
|
||||
result = iterator.next()
|
||||
if not result:
|
||||
iterator.close()
|
||||
break
|
||||
results += result
|
||||
return results
|
||||
def get_chunk_by_id(self, collection_name, database_name,id):
|
||||
"""
|
||||
return List :
|
||||
"""
|
||||
try:
|
||||
self.client.using_database(db_name=database_name)
|
||||
if not self.client.has_collection(collection_name):
|
||||
return []
|
||||
else:
|
||||
self.client.load_collection(collection_name)
|
||||
iterator = self.client.query_iterator(
|
||||
collection_name=collection_name,
|
||||
batch_size=5,
|
||||
filter=f"id=={id}",
|
||||
output_fields=["id","text","resource"],
|
||||
)
|
||||
|
||||
results = []
|
||||
|
||||
while True:
|
||||
result = iterator.next()
|
||||
if not result:
|
||||
iterator.close()
|
||||
break
|
||||
results += result
|
||||
return results
|
||||
finally:
|
||||
# 3. 无论成功还是失败,最终释放集合
|
||||
try:
|
||||
self.client.release_collection(collection_name)
|
||||
except Exception as e:
|
||||
print(f"释放集合失败: {e}")
|
||||
|
||||
def get_chunk(collection_name,database_name,resource,text,id):
|
||||
milvus_db=Milvus_Database()
|
||||
# results = milvus_db.get_text_by_resource(collection_name=collection_name,database_name=database_name,resource=resource,text=text)
|
||||
results = milvus_db.get_chunk_by_id(collection_name=collection_name,database_name=database_name,id=id)
|
||||
return {"code": 200, "message": "success", "data": results}
|
||||
|
||||
if __name__ == "__main__":
|
||||
text = "你好"
|
||||
# text = "应,进一步放大风险的危害,而传播环境的深刻变革对触发和放大生成式人工智 能传播风险也具有重要影响。生成式人工智能传播风险打破了独立、单层、定向的风险路径, 通过机内交互、人机交互与人人交互的复杂网络结构形成了联动、多层、多向的风险路径, 极易形成多领域共振传导。"
|
||||
result = get_chunk(database_name="XIAN",collection_name="ye85a36d295a199730340c5c755c1e2b",id="459963917155397121",text=text,resource="test.md")
|
||||
print(result)
|
||||
76
app_pkg/examples/search_vector.py
Normal file
76
app_pkg/examples/search_vector.py
Normal file
@ -0,0 +1,76 @@
|
||||
from pymilvus import connections, Collection, db
|
||||
import pandas as pd
|
||||
from openpyxl import Workbook
|
||||
|
||||
# 连接 Milvus
|
||||
connections.connect(
|
||||
uri="http://172.18.30.165:19530",
|
||||
token="root:Milvus"
|
||||
)
|
||||
|
||||
# 使用数据库
|
||||
db.using_database("XIAN")
|
||||
|
||||
# 获取集合
|
||||
collection = Collection("ce_shi_3_an_zi_duan_token_fen_a1da59")
|
||||
collection.load() # 确保集合已加载
|
||||
|
||||
# 创建查询迭代器
|
||||
iterator = collection.query_iterator(
|
||||
batch_size=10, # 可以调整批量大小以提高效率
|
||||
expr="id > 0",
|
||||
output_fields=["*"] # 指定要获取的字段
|
||||
)
|
||||
|
||||
results = []
|
||||
|
||||
# 遍历所有结果
|
||||
try:
|
||||
while True:
|
||||
# 获取下一批结果
|
||||
res = iterator.next()
|
||||
# print(res)
|
||||
|
||||
# 如果没有更多结果,退出循环
|
||||
if not res:
|
||||
break
|
||||
|
||||
# 将结果添加到列表中
|
||||
results.extend(res)
|
||||
print(f"已获取 {len(results)} 条记录")
|
||||
|
||||
finally:
|
||||
# 确保迭代器被关闭
|
||||
iterator.close()
|
||||
# 将结果转换为DataFrame
|
||||
df = pd.DataFrame(results)
|
||||
|
||||
# 检查是否有数据
|
||||
if not df.empty:
|
||||
# 写入Excel文件
|
||||
excel_file = "milvus_query_results.xlsx"
|
||||
|
||||
# 使用openpyxl引擎以获得更好的格式控制
|
||||
with pd.ExcelWriter(excel_file, engine='openpyxl') as writer:
|
||||
df.to_excel(writer, index=False, sheet_name='查询结果')
|
||||
|
||||
# 获取工作簿和工作表对象以进行格式设置
|
||||
workbook = writer.book
|
||||
worksheet = writer.sheets['查询结果']
|
||||
|
||||
# 设置列宽
|
||||
for column in worksheet.columns:
|
||||
max_length = 0
|
||||
column_letter = column[0].column_letter
|
||||
for cell in column:
|
||||
try:
|
||||
if len(str(cell.value)) > max_length:
|
||||
max_length = len(cell.value)
|
||||
except:
|
||||
pass
|
||||
adjusted_width = (max_length + 2) * 1.2
|
||||
worksheet.column_dimensions[column_letter].width = adjusted_width
|
||||
|
||||
print(f"结果已成功写入 {excel_file}")
|
||||
else:
|
||||
print("没有查询到任何结果")
|
||||
4
app_pkg/examples/str2url.py
Normal file
4
app_pkg/examples/str2url.py
Normal file
@ -0,0 +1,4 @@
|
||||
import urllib.parse
|
||||
resource = "邱茜茜.pdf"
|
||||
encoded_resource = urllib.parse.quote(resource)
|
||||
print(encoded_resource)
|
||||
57
app_pkg/examples/text_match.py
Normal file
57
app_pkg/examples/text_match.py
Normal file
@ -0,0 +1,57 @@
|
||||
from pymilvus import (
|
||||
connections,
|
||||
utility,
|
||||
FieldSchema,
|
||||
CollectionSchema,
|
||||
DataType,
|
||||
Collection,
|
||||
MilvusClient,
|
||||
db,
|
||||
Function,
|
||||
FunctionType,
|
||||
AnnSearchRequest,
|
||||
WeightedRanker
|
||||
)
|
||||
import os
|
||||
|
||||
class Milvus_Database:
|
||||
def __init__(self, user=os.getenv("Milvus_USER", "root"), password=os.getenv("Milvus_PASSWORD", "Milvus"), uri=os.getenv("Milvus_URI")):
|
||||
self.client = MilvusClient(
|
||||
uri="http://172.18.30.165:19530",
|
||||
token=f"{user}:{password}"
|
||||
)
|
||||
|
||||
def search_resource(self, collection_name, database_name,resource):
|
||||
self.client.using_database(db_name=database_name)
|
||||
iterator = self.client.query_iterator(
|
||||
collection_name=collection_name,
|
||||
batch_size=5,
|
||||
filter=f"resource like \"{resource}\"",
|
||||
output_fields=["text","resource"],
|
||||
)
|
||||
|
||||
results = []
|
||||
|
||||
while True:
|
||||
result = iterator.next()
|
||||
if not result:
|
||||
iterator.close()
|
||||
break
|
||||
|
||||
print(result)
|
||||
results += result
|
||||
return results
|
||||
# res = self.client.query(
|
||||
# collection_name=collection_name,
|
||||
# filter=f"resource like \"{resource}\"",
|
||||
# output_fields=["resource"],
|
||||
# limit=None # 先限制数量测试
|
||||
# )
|
||||
# return res
|
||||
|
||||
if __name__ == "__main__":
|
||||
db = Milvus_Database()
|
||||
res = db.search_resource("DockerTest", database_name="XIAN",resource="测试.pdf")
|
||||
for i,re in enumerate(res):
|
||||
print(i)
|
||||
# print(res)
|
||||
76
app_pkg/examples/text_split_markdown.py
Normal file
76
app_pkg/examples/text_split_markdown.py
Normal file
@ -0,0 +1,76 @@
|
||||
|
||||
from langchain_text_splitters import CharacterTextSplitter,TokenTextSplitter,RecursiveCharacterTextSplitter
|
||||
from markdown_splits import convert_numbered_headings,parse_markdown,header_postprocess
|
||||
|
||||
def split_01(md_Document,chunk_length=256,chunk_overlap=50):
|
||||
md_header_json = parse_markdown(md_Document)
|
||||
# print(md_header_json)
|
||||
#合并处理
|
||||
md_header_list = header_postprocess(md_header_json)
|
||||
for i ,md in enumerate(md_header_list):
|
||||
if md["text"]:
|
||||
header_1 = md["Header_1"]
|
||||
header_2 = md["Header_2"]
|
||||
header_3 = md["Header_3"]
|
||||
text = md["text"]
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=chunk_length,
|
||||
chunk_overlap=chunk_overlap,
|
||||
length_function=len,
|
||||
separators=[" "]
|
||||
)
|
||||
texts = text_splitter.split_text(text)
|
||||
for txt in texts:
|
||||
print("Header_1:",header_1)
|
||||
print("Header_2:",header_2)
|
||||
print("Header_3:",header_3)
|
||||
print("text:",txt)
|
||||
print("======================================")
|
||||
return md_header_list
|
||||
|
||||
|
||||
def split_02(md_Document,chunk_length=256,chunk_overlap=50):
|
||||
md_header_json = parse_markdown(md_Document)
|
||||
# print(md_header_json)
|
||||
#合并处理
|
||||
md_header_list = header_postprocess(md_header_json)
|
||||
for i ,md in enumerate(md_header_list):
|
||||
if md["text"]:
|
||||
header_1 = md["Header_1"]
|
||||
header_2 = md["Header_2"]
|
||||
header_3 = md["Header_3"]
|
||||
text = md["text"]
|
||||
print("Header_1:",header_1)
|
||||
print("Header_2:",header_2)
|
||||
print("Header_3:",header_3)
|
||||
print("text:",text)
|
||||
print("======================================")
|
||||
return md_header_list
|
||||
|
||||
def split_03(md_Document,chunk_length=256,chunk_overlap=50):
|
||||
# text_splitter = text_splitter = CharacterTextSplitter(
|
||||
# separator="\n\n",
|
||||
# chunk_size=chunk_length,
|
||||
# chunk_overlap=chunk_overlap,
|
||||
# length_function=len,
|
||||
# is_separator_regex=False,
|
||||
# )
|
||||
text_splitter = TokenTextSplitter(chunk_size=chunk_length, chunk_overlap=chunk_overlap)
|
||||
# text_splitter = RecursiveCharacterTextSplitter(
|
||||
# chunk_size=chunk_length,
|
||||
# chunk_overlap=chunk_overlap,
|
||||
# length_function=len,
|
||||
# separators=[""]
|
||||
# )
|
||||
texts = text_splitter.split_text(md_Document)
|
||||
return texts
|
||||
|
||||
if __name__ == "__main__":
|
||||
with open("邱茜茜.md","r",encoding="utf-8") as f:
|
||||
md_Document = f.read()
|
||||
# texts = split_03(md_Document)
|
||||
# for text in texts:
|
||||
# print("______________")
|
||||
# print(text)
|
||||
# split_02(md_Document)
|
||||
split_01(md_Document)
|
||||
1
app_pkg/examples/毕设.md
Normal file
1
app_pkg/examples/毕设.md
Normal file
File diff suppressed because one or more lines are too long
98
app_pkg/examples/邱茜茜.md
Normal file
98
app_pkg/examples/邱茜茜.md
Normal file
@ -0,0 +1,98 @@
|
||||
# 经挠动脉和股动脉穿刺行支架置入术治疗颅内动脉狭窄的效果比较
|
||||
|
||||
邱茜茜 蔡直铎 徐晓敏(焦煤集团中央医院 神经内科 河南 焦作454000)
|
||||
|
||||
摘要:目的比较经桡动脉和股动脉穿刺行支架置入术治疗颅内动脉狭窄的效果。方法选取2019 年12月至2021年1月焦煤集团中央医院收治的60例颅内动脉狭窄患者作为研究对象根据治疗方式分为桡动脉穿刺组(30例)和股动脉穿刺组(30例)。挠动脉穿刺组接受经挠动脉穿刺行支架植入术股动脉穿刺组接受经股动脉穿刺行支架植入术。比较两组患者手术指标(手术时间、肢体制动时间、止血包扎时间、鞘管置入时间、手术成功率及穿刺成功率)以及并发症发生情况。结果桡动脉穿刺组肢体制动时间、止血包扎时间短于股动脉穿刺组鞘管置入时间、手术时间长于股动脉穿刺组( $P<0.\ 05)$ )。两组患者手术成功率及穿刺成功率比较差异无统计学意义( $P>0.\,05)$ 。股动脉穿刺组肺栓塞、术后尿留、动脉闭塞、穿刺部位血肿、动脉痉挛等并发症发生率高于动脉穿刺组( $P<0.\;05`$ )。结论经桡动脉穿刺行支架植入术治疗颅内动脉狭窄可优化手术指标 降低并发症发生率。
|
||||
|
||||
R615 doi 10. 3969 /j. issn. 1004 - 437X. 2021. 34. 039
|
||||
|
||||
# Comparison of Stenting Through Radial Artery and Femoral Artery in the Treatment of Intracranial Artery Stenosis
|
||||
|
||||
QIU Xixi CAI Zhiduo XU Xiaomin
|
||||
|
||||
Department of Neurology the Coking Coal Group Central Hospital Jiaozuo 454000 China
|
||||
|
||||
Abstract Objective To compare the effect of stent implantation through radial artery and femoral artery in the treatment of intracranial artery stenosis. Methods A total of 60 patients with intracranial artery stenosis treated in the Coking Coal Group Central Hospital from December 2019 to January 2021 were selected as the research objects. According to the treatment methods they were divided into radial artery puncture group 30 cases and femoral artery puncture group 30 cases . The radial artery puncture group received stent implantation through radial artery puncture and the femoral artery puncture group received stent implantation through femoral artery puncture. The operation indexes operation time limb braking time hemostatic dressing time sheath insertion time operation success rate and puncture success rate and complications were compared between the two groups. Results Compared with the radial artery puncture group the limb braking time hemostasis and bandage time were shorter in the femoral artery puncture group and the sheath insertion time and operation time were longer in the femoral artery puncture group $P<0.\;05)$ . There was no statistical significance in the success rate of operation and puncture between the two groups $P>0.\,05$ . The incidence of complications such as pulmonary embolism postoperative urinary retention arterial occlusion puncture site hematoma and arterial spasm in femoral artery puncture group were higher than those in radial artery puncture group $P<0.\,05)$ . Conclusion Stent implantation through radial artery puncture in the treatment of intracranial artery stenosis can optimize the surgical indexes and reduce the incidence of complications.
|
||||
|
||||
nial artery stenosis stent implantation radial artery puncture femoral artery
|
||||
|
||||
目前经股动脉穿刺支架植入术为临床治疗血管狭窄的传统术式[-2]。但是部分患者存在主动脉钙化严重、主动脉弓迁曲、股动脉闭塞等情况,影响治疗[3]。经桡动脉穿刺较为安全且患者痛苦小,易于接受,已成为实施血管造影技术的常用方式[4-6]。但是经挠动脉穿刺行支架置入术治疗颅内动脉狭窄的相关研究较少本研究旨在比较经桡动脉和股动脉穿刺行支架置入术治疗颅内动脉狭窄的效果。
|
||||
|
||||
# 1 资料与方法
|
||||
|
||||
1. 1 2019 12 2021 1煤集团中央医院收治的60例颅内动脉狭窄患者作为研究对象根据治疗方式分为桡动脉穿刺组(30例)和30 。 16 1445 ~ 68 $55.\,46\pm4.\,12)$ 岁;疾病类型基底动脉狭窄5例,椎动脉狭窄12例,颈内动脉狭窄13 body mass index BMI 20. 78 ~
|
||||
|
||||
$24.\,01\,\mathrm{\kg\cdot{m}^{\mathrm{~-~}2}}$ ,平均 $(\ 22.\ 47\ \pm2.\ 56)\ \mathrm{kg}\ ^{\bullet}\ \mathrm{m}\,^{-2}$ 。股动脉穿刺组男18例,女12例;年龄 $42\,\sim70\$ 岁,平均$56.\,11\pm4.\,16)$ 岁;疾病类型基底动脉狭窄6例,椎动14 10 BMI 20. 41 ~$24.\:32\:\mathrm{~kg~}\cdot\mathrm{m~}^{-2}$ 平均 $(~22.~76~\pm~2.~62)~\mathrm{kg}~\cdot~\mathrm{m}^{\mathrm{~-~2~}}$ 。两组患者一般资料比较,差异无统计学意义( $P>0.\;05\,)$ 。本研究经焦煤集团中央医院医学伦理委员会审核批准。患者对本研究内容知情自愿签署知情同意书。
|
||||
|
||||
# 1. 2
|
||||
|
||||
1. 2. 1 1 CT、MRI动脉狭窄;(2)无手术禁忌证;(3)临床资料完整;4 。
|
||||
|
||||
1. 2. 2 1 、病;(2)肝、肾功能异常;(3)合并血管瘤、恶性肿瘤。
|
||||
|
||||
1. 3
|
||||
|
||||
1.3.1桡动脉穿刺组术前口服氯吡格雷和阿司匹林 ${\geqslant}5$ d Seldinger入术[7-8]。患者取仰卧位,垫高腕部,平伸外展手臂,选取桡骨茎突上 $1.\,0\sim2.\,0\,\,\mathrm{{cm}}$ 动脉搏动最强处穿刺;TERUMO 6 F 1 F周长 $1\ \mathrm{mm}$ CORDIS 6 F持一致的 $30^{\circ}$ 方向进针。穿刺成功后将 $50~\mathrm{U}\cdot\mathrm{kg}^{\mathrm{~-~}1}$ 普通肝素经鞘管注入,若手术时间 $>\,1\,\mathrm{~h~}$ ,再追加$1\ 000\ \mathrm{U}\cdot\mathrm{h}^{\mathrm{~-~}1}$ 。穿刺成功后给予全身麻醉经桡动脉鞘将引l导管引|进远端动脉,使用购自美国BostonSci-entiufic Winspan Gateway导管或购自微创医疗器械上海有限公司的Apollo球囊扩张式支架系统重建狭窄段血管。术后通过动脉造影评估手术效果将导引导管撤出至锁骨下动脉通过造影检查颅内动脉入路是否受损。
|
||||
|
||||
1.3.2股动脉穿刺组术前口服氯吡格雷和阿司匹林 ${\geqslant}5\mathrm{~d~}$ 通过Seldinger穿刺法经股动脉穿刺行支架植入术。常规消毒腹股沟部位,置入 $8\,\mathrm{~F~}$ 动脉鞘,并将$50~\mathrm{U}\cdot\mathrm{kg}^{\mathrm{~-~}1}$ 普通肝素注入鞘内,后续步骤同桡动脉穿刺组。
|
||||
|
||||
1. 3. 3 1 TERUMO挠动脉压迫器进行止血,以既能扪及远端挠动脉搏动又能止血为佳术后 $2\,\,\mathrm{h}$ 减压1次,术后 $8\sim12\,\mathrm{~h~}$ 解除压迫器。(2)股动脉穿刺组于术后 $4~\mathrm{{h}}$ 拔除鞘管局部压迫 $15\sim30\,\mathrm{\min}$ 后加压包扎,局部制动 $\mathcal{J}4\,\mathrm{~h~}$ 后可下床活动。(3)术后严密监测两组患者生命体征连续3 d 4 000 U $12\,\mathrm{~h~}$ 注射1次;口服$75~\mathrm{mg}$ 氯吡格雷、 $100~\mathrm{{mg}}$ 阿司匹林。
|
||||
|
||||
1.4手术成功标准(1)入路动脉无急性血栓形成、血管穿破及血管夹层等并发症;(2)术前血管造影检查显示患者血管狭窄率达 $70\%$ ,术后残余狭窄率不足$20\%$ ,且患者溶栓血流分级达3 级[3]。
|
||||
|
||||
1.5观察指标(1)手术指标。手术时间、肢体制动时间、止血包扎时间、鞘管置入时间、手术成功率及穿刺成功率。(2)并发症发生情况。肺栓塞、术后尿潴留、动脉闭塞、穿刺部位血肿、动脉痉挛。
|
||||
|
||||
1. 6 SPSS 22. 0据。手术时间、肢体制动时间、止血包扎时间、鞘管置入时间等计量资料,以均数 $\pm$ 标准差( $\bar{\boldsymbol{x}}\pm\boldsymbol{s})$ 表示采用$t$ 检验比较;并发症发生率、手术成功率及穿刺成功率为计数资料,以频数 $\left(\begin{array}{l}{n}\end{array}\right)$ 和率 $(\,\%\,)$ )表示,采用 $\chi^{2}$ 检验比较。 $P<0.\ 05$ 为差异有统计学意义。
|
||||
|
||||
# 2 结果
|
||||
|
||||
2.1手术指标桡动脉穿刺组手术成功率为90. $00\%$ 27 /30 $93.\,33\%$ 28 /30动脉穿刺组手术成功率为 $96.\ 67\%\ (\ 29\,/30)$ 穿刺成功率为 $100.\:00\%\:(\:30\,/\,30)$ 。两组患者手术成功率和穿刺成功率比较,差异无统计学意义( $P>0.\,05)$ 。桡动脉穿刺组肢体制动时间、止血包扎时间短于股动脉穿刺组鞘管置入时间、手术时间长于股动脉穿刺组( $P<$ 0. 05 。 1。
|
||||
|
||||
表1两组手术指标比较( $\bar{\boldsymbol{x}}\pm\boldsymbol{s})$
|
||||
|
||||
|
||||
<html><body><table><tr><td>组别</td><td>例数</td><td>手术时间/ min</td><td>肢体制动 时间/h</td><td>止血包扎 时间/min</td><td>鞘管置入 时间/min</td></tr><tr><td>桡动脉穿刺组30</td><td></td><td>45.67 ±4.21</td><td>4.78 ±1.23</td><td>1.57 ±0.56</td><td>6.01 ±1.57</td></tr><tr><td>股动脉穿刺组30</td><td></td><td>34.68 ±3.51</td><td>22.46 ±2.78</td><td>20.64 ±2.46</td><td>4.23 ±1.21</td></tr><tr><td>t/x2</td><td></td><td>10.982</td><td>31.855</td><td>41. 400</td><td>4.919</td></tr><tr><td>P</td><td></td><td><0.001</td><td><0.001</td><td><0.001</td><td><0.001</td></tr></table></body></html>
|
||||
|
||||
2.2并发症发生情况股动脉穿刺组肺栓塞、术后尿潴留、动脉闭塞、穿刺部位血肿、动脉痉挛等并发症发生率高于桡动脉穿刺组( $P<0.\,05)$ 。 2。
|
||||
|
||||
表2两组术后并发症发生情况比较
|
||||
|
||||
|
||||
<html><body><table><tr><td>组别</td><td>例数</td><td>肺栓塞</td><td>术后 尿潴留</td><td>动脉 闭塞</td><td>穿刺部位</td><td>动脉 痉挛</td><td>总发生率/</td></tr><tr><td>桡动脉穿刺组</td><td>30</td><td>0</td><td>0</td><td>0</td><td>0</td><td>1</td><td>% 3.33</td></tr><tr><td>股动脉穿刺组</td><td>30</td><td>1</td><td>2</td><td>1</td><td>3</td><td>1</td><td>26.67</td></tr><tr><td></td><td></td><td></td><td></td><td></td><td></td><td></td><td>4.706</td></tr><tr><td>x² P</td><td></td><td></td><td></td><td></td><td></td><td></td><td>0.030</td></tr></table></body></html>
|
||||
|
||||
# 3 讨论
|
||||
|
||||
经股动脉穿刺为介入治疗的传统途径[9-10]。但是由于股动脉解剖部位较深,穿刺具有局限性特别是接受抗凝治疗、强化抗血小板治疗及肥胖等患者术后局部压迫止血较为困难,可诱发尿潴留、动静脉瘘、体位性低血压、心动过缓、血肿及出血等并发症。支架置入术后患者出现严重迷走反射可引发血栓形成,严重
|
||||
|
||||
11 - 12 。肺栓塞、深静脉血栓形成等。刘宇凡[1} 研究结果显示经挠动脉穿刺支架置入术效果较好。
|
||||
|
||||
目前经股、桡动脉穿刺技术已趋于成熟。经桡动脉穿刺优点:(1)适于平卧患者:(2)恢复快可有效减少护理量;(3)手术后不需要肢体制动;(4)血管表浅,易止血。经挠动脉穿刺缺点:挠动脉血管管径较小操作较为困难。经股动脉穿刺优点:(1)可通过股动脉穿刺检测患者血流动力学指标;(2)穿刺部位易处理;(3)可耐受较长时间的手术操作;(4)导管及鞘管摆放较为容易:(5)穿刺成功率较高;(6)股动脉血管较粗大有利于器械操作。经股动脉穿刺缺点:(1)可引发腹膜后血肿等严重并发症;(2)术后需长时间下肢制动;(3)术后需人工按压穿刺点耗费手术时间影响14 。
|
||||
|
||||
Ha 等[15] 研究显示 1 例颈内动脉夹层患者由于腹主动脉闭塞而接受经挠动脉入路行支架形成术,但效。Montorsi 16动脉内支架成形术较为安全、可靠,可减少血管并发症。Daou 等[1"报道,经桡动脉入路治疗脑血管病变、 。 Mendiz 18研究中分别经股动脉入路、桡动脉入路行颈内动脉支架形成术均取得成功,且心血管事件发生率比较差异无统计学意义。本研究结果显示,桡动脉穿刺组肢体制动时间、止血包扎时间短于股动脉穿刺组鞘管置入时间、手术时间长于股动脉穿刺组。两组患者手术成功率和穿刺成功率比较差异无统计学意义。以上结果说明经桡动脉穿刺治疗效果更好,与 Daou 等[17]研究结果基本一致。
|
||||
|
||||
动脉痉挛为介入术中的一项常见并发症,其发生与术中是否使用扩血管药物、鞘管直径及挠动脉直径有关。本研究结果显示股动脉穿刺组肺栓塞、术后尿潴留、动脉闭塞、穿刺部位血肿、动脉痉挛等并发症发生率高于动脉穿刺组,说明经桡动脉穿刺行支架植入术治疗颅内动脉狭窄可有效减少并发症的发生。动脉痉挛发生率降低可能与术中操作轻柔、术前通过宣教减轻患者紧张情绪有关。动脉闭塞的发生可能与长时间阻断动脉血流、局部血栓形成率增加有关。有研究提出,保持动脉血流在术后压迫止血时顺畅可有效预防动脉闭塞[2},但是本研究中股动脉穿刺点局部压迫时间较长易影响局部血流通畅,而经桡动脉穿刺对桡动脉损伤较小压迫止血时间较短且能够保持动脉通畅。
|
||||
|
||||
综上所述经桡动脉穿刺行支架置入术治疗颅内动脉狭窄可优化手术指标,降低并发症发生率。本研究还需增加样本量、延长随访时间,以探究经桡动脉穿刺行支架植入术的长期疗效验证研究结果可信度,为临床提供更多参考。
|
||||
|
||||
# 参考文献
|
||||
|
||||
1 . Angio - seal置与手工压迫的短期并发症比较[J].心血管康复医学杂志,2020 29 5 575 - 578.
|
||||
[2]莫伟,向华 阳秀春 等.股动脉穿刺介入术后制动时间的循证证J . 2019 28 1 85 - 88.
|
||||
[3]刘恋高峰 莫大鹏 等.经动脉入路支架成形术治疗颅内椎-基底动脉狭窄[J].中华神经外科杂志 $2015\ 31(\,9):899\,-902.$
|
||||
[4]王辉 彭文近,刘艳红 等.老年患者经远端桡动脉与经典桡动脉行冠状动脉诊疗的临床疗效及安全性比较[J].中华老年心脑血2020 22 4 47 - 50.
|
||||
[5]孔杰,占大权.经股动脉入路和经桡动脉入路行全脑血管造影术J . 2020 15 10 26 - 28.
|
||||
[6]冯文勇,罗宁 高春萍 等.两种导管在经桡动脉入路脑血管造影J . 2020 17 11641 - 647.
|
||||
[7]臧君仁侯国锐,成瑞博 等.经桡动脉入路冠状动脉造影及支架植入术治疗老年冠心病患者的效果[J].河南医学研究,2019,28 23 4260 - 4261.
|
||||
[8]李五记.经桡动脉冠状动脉造影后即行支架植入应用于老年冠心病心绞痛的效果[J].现代临床医学 $2020\ A6(\ 1):13-14$ .
|
||||
9 NOORI V J JENS E J. A systematic review of vascular closure devicesfor femoral artery puncture sites J . J Vasc Surg 2018 68 3 887 -899.
|
||||
[10]敬子洋.经超声引导下股动脉穿刺建立介入治疗入路的临床应用分析[J].影像研究与医学应用 $2020\ {\cal A}(\ 24):205-206$ .
|
||||
[11] 庞猛.脑血管介入术后股动脉压迫止血致迷走神经反射的相关因J . 2019 19 A2 46.
|
||||
[12] 丁翔,王志权 ,朱启明 等.经动脉和股动脉途径行冠状动脉介932 J . 2015 31 7742 - 745.
|
||||
[13]刘宇凡.经桡动脉穿刺进行椎动脉狭窄支架成形术的有效性和安J . 2020 39 6 85 - 87.
|
||||
[14]梁珂,黄志志 李绍发 等.经桡动脉入路弓上动脉内支架成形术J . 2018 27 3 207 - 210.
|
||||
15 HA J H KIM H PARK I S et al. Trans - radial carotid artery steningin a patient with abdominal aortic occlusion J . J Cerebrovasc Endo-vasc Neurosurg 2016 18 2 129 - 134.
|
||||
16 MONTORSI P GALLI S RAVAGNANI P M et al. Carotid ArteryStenting With Proximal Embolic Protection via a Transradial or Trans-brachial Approach Pushing the Boundaries of the Technique WhileMaintaining Safety and Efficacy J . J Endovasc Ther 2016 23 4549 - 560.
|
||||
17 DAOU B CHALOUHI N TJOUMAKARIS S et al. Altenmative ac-cess for endovascular treatment of cerebrpvascular diseases J . ClinNeurol Neurosurg 2016 145 251 89 - 95.
|
||||
18 MENDIZ O A FAVA C LEV G et al. Transradial Versus TransfemoralCarotid Artery Stenting A 16 - Year Single - Center Experience J . JInterv Cardiol 2016 29 6 588 - 593.
|
||||
|
||||
2021 - 06 - 18
|
||||
BIN
app_pkg/examples/邱茜茜.pdf
Normal file
BIN
app_pkg/examples/邱茜茜.pdf
Normal file
Binary file not shown.
0
app_pkg/flow/__init__.py
Normal file
0
app_pkg/flow/__init__.py
Normal file
15
app_pkg/flow/main.py
Normal file
15
app_pkg/flow/main.py
Normal file
@ -0,0 +1,15 @@
|
||||
from database.milvus_api import Milvus_Database
|
||||
from modelsAPI.model_api import get_embeddings
|
||||
|
||||
async def autorag_flow(query,milvus_db,collection_name,top_k=5):
|
||||
query_vector = get_embeddings(list(query))
|
||||
search_results = milvus_db.search(collection_name=collection_name,query_vector=query_vector, top_k=top_k)
|
||||
|
||||
search_prompt="""你作为一个由中科院的资深文档问答专家,请你从以下文本信息中选取与用户问题相关的内容作为参考:
|
||||
-----
|
||||
文本信息:{context}
|
||||
-----
|
||||
请你仔细思考后用中文回答用户的问题,如果文本中没有关于问题的答案,则直接回答 文档中没有相关内容
|
||||
用户提问:{query}
|
||||
你的回答:"""
|
||||
|
||||
0
app_pkg/markdown_process/__init__.py
Normal file
0
app_pkg/markdown_process/__init__.py
Normal file
119
app_pkg/markdown_process/breakdown_txt.py
Normal file
119
app_pkg/markdown_process/breakdown_txt.py
Normal file
@ -0,0 +1,119 @@
|
||||
|
||||
def force_breakdown(txt, limit, get_token_fn):
|
||||
""" 当无法用标点、空行分割时,我们用最暴力的方法切割
|
||||
"""
|
||||
for i in reversed(range(len(txt))):
|
||||
if get_token_fn(txt[:i]) < limit:
|
||||
return txt[:i], txt[i:]
|
||||
return "Tiktoken未知错误", "Tiktoken未知错误"
|
||||
|
||||
|
||||
def maintain_storage(remain_txt_to_cut, remain_txt_to_cut_storage):
|
||||
""" 为了加速计算,我们采样一个特殊的手段。当 remain_txt_to_cut > `_max` 时, 我们把 _max 后的文字转存至 remain_txt_to_cut_storage
|
||||
当 remain_txt_to_cut < `_min` 时,我们再把 remain_txt_to_cut_storage 中的部分文字取出
|
||||
"""
|
||||
_min = int(5e4)
|
||||
_max = int(1e5)
|
||||
# print(len(remain_txt_to_cut), len(remain_txt_to_cut_storage))
|
||||
if len(remain_txt_to_cut) < _min and len(remain_txt_to_cut_storage) > 0:
|
||||
remain_txt_to_cut = remain_txt_to_cut + remain_txt_to_cut_storage
|
||||
remain_txt_to_cut_storage = ""
|
||||
if len(remain_txt_to_cut) > _max:
|
||||
remain_txt_to_cut_storage = remain_txt_to_cut[_max:] + remain_txt_to_cut_storage
|
||||
remain_txt_to_cut = remain_txt_to_cut[:_max]
|
||||
return remain_txt_to_cut, remain_txt_to_cut_storage
|
||||
|
||||
|
||||
def cut(limit, get_token_fn, txt_tocut, must_break_at_empty_line, break_anyway=False):
|
||||
""" 文本切分
|
||||
"""
|
||||
res = []
|
||||
total_len = len(txt_tocut)
|
||||
fin_len = 0
|
||||
remain_txt_to_cut = txt_tocut
|
||||
remain_txt_to_cut_storage = ""
|
||||
# 为了加速计算,我们采样一个特殊的手段。当 remain_txt_to_cut > `_max` 时, 我们把 _max 后的文字转存至 remain_txt_to_cut_storage
|
||||
remain_txt_to_cut, remain_txt_to_cut_storage = maintain_storage(remain_txt_to_cut, remain_txt_to_cut_storage)
|
||||
|
||||
while True:
|
||||
if get_token_fn(remain_txt_to_cut) <= limit:
|
||||
# 如果剩余文本的token数小于限制,那么就不用切了
|
||||
res.append(remain_txt_to_cut);
|
||||
fin_len += len(remain_txt_to_cut)
|
||||
break
|
||||
else:
|
||||
# 如果剩余文本的token数大于限制,那么就切
|
||||
lines = remain_txt_to_cut.split('\n')
|
||||
|
||||
# 估计一个切分点
|
||||
estimated_line_cut = limit / get_token_fn(remain_txt_to_cut) * len(lines)
|
||||
estimated_line_cut = int(estimated_line_cut)
|
||||
|
||||
# 开始查找合适切分点的偏移(cnt)
|
||||
cnt = 0
|
||||
for cnt in reversed(range(estimated_line_cut)):
|
||||
if must_break_at_empty_line:
|
||||
# 首先尝试用双空行(\n\n)作为切分点
|
||||
if lines[cnt] != "":
|
||||
continue
|
||||
prev = "\n".join(lines[:cnt])
|
||||
post = "\n".join(lines[cnt:])
|
||||
if get_token_fn(prev) < limit:
|
||||
break
|
||||
|
||||
if cnt == 0:
|
||||
# 如果没有找到合适的切分点
|
||||
if break_anyway:
|
||||
# 是否允许暴力切分
|
||||
prev, post = force_breakdown(remain_txt_to_cut, limit, get_token_fn)
|
||||
else:
|
||||
# 不允许直接报错
|
||||
raise RuntimeError(f"存在一行极长的文本!{remain_txt_to_cut}")
|
||||
|
||||
# 追加列表
|
||||
res.append(prev);
|
||||
fin_len += len(prev)
|
||||
# 准备下一次迭代
|
||||
remain_txt_to_cut = post
|
||||
remain_txt_to_cut, remain_txt_to_cut_storage = maintain_storage(remain_txt_to_cut,
|
||||
remain_txt_to_cut_storage)
|
||||
process = fin_len / total_len
|
||||
# print(f'正在文本切分 {int(process * 100)}%')
|
||||
if len(remain_txt_to_cut.strip()) == 0:
|
||||
break
|
||||
return res
|
||||
|
||||
def breakdown_text_to_satisfy_token_limit(txt, limit, llm_model="gpt-3.5-turbo"):
|
||||
""" 使用多种方式尝试切分文本,以满足 token 限制
|
||||
"""
|
||||
# from tools.bridge_all import LazyloadTiktoken
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
|
||||
# tokenizer_gpt35 = LazyloadTiktoken("gpt-3.5-turbo")
|
||||
tokenizer_gpt35 = AutoTokenizer.from_pretrained('./pretrained_models/chinese-llama-2-7b-64k')
|
||||
|
||||
#enc = model_info[llm_model]['tokenizer']
|
||||
enc = tokenizer_gpt35
|
||||
def get_token_fn(txt):
|
||||
return len(enc.encode(txt))
|
||||
try:
|
||||
# 第1次尝试,将双空行(\n\n)作为切分点
|
||||
return cut(limit, get_token_fn, txt, must_break_at_empty_line=True)
|
||||
except RuntimeError:
|
||||
try:
|
||||
# 第2次尝试,将单空行(\n)作为切分点
|
||||
return cut(limit, get_token_fn, txt, must_break_at_empty_line=False)
|
||||
except RuntimeError:
|
||||
try:
|
||||
# 第3次尝试,将英文句号(.)作为切分点
|
||||
res = cut(limit, get_token_fn, txt.replace('.', '。\n'), must_break_at_empty_line=False) # 这个中文的句号是故意的,作为一个标识而存在
|
||||
return [r.replace('。\n', '.') for r in res]
|
||||
except RuntimeError as e:
|
||||
try:
|
||||
# 第4次尝试,将中文句号(。)作为切分点
|
||||
res = cut(limit, get_token_fn, txt.replace('。', '。。\n'), must_break_at_empty_line=False)
|
||||
return [r.replace('。。\n', '。') for r in res]
|
||||
except RuntimeError as e:
|
||||
# 第5次尝试,没办法了,随便切一下吧
|
||||
return cut(limit, get_token_fn, txt, must_break_at_empty_line=False, break_anyway=True)
|
||||
63
app_pkg/markdown_process/chunk_postprocess.py
Normal file
63
app_pkg/markdown_process/chunk_postprocess.py
Normal file
@ -0,0 +1,63 @@
|
||||
from modelsAPI.model_api import OpenaiAPI
|
||||
import os
|
||||
def context_chunk_process(markdown_context,chunks):
|
||||
def prepare_chunk(WHOLE_DOCUMENT,CHUNK_CONTENT):
|
||||
prompt = f"""
|
||||
<document>
|
||||
{WHOLE_DOCUMENT}
|
||||
</document>
|
||||
Here is the chunk we want to situate within the whole document
|
||||
<chunk>
|
||||
{CHUNK_CONTENT}
|
||||
</chunk>
|
||||
Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else.
|
||||
"""
|
||||
chinese_prompt = f"""
|
||||
<文档>
|
||||
{WHOLE_DOCUMENT}
|
||||
</文档>
|
||||
以下是需要定位到整体文档中的片段
|
||||
<片段>
|
||||
{CHUNK_CONTENT}
|
||||
</片段>
|
||||
请简要说明该片段在整体文档中的上下文关系,以提升片段搜索效果。只需给出简洁的上下文说明,无需其他内容。
|
||||
"""
|
||||
return chinese_prompt
|
||||
for chunk in chunks:
|
||||
text = chunk.get("text","")
|
||||
CONTEXT_CHUNK = prepare_chunk(markdown_context,text)
|
||||
NEW_CHUNK = OpenaiAPI().open_api_caht_without_thinking(query=CONTEXT_CHUNK,model=os.getenv("OPENAI_MODEL","qwen3-32b"))
|
||||
chunk["text"] = NEW_CHUNK + text
|
||||
chunk["origin_text"] = text
|
||||
return chunks
|
||||
|
||||
def abstract_chunk(markdown_context,chunks):
|
||||
def prepare_chunk(WHOLE_DOCUMENT,CHUNK_CONTENT):
|
||||
prompt = f"""
|
||||
<document>
|
||||
{WHOLE_DOCUMENT}
|
||||
</document>
|
||||
Here is the chunk we want to situate within the whole document
|
||||
<chunk>
|
||||
{CHUNK_CONTENT}
|
||||
</chunk>
|
||||
Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else.
|
||||
"""
|
||||
chinese_prompt = f"""
|
||||
<文档>
|
||||
{WHOLE_DOCUMENT}
|
||||
</文档>
|
||||
以下是需要定位到整体文档中的片段
|
||||
<片段>
|
||||
{CHUNK_CONTENT}
|
||||
</片段>
|
||||
请简要说明该片段在整体文档中的上下文关系,以提升片段搜索效果。只需给出简洁的上下文说明,无需其他内容。
|
||||
"""
|
||||
return chinese_prompt
|
||||
for chunk in chunks:
|
||||
text = chunk.get("text","")
|
||||
CONTEXT_CHUNK = prepare_chunk(markdown_context,text)
|
||||
NEW_CHUNK = OpenaiAPI().open_api_caht_without_thinking(query=CONTEXT_CHUNK,model=os.getenv("OPENAI_MODEL","qwen3-32b"))
|
||||
chunk["text"] = NEW_CHUNK
|
||||
chunk["origin_text"] = NEW_CHUNK + text
|
||||
return chunks
|
||||
326
app_pkg/markdown_process/convert_header.py
Normal file
326
app_pkg/markdown_process/convert_header.py
Normal file
@ -0,0 +1,326 @@
|
||||
import copy
|
||||
import re
|
||||
|
||||
def convert_numbered_headings(text):
|
||||
"""
|
||||
将编号转换为对应的Markdown标题:
|
||||
- `1.2` → `## 1.2`
|
||||
- `1.2.1` → `### 1.2.1`
|
||||
- `1. 3. 1` → `### 1.3.1`(去掉空格)
|
||||
"""
|
||||
# 替换全角 `.` 为半角 `.`
|
||||
text = text.replace('.', '.')
|
||||
|
||||
# 遍历每一行,处理标题
|
||||
lines = text.split('\n')
|
||||
|
||||
for i in range(len(lines)):
|
||||
line = lines[i].strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
# 只匹配 `1.2` 或 `1.2.1` 形式的编号,允许 `1. 3. 1` 这种格式
|
||||
match = re.match(r'^(#*\s*)(\d+(?:\s*\.\s*\d+)+)', line)
|
||||
if match:
|
||||
prefix, num = match.groups()
|
||||
level = num.count('.') # `.` 的个数决定层级
|
||||
new_prefix = '#' * (level + 1) # 1.2 → ##, 1.2.1 → ###
|
||||
|
||||
# 去掉数字和 `.` 之间的空格
|
||||
clean_num = re.sub(r'\s*\.\s*', '.', num)
|
||||
|
||||
# 替换标题
|
||||
lines[i] = f"{new_prefix} {clean_num}{line[match.end():]}"
|
||||
|
||||
return '\n'.join(lines)
|
||||
|
||||
|
||||
|
||||
def custom_convert_numbered_headings(text, max_level=9):
|
||||
"""
|
||||
将编号转换为对应的Markdown标题,支持自定义最大分级层数:
|
||||
- `1.2` → `## 1.2` (当max_level>=2)
|
||||
- `1.2.1` → `### 1.2.1` (当max_level>=3)
|
||||
- `1. 3. 1` → `### 1.3.1`(去掉空格)
|
||||
|
||||
参数:
|
||||
text: 要处理的文本
|
||||
max_level: 最大支持的分级层数(默认为3)
|
||||
"""
|
||||
# 替换全角 `.` 为半角 `.`
|
||||
text = text.replace('.', '.')
|
||||
|
||||
# 遍历每一行,处理标题
|
||||
lines = text.split('\n')
|
||||
|
||||
for i in range(len(lines)):
|
||||
line = lines[i].strip()
|
||||
if not line:
|
||||
continue
|
||||
|
||||
# 只匹配 `1.2` 或 `1.2.1` 形式的编号,允许 `1. 3. 1` 这种格式
|
||||
match = re.match(r'^(#*\s*)(\d+(?:\s*\.\s*\d+)+)', line)
|
||||
if match:
|
||||
prefix, num = match.groups()
|
||||
level = num.count('.') # `.` 的个数决定层级
|
||||
|
||||
# 只处理不超过max_level的标题
|
||||
if level <= max_level:
|
||||
new_prefix = '#' * (level + 1) # 1.2 → ##, 1.2.1 → ###
|
||||
|
||||
# 去掉数字和 `.` 之间的空格
|
||||
clean_num = re.sub(r'\s*\.\s*', '.', num)
|
||||
|
||||
# 替换标题
|
||||
lines[i] = f"{new_prefix} {clean_num}{line[match.end():]}"
|
||||
else:
|
||||
# 超过max_level的标题保持不变
|
||||
lines[i] = line
|
||||
|
||||
return lines #'\n'.join(lines)
|
||||
|
||||
def parse_markdown(md_text, max_level=6):
|
||||
# md_text = convert_numbered_headings(md_text)
|
||||
md_text = custom_convert_numbered_headings(md_text, max_level=max_level)
|
||||
# lines = md_text.strip().split("\n")
|
||||
lines = md_text
|
||||
results = []
|
||||
current_headers = [None] * max_level # 动态存储各级标题(索引0=H1,1=H2...)
|
||||
current_text = []
|
||||
|
||||
def add_entry():
|
||||
nonlocal current_text
|
||||
if current_text or any(current_headers):
|
||||
# 拼接有效标题到 text 中(跳过 None)
|
||||
full_text = []
|
||||
for level in range(max_level):
|
||||
if current_headers[level] is not None:
|
||||
full_text.append(f"{'#' * (level + 1)} {current_headers[level]}")
|
||||
full_text.extend(current_text)
|
||||
|
||||
# 生成结果(动态分配 Header_1, Header_2...)
|
||||
entry = {"text": "".join(full_text)}
|
||||
for level in range(max_level):
|
||||
entry[f"Header_{level + 1}"] = current_headers[level]
|
||||
results.append(entry)
|
||||
current_text = []
|
||||
|
||||
# 动态生成正则表达式,匹配 1 到 max_level 个 #
|
||||
header_pattern = re.compile(r'^(#{1,%d})\s*(.*)' % max_level)
|
||||
|
||||
for line in lines:
|
||||
line = line
|
||||
header_match = header_pattern.match(line)
|
||||
|
||||
if header_match:
|
||||
add_entry()
|
||||
level = len(header_match.group(1)) - 1 # 转换为 0-based 索引
|
||||
title = header_match.group(2)
|
||||
|
||||
if level < max_level:
|
||||
# 更新当前标题,并清空更低层级的标题
|
||||
current_headers[level] = title
|
||||
for l in range(level + 1, max_level):
|
||||
current_headers[l] = None
|
||||
elif line:
|
||||
current_text.append(line)
|
||||
|
||||
add_entry()
|
||||
return results
|
||||
|
||||
# def parse_markdown(md_text,max_level=3):
|
||||
# # md_text = convert_numbered_headings(md_text)
|
||||
# md_text = custom_convert_numbered_headings(md_text,max_level=max_level)
|
||||
# lines = md_text.strip().split("\n")
|
||||
# results = []
|
||||
# current_h1, current_h2, current_h3 = None, None, None
|
||||
# current_text = []
|
||||
|
||||
# def add_entry():
|
||||
# if current_h1 or current_h2 or current_h3:
|
||||
# results.append({
|
||||
# "Header_1": current_h1,
|
||||
# "Header_2": current_h2,
|
||||
# "Header_3": current_h3,
|
||||
# "text": "\n".join(current_text).strip()
|
||||
# })
|
||||
# current_text.clear()
|
||||
|
||||
# for line in lines:
|
||||
# line = line.strip()
|
||||
# header_match = re.match(r'^(#{1,3})\s*(.*)', line)
|
||||
|
||||
# if header_match:
|
||||
# add_entry()
|
||||
# level = len(header_match.group(1))
|
||||
# title = header_match.group(2).strip()
|
||||
|
||||
# if level == 1:
|
||||
# current_h1, current_h2, current_h3 = title, None, None
|
||||
# elif level == 2:
|
||||
# current_h2, current_h3 = title, None
|
||||
# elif level == 3:
|
||||
# current_h3 = title
|
||||
# else:
|
||||
# current_text.append(line)
|
||||
|
||||
# add_entry()
|
||||
# return results
|
||||
|
||||
def header_postprocess(md_header_json):
|
||||
processed_results = []
|
||||
buffer_entry = None
|
||||
|
||||
for md in md_header_json:
|
||||
headers = {k: md.get(k) for k in md if k.startswith("Header_")}
|
||||
text = md.get("text", "")
|
||||
|
||||
md["origin_text"] = copy.deepcopy(text) # 确保完全独立
|
||||
|
||||
# 修复后的判断逻辑
|
||||
is_header_only = (
|
||||
not text.strip() or
|
||||
any(text.strip() == headers.get(f"Header_{i}", "") for i in range(1, 5)
|
||||
))
|
||||
|
||||
if is_header_only:
|
||||
if buffer_entry is None:
|
||||
buffer_entry = md.copy()
|
||||
else:
|
||||
# 合并标题,保留最高级标题
|
||||
for level in range(1, 5):
|
||||
key = f"Header_{level}"
|
||||
if key in md and md[key] is not None:
|
||||
buffer_entry[key] = md[key]
|
||||
else:
|
||||
if buffer_entry is not None:
|
||||
# 合并缓存的标题和当前内容
|
||||
new_text = []
|
||||
for level in range(1, 5):
|
||||
key = f"Header_{level}"
|
||||
if buffer_entry.get(key):
|
||||
new_text.append(f"{'#' * level} {buffer_entry[key]}")
|
||||
new_text.append(text)
|
||||
md["text"] = "".join(new_text)
|
||||
|
||||
for level in range(1, 5):
|
||||
key = f"Header_{level}"
|
||||
md[key] = buffer_entry.get(key)
|
||||
|
||||
buffer_entry = None
|
||||
processed_results.append(md)
|
||||
|
||||
if buffer_entry is not None:
|
||||
processed_results.append(buffer_entry)
|
||||
|
||||
return processed_results
|
||||
|
||||
def postprocess_chunks(chunks):
|
||||
"""合并被分割的原型图上下文(适配_process_chunk生成的结构)"""
|
||||
new_chunks = []
|
||||
i = 0
|
||||
|
||||
# 修正后的图片匹配正则表达式(确保括号配对)
|
||||
image_pattern = re.compile(
|
||||
r'(!\[.*?\]\(.*?\))|' # Markdown图片 
|
||||
r'(<img.*?src=["\'].*?["\'].*?>)|' # HTML <img>标签
|
||||
r'(\[.*?\]\(.*?\))|' # Markdown链接 [text](url)
|
||||
r'(https?://\S+\.(?:jpe?g|png|gif|bmp|svg|webp)(?:\?\S*)?)', # 纯图片URL
|
||||
re.IGNORECASE
|
||||
)
|
||||
|
||||
# 修正后的标题匹配正则表达式
|
||||
title_pattern = re.compile(
|
||||
r'(#\s*[0-9一二三四五六七八九十]+[\..、::]?\s*[原形图圖示示])|'
|
||||
r'(prototype|diagram|screenshot|界面图)',
|
||||
re.IGNORECASE
|
||||
)
|
||||
|
||||
while i < len(chunks):
|
||||
current_chunk = chunks[i]
|
||||
current_text = current_chunk.get("text", "")
|
||||
|
||||
# 检查是否是原型图标题块
|
||||
is_prototype_title = title_pattern.search(current_text)
|
||||
|
||||
# 情况1:当前是原型图标题但未包含图片
|
||||
if is_prototype_title and not image_pattern.search(current_text):
|
||||
# 向后查找最多3个块寻找图片
|
||||
for j in range(i+1, min(i+4, len(chunks))):
|
||||
next_chunk = chunks[j]
|
||||
next_text = next_chunk.get("text", "")
|
||||
|
||||
if image_pattern.search(next_text):
|
||||
# 合并元数据
|
||||
merged_chunk = {
|
||||
"id": current_chunk["id"],
|
||||
"Header_1": current_chunk.get("Header_1", "") or next_chunk.get("Header_1", ""),
|
||||
"Header_2": current_chunk.get("Header_2", "") or next_chunk.get("Header_2", ""),
|
||||
"Header_3": current_chunk.get("Header_3", "") or next_chunk.get("Header_3", ""),
|
||||
"text": current_text + next_text,
|
||||
"resource": current_chunk["resource"],
|
||||
"img_path": ",".join(filter(None, [
|
||||
current_chunk.get("img_path", ""),
|
||||
next_chunk.get("img_path", "")
|
||||
])) or None
|
||||
}
|
||||
new_chunks.append(merged_chunk)
|
||||
i = j + 1 # 跳过已合并的块
|
||||
break
|
||||
else:
|
||||
# 没找到图片则保留原块
|
||||
new_chunks.append(current_chunk)
|
||||
i += 1
|
||||
else:
|
||||
# 情况2:普通块处理
|
||||
new_chunks.append(current_chunk)
|
||||
i += 1
|
||||
|
||||
return new_chunks
|
||||
|
||||
# def header_postprocess(md_header_json):
|
||||
# processed_results = []
|
||||
# buffer_entry = None
|
||||
|
||||
# for md in md_header_json:
|
||||
# header_1 = md.get("Header_1")
|
||||
# header_2 = md.get("Header_2")
|
||||
# header_3 = md.get("Header_3")
|
||||
# text = md.get("text", "")
|
||||
|
||||
# # 如果当前条目只有标题没有正文,则缓存起来
|
||||
# if not text.strip() or text.strip() == header_1 or text.strip() == header_2 or text.strip() == header_3:
|
||||
# if buffer_entry is None:
|
||||
# buffer_entry = md.copy()
|
||||
# else:
|
||||
# # 合并标题
|
||||
# for level in range(1, 4):
|
||||
# key = f"Header_{level}"
|
||||
# if key in md and md[key] is not None:
|
||||
# buffer_entry[key] = md[key]
|
||||
# else:
|
||||
# # 如果有正文内容,检查是否有缓存的标题
|
||||
# if buffer_entry is not None:
|
||||
# # 合并缓存的标题和当前内容
|
||||
# new_text = []
|
||||
# for level in range(1, 4):
|
||||
# key = f"Header_{level}"
|
||||
# if buffer_entry.get(key):
|
||||
# new_text.append(f"{'#' * level} {buffer_entry[key]}")
|
||||
# new_text.append(text)
|
||||
# md["text"] = "\n".join(new_text)
|
||||
|
||||
# # 更新标题信息
|
||||
# for level in range(1, 4):
|
||||
# key = f"Header_{level}"
|
||||
# md[key] = buffer_entry.get(key)
|
||||
|
||||
# buffer_entry = None
|
||||
# processed_results.append(md)
|
||||
|
||||
# # 处理最后可能剩余的缓存条目
|
||||
# if buffer_entry is not None:
|
||||
# # 如果最后只剩下标题没有正文,单独作为一个条目
|
||||
# processed_results.append(buffer_entry)
|
||||
|
||||
# return processed_results
|
||||
297
app_pkg/markdown_process/markdown_preprocess.py
Normal file
297
app_pkg/markdown_process/markdown_preprocess.py
Normal file
@ -0,0 +1,297 @@
|
||||
from fastapi import UploadFile
|
||||
import httpx
|
||||
import os
|
||||
from typing import Optional
|
||||
import tempfile
|
||||
import requests
|
||||
import time
|
||||
from tqdm import tqdm
|
||||
from pathlib import Path
|
||||
import json
|
||||
from io import BytesIO
|
||||
from urllib.parse import urljoin
|
||||
|
||||
|
||||
async def process_pdf_logic(pdf_file):
|
||||
"""
|
||||
接收 PDF 文件并转发到解析服务
|
||||
"""
|
||||
# 验证文件类型
|
||||
if not pdf_file.content_type.startswith("application/pdf"):
|
||||
raise Exception(400, "仅支持 PDF 文件")
|
||||
|
||||
try:
|
||||
# 读取文件内容
|
||||
# file_content = await pdf_file.read()
|
||||
# 配置目标接口参数
|
||||
target_url = os.getenv("Mineru_URL", "http://172.18.30.122:9988")
|
||||
target_url = target_url + "/analyze-pdf"
|
||||
# target_url = os.getenv("Mineru_URL","http://172.18.127.124:9088/pdf_parse")
|
||||
# target_url = "http://172.18.30.122:8998/pdf_parse"
|
||||
# params = {
|
||||
# "parse_method": "auto",
|
||||
# "is_json_md_dump": "true",
|
||||
# "output_dir": "output"
|
||||
# }
|
||||
|
||||
# 构建文件上传格式
|
||||
files = {
|
||||
"file": (pdf_file.filename, pdf_file.file, "application/pdf")
|
||||
# "pdf_file": (pdf_file.filename, pdf_file.file, "application/pdf")
|
||||
}
|
||||
|
||||
# 发送请求到解析服务
|
||||
async with httpx.AsyncClient(timeout=1200.0) as client:
|
||||
response = await client.post(
|
||||
target_url,
|
||||
# params=params,
|
||||
files=files,
|
||||
headers={"Accept": "application/json"}
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
content_chunk_list = []
|
||||
full_content = result.get("data", {}).get("full_content", "")
|
||||
content_list = result.get("data", {}).get("content_list", [])
|
||||
content_chunk_list = result.get("data", {}).get("content_middle", [])
|
||||
print(content_chunk_list)
|
||||
# combined_json = ''.join(content_list)
|
||||
# content_data = json.loads(combined_json)
|
||||
# for idx, item in enumerate(content_data):
|
||||
# if isinstance(item, dict): # 确保是字典类型
|
||||
# text_type = item.get("type", "")
|
||||
# text_level = item.get("text_level", "")
|
||||
# text = item.get("text", "")
|
||||
# page_idx = item.get("page_idx", -1) # 默认值 -1 表示缺失
|
||||
# table_caption = item.get("table_caption", "")
|
||||
# table_body = item.get("table_body", "")
|
||||
# image_path = item.get("img_path", "")
|
||||
# full_url = ""
|
||||
# if image_path:
|
||||
# base_url = os.getenv("HOST_IP", "http://172.18.30.122:9988") # 基础URL(不带/outputs)
|
||||
# outputs_url = urljoin(base_url + "/", "output/")
|
||||
# full_url = urljoin(outputs_url, image_path) # 最终拼接image_path
|
||||
# if not text and not table_body and not image_path:
|
||||
# continue
|
||||
# image_caption = item.get("img_caption", "")
|
||||
# content_chunk_list.append(
|
||||
# {"id": idx, "type": text_type, "text_level": text_level, "text": text, "page_idx": page_idx,
|
||||
# "table_caption": table_caption, "table_body": table_body, "image_path": full_url,
|
||||
# "image_caption": image_caption})
|
||||
# markdown_content =""
|
||||
# for content in content_list:
|
||||
# text = content["text"]
|
||||
|
||||
return full_content, content_chunk_list
|
||||
|
||||
except httpx.RequestError as e:
|
||||
raise Exception(503, f"无法连接到解析服务: {str(e)}")
|
||||
except Exception as e:
|
||||
raise Exception(500, f"服务器内部错误: {str(e)}")
|
||||
|
||||
|
||||
async def process_doc_logic(doc_file):
|
||||
"""
|
||||
接收Word文档(doc/docx)并转发到解析服务
|
||||
"""
|
||||
# 验证文件类型
|
||||
allowed_extensions = {'.doc', '.docx'}
|
||||
file_ext = os.path.splitext(doc_file.filename)[1].lower()
|
||||
|
||||
if file_ext not in allowed_extensions:
|
||||
raise Exception(400, "仅支持doc或docx格式文件")
|
||||
|
||||
try:
|
||||
# 配置目标接口参数
|
||||
target_url = os.getenv("MINERU_URL", "http://172.18.30.122:9988")
|
||||
target_url = target_url + "/analyze-doc"
|
||||
# 根据文件类型设置正确的MIME类型
|
||||
mime_type = "application/msword" if file_ext == '.doc' else "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||||
|
||||
# 构建文件上传格式
|
||||
files = {
|
||||
"file": (doc_file.filename, doc_file.file, mime_type)
|
||||
}
|
||||
|
||||
# 发送请求到解析服务
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
response = await client.post(
|
||||
target_url,
|
||||
files=files,
|
||||
headers={"Accept": "application/json"}
|
||||
)
|
||||
response.raise_for_status()
|
||||
result = response.json()
|
||||
|
||||
# 假设返回数据结构与PDF不同,需要相应调整
|
||||
# 这里根据实际API返回结构调整
|
||||
|
||||
return result["data"]["full_content"]
|
||||
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
raise Exception(e.response.status_code, f"文档解析服务返回错误: {str(e)}")
|
||||
except httpx.RequestError as e:
|
||||
raise Exception(503, f"无法连接到解析服务: {str(e)}")
|
||||
except Exception as e:
|
||||
raise Exception(500, f"服务器内部错误: {str(e)}")
|
||||
|
||||
|
||||
class ProgressFileReader:
|
||||
"""带有进度条的文件读取器"""
|
||||
|
||||
def __init__(self, file_path, progress_bar):
|
||||
self.file = open(file_path, 'rb')
|
||||
self.progress_bar = progress_bar
|
||||
|
||||
def read(self, size=-1):
|
||||
data = self.file.read(size)
|
||||
if data:
|
||||
self.progress_bar.update(len(data))
|
||||
return data
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.file.close()
|
||||
|
||||
|
||||
async def convert_doc_to_pdf(api_url: str, file: UploadFile, timeout: int = 300) -> Optional[str]:
|
||||
"""
|
||||
将DOC文件转换为pdf
|
||||
|
||||
参数:
|
||||
api_url: 转换服务的URL (如 "http://172.18.127.124:8000")
|
||||
file: 要转换的UploadFile对象
|
||||
timeout: 请求超时时间(秒)
|
||||
"""
|
||||
try:
|
||||
# 创建临时文件保存上传的DOC
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".doc") as tmp_doc:
|
||||
content = await file.read()
|
||||
tmp_doc.write(content)
|
||||
tmp_path = tmp_doc.name
|
||||
|
||||
# 获取文件信息
|
||||
file_name = file.filename
|
||||
file_size = os.path.getsize(tmp_path)
|
||||
output_name = Path(file_name).stem + ".pdf"
|
||||
|
||||
# 使用带进度条的文件读取器
|
||||
with tqdm(total=file_size, unit='B', unit_scale=True, desc="上传进度") as pbar:
|
||||
with ProgressFileReader(tmp_path, pbar) as file_reader:
|
||||
files = {'file': (file_name, file_reader)}
|
||||
params = {'src': 'doc', 'to': 'pdf'}
|
||||
|
||||
response = requests.post(
|
||||
f"{api_url}/transfer",
|
||||
files=files,
|
||||
params=params,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
# 处理响应
|
||||
if response.status_code == 200:
|
||||
# 创建临时文件保存转换后的DOCX
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_docx:
|
||||
tmp_docx.write(response.content)
|
||||
output_path = tmp_docx.name
|
||||
|
||||
# 清理临时文件
|
||||
os.unlink(tmp_path)
|
||||
|
||||
return output_path
|
||||
else:
|
||||
raise Exception(f"转换失败: {response.text}")
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
raise Exception(f"请求超时(超过{timeout}秒)")
|
||||
except requests.exceptions.RequestException as e:
|
||||
raise Exception(f"请求错误: {str(e)}")
|
||||
except Exception as e:
|
||||
# 确保临时文件被清理
|
||||
if 'tmp_path' in locals() and os.path.exists(tmp_path):
|
||||
os.unlink(tmp_path)
|
||||
if 'output_path' in locals() and os.path.exists(output_path):
|
||||
os.unlink(output_path)
|
||||
raise Exception(f"转换过程中出错: {str(e)}")
|
||||
|
||||
|
||||
async def convert_docx_to_pdf(api_url: str, file: UploadFile, timeout: int = 300) -> Optional[str]:
|
||||
"""
|
||||
将DOCX文件转换为PDF(解决中文缺失问题)
|
||||
|
||||
参数:
|
||||
api_url: 转换服务的URL (如 "http://172.18.127.124:8000")
|
||||
file: 要转换的UploadFile对象
|
||||
timeout: 请求超时时间(秒)
|
||||
"""
|
||||
try:
|
||||
# 创建临时文件保存上传的DOCX
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".docx") as tmp_docx:
|
||||
content = await file.read()
|
||||
tmp_docx.write(content)
|
||||
tmp_path = tmp_docx.name
|
||||
|
||||
# 获取文件信息
|
||||
file_name = file.filename
|
||||
file_size = os.path.getsize(tmp_path)
|
||||
output_name = Path(file_name).stem + ".pdf"
|
||||
|
||||
print(f"正在转换: {file_name} ({file_size / 1024:.1f} KB) -> PDF")
|
||||
start_time = time.time()
|
||||
|
||||
# 使用带进度条的文件读取器
|
||||
with tqdm(total=file_size, unit='B', unit_scale=True, desc="上传进度") as pbar:
|
||||
with ProgressFileReader(tmp_path, pbar) as file_reader:
|
||||
files = {'file': (file_name, file_reader)}
|
||||
params = {
|
||||
'src': 'docx',
|
||||
'to': 'pdf',
|
||||
'font': 'SimSun', # 指定中文字体
|
||||
'encoding': 'utf-8' # 明确编码
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
f"{api_url}/transfer",
|
||||
files=files,
|
||||
params=params,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
# 处理响应
|
||||
if response.status_code == 200:
|
||||
# 创建临时文件保存转换后的PDF
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_pdf:
|
||||
tmp_pdf.write(response.content)
|
||||
output_path = tmp_pdf.name
|
||||
|
||||
# 验证PDF是否包含文字
|
||||
if os.path.getsize(output_path) < 1024: # 如果PDF文件过小
|
||||
raise Exception("生成的PDF可能不包含有效内容")
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"\n转换成功! 保存为: {output_path}")
|
||||
print(f"总耗时: {elapsed:.2f}秒")
|
||||
|
||||
# 清理临时文件
|
||||
os.unlink(tmp_path)
|
||||
|
||||
return output_path
|
||||
else:
|
||||
raise Exception(f"转换失败: {response.text}")
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
raise Exception(f"请求超时(超过{timeout}秒)")
|
||||
except requests.exceptions.RequestException as e:
|
||||
raise Exception(f"请求错误: {str(e)}")
|
||||
except Exception as e:
|
||||
# 确保临时文件被清理
|
||||
if 'tmp_path' in locals() and os.path.exists(tmp_path):
|
||||
os.unlink(tmp_path)
|
||||
if 'output_path' in locals() and os.path.exists(output_path):
|
||||
os.unlink(output_path)
|
||||
raise Exception(f"转换过程中出错: {str(e)}")
|
||||
|
||||
20
app_pkg/markdown_process/save_markdown.py
Normal file
20
app_pkg/markdown_process/save_markdown.py
Normal file
@ -0,0 +1,20 @@
|
||||
from database.milvus_api import Milvus_Database
|
||||
import os
|
||||
|
||||
class Save_Markdown_to_Milvus():
|
||||
def __init__(self,database_name, collection_name):
|
||||
self.database_name = database_name
|
||||
self.collection_name = collection_name
|
||||
# self.milvus_database = Milvus_Database(database_name=self.database_name).client
|
||||
self.milvus_database = Milvus_Database(uri=os.getenv("Milvus_URI"))
|
||||
def save_md(self,data):
|
||||
# milvus_database = Milvus_Database("JXTest")
|
||||
self.milvus_database.using_database(self.database_name)
|
||||
result = self.milvus_database.insert(collection_name=self.collection_name,params=data)
|
||||
return result
|
||||
def save_hybrid_md(self,data):
|
||||
self.milvus_database.using_database(self.database_name)
|
||||
result = self.milvus_database.hybrid_insert(collection_name=self.collection_name,params=data,database_name=self.database_name)
|
||||
return result
|
||||
def close(self):
|
||||
self.milvus_database.close()
|
||||
717
app_pkg/markdown_process/split_markdown.py
Normal file
717
app_pkg/markdown_process/split_markdown.py
Normal file
@ -0,0 +1,717 @@
|
||||
import json
|
||||
import re
|
||||
from typing import Dict, List
|
||||
from langchain_text_splitters import CharacterTextSplitter, MarkdownHeaderTextSplitter,MarkdownTextSplitter,TokenTextSplitter,RecursiveCharacterTextSplitter
|
||||
from modelsAPI.model_api import OpenaiAPI
|
||||
from markdown_process.save_markdown import Save_Markdown_to_Milvus
|
||||
from markdown_process.convert_header import convert_numbered_headings, custom_convert_numbered_headings,parse_markdown,header_postprocess,postprocess_chunks
|
||||
from markdown_it import MarkdownIt
|
||||
from markdown_it.tree import SyntaxTreeNode
|
||||
def reset_chunk(chunks):
|
||||
id_list = []
|
||||
for res in chunks:
|
||||
if res['Header_1'] == res['text']:
|
||||
id = res['id']
|
||||
id_list.append(id)
|
||||
# print(id_list)
|
||||
for iditem in id_list:
|
||||
result = [item for item in chunks if item['id'] == iditem][0]
|
||||
result_ = [item for item in chunks if item['id'] == iditem + 1][0]
|
||||
result_['Header_2'] = result_['Header_1']
|
||||
result_['Header_1'] = result['Header_1']
|
||||
result_['text'] = result['text'] + '\n' + result_['text']
|
||||
result_devide = [record for record in chunks if not record['Header_1'] == record['text']]
|
||||
for index, item in enumerate(result_devide):
|
||||
item['id'] = index + 1
|
||||
return result_devide
|
||||
def remove_markdown_images(text):
|
||||
"""
|
||||
移除Markdown格式的图像链接
|
||||
|
||||
参数:
|
||||
text: 包含Markdown图像链接的原始字符串
|
||||
|
||||
返回:
|
||||
移除图像链接后的文本
|
||||
"""
|
||||
# 匹配  或 [](path) 格式
|
||||
pattern = r'!?\[\]\(/root/output/[^/]+/images/[a-f0-9]+\.(?:jpg|png|gif|bmp)\)'
|
||||
return re.sub(pattern, '', text)
|
||||
def remove_markdown_images(text):
|
||||
"""
|
||||
移除Markdown格式的图像链接
|
||||
|
||||
参数:
|
||||
text: 包含Markdown图像链接的原始字符串
|
||||
|
||||
返回:
|
||||
移除图像链接后的文本
|
||||
"""
|
||||
# 匹配  或 [](path) 格式
|
||||
pattern = r'http://\d{0,9}\.\d{0,9}\.\d{0,9}\.\d{0,9}:\d+/[^\s\'"<>]+'
|
||||
return re.sub(pattern, '', text)
|
||||
|
||||
def extract_urls(text):
|
||||
# 定义匹配 URL 的正则表达式模式
|
||||
pattern = r'http://\d{0,9}\.\d{0,9}\.\d{0,9}\.\d{0,9}:\d{0,9}/output/images/[a-f0-9]+\.(?:jpg|jpeg|png|gif)'
|
||||
|
||||
# 查找所有匹配的 URL
|
||||
urls = re.findall(pattern, text)
|
||||
|
||||
# # 可选:过滤只保留特定扩展名的文件(如.jpg, .png等)
|
||||
# image_urls = [url for url in urls if url.lower().endswith(('.jpg', '.jpeg', '.png', '.gif'))]
|
||||
|
||||
return urls
|
||||
|
||||
def extract_and_convert_image_paths(text):
|
||||
"""
|
||||
从字符串中提取图像路径并转换为新格式
|
||||
|
||||
参数:
|
||||
text: 包含图像路径的原始字符串
|
||||
|
||||
返回:
|
||||
包含原始路径和转换后路径的字典列表
|
||||
"""
|
||||
# 正则表达式匹配图像路径
|
||||
pattern = r'(/root/output/[^/]+/images/[a-f0-9]+\.(?:jpg|png|gif|bmp))'
|
||||
matches = re.findall(pattern, text)
|
||||
|
||||
converted_paths = []
|
||||
|
||||
for original_path in matches:
|
||||
# 分解原始路径
|
||||
path_parts = original_path.split('/')
|
||||
|
||||
# 提取关键部分
|
||||
# 格式: /root/output/<文件夹>/images/<图片名>
|
||||
folder_name = path_parts[3] # 获取"生成式人工智能传播风险:理论内涵、形成机理与治理策略_景奎"这部分
|
||||
image_name = path_parts[-1] # 获取图片文件名
|
||||
|
||||
# 构建新路径
|
||||
new_path = f"mineru/{folder_name}/images/{image_name}"
|
||||
|
||||
converted_paths.append({
|
||||
'original_path': original_path,
|
||||
'converted_path': new_path
|
||||
})
|
||||
|
||||
return converted_paths
|
||||
|
||||
class MarkdownAwareTextSplitter(RecursiveCharacterTextSplitter):
|
||||
def __init__(self, **kwargs):
|
||||
# 优先用Markdown的分隔符
|
||||
separators = [
|
||||
# 保持表格完整 (匹配Markdown表格)
|
||||
r'\n\s*\|', # 表格行
|
||||
r'\n#+\s+', # 标题
|
||||
r'\n\*\*\s+', # 分隔线
|
||||
r'\n\n', # 双换行
|
||||
r'\n', # 单换行
|
||||
r'\.\s+', # 句号+空格
|
||||
r'\s+', # 任意空白
|
||||
]
|
||||
super().__init__(separators=separators, **kwargs)
|
||||
|
||||
def split_text(self, text: str) -> List[str]:
|
||||
# 预处理:保护表格结构
|
||||
protected_tables = []
|
||||
def _protect_tables(match):
|
||||
protected_tables.append(match.group(0))
|
||||
return f"__TABLE_{len(protected_tables)-1}__"
|
||||
|
||||
# 用临时标记替换表格
|
||||
table_pattern = r'(\n\|.+\|.+\n(\|.+\|.+\n)+)'
|
||||
text_with_protected_tables = re.sub(
|
||||
table_pattern,
|
||||
_protect_tables,
|
||||
text,
|
||||
flags=re.MULTILINE
|
||||
)
|
||||
|
||||
# 常规切分
|
||||
chunks = super().split_text(text_with_protected_tables)
|
||||
|
||||
# 恢复表格
|
||||
final_chunks = []
|
||||
for chunk in chunks:
|
||||
for i, table in enumerate(protected_tables):
|
||||
chunk = chunk.replace(f"__TABLE_{i}__", table)
|
||||
final_chunks.append(chunk)
|
||||
|
||||
return final_chunks
|
||||
|
||||
def split_regular(resource,mdDocs,logger,database_name,collection_name,chunk_type,chunk_size,chunk_overlap=50,max_level=3):
|
||||
|
||||
if chunk_type =="03":
|
||||
text_splitter =TokenTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||||
texts = text_splitter.split_text(mdDocs)
|
||||
milvus = Save_Markdown_to_Milvus(database_name=database_name,collection_name=collection_name)
|
||||
for text in texts:
|
||||
convert_path_list=extract_and_convert_image_paths(text)
|
||||
text = remove_markdown_images(text)
|
||||
if convert_path_list:
|
||||
img_path = ",".join(convert_path_list)
|
||||
vector = OpenaiAPI.get_embeddings(text)
|
||||
data = [[text],[vector],[resource],[img_path]]
|
||||
milvus.save_hybrid_md(data=data)
|
||||
results.append({"id":j+i+1,"text":text,"resource":resource,"img_path":img_path})
|
||||
milvus.close()
|
||||
else:
|
||||
#自定义分解
|
||||
md_header_json = parse_markdown(mdDocs,max_level=max_level)
|
||||
# print(md_header_json)
|
||||
#合并处理
|
||||
md_header_list = header_postprocess(md_header_json)
|
||||
logger.info("markdown splitting...")
|
||||
|
||||
milvus = Save_Markdown_to_Milvus(database_name=database_name,collection_name=collection_name)
|
||||
results = []
|
||||
for i ,md in enumerate(md_header_list):
|
||||
if md["text"]:
|
||||
header_1 = md["Header_1"]
|
||||
header_2 = md["Header_2"]
|
||||
header_3 = md["Header_3"]
|
||||
text = md["text"]
|
||||
|
||||
# print("Header_1:",header_1)
|
||||
# print("Header_2:",header_2)
|
||||
# print("Header_3:",header_3)
|
||||
# print("text:",text_)
|
||||
# print("======================================")
|
||||
if chunk_type =="01":
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
length_function=len,
|
||||
separators=[" ","\n\n"]
|
||||
)
|
||||
texts = text_splitter.split_text(text)
|
||||
for j,txt in enumerate(texts):
|
||||
img_list = extract_urls(txt)
|
||||
txt = remove_markdown_images(txt)
|
||||
if img_list:
|
||||
img_path = ",".join(img_list)
|
||||
# convert_path_list=extract_and_convert_image_paths(text)
|
||||
# text = remove_markdown_images(text)
|
||||
# if convert_path_list:
|
||||
# img_path = ",".join(convert_path_list)
|
||||
# print(j)
|
||||
# print(texts)
|
||||
# print("======================================")
|
||||
vector = OpenaiAPI.get_embeddings(txt)
|
||||
data = [[header_1],[header_2],[header_3],[txt],[vector],[resource],[img_path]]
|
||||
# milvus.save_md(data=data)
|
||||
milvus.save_hybrid_md(data=data)
|
||||
results.append({"id":j+i+1,"Header_1":header_1,"Header_2":header_2,"Header_3":header_3,"text":txt,"resource":resource,"img_path":img_path})
|
||||
|
||||
else:
|
||||
img_list = extract_urls(txt)
|
||||
txt = remove_markdown_images(txt)
|
||||
if img_list:
|
||||
img_path = ",".join(img_list)
|
||||
# convert_path_list=extract_and_convert_image_paths(text)
|
||||
# text = remove_markdown_images(text)
|
||||
# if convert_path_list:
|
||||
# img_path = ",".join(convert_path_list)
|
||||
vector = OpenaiAPI.get_embeddings(text)
|
||||
data = [[header_1],[header_2],[header_3],[text],[vector],[resource],[img_path]]
|
||||
milvus.save_hybrid_md(data=data)
|
||||
results.append({"id":i+1,"Header_1":header_1,"Header_2":header_2,"Header_3":header_3,"text":txt,"resource":resource,"img_path":img_path})
|
||||
milvus.close()
|
||||
|
||||
return results
|
||||
|
||||
def split_text(resource, mdDocs, logger, chunk_type, chunk_size, chunk_overlap=50,max_level=3):
|
||||
"""
|
||||
拆分文本的函数
|
||||
|
||||
Args:
|
||||
resource: 资源标识
|
||||
mdDocs: Markdown文档内容
|
||||
logger: 日志记录器
|
||||
chunk_type: 拆分类型("01"或"03")
|
||||
"01": 智能分割(基于递归字符分割)
|
||||
"03": 按字符分割(基于token分割)
|
||||
chunk_size: 块长度
|
||||
chunk_overlap: 块重叠大小(默认为50)
|
||||
|
||||
Returns:
|
||||
list: 拆分后的文本块列表,每个元素是包含文本和元数据的字典
|
||||
"""
|
||||
#print(2222222222222222222222222222222222)
|
||||
if chunk_type == "03":
|
||||
logger.info(f"chunk size: {chunk_size},chunk_overlap: {chunk_overlap},chunk_type: {chunk_type}")
|
||||
return _split_by_token(resource, mdDocs, logger, chunk_size, chunk_overlap)
|
||||
else:
|
||||
return _split_custom(resource, mdDocs, logger, chunk_type, chunk_size, chunk_overlap,max_level)
|
||||
|
||||
def split_page_chunk(resource,content_chunk_list,chunk_size,chunk_overlap,logger,max_level,chunk_type):
|
||||
def process_content_list(parsed_data: List[Dict],max_level) -> List[Dict]:
|
||||
"""
|
||||
处理 parsed_data:
|
||||
1. 将 text_level=1/2 的标题添加到后续无 text_level 的文本前
|
||||
2. 合并 type=image 的项到前一个 chunk
|
||||
"""
|
||||
processed_data = []
|
||||
current_titles = {level: "" for level in range(1, max_level + 1)} # 各级标题缓存
|
||||
|
||||
if isinstance(max_level, str):
|
||||
try:
|
||||
max_level = int(max_level)
|
||||
except ValueError:
|
||||
max_level = 3 # 默认值
|
||||
|
||||
for item in parsed_data:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
|
||||
raw_text_level = item.get("text_level")
|
||||
text_level = None
|
||||
if raw_text_level is not None:
|
||||
if isinstance(raw_text_level, str):
|
||||
try:
|
||||
text_level = int(raw_text_level.strip()) if raw_text_level.strip() else None
|
||||
except ValueError:
|
||||
text_level = None
|
||||
elif isinstance(raw_text_level, int):
|
||||
text_level = raw_text_level
|
||||
|
||||
item_type = item.get("type", "text")
|
||||
text = item.get("text", "")
|
||||
page_idx = item.get("page_idx", -1)
|
||||
id = item.get("id")
|
||||
|
||||
if text_level is not None and 1 <= text_level <= max_level:
|
||||
# 更新当前层级及更低层级的标题(例如:遇到2级标题会清除3级的缓存)
|
||||
for level in range(text_level, max_level + 1):
|
||||
current_titles[level] = text + "\n" if level == text_level else ""
|
||||
continue # 标题本身不单独生成 chunk
|
||||
|
||||
# 合并图片逻辑
|
||||
if item_type == "image":
|
||||
image_path = item.get("image_path","")
|
||||
image_caption = item.get("image_caption","")
|
||||
if image_path:
|
||||
text = f""
|
||||
if processed_data:
|
||||
processed_data[-1]["text"] += f""
|
||||
continue
|
||||
elif item_type =="table":
|
||||
table_caption = "".join(item.get("table_caption", ""))
|
||||
table_body = item.get("table_body", "")
|
||||
text = table_caption + "\n" + table_body
|
||||
# 普通文本:添加当前所有层级的标题(从高到低)
|
||||
combined_title = "".join(
|
||||
current_titles[level]
|
||||
for level in range(1, max_level + 1)
|
||||
if current_titles[level]
|
||||
)
|
||||
full_text = combined_title + text
|
||||
|
||||
# 添加到结果(仅当有内容时)
|
||||
if full_text.strip():
|
||||
processed_data.append({
|
||||
"id":id,
|
||||
"text": full_text,
|
||||
"page_idx": page_idx+1,
|
||||
"type": item_type,
|
||||
"image_path":image_path if item_type == "image" else ""
|
||||
})
|
||||
|
||||
return processed_data
|
||||
results = []
|
||||
processed_data = process_content_list(content_chunk_list,max_level=max_level)
|
||||
for i,data in enumerate(processed_data):
|
||||
if chunk_type == "01" or "02":
|
||||
chunks = _split_recursively(data["text"], chunk_size, chunk_overlap)
|
||||
else:
|
||||
chunks = _split_by_token(resource=resource,md_text=data.get["text"],logger=logger,chunk_size=chunk_size,chunk_overlap=chunk_overlap)
|
||||
for j, chunk in enumerate(chunks):
|
||||
result = _process_chunk(chunk, resource, {}, str(data.get("id")) +f"-{j}", logger,page_idx=data.get("page_idx"))
|
||||
results.append(result)
|
||||
print(111111111111111111)
|
||||
return results
|
||||
|
||||
def _split_by_token(resource, md_text,logger, chunk_size, chunk_overlap):
|
||||
"""
|
||||
使用markdown解析器保持结构完整的切分方法,支持chunk重叠
|
||||
|
||||
参数:
|
||||
resource: 关联的资源标识
|
||||
md_text: 要切分的Markdown文本
|
||||
chunk_size: 每个chunk的目标大小
|
||||
chunk_overlap: chunk之间的重叠大小
|
||||
|
||||
返回:
|
||||
包含text和resource字段的字典列表
|
||||
"""
|
||||
logger.info("开始按字符分片")
|
||||
if chunk_overlap >= chunk_size:
|
||||
raise ValueError("chunk_overlap必须小于chunk_size")
|
||||
|
||||
md = MarkdownIt()
|
||||
tokens = md.parse(md_text)
|
||||
root = SyntaxTreeNode(tokens)
|
||||
if chunk_size >= 8192:
|
||||
raise ValueError("chunk_size不能大于4096")
|
||||
chunks = []
|
||||
current_chunk = ""
|
||||
overlap_buffer = ""
|
||||
|
||||
def get_node_text(node):
|
||||
"""递归获取节点的完整文本内容"""
|
||||
if node.type == "text":
|
||||
return node.content
|
||||
elif node.type == "inline" and node.children:
|
||||
return "".join(get_node_text(child) for child in node.children)
|
||||
elif node.children:
|
||||
# 对于块级元素,添加适当的换行
|
||||
contents = []
|
||||
for child in node.children:
|
||||
content = get_node_text(child)
|
||||
if content:
|
||||
contents.append(content)
|
||||
return "\n".join(contents)
|
||||
return ""
|
||||
|
||||
def process_node(node):
|
||||
nonlocal current_chunk, overlap_buffer
|
||||
content = get_node_text(node)
|
||||
|
||||
# 处理当前内容
|
||||
remaining_content = content
|
||||
while remaining_content:
|
||||
# 计算还能添加到当前chunk的内容量
|
||||
remaining_space = chunk_size - len(current_chunk)
|
||||
|
||||
if remaining_space <= 0:
|
||||
# 当前chunk已满,保存并准备重叠部分
|
||||
chunks.append(current_chunk)
|
||||
overlap_buffer = current_chunk[-chunk_overlap:] if chunk_overlap > 0 else ""
|
||||
current_chunk = overlap_buffer + remaining_content[:remaining_space]
|
||||
remaining_content = remaining_content[remaining_space:]
|
||||
else:
|
||||
# 添加到当前chunk
|
||||
add_content = remaining_content[:remaining_space]
|
||||
current_chunk += add_content
|
||||
remaining_content = remaining_content[remaining_space:]
|
||||
|
||||
# 处理所有节点
|
||||
for child in root.children:
|
||||
process_node(child)
|
||||
|
||||
# 添加最后一个chunk
|
||||
if current_chunk:
|
||||
chunks.append(current_chunk)
|
||||
|
||||
# 过滤空chunk并添加resource信息
|
||||
return [{"text": text.strip(), "resource": resource}
|
||||
for text in chunks if text.strip()]
|
||||
|
||||
# def _split_by_token(resource, text, logger, chunk_size, chunk_overlap):
|
||||
# """按token分割文本"""
|
||||
# logger.info("开始按字符分割")
|
||||
# text_splitter = TokenTextSplitter(chunk_size=1024, chunk_overlap=100)
|
||||
# texts = text_splitter.split_text(text)
|
||||
# return [{"text": text, "resource": resource} for text in texts]
|
||||
|
||||
|
||||
def _split_custom(resource, mdDocs, logger, chunk_type, chunk_size, chunk_overlap,max_level):
|
||||
"""自定义分割处理"""
|
||||
logger.info("开始按自定义分割")
|
||||
results = []
|
||||
# if chunk_type=="02":
|
||||
if chunk_type=="02" or chunk_type=="01":
|
||||
# md_text = custom_convert_numbered_headings(mdDocs, max_level)
|
||||
# md_text = '\n'.join(md_text)
|
||||
chunks = _split_markdown(mdDocs,max_level)
|
||||
# chunks = _split_markdown(md_text,max_level)
|
||||
for idx, doc in enumerate(chunks, start=1):
|
||||
# 确保所有Header字段都存在
|
||||
headers = {f"Header_{i}": doc.metadata.get(f"Header_{i}", "")
|
||||
for i in range(1, max_level + 1)}
|
||||
# if headers.get("Header_1") =="目 录" or headers.get("Header_1") =="目录" or headers.get("Header_1") =="目 录":
|
||||
# continue
|
||||
result = _process_chunk(doc.page_content, resource, headers, idx, logger)
|
||||
results.append(result)
|
||||
results = reset_chunk(results)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def _split_markdown(text,max_level,chunk_size=2048,chunk_overlap=50):
|
||||
"""
|
||||
根据 max_level 动态生成 headers_to_split_on 配置
|
||||
规则:
|
||||
- max_level=1 → 只按 # 分割
|
||||
- max_level=2 → 按 # 和 ## 分割
|
||||
- 以此类推
|
||||
"""
|
||||
# 动态生成 headers_to_split_on 配置
|
||||
headers_to_split_on = [
|
||||
(f"{'#' * i}", f"Header_{i}") for i in range(1, max_level + 1)
|
||||
]
|
||||
|
||||
# 初始化分割器
|
||||
splitter = MarkdownHeaderTextSplitter(
|
||||
headers_to_split_on=headers_to_split_on,
|
||||
strip_headers=False # 废弃标题文本
|
||||
)
|
||||
chunks = splitter.split_text(text)
|
||||
char_splitter = CharacterTextSplitter(
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
length_function=len,
|
||||
separator="" # 优先按段落分割
|
||||
)
|
||||
# 检查并处理过大的分块
|
||||
|
||||
final_chunks = []
|
||||
header_chunks = []
|
||||
for chunk in chunks:
|
||||
# content = chunk.page_content # 只获取文本内容,忽略metadata
|
||||
if len(chunk.page_content) <=8000:
|
||||
final_chunks.append(chunk)
|
||||
header_chunks.append(chunk.metadata) #dict
|
||||
else:
|
||||
# 直接分割文本内容
|
||||
split_texts = char_splitter.split_documents([chunk])
|
||||
final_chunks.extend(split_texts)
|
||||
|
||||
return final_chunks
|
||||
|
||||
# 执行分割
|
||||
# return splitter.split_text(text)
|
||||
|
||||
def _split_recursively(text, chunk_size, chunk_overlap):
|
||||
"""递归字符分割"""
|
||||
text_splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=1024,
|
||||
chunk_overlap=100,
|
||||
length_function=len,
|
||||
# separators=["\n\n\n","(?<!#)\n\n"," "],
|
||||
separators=[""],
|
||||
keep_separator=True,
|
||||
is_separator_regex=True
|
||||
)
|
||||
semantic_chunks = text_splitter.split_text(text)
|
||||
|
||||
# 2. 再用 CharacterTextSplitter 按字符数精确分割
|
||||
char_splitter = CharacterTextSplitter(
|
||||
chunk_size=1024,#2048
|
||||
chunk_overlap=50,
|
||||
length_function=len,
|
||||
separator="" # 按纯字符切分
|
||||
)
|
||||
|
||||
final_chunks = []
|
||||
for chunk in semantic_chunks:
|
||||
final_chunks.extend(char_splitter.split_text(chunk))
|
||||
|
||||
return final_chunks
|
||||
# return text_splitter.split_text(text)
|
||||
|
||||
|
||||
def _split_with_images(resource, text, logger, chunk_size, chunk_overlap):
|
||||
"""处理没有章节信息的文本分割"""
|
||||
results = []
|
||||
chunks = _split_recursively(text, chunk_size, chunk_overlap)
|
||||
|
||||
for i, chunk in enumerate(chunks):
|
||||
img_list = extract_urls(chunk)
|
||||
logger.info(f"img_list: {img_list}")
|
||||
# clean_text = remove_markdown_images(chunk)
|
||||
img_path = ",".join(img_list) if img_list else ""
|
||||
|
||||
results.append({
|
||||
"id": i+1,
|
||||
"text": text,#clean_text,
|
||||
"resource": resource,
|
||||
"img_path": img_path
|
||||
})
|
||||
logger.info(f"results: {results}")
|
||||
return results
|
||||
|
||||
|
||||
def _process_chunk(text, resource, headers, chunk_id, logger,page_idx=None,origin_text=None):
|
||||
"""处理单个文本块,提取图片和清理文本"""
|
||||
img_list = extract_urls(text)
|
||||
# logger.info(f"img_list: {img_list}")
|
||||
# clean_text = remove_markdown_images(text)
|
||||
img_path = ",".join(img_list) if img_list else ""
|
||||
if img_path:
|
||||
logger.info(f"Get Image Path: {img_path}")
|
||||
|
||||
return {
|
||||
"id": chunk_id,
|
||||
"Header_1": headers.get("Header_1", ""),
|
||||
"Header_2": headers.get("Header_2", ""),
|
||||
"Header_3": headers.get("Header_3", ""),
|
||||
"text":text.replace("# ",""),
|
||||
"resource": resource,
|
||||
"img_path": img_path,
|
||||
"page":page_idx,
|
||||
"origin_text":text
|
||||
}
|
||||
|
||||
def save_md(results, collection_name,database_name):
|
||||
"""
|
||||
保存Markdown到Milvus数据库的函数
|
||||
:param results: 拆分后的文本块列表
|
||||
:param database_name: 数据库名称
|
||||
:param collection_name: 集合名称
|
||||
:return: 保存后的结果列表
|
||||
"""
|
||||
texts = [item.get("text") for item in results]
|
||||
# vectors = OpenaiAPI.get_embeddings(texts)
|
||||
vectors = OpenaiAPI.batch_embeddings(texts)
|
||||
|
||||
milvus = Save_Markdown_to_Milvus(database_name=database_name, collection_name=collection_name)
|
||||
all_data = []
|
||||
for idx, item in enumerate(results):
|
||||
text = item.get("text")
|
||||
resource = item.get("resource")
|
||||
img_path = item.get("img_path", "null")
|
||||
vector = vectors[idx].get("embedding") # 获取对应的嵌入向量
|
||||
|
||||
data_item = {
|
||||
"text": text,
|
||||
"vector": vector,
|
||||
"resource": resource,
|
||||
"img_path": img_path,
|
||||
"page_idx": item.get("page_idx"),
|
||||
"origin_id": str(item.get("id")),
|
||||
"origin_text":item.get("origin_text")
|
||||
}
|
||||
|
||||
# 添加标题信息(如果有)
|
||||
if "Header_1" in item:
|
||||
data_item.update({
|
||||
"Header_1": item.get("Header_1"),
|
||||
"Header_2": item.get("Header_2"),
|
||||
"Header_3": item.get("Header_3")
|
||||
})
|
||||
all_data.append(data_item)
|
||||
|
||||
# milvus.save_hybrid_md(data=[data_item])
|
||||
milvus.save_hybrid_md(data=all_data)
|
||||
|
||||
milvus.close()
|
||||
return results
|
||||
# milvus = Save_Markdown_to_Milvus(database_name=database_name, collection_name=collection_name)
|
||||
|
||||
# for item in results:
|
||||
# text = item.get("text")
|
||||
# resource = item.get("resource")
|
||||
# img_path = item.get("img_path", "null")
|
||||
# vector = OpenaiAPI.get_embeddings(text)
|
||||
|
||||
# if "Header_1" in item: # 有标题的情况
|
||||
# data = [{"text":text,
|
||||
# "vector":vector,
|
||||
# "resource":resource,
|
||||
# "img_path":img_path,
|
||||
# "Header_1":item.get("Header_1"),
|
||||
# "Header_2":item.get("Header_2"),
|
||||
# "Header_3":item.get("Header_3"),
|
||||
# "page_idx":item.get("page_idx"),
|
||||
# "origin_id":str(item.get("id"))}]
|
||||
|
||||
# else: # 没有标题的情况
|
||||
# data = [{"text":text,
|
||||
# "vector":vector,
|
||||
# "resource":resource,
|
||||
# "img_path":img_path,
|
||||
# "page_idx":item.get("page_idx"),
|
||||
# "origin_id":str(item.get("id"))
|
||||
# }]
|
||||
|
||||
# milvus.save_hybrid_md(data=data)
|
||||
|
||||
# milvus.close()
|
||||
# return results
|
||||
|
||||
def hierarchical_save(results,collection_name,database_name):
|
||||
texts = [item.get("text") for item in results]
|
||||
# vectors = OpenaiAPI.get_embeddings(texts)
|
||||
vectors = OpenaiAPI.batch_embeddings(texts)
|
||||
|
||||
milvus = Save_Markdown_to_Milvus(database_name=database_name, collection_name=collection_name)
|
||||
all_data = []
|
||||
for idx, item in enumerate(results):
|
||||
text = item.get("text")
|
||||
resource = item.get("resource")
|
||||
img_path = item.get("img_path", "null")
|
||||
vector = vectors[idx].get("embedding") # 获取对应的嵌入向量
|
||||
|
||||
data_item = {
|
||||
"text": text,
|
||||
"vector": vector,
|
||||
"resource": resource,
|
||||
"img_path": img_path,
|
||||
"page_idx": item.get("page_idx"),
|
||||
"origin_id": str(item.get("id")),
|
||||
"origin_text": item.get("origin_text"),
|
||||
}
|
||||
|
||||
# 添加标题信息(如果有)
|
||||
if "Header_1" in item:
|
||||
data_item.update({
|
||||
"Header_1": item.get("Header_1"),
|
||||
"Header_2": item.get("Header_2"),
|
||||
"Header_3": item.get("Header_3")
|
||||
})
|
||||
all_data.append(data_item)
|
||||
|
||||
# milvus.save_hybrid_md(data=[data_item])
|
||||
milvus.hierarchical_insert(data=all_data,collection_name=collection_name,database_name=database_name)
|
||||
|
||||
milvus.close()
|
||||
return results
|
||||
|
||||
def split_insert(resource, mdDocs, logger, database_name, collection_name, chunk_type, chunk_length, chunk_overlap=50,max_level=3):
|
||||
"""
|
||||
原始函数的包装函数,保持接口不变
|
||||
"""
|
||||
# 先拆分文本
|
||||
results = split_text(resource, mdDocs, logger, chunk_type, chunk_length, chunk_overlap,max_level)
|
||||
# 然后保存到数据库
|
||||
return save_md(results, database_name, collection_name)
|
||||
|
||||
|
||||
def split_by_md_header(mdDocs,logger,database_name,collection_name):
|
||||
"""
|
||||
Split with metadata like "Metadata (("Header_1", "xxx"), ("Header_2", "xxxx")) page_content="xxx..."
|
||||
:param mdDocs: markdown text
|
||||
:return:
|
||||
"""
|
||||
mdDocs = convert_numbered_headings(mdDocs)
|
||||
headers_to_split_on = [("#", "Header_1"), ("##", "Header_2"), ("###", "Header_3")]
|
||||
md_header_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on, return_each_line=True)
|
||||
md_header_splits = md_header_splitter.split_text(mdDocs)
|
||||
# md_text_splitter = MarkdownTextSplitter(chunk_size=512,)
|
||||
#存入向量数据库
|
||||
milvus = Save_Markdown_to_Milvus(database_name=database_name,collection_name=collection_name)
|
||||
final_splits = []
|
||||
for _, i in enumerate(md_header_splits):
|
||||
logger.info(f"Index: {_}, Metadata: {i.metadata} Content: {i.page_content}")
|
||||
|
||||
if i.metadata:
|
||||
header_1 = i.metadata["Header_1"] if "Header_1" in i.metadata.keys() else None
|
||||
header_2 = i.metadata["Header_2"] if "Header_2" in i.metadata.keys() else None
|
||||
header_3 = i.metadata["Header_3"] if "Header_3" in i.metadata.keys() else None
|
||||
text = i.page_content
|
||||
# metadata = {"title": header_1, "sub_title": header_2, "sub_sub_title": header_3,"text":text}
|
||||
vector = OpenaiAPI.get_embeddings(text)
|
||||
if vector is None:
|
||||
logger.error("向量生成失败")
|
||||
raise ValueError("向量生成失败")
|
||||
|
||||
data = [[header_1],[header_2],[header_3],[text],[vector]]
|
||||
logger.info(f"Index:{_},header_1:{header_1},header_2:{header_2},header_3:{header_3},text:{text}")
|
||||
# data={"text":text,"header_1":header_1,"header_2":header_2,"header_3":header_3,"vector":vector}
|
||||
result = milvus.save_md(data=data)
|
||||
# print(f"Index:{_},header_1:{header_1},header_2:{header_2},header_3:{header_3},text:{text}")
|
||||
# print("---------------------")
|
||||
return md_header_splits
|
||||
0
app_pkg/milvus_api.py
Normal file
0
app_pkg/milvus_api.py
Normal file
0
app_pkg/modelsAPI/__init__.py
Normal file
0
app_pkg/modelsAPI/__init__.py
Normal file
525
app_pkg/modelsAPI/model_api.py
Normal file
525
app_pkg/modelsAPI/model_api.py
Normal file
@ -0,0 +1,525 @@
|
||||
import httpx
|
||||
from openai import OpenAI
|
||||
import os
|
||||
import requests
|
||||
import json
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from typing import List, Dict, Any
|
||||
from tqdm import tqdm # 进度条支持(可选)
|
||||
import re
|
||||
|
||||
class OpenaiAPI:
|
||||
# def __init__(self):
|
||||
# self.openai_embedding_api_base = os.getenv("OPENAI_API_EMBEDDING_BASE")
|
||||
# self.openai_embedding_model = os.getenv("OPENAI_EMBEDDING_MODEL")
|
||||
# self.openai_chat_api_base = os.getenv("OPENAI_API_BASE")
|
||||
# self.openai_chat_model = os.getenv("OPENAI_MODEL")
|
||||
# self.openai_api_key = os.getenv("OPENAI_API_KEY")
|
||||
# self.openai_reranker_api_base = os.getenv("OPENAI_API_RERANKER_BASE")
|
||||
# self.openai_reranker_model = os.getenv("OPENAI_RERANKER_MODEL")
|
||||
|
||||
# self.embedding_client = OpenAI(
|
||||
# api_key=self.openai_api_key,
|
||||
# base_url=self.openai_embedding_api_base,
|
||||
# )
|
||||
# self.chat_client = OpenAI(
|
||||
# api_key=self.openai_api_key,
|
||||
# base_url=self.openai_chat_api_base,
|
||||
# )
|
||||
# self.rerank_client = OpenAI(
|
||||
# api_key=self.openai_api_key,
|
||||
# base_url=self.openai_reranker_api_base,
|
||||
# )
|
||||
@staticmethod
|
||||
def get_embeddings(embedding_texts: List[str]):
|
||||
base_url = os.getenv("OPENAI_API_EMBEDDING_BASE")
|
||||
embedding_base_url = os.path.join(base_url, "embeddings")
|
||||
with httpx.Client(timeout=60) as client:
|
||||
response = client.post(
|
||||
embedding_base_url,
|
||||
json={"model": "bge-m3", "input": embedding_texts},
|
||||
headers={"Authorization": os.getenv("OPENAI_API_KEY")},
|
||||
)
|
||||
embedding_list = response.json()["data"][0]["embedding"]
|
||||
# response = OpenAI(
|
||||
# api_key=os.getenv("OPENAI_API_KEY","EMPTY"),
|
||||
# base_url=os.getenv("OPENAI_API_EMBEDDING_BASE","http://172.18.30.122:40052/v1"),
|
||||
# ).embeddings.create(
|
||||
# model=os.getenv("OPENAI_EMBEDDING_MODEL","bge-m3"),
|
||||
# input=embedding_texts,
|
||||
# )
|
||||
# embedding_list = response.data[0].embedding
|
||||
return embedding_list
|
||||
|
||||
@staticmethod
|
||||
def batch_embeddings(embedding_texts: List[str]):
|
||||
base_url = os.getenv("OPENAI_API_EMBEDDING_BASE")
|
||||
embedding_base_url = os.path.join(base_url, "embeddings")
|
||||
with httpx.Client(timeout=60) as client:
|
||||
response = client.post(
|
||||
embedding_base_url,
|
||||
json={"model": "bge-m3", "input": embedding_texts},
|
||||
headers={"Authorization": os.getenv("OPENAI_API_KEY")},
|
||||
)
|
||||
embedding_list = response.json()["data"]
|
||||
|
||||
return embedding_list
|
||||
|
||||
@staticmethod
|
||||
def rerank_query(query: str, documents: List[str], top_n: str = 3):
|
||||
"""
|
||||
发送重排序请求到 GPUStack API。
|
||||
|
||||
:param query: 查询字符串
|
||||
:param documents: 文档列表(字符串列表)
|
||||
:param top_n: 返回的文档数量
|
||||
:param model: 使用的模型名称
|
||||
:param api_key: API 密钥
|
||||
:param server_url: API 服务器地址
|
||||
:return: API 的响应结果(JSON 格式)
|
||||
|
||||
return :
|
||||
{
|
||||
"id": "rerank-1446b27f0d8642288578ee26707253a1",
|
||||
"model": "bge-reranker-v2-m3",
|
||||
"usage": {
|
||||
"total_tokens": 53
|
||||
},
|
||||
"results": [
|
||||
{
|
||||
"index": 2,
|
||||
"document": {
|
||||
"text": "I have a dog named Dollor"
|
||||
},
|
||||
"relevance_score": 0.297119140625
|
||||
},
|
||||
{
|
||||
"index": 1,
|
||||
"document": {
|
||||
"text": "hi"
|
||||
},
|
||||
"relevance_score": 0.0009255409240722656
|
||||
},
|
||||
{
|
||||
"index": 0,
|
||||
"document": {
|
||||
"text": "what is panda?"
|
||||
},
|
||||
"relevance_score": 0.0008897781372070312
|
||||
}
|
||||
]
|
||||
}
|
||||
"""
|
||||
|
||||
# 请求头
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': f'Bearer {os.getenv("OPENAI_API_KEY")}'
|
||||
}
|
||||
|
||||
# 请求体
|
||||
data = {
|
||||
"model": os.getenv("OPENAI_RERANKER_MODEL"),
|
||||
"query": query,
|
||||
"top_n": top_n,
|
||||
"documents": documents
|
||||
}
|
||||
|
||||
# 发送 POST 请求
|
||||
response = requests.post(os.getenv("OPENAI_API_RERANKER_BASE") + "/rerank", headers=headers,
|
||||
data=json.dumps(data))
|
||||
results_list = []
|
||||
|
||||
# 检查响应状态码
|
||||
if response.status_code == 200:
|
||||
result_json = response.json() # 返回解析后的 JSON 数据
|
||||
for result in result_json["results"]:
|
||||
results_list.append(
|
||||
{"index": result["index"], "text": result["document"]["text"], "scores": result["relevance_score"]})
|
||||
# results_list.append(result["document"]["text"])
|
||||
# relevance_scores.append(result["relevance_score"])
|
||||
return results_list
|
||||
|
||||
else:
|
||||
raise Exception(f"Error: {response.status_code}\n{response.text}") # 抛出异常
|
||||
|
||||
@staticmethod
|
||||
def embedding_cosine_similarity(text_1, text_2):
|
||||
# 请求头
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': f'Bearer {os.getenv("OPENAI_API_KEY")}'
|
||||
}
|
||||
|
||||
# 请求体
|
||||
data = {
|
||||
"text_1": text_1,
|
||||
"text_2": text_2,
|
||||
"model": "bge-rerank"
|
||||
}
|
||||
|
||||
# # 发送 POST 请求
|
||||
# response = requests.post(os.getenv("OPENAI_API_EMBEDDING_BASE") + "/score", headers=headers, data=json.dumps(data))
|
||||
response = requests.post(os.getenv("OPENAI_API_RERANKER_BASE") + "/score", headers=headers,
|
||||
data=json.dumps(data))
|
||||
if response.status_code == 200:
|
||||
result_json = response.json()
|
||||
score = result_json["data"][0]["score"]
|
||||
return score
|
||||
else:
|
||||
raise Exception(f"Error: {response.status_code}\n{response.text}") # 抛出异常
|
||||
|
||||
@staticmethod
|
||||
def open_api_chat(query: str, model: str):
|
||||
response = OpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
base_url=os.getenv("OPENAI_API_BASE"),
|
||||
).chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": query},
|
||||
],
|
||||
temperature=0.6,
|
||||
stream=False
|
||||
)
|
||||
return response.choices[0].message.content
|
||||
|
||||
@staticmethod
|
||||
def open_api_caht_without_thinking(query: str, model: str):
|
||||
response = OpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
base_url=os.getenv("OPENAI_API_BASE"),
|
||||
).chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": query}, ],
|
||||
|
||||
temperature=0.1,
|
||||
stream=False,
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}})
|
||||
return response.choices[0].message.content
|
||||
|
||||
@staticmethod
|
||||
def open_api_chat_stream(query: str, model: str):
|
||||
response = OpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
base_url=os.getenv("OPENAI_API_BASE"),
|
||||
).chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": query},
|
||||
],
|
||||
temperature=0.6,
|
||||
stream=True
|
||||
)
|
||||
for chunk in response:
|
||||
content = chunk.choices[0].delta.content
|
||||
if content:
|
||||
yield content.encode('utf-8')
|
||||
|
||||
@staticmethod
|
||||
def context_chunk_process(markdown_context: str, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""同步调用入口"""
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
result = loop.run_until_complete(
|
||||
batch_chunk_process(markdown_context, chunks)
|
||||
)
|
||||
finally:
|
||||
loop.close()
|
||||
return result
|
||||
|
||||
|
||||
async def batch_chunk_process(
|
||||
markdown_context: str,
|
||||
chunks: List[Dict[str, Any]],
|
||||
batch_size: int = 5,
|
||||
timeout: int = 30
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
批量处理文档块(异步并发)
|
||||
|
||||
Args:
|
||||
markdown_context: 完整文档内容
|
||||
chunks: 文档块列表,每个块需包含"text"键
|
||||
batch_size: 并发请求数 (默认5)
|
||||
timeout: 单个请求超时时间(秒)
|
||||
|
||||
Returns:
|
||||
处理后的文档块列表
|
||||
"""
|
||||
# 初始化参数
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
base_url = os.getenv("OPENAI_API_BASE")
|
||||
model = os.getenv("OPENAI_MODEL", "qwen3-32b")
|
||||
|
||||
if not api_key or not base_url:
|
||||
raise ValueError("Missing OpenAI API configuration")
|
||||
|
||||
async def process_single_chunk(
|
||||
session: aiohttp.ClientSession,
|
||||
chunk: Dict[str, Any],
|
||||
semaphore: asyncio.Semaphore
|
||||
) -> Dict[str, Any]:
|
||||
"""处理单个chunk的异步函数"""
|
||||
async with semaphore:
|
||||
text = chunk.get("text", "").strip()
|
||||
if not text:
|
||||
return chunk
|
||||
|
||||
prompt = f"""
|
||||
<document>{markdown_context}</document>
|
||||
<chunk>{text}</chunk>
|
||||
Please provide a succinct context for this chunk within the document.
|
||||
Respond ONLY with the context text, no additional commentary.
|
||||
"""
|
||||
chinese_prompt = f"""
|
||||
<文档>
|
||||
{markdown_context}
|
||||
</文档>
|
||||
以下是需要定位到整体文档中的片段
|
||||
<片段>
|
||||
{text}
|
||||
</片段>
|
||||
请简要说明该片段在整体文档中的上下文关系,以提升片段搜索效果。只需给出简洁的上下文说明,无需其他内容。
|
||||
"""
|
||||
|
||||
try:
|
||||
async with session.post(
|
||||
f"{base_url}/chat/completions",
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
json={
|
||||
"model": model,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": chinese_prompt}
|
||||
],
|
||||
"temperature": 0.1,
|
||||
"extra_body": {"chat_template_kwargs": {"enable_thinking": False}}
|
||||
},
|
||||
timeout=timeout
|
||||
) as response:
|
||||
if response.status == 200:
|
||||
data = await response.json()
|
||||
if data.get("choices"):
|
||||
chunk["text"] = data["choices"][0]["message"]["content"].strip()
|
||||
else:
|
||||
print(f"API error for chunk: {await response.text()}")
|
||||
except Exception as e:
|
||||
print(f"Error processing chunk: {e}")
|
||||
|
||||
return chunk
|
||||
|
||||
# 创建信号量控制并发数
|
||||
semaphore = asyncio.Semaphore(batch_size)
|
||||
|
||||
# 异步处理所有chunks
|
||||
async with aiohttp.ClientSession() as session:
|
||||
tasks = [
|
||||
process_single_chunk(session, chunk, semaphore)
|
||||
for chunk in chunks
|
||||
]
|
||||
|
||||
# 使用tqdm显示进度(可选)
|
||||
processed_chunks = []
|
||||
for future in tqdm(asyncio.as_completed(tasks), total=len(tasks)):
|
||||
processed_chunks.append(await future)
|
||||
|
||||
return processed_chunks
|
||||
|
||||
|
||||
def prepare_rag_prompt(query, documents):
|
||||
docs = ""
|
||||
for i, doc in enumerate(documents):
|
||||
docs += f"文本信息{i + 1}:{doc}\n"
|
||||
return f"""你作为一个文档检索问答专家,请你从以下文本信息中选取有用的内容作为参考,仔细思考后用中文回答用户的问题,如果文本中没有关于问题的答案,则回答不知道:
|
||||
-----
|
||||
{docs}
|
||||
-----
|
||||
用户提问:{query}
|
||||
你的回答:"""
|
||||
|
||||
|
||||
def query_rewrite(query):
|
||||
prompt = f"""
|
||||
你是一个专业的搜索查询优化助手。你的任务是根据用户输入的原始查询,生成更清晰、更易被搜索引擎理解的改写版本。请遵循以下规则:
|
||||
|
||||
1. **保持原意**:改写后的查询必须与用户意图一致,不能曲解原意。
|
||||
2. **简洁清晰**:去掉冗余词,优化表达,但不要过度缩写。
|
||||
3. **结构化**:如果查询适合拆解,可以适当结构化(如添加隐含的限定条件)。
|
||||
4. **自然语言**:确保改写后的查询仍然符合自然语言习惯。
|
||||
5. **多语言支持**:如果查询是中文,优先用中文改写;如果是英文,用英文优化。
|
||||
|
||||
请改写以下查询,直接返回优化后的内容,不要额外解释:
|
||||
【用户查询开始】
|
||||
{query}
|
||||
【用户查询结束】
|
||||
"""
|
||||
return OpenaiAPI.open_api_chat(prompt)
|
||||
|
||||
|
||||
def RAG_PROMPT(QUERY, CONTEXT):
|
||||
PROMPT = f"""
|
||||
### Task:
|
||||
Respond to the user query using the provided context, incorporating inline citations in the format [source_id] **only when the <source_id> tag is explicitly provided** in the context.
|
||||
|
||||
### Guidelines:
|
||||
- If you don't know the answer, clearly state that.
|
||||
- If uncertain, ask the user for clarification.
|
||||
- Respond in the same language as the user's query.
|
||||
- If the context is unreadable or of poor quality, inform the user and provide the best possible answer.
|
||||
- If the answer isn't present in the context but you possess the knowledge, explain this to the user and provide the answer using your own understanding.
|
||||
- When you find corresponding content in the provided sources, try to quote it verbatim as much as possible.
|
||||
- When quoting from sources, reproduce the original text verbatim—do not omit, paraphrase, or add content.
|
||||
- **Only include inline citations using [source_id] (e.g., [1], [2]) when a `<source_id>` tag is explicitly provided in the context.**
|
||||
- If the source file has a .glb extension, directly mark the corresponding content with the source_id in the relevant position.
|
||||
- Do not cite if the <source_id> or tag is not provided in the context.
|
||||
- Do not use XML tags in your response.
|
||||
- **If none of the provided sources contain relevant information to answer the question, respond only with: "没有找到相关来源" (No relevant sources found). Do not add any other content.
|
||||
- Ensure citations are concise and directly related to the information provided.
|
||||
|
||||
### Example of Citation:
|
||||
If the user asks about a specific topic and the information is found in "whitepaper.pdf" with a provided <source_id> , the response should include the citation like so:
|
||||
* "According to the study, the proposed method increases efficiency by 20% [whitepaper.pdf]."
|
||||
If no <source_id> is present, the response should omit the citation.
|
||||
|
||||
### Output:
|
||||
Provide a clear and direct response to the user's query, including inline citations in the format [source_id] only when the <source_id> tag is present in the context.
|
||||
|
||||
<context>
|
||||
{CONTEXT}
|
||||
</context>
|
||||
|
||||
<user_query>
|
||||
{QUERY}
|
||||
</user_query>
|
||||
"""
|
||||
return PROMPT
|
||||
|
||||
# for chunk in response:
|
||||
# content = chunk.choices[0].delta.content
|
||||
# if content:
|
||||
# print(content, end='', flush=True)
|
||||
# print('\n')
|
||||
# return response
|
||||
def split_by_period(text):
|
||||
sentences = re.split('(?<=。)', text)
|
||||
sentences = [sentence.strip() for sentence in sentences if sentence.strip()]
|
||||
return sentences
|
||||
|
||||
|
||||
def rerank_text(query,data):
|
||||
for datainfo in data:
|
||||
splitdata = split_by_period(datainfo['text'])
|
||||
current_score = datainfo['score']
|
||||
text_1 = query
|
||||
url = os.getenv("OPENAI_API_RERANKER_BASE") + "/score"
|
||||
headers = {
|
||||
'Authorization': 'Bearer 123456',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
parse_data = {
|
||||
"model": "bge-rerank",
|
||||
"text_1": text_1,
|
||||
"text_2": splitdata
|
||||
}
|
||||
response = requests.post(url, headers=headers, json=parse_data)
|
||||
parseresult = response.json()
|
||||
scores = [item['score'] for item in parseresult['data'] if 'score' in item]
|
||||
|
||||
max_score = max(scores)
|
||||
socore = 0.1 * current_score + 0.9 * max_score
|
||||
datainfo['score'] = socore
|
||||
return data
|
||||
|
||||
def rerank_3d(query,data):
|
||||
for recall_result in data:
|
||||
text_2 = recall_result["text"]
|
||||
current_score = recall_result["score"]
|
||||
text_1 = query
|
||||
url = os.getenv("OPENAI_API_RERANKER_BASE") + "/score"
|
||||
headers = {
|
||||
'Authorization': 'Bearer 123456',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
parse_data = {
|
||||
"model": "bge-rerank",
|
||||
"text_1": text_1,
|
||||
"text_2": text_2
|
||||
}
|
||||
response = requests.post(url, headers=headers, json=parse_data)
|
||||
parseresult = response.json()
|
||||
rerank_score = parseresult['data'][0]['score']
|
||||
score = current_score * 0.2 + rerank_score * 0.8
|
||||
recall_result["score"] = score
|
||||
result_dict = {}
|
||||
for item in data:
|
||||
resource = item['resource']
|
||||
# 如果该 resource 不存在,或当前 score 更高,则替换
|
||||
if resource not in result_dict or item['score'] > result_dict[resource]['score']:
|
||||
result_dict[resource] = item
|
||||
|
||||
# 转换为列表
|
||||
filtered_result = list(result_dict.values())
|
||||
for new_index, item in enumerate(filtered_result, start=1):
|
||||
item['index'] = new_index
|
||||
|
||||
return filtered_result
|
||||
|
||||
|
||||
class OPENAIAPI:
|
||||
@staticmethod
|
||||
def openai_chat(url,query, model):
|
||||
"""调用OpenAI API"""
|
||||
try:
|
||||
client = OpenAI(
|
||||
api_key="none",
|
||||
base_url=url,
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "system",
|
||||
"content": "你是一个专业的武器类查询改写助手,将用户问题改写为一个容易理解、且明确的查询。"},
|
||||
{"role": "user", "content": query},
|
||||
],
|
||||
temperature=0.1,
|
||||
stream=False,
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}}
|
||||
)
|
||||
return response.choices[0].message.content
|
||||
except Exception as e:
|
||||
print(f"调用OpenAI API时出错: {e}")
|
||||
return None
|
||||
prompt = """
|
||||
你是一个查询重写工具。请将以下用户查询改写为一个清晰、具体、完整、无歧义的自然语言问题。
|
||||
要求:
|
||||
- 保留原始意图,不添加未提及的假设;
|
||||
- 若查询模糊(如缩写、代词、片段),基于常识明确实体和领域(例如“AR50”应理解为“AR-50狙击步枪”);
|
||||
- 使用完整陈述句,聚焦信息需求(如定义等);
|
||||
- **仅输出改写后的问题本身,不要任何其他文字、标点前缀、问候、反问或说明。**
|
||||
- 不要进行任何思考,直接输出改写后查询
|
||||
- 请按照下述示例进行改写
|
||||
原始查询:ar50
|
||||
改写后查询:AR-50狙击步枪是什么?
|
||||
|
||||
原始查询:变速器?
|
||||
改写后查询:变速器是什么?
|
||||
|
||||
原始查询:{user_query}
|
||||
"""
|
||||
def resetQuery(query:str):
|
||||
url = os.getenv("OPENAI_API_MODEL_BASE"+"/v1","http://192.168.0.46:59800/v1")
|
||||
model = os.getenv("MODEL_NAME","Qwen3-32B")
|
||||
query_prompt = prompt.replace("{user_query}",query)
|
||||
answer = OPENAIAPI.openai_chat(url,query_prompt, model)
|
||||
return answer
|
||||
|
||||
475
app_pkg/modelsAPI/model_api.py11
Normal file
475
app_pkg/modelsAPI/model_api.py11
Normal file
@ -0,0 +1,475 @@
|
||||
import httpx
|
||||
from openai import OpenAI
|
||||
import os
|
||||
import requests
|
||||
import json
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from typing import List, Dict, Any
|
||||
from tqdm import tqdm # 进度条支持(可选)
|
||||
import re
|
||||
|
||||
class OpenaiAPI:
|
||||
# def __init__(self):
|
||||
# self.openai_embedding_api_base = os.getenv("OPENAI_API_EMBEDDING_BASE")
|
||||
# self.openai_embedding_model = os.getenv("OPENAI_EMBEDDING_MODEL")
|
||||
# self.openai_chat_api_base = os.getenv("OPENAI_API_BASE")
|
||||
# self.openai_chat_model = os.getenv("OPENAI_MODEL")
|
||||
# self.openai_api_key = os.getenv("OPENAI_API_KEY")
|
||||
# self.openai_reranker_api_base = os.getenv("OPENAI_API_RERANKER_BASE")
|
||||
# self.openai_reranker_model = os.getenv("OPENAI_RERANKER_MODEL")
|
||||
|
||||
# self.embedding_client = OpenAI(
|
||||
# api_key=self.openai_api_key,
|
||||
# base_url=self.openai_embedding_api_base,
|
||||
# )
|
||||
# self.chat_client = OpenAI(
|
||||
# api_key=self.openai_api_key,
|
||||
# base_url=self.openai_chat_api_base,
|
||||
# )
|
||||
# self.rerank_client = OpenAI(
|
||||
# api_key=self.openai_api_key,
|
||||
# base_url=self.openai_reranker_api_base,
|
||||
# )
|
||||
@staticmethod
|
||||
def get_embeddings(embedding_texts: List[str]):
|
||||
base_url = os.getenv("OPENAI_API_EMBEDDING_BASE")
|
||||
embedding_base_url = os.path.join(base_url, "embeddings")
|
||||
with httpx.Client(timeout=60) as client:
|
||||
response = client.post(
|
||||
embedding_base_url,
|
||||
json={"model": "bge-m3", "input": embedding_texts},
|
||||
headers={"Authorization": os.getenv("OPENAI_API_KEY")},
|
||||
)
|
||||
embedding_list = response.json()["data"][0]["embedding"]
|
||||
# response = OpenAI(
|
||||
# api_key=os.getenv("OPENAI_API_KEY","EMPTY"),
|
||||
# base_url=os.getenv("OPENAI_API_EMBEDDING_BASE","http://172.18.30.122:40052/v1"),
|
||||
# ).embeddings.create(
|
||||
# model=os.getenv("OPENAI_EMBEDDING_MODEL","bge-m3"),
|
||||
# input=embedding_texts,
|
||||
# )
|
||||
# embedding_list = response.data[0].embedding
|
||||
return embedding_list
|
||||
|
||||
@staticmethod
|
||||
def batch_embeddings(embedding_texts: List[str]):
|
||||
base_url = os.getenv("OPENAI_API_EMBEDDING_BASE")
|
||||
embedding_base_url = os.path.join(base_url, "embeddings")
|
||||
with httpx.Client(timeout=60) as client:
|
||||
response = client.post(
|
||||
embedding_base_url,
|
||||
json={"model": "bge-m3", "input": embedding_texts},
|
||||
headers={"Authorization": os.getenv("OPENAI_API_KEY")},
|
||||
)
|
||||
embedding_list = response.json()["data"]
|
||||
|
||||
return embedding_list
|
||||
|
||||
@staticmethod
|
||||
def rerank_query(query: str, documents: List[str], top_n: str = 3):
|
||||
"""
|
||||
发送重排序请求到 GPUStack API。
|
||||
|
||||
:param query: 查询字符串
|
||||
:param documents: 文档列表(字符串列表)
|
||||
:param top_n: 返回的文档数量
|
||||
:param model: 使用的模型名称
|
||||
:param api_key: API 密钥
|
||||
:param server_url: API 服务器地址
|
||||
:return: API 的响应结果(JSON 格式)
|
||||
|
||||
return :
|
||||
{
|
||||
"id": "rerank-1446b27f0d8642288578ee26707253a1",
|
||||
"model": "bge-reranker-v2-m3",
|
||||
"usage": {
|
||||
"total_tokens": 53
|
||||
},
|
||||
"results": [
|
||||
{
|
||||
"index": 2,
|
||||
"document": {
|
||||
"text": "I have a dog named Dollor"
|
||||
},
|
||||
"relevance_score": 0.297119140625
|
||||
},
|
||||
{
|
||||
"index": 1,
|
||||
"document": {
|
||||
"text": "hi"
|
||||
},
|
||||
"relevance_score": 0.0009255409240722656
|
||||
},
|
||||
{
|
||||
"index": 0,
|
||||
"document": {
|
||||
"text": "what is panda?"
|
||||
},
|
||||
"relevance_score": 0.0008897781372070312
|
||||
}
|
||||
]
|
||||
}
|
||||
"""
|
||||
|
||||
# 请求头
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': f'Bearer {os.getenv("OPENAI_API_KEY")}'
|
||||
}
|
||||
|
||||
# 请求体
|
||||
data = {
|
||||
"model": os.getenv("OPENAI_RERANKER_MODEL"),
|
||||
"query": query,
|
||||
"top_n": top_n,
|
||||
"documents": documents
|
||||
}
|
||||
|
||||
# 发送 POST 请求
|
||||
response = requests.post(os.getenv("OPENAI_API_RERANKER_BASE") + "/rerank", headers=headers,
|
||||
data=json.dumps(data))
|
||||
results_list = []
|
||||
|
||||
# 检查响应状态码
|
||||
if response.status_code == 200:
|
||||
result_json = response.json() # 返回解析后的 JSON 数据
|
||||
for result in result_json["results"]:
|
||||
results_list.append(
|
||||
{"index": result["index"], "text": result["document"]["text"], "scores": result["relevance_score"]})
|
||||
# results_list.append(result["document"]["text"])
|
||||
# relevance_scores.append(result["relevance_score"])
|
||||
return results_list
|
||||
|
||||
else:
|
||||
raise Exception(f"Error: {response.status_code}\n{response.text}") # 抛出异常
|
||||
|
||||
@staticmethod
|
||||
def embedding_cosine_similarity(text_1, text_2):
|
||||
# 请求头
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
'Authorization': f'Bearer {os.getenv("OPENAI_API_KEY")}'
|
||||
}
|
||||
|
||||
# 请求体
|
||||
data = {
|
||||
"text_1": text_1,
|
||||
"text_2": text_2,
|
||||
"model": "bge-rerank"
|
||||
}
|
||||
|
||||
# # 发送 POST 请求
|
||||
# response = requests.post(os.getenv("OPENAI_API_EMBEDDING_BASE") + "/score", headers=headers, data=json.dumps(data))
|
||||
response = requests.post(os.getenv("OPENAI_API_RERANKER_BASE") + "/score", headers=headers,
|
||||
data=json.dumps(data))
|
||||
if response.status_code == 200:
|
||||
result_json = response.json()
|
||||
score = result_json["data"][0]["score"]
|
||||
return score
|
||||
else:
|
||||
raise Exception(f"Error: {response.status_code}\n{response.text}") # 抛出异常
|
||||
|
||||
@staticmethod
|
||||
def open_api_chat(query: str, model: str):
|
||||
response = OpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
base_url=os.getenv("OPENAI_API_BASE"),
|
||||
).chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": query},
|
||||
],
|
||||
temperature=0.6,
|
||||
stream=False
|
||||
)
|
||||
return response.choices[0].message.content
|
||||
|
||||
@staticmethod
|
||||
def open_api_caht_without_thinking(query: str, model: str):
|
||||
response = OpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
base_url=os.getenv("OPENAI_API_BASE"),
|
||||
).chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": query}, ],
|
||||
|
||||
temperature=0.1,
|
||||
stream=False,
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}})
|
||||
return response.choices[0].message.content
|
||||
|
||||
@staticmethod
|
||||
def open_api_chat_stream(query: str, model: str):
|
||||
response = OpenAI(
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
base_url=os.getenv("OPENAI_API_BASE"),
|
||||
).chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": query},
|
||||
],
|
||||
temperature=0.6,
|
||||
stream=True
|
||||
)
|
||||
for chunk in response:
|
||||
content = chunk.choices[0].delta.content
|
||||
if content:
|
||||
yield content.encode('utf-8')
|
||||
|
||||
@staticmethod
|
||||
def context_chunk_process(markdown_context: str, chunks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""同步调用入口"""
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
result = loop.run_until_complete(
|
||||
batch_chunk_process(markdown_context, chunks)
|
||||
)
|
||||
finally:
|
||||
loop.close()
|
||||
return result
|
||||
|
||||
|
||||
async def batch_chunk_process(
|
||||
markdown_context: str,
|
||||
chunks: List[Dict[str, Any]],
|
||||
batch_size: int = 5,
|
||||
timeout: int = 30
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
批量处理文档块(异步并发)
|
||||
|
||||
Args:
|
||||
markdown_context: 完整文档内容
|
||||
chunks: 文档块列表,每个块需包含"text"键
|
||||
batch_size: 并发请求数 (默认5)
|
||||
timeout: 单个请求超时时间(秒)
|
||||
|
||||
Returns:
|
||||
处理后的文档块列表
|
||||
"""
|
||||
# 初始化参数
|
||||
api_key = os.getenv("OPENAI_API_KEY")
|
||||
base_url = os.getenv("OPENAI_API_BASE")
|
||||
model = os.getenv("OPENAI_MODEL", "qwen3-32b")
|
||||
|
||||
if not api_key or not base_url:
|
||||
raise ValueError("Missing OpenAI API configuration")
|
||||
|
||||
async def process_single_chunk(
|
||||
session: aiohttp.ClientSession,
|
||||
chunk: Dict[str, Any],
|
||||
semaphore: asyncio.Semaphore
|
||||
) -> Dict[str, Any]:
|
||||
"""处理单个chunk的异步函数"""
|
||||
async with semaphore:
|
||||
text = chunk.get("text", "").strip()
|
||||
if not text:
|
||||
return chunk
|
||||
|
||||
prompt = f"""
|
||||
<document>{markdown_context}</document>
|
||||
<chunk>{text}</chunk>
|
||||
Please provide a succinct context for this chunk within the document.
|
||||
Respond ONLY with the context text, no additional commentary.
|
||||
"""
|
||||
chinese_prompt = f"""
|
||||
<文档>
|
||||
{markdown_context}
|
||||
</文档>
|
||||
以下是需要定位到整体文档中的片段
|
||||
<片段>
|
||||
{text}
|
||||
</片段>
|
||||
请简要说明该片段在整体文档中的上下文关系,以提升片段搜索效果。只需给出简洁的上下文说明,无需其他内容。
|
||||
"""
|
||||
|
||||
try:
|
||||
async with session.post(
|
||||
f"{base_url}/chat/completions",
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
json={
|
||||
"model": model,
|
||||
"messages": [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": chinese_prompt}
|
||||
],
|
||||
"temperature": 0.1,
|
||||
"extra_body": {"chat_template_kwargs": {"enable_thinking": False}}
|
||||
},
|
||||
timeout=timeout
|
||||
) as response:
|
||||
if response.status == 200:
|
||||
data = await response.json()
|
||||
if data.get("choices"):
|
||||
chunk["text"] = data["choices"][0]["message"]["content"].strip()
|
||||
else:
|
||||
print(f"API error for chunk: {await response.text()}")
|
||||
except Exception as e:
|
||||
print(f"Error processing chunk: {e}")
|
||||
|
||||
return chunk
|
||||
|
||||
# 创建信号量控制并发数
|
||||
semaphore = asyncio.Semaphore(batch_size)
|
||||
|
||||
# 异步处理所有chunks
|
||||
async with aiohttp.ClientSession() as session:
|
||||
tasks = [
|
||||
process_single_chunk(session, chunk, semaphore)
|
||||
for chunk in chunks
|
||||
]
|
||||
|
||||
# 使用tqdm显示进度(可选)
|
||||
processed_chunks = []
|
||||
for future in tqdm(asyncio.as_completed(tasks), total=len(tasks)):
|
||||
processed_chunks.append(await future)
|
||||
|
||||
return processed_chunks
|
||||
|
||||
|
||||
def prepare_rag_prompt(query, documents):
|
||||
docs = ""
|
||||
for i, doc in enumerate(documents):
|
||||
docs += f"文本信息{i + 1}:{doc}\n"
|
||||
return f"""你作为一个文档检索问答专家,请你从以下文本信息中选取有用的内容作为参考,仔细思考后用中文回答用户的问题,如果文本中没有关于问题的答案,则回答不知道:
|
||||
-----
|
||||
{docs}
|
||||
-----
|
||||
用户提问:{query}
|
||||
你的回答:"""
|
||||
|
||||
|
||||
def query_rewrite(query):
|
||||
prompt = f"""
|
||||
你是一个专业的搜索查询优化助手。你的任务是根据用户输入的原始查询,生成更清晰、更易被搜索引擎理解的改写版本。请遵循以下规则:
|
||||
|
||||
1. **保持原意**:改写后的查询必须与用户意图一致,不能曲解原意。
|
||||
2. **简洁清晰**:去掉冗余词,优化表达,但不要过度缩写。
|
||||
3. **结构化**:如果查询适合拆解,可以适当结构化(如添加隐含的限定条件)。
|
||||
4. **自然语言**:确保改写后的查询仍然符合自然语言习惯。
|
||||
5. **多语言支持**:如果查询是中文,优先用中文改写;如果是英文,用英文优化。
|
||||
|
||||
请改写以下查询,直接返回优化后的内容,不要额外解释:
|
||||
【用户查询开始】
|
||||
{query}
|
||||
【用户查询结束】
|
||||
"""
|
||||
return OpenaiAPI.open_api_chat(prompt)
|
||||
|
||||
|
||||
def RAG_PROMPT(QUERY, CONTEXT):
|
||||
PROMPT = f"""
|
||||
### Task:
|
||||
Respond to the user query using the provided context, incorporating inline citations in the format [source_id] **only when the <source_id> tag is explicitly provided** in the context.
|
||||
|
||||
### Guidelines:
|
||||
- If you don't know the answer, clearly state that.
|
||||
- If uncertain, ask the user for clarification.
|
||||
- Respond in the same language as the user's query.
|
||||
- If the context is unreadable or of poor quality, inform the user and provide the best possible answer.
|
||||
- If the answer isn't present in the context but you possess the knowledge, explain this to the user and provide the answer using your own understanding.
|
||||
- When you find corresponding content in the provided sources, try to quote it verbatim as much as possible.
|
||||
- When quoting from sources, reproduce the original text verbatim—do not omit, paraphrase, or add content.
|
||||
- **Only include inline citations using [source_id] (e.g., [1], [2]) when a `<source_id>` tag is explicitly provided in the context.**
|
||||
- If the source file has a .glb extension, directly mark the corresponding content with the source_id in the relevant position.
|
||||
- Do not cite if the <source_id> or tag is not provided in the context.
|
||||
- Do not use XML tags in your response.
|
||||
- **If none of the provided sources contain relevant information to answer the question, respond only with: "没有找到相关来源" (No relevant sources found). Do not add any other content.
|
||||
- Ensure citations are concise and directly related to the information provided.
|
||||
|
||||
### Example of Citation:
|
||||
If the user asks about a specific topic and the information is found in "whitepaper.pdf" with a provided <source_id> , the response should include the citation like so:
|
||||
* "According to the study, the proposed method increases efficiency by 20% [whitepaper.pdf]."
|
||||
If no <source_id> is present, the response should omit the citation.
|
||||
|
||||
### Output:
|
||||
Provide a clear and direct response to the user's query, including inline citations in the format [source_id] only when the <source_id> tag is present in the context.
|
||||
|
||||
<context>
|
||||
{CONTEXT}
|
||||
</context>
|
||||
|
||||
<user_query>
|
||||
{QUERY}
|
||||
</user_query>
|
||||
"""
|
||||
return PROMPT
|
||||
|
||||
# for chunk in response:
|
||||
# content = chunk.choices[0].delta.content
|
||||
# if content:
|
||||
# print(content, end='', flush=True)
|
||||
# print('\n')
|
||||
# return response
|
||||
def split_by_period(text):
|
||||
sentences = re.split('(?<=。)', text)
|
||||
sentences = [sentence.strip() for sentence in sentences if sentence.strip()]
|
||||
return sentences
|
||||
|
||||
|
||||
def rerank_text(query,data):
|
||||
for datainfo in data:
|
||||
splitdata = split_by_period(datainfo['text'])
|
||||
current_score = datainfo['score']
|
||||
text_1 = query
|
||||
url = os.getenv("OPENAI_API_RERANKER_BASE") + "/score"
|
||||
headers = {
|
||||
'Authorization': 'Bearer 123456',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
parse_data = {
|
||||
"model": "bge-rerank",
|
||||
"text_1": text_1,
|
||||
"text_2": splitdata
|
||||
}
|
||||
response = requests.post(url, headers=headers, json=parse_data)
|
||||
parseresult = response.json()
|
||||
scores = [item['score'] for item in parseresult['data'] if 'score' in item]
|
||||
|
||||
max_score = max(scores)
|
||||
socore = 0.1 * current_score + 0.9 * max_score
|
||||
datainfo['score'] = socore
|
||||
return data
|
||||
|
||||
def rerank_3d(query,data):
|
||||
for recall_result in data:
|
||||
text_2 = recall_result["text"]
|
||||
current_score = recall_result["score"]
|
||||
text_1 = query
|
||||
url = os.getenv("OPENAI_API_RERANKER_BASE") + "/score"
|
||||
headers = {
|
||||
'Authorization': 'Bearer 123456',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
parse_data = {
|
||||
"model": "bge-rerank",
|
||||
"text_1": text_1,
|
||||
"text_2": text_2
|
||||
}
|
||||
response = requests.post(url, headers=headers, json=parse_data)
|
||||
parseresult = response.json()
|
||||
rerank_score = parseresult['data'][0]['score']
|
||||
score = current_score * 0.2 + rerank_score * 0.8
|
||||
recall_result["score"] = score
|
||||
result_dict = {}
|
||||
for item in data:
|
||||
resource = item['resource']
|
||||
# 如果该 resource 不存在,或当前 score 更高,则替换
|
||||
if resource not in result_dict or item['score'] > result_dict[resource]['score']:
|
||||
result_dict[resource] = item
|
||||
|
||||
# 转换为列表
|
||||
filtered_result = list(result_dict.values())
|
||||
for new_index, item in enumerate(filtered_result, start=1):
|
||||
item['index'] = new_index
|
||||
|
||||
return filtered_result
|
||||
|
||||
BIN
app_pkg/pdf/test.pdf
Normal file
BIN
app_pkg/pdf/test.pdf
Normal file
Binary file not shown.
17
app_pkg/requirements.txt
Normal file
17
app_pkg/requirements.txt
Normal file
@ -0,0 +1,17 @@
|
||||
openai
|
||||
fastapi
|
||||
aiohttp
|
||||
uvicorn
|
||||
python-multipart
|
||||
langchain
|
||||
httpx
|
||||
pymilvus
|
||||
pypinyin
|
||||
python-dotenv
|
||||
tiktoken
|
||||
pdfplumber
|
||||
python-docx
|
||||
markitdown[all]
|
||||
markdown-it-py
|
||||
fuzzywuzzy
|
||||
python-Levenshtein
|
||||
19
app_pkg/test.py
Normal file
19
app_pkg/test.py
Normal file
@ -0,0 +1,19 @@
|
||||
import subprocess
|
||||
import os
|
||||
|
||||
def excel_to_pdf(excel_path, output_dir):
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
subprocess.run(
|
||||
[
|
||||
"soffice",
|
||||
"--headless",
|
||||
"--convert-to", "pdf",
|
||||
excel_path,
|
||||
"--outdir", output_dir
|
||||
],
|
||||
check=True
|
||||
)
|
||||
|
||||
excel_to_pdf("test.xlsx", "./pdf")
|
||||
|
||||
BIN
app_pkg/test.xlsx
Normal file
BIN
app_pkg/test.xlsx
Normal file
Binary file not shown.
465
batch_kg_build.py
Normal file
465
batch_kg_build.py
Normal file
@ -0,0 +1,465 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
批量知识图谱构建脚本 —— 三阶段 ETL
|
||||
|
||||
阶段1 (extract): 对文件夹内所有文件做实体抽取,每个文件结果存为 JSON
|
||||
阶段2 (merge): 汇总所有 JSON,去重,生成 embedding,写入 Neo4j
|
||||
|
||||
用法:
|
||||
# 只跑抽取(4个容器各跑各的文件夹)
|
||||
python batch_kg_build.py --folder /app/files --output_dir /app/kg_output --phase extract --workers 4
|
||||
|
||||
# 只跑入库(所有容器跑完后,跑一次)
|
||||
python batch_kg_build.py --output_dir /app/kg_output --phase merge
|
||||
|
||||
#抽取好的实体关系直接进行存储
|
||||
python batch_kg_build.py --phase load_snapshot --snapshot_path /path/to/my_custom_snapshot.json
|
||||
|
||||
# 全流程一次跑完
|
||||
python batch_kg_build.py --folder /app/files --output_dir /app/kg_output --phase all --workers 4
|
||||
"""
|
||||
|
||||
from default_ontology_config import (
|
||||
ONTOLOGY as _BASE_ONTOLOGY,
|
||||
ONTOLOGY_DESCRIPTION as _BASE_ONTOLOGY_DESC,
|
||||
NODE_TYPES as _BASE_NODE_TYPES,
|
||||
RELATIONSHIPS as RELATIONSHIPS,
|
||||
RELATIONSHIP_TYPES as RELATIONSHIP_TYPES,
|
||||
)
|
||||
import entity_registry as _entity_registry
|
||||
import importlib
|
||||
from modelsAPI.model_api import OpenaiAPI
|
||||
from kg_build.kg_save import MilitaryKnowledgeGraph
|
||||
from kg_build.filter_node_relation import filter_node_relationship
|
||||
from filename_proceess_and_kgquery import get_entity
|
||||
from kg_build.extract_filtertext import get_filtertext_node_relation
|
||||
from extract_text_node_relation import process_file_content, check_filename_ocr
|
||||
import os
|
||||
import json
|
||||
import logging
|
||||
import argparse
|
||||
import asyncio
|
||||
import traceback
|
||||
import time
|
||||
from datetime import datetime
|
||||
from multiprocessing import Process, Queue
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import List, Dict, Tuple, Any, Optional
|
||||
|
||||
EXTRACT_TIMEOUT = 600 # 单文件抽取超时秒数
|
||||
EMBED_WORKERS = 8 # Embedding 并发线程数
|
||||
FAULT_DIR = "/app/fault_wenjian"
|
||||
|
||||
|
||||
def _record_fault(file_name: str):
|
||||
"""将超时文件名追加到 /app/fault_wenjian/timeout_YYYYMMDD.txt"""
|
||||
os.makedirs(FAULT_DIR, exist_ok=True)
|
||||
date_str = datetime.now().strftime("%Y%m%d")
|
||||
fault_path = os.path.join(FAULT_DIR, f"timeout_{date_str}.txt")
|
||||
with open(fault_path, "a", encoding="utf-8") as f:
|
||||
f.write(file_name + "\n")
|
||||
logger.warning(f"[FAULT] 已记录超时文件: {fault_path}")
|
||||
|
||||
|
||||
# ── 项目内模块 ──────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _build_ontology_config():
|
||||
"""
|
||||
与 app.py build_kg_config 保持一致:
|
||||
default_ontology_config 的静态配置 + entity_registry.json 动态注册的实体类型合并。
|
||||
"""
|
||||
importlib.reload(_entity_registry)
|
||||
reg_list = _entity_registry.get_ontology_list()
|
||||
reg_desc = _entity_registry.get_ontology_description()
|
||||
reg_nodes = _entity_registry.get_node_types()
|
||||
|
||||
# 合并:静态配置优先(手写配置不被覆盖)
|
||||
ontology = list(dict.fromkeys(_BASE_ONTOLOGY + reg_list)) # 去重保序
|
||||
ontology_desc = {**reg_desc, **_BASE_ONTOLOGY_DESC}
|
||||
node_types = {**reg_nodes, **_BASE_NODE_TYPES}
|
||||
return ontology, ontology_desc, node_types
|
||||
|
||||
|
||||
# 模块加载时构建一次(抽取阶段开始前已有 registry 内容)
|
||||
ONTOLOGY, ONTOLOGY_DESC, NODE_TYPES = _build_ontology_config()
|
||||
|
||||
# ── 配置(与 app.py 保持一致,按需修改) ────────────────────────────────────
|
||||
NEO4J_URI = os.getenv("NEO4J_URI", "bolt://localhost:7687")
|
||||
NEO4J_USER = os.getenv("NEO4J_USER", "neo4j")
|
||||
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD", "zdht123@")
|
||||
MINERU_INSTANCES = [
|
||||
# ("http://192.168.0.64:18000/analyze-pdf"),
|
||||
#("http://192.168.0.111:9977/analyze-pdf/"),
|
||||
# ("http://192.168.0.111:9979/analyze-pdf"),
|
||||
("http://192.168.1.64:18000/analyze-pdf/"),
|
||||
]
|
||||
PREFIX_URL = os.getenv("PREFIX_URL", "http://192.168.1.64:9085")
|
||||
|
||||
SUPPORTED_EXT = {".pdf", ".md", ".docx", ".xlsx"}
|
||||
|
||||
# ── 日志 ────────────────────────────────────────────────────────────────────
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||||
handlers=[
|
||||
logging.StreamHandler(),
|
||||
logging.FileHandler("batch_kg_build.log", encoding="utf-8"),
|
||||
],
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ════════════════════════════════════════════════════════════════════════════
|
||||
# 阶段 1:单文件抽取
|
||||
# ════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
|
||||
def extract_single_file(file_path: str, output_dir: str, api_url: str = MINERU_INSTANCES[0]) -> Dict:
|
||||
"""
|
||||
对单个文件做实体/关系抽取,结果保存到 output_dir/<filename>.json。
|
||||
已存在 JSON 则跳过(断点续跑)。
|
||||
返回 {"file": ..., "status": "ok"/"skip"/"error", "entity_count": ..., "rel_count": ...}
|
||||
"""
|
||||
logger.info(f'[api_url] {api_url}')
|
||||
file_name = os.path.basename(file_path)
|
||||
ext = os.path.splitext(file_name)[1].lower()
|
||||
out_path = os.path.join(output_dir, file_name + ".json")
|
||||
|
||||
# ── 断点续跑:已有结果直接跳过 ──────────────────────────────────────────
|
||||
if os.path.exists(out_path):
|
||||
logger.info(f"[SKIP] 已有结果,跳过: {file_name}")
|
||||
return {"file": file_name, "status": "skip"}
|
||||
|
||||
if ext not in SUPPORTED_EXT:
|
||||
logger.warning(f"[SKIP] 不支持的文件类型: {file_name}")
|
||||
return {"file": file_name, "status": "skip"}
|
||||
|
||||
t0 = time.time()
|
||||
entities: List[Dict] = []
|
||||
relations: List[Dict] = []
|
||||
lower_entity_device: Optional[str] = None
|
||||
lower_entity_type: Optional[str] = None
|
||||
|
||||
try:
|
||||
# ── 1. 文件名实体提取 ────────────────────────────────────────────────
|
||||
try:
|
||||
fn_result = get_entity(file_name)
|
||||
fn_entities = fn_result.get("entities", [])
|
||||
fn_relations = fn_result.get("relationships", [])
|
||||
lower_entity = fn_result.get("low_level")
|
||||
for e in fn_entities:
|
||||
if e.get("properties", {}).get("名称") == lower_entity:
|
||||
lower_entity_device = e["properties"].get("名称")
|
||||
lower_entity_type = e.get("type")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.warning(f"[{file_name}] 文件名实体提取失败: {e}")
|
||||
fn_entities, fn_relations = [], []
|
||||
|
||||
# ── 2. 正文实体抽取 ──────────────────────────────────────────────────
|
||||
check = check_filename_ocr(file_name)
|
||||
|
||||
if check == "ocr_true":
|
||||
entities, relations = process_file_content(
|
||||
file_name=file_name,
|
||||
file_path=file_path,
|
||||
task_id="batch",
|
||||
data_dir=os.path.dirname(file_path),
|
||||
ontology=ONTOLOGY,
|
||||
ontology_desc=ONTOLOGY_DESC,
|
||||
relationships=RELATIONSHIPS,
|
||||
node_types=NODE_TYPES,
|
||||
relationship_types=RELATIONSHIP_TYPES,
|
||||
prefix_url=PREFIX_URL,
|
||||
api_url=api_url,
|
||||
check_cancelled_callback=lambda: False,
|
||||
device=lower_entity_device,
|
||||
device_type=lower_entity_type,
|
||||
)
|
||||
|
||||
elif check == "ocr_false":
|
||||
async def _run_filtertext():
|
||||
return await get_filtertext_node_relation(
|
||||
input_pdf_path=file_path,
|
||||
data=[],
|
||||
prefix_url=PREFIX_URL,
|
||||
device=lower_entity_device,
|
||||
)
|
||||
|
||||
result = asyncio.run(_run_filtertext())
|
||||
if isinstance(result, tuple) and len(result) >= 2:
|
||||
entities, relations = result[0], result[1]
|
||||
else:
|
||||
entities, relations = [], []
|
||||
logger.info("无需进行ocr解析,直接构建实体")
|
||||
else:
|
||||
logger.warning(f"[{file_name}] check_filename_ocr 返回 None,跳过正文抽取")
|
||||
|
||||
# ── 3. 合并文件名实体 ────────────────────────────────────────────────
|
||||
all_entities = entities + fn_entities
|
||||
all_relations = relations + fn_relations
|
||||
# 过滤掉没有名称的实体
|
||||
all_entities = [e for e in all_entities if e.get("properties", {}).get("名称")]
|
||||
|
||||
# ── 4. 保存 JSON ─────────────────────────────────────────────────────
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
payload = {"entities": all_entities, "relationships": all_relations}
|
||||
with open(out_path, "w", encoding="utf-8") as f:
|
||||
json.dump(payload, f, ensure_ascii=False, indent=2)
|
||||
|
||||
cost = time.time() - t0
|
||||
logger.info(f"[OK] {file_name} 实体:{len(all_entities)} 关系:{len(all_relations)} 耗时:{cost:.1f}s")
|
||||
return {
|
||||
"file": file_name,
|
||||
"status": "ok",
|
||||
"entity_count": len(all_entities),
|
||||
"rel_count": len(all_relations),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[ERROR] {file_name}: {e}\n{traceback.format_exc()}")
|
||||
return {"file": file_name, "status": "error", "error": str(e)}
|
||||
|
||||
|
||||
def _extract_worker(file_path: str, output_dir: str, result_queue: Queue, api_url: str = MINERU_INSTANCES[0]):
|
||||
"""子进程入口:抽取单文件,结果放入队列"""
|
||||
try:
|
||||
result = extract_single_file(file_path, output_dir, api_url)
|
||||
result_queue.put(result)
|
||||
except Exception as e:
|
||||
result_queue.put({"file": os.path.basename(file_path), "status": "error", "error": str(e)})
|
||||
|
||||
|
||||
def _run_file_with_timeout(args: Tuple) -> Dict:
|
||||
"""在线程中启动子进程并监控超时,实现滚动窗口并发。"""
|
||||
fp, output_dir, inst = args
|
||||
file_name = os.path.basename(fp)
|
||||
q = Queue()
|
||||
p = Process(target=_extract_worker, args=(fp, output_dir, q, inst))
|
||||
p.start()
|
||||
p.join(timeout=EXTRACT_TIMEOUT)
|
||||
if p.is_alive():
|
||||
p.kill()
|
||||
p.join()
|
||||
logger.warning(f"[TIMEOUT] {file_name} 超过 {EXTRACT_TIMEOUT}s,已跳过")
|
||||
_record_fault(file_name)
|
||||
return {"file": file_name, "status": "timeout"}
|
||||
try:
|
||||
return q.get_nowait()
|
||||
except Exception:
|
||||
return {"file": file_name, "status": "error"}
|
||||
|
||||
|
||||
def phase_extract(folder: str, output_dir: str, workers: int):
|
||||
"""阶段1:滚动窗口多进程抽取,每完成一个立即补充下一个。"""
|
||||
files = [
|
||||
os.path.join(folder, f)
|
||||
for f in os.listdir(folder)
|
||||
if os.path.isfile(os.path.join(folder, f))
|
||||
and os.path.splitext(f)[1].lower() in SUPPORTED_EXT
|
||||
]
|
||||
logger.info(f"发现 {len(files)} 个文件,并发数: {workers}")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
tasks = [
|
||||
(fp, output_dir, MINERU_INSTANCES[i % len(MINERU_INSTANCES)])
|
||||
for i, fp in enumerate(files)
|
||||
]
|
||||
|
||||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||||
results = list(executor.map(_run_file_with_timeout, tasks))
|
||||
|
||||
ok = sum(1 for r in results if r["status"] == "ok")
|
||||
skip = sum(1 for r in results if r["status"] == "skip")
|
||||
error = sum(1 for r in results if r["status"] == "error")
|
||||
timeout = sum(1 for r in results if r["status"] == "timeout")
|
||||
logger.info(f"抽取完成 成功:{ok} 跳过:{skip} 失败:{error} 超时:{timeout}")
|
||||
return results
|
||||
|
||||
|
||||
# ════════════════════════════════════════════════════════════════════════════
|
||||
# 阶段 2:汇总 → 去重 → Embedding → 写 Neo4j
|
||||
# ════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
|
||||
def phase_merge(output_dir: str):
|
||||
"""阶段2:读所有 JSON → 去重 → 向量化 → 写 Neo4j"""
|
||||
|
||||
# ── 1. 读取所有抽取结果 ──────────────────────────────────────────────────
|
||||
all_entities: List[Dict] = []
|
||||
all_relations: List[Dict] = []
|
||||
|
||||
json_files = [f for f in os.listdir(output_dir) if f.endswith(".json")]
|
||||
logger.info(f"读取 {len(json_files)} 个抽取结果文件")
|
||||
|
||||
for fname in json_files:
|
||||
fpath = os.path.join(output_dir, fname)
|
||||
try:
|
||||
with open(fpath, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
all_entities.extend(data.get("entities", []))
|
||||
all_relations.extend(data.get("relationships", []))
|
||||
except Exception as e:
|
||||
logger.warning(f"读取失败 {fname}: {e}")
|
||||
|
||||
logger.info(f"汇总: 实体 {len(all_entities)} 关系 {len(all_relations)}")
|
||||
|
||||
# ── 2. 去重(filter_node_relationship 只跑这一次) ───────────────────────
|
||||
merged_input = {"entities": all_entities, "relationships": all_relations}
|
||||
logger.info("开始全局去重(filter_node_relationship)...")
|
||||
t0 = time.time()
|
||||
try:
|
||||
filter_data = filter_node_relationship(
|
||||
URI=NEO4J_URI,
|
||||
USERNAME=NEO4J_USER,
|
||||
PASSWORD=NEO4J_PASSWORD,
|
||||
merged_output_data=merged_input,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"去重失败,使用原始数据: {e}")
|
||||
filter_data = merged_input
|
||||
logger.info(f"去重完成 实体:{len(filter_data['entities'])} 关系:{len(filter_data['relationships'])} 耗时:{time.time() - t0:.1f}s")
|
||||
|
||||
# ── 3. 生成 Embedding(并行) ────────────────────────────────────────────
|
||||
logger.info(f"生成 Embedding(并发 {EMBED_WORKERS} 线程)...")
|
||||
t0 = time.time()
|
||||
|
||||
def _embed_one(entity: Dict) -> int:
|
||||
name = str(entity.get("properties", {}).get("名称", "") or "")
|
||||
if not name or "未知" in name:
|
||||
return 0
|
||||
try:
|
||||
entity["properties"]["embedding"] = OpenaiAPI.generate_embedding(name)
|
||||
return 1
|
||||
except Exception as e:
|
||||
logger.warning(f"Embedding 失败 [{name}]: {e}")
|
||||
return 0
|
||||
|
||||
with ThreadPoolExecutor(max_workers=EMBED_WORKERS) as executor:
|
||||
embedded = sum(executor.map(_embed_one, filter_data["entities"]))
|
||||
logger.info(f"Embedding 完成 {embedded} 个实体 耗时:{time.time() - t0:.1f}s")
|
||||
|
||||
# ── 3.5 本地保存带 Embedding 的实体关系 ──────────────────────────────────
|
||||
snapshot_path = os.path.join(os.path.dirname(output_dir) if output_dir != "/" else "/app", "kg_snapshot.json")
|
||||
try:
|
||||
with open(snapshot_path, "w", encoding="utf-8") as f:
|
||||
json.dump(filter_data, f, ensure_ascii=False, indent=2)
|
||||
logger.info(f"本地快照已保存: {snapshot_path} 实体:{len(filter_data['entities'])} 关系:{len(filter_data['relationships'])}")
|
||||
except Exception as e:
|
||||
logger.warning(f"本地快照保存失败: {e}")
|
||||
|
||||
# ── 4. 写入 Neo4j ─────────────────────────────────────────────────────────
|
||||
logger.info("写入 Neo4j...")
|
||||
# ontology, ontology_desc, node_types = _build_ontology_config()
|
||||
|
||||
# 从实体数据中提取动态类型,追加到静态 ontology 末尾(静态配置优先,数据类型补充)
|
||||
data_types = list(dict.fromkeys(
|
||||
e["type"] for e in filter_data["entities"] if e.get("type")
|
||||
))
|
||||
merged_ontology = list(dict.fromkeys( data_types))
|
||||
# new_types = [t for t in data_types if t not in set(ontology)]
|
||||
# if new_types:
|
||||
# logger.info(f"从实体数据补充新类型 {len(new_types)} 个: {new_types}")
|
||||
|
||||
t0 = time.time()
|
||||
kg = MilitaryKnowledgeGraph(
|
||||
uri=NEO4J_URI,
|
||||
auth=(NEO4J_USER, NEO4J_PASSWORD),
|
||||
Ontology=merged_ontology,
|
||||
)
|
||||
try:
|
||||
entity_count, rel_count = kg.process_data(filter_data)
|
||||
logger.info(f"Neo4j 写入完成 实体:{entity_count} 关系:{rel_count} 耗时:{time.time() - t0:.1f}s")
|
||||
finally:
|
||||
kg.close()
|
||||
|
||||
|
||||
def phase_load_snapshot(snapshot_path: str):
|
||||
"""
|
||||
Directly read kg_snapshot.json and store entities/relations into Neo4j.
|
||||
Skips extraction and embedding generation steps.
|
||||
"""
|
||||
if not os.path.exists(snapshot_path):
|
||||
logger.error(f"Snapshot file not found: {snapshot_path}")
|
||||
return
|
||||
|
||||
logger.info(f"Loading snapshot from: {snapshot_path}")
|
||||
t0 = time.time()
|
||||
|
||||
try:
|
||||
with open(snapshot_path, "r", encoding="utf-8") as f:
|
||||
filter_data = json.load(f)
|
||||
|
||||
entities = filter_data.get("entities", [])
|
||||
relations = filter_data.get("relationships", [])
|
||||
logger.info(f"Loaded snapshot: {len(entities)} entities, {len(relations)} relations")
|
||||
|
||||
if not entities and not relations:
|
||||
logger.warning("Snapshot contains no data to load.")
|
||||
return
|
||||
|
||||
# Extract dynamic types from entities to ensure Ontology coverage
|
||||
data_types = list(dict.fromkeys(
|
||||
e["type"] for e in entities if e.get("type")
|
||||
))
|
||||
merged_ontology = list(dict.fromkeys(data_types))
|
||||
|
||||
logger.info(f"Initializing MilitaryKnowledgeGraph with ontology types: {merged_ontology[:5]}...") # Log first few for brevity
|
||||
|
||||
kg = MilitaryKnowledgeGraph(
|
||||
uri=NEO4J_URI,
|
||||
auth=(NEO4J_USER, NEO4J_PASSWORD),
|
||||
Ontology=merged_ontology,
|
||||
)
|
||||
|
||||
try:
|
||||
entity_count, rel_count = kg.process_data(filter_data)
|
||||
elapsed = time.time() - t0
|
||||
logger.info(f"Neo4j write completed from snapshot. Entities: {entity_count}, Relations: {rel_count}. Time: {elapsed:.1f}s")
|
||||
finally:
|
||||
kg.close()
|
||||
|
||||
except json.JSONDecodeError:
|
||||
logger.error(f"Failed to decode JSON from {snapshot_path}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error loading snapshot: {e}", exc_info=True)
|
||||
|
||||
# ════════════════════════════════════════════════════════════════════════════
|
||||
# 入口
|
||||
# ════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="批量知识图谱构建")
|
||||
parser.add_argument("--folder", default="", help="待处理文件夹(阶段1必填)")
|
||||
parser.add_argument("--output_dir", default="/app/kg_output", help="JSON 输出/读取目录")
|
||||
# Update choices to include 'load_snapshot'
|
||||
parser.add_argument("--phase", default="all", choices=["extract", "merge", "load_snapshot", "all"],
|
||||
help="extract=只抽取 merge=只入库(含抽取结果汇总) load_snapshot=直接加载快照 all=全流程")
|
||||
parser.add_argument("--workers", type=int, default=1, help="阶段1并发线程数")
|
||||
parser.add_argument("--snapshot_path", default="", help="快照文件路径 (用于 load_snapshot 阶段)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Determine default snapshot path if not provided and phase is load_snapshot
|
||||
if args.phase == "load_snapshot" and not args.snapshot_path:
|
||||
# Default to the location where phase_merge saves it
|
||||
args.snapshot_path = os.path.join(os.path.dirname(args.output_dir) if args.output_dir != "/" else "/app", "kg_snapshot.json")
|
||||
|
||||
if args.phase in ("extract", "all"):
|
||||
if not args.folder:
|
||||
parser.error("--phase extract/all 需要指定 --folder")
|
||||
phase_extract(args.folder, args.output_dir, args.workers)
|
||||
|
||||
if args.phase in ("merge", "all"):
|
||||
phase_merge(args.output_dir)
|
||||
|
||||
if args.phase == "load_snapshot":
|
||||
phase_load_snapshot(args.snapshot_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
# ontology, ontology_desc, node_types = _build_ontology_config()
|
||||
# print(ontology)
|
||||
|
||||
422
batch_kg_build_timeout.py
Normal file
422
batch_kg_build_timeout.py
Normal file
@ -0,0 +1,422 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
批量知识图谱构建脚本 —— 三阶段 ETL
|
||||
|
||||
阶段1 (extract): 对文件夹内所有文件做实体抽取,每个文件结果存为 JSON
|
||||
阶段2 (merge): 汇总所有 JSON,去重,生成 embedding,写入 Neo4j
|
||||
|
||||
用法:
|
||||
# 只跑抽取(4个容器各跑各的文件夹)
|
||||
python batch_kg_build.py --folder /app/files --output_dir /app/kg_output --phase extract --workers 4
|
||||
|
||||
# 只跑入库(所有容器跑完后,跑一次)
|
||||
python batch_kg_build.py --output_dir /app/kg_output --phase merge
|
||||
|
||||
# 全流程一次跑完
|
||||
python batch_kg_build.py --folder /app/files --output_dir /app/kg_output --phase all --workers 4
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import logging
|
||||
import argparse
|
||||
import asyncio
|
||||
import traceback
|
||||
import time
|
||||
from datetime import datetime
|
||||
from multiprocessing import Process, Queue
|
||||
from typing import List, Dict, Tuple, Any, Optional
|
||||
|
||||
EXTRACT_TIMEOUT = 600 # 单文件抽取超时秒数
|
||||
FAULT_DIR = "/app/fault_wenjian"
|
||||
|
||||
|
||||
def _record_fault(file_name: str):
|
||||
"""将超时文件名追加到 /app/fault_wenjian/timeout_YYYYMMDD.txt"""
|
||||
os.makedirs(FAULT_DIR, exist_ok=True)
|
||||
date_str = datetime.now().strftime("%Y%m%d")
|
||||
fault_path = os.path.join(FAULT_DIR, f"timeout_{date_str}.txt")
|
||||
with open(fault_path, "a", encoding="utf-8") as f:
|
||||
f.write(file_name + "\n")
|
||||
logger.warning(f"[FAULT] 已记录超时文件: {fault_path}")
|
||||
|
||||
# ── 项目内模块 ──────────────────────────────────────────────────────────────
|
||||
from extract_text_node_relation import process_file_content, check_filename_ocr
|
||||
from kg_build.extract_filtertext import get_filtertext_node_relation
|
||||
from filename_proceess_and_kgquery import get_entity
|
||||
from kg_build.filter_node_relation import filter_node_relationship
|
||||
from kg_build.kg_save import MilitaryKnowledgeGraph
|
||||
from modelsAPI.model_api import OpenaiAPI
|
||||
import importlib
|
||||
import entity_registry as _entity_registry
|
||||
from default_ontology_config import (
|
||||
ONTOLOGY as _BASE_ONTOLOGY,
|
||||
ONTOLOGY_DESCRIPTION as _BASE_ONTOLOGY_DESC,
|
||||
NODE_TYPES as _BASE_NODE_TYPES,
|
||||
RELATIONSHIPS as RELATIONSHIPS,
|
||||
RELATIONSHIP_TYPES as RELATIONSHIP_TYPES,
|
||||
)
|
||||
|
||||
|
||||
def _build_ontology_config():
|
||||
"""
|
||||
与 app.py build_kg_config 保持一致:
|
||||
default_ontology_config 的静态配置 + entity_registry.json 动态注册的实体类型合并。
|
||||
"""
|
||||
importlib.reload(_entity_registry)
|
||||
reg_list = _entity_registry.get_ontology_list()
|
||||
reg_desc = _entity_registry.get_ontology_description()
|
||||
reg_nodes = _entity_registry.get_node_types()
|
||||
|
||||
# 合并:静态配置优先(手写配置不被覆盖)
|
||||
ontology = list(dict.fromkeys(_BASE_ONTOLOGY + reg_list)) # 去重保序
|
||||
ontology_desc = {**reg_desc, **_BASE_ONTOLOGY_DESC}
|
||||
node_types = {**reg_nodes, **_BASE_NODE_TYPES}
|
||||
return ontology, ontology_desc, node_types
|
||||
|
||||
|
||||
# 模块加载时构建一次(抽取阶段开始前已有 registry 内容)
|
||||
ONTOLOGY, ONTOLOGY_DESC, NODE_TYPES = _build_ontology_config()
|
||||
|
||||
# ── 配置(与 app.py 保持一致,按需修改) ────────────────────────────────────
|
||||
NEO4J_URI = os.getenv("NEO4J_URI", "bolt://localhost:7687")
|
||||
NEO4J_USER = os.getenv("NEO4J_USER", "neo4j")
|
||||
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD", "zdht123@")
|
||||
MINERU_ANALYZE = os.getenv("MINERU_ANALYZE", "http://192.168.1.64:18000/analyze-pdf")
|
||||
MINERU_CACHE = os.getenv("MINERU_CACHE", "http://192.168.1.64:18000/get_content_list")
|
||||
PREFIX_URL = os.getenv("PREFIX_URL", "http://192.168.1.64:9085")
|
||||
|
||||
SUPPORTED_EXT = {".pdf", ".md", ".docx", ".xlsx"}
|
||||
|
||||
# ── 日志 ────────────────────────────────────────────────────────────────────
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||||
handlers=[
|
||||
logging.StreamHandler(),
|
||||
logging.FileHandler("batch_kg_build.log", encoding="utf-8"),
|
||||
],
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# ════════════════════════════════════════════════════════════════════════════
|
||||
# 阶段 1:单文件抽取
|
||||
# ════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def extract_single_file(file_path: str, output_dir: str) -> Dict:
|
||||
"""
|
||||
对单个文件做实体/关系抽取,结果保存到 output_dir/<filename>.json。
|
||||
已存在 JSON 则跳过(断点续跑)。
|
||||
返回 {"file": ..., "status": "ok"/"skip"/"error", "entity_count": ..., "rel_count": ...}
|
||||
"""
|
||||
file_name = os.path.basename(file_path)
|
||||
ext = os.path.splitext(file_name)[1].lower()
|
||||
out_path = os.path.join(output_dir, file_name + ".json")
|
||||
|
||||
# ── 断点续跑:已有结果直接跳过 ──────────────────────────────────────────
|
||||
if os.path.exists(out_path):
|
||||
logger.info(f"[SKIP] 已有结果,跳过: {file_name}")
|
||||
return {"file": file_name, "status": "skip"}
|
||||
|
||||
if ext not in SUPPORTED_EXT:
|
||||
logger.warning(f"[SKIP] 不支持的文件类型: {file_name}")
|
||||
return {"file": file_name, "status": "skip"}
|
||||
|
||||
t0 = time.time()
|
||||
entities: List[Dict] = []
|
||||
relations: List[Dict] = []
|
||||
lower_entity_device: Optional[str] = None
|
||||
lower_entity_type: Optional[str] = None
|
||||
|
||||
try:
|
||||
# ── 1. 文件名实体提取 ────────────────────────────────────────────────
|
||||
try:
|
||||
fn_result = get_entity(file_name)
|
||||
fn_entities = fn_result.get("entities", [])
|
||||
fn_relations = fn_result.get("relationships", [])
|
||||
lower_entity = fn_result.get("low_level")
|
||||
for e in fn_entities:
|
||||
if e.get("properties", {}).get("名称") == lower_entity:
|
||||
lower_entity_device = e["properties"].get("名称")
|
||||
lower_entity_type = e.get("type")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.warning(f"[{file_name}] 文件名实体提取失败: {e}")
|
||||
fn_entities, fn_relations = [], []
|
||||
|
||||
# ── 2. 正文实体抽取 ──────────────────────────────────────────────────
|
||||
check = check_filename_ocr(file_name)
|
||||
|
||||
if check == "ocr_true":
|
||||
entities, relations = process_file_content(
|
||||
file_name = file_name,
|
||||
file_path = file_path,
|
||||
task_id = "batch",
|
||||
data_dir = os.path.dirname(file_path),
|
||||
ontology = ONTOLOGY,
|
||||
ontology_desc = ONTOLOGY_DESC,
|
||||
relationships = RELATIONSHIPS,
|
||||
node_types = NODE_TYPES,
|
||||
relationship_types = RELATIONSHIP_TYPES,
|
||||
prefix_url = PREFIX_URL,
|
||||
api_url = MINERU_ANALYZE,
|
||||
check_cancelled_callback = lambda: False,
|
||||
device = lower_entity_device,
|
||||
device_type = lower_entity_type,
|
||||
get_contentlist_url = MINERU_CACHE,
|
||||
)
|
||||
|
||||
elif check == "ocr_false":
|
||||
async def _run_filtertext():
|
||||
return await get_filtertext_node_relation(
|
||||
input_pdf_path = file_path,
|
||||
data = [],
|
||||
prefix_url = PREFIX_URL,
|
||||
device = lower_entity_device,
|
||||
)
|
||||
result = asyncio.run(_run_filtertext())
|
||||
if isinstance(result, tuple) and len(result) >= 2:
|
||||
entities, relations = result[0], result[1]
|
||||
else:
|
||||
entities, relations = [], []
|
||||
|
||||
else:
|
||||
logger.warning(f"[{file_name}] check_filename_ocr 返回 None,跳过正文抽取")
|
||||
|
||||
# ── 3. 合并文件名实体 ────────────────────────────────────────────────
|
||||
all_entities = entities + fn_entities
|
||||
all_relations = relations + fn_relations
|
||||
# 过滤掉没有名称的实体
|
||||
all_entities = [e for e in all_entities if e.get("properties", {}).get("名称")]
|
||||
|
||||
# ── 4. 保存 JSON ─────────────────────────────────────────────────────
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
payload = {"entities": all_entities, "relationships": all_relations}
|
||||
with open(out_path, "w", encoding="utf-8") as f:
|
||||
json.dump(payload, f, ensure_ascii=False, indent=2)
|
||||
|
||||
cost = time.time() - t0
|
||||
logger.info(
|
||||
f"[OK] {file_name} 实体:{len(all_entities)} 关系:{len(all_relations)} 耗时:{cost:.1f}s"
|
||||
)
|
||||
return {
|
||||
"file": file_name,
|
||||
"status": "ok",
|
||||
"entity_count": len(all_entities),
|
||||
"rel_count": len(all_relations),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[ERROR] {file_name}: {e}\n{traceback.format_exc()}")
|
||||
return {"file": file_name, "status": "error", "error": str(e)}
|
||||
|
||||
|
||||
def _extract_worker(file_path: str, output_dir: str, result_queue: Queue):
|
||||
"""子进程入口:抽取单文件,结果放入队列"""
|
||||
try:
|
||||
result = extract_single_file(file_path, output_dir)
|
||||
result_queue.put(result)
|
||||
except Exception as e:
|
||||
result_queue.put({"file": os.path.basename(file_path), "status": "error", "error": str(e)})
|
||||
|
||||
|
||||
def phase_extract(folder: str, output_dir: str, workers: int):
|
||||
"""阶段1:多进程抽取,每个文件独立进程,超时后真正 kill"""
|
||||
files = [
|
||||
os.path.join(folder, f)
|
||||
for f in os.listdir(folder)
|
||||
if os.path.isfile(os.path.join(folder, f))
|
||||
and os.path.splitext(f)[1].lower() in SUPPORTED_EXT
|
||||
]
|
||||
logger.info(f"发现 {len(files)} 个文件,并发数: {workers}")
|
||||
|
||||
results = []
|
||||
# 按 workers 大小分批,每批同时跑 workers 个进程
|
||||
for i in range(0, len(files), workers):
|
||||
batch = files[i:i + workers]
|
||||
procs = []
|
||||
for fp in batch:
|
||||
q = Queue()
|
||||
p = Process(target=_extract_worker, args=(fp, output_dir, q))
|
||||
p.start()
|
||||
procs.append((p, fp, q))
|
||||
|
||||
for p, fp, q in procs:
|
||||
file_name = os.path.basename(fp)
|
||||
p.join(timeout=EXTRACT_TIMEOUT)
|
||||
if p.is_alive():
|
||||
p.kill()
|
||||
p.join()
|
||||
logger.warning(f"[TIMEOUT] {file_name} 超过 {EXTRACT_TIMEOUT}s,已跳过")
|
||||
_record_fault(file_name)
|
||||
results.append({"file": file_name, "status": "timeout"})
|
||||
else:
|
||||
try:
|
||||
results.append(q.get_nowait())
|
||||
except Exception:
|
||||
results.append({"file": file_name, "status": "error"})
|
||||
|
||||
ok = sum(1 for r in results if r["status"] == "ok")
|
||||
skip = sum(1 for r in results if r["status"] == "skip")
|
||||
error = sum(1 for r in results if r["status"] == "error")
|
||||
timeout = sum(1 for r in results if r["status"] == "timeout")
|
||||
logger.info(f"抽取完成 成功:{ok} 跳过:{skip} 失败:{error} 超时:{timeout}")
|
||||
return results
|
||||
|
||||
|
||||
# ════════════════════════════════════════════════════════════════════════════
|
||||
# 阶段 2:汇总 → 去重 → Embedding → 写 Neo4j
|
||||
# ════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def phase_merge(output_dir: str):
|
||||
"""阶段2:读所有 JSON → 去重 → 向量化 → 写 Neo4j"""
|
||||
|
||||
# ── 1. 读取所有抽取结果 ──────────────────────────────────────────────────
|
||||
all_entities: List[Dict] = []
|
||||
all_relations: List[Dict] = []
|
||||
|
||||
json_files = [f for f in os.listdir(output_dir) if f.endswith(".json")]
|
||||
logger.info(f"读取 {len(json_files)} 个抽取结果文件")
|
||||
|
||||
for fname in json_files:
|
||||
fpath = os.path.join(output_dir, fname)
|
||||
try:
|
||||
with open(fpath, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
all_entities.extend(data.get("entities", []))
|
||||
all_relations.extend(data.get("relationships", []))
|
||||
except Exception as e:
|
||||
logger.warning(f"读取失败 {fname}: {e}")
|
||||
|
||||
logger.info(f"汇总: 实体 {len(all_entities)} 关系 {len(all_relations)}")
|
||||
|
||||
# ── 2. 去重(filter_node_relationship 只跑这一次) ───────────────────────
|
||||
merged_input = {"entities": all_entities, "relationships": all_relations}
|
||||
logger.info("开始全局去重(filter_node_relationship)...")
|
||||
t0 = time.time()
|
||||
try:
|
||||
filter_data = filter_node_relationship(
|
||||
URI = NEO4J_URI,
|
||||
USERNAME = NEO4J_USER,
|
||||
PASSWORD = NEO4J_PASSWORD,
|
||||
merged_output_data = merged_input,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"去重失败,使用原始数据: {e}")
|
||||
filter_data = merged_input
|
||||
logger.info(f"去重完成 实体:{len(filter_data['entities'])} 关系:{len(filter_data['relationships'])} 耗时:{time.time()-t0:.1f}s")
|
||||
|
||||
# ── 3. 生成 Embedding ────────────────────────────────────────────────────
|
||||
logger.info("生成 Embedding...")
|
||||
t0 = time.time()
|
||||
embedded = 0
|
||||
for entity in filter_data["entities"]:
|
||||
name = entity.get("properties", {}).get("名称", "")
|
||||
if not name or "未知" in name:
|
||||
continue
|
||||
try:
|
||||
entity["properties"]["embedding"] = OpenaiAPI.generate_embedding(name)
|
||||
embedded += 1
|
||||
except Exception as e:
|
||||
logger.warning(f"Embedding 失败 [{name}]: {e}")
|
||||
logger.info(f"Embedding 完成 {embedded} 个实体 耗时:{time.time()-t0:.1f}s")
|
||||
|
||||
# ── 4. 写入 Neo4j ─────────────────────────────────────────────────────────
|
||||
logger.info("写入 Neo4j...")
|
||||
t0 = time.time()
|
||||
kg = MilitaryKnowledgeGraph(
|
||||
uri = NEO4J_URI,
|
||||
auth = (NEO4J_USER, NEO4J_PASSWORD),
|
||||
Ontology = ONTOLOGY,
|
||||
)
|
||||
try:
|
||||
entity_count, rel_count = kg.process_data(filter_data)
|
||||
logger.info(
|
||||
f"Neo4j 写入完成 实体:{entity_count} 关系:{rel_count} 耗时:{time.time()-t0:.1f}s"
|
||||
)
|
||||
finally:
|
||||
kg.close()
|
||||
|
||||
|
||||
# ════════════════════════════════════════════════════════════════════════════
|
||||
# 入口
|
||||
# ════════════════════════════════════════════════════════════════════════════
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="批量知识图谱构建")
|
||||
parser.add_argument("--folder", default="", help="待处理文件夹(阶段1必填)")
|
||||
parser.add_argument("--output_dir", default="/app/kg_output", help="JSON 输出/读取目录")
|
||||
parser.add_argument("--phase", default="all", choices=["extract", "merge", "all"],
|
||||
help="extract=只抽取 merge=只入库 all=全流程")
|
||||
parser.add_argument("--workers", type=int, default=1, help="阶段1并发线程数")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.phase in ("extract", "all"):
|
||||
if not args.folder:
|
||||
parser.error("--phase extract/all 需要指定 --folder")
|
||||
phase_extract(args.folder, args.output_dir, args.workers)
|
||||
|
||||
if args.phase in ("merge", "all"):
|
||||
phase_merge(args.output_dir)
|
||||
|
||||
|
||||
def _fake_worker(file_path: str, output_dir: str, result_queue: Queue):
|
||||
"""测试用子进程:file_B 睡 20s 模拟超时,其余 2s"""
|
||||
file_name = os.path.basename(file_path)
|
||||
sleep_time = 20 if file_name == "file_B" else 2
|
||||
time.sleep(sleep_time)
|
||||
out_path = os.path.join(output_dir, file_name + ".json")
|
||||
with open(out_path, "w", encoding="utf-8") as f:
|
||||
json.dump({"file": file_name}, f)
|
||||
result_queue.put({"file": file_name, "status": "ok"})
|
||||
|
||||
|
||||
def _test_timeout():
|
||||
"""
|
||||
模拟超时跳过逻辑,验证完删掉此函数及 _fake_worker、下面的调用。
|
||||
file_A 正常(2s),file_B 超时(20s),file_C 正常(2s)
|
||||
期望总耗时约 TIMEOUT+2s,file_B 不写 JSON
|
||||
"""
|
||||
import tempfile
|
||||
|
||||
global EXTRACT_TIMEOUT
|
||||
EXTRACT_TIMEOUT = 5 # 测试用 5 秒
|
||||
|
||||
output_dir = tempfile.mkdtemp()
|
||||
logger.info(f"测试输出目录: {output_dir}")
|
||||
|
||||
files = ["file_A", "file_B", "file_C"]
|
||||
results = []
|
||||
t_start = time.time()
|
||||
|
||||
for i in range(0, len(files), 2):
|
||||
batch = files[i:i + 2]
|
||||
procs = []
|
||||
for f in batch:
|
||||
q = Queue()
|
||||
p = Process(target=_fake_worker, args=(f, output_dir, q))
|
||||
p.start()
|
||||
procs.append((p, f, q))
|
||||
|
||||
for p, f, q in procs:
|
||||
p.join(timeout=EXTRACT_TIMEOUT)
|
||||
if p.is_alive():
|
||||
p.kill()
|
||||
p.join()
|
||||
logger.warning(f"[TIMEOUT] {f} 超时,进程已 kill,不写 JSON")
|
||||
_record_fault(f)
|
||||
results.append({"file": f, "status": "timeout"})
|
||||
else:
|
||||
results.append(q.get_nowait())
|
||||
|
||||
logger.info(f"总耗时: {time.time()-t_start:.1f}s")
|
||||
logger.info(f"结果: {results}")
|
||||
logger.info(f"output_dir 中的文件: {os.listdir(output_dir)}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
_test_timeout()
|
||||
# main()
|
||||
|
||||
|
||||
1181
chunk_text.py
Normal file
1181
chunk_text.py
Normal file
File diff suppressed because it is too large
Load Diff
41
cleardatabase.py
Normal file
41
cleardatabase.py
Normal file
@ -0,0 +1,41 @@
|
||||
import re
|
||||
from langchain_community.graphs.graph_document import GraphDocument, Node, Relationship
|
||||
from langchain_core.documents import Document
|
||||
|
||||
from langchain_community.graphs import Neo4jGraph
|
||||
from langchain_core.documents import Document
|
||||
import os
|
||||
os.environ["NEO4J_URI"] = "bolt://192.168.0.46:57687"
|
||||
os.environ["NEO4J_USERNAME"] = "neo4j"
|
||||
os.environ["NEO4J_PASSWORD"] = "zdht123@"
|
||||
|
||||
graph = Neo4jGraph(refresh_schema=False)
|
||||
|
||||
|
||||
#清空数据库内容
|
||||
graph.query("MATCH (n) DETACH DELETE n")
|
||||
|
||||
|
||||
|
||||
#查询某个节点的相关节点
|
||||
# selectsql = """
|
||||
# MATCH (e:__Entity__ {id: "进气系统"})-[r]-(related)
|
||||
# RETURN related, type(r) AS relationship_type
|
||||
# """
|
||||
# results = graph.query(selectsql)
|
||||
# #
|
||||
# # # 打印结果
|
||||
# print(f"共找到 {len(results)} 个关联节点:\n")
|
||||
# for i, record in enumerate(results, 1):
|
||||
# related_node = record["related"]
|
||||
# rel_type = record["relationship_type"]
|
||||
|
||||
# # 提取 related 节点的 id(或其他属性)
|
||||
# node_id = related_node.get("id", "N/A")
|
||||
# labels = related_node.get("labels", []) # 注意:LangChain 返回的节点可能包含 labels
|
||||
|
||||
# print(f"{i}. 关联节点 ID: {node_id}")
|
||||
# print(f" 标签: {labels}")
|
||||
# print(f" 关系类型: {rel_type}")
|
||||
# print(f" 完整节点数据: {related_node}")
|
||||
# print("-" * 50)
|
||||
49
config.py
Normal file
49
config.py
Normal file
@ -0,0 +1,49 @@
|
||||
# app/model_config.py
|
||||
|
||||
import os
|
||||
|
||||
# ==================== 大模型配置 ====================
|
||||
LLM_CONFIG = {
|
||||
"model": "46-qwen3.5-35B",
|
||||
"base_url": "https://openai.zkzdht.com/v1",
|
||||
"api_key": "gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5",
|
||||
"temperature": 0.1,
|
||||
"max_tokens": 25096,
|
||||
"timeout": 30
|
||||
}
|
||||
|
||||
# ==================== 嵌入模型配置 ====================
|
||||
EMBEDDING_CONFIG = {
|
||||
"model": "bge-m3",
|
||||
"base_url": "http://embedding:9700/v1/embeddings",
|
||||
"api_key": "gpustack_dee9ca823290886c_5edfc86aeeeceb1e9ee5162941cb2cb5",
|
||||
"base_url_v2": "http://embedding:9700/v1",
|
||||
}
|
||||
|
||||
# ==================== Neo4j 图数据库配置 ====================
|
||||
NEO4J_CONFIG = {
|
||||
# 优先读取环境变量,如果没有则使用默认值
|
||||
# 注意:本地开发通常用 localhost,部署时用服务名 neo4j
|
||||
"uri": "bolt://neo4j:7687",
|
||||
"username": "neo4j",
|
||||
"password": "zdht123@",
|
||||
"database": "neo4j", # 新增 database 配置
|
||||
}
|
||||
|
||||
# ==================== 检索与索引配置 ====================
|
||||
SEARCH_CONFIG = {
|
||||
"name_property": "名称",
|
||||
"fulltext_property": "fulltext",
|
||||
"embedding_property": "embedding",
|
||||
"search_label": "Searchable",
|
||||
"vector_index_name": "global_searchable_embedding",
|
||||
"fulltext_index_name": "global_searchable_content_search",
|
||||
"fulltext_field_index_name": "searchable_fulltext",
|
||||
"excluded_business_labels": [
|
||||
"Entity", "Chunk", "Document", "_Bloom_Perspective_", "Searchable"
|
||||
],
|
||||
"vector_dimension": 1024,
|
||||
"jieba_pos_whitelist": ["n", "nr", "ns", "nt", "nz", "vn", "eng"],
|
||||
}
|
||||
|
||||
|
||||
137
dataprocess.py
Normal file
137
dataprocess.py
Normal file
@ -0,0 +1,137 @@
|
||||
from langchain_text_splitters import MarkdownHeaderTextSplitter, CharacterTextSplitter
|
||||
import requests
|
||||
import json
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
def _split_markdown(text, chunk_size=2048, chunk_overlap=50):
|
||||
"""
|
||||
根据 max_level 动态生成 headers_to_split_on 配置
|
||||
规则:
|
||||
- max_level=1 → 只按 # 分割
|
||||
- max_level=2 → 按 # 和 ## 分割
|
||||
- 以此类推
|
||||
"""
|
||||
# 动态生成 headers_to_split_on 配置(目前固定为 Header_1)
|
||||
headers_to_split_on = [("#", "Header_1")]
|
||||
|
||||
# 初始化分割器
|
||||
splitter = MarkdownHeaderTextSplitter(
|
||||
headers_to_split_on=headers_to_split_on,
|
||||
strip_headers=True
|
||||
)
|
||||
chunks = splitter.split_text(text)
|
||||
|
||||
char_splitter = CharacterTextSplitter(
|
||||
chunk_size=chunk_size,
|
||||
chunk_overlap=chunk_overlap,
|
||||
length_function=len,
|
||||
separator="\n\n" # 更合理的分隔符
|
||||
)
|
||||
|
||||
final_chunks = []
|
||||
start_chars = ["1", "2", "3", "4", "5", "6", "7", "8", "9", "0", "第", "十",
|
||||
"一", "二", "三", "四", "五", "六", "七", "八", "九"]
|
||||
new_header = []
|
||||
|
||||
# 提取所有 Header_1 并筛选出“有效标题”
|
||||
for chunk in chunks:
|
||||
header_1 = chunk.metadata.get('Header_1', '')
|
||||
if header_1 and header_1[0] in start_chars:
|
||||
new_header.append(header_1)
|
||||
|
||||
new_header_set = set(new_header)
|
||||
|
||||
# 构建带标题前缀的 final_chunks
|
||||
for chunk in chunks:
|
||||
header_1 = chunk.metadata.get('Header_1', '')
|
||||
content = chunk.page_content
|
||||
if len(content) <= 20000:
|
||||
chunk_with_header = f"{header_1}\n{content}" if header_1 else content
|
||||
final_chunks.append(chunk_with_header)
|
||||
else:
|
||||
split_texts = char_splitter.split_documents([chunk])
|
||||
for split_chunk in split_texts:
|
||||
chunk_with_header = f"{header_1}\n{split_chunk.page_content}" if header_1 else split_chunk.page_content
|
||||
final_chunks.append(chunk_with_header)
|
||||
|
||||
# ===== 重新分组,并记录每个组的 header =====
|
||||
rechunked = []
|
||||
rechunked_headers = [] # 新增:每个 rechunked 块对应的 header
|
||||
current = ""
|
||||
current_header = "" # 当前正在构建的块所属的标题
|
||||
|
||||
for chunk in final_chunks:
|
||||
is_new_section = False
|
||||
matched_header = ""
|
||||
|
||||
# 检查是否以某个 new_header 开头
|
||||
for h in new_header_set:
|
||||
if chunk.startswith(h + "\n"):
|
||||
is_new_section = True
|
||||
matched_header = h
|
||||
break
|
||||
|
||||
if is_new_section:
|
||||
# 保存上一段
|
||||
if current:
|
||||
rechunked.append(current)
|
||||
rechunked_headers.append(current_header)
|
||||
# 开启新段
|
||||
current = chunk
|
||||
current_header = matched_header
|
||||
else:
|
||||
# 合并到当前段
|
||||
if current:
|
||||
current += "\n" + chunk
|
||||
else:
|
||||
current = chunk
|
||||
current_header = "" # 如果开头无标题,则 header 为空
|
||||
|
||||
# 添加最后一段
|
||||
if current:
|
||||
rechunked.append(current)
|
||||
rechunked_headers.append(current_header)
|
||||
|
||||
return rechunked, rechunked_headers
|
||||
|
||||
def read_markdown_file(file_path):
|
||||
"""
|
||||
读取 Markdown 文件并返回其全部内容作为字符串。
|
||||
|
||||
参数:
|
||||
file_path (str): Markdown 文件的路径(支持相对或绝对路径)
|
||||
|
||||
返回:
|
||||
str 或 None: 成功时返回文件内容;出错时返回 None
|
||||
"""
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
content = f.read()
|
||||
return content
|
||||
except FileNotFoundError:
|
||||
print(f"错误:文件未找到 - {file_path}")
|
||||
except PermissionError:
|
||||
print(f"错误:没有权限读取文件 - {file_path}")
|
||||
except UnicodeDecodeError:
|
||||
print(f"错误:文件编码不是 UTF-8,无法读取 - {file_path}")
|
||||
except Exception as e:
|
||||
print(f"读取文件时发生未知错误: {e}")
|
||||
return None
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
# if __name__ == "__main__":
|
||||
|
||||
|
||||
# # 测试代码
|
||||
# file_path = "F:\zklnlp\HJ\code\output.md" # 替换为你的 Markdown 文件路径
|
||||
# markdown_content = read_markdown_file(file_path)
|
||||
# content, headers = _split_markdown(markdown_content, chunk_size=3000, chunk_overlap=100)
|
||||
# for ins in content:
|
||||
# print(ins)
|
||||
# print(1111111111111111111111111)
|
||||
|
||||
152
default_ontology_config.py
Normal file
152
default_ontology_config.py
Normal file
@ -0,0 +1,152 @@
|
||||
|
||||
# ========== 1. 本体类列表 ==========
|
||||
ONTOLOGY = [
|
||||
"舰艇", "系统", "子系统", "设备", "维修项目", "维修工作", "操作项目",
|
||||
"故障模式", "安全警告", "工作原理",
|
||||
"组成", "技术指标参数", "布置情况", "保障设施保障设备工具信息汇总表",
|
||||
"消耗材料信息汇总表", "维修安全防护及注意事项汇总表", "图册", "安全使用操作",
|
||||
"调试", "维护保养", "常见问题分析及排除", "功能", "接口",
|
||||
"性能指标", "操作使用","环境条件","设计单位","操作程序"
|
||||
]
|
||||
|
||||
# ========== 2. 本体描述 ==========
|
||||
ONTOLOGY_DESCRIPTION = {
|
||||
"舰艇": ["代表一艘完整海上作战或保障平台的实体,是所有系统、设备和零部件的最高层级载体。区别于“系统”或“设备”,舰艇具有独立的型号、弦号和部队归属,不可作为其他舰艇的组成部分。"],
|
||||
"系统": ["舰艇中实现某一类综合功能(如动力、导航、通信)的集成化功能单元,由多个子系统或设备协同组成。与“设备”不同,系统本身侧重于逻辑/功能层面的集合与统筹;与“舰艇”不同,系统不具备独立航行能力或部队归属。"],
|
||||
"子系统": ["系统内部为实现特定子功能而划分的独立功能模块,是系统的直接组成部分。子系统既具备相对独立的功能完整性(可单独测试或隔离),又必须依赖所属系统才能发挥最终战术效能,是连接系统逻辑与物理设备的中间层级。"],
|
||||
"设备": ["构成子系统或直接隶属于系统的具有独立物理形态和功能的具体装置。设备是执行具体战术动作或保障任务的最小功能实体(如雷达天线座、柴油发电机、水泵),由多个部件组装而成,可直接进行整体拆装和替换。"],
|
||||
"维修项目": ["针对特定装备或系统设定的标准化维护任务类别,定义了“需要做什么”(如定期更换滤芯、年度校准、预防性检查)。维修项目是制定维修计划和工单的依据,侧重于周期性和预防性,不直接包含具体的执行步骤细节。"],
|
||||
"维修工作": ["执行具体“维修项目”时实际开展的一系列技术操作过程的总和。维修工作包含了从故障诊断、部件拆卸、修理/更换、到重新安装测试的全流程实际操作,是维修人员消耗工时和资源的直接体现,侧重于“正在做”的过程。"],
|
||||
"操作项目": ["舰员在日常使用、值更或训练过程中,对装备执行的具体动作指令或检查任务(如启动、切换模式、参数读取、日常擦拭)。操作项目侧重于装备的正确使用与状态监控,区别于以修复或深度维护为目的的“维修工作”。"],
|
||||
"故障模式": ["对舰艇装备在运行过程中表现出的异常状态、性能退化或功能丧失现象的标准化描述。故障模式明确了“发生了什么错误”(如断裂、短路、泄漏、卡滞),是关联故障现象与维修方案的桥梁,不包含原因分析或解决措施。"],
|
||||
"安全警告": ["针对发动机及相关设备在操作、维修过程中的人身安全防护、设备损坏风险及应急处置的强制性规范和提醒。"],
|
||||
"工作原理": ["阐述设备或系统如何通过能量转换、机械运动或逻辑控制实现其预定功能的动态过程描述。"],
|
||||
"组成": ["定义设备的物理或逻辑构成,包括系统、子系统及零部件之间的层级隶属关系。"],
|
||||
"技术指标参数": ["设备的额定功率、转速、尺寸、重量等可量化的物理性能参数。"],
|
||||
"布置情况": ["描述设备及其各组成部分在船体或机舱内的安装位置与空间分布。"],
|
||||
"保障设施保障设备工具信息汇总表": ["执行维修或保障任务时所需的专用或通用工具、测量仪器及辅助设施清单。"],
|
||||
"消耗材料信息汇总表": ["执行维修或保障任务时所需消耗的物料,包括润滑油、密封件、清洗剂等耗材清单。"],
|
||||
"维修安全防护及注意事项汇总表": ["在执行具体维修项目时,为确保作业安全而必须遵守的防护措施、注意事项及环境要求汇总。"],
|
||||
"图册": ["图录是对舰艇及其系统、设备的图片资料进行分类、描述和索引的集合。图录中的每个条目通常包含图片文件、名称、描述、拍摄角度、相关系统或设备的关联信息等。图录的核心作用在于提供直观的视觉参考,辅助理解设备结构、维修步骤和操作方法等内容,是维修手册和操作指南的重要补充。"],
|
||||
"安全使用操作": ["安全手册主要描述设备主要安全防护、安装方法、启动及运行过程中的操作、注意事项及防范措施等的基本操作规范"],
|
||||
"调试": ["描述设备的安装调试内容"],
|
||||
"维护保养": ["描述设备的维护保养步骤以及流程"],
|
||||
"常见问题分析及排除": ["描述设备的常见问题分析及排除流程"],
|
||||
"接口": ["描述设备的接口"],
|
||||
"性能指标": ["描述设备的性能指标"],
|
||||
"操作使用": ["主要描述设备主要安装方法、启动及运行过程中的操作、注意事项及防范措施等的基本操作规范"],
|
||||
"功能": ["主要描述设备使用功能"],
|
||||
"环境条件": ["主要描述设备运行的环境条件"],
|
||||
"器材保障" : ["设备的器材保障"],
|
||||
"设计单位" : ["设备来源设计单位设计"],
|
||||
"操作程序" : ["设备的操作使用程序"]
|
||||
}
|
||||
|
||||
# ========== 3. 节点类型属性定义 ==========
|
||||
NODE_TYPES = {
|
||||
"舰艇": ["名称", "型号", "弦号", "设计单位", "总装工厂", "所属部队"],
|
||||
"系统": ["名称", "型号", "生产单位", "等级", "功能"],
|
||||
"子系统": ["名称", "型号", "生产单位", "等级", "功能"],
|
||||
"维修工作": ["名称", "故障原因", "故障影响", "故障类型", "发生时间", "处置情况", "修复情况", "维修方案"],
|
||||
"故障模式": ["名称", "故障原因", "故障影响", "故障类型", "发生时间", "处置情况", "修复情况", "维修方案"],
|
||||
"维修项目": ["名称", "维修项目编号", "维修间隔期", "安全注意事项", "保障设备(工具)", "备品备件", "消耗材料", "人力人员","修前判据", "修理步骤及工艺", "验收标准"],
|
||||
"操作项目": ["名称", "操作项目编号", "安全防护", "注意事项", "保障设施、保障设备(工具)", "人力人员", "操作步骤"],
|
||||
"安全警告": ["名称", "警告", "注意", "原文"],
|
||||
"工作原理": [
|
||||
"名称", "主发动机工作原理", "机座与机架组件", "曲柄连杆机构",
|
||||
"燃烧室部件", "排气阀驱动系统", "燃油喷射系统", "扫气与增压系统", "原文"
|
||||
],
|
||||
"组成": ["名称", "描述", "原文"],
|
||||
"技术指标参数": ["名称", "功能特性", "能力限制", "额定输出", "环境特性", "重量", "其它指标", "原文"],
|
||||
"布置情况": ["名称", "描述", "原文"],
|
||||
"保障设施保障设备工具信息汇总表": [
|
||||
"名称", "房号", "规格/型号", "生产单位", "通用/专用", "重量",
|
||||
"参考价(元)", "存放位置", "备注", "原文"
|
||||
],
|
||||
"消耗材料信息汇总表": ["名称", "牌号/品级/规格", "生产单位", "备注", "原文"],
|
||||
"维修安全防护及注意事项汇总表": ["名称", "安全防护", "注意事项", "原文"],
|
||||
"图册": ["名称", "描述", "相关系统或设备", "拍摄角度", "原文"],
|
||||
"安全使用操作": ["名称", "内容", "使用前的准备和检查"],
|
||||
"调试": ["名称", "内容", "原文"],
|
||||
"维护保养": ["名称", "内容", "原文"],
|
||||
"常见问题分析及排除": ["名称", "内容", "原文"],
|
||||
"接口": ["名称", "内容", "原文"],
|
||||
"性能指标": ["名称", "内容", "原文"],
|
||||
"操作使用": ["名称", "内容", "原文"],
|
||||
"功能": ["名称", "内容", "原文"],
|
||||
"环境条件": ["名称", "内容", "原文"],
|
||||
"设计单位": ["名称", "内容", "原文"],
|
||||
"操作程序": ["名称", "内容", "原文"],
|
||||
|
||||
}
|
||||
|
||||
# ========== 4. 关系类型列表 ==========
|
||||
RELATIONSHIPS = [
|
||||
"包含", "维修工作中包含故障的描述", "发生故障使用维修工作", "操作时使用",
|
||||
"维修时使用", "指导", "需求", "导致", "描述", "具有","配备","包含故障","实现","适用","配套","调试","支持操作","依赖","安全警告","设计","使用"
|
||||
]
|
||||
|
||||
# ========== 5. 关系类型定义 ==========
|
||||
RELATIONSHIP_TYPES = {
|
||||
"包含": [
|
||||
"舰艇包含系统", "系统包含子系统", "子系统包含设备","设备包含组成"
|
||||
],
|
||||
"维修工作中包含故障的描述": ["维修工作中包含故障的描述"],
|
||||
"维修时使用": ["设备维修时使用维修项目"],
|
||||
"发生故障使用维修工作": ["设备发生故障使用维修工作指导"],
|
||||
"操作时使用": ["操作程序操作时使用操作项目"],
|
||||
"指导": ["安全警告指导维修安全防护及注意事项汇总表", "工作原理指导基本情况"],
|
||||
"导致": ["非正常运行条件导致故障后果"],
|
||||
"需求": [
|
||||
"维修安全防护及注意事项汇总表需求保障设施保障设备工具信息汇总表",
|
||||
"维修安全防护及注意事项汇总表需求消耗材料信息汇总表"
|
||||
],
|
||||
"具有": [
|
||||
"设备具有常见问题分析及排除",
|
||||
"设备具有工作原理","设备具有技术指标参数"
|
||||
],
|
||||
"包含故障" : ["维修工作包含故障"],
|
||||
"配备" : ["设备配备接口"],
|
||||
"实现" : ["设备实现功能"],
|
||||
"适用" : ["设备适用环境条件"],
|
||||
"配套" : ["设备配套图录"],
|
||||
"调试" : ["设备调试"],
|
||||
"安全警告" : ["设备安全警告"],
|
||||
"支持操作" : ["设备支持操作使用"],
|
||||
"依赖" : ["设备依赖器材保障"],
|
||||
"设计" : ["设计单位设计设备"],
|
||||
"使用" : ["操作使用操作程序"]
|
||||
}
|
||||
|
||||
# ========== 6. 便捷函数:获取完整配置 ==========
|
||||
def get_ontology_config():
|
||||
"""返回完整的本体配置字典"""
|
||||
return {
|
||||
"ontology": ONTOLOGY,
|
||||
"ontology_description": ONTOLOGY_DESCRIPTION,
|
||||
"node_types": NODE_TYPES,
|
||||
"relationships": RELATIONSHIPS,
|
||||
"relationship_types": RELATIONSHIP_TYPES
|
||||
}
|
||||
|
||||
def get_ontology_list():
|
||||
"""获取本体类列表"""
|
||||
return ONTOLOGY.copy()
|
||||
|
||||
def get_ontology_description():
|
||||
"""获取本体描述字典"""
|
||||
return ONTOLOGY_DESCRIPTION.copy()
|
||||
|
||||
def get_node_types():
|
||||
"""获取节点类型定义"""
|
||||
return NODE_TYPES.copy()
|
||||
|
||||
def get_relationships():
|
||||
"""获取关系类型列表"""
|
||||
return RELATIONSHIPS.copy()
|
||||
|
||||
def get_relationship_types():
|
||||
"""获取关系类型定义"""
|
||||
return RELATIONSHIP_TYPES.copy()
|
||||
|
||||
|
||||
263
doc2pdf.py
Normal file
263
doc2pdf.py
Normal file
@ -0,0 +1,263 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
文档转 PDF 工具
|
||||
使用 LibreOffice 将各种文档格式转换为 PDF
|
||||
|
||||
支持格式:
|
||||
- Word: .doc, .docx, .rtf
|
||||
- Excel: .xls, .xlsx
|
||||
- PowerPoint: .ppt, .pptx
|
||||
- OpenDocument: .odt, .ods, .odp
|
||||
- 文本: .txt, .csv
|
||||
- 图片: .jpg, .png (会嵌入PDF)
|
||||
"""
|
||||
|
||||
import subprocess
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional, List
|
||||
import concurrent.futures
|
||||
|
||||
|
||||
class Doc2PDF:
|
||||
"""文档转 PDF 转换器"""
|
||||
|
||||
SUPPORTED_FORMATS = {
|
||||
# 文档
|
||||
'.doc', '.docx', '.rtf', '.odt', '.txt',
|
||||
# 表格
|
||||
'.xls', '.xlsx', '.ods', '.csv',
|
||||
# 演示文稿
|
||||
'.ppt', '.pptx', '.odp',
|
||||
# 图片
|
||||
'.jpg', '.jpeg', '.png', '.bmp', '.gif',
|
||||
# 其他
|
||||
'.html', '.htm', '.xml'
|
||||
}
|
||||
|
||||
def __init__(self, libreoffice_path: str = 'libreoffice'):
|
||||
self.libreoffice_path = libreoffice_path
|
||||
self._verify_libreoffice()
|
||||
|
||||
def _verify_libreoffice(self):
|
||||
"""验证 LibreOffice 是否可用"""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[self.libreoffice_path, '--version'],
|
||||
capture_output=True, text=True
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError("LibreOffice 不可用")
|
||||
print(f"✓ 检测到 {result.stdout.strip().split(chr(10))[0]}")
|
||||
except FileNotFoundError:
|
||||
raise RuntimeError("未找到 LibreOffice,请先安装")
|
||||
|
||||
def convert(
|
||||
self,
|
||||
input_file: str,
|
||||
output_dir: Optional[str] = None,
|
||||
timeout: int = 120
|
||||
) -> str:
|
||||
"""
|
||||
转换单个文件为 PDF
|
||||
|
||||
Args:
|
||||
input_file: 输入文件路径
|
||||
output_dir: 输出目录(默认为输入文件所在目录)
|
||||
timeout: 超时时间(秒)
|
||||
|
||||
Returns:
|
||||
生成的 PDF 文件路径
|
||||
"""
|
||||
input_path = Path(input_file).resolve()
|
||||
|
||||
if not input_path.exists():
|
||||
raise FileNotFoundError(f"文件不存在: {input_file}")
|
||||
|
||||
suffix = input_path.suffix.lower()
|
||||
if suffix not in self.SUPPORTED_FORMATS:
|
||||
raise ValueError(f"不支持的格式: {suffix}")
|
||||
|
||||
if output_dir is None:
|
||||
output_dir = str(input_path.parent)
|
||||
else:
|
||||
output_dir = str(Path(output_dir).resolve())
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
cmd = [
|
||||
self.libreoffice_path,
|
||||
'--headless',
|
||||
'--invisible',
|
||||
'--nologo',
|
||||
'--nofirststartwizard',
|
||||
'--convert-to', 'pdf',
|
||||
'--outdir', output_dir,
|
||||
str(input_path)
|
||||
]
|
||||
|
||||
print(f" 转换中: {input_path.name} ...", end=' ', flush=True)
|
||||
|
||||
try:
|
||||
result = subprocess.run(
|
||||
cmd,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=timeout
|
||||
)
|
||||
|
||||
if result.returncode != 0:
|
||||
print("✗")
|
||||
raise RuntimeError(f"转换失败: {result.stderr}")
|
||||
|
||||
output_file = os.path.join(output_dir, f"{input_path.stem}.pdf")
|
||||
|
||||
if os.path.exists(output_file):
|
||||
print("✓")
|
||||
return output_file
|
||||
else:
|
||||
print("✗")
|
||||
raise RuntimeError("转换完成但未找到输出文件")
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
print("✗ (超时)")
|
||||
raise RuntimeError(f"转换超时 ({timeout}秒)")
|
||||
|
||||
def convert_batch(
|
||||
self,
|
||||
input_files: List[str],
|
||||
output_dir: Optional[str] = None,
|
||||
max_workers: int = 1 # LibreOffice 并发可能有问题,建议设为1
|
||||
) -> List[dict]:
|
||||
"""
|
||||
批量转换文件
|
||||
|
||||
Args:
|
||||
input_files: 输入文件列表
|
||||
output_dir: 输出目录
|
||||
max_workers: 并发数
|
||||
|
||||
Returns:
|
||||
转换结果列表
|
||||
"""
|
||||
results = []
|
||||
|
||||
print(f"\n开始批量转换 {len(input_files)} 个文件...\n")
|
||||
|
||||
for i, input_file in enumerate(input_files, 1):
|
||||
print(f"[{i}/{len(input_files)}]", end='')
|
||||
try:
|
||||
output_file = self.convert(input_file, output_dir)
|
||||
results.append({
|
||||
'input': input_file,
|
||||
'output': output_file,
|
||||
'success': True,
|
||||
'error': None
|
||||
})
|
||||
except Exception as e:
|
||||
results.append({
|
||||
'input': input_file,
|
||||
'output': None,
|
||||
'success': False,
|
||||
'error': str(e)
|
||||
})
|
||||
|
||||
# 打印统计
|
||||
success_count = sum(1 for r in results if r['success'])
|
||||
print(f"\n转换完成: {success_count}/{len(results)} 成功")
|
||||
|
||||
return results
|
||||
|
||||
def convert_directory(
|
||||
self,
|
||||
input_dir: str,
|
||||
output_dir: Optional[str] = None,
|
||||
recursive: bool = False
|
||||
) -> List[dict]:
|
||||
"""
|
||||
转换目录中的所有支持文件
|
||||
|
||||
Args:
|
||||
input_dir: 输入目录
|
||||
output_dir: 输出目录
|
||||
recursive: 是否递归处理子目录
|
||||
"""
|
||||
input_path = Path(input_dir)
|
||||
|
||||
if not input_path.is_dir():
|
||||
raise NotADirectoryError(f"不是目录: {input_dir}")
|
||||
|
||||
# 收集所有支持的文件
|
||||
files = []
|
||||
pattern = '**/*' if recursive else '*'
|
||||
|
||||
for file_path in input_path.glob(pattern):
|
||||
if file_path.is_file() and file_path.suffix.lower() in self.SUPPORTED_FORMATS:
|
||||
files.append(str(file_path))
|
||||
|
||||
if not files:
|
||||
print("未找到支持的文件")
|
||||
return []
|
||||
|
||||
return self.convert_batch(files, output_dir)
|
||||
|
||||
|
||||
def main():
|
||||
"""命令行入口"""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
description='使用 LibreOffice 将文档转换为 PDF',
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog='''
|
||||
示例:
|
||||
python doc2pdf.py document.docx
|
||||
python doc2pdf.py document.xlsx -o ./output
|
||||
python doc2pdf.py ./documents/ -o ./pdfs -r
|
||||
'''
|
||||
)
|
||||
|
||||
parser.add_argument('input', help='输入文件或目录')
|
||||
parser.add_argument('-o', '--output', help='输出目录')
|
||||
parser.add_argument('-r', '--recursive', action='store_true',
|
||||
help='递归处理子目录')
|
||||
parser.add_argument('-t', '--timeout', type=int, default=120,
|
||||
help='单个文件转换超时时间(秒)')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
converter = Doc2PDF()
|
||||
|
||||
input_path = Path(args.input)
|
||||
|
||||
if input_path.is_file():
|
||||
# 单文件转换
|
||||
try:
|
||||
output = converter.convert(args.input, args.output, args.timeout)
|
||||
print(f"\n输出文件: {output}")
|
||||
except Exception as e:
|
||||
print(f"\n错误: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
elif input_path.is_dir():
|
||||
# 目录转换
|
||||
results = converter.convert_directory(
|
||||
args.input,
|
||||
args.output,
|
||||
args.recursive
|
||||
)
|
||||
|
||||
# 显示失败的文件
|
||||
failures = [r for r in results if not r['success']]
|
||||
if failures:
|
||||
print("\n失败的文件:")
|
||||
for f in failures:
|
||||
print(f" - {f['input']}: {f['error']}")
|
||||
|
||||
else:
|
||||
print(f"错误: 路径不存在 - {args.input}")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
BIN
docker-images/kgrag_arm64.docker-save.tar.gz
(Stored with Git LFS)
Normal file
BIN
docker-images/kgrag_arm64.docker-save.tar.gz
(Stored with Git LFS)
Normal file
Binary file not shown.
129
entity_registry.json
Normal file
129
entity_registry.json
Normal file
@ -0,0 +1,129 @@
|
||||
{
|
||||
"安全警告": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"概述",
|
||||
"警告",
|
||||
"注意",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"基本情况": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"主要功能",
|
||||
"工作条件",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"组成": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"工作原理": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"技术指标": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"接口关系": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"布置情况": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"器材保障": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"保障设施、保障设备(工具)信息汇总表",
|
||||
"消耗材料信息汇总表",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"维修安全防护及注意事项汇总表": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"概述",
|
||||
"安全防护",
|
||||
"注意事项",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"接口情况": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"操作使用": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"操作使用项目索引表",
|
||||
"操作程序",
|
||||
"战位操作表",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"使用期间安全防护及注意事项": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"安全防护",
|
||||
"注意事项",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"技术指标参数": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
},
|
||||
"接口": {
|
||||
"description": "",
|
||||
"properties": [
|
||||
"名称",
|
||||
"内容",
|
||||
"knowledge_source"
|
||||
]
|
||||
}
|
||||
}
|
||||
74
entity_registry.py
Normal file
74
entity_registry.py
Normal file
@ -0,0 +1,74 @@
|
||||
"""
|
||||
entity_registry.py
|
||||
实体类型注册中心。
|
||||
|
||||
- extract_filtertext_v8.py 在生成实体时调用 register() 写入
|
||||
- app.py 启动时调用 get_ontology_description() / get_node_types() 合并进配置
|
||||
- 注册信息持久化到 entity_registry.json,跨进程保留
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
_REGISTRY_DIR = os.path.join(os.path.dirname(__file__), "registry")
|
||||
_REGISTRY_FILE = os.path.join(_REGISTRY_DIR, "entity_registry.json")
|
||||
|
||||
os.makedirs(_REGISTRY_DIR, exist_ok=True)
|
||||
|
||||
# 内存缓存,避免每次读文件
|
||||
_cache: dict = {}
|
||||
|
||||
|
||||
def _load():
|
||||
global _cache
|
||||
if _cache:
|
||||
return
|
||||
if os.path.exists(_REGISTRY_FILE):
|
||||
with open(_REGISTRY_FILE, encoding="utf-8") as f:
|
||||
_cache = json.load(f)
|
||||
else:
|
||||
_cache = {}
|
||||
|
||||
|
||||
def _save():
|
||||
with open(_REGISTRY_FILE, "w", encoding="utf-8") as f:
|
||||
json.dump(_cache, f, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
def register(entity_type: str, description: str = "", properties: list = None):
|
||||
"""
|
||||
注册一个实体类型。已存在则跳过,不覆盖。
|
||||
|
||||
参数:
|
||||
entity_type: 实体类型名,如 "维修工作指导"
|
||||
description: 对该类型的描述(可为空)
|
||||
properties: 该类型的属性列表,如 ["名称", "内容", "knowledge_source"]
|
||||
"""
|
||||
_load()
|
||||
if entity_type in _cache:
|
||||
return # 已注册,不覆盖
|
||||
_cache[entity_type] = {
|
||||
"description": description,
|
||||
"properties": properties or ["名称", "内容", "knowledge_source"]
|
||||
}
|
||||
_save()
|
||||
print(f" [registry] 注册新实体类型: 【{entity_type}】")
|
||||
|
||||
|
||||
def get_ontology_description() -> dict:
|
||||
"""返回 {type: [description]} 格式,可直接合并进 Ontology_DESCRIPTION。"""
|
||||
_load()
|
||||
return {k: [v["description"]] for k, v in _cache.items()}
|
||||
|
||||
|
||||
def get_node_types() -> dict:
|
||||
"""返回 {type: [properties]} 格式,可直接合并进 node_types。"""
|
||||
_load()
|
||||
return {k: v["properties"] for k, v in _cache.items()}
|
||||
|
||||
|
||||
def get_ontology_list() -> list:
|
||||
"""返回所有已注册的实体类型名列表,可合并进 ontology。"""
|
||||
_load()
|
||||
return list(_cache.keys())
|
||||
|
||||
658
excel_process.py
Normal file
658
excel_process.py
Normal file
@ -0,0 +1,658 @@
|
||||
|
||||
from collections import defaultdict
|
||||
import re
|
||||
from typing import Optional
|
||||
import traceback
|
||||
import logging
|
||||
import pandas as pd
|
||||
from openpyxl import load_workbook
|
||||
import json
|
||||
import os
|
||||
# import xlrd
|
||||
# =====================================================
|
||||
# 行类型判断
|
||||
# =====================================================
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def is_data_row(row):
|
||||
"""
|
||||
判断是否为数据行
|
||||
数据行的特征:第1列包含数字
|
||||
|
||||
Args:
|
||||
row (dict): 包含行数据的字典,格式为 {"content": {"列号": "值", ...}}
|
||||
|
||||
Returns:
|
||||
bool: 如果是数据行返回True,否则返回False
|
||||
"""
|
||||
try:
|
||||
v = row["content"].get("1", "")
|
||||
result = isinstance(v, str) and v.strip().isdigit()
|
||||
logger.debug(f"is_data_row: value='{v}', result={result}")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error in is_data_row: {e}, row: {row}")
|
||||
return False
|
||||
|
||||
|
||||
def is_field_row(row):
|
||||
"""
|
||||
判断是否为字段行
|
||||
字段行的特征:至少有5个值,且超过90%的值不是数字,根据环境处理
|
||||
|
||||
Args:
|
||||
row (dict): 包含行数据的字典
|
||||
|
||||
Returns:
|
||||
bool: 如果是字段行返回True,否则返回False
|
||||
"""
|
||||
try:
|
||||
values = list(row["content"].values())
|
||||
if len(values) < 5:
|
||||
logger.debug(f"is_field_row: less than 5 values, returning False")
|
||||
return False
|
||||
non_digit = sum(1 for v in values if isinstance(v, str) and not v.strip().isdigit())
|
||||
result = non_digit / len(values) > 0.9
|
||||
logger.debug(f"is_field_row: non_digit_ratio={non_digit / len(values)}, result={result}")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error in is_field_row: {e}, row: {row}")
|
||||
return False
|
||||
|
||||
|
||||
def is_header_row(row):
|
||||
"""
|
||||
判断是否为表头行
|
||||
表头行的特征:列数不超过4个,且不是字段行或数据行,根据环境处理
|
||||
|
||||
Args:
|
||||
row (dict): 包含行数据的字典
|
||||
|
||||
Returns:
|
||||
bool: 如果是表头行返回True,否则返回False
|
||||
"""
|
||||
try:
|
||||
content_len = len(row["content"])
|
||||
field_result = is_field_row(row)
|
||||
data_result = is_data_row(row)
|
||||
|
||||
result = content_len <= 4 and not field_result and not data_result
|
||||
logger.debug(f"is_header_row: content_len={content_len}, is_field={field_result}, is_data={data_result}, result={result}")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error in is_header_row: {e}, row: {row}")
|
||||
return False
|
||||
|
||||
# =====================================================
|
||||
# 多级表头
|
||||
# =====================================================
|
||||
|
||||
def collect_field_rows(rows, start_idx):
|
||||
"""
|
||||
从指定索引开始收集连续的字段行
|
||||
|
||||
Args:
|
||||
rows (list): 行数据列表
|
||||
start_idx (int): 开始收集的索引
|
||||
|
||||
Returns:
|
||||
tuple: (字段行列表, 下一个非字段行的索引)
|
||||
"""
|
||||
try:
|
||||
field_rows = []
|
||||
i = start_idx
|
||||
initial_i = start_idx
|
||||
while i < len(rows) and is_field_row(rows[i]):
|
||||
field_rows.append(rows[i]["content"])
|
||||
i += 1
|
||||
logger.info(f"collect_field_rows: collected {len(field_rows)} field rows from index {initial_i} to {i}")
|
||||
return field_rows, i
|
||||
except Exception as e:
|
||||
logger.error(f"Error in collect_field_rows: {e}")
|
||||
return [], start_idx
|
||||
|
||||
def dedup_path(path):
|
||||
"""
|
||||
去除路径中的重复元素,保持顺序
|
||||
|
||||
Args:
|
||||
path (list): 路径列表
|
||||
|
||||
Returns:
|
||||
list: 去重后的路径列表
|
||||
"""
|
||||
try:
|
||||
new_path = []
|
||||
for p in path:
|
||||
if p and (not new_path or new_path[-1] != p):
|
||||
new_path.append(p)
|
||||
logger.debug(f"dedup_path: original={path}, deduplicated={new_path}")
|
||||
return new_path
|
||||
except Exception as e:
|
||||
logger.error(f"Error in dedup_path: {e}, path: {path}")
|
||||
return []
|
||||
|
||||
def build_field_paths(field_rows):
|
||||
"""
|
||||
构建字段路径映射,将列号映射到字段路径
|
||||
|
||||
Args:
|
||||
field_rows (list): 字段行列表
|
||||
|
||||
Returns:
|
||||
defaultdict: 列号到字段路径的映射
|
||||
"""
|
||||
try:
|
||||
field_paths = defaultdict(list)
|
||||
for row_idx, row in enumerate(field_rows):
|
||||
for col, name in row.items():
|
||||
if name:
|
||||
field_paths[col].append(name)
|
||||
|
||||
for col in field_paths:
|
||||
field_paths[col] = dedup_path(field_paths[col])
|
||||
|
||||
logger.info(f"build_field_paths: processed {len(field_paths)} columns from {len(field_rows)} rows")
|
||||
return field_paths
|
||||
except Exception as e:
|
||||
logger.error(f"Error in build_field_paths: {e}")
|
||||
return defaultdict(list)
|
||||
|
||||
|
||||
# =====================================================
|
||||
# 值 & 结构处理
|
||||
# =====================================================
|
||||
|
||||
def normalize_value(value):
|
||||
"""
|
||||
标准化值,将特殊值转换为对应格式
|
||||
|
||||
Args:
|
||||
value: 待处理的值
|
||||
|
||||
Returns:
|
||||
处理后的值
|
||||
"""
|
||||
try:
|
||||
if value == "√":
|
||||
logger.debug(f"normalize_value: converting '√' to True")
|
||||
return True
|
||||
if value in ("", None):
|
||||
logger.debug(f"normalize_value: empty value converted to empty string")
|
||||
return ""
|
||||
logger.debug(f"normalize_value: value={value} unchanged")
|
||||
return value
|
||||
except Exception as e:
|
||||
logger.error(f"Error in normalize_value: {e}, value: {value}")
|
||||
return ""
|
||||
|
||||
|
||||
def insert_by_path(target, path, value):
|
||||
"""
|
||||
根据路径在目标对象中插入值
|
||||
|
||||
Args:
|
||||
target (dict): 目标字典
|
||||
path (list): 插入路径
|
||||
value: 待插入的值
|
||||
"""
|
||||
try:
|
||||
if not path:
|
||||
logger.warning("insert_by_path: empty path provided")
|
||||
return
|
||||
|
||||
cur = target
|
||||
for key in path[:-1]:
|
||||
cur = cur.setdefault(key, {})
|
||||
cur[path[-1]] = value
|
||||
logger.debug(f"insert_by_path: inserted value '{value}' at path {path}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error in insert_by_path: {e}, target: {target}, path: {path}, value: {value}")
|
||||
|
||||
|
||||
def flatten_singleton_dict(obj):
|
||||
"""
|
||||
扁平化单元素字典,如果字典只有一个键且键名与值中的键名相同,则合并
|
||||
|
||||
Args:
|
||||
obj: 待处理的对象
|
||||
|
||||
Returns:
|
||||
处理后的对象
|
||||
"""
|
||||
try:
|
||||
if not isinstance(obj, dict):
|
||||
return obj
|
||||
|
||||
for k in list(obj.keys()):
|
||||
v = obj[k]
|
||||
if isinstance(v, dict):
|
||||
v = flatten_singleton_dict(v)
|
||||
if len(v) == 1 and k in v:
|
||||
obj[k] = v[k]
|
||||
logger.debug(f"flatten_singleton_dict: flattened singleton dict for key '{k}'")
|
||||
else:
|
||||
obj[k] = v
|
||||
return obj
|
||||
except Exception as e:
|
||||
logger.error(f"Error in flatten_singleton_dict: {e}, obj: {obj}")
|
||||
return obj
|
||||
|
||||
|
||||
# def build_content(field_paths, row_content):
|
||||
|
||||
# try:
|
||||
# result = {}
|
||||
# for col, path in field_paths.items():
|
||||
# value = normalize_value(row_content.get(col, ""))
|
||||
# insert_by_path(result, path, value)
|
||||
|
||||
# flatten_singleton_dict(result)
|
||||
# logger.debug(f"build_content: built content with {len(result)} top-level keys")
|
||||
# return result
|
||||
# except Exception as e:
|
||||
# logger.error(f"Error in build_content: {e}, field_paths: {field_paths}, row_content: {row_content}")
|
||||
# return {}
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
def flatten_dict(d, parent_key='', sep='_'):
|
||||
"""
|
||||
将 {"A": {"B": "1"}} 转换为 {"A_B": "1"}
|
||||
"""
|
||||
items = []
|
||||
for k, v in d.items():
|
||||
new_key = f"{parent_key}{sep}{k}" if parent_key else k
|
||||
if isinstance(v, dict):
|
||||
# 递归展开
|
||||
items.extend(flatten_dict(v, new_key, sep=sep).items())
|
||||
else:
|
||||
items.append((new_key, v))
|
||||
return dict(items)
|
||||
|
||||
def build_content(field_paths, row_content):
|
||||
"""
|
||||
根据字段路径构建内容结构
|
||||
|
||||
Args:
|
||||
field_paths (dict): 字段路径映射
|
||||
row_content (dict): 行内容数据
|
||||
|
||||
Returns:
|
||||
dict: 构建的内容结构
|
||||
"""
|
||||
result = {}
|
||||
|
||||
for col, path in field_paths.items():
|
||||
value = normalize_value(row_content.get(col, ""))
|
||||
insert_by_path(result, path, value)
|
||||
|
||||
# 1. 先执行你原有的单级压缩 (解决 A -> {A: x})
|
||||
flatten_singleton_dict(result)
|
||||
|
||||
# 2. 执行深度展平 (将嵌套字典转为 "键_子键" 格式)
|
||||
# 我们遍历第一层,如果值是字典,则展开它
|
||||
final_result = {}
|
||||
for k, v in result.items():
|
||||
if isinstance(v, dict):
|
||||
# 展开子字典,并带上父级的键名作为前缀
|
||||
flattened_sub = flatten_dict(v, parent_key=k, sep='_')
|
||||
final_result.update(flattened_sub)
|
||||
else:
|
||||
final_result[k] = v
|
||||
|
||||
return final_result
|
||||
|
||||
|
||||
# =====================================================
|
||||
# 主解析逻辑
|
||||
# =====================================================
|
||||
|
||||
def transform_excel_json(rows):
|
||||
"""
|
||||
将Excel行数据转换为JSON格式
|
||||
|
||||
Args:
|
||||
rows (list): 行数据列表
|
||||
|
||||
Returns:
|
||||
list: 转换后的JSON数据列表
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Starting transform_excel_json with {len(rows)} rows")
|
||||
output = []
|
||||
i = 0
|
||||
processed_count = 0
|
||||
|
||||
while i < len(rows):
|
||||
row = rows[i]
|
||||
row_type = None
|
||||
|
||||
if is_header_row(row):
|
||||
output.append({"header": row["content"]})
|
||||
row_type = "header"
|
||||
i += 1
|
||||
elif is_field_row(row):
|
||||
field_rows, next_idx = collect_field_rows(rows, i)
|
||||
field_paths = build_field_paths(field_rows)
|
||||
i = next_idx
|
||||
row_type = "field"
|
||||
|
||||
while i < len(rows) and is_data_row(rows[i]):
|
||||
content = build_content(field_paths, rows[i]["content"])
|
||||
output.append({"content": content})
|
||||
i += 1
|
||||
processed_count += 1
|
||||
else:
|
||||
i += 1
|
||||
continue # Skip non-header, non-field, non-data rows
|
||||
|
||||
logger.debug(f"Processed row {i-1}, type: {row_type}")
|
||||
|
||||
logger.info(f"Completed transform_excel_json: {len(output)} items, {processed_count} data rows processed")
|
||||
return output
|
||||
except Exception as e:
|
||||
logger.error(f"Error in transform_excel_json: {e}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise
|
||||
|
||||
|
||||
# =====================================================
|
||||
# Excel 读取(含合并单元格)
|
||||
# =====================================================
|
||||
|
||||
def read_excel_with_merged_cells(file_path: str, sheet_name: Optional[str] = None):
|
||||
"""
|
||||
读取Excel文件,包括合并单元格的处理
|
||||
|
||||
Args:
|
||||
file_path (str): Excel文件路径
|
||||
sheet_name (Optional[str]): 工作表名称,如果为None则使用默认表
|
||||
|
||||
Returns:
|
||||
pd.DataFrame: 包含Excel数据的DataFrame
|
||||
"""
|
||||
try:
|
||||
_, ext = os.path.splitext(file_path)
|
||||
ext = ext.lower()
|
||||
if ext == ".xlsx":
|
||||
logger.info(f"Reading Excel file: {file_path}, sheet: {sheet_name}")
|
||||
wb = load_workbook(file_path, data_only=True)
|
||||
ws = wb[sheet_name] if sheet_name else wb.active
|
||||
|
||||
max_row, max_col = ws.max_row, ws.max_column
|
||||
logger.info(f"Excel sheet dimensions: {max_row} rows x {max_col} columns")
|
||||
|
||||
data = [["" for _ in range(max_col)] for _ in range(max_row)]
|
||||
|
||||
for r in range(1, max_row + 1):
|
||||
for c in range(1, max_col + 1):
|
||||
val = ws.cell(row=r, column=c).value
|
||||
data[r - 1][c - 1] = "" if val is None else str(val).strip()
|
||||
|
||||
# 处理合并单元格
|
||||
merged_count = 0
|
||||
for merged in ws.merged_cells.ranges:
|
||||
tl = data[merged.min_row - 1][merged.min_col - 1]
|
||||
for r in range(merged.min_row - 1, merged.max_row):
|
||||
for c in range(merged.min_col - 1, merged.max_col):
|
||||
data[r][c] = tl
|
||||
merged_count += 1
|
||||
|
||||
logger.info(f"Processed {merged_count} merged cell ranges")
|
||||
df = pd.DataFrame(data)
|
||||
logger.info(f"Created DataFrame with shape {df.shape}")
|
||||
return df
|
||||
|
||||
else:
|
||||
raise ValueError("Unsupported file format")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in read_excel_with_merged_cells: {e}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise
|
||||
|
||||
|
||||
|
||||
def dataframe_to_row_json(df: pd.DataFrame):
|
||||
"""
|
||||
将DataFrame转换为行JSON格式
|
||||
|
||||
Args:
|
||||
df (pd.DataFrame): 源DataFrame
|
||||
|
||||
Returns:
|
||||
list: 行JSON格式的列表
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Converting DataFrame to row JSON, shape: {df.shape}")
|
||||
rows = []
|
||||
for idx, row in df.iterrows():
|
||||
content = {
|
||||
str(col_idx + 1): val
|
||||
for col_idx, val in enumerate(row)
|
||||
if isinstance(val, str) and val != ""
|
||||
}
|
||||
if content:
|
||||
rows.append({
|
||||
"row_index": idx + 1,
|
||||
"content": content
|
||||
})
|
||||
|
||||
logger.info(f"Converted DataFrame to {len(rows)} JSON rows")
|
||||
return rows
|
||||
except Exception as e:
|
||||
logger.error(f"Error in dataframe_to_row_json: {e}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
# =====================================================
|
||||
# Markdown 转换
|
||||
# =====================================================
|
||||
|
||||
# def json_to_markdown(data):
|
||||
# markdown_output = []
|
||||
|
||||
# for item in data:
|
||||
# # 处理标题行 (Header)
|
||||
# if "header" in item:
|
||||
# # 提取 header 中的描述信息
|
||||
# header_vals = item["header"].values()
|
||||
# title = " ".join([v for v in header_vals if ";" in v or v.startswith("CPP") or "清单" in v])
|
||||
# # 如果 header 中包含设备名称和型号,进行拼接
|
||||
# if not title:
|
||||
# title = " ".join(header_vals)
|
||||
|
||||
# markdown_output.append(f"\n## {title.replace(';', ' ')}\n")
|
||||
|
||||
# # 处理内容行 (Content)
|
||||
# elif "content" in item:
|
||||
# content = item["content"]
|
||||
# # 如果是该表格的第一行,生成表头
|
||||
# if not any("|" in line for line in markdown_output[-2:]):
|
||||
# columns = list(content.keys())
|
||||
# header_row = "| " + " | ".join(columns) + " |"
|
||||
# separator_row = "| " + " | ".join(["---"] * len(columns)) + " |"
|
||||
# markdown_output.append(header_row)
|
||||
# markdown_output.append(separator_row)
|
||||
|
||||
# # 生成数据行
|
||||
# values = [str(v) if v is not None else "" for v in content.values()]
|
||||
# # 处理布尔值显示
|
||||
# values = ["是" if v == "True" else "否" if v == "False" else v for v in values]
|
||||
# markdown_output.append("| " + " | ".join(values) + " |")
|
||||
|
||||
# return "\n".join(markdown_output)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
import json
|
||||
|
||||
def json_to_markdown(json_data):
|
||||
markdown_lines = []
|
||||
current_table_data = []
|
||||
|
||||
for entry in json_data:
|
||||
# 1. 处理标题行 (Header)
|
||||
if "header" in entry:
|
||||
# 如果之前已经有累积的数据,先输出旧表格(防止连续多个表格的情况)
|
||||
if current_table_data:
|
||||
markdown_lines.extend(build_table(current_table_data))
|
||||
current_table_data = []
|
||||
|
||||
# 提取 Header 内容并格式化为标题
|
||||
header_dict = entry["header"]
|
||||
# 按照您的要求,处理设备名称和型号的拼接显示
|
||||
title_parts = [v.replace(';', ' ') for v in header_dict.values()]
|
||||
title = " ".join(title_parts)
|
||||
markdown_lines.append(f"\n## {title}\n")
|
||||
|
||||
# 2. 处理内容行 (Content)
|
||||
elif "content" in entry:
|
||||
current_table_data.append(entry["content"])
|
||||
|
||||
# 3. 输出最后一个表格
|
||||
if current_table_data:
|
||||
markdown_lines.extend(build_table(current_table_data))
|
||||
|
||||
return "\n".join(markdown_lines)
|
||||
|
||||
def build_table(data_list):
|
||||
"""将字典列表转换为 Markdown 表格字符串"""
|
||||
if not data_list:
|
||||
return []
|
||||
|
||||
lines = []
|
||||
# 提取所有键作为表头
|
||||
headers = list(data_list[0].keys())
|
||||
|
||||
# 构造表头行
|
||||
header_line = "| " + " | ".join(headers) + " |"
|
||||
# 构造分割行
|
||||
separator_line = "| " + " | ".join(["---"] * len(headers)) + " |"
|
||||
|
||||
lines.append(header_line)
|
||||
lines.append(separator_line)
|
||||
|
||||
# 构造数据行
|
||||
for item in data_list:
|
||||
row_values = []
|
||||
for h in headers:
|
||||
val = item.get(h, "")
|
||||
# 转换布尔值为中文,处理 None 值
|
||||
if val is True: val = "是"
|
||||
elif val is False or val == "": val = " "
|
||||
row_values.append(str(val).replace("\n", " ")) # 避免数据内换行破坏表格结构
|
||||
|
||||
lines.append("| " + " | ".join(row_values) + " |")
|
||||
|
||||
return lines
|
||||
|
||||
# ---------------------------------------------------------
|
||||
# 执行转换
|
||||
# ---------------------------------------------------------
|
||||
# 假设 source_json 是您提供的那段原始数据
|
||||
try:
|
||||
# 这里的 source_json 代表您贴出的完整 JSON 列表
|
||||
# md_output = json_to_markdown_tables(source_json)
|
||||
# print(md_output)
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"解析错误: {e}")
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
# def flatten_dict(d, parent_key='', sep='(', end_sep=')'):
|
||||
# """
|
||||
# 递归扁平化嵌套字典,生成 '父字段(子字段)(孙字段)' 格式的键。
|
||||
# 支持任意深度嵌套。
|
||||
# """
|
||||
# items = []
|
||||
# for k, v in d.items():
|
||||
# new_key = f"{parent_key}{sep}{k}{end_sep}" if parent_key else k
|
||||
# if isinstance(v, dict):
|
||||
# items.extend(flatten_dict(v, new_key, sep, end_sep).items())
|
||||
# else:
|
||||
# items.append((new_key, v))
|
||||
# return dict(items)
|
||||
|
||||
# def extract_all_content_records(data):
|
||||
# """
|
||||
# 从任意结构的 JSON 数据中提取所有顶层包含 "content" 键的对象。
|
||||
# 假设数据是一个列表,每个元素可能是 header 或 content。
|
||||
# 如果不是列表,也尝试遍历。
|
||||
# """
|
||||
# records = []
|
||||
|
||||
# def recurse(obj):
|
||||
# if isinstance(obj, list):
|
||||
# for item in obj:
|
||||
# recurse(item)
|
||||
# elif isinstance(obj, dict):
|
||||
# if "content" in obj and isinstance(obj["content"], dict):
|
||||
# records.append(obj["content"])
|
||||
# else:
|
||||
# for value in obj.values():
|
||||
# recurse(value)
|
||||
|
||||
# recurse(data)
|
||||
# return records
|
||||
|
||||
# def build_markdown_table(records):
|
||||
# if not records:
|
||||
# return "> 未找到任何 content 记录。\n"
|
||||
|
||||
# flat_records = [flatten_dict(record) for record in records]
|
||||
|
||||
# # 收集所有字段(保持顺序)
|
||||
# all_keys = {}
|
||||
# for fr in flat_records:
|
||||
# for k in fr:
|
||||
# all_keys[k] = None
|
||||
# all_keys = list(all_keys.keys())
|
||||
|
||||
# # 构建表格
|
||||
# header = "| " + " | ".join(all_keys) + " |"
|
||||
# separator = "| " + " | ".join(["---"] * len(all_keys)) + " |"
|
||||
# rows = [
|
||||
# "| " + " | ".join(str(fr.get(k, "")) for k in all_keys) + " |"
|
||||
# for fr in flat_records
|
||||
# ]
|
||||
|
||||
# return "\n".join([header, separator] + rows)
|
||||
|
||||
# def json_to_markdown(json_data):
|
||||
# """
|
||||
# 将任意 JSON 文件中所有 'content' 对象转换为 Markdown 表格。
|
||||
# """
|
||||
|
||||
|
||||
# records = extract_all_content_records(json_data)
|
||||
# md_table = build_markdown_table(records)
|
||||
# return md_table
|
||||
|
||||
|
||||
|
||||
|
||||
521
excel_to_json.py
Normal file
521
excel_to_json.py
Normal file
@ -0,0 +1,521 @@
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
import requests
|
||||
from collections import defaultdict
|
||||
|
||||
import traceback
|
||||
import logging
|
||||
import pandas as pd
|
||||
from openpyxl import load_workbook
|
||||
import json
|
||||
import os
|
||||
# ======================
|
||||
# 日志
|
||||
# ======================
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
import re
|
||||
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def is_data_row(row):
|
||||
"""
|
||||
判断是否为数据行
|
||||
数据行的特征:第1列包含数字
|
||||
|
||||
Args:
|
||||
row (dict): 包含行数据的字典,格式为 {"content": {"列号": "值", ...}}
|
||||
|
||||
Returns:
|
||||
bool: 如果是数据行返回True,否则返回False
|
||||
"""
|
||||
try:
|
||||
v = row["content"].get("1", "")
|
||||
result = isinstance(v, str) and v.strip().isdigit()
|
||||
logger.debug(f"is_data_row: value='{v}', result={result}")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error in is_data_row: {e}, row: {row}")
|
||||
return False
|
||||
|
||||
|
||||
def is_field_row(row):
|
||||
"""
|
||||
判断是否为字段行
|
||||
字段行的特征:至少有5个值,且超过90%的值不是数字,根据环境处理
|
||||
|
||||
Args:
|
||||
row (dict): 包含行数据的字典
|
||||
|
||||
Returns:
|
||||
bool: 如果是字段行返回True,否则返回False
|
||||
"""
|
||||
try:
|
||||
values = list(row["content"].values())
|
||||
if len(values) < 5:
|
||||
logger.debug(f"is_field_row: less than 5 values, returning False")
|
||||
return False
|
||||
non_digit = sum(1 for v in values if isinstance(v, str) and not v.strip().isdigit())
|
||||
result = non_digit / len(values) > 0.9
|
||||
logger.debug(f"is_field_row: non_digit_ratio={non_digit / len(values)}, result={result}")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error in is_field_row: {e}, row: {row}")
|
||||
return False
|
||||
|
||||
|
||||
def is_header_row(row):
|
||||
"""
|
||||
判断是否为表头行
|
||||
表头行的特征:列数不超过4个,且不是字段行或数据行,根据环境处理
|
||||
|
||||
Args:
|
||||
row (dict): 包含行数据的字典
|
||||
|
||||
Returns:
|
||||
bool: 如果是表头行返回True,否则返回False
|
||||
"""
|
||||
try:
|
||||
content_len = len(row["content"])
|
||||
field_result = is_field_row(row)
|
||||
data_result = is_data_row(row)
|
||||
|
||||
result = content_len <= 4 and not field_result and not data_result
|
||||
logger.debug(f"is_header_row: content_len={content_len}, is_field={field_result}, is_data={data_result}, result={result}")
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error in is_header_row: {e}, row: {row}")
|
||||
return False
|
||||
|
||||
# =====================================================
|
||||
# 多级表头
|
||||
# =====================================================
|
||||
|
||||
def collect_field_rows(rows, start_idx):
|
||||
"""
|
||||
从指定索引开始收集连续的字段行
|
||||
|
||||
Args:
|
||||
rows (list): 行数据列表
|
||||
start_idx (int): 开始收集的索引
|
||||
|
||||
Returns:
|
||||
tuple: (字段行列表, 下一个非字段行的索引)
|
||||
"""
|
||||
try:
|
||||
field_rows = []
|
||||
i = start_idx
|
||||
initial_i = start_idx
|
||||
while i < len(rows) and is_field_row(rows[i]):
|
||||
field_rows.append(rows[i]["content"])
|
||||
i += 1
|
||||
logger.info(f"collect_field_rows: collected {len(field_rows)} field rows from index {initial_i} to {i}")
|
||||
return field_rows, i
|
||||
except Exception as e:
|
||||
logger.error(f"Error in collect_field_rows: {e}")
|
||||
return [], start_idx
|
||||
|
||||
def dedup_path(path):
|
||||
"""
|
||||
去除路径中的重复元素,保持顺序
|
||||
|
||||
Args:
|
||||
path (list): 路径列表
|
||||
|
||||
Returns:
|
||||
list: 去重后的路径列表
|
||||
"""
|
||||
try:
|
||||
new_path = []
|
||||
for p in path:
|
||||
if p and (not new_path or new_path[-1] != p):
|
||||
new_path.append(p)
|
||||
logger.debug(f"dedup_path: original={path}, deduplicated={new_path}")
|
||||
return new_path
|
||||
except Exception as e:
|
||||
logger.error(f"Error in dedup_path: {e}, path: {path}")
|
||||
return []
|
||||
|
||||
def build_field_paths(field_rows):
|
||||
"""
|
||||
构建字段路径映射,将列号映射到字段路径
|
||||
|
||||
Args:
|
||||
field_rows (list): 字段行列表
|
||||
|
||||
Returns:
|
||||
defaultdict: 列号到字段路径的映射
|
||||
"""
|
||||
try:
|
||||
field_paths = defaultdict(list)
|
||||
for row_idx, row in enumerate(field_rows):
|
||||
for col, name in row.items():
|
||||
if name:
|
||||
field_paths[col].append(name)
|
||||
|
||||
for col in field_paths:
|
||||
field_paths[col] = dedup_path(field_paths[col])
|
||||
|
||||
logger.info(f"build_field_paths: processed {len(field_paths)} columns from {len(field_rows)} rows")
|
||||
return field_paths
|
||||
except Exception as e:
|
||||
logger.error(f"Error in build_field_paths: {e}")
|
||||
return defaultdict(list)
|
||||
|
||||
|
||||
# =====================================================
|
||||
# 值 & 结构处理
|
||||
# =====================================================
|
||||
|
||||
def normalize_value(value):
|
||||
"""
|
||||
标准化值,将特殊值转换为对应格式
|
||||
|
||||
Args:
|
||||
value: 待处理的值
|
||||
|
||||
Returns:
|
||||
处理后的值
|
||||
"""
|
||||
try:
|
||||
if value == "√":
|
||||
logger.debug(f"normalize_value: converting '√' to True")
|
||||
return True
|
||||
if value in ("", None):
|
||||
logger.debug(f"normalize_value: empty value converted to empty string")
|
||||
return ""
|
||||
logger.debug(f"normalize_value: value={value} unchanged")
|
||||
return value
|
||||
except Exception as e:
|
||||
logger.error(f"Error in normalize_value: {e}, value: {value}")
|
||||
return ""
|
||||
|
||||
|
||||
def insert_by_path(target, path, value):
|
||||
"""
|
||||
根据路径在目标对象中插入值
|
||||
|
||||
Args:
|
||||
target (dict): 目标字典
|
||||
path (list): 插入路径
|
||||
value: 待插入的值
|
||||
"""
|
||||
try:
|
||||
if not path:
|
||||
logger.warning("insert_by_path: empty path provided")
|
||||
return
|
||||
|
||||
cur = target
|
||||
for key in path[:-1]:
|
||||
cur = cur.setdefault(key, {})
|
||||
cur[path[-1]] = value
|
||||
logger.debug(f"insert_by_path: inserted value '{value}' at path {path}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error in insert_by_path: {e}, target: {target}, path: {path}, value: {value}")
|
||||
|
||||
|
||||
def flatten_singleton_dict(obj):
|
||||
"""
|
||||
扁平化单元素字典,如果字典只有一个键且键名与值中的键名相同,则合并
|
||||
|
||||
Args:
|
||||
obj: 待处理的对象
|
||||
|
||||
Returns:
|
||||
处理后的对象
|
||||
"""
|
||||
try:
|
||||
if not isinstance(obj, dict):
|
||||
return obj
|
||||
|
||||
for k in list(obj.keys()):
|
||||
v = obj[k]
|
||||
if isinstance(v, dict):
|
||||
v = flatten_singleton_dict(v)
|
||||
if len(v) == 1 and k in v:
|
||||
obj[k] = v[k]
|
||||
logger.debug(f"flatten_singleton_dict: flattened singleton dict for key '{k}'")
|
||||
else:
|
||||
obj[k] = v
|
||||
return obj
|
||||
except Exception as e:
|
||||
logger.error(f"Error in flatten_singleton_dict: {e}, obj: {obj}")
|
||||
return obj
|
||||
|
||||
|
||||
|
||||
|
||||
def flatten_dict(d, parent_key='', sep='_'):
|
||||
"""
|
||||
将 {"A": {"B": "1"}} 转换为 {"A_B": "1"}
|
||||
"""
|
||||
items = []
|
||||
for k, v in d.items():
|
||||
new_key = f"{parent_key}{sep}{k}" if parent_key else k
|
||||
if isinstance(v, dict):
|
||||
# 递归展开
|
||||
items.extend(flatten_dict(v, new_key, sep=sep).items())
|
||||
else:
|
||||
items.append((new_key, v))
|
||||
return dict(items)
|
||||
|
||||
def build_content(field_paths, row_content):
|
||||
"""
|
||||
根据字段路径构建内容结构
|
||||
|
||||
Args:
|
||||
field_paths (dict): 字段路径映射
|
||||
row_content (dict): 行内容数据
|
||||
|
||||
Returns:
|
||||
dict: 构建的内容结构
|
||||
"""
|
||||
result = {}
|
||||
|
||||
for col, path in field_paths.items():
|
||||
value = normalize_value(row_content.get(col, ""))
|
||||
insert_by_path(result, path, value)
|
||||
|
||||
# 1. 先执行你原有的单级压缩 (解决 A -> {A: x})
|
||||
flatten_singleton_dict(result)
|
||||
|
||||
# 2. 执行深度展平 (将嵌套字典转为 "键_子键" 格式)
|
||||
# 我们遍历第一层,如果值是字典,则展开它
|
||||
final_result = {}
|
||||
for k, v in result.items():
|
||||
if isinstance(v, dict):
|
||||
# 展开子字典,并带上父级的键名作为前缀
|
||||
flattened_sub = flatten_dict(v, parent_key=k, sep='_')
|
||||
final_result.update(flattened_sub)
|
||||
else:
|
||||
final_result[k] = v
|
||||
|
||||
return final_result
|
||||
|
||||
|
||||
# =====================================================
|
||||
# 主解析逻辑
|
||||
# =====================================================
|
||||
|
||||
def transform_excel_json(rows):
|
||||
"""
|
||||
将Excel行数据转换为JSON格式
|
||||
|
||||
Args:
|
||||
rows (list): 行数据列表
|
||||
|
||||
Returns:
|
||||
list: 转换后的JSON数据列表
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Starting transform_excel_json with {len(rows)} rows")
|
||||
output = []
|
||||
i = 0
|
||||
processed_count = 0
|
||||
|
||||
while i < len(rows):
|
||||
row = rows[i]
|
||||
row_type = None
|
||||
|
||||
if is_header_row(row):
|
||||
output.append({"header": row["content"]})
|
||||
row_type = "header"
|
||||
i += 1
|
||||
elif is_field_row(row):
|
||||
field_rows, next_idx = collect_field_rows(rows, i)
|
||||
field_paths = build_field_paths(field_rows)
|
||||
i = next_idx
|
||||
row_type = "field"
|
||||
|
||||
while i < len(rows) and is_data_row(rows[i]):
|
||||
content = build_content(field_paths, rows[i]["content"])
|
||||
output.append({"content": content})
|
||||
i += 1
|
||||
processed_count += 1
|
||||
else:
|
||||
i += 1
|
||||
continue # Skip non-header, non-field, non-data rows
|
||||
|
||||
logger.debug(f"Processed row {i-1}, type: {row_type}")
|
||||
|
||||
logger.info(f"Completed transform_excel_json: {len(output)} items, {processed_count} data rows processed")
|
||||
return output
|
||||
except Exception as e:
|
||||
logger.error(f"Error in transform_excel_json: {e}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise
|
||||
|
||||
|
||||
# =====================================================
|
||||
# Excel 读取(含合并单元格)
|
||||
# =====================================================
|
||||
|
||||
def read_excel_with_merged_cells(file_path: str, sheet_name: Optional[str] = None):
|
||||
"""
|
||||
读取Excel文件,包括合并单元格的处理
|
||||
|
||||
Args:
|
||||
file_path (str): Excel文件路径
|
||||
sheet_name (Optional[str]): 工作表名称,如果为None则使用默认表
|
||||
|
||||
Returns:
|
||||
pd.DataFrame: 包含Excel数据的DataFrame
|
||||
"""
|
||||
try:
|
||||
_, ext = os.path.splitext(file_path)
|
||||
ext = ext.lower()
|
||||
if ext == ".xlsx":
|
||||
logger.info(f"Reading Excel file: {file_path}, sheet: {sheet_name}")
|
||||
wb = load_workbook(file_path, data_only=True)
|
||||
ws = wb[sheet_name] if sheet_name else wb.active
|
||||
|
||||
max_row, max_col = ws.max_row, ws.max_column
|
||||
logger.info(f"Excel sheet dimensions: {max_row} rows x {max_col} columns")
|
||||
|
||||
data = [["" for _ in range(max_col)] for _ in range(max_row)]
|
||||
|
||||
for r in range(1, max_row + 1):
|
||||
for c in range(1, max_col + 1):
|
||||
val = ws.cell(row=r, column=c).value
|
||||
data[r - 1][c - 1] = "" if val is None else str(val).strip()
|
||||
|
||||
# 处理合并单元格
|
||||
merged_count = 0
|
||||
for merged in ws.merged_cells.ranges:
|
||||
tl = data[merged.min_row - 1][merged.min_col - 1]
|
||||
for r in range(merged.min_row - 1, merged.max_row):
|
||||
for c in range(merged.min_col - 1, merged.max_col):
|
||||
data[r][c] = tl
|
||||
merged_count += 1
|
||||
|
||||
logger.info(f"Processed {merged_count} merged cell ranges")
|
||||
df = pd.DataFrame(data)
|
||||
logger.info(f"Created DataFrame with shape {df.shape}")
|
||||
return df
|
||||
|
||||
else:
|
||||
raise ValueError("Unsupported file format")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in read_excel_with_merged_cells: {e}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise
|
||||
|
||||
|
||||
|
||||
def dataframe_to_row_json(df: pd.DataFrame):
|
||||
"""
|
||||
将DataFrame转换为行JSON格式
|
||||
|
||||
Args:
|
||||
df (pd.DataFrame): 源DataFrame
|
||||
|
||||
Returns:
|
||||
list: 行JSON格式的列表
|
||||
"""
|
||||
try:
|
||||
logger.info(f"Converting DataFrame to row JSON, shape: {df.shape}")
|
||||
rows = []
|
||||
for idx, row in df.iterrows():
|
||||
content = {
|
||||
str(col_idx + 1): val
|
||||
for col_idx, val in enumerate(row)
|
||||
if isinstance(val, str) and val != ""
|
||||
}
|
||||
if content:
|
||||
rows.append({
|
||||
"row_index": idx + 1,
|
||||
"content": content
|
||||
})
|
||||
|
||||
logger.info(f"Converted DataFrame to {len(rows)} JSON rows")
|
||||
return rows
|
||||
except Exception as e:
|
||||
logger.error(f"Error in dataframe_to_row_json: {e}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise
|
||||
|
||||
|
||||
def json_to_markdown(json_data):
|
||||
markdown_lines = []
|
||||
current_table_data = []
|
||||
|
||||
for entry in json_data:
|
||||
# 1. 处理标题行 (Header)
|
||||
if "header" in entry:
|
||||
# 如果之前已经有累积的数据,先输出旧表格(防止连续多个表格的情况)
|
||||
if current_table_data:
|
||||
markdown_lines.extend(build_table(current_table_data))
|
||||
current_table_data = []
|
||||
|
||||
# 提取 Header 内容并格式化为标题
|
||||
header_dict = entry["header"]
|
||||
# 按照您的要求,处理设备名称和型号的拼接显示
|
||||
title_parts = [v.replace(';', ' ') for v in header_dict.values()]
|
||||
title = " ".join(title_parts)
|
||||
markdown_lines.append(f"\n## {title}\n")
|
||||
|
||||
# 2. 处理内容行 (Content)
|
||||
elif "content" in entry:
|
||||
current_table_data.append(entry["content"])
|
||||
|
||||
# 3. 输出最后一个表格
|
||||
if current_table_data:
|
||||
markdown_lines.extend(build_table(current_table_data))
|
||||
|
||||
return "\n".join(markdown_lines)
|
||||
|
||||
def build_table(data_list):
|
||||
"""将字典列表转换为 Markdown 表格字符串"""
|
||||
if not data_list:
|
||||
return []
|
||||
|
||||
lines = []
|
||||
# 提取所有键作为表头
|
||||
headers = list(data_list[0].keys())
|
||||
|
||||
# 构造表头行
|
||||
header_line = "| " + " | ".join(headers) + " |"
|
||||
# 构造分割行
|
||||
separator_line = "| " + " | ".join(["---"] * len(headers)) + " |"
|
||||
|
||||
lines.append(header_line)
|
||||
lines.append(separator_line)
|
||||
|
||||
# 构造数据行
|
||||
for item in data_list:
|
||||
row_values = []
|
||||
for h in headers:
|
||||
val = item.get(h, "")
|
||||
# 转换布尔值为中文,处理 None 值
|
||||
if val is True: val = "是"
|
||||
elif val is False or val == "": val = " "
|
||||
row_values.append(str(val).replace("\n", " ")) # 避免数据内换行破坏表格结构
|
||||
|
||||
lines.append("| " + " | ".join(row_values) + " |")
|
||||
|
||||
return lines
|
||||
|
||||
|
||||
def excel_to_jsondata(excel_path: str,sheet_name: Optional[str] = None):
|
||||
|
||||
logger.info(f"Reading Excel: {excel_path}")
|
||||
df = read_excel_with_merged_cells(excel_path, sheet_name)
|
||||
|
||||
rows = dataframe_to_row_json(df)
|
||||
table_json = transform_excel_json(rows)
|
||||
return table_json
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
605
extract_excel_node_relation.py
Normal file
605
extract_excel_node_relation.py
Normal file
@ -0,0 +1,605 @@
|
||||
from kg_build_v2.excel_to_json import excel_to_jsondata
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from typing import Optional,Dict
|
||||
import requests
|
||||
from collections import defaultdict
|
||||
|
||||
import traceback
|
||||
import logging
|
||||
import pandas as pd
|
||||
from openpyxl import load_workbook
|
||||
import json
|
||||
import os
|
||||
from kg_build_v2.Prompt_excel import prompt_beipinbeijian,prompt_type_excel,prompt_guzhanghuizong
|
||||
from openai import OpenAI
|
||||
import re
|
||||
import time
|
||||
|
||||
# ======================
|
||||
# 日志
|
||||
# ======================
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
System_Name_Replace_Dict={
|
||||
"00": {
|
||||
"code": "00",
|
||||
"name": "总体及综合保障",
|
||||
"level": 1,
|
||||
"description":None,
|
||||
"children": None
|
||||
},
|
||||
|
||||
"01": {
|
||||
"code": "01",
|
||||
"name": "船体结构",
|
||||
"level": 1,
|
||||
"description": "由各种结构类型组成的JC船体,包括船体壳体、舱壁、甲板、平台、上层建筑、舷台、专用结构、复合材料结构、基座及船体附体等。",
|
||||
"children": [
|
||||
{ "code": "00", "name": "总论", "level": 2,"description": "通用要求"},
|
||||
{ "code": "10", "name": "主船体壳体结构", "level": 2,"description": "构成水面JC主船体的壳体结构和潜艇的耐压壳体、非耐压壳体结构" },
|
||||
{ "code": "20", "name": "船体舱壁结构", "level": 2,"description": "构成水面JC船体的主横舱壁、主纵舱壁、围壁、轴隧、围井及其他专用舱壁等船体舱壁结构。" },
|
||||
{ "code": "30", "name": "船体甲板及平台结构", "level": 2,"description": "将水面JC主船体分为若干层的连续甲板结构(包括甲板板和骨架)及非连续甲板结构(包括平台板和骨架)。" },
|
||||
{ "code": "40", "name": "上层建筑(或甲板室)结构", "level": 2,"description": "水面JC主船体l甲板以上的各层上层建筑(或甲板室)结构以及潜艇上层建筑、围壳结构等。" },
|
||||
{ "code": "50", "name": "舷台结构", "level": 2,"description": "大型水面JC为扩展甲板面积而在主船体两舷外侧设置的特殊结构,以主船体外板为边界,包括与主船体相对应而又相对独立的多层甲板、舱壁及外板等结构。" },
|
||||
{ "code": "60", "name": "专用结构", "level": 2,"description": "安装在主船体或上层建筑上为特殊用途设置的各种结构。" },
|
||||
{ "code": "70", "name": "复合材料结构", "level": 2,"description": "应用复合材料代替金属材料的船体局部结构或全艇体结构。"},
|
||||
{ "code": "80", "name": "基座结构", "level": 2,"description": "装焊于JC上用于安装和支承各种设备的底座结构。其中包括基座减振的弹性连接件。不包括随设备的弹性连接件、紧固件以及设备或机组的机架"}
|
||||
]
|
||||
},
|
||||
|
||||
"02": {
|
||||
"code": "02",
|
||||
"name": "推进系统",
|
||||
"level": 1,
|
||||
"description":"为JC推进提供动力的系统,包括能源发生系统(含H能和非H能)、动力机组及推进装置、功率传递系统和推进器、推进保障系统、推进监测和控制系统。",
|
||||
"children": [
|
||||
{ "code": "00", "name": "总论", "level": 2,"description": "通用要求" },
|
||||
{ "code": "10", "name": "H能源发生系统", "level": 2,"description": "通过反应堆的可控H反应产生H能,并将其转化为热能的系统。包括反应堆系统、蒸汽发生器、反应堆中间回路系统、反应堆冷却剂系统、反应堆专设安全设施、反应堆保障系统、辐射监测系统及屏蔽等。"},
|
||||
{ "code": "20", "name": "非H能源发生系统", "level": 2,"description": "H能源以外的JC推进用能源发生系统。包括常压锅炉、增压锅炉等。"},
|
||||
{ "code": "30", "name": "动力机组及推进装置", "level": 2,"description": "为JC提供推进动力或为发电机提供原动力的动力机组或装置。包括蒸汽轮机、柴油机、燃气轮机、主电力推进装置、辅助电力推进装置、潜艇应急推进装置、联合推进装置等。"},
|
||||
{ "code": "40", "name": "功率传递系统和推进器", "level": 2,"description": "将推进机组发出的功率传递到推进器推动JCHX的系统及设备。包括齿轮传动装置、离合器和联轴器、轴、轴承和轴封、推进器等。" },
|
||||
{ "code": "50", "name": "推进保障系统", "level": 2,"description": "推进系统中为主锅炉、推进机组、功率传递系统服务的系统、机械设备及附件。包括日用燃油系统、滑油系统、蒸汽系统、凝水和给水系统、循环水和冷却水系统、起动系统、进排气系统、锅炉鼓风系统和排污系统、燃气轮机引气和清洗系统等。" },
|
||||
{ "code": "60", "name": "推进监测和控制系统", "level": 2,"description": "对各种推进动力装置及其保障系统实施手动或自动控制(含集中控制或机旁控制)并进行显示、报警、安全保护等项目" }
|
||||
]
|
||||
},
|
||||
|
||||
"03": {
|
||||
"code": "03",
|
||||
"name": "电力系统",
|
||||
"level": 1,
|
||||
"children": [
|
||||
{ "code": "00", "name": "总论", "level": 2,"description": "通用要求" },
|
||||
{ "code": "10", "name": "供电系统", "level": 2,"description": "电力系统中提供各种电能的系统。包括柴油发电机组、蒸汽轮机发电机组、燃气轮机发电机组、热气机发电机组、电池组及其附属设备、变电设备等。" },
|
||||
{ "code": "20", "name": "配电系统", "level": 2,"description": "电力系统中对电源设备产生的电能进行分配、检测,以向负载供电的系统。包括主配电系统、应急配电系统、事故配电系统、专用配电系统和接岸电设施等。" },
|
||||
{ "code": "30", "name": "照明系统", "level": 2,"description": "由各种灯具、照明配电板、配电和控制电路构成的用于JC照明的系统。包括主照明系统、应急照明系统、低照度照明系统和专用灯照明系统等。" },
|
||||
{ "code": "40", "name": "电气监测和控制系统", "level": 2,"description": "对JC电站、电气传动设备等实施监测和控制的系统。包括电站、电气传动、配电系统、照明系统和电池等的监测和控制。" },
|
||||
{ "code": "50", "name": "发电机组保障系统", "level": 2,"description": "对JC发电机组的正常运行起保障作用的设备和动力管路系统的设计、制造、布置、安装、材料和试验验收等。包括柴油发电机组、蒸汽轮机发电机组、燃气轮机发电机组和热气机发电机组等的保障系统。" },
|
||||
{ "code": "60", "name": "电缆敷设", "level": 2,"description": "包括全船电力电缆、信号电缆、控制电缆的敷设。不包括消磁电缆的敷设。" },
|
||||
{ "code": "70", "name": "综合电缆系统", "level": 2,"description": "JC发电、变电、配电、监控、照明、储能、电气传动及电力推进等部分综合集成的电力系统,它可实现艇电力推进与日常负载之间能量的灵活调配。包括供电系统、配电系统、电力推进系统、能量管理系统及储能系统等。" },
|
||||
{ "code": "80", "name": "电场、磁场防护系统", "level": 2,"description": "通过对JC磁场、电场的监测、控制和管理,确保JC电磁场隐身性能的系统。包括消磁系统和电场防护系统。" },
|
||||
{ "code": "90", "name": "专用项目", "level": 2,"description": "包括避雷装置、日用生活电器、电力系统工作液体、电力系统备品备件和专用工具。" }
|
||||
]
|
||||
},
|
||||
|
||||
"04": {
|
||||
"code": "04",
|
||||
"name": "ZZ系统",
|
||||
"level": 1,
|
||||
"children": [
|
||||
{ "code": "00", "name": "总论", "level": 2,"description": "通用要求" },
|
||||
{ "code": "10", "name": "信息感知", "level": 2,"description": "JC上用于获取目标信息,对信息进行综合分析和处理,概略测定其位置和判定其敌我属性的设备。包括射频前端设备、光频前端设备、声频设备等。" },
|
||||
{ "code": "20", "name": "本JZZ指挥控制", "level": 2,"description": "用于实现单J(艇)ZZ指挥控制功能的系统、设备、模块等的总称。本组按照ZZ任务不同,主要包括ZZ态势生成,综合ZZ指挥控制,对空、对海、对水下、对陆方面ZZ综合控制,两栖ZZ、航空兵ZZ指挥控制,航空管制,无人ZBZZ控制等功能组成。" },
|
||||
{ "code": "30", "name": "编队ZZ指挥", "level": 2,"description": "用于基于情报处理功能,指挥员根据分配的任务对不同任务编队实施ZZ指挥所必需的设备、软件和人员的总称。包括编队ZZ筹划、编队情报保障、编队协同ZZ指挥和编队ZZ保障等。" },
|
||||
{ "code": "40", "name": "WQ控制及保障", "level": 2,"description": "主要包括雷弹发射装置,控制WQ对目标进行GJ的火力控制系统以及雷弹转运与贮存、保障系统和有关设施。" },
|
||||
{ "code": "50", "name": "综合通信", "level": 2,"description": "用于实现JC内部和外部信息传输的各类通信系统以及配套的设备、设施等。具体包括内部通信保障设备、通信控制管理设备、无线通信设备、数据链设备、全J综合视频JS设备、应急遇险救生通信设备等。" },
|
||||
{ "code": "60", "name": "综合导航", "level": 2,"description": "引导和保障JC平台按ZZ要求HX并提供JC运动状态信息支援的JZ系统、设备和器材的统称" },
|
||||
{ "code": "70", "name": "大气海洋环境探测", "level": 2,"description": "探测、测量温度、湿度、气压、风速、海水温度、盐度、密度,以及云、降水、海雾、冰雹、雷电、浪涌等大气海洋环境要素和现象的设备,并对大气海洋环境信息接收、处理及提供相应保障服务的设备。包括大气环境探测设备、海洋环境探测设备、大气海洋环境信息接收设备、大气海洋环境信息处理及服务设备等。" },
|
||||
{ "code": "80", "name": "基础保障", "level": 2,"description": "包括公共计算服务设备、密码BM设备、信息安全设备和对准保障设备等。" },
|
||||
{ "code": "90", "name": "专用项目", "level": 2,"description": "包括电子试验、检测与监控设备、保障管理系统(政工保障和后勤保障)以及ZZ系统备品备件和专用工具等。" }
|
||||
]
|
||||
},
|
||||
|
||||
"05": {
|
||||
"code": "05",
|
||||
"name": "船舶保障系统",
|
||||
"level": 1,
|
||||
"description": "为JC控制、安全、供应和居住所要求的系统,包括舱室大气环境控制、消防系统、日用水系统、HSH防护系统、日用压缩空气和其他流体系统、生活保障系统、环境污染控制系统、损害管制监测与控制系统、综合平台管理系统、JC姿态控制系统等。",
|
||||
"children": [
|
||||
{ "code": "00", "name": "总论", "level": 2,"description": "通用要求" },
|
||||
{ "code": "10", "name": "舱室大气环境控制系统", "level": 2,"description": "指控制舱室温度、湿度和净化空气、改善舱室环境的系统,包括日用蒸汽和供暖系统、通风系统、空气调节系统、空气再生与净化系统、冷藏系统及空气分析系统等。" },
|
||||
{ "code": "20", "name": "消防系统", "level": 2,"description": "包括水消防系统、气体灭火系统、泡沫灭火系统、防爆系统及消防用品等。" },
|
||||
{ "code": "30", "name": "日用水系统", "level": 2,"description": "提供JY生活或设备正常工作的海水和淡水系统,包括饮用水和洗涤水系统、甲板排水系统、日用海水系统、海水淡化装置及其系统、淡水冷却系统、机械和管路的泄水系统及舱底水系统等。" },
|
||||
{ "code": "40", "name": "HSH防护系统", "level": 2,"description": "为保证HSHZZ条件下JY和ZB的ZZ能力而设置的综合防御系统,能够对HWQ、化学WQ、生物WQ等自身、载体及次生造成的危害,进行探测、防护和洗消,包括HSH监控系统、集体防护系统、水幕洗消系统和其他防护器材。" },
|
||||
{ "code": "50", "name": "日用压缩空气和其他流体系统", "level": 2,"description": "包括日用压缩空气系统、其他压缩气体系统、气幕降噪系统、液压系统和专用管路系统等。" },
|
||||
{ "code": "60", "name": "生活保障系统", "level": 2,"description": "为保障JY工作、生活需要而设置的居住、卫生、物品贮藏、膳食保障、休息活动以及医疗保障的舱室和设施,包括居住设施、膳食设施、服务设施、医疗设施和工作舱室设施。本组还包括船用家具和JC上小五金用品的基本要求。" },
|
||||
{ "code": "70", "name": "环境污染控制系统", "level": 2,"description": "为防止海洋污染而设置的对油污水、生活污水、生活垃圾、医疗垃圾、餐厨垃圾、压载水进行处理的综合控制系统,包括油污水处理系统、生活污水处理系统、固体垃圾处理系统、餐厨垃圾处理系统、潜艇废物抛出装置和压载水处理系统。" },
|
||||
{ "code": "80", "name": "损害管制监测与控制系统", "level": 2,"description": "包括对重要舱室的状态、舱室浸水、失火、通风及其他辅机进行集中监测、报警和控制。由损害管制控制台、综合报警器以及各种传感器等组成。" },
|
||||
{ "code": "90", "name": "综合平台管理系统", "level": 2,"description": "ZB于JC,将平台上系统和设备的监测、控制、管理及辅助决策等功能集成为一个整体,形成的监控与管理信息系统,一般包括推进监控设备、电力监控设备、损管综合监控设备、综合J桥设备、综合保障管理系统、状态监测及维护维修设备、综合管理与决策设备以及系统共用基础设施等。" },
|
||||
{ "code": "A0", "name": "JC姿态控制系统", "level": 2,"description": "利用JC上ZB的系统、设备,在JC破损、非破损状态下,为JC、JZ机及J上其他系统、ZB的正常使用提供姿态保障、提供舷侧防护及保证JC基本浮性稳性的系统,一般包括水面JC的纵倾平衡系统、横倾平衡系统、应急排水系统、压载水系统,以及潜艇潜浮系统、均衡系统、深度控制系统和操纵控制系统。" },
|
||||
{ "code": "B0", "name": "专用项目", "level": 2,"description": "包括船舶保障系统工作液体、船舶保障系统备品备件和专用工具。" }
|
||||
]
|
||||
},
|
||||
|
||||
"06": {
|
||||
"code": "06",
|
||||
"name": "航空保障系统",
|
||||
"level": 1,
|
||||
"description": "JC上用于保障各型JZ机实现高效出动、安全回收、塔台飞行指挥、J面调运、油气电液供给、DY贮运加挂、维修任务支援等各类航空作业,保障JZ机安全上J并形成出动回收能力的系统,包括航空保障指挥管理系统,起飞和着J引导系统,JZ机放飞、回收及贮运系统,机载WQ保障系统,喷气燃料系统,航空电源系统,航空供气系统,JZ机辅助保障设施、JZ机维修及任务支援系统等。",
|
||||
"children": [
|
||||
{ "code": "00", "name": "总论", "level": 2,"description": "通用要求" },
|
||||
{ "code": "10", "name": "航空保障指挥管理系统", "level": 2,"description": "为航空保障作业提供塔台飞行指挥、机务勤务综合指挥调度、飞行作训管理航空ZB和航空保障ZB信息管理、综合态势服务等指挥管理功能的系统。包括塔台飞行指挥设施、机务勤务综合指挥调度设施、飞行作训管理设施、航空ZB及资源信息管理设施、航空保障综合态势服务设施、航空保障作业能力评估设施等。" },
|
||||
{ "code": "20", "name": "起飞和着J引导系统", "level": 2,"description": "用于指挥引导JZ机在JC上起飞和着J的系统,其采用光学、惯性、卫星等技术手段为JZ机提供位置、角度及目视引导信息,实现JZ机的安全起飞和着J。包括起飞和着J引导指挥设施、 雷达着J引导设备、光学助降设备、起降电视JS设备、微波着J引导设备、光电着J引导设备、卫星着J引导设备、甲板灯光设施等。" },
|
||||
{ "code": "30", "name": "JZ机放飞系统", "level": 2,"description": "JC上保障JZ机起飞的系统。包括电磁弹射装置、飞机止动装置、喷气偏流板装置等。" },
|
||||
{ "code": "40", "name": "JZ机回收系统", "level": 2,"description": "JC上保障JZ机回收的系统。包括电磁阻拦装置、液压阻拦装置、阻拦网设施、格栅装置、集成抓曳装置等。" },
|
||||
{ "code": "50", "name": "JZ机贮运系统", "level": 2,"description": "用于JZ机在JC上贮存和转运的系统。包括贮运指挥设施、JZ机升降机、机库大门、 牵引设施、系留设施、调向转盘等。" },
|
||||
{ "code": "60", "name": "机载WQ保障系统", "level": 2,"description": "JC上用于完成机载WQ管理、调度、贮存、转运、装配、维护等保障功能的设备与设施统称。包括机载WQ管理调度设施、机载WQ贮存装置、机载WQ转运设备、机载WQ挂卸设备、机载WQ装配维护设备等。" },
|
||||
{ "code": "70", "name": "喷气燃料系统", "level": 2,"description": "为JZ机提供喷气燃料保障的系统。包括喷气燃料动力设施、喷气燃料监控设施、集成加油装置、喷气燃料安全保障设施、喷气燃料维护保障设施等。" },
|
||||
{ "code": "80", "name": "航空电源系统", "level": 2,"description": "为JZ机提供航空电源保障的系统。包括航空电源供电设施、航空电源辅助保障设备、航空电源监控设施等。" },
|
||||
{ "code": "90", "name": "航空供气系统", "level": 2,"description": "为JZ机提供航空氧气、氮气、压缩空气以及空调通风气体等各种气体保障的系统。包括航空气体制备设施、航空气体储用设施、压缩空气设备、航空TZ空调装置、航空供气监控设施等。" },
|
||||
{ "code": "A0", "name": "JZ机辅助保障设施", "level": 2,"description": "为JZ机提供液压、冲洗、应急救援、人员保障、跑道保障等辅助保障作业的设施。包括液压保障设施、液冷保障设施、冲洗保障设施、JZ机应急救援设施、航空人员专用保障设施、跑道保障设施以及其它保障设施等。" },
|
||||
{ "code": "B0", "name": "JZ机维修及任务支援系统", "level": 2,"description": "为JZ机在JC上的使用、维护、修理、训练以及数据分析、处理等各项作业提供直接保障的软、硬件设备的总称。包括维修保障设备、J基训练系统、远程维修支援系统、任务规划支持系统等。" }
|
||||
]
|
||||
},
|
||||
|
||||
"07": {
|
||||
"code": "07",
|
||||
"name": "TZ保障系统",
|
||||
"level": 1,
|
||||
"description":"JC为了保障履行陆特Z、补给、救捞、医疗卫勤、维修、海洋测量、ZC等TZ任务而配置的系统,包括陆特ZZB保障系统、无人ZB保障系统、海上补给与接收系统、救援和打捞系统、卫生勤务保障系统、维修系统、海洋测量系统、ZC系统等。",
|
||||
"children": [
|
||||
{ "code": "00", "name": "总论", "level": 2,"description": "通用要求" },
|
||||
{ "code": "10", "name": "陆特ZZB保障系统", "level": 2,"description": "为登陆和TZZZZB提供ZZ支援保障的系统。包括陆特ZZB保障任务规划及指挥管理系统、陆特ZZB收放贮运系统、陆特ZZB辅助引导系统、陆特ZZB资源供应保障系统、陆特ZZB维护保障系统、陆特ZZBDY保障系统、陆特ZZB保障作业支持系统等。" },
|
||||
{ "code": "20", "name": "无人ZB保障系统", "level": 2,"description": "为JZ无人ZB(JZ无人机保障见6组)提供ZZ支援保障的系统。包括无人ZB保障指挥调度及驻J管理系统、无人ZB收放贮运系统、无人ZB勤务保障系统、无人ZBDY保障系统、无人ZB辅助保障设施等。" },
|
||||
{ "code": "30", "name": "海上补给与接收系统", "level": 2,"description": "JC海上HX时的补给和接收系统。包括JC之间的横向、纵向、靠帮补给和接收系统、直升机与JC之间的垂直补给系统、JC上的干货、液货存储转运系统,以及补给船的补给作业指挥管理系统。" },
|
||||
{ "code": "40", "name": "救援和打捞系统", "level": 2,"description": "救生JC对水面事故(ZS)JC、ZB及落水人员等进行搜寻救助,减少其损失及伤害;对水下事故潜艇进行搜索探测,为艇员提供生命支持,协助艇员离艇脱险;对沉没于水中的JC、ZB或货物等进行打捞处置的系统。包括救援指挥管理系统、水面救助系统、水下救援系统和打捞系统等。" },
|
||||
{ "code": "50", "name": "卫生勤务保障系统", "level": 2,"description": "指JC上执行本J及海上人员平时诊疗、Z时救治、转运等卫生勤务保障任务的系统。包括医疗信息管理系统、伤病员换乘系统、常规医疗系统、特勤医疗系统、JC防疫与防护设施和医疗保障系统等。" },
|
||||
{ "code": "60", "name": "维修系统", "level": 2,"description": "JC上配置的对目标JC执行修理抢修操作,阻止其损害,恢复部分系统功能和HX能力的保障任务系统。" },
|
||||
{ "code": "70", "name": "海洋测量系统", "level": 2,"description": "对海洋气象水文、潮汐洋流、地形地貌、温盐密度等进行测量记录JS的系统。包括海洋测量指挥管理系统、海洋测量设备及其收放贮运、资源供应、信息保障、维护保障等系统。" },
|
||||
{ "code": "80", "name": "ZC系统", "level": 2,"description": "利用JC上配置的专用通信、雷达、光电、水声等ZC手段对外界目标环境等进行ZC,并对ZC得到的情报进行综合分析处理的系统。包括通信ZC系统、雷达ZC系统、光电ZC系统、水声ZC系统、情报综合处理系统和综合管理系统等。" },
|
||||
{ "code": "90", "name": "专用项目", "level": 2,"description": "TZ保障系统备品备件和专用工具。" }
|
||||
]
|
||||
},
|
||||
|
||||
"08": {
|
||||
"code": "08",
|
||||
"name": "船舶装置",
|
||||
"level": 1,
|
||||
"description":"包括JC运动控制系统,机械式转运系统,门、窗、梯、盖等船体属具,舱室分隔和表面覆盖,防腐蚀与防污损和其他专用项目。",
|
||||
"children": [
|
||||
{ "code": "00", "name": "总论", "level": 2,"description": "通用要求" },
|
||||
{ "code": "10", "name": "JC运动控制系统", "level": 2,"description": "为控制JC运动而设置的系统,包括操舵控制系统及操舵装置、舵、综合J桥设备、减摇系统、动力定位系统、主动转向装置、水翼收放控制装置和艉板可调装置等。" },
|
||||
{ "code": "20", "name": "机械式转运系统", "level": 2,"description": "采用机械设施进行转运、操纵的系统。包括锚及其存放与操纵系统、系泊和拖曳系统、小艇及其收放系统、升降装置及其附属系统、货物转运系统等。" },
|
||||
{ "code": "30", "name": "船体属具", "level": 2,"description": "JC上主要供人员通行又未列入其他各系统的门、窗、梯、盖、桅和桅樯信号,以及为保证人员安全、生活或JC运行所必需的索具和帆布制品、号型和号旗、彩旗、彩灯悬挂装置、栏杆、扶手和天幕装置等。" },
|
||||
{ "code": "40", "name": "舱室分隔和表面覆盖", "level": 2,"description": "作为舱室分隔的非结构性舱壁,活动铺板和格栅,除涂料和金属镀层外的一切敷盖在船体和上层建筑舱室内用作隔热绝缘、防火、减振和隐身用的材料等。" },
|
||||
{ "code": "50", "name": "防腐蚀与防污损", "level": 2,"description": "以防腐蚀为目的的涂料保护、金属涂层和电化保护,以及防止海洋生物附着于船体外表面造成污底的防污损。" },
|
||||
{ "code": "60", "name": "专用项目", "level": 2,"description": "包括船体标识和标志、吃水标志、锁、钥匙和标签、防鼠和防害虫、救生设施、抗沉堵漏器材、拖曳体收放装置、浮筏及隔振装置、船舶装置工作液体、船舶装置备品备件和专用工具。" }
|
||||
]
|
||||
},
|
||||
|
||||
"09": {
|
||||
"code": "09",
|
||||
"name": "专用ZB系统",
|
||||
"level": 1,
|
||||
"description":"提供TZ任务能力的系统或设备。",
|
||||
"children": [
|
||||
{ "code": "00", "name": "总论", "level": 2,"description": "通用要求" },
|
||||
{ "code": "10", "name": "TZ类型设备", "level": 2,"description": "为实现某个TZ任务能力配置的TZ设备" },
|
||||
{ "code": "20", "name": "TZ修复设备", "level": 2,"description": "用来恢复设备性能的TZ修复设备(硬件/软件),包括起重机和牵引支架设备等" },
|
||||
{ "code": "30", "name": "TZ用途设备", "level": 2,"description": "系统中与JCZB相配合,以实现特殊任务目的的部分" }
|
||||
]
|
||||
},
|
||||
|
||||
"10": {
|
||||
"code": "10",
|
||||
"name": "训练系统",
|
||||
"level": 1,
|
||||
"description":"提供训练功能的系统或设备。",
|
||||
"children": [
|
||||
{ "code": "00", "name": "总论", "level": 2,"description": "通用要求" },
|
||||
{ "code": "10", "name": "装置/附件/帮助", "level": 2,"description": "用来使培训人员获得足够的知识和技能,从而能够以最大的效率操作和维护系统的那些可交付使用的训练装置/附件/帮助。" },
|
||||
{ "code": "20", "name": "设备", "level": 2,"description": "用来使培训人员获得足够的知识和技能,从而能够以最大的效率操作和维护系统的那些可交付使用的训练设备。" },
|
||||
{ "code": "30", "name": "设施", "level": 2,"description": "用来使培训人员获得足够的知识和技能,从而能够以最大的效率操作和维护系统的那些可交付使用的训练设施。" }
|
||||
]
|
||||
},
|
||||
|
||||
"11": {
|
||||
"code": "11",
|
||||
"name": "维修、测试和保障系统",
|
||||
"level": 1,
|
||||
"description":"用于ZB监测、维护保养和维修的保障系统、设备或设施。",
|
||||
"children": [
|
||||
{ "code": "00", "name": "总论", "level": 2,"description": "通用要求" },
|
||||
{ "code": "10", "name": "通用维修设备", "level": 2,"description": "用于使系统恢复到可使用状态的设施。" },
|
||||
{ "code": "20", "name": "通用保障设备", "level": 2,"description": "用来保障和维护系统或子系统正常运行而不直接参与执行任务的设备。" },
|
||||
{ "code": "30", "name": "测试和测量设备", "level": 2,"description": "实施特殊诊断、筛选或质量保证工作的过程中用来评估系统或设备的运行状态的通用测试和测量设备" },
|
||||
{ "code": "40", "name": "保障和装卸设备", "level": 2,"description": "用于保障任务系统的可交付使用的工具和装卸设备" }
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
_LEVEL2_NAME_MAP: Dict[str, Dict[str, str]] = {
|
||||
level1_code: {c["code"]: c["name"] for c in (level1.get("children") or [])}
|
||||
for level1_code, level1 in System_Name_Replace_Dict.items()
|
||||
}
|
||||
|
||||
|
||||
File_Dict={
|
||||
"A": {
|
||||
"name": "全船性完工文件",
|
||||
"files": {
|
||||
"01": "全船性完工文件目录",
|
||||
"02": "图样和图表",
|
||||
"03": "说明书",
|
||||
"04": "计算书",
|
||||
"05": "供应品清单",
|
||||
"06": "备品备件、工具等清单",
|
||||
"07": "履历簿",
|
||||
"08": "培训教材",
|
||||
"09": "多媒体",
|
||||
"10": "试验类文件",
|
||||
"11": "使用保养规程建议稿",
|
||||
"12": "质量证明类",
|
||||
"13": "其他文件"
|
||||
}
|
||||
},
|
||||
"B": {
|
||||
"name": "系统设备随机文件",
|
||||
"files": {
|
||||
"01": "随机文件目录",
|
||||
"02": "技术说明书",
|
||||
"03": "操作使用手册",
|
||||
"04": "BD级维修手册",
|
||||
"05": "零部件图解目录",
|
||||
"06": "图样及图表",
|
||||
"07": "备品备件、工具等清单",
|
||||
"08": "试验类文件",
|
||||
"09": "履历簿(书)",
|
||||
"10": "质量证明类",
|
||||
"11": "软件用户手册",
|
||||
"12": "成套设备汇总表",
|
||||
"13": "元器件明细表(包括自制及外购)",
|
||||
"14": "精度检查表",
|
||||
"15": "培训教材",
|
||||
"16": "多媒体",
|
||||
"17": "其他文件"
|
||||
}
|
||||
},
|
||||
"C": {
|
||||
"name": "JD级技术资料",
|
||||
"files": {
|
||||
"01": "预防性维修大纲",
|
||||
"02": "等级修理基准工程单",
|
||||
"03": "JD级维修手册",
|
||||
"04": "修理技术要求",
|
||||
"05": "修理原则工艺",
|
||||
"06": "关键件及寿命件检查方案及更换标准",
|
||||
"07": "计量校准保障方案",
|
||||
"08": "调试细则",
|
||||
"09": "大型设备出舱技术方案",
|
||||
"10": "修理结构方案",
|
||||
"11": "其他文件"
|
||||
}
|
||||
},
|
||||
"D": {
|
||||
"name": "非随装技术资料",
|
||||
"files": {
|
||||
"01": "ZB性能底数报告",
|
||||
"02": "ZB效能底数报告",
|
||||
"03": "ZB使用保养规程",
|
||||
"04": "ZB故障案例",
|
||||
"05": "ZBZZ使用手册"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
def openai_chat(query: str, model="Qwen3.5-35B-A3B"):
|
||||
"""调用OpenAI API"""
|
||||
try:
|
||||
client = OpenAI(
|
||||
api_key="none",
|
||||
base_url="http://192.168.0.46:59800/v1"
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "user", "content": query},
|
||||
],
|
||||
temperature=0.1,
|
||||
stream=False,
|
||||
response_format={"type": "json_object"},
|
||||
extra_body={"chat_template_kwargs": {"enable_thinking": False}}
|
||||
)
|
||||
return response.choices[0].message.content
|
||||
except Exception as e:
|
||||
print(f"调用OpenAI API时出错: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def safe_json_loads(s: str):
|
||||
try:
|
||||
return json.loads(s)
|
||||
except json.JSONDecodeError as e:
|
||||
# 尝试修复常见问题
|
||||
s = s.replace("\\n", "\\\n").replace("\r", "\\r").replace("\t", "\\t")
|
||||
# 或者直接清理非法字符
|
||||
s = re.sub(r'[\x00-\x1f\x7f]', '', s) # 移除控制字符
|
||||
try:
|
||||
return json.loads(s)
|
||||
except:
|
||||
raise e
|
||||
|
||||
|
||||
|
||||
def extract_entity_rela(prompt= str) -> dict:
|
||||
content = openai_chat(prompt)
|
||||
try:
|
||||
cleaned = content.strip()
|
||||
if cleaned.startswith("```json"):
|
||||
cleaned = cleaned[7:].lstrip()
|
||||
if cleaned.endswith("```"):
|
||||
cleaned = cleaned[:-3].rstrip()
|
||||
|
||||
json_result = safe_json_loads(cleaned)
|
||||
|
||||
return json_result
|
||||
except Exception as e:
|
||||
return []
|
||||
|
||||
|
||||
# def get_excel_node_relation(
|
||||
# excel_path: str,
|
||||
# sheet_name: Optional[str] = None,
|
||||
# batch_size: int = 3):
|
||||
# """从Excel提取实体和关系"""
|
||||
# all_entities = []
|
||||
# all_relationships = []
|
||||
|
||||
# try:
|
||||
# # 1. 读取Excel数据
|
||||
# table_json = excel_to_jsondata(excel_path=excel_path, sheet_name=sheet_name)
|
||||
# if not table_json:
|
||||
# logger.warning("Excel文件为空")
|
||||
# return all_entities, all_relationships
|
||||
|
||||
# # 2. 判断表格类型
|
||||
# data = json.dumps(table_json, ensure_ascii=False, indent=2)
|
||||
# prompt_type = prompt_type_excel.replace("{text}", data)
|
||||
# typeresult = openai_chat(prompt_type)
|
||||
|
||||
# if not typeresult:
|
||||
# logger.error("类型判断失败")
|
||||
# return all_entities, all_relationships
|
||||
|
||||
# # 3. 解析类型
|
||||
# try:
|
||||
# cleaned = typeresult.strip()
|
||||
# if cleaned.startswith("```json"):
|
||||
# cleaned = cleaned[7:].lstrip()
|
||||
# if cleaned.endswith("```"):
|
||||
# cleaned = cleaned[:-3].rstrip()
|
||||
|
||||
# json_result = safe_json_loads(cleaned)
|
||||
# table_type = json_result.get("type", "")
|
||||
|
||||
# if not table_type:
|
||||
# logger.error("类型字段缺失")
|
||||
# return all_entities, all_relationships
|
||||
|
||||
# except Exception as e:
|
||||
# logger.error(f"解析类型失败: {e}")
|
||||
# return all_entities, all_relationships
|
||||
|
||||
# # 4. 根据类型处理
|
||||
# if table_type == "备品备件":
|
||||
# logger.info("提取备品备件...")
|
||||
# print(data)
|
||||
# prompt = prompt_beipinbeijian.replace("text", data)
|
||||
|
||||
# try:
|
||||
# result = extract_entity_rela(prompt)
|
||||
|
||||
# if not result or not isinstance(result, dict):
|
||||
# return all_entities, all_relationships
|
||||
|
||||
# entities = result.get("entities", [])
|
||||
# relationships = result.get("relationships", [])
|
||||
|
||||
# for entity in entities:
|
||||
# if "properties" not in entity:
|
||||
# entity["properties"] = {}
|
||||
# if isinstance(entity["properties"], dict):
|
||||
# entity["properties"]["切片"] = data
|
||||
|
||||
# all_entities.extend(entities)
|
||||
# all_relationships.extend(relationships)
|
||||
# logger.info(f"完成: {len(entities)}实体, {len(relationships)}关系")
|
||||
|
||||
# except Exception as e:
|
||||
# logger.error(f"备品备件提取出错: {e}")
|
||||
|
||||
# elif table_type == "故障汇总":
|
||||
# logger.info(f"提取故障汇总 (共{len(table_json)}条, 批量{batch_size})...")
|
||||
|
||||
# for batch_idx in range(0, len(table_json), batch_size):
|
||||
# batch = table_json[batch_idx:batch_idx + batch_size]
|
||||
|
||||
# try:
|
||||
# batch_data = json.dumps(batch, ensure_ascii=False, indent=2)
|
||||
# prompt = prompt_guzhanghuizong.replace("text", batch_data)
|
||||
|
||||
# result = extract_entity_rela(prompt)
|
||||
|
||||
# if not result or not isinstance(result, dict):
|
||||
# continue
|
||||
|
||||
# entities = result.get("entities", [])
|
||||
# relationships = result.get("relationships", [])
|
||||
|
||||
# for entity in entities:
|
||||
# if "properties" not in entity:
|
||||
# entity["properties"] = {}
|
||||
# if isinstance(entity["properties"], dict):
|
||||
# entity["properties"]["切片"] = batch_data
|
||||
|
||||
# all_entities.extend(entities)
|
||||
# all_relationships.extend(relationships)
|
||||
|
||||
# except Exception as e:
|
||||
# logger.error(f"批次{batch_idx//batch_size + 1}出错: {e}")
|
||||
# continue
|
||||
|
||||
# logger.info(f"完成: {len(all_entities)}实体, {len(all_relationships)}关系")
|
||||
|
||||
# else:
|
||||
# logger.warning(f"未知类型: {table_type}")
|
||||
|
||||
# except Exception as e:
|
||||
# logger.error(f"执行失败: {e}")
|
||||
|
||||
# return all_entities, all_relationships
|
||||
|
||||
|
||||
def get_excel_node_relation(
|
||||
excel_path: str,
|
||||
sheet_name: Optional[str] = None,
|
||||
batch_size: int = 3):
|
||||
"""从Excel提取实体和关系"""
|
||||
all_entities = []
|
||||
all_relationships = []
|
||||
|
||||
try:
|
||||
# 1. 读取Excel数据
|
||||
table_json = excel_to_jsondata(excel_path=excel_path, sheet_name=sheet_name)
|
||||
if not table_json:
|
||||
logger.warning("Excel文件为空")
|
||||
return all_entities, all_relationships
|
||||
|
||||
# 2. 判断表格类型
|
||||
data = json.dumps(table_json, ensure_ascii=False, indent=2)
|
||||
prompt_type = prompt_type_excel.replace("{text}", data)
|
||||
typeresult = openai_chat(prompt_type)
|
||||
|
||||
if not typeresult:
|
||||
logger.error("类型判断失败")
|
||||
return all_entities, all_relationships
|
||||
|
||||
# 3. 解析类型
|
||||
try:
|
||||
cleaned = typeresult.strip()
|
||||
if cleaned.startswith("```json"):
|
||||
cleaned = cleaned[7:].lstrip()
|
||||
if cleaned.endswith("```"):
|
||||
cleaned = cleaned[:-3].rstrip()
|
||||
|
||||
json_result = safe_json_loads(cleaned)
|
||||
table_type = json_result.get("type", "")
|
||||
|
||||
if not table_type:
|
||||
logger.error("类型字段缺失")
|
||||
return all_entities, all_relationships
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"解析类型失败: {e}")
|
||||
return all_entities, all_relationships
|
||||
|
||||
# 4. 根据类型处理
|
||||
if table_type == "备品备件":
|
||||
logger.info("提取备品备件...")
|
||||
|
||||
total_batches = (len(table_json) + batch_size - 1) // batch_size
|
||||
|
||||
for batch_idx in range(0, len(table_json), batch_size):
|
||||
batch = table_json[batch_idx:batch_idx + batch_size]
|
||||
current_batch_num = batch_idx // batch_size + 1
|
||||
|
||||
try:
|
||||
# 将当前批次转为JSON字符串
|
||||
batch_data = json.dumps(batch, ensure_ascii=False, indent=2)
|
||||
|
||||
# 使用对应类型的prompt模板
|
||||
prompt = prompt = prompt_beipinbeijian.replace("text", data)
|
||||
|
||||
# 调用实体关系提取
|
||||
result = extract_entity_rela(prompt)
|
||||
|
||||
if not result or not isinstance(result, dict):
|
||||
logger.warning(f"批次{current_batch_num}/{total_batches}未返回有效结果")
|
||||
continue
|
||||
|
||||
entities = result.get("entities", [])
|
||||
relationships = result.get("relationships", [])
|
||||
|
||||
# 为每个实体添加切片属性
|
||||
for entity in entities:
|
||||
if "properties" not in entity:
|
||||
entity["properties"] = {}
|
||||
if isinstance(entity["properties"], dict):
|
||||
entity["properties"]["切片"] = batch_data
|
||||
|
||||
# 累加到总结果中
|
||||
all_entities.extend(entities)
|
||||
all_relationships.extend(relationships)
|
||||
|
||||
logger.info(
|
||||
f"批次{current_batch_num}/{total_batches}完成: "
|
||||
f"{len(entities)}实体, {len(relationships)}关系"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"批次{current_batch_num}/{total_batches}出错: {e}")
|
||||
continue
|
||||
|
||||
elif table_type == "故障汇总":
|
||||
logger.info(f"提取故障汇总 (共{len(table_json)}条, 批量{batch_size})...")
|
||||
|
||||
for batch_idx in range(0, len(table_json), batch_size):
|
||||
batch = table_json[batch_idx:batch_idx + batch_size]
|
||||
|
||||
try:
|
||||
batch_data = json.dumps(batch, ensure_ascii=False, indent=2)
|
||||
prompt = prompt_guzhanghuizong.replace("text", batch_data)
|
||||
|
||||
result = extract_entity_rela(prompt)
|
||||
|
||||
if not result or not isinstance(result, dict):
|
||||
continue
|
||||
|
||||
entities = result.get("entities", [])
|
||||
relationships = result.get("relationships", [])
|
||||
|
||||
for entity in entities:
|
||||
if "properties" not in entity:
|
||||
entity["properties"] = {}
|
||||
if isinstance(entity["properties"], dict):
|
||||
entity["properties"]["切片"] = batch_data
|
||||
|
||||
all_entities.extend(entities)
|
||||
all_relationships.extend(relationships)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"批次{batch_idx//batch_size + 1}出错: {e}")
|
||||
continue
|
||||
|
||||
logger.info(f"完成: {len(all_entities)}实体, {len(all_relationships)}关系")
|
||||
|
||||
else:
|
||||
logger.warning(f"未知类型: {table_type}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"执行失败: {e}")
|
||||
|
||||
return all_entities, all_relationships
|
||||
|
||||
|
||||
# if __name__ == "__main__":
|
||||
# strat_time = time.time()
|
||||
# excel_file = "F:\zklnlp\HJ\hj_kg_code\kgrag\extract_excel\螺旋桨-备品备件.xlsx" # ← 你的 Excel
|
||||
# all_entities, all_relationships = get_excel_node_relation(excel_file)
|
||||
# end_time = time.time()
|
||||
# print(len(all_entities))
|
||||
# print(len(all_relationships))
|
||||
# print(end_time-strat_time)
|
||||
|
||||
518
extract_text_node_relation.py
Normal file
518
extract_text_node_relation.py
Normal file
@ -0,0 +1,518 @@
|
||||
import os
|
||||
import json
|
||||
from typing import List, Tuple, Dict, Any, Optional
|
||||
import re
|
||||
from kg_build.extract_compose_table import get_compose_node_relation
|
||||
from kg_build.extract_compose_tujie_table import get_tujiecompose_node_relation
|
||||
from kg_build.extract_guzhang_table import get_guzhang_node_relation
|
||||
from kg_build.extract_operation_table import get_operation_node_relation
|
||||
from kg_build.extract_repair_table import get_repair_node_relation
|
||||
from kg_build.extract_wxanli_table import get_wxanli_node_relation
|
||||
from kg_build.get_text_node_reala import extract_text_node_relationship
|
||||
from kg_build.extract_beipinbeijian_table import get_beipinbeijian_node_relation
|
||||
from kg_build.extract_filtertext import get_filtertext_node_relation
|
||||
import requests
|
||||
from doc2pdf import Doc2PDF
|
||||
import traceback
|
||||
from extract_excel_node_relation import get_excel_node_relation
|
||||
|
||||
import logging
|
||||
import asyncio # ← 新增导入
|
||||
import time
|
||||
# --- 日志配置 ---
|
||||
logger = logging.getLogger(__name__)
|
||||
if not logger.handlers:
|
||||
handler = logging.StreamHandler()
|
||||
fmt = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||||
handler.setFormatter(fmt)
|
||||
logger.addHandler(handler)
|
||||
logger.setLevel(logging.INFO)
|
||||
|
||||
converter = Doc2PDF()
|
||||
|
||||
def build_repair_relationships(entities: List[Dict]) -> Tuple[List[Dict], List[Dict]]:
|
||||
"""
|
||||
根据维修项目编号 OR 名称在故障维修方案中的出现情况,建立“修理”关系。
|
||||
匹配规则:只要编号或名称任意一个出现在维修方案文本中,即视为匹配。
|
||||
"""
|
||||
# 1. 数据分类:分离故障模式和维修项目
|
||||
faults = [e for e in entities if e.get("type") == "故障模式"]
|
||||
maintenances = [e for e in entities if e.get("type") == "维修项目"]
|
||||
devices = [e for e in entities if e.get("type") == "设备"]
|
||||
|
||||
|
||||
relationships = []
|
||||
entitys = []
|
||||
# 2. 遍历每一个维修项目
|
||||
for mt in maintenances:
|
||||
props = mt.get("properties", {})
|
||||
mt_code = props.get("编号")
|
||||
mt_name = props.get("名称")
|
||||
|
||||
# 如果连名称和编号都没有,无法进行任何匹配,直接跳过
|
||||
if not mt_code and not mt_name:
|
||||
continue
|
||||
|
||||
# 3. 遍历每一个故障模式
|
||||
for fault in faults:
|
||||
f_props = fault.get("properties", {})
|
||||
fault_name = f_props.get("名称")
|
||||
repair_text = f_props.get("维修方案", "")
|
||||
|
||||
if not fault_name or not repair_text:
|
||||
continue
|
||||
|
||||
# 4. 核心匹配逻辑修改 (OR 关系)
|
||||
match_by_code = mt_code in repair_text if mt_code else False
|
||||
match_by_name = mt_name in repair_text if mt_name else False
|
||||
|
||||
# 只要满足任一条件,即建立关系
|
||||
if match_by_code or match_by_name:
|
||||
relationships.append({
|
||||
"type": "修理",
|
||||
"from_entity": mt_name, # 关系起点:维修项目名称
|
||||
"to_entity": fault_name # 关系终点:故障名称
|
||||
})
|
||||
for mt in devices:
|
||||
props = mt.get("properties", {})
|
||||
mt_code = props.get("维修项目名称")
|
||||
mt_name = props.get("维修项目编号")
|
||||
devicename = props.get("名称")
|
||||
# 如果连名称和编号都没有,无法进行任何匹配,直接跳过
|
||||
if not mt_code and not mt_name:
|
||||
continue
|
||||
|
||||
# 3. 遍历每一个故障模式
|
||||
for fault in faults:
|
||||
f_props = fault.get("properties", {})
|
||||
fault_name = f_props.get("名称")
|
||||
repair_text = f_props.get("维修方案", "")
|
||||
|
||||
if not fault_name or not repair_text:
|
||||
continue
|
||||
|
||||
# 4. 核心匹配逻辑修改 (OR 关系)
|
||||
match_by_code = mt_code in repair_text if mt_code else False
|
||||
match_by_name = mt_name in repair_text if mt_name else False
|
||||
|
||||
# 只要满足任一条件,即建立关系
|
||||
if match_by_code or match_by_name:
|
||||
entitys.append({
|
||||
"type": "维修工作",
|
||||
"properties": {"名称": f"{devicename}的维修工作"}
|
||||
})
|
||||
relationships.append({
|
||||
"type": "发生故障使用维修工作",
|
||||
"from_entity": devicename,
|
||||
"to_entity": f"{devicename}的维修工作"
|
||||
})
|
||||
relationships.append({
|
||||
"type": "包含故障",
|
||||
"from_entity": f"{devicename}的维修工作",
|
||||
"to_entity": fault_name
|
||||
})
|
||||
relationships.append({
|
||||
"type": "修理",
|
||||
"from_entity": mt_name, # 关系起点:维修项目名称
|
||||
"to_entity": fault_name
|
||||
})
|
||||
|
||||
|
||||
|
||||
return entitys,relationships
|
||||
|
||||
async def get_md_node_relation(
|
||||
content_list: str,
|
||||
input_pdf_path:str,
|
||||
Ontology: List[str],
|
||||
Ontology_DESCRIPTION: Dict[str, List[str]],
|
||||
Relationships: List[str],
|
||||
RELATIONSHIP_TYPES: Dict[str, List[str]],
|
||||
NODE_TYPES: Dict[str, List[str]],
|
||||
PREFIX_URL : str = "http://192.168.0.111:9085",
|
||||
device : str = "",
|
||||
device_type : str = ""
|
||||
) -> Tuple[List[Dict], List[Dict]]:
|
||||
"""
|
||||
核心函数:协调各模块完成知识抽取、清洗、向量化。
|
||||
"""
|
||||
start_time = time.time()
|
||||
final_all_entity = []
|
||||
final_all_relation = []
|
||||
|
||||
if not content_list:
|
||||
logger.error("输入的 Markdown 内容为空")
|
||||
return [], []
|
||||
if "维修手册" in input_pdf_path or "维修说明书" in input_pdf_path or "培训手册" in input_pdf_path:
|
||||
t = time.time()
|
||||
all_entity_results_guzhang ,all_relation_results_guzhang =await get_guzhang_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device=device)
|
||||
length_entity_guzhang = len(all_entity_results_guzhang)
|
||||
length_relation_guzhang = len(all_relation_results_guzhang)
|
||||
logger.info(f"故障抽取 - 实体: {length_entity_guzhang}, 关系: {length_relation_guzhang}, 耗时: {time.time()-t:.2f}s")
|
||||
if length_entity_guzhang > 0 :
|
||||
final_all_entity.extend(all_entity_results_guzhang)
|
||||
if length_relation_guzhang > 0:
|
||||
final_all_relation.extend(all_relation_results_guzhang)
|
||||
|
||||
t = time.time()
|
||||
all_entity_results_repair ,all_relation_results_repair =await get_repair_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL)
|
||||
length_entity_repair = len(all_entity_results_repair)
|
||||
length_relation_repair = len(all_relation_results_repair)
|
||||
logger.info(f"维修项目表抽取 - 实体: {length_entity_repair}, 关系: {length_relation_repair}, 耗时: {time.time()-t:.2f}s")
|
||||
if length_entity_repair > 0 :
|
||||
final_all_entity.extend(all_entity_results_repair)
|
||||
if length_relation_repair > 0:
|
||||
final_all_relation.extend(all_relation_results_repair)
|
||||
|
||||
t = time.time()
|
||||
all_entity_results_text ,all_relation_results_text =await get_filtertext_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device= device)
|
||||
length_entity_text= len(all_entity_results_text)
|
||||
length_relation_text = len(all_relation_results_text)
|
||||
logger.info(f"文本抽取 - 实体: {length_entity_text}, 关系: {length_relation_text}, 耗时: {time.time()-t:.2f}s")
|
||||
if length_entity_text > 0 :
|
||||
final_all_entity.extend(all_entity_results_text)
|
||||
if length_relation_text > 0:
|
||||
final_all_relation.extend(all_relation_results_text)
|
||||
elif "图解目录" in input_pdf_path:
|
||||
# 2. 处理图解组成表
|
||||
t = time.time()
|
||||
all_entity_results_tujiecompose ,all_relation_results_tujiecompose =await get_tujiecompose_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device=device,lower_entity_type=device_type)
|
||||
length_entity_tujiecompose = len(all_entity_results_tujiecompose)
|
||||
length_relation_tujiecompose = len(all_relation_results_tujiecompose)
|
||||
logger.info(f"图解组成表抽取 - 实体: {length_entity_tujiecompose}, 关系: {length_relation_tujiecompose}, 耗时: {time.time()-t:.2f}s")
|
||||
if length_entity_tujiecompose > 0 :
|
||||
final_all_entity.extend(all_entity_results_tujiecompose)
|
||||
if length_relation_tujiecompose > 0:
|
||||
final_all_relation.extend(all_relation_results_tujiecompose)
|
||||
|
||||
elif "操作使用手册" in input_pdf_path:
|
||||
# 6. 处理操作项目表
|
||||
t = time.time()
|
||||
all_entity_results_operation ,all_relation_results_goperation =await get_operation_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device=device)
|
||||
length_entity_operation= len(all_entity_results_operation)
|
||||
length_relation_goperation = len(all_relation_results_goperation)
|
||||
logger.info(f"操作项目表抽取 - 实体: {length_entity_operation}, 关系: {length_relation_goperation}, 耗时: {time.time()-t:.2f}s")
|
||||
if length_entity_operation > 0 :
|
||||
final_all_entity.extend(all_entity_results_operation)
|
||||
if length_relation_goperation > 0:
|
||||
final_all_relation.extend(all_relation_results_goperation)
|
||||
|
||||
t = time.time()
|
||||
all_entity_results_text ,all_relation_results_text =await get_filtertext_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device= device)
|
||||
length_entity_text= len(all_entity_results_text)
|
||||
length_relation_text = len(all_relation_results_text)
|
||||
logger.info(f"文本抽取 - 实体: {length_entity_text}, 关系: {length_relation_text}, 耗时: {time.time()-t:.2f}s")
|
||||
if length_entity_text > 0 :
|
||||
final_all_entity.extend(all_entity_results_text)
|
||||
if length_relation_text > 0:
|
||||
final_all_relation.extend(all_relation_results_text)
|
||||
|
||||
elif "备品备件" in input_pdf_path:
|
||||
# 5. 处理备品备件表
|
||||
t = time.time()
|
||||
all_entity_results_beipinbeijian ,all_relation_results_beipinbeijian =await get_filtertext_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device=device)
|
||||
length_entity_beipinbeijian= len(all_entity_results_beipinbeijian)
|
||||
length_relation_beipinbeijian = len(all_relation_results_beipinbeijian)
|
||||
logger.info(f"备品备件表抽取 - 实体: {length_entity_beipinbeijian}, 关系: {length_relation_beipinbeijian}, 耗时: {time.time()-t:.2f}s")
|
||||
if length_entity_beipinbeijian > 0 :
|
||||
final_all_entity.extend(all_entity_results_beipinbeijian)
|
||||
if length_relation_beipinbeijian > 0:
|
||||
final_all_relation.extend(all_relation_results_beipinbeijian)
|
||||
elif "故障案例" in input_pdf_path:
|
||||
|
||||
# 7. 处理维修案例
|
||||
t = time.time()
|
||||
all_entity_results_wxanli ,all_relation_results_wxanli =await get_wxanli_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL)
|
||||
length_entity_wxanli= len(all_entity_results_wxanli)
|
||||
length_relation_wxanli = len(all_relation_results_wxanli)
|
||||
logger.info(f"维修案例抽取 - 实体: {length_entity_wxanli}, 关系: {length_relation_wxanli}, 耗时: {time.time()-t:.2f}s")
|
||||
if length_entity_wxanli > 0 :
|
||||
final_all_entity.extend(all_entity_results_wxanli)
|
||||
if length_relation_wxanli > 0:
|
||||
final_all_relation.extend(all_relation_results_wxanli)
|
||||
else:
|
||||
logger.warning("未识别的文档类型,默认进行文本抽取")
|
||||
t = time.time()
|
||||
all_entity_results_text ,all_relation_results_text =await get_filtertext_node_relation(input_pdf_path=input_pdf_path,data=content_list,prefix_url=PREFIX_URL,device= device)
|
||||
length_entity_text= len(all_entity_results_text)
|
||||
length_relation_text = len(all_relation_results_text)
|
||||
logger.info(f"文本抽取 - 实体: {length_entity_text}, 关系: {length_relation_text}, 耗时: {time.time()-t:.2f}s")
|
||||
if length_entity_text > 0 :
|
||||
final_all_entity.extend(all_entity_results_text)
|
||||
if length_relation_text > 0:
|
||||
final_all_relation.extend(all_relation_results_text)
|
||||
|
||||
|
||||
# 8. 建立修复关系
|
||||
reparientitys,repair_relations = build_repair_relationships(final_all_entity)
|
||||
final_all_entity.extend(reparientitys)
|
||||
final_all_relation.extend(repair_relations)
|
||||
logger.info(f"构建故障修复实体数量: {len(reparientitys)}")
|
||||
logger.info(f"构建故障修复关系数量: {len(repair_relations)}")
|
||||
|
||||
|
||||
|
||||
logger.info(f"总实体: {len(final_all_entity)}, 总关系: {len(final_all_relation)}")
|
||||
result = {
|
||||
"entities": final_all_entity,
|
||||
"relations": final_all_relation
|
||||
}
|
||||
output_json_path= "/app/wxproject.json"
|
||||
# 保存为 JSON 文件
|
||||
try:
|
||||
with open(output_json_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(result, f, ensure_ascii=False, indent=4)
|
||||
print(f"结果已保存至:{output_json_path}")
|
||||
except Exception as e:
|
||||
print(f"保存 JSON 文件时出错:{e}")
|
||||
end_time = time.time()
|
||||
print(f"总耗时: {end_time - start_time:.2f} 秒")
|
||||
return final_all_entity, final_all_relation
|
||||
|
||||
def get_content_list_json(filename, folder_path):
|
||||
"""
|
||||
根据传入的文件名,在指定文件夹下查找对应的 _content_list.json 文件并读取内容。
|
||||
|
||||
参数:
|
||||
filename (str): 原始文件名,例如 "163-06A0015-B01001_发动机-操作使用手册.pdf"
|
||||
folder_path (str): 目标文件夹路径
|
||||
|
||||
返回:
|
||||
dict or None: 如果找到文件则返回 JSON 内容字典,否则返回 None
|
||||
"""
|
||||
# 提取文件名主体(去掉扩展名)
|
||||
base_name = os.path.splitext(filename)[0]
|
||||
|
||||
# 构造目标 JSON 文件名
|
||||
target_json_filename = f"{base_name}_content_list.json"
|
||||
|
||||
# 构造完整路径
|
||||
target_path = os.path.join(folder_path, target_json_filename)
|
||||
|
||||
# 检查文件是否存在
|
||||
if not os.path.exists(target_path):
|
||||
print(f"未找到文件: {target_path}")
|
||||
return None
|
||||
|
||||
# 读取并解析 JSON
|
||||
try:
|
||||
with open(target_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
return data
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"JSON 解析错误: {e}")
|
||||
return None
|
||||
except Exception as e:
|
||||
print(f"读取文件时发生错误: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def process_file_content(
|
||||
file_name: str,
|
||||
file_path: str,
|
||||
task_id: str,
|
||||
data_dir: str,
|
||||
ontology: Any,
|
||||
ontology_desc: str,
|
||||
relationships: List,
|
||||
node_types: List,
|
||||
relationship_types: List,
|
||||
prefix_url: str,
|
||||
api_url: str,
|
||||
check_cancelled_callback: callable,
|
||||
device : str,
|
||||
device_type : str,
|
||||
) -> Tuple[List[Dict], List[Dict]]:
|
||||
"""
|
||||
根据文件类型解析内容并抽取实体和关系。
|
||||
|
||||
Returns:
|
||||
tuple: (entities, relationships)
|
||||
"""
|
||||
file_extension = os.path.splitext(file_name)[1].lower()
|
||||
merged_output_data = {"entities": [], "relationships": []}
|
||||
from fileparse_util import process_document
|
||||
|
||||
# --- 处理 PDF 和 MD ---
|
||||
if file_extension in ['.pdf', '.md']:
|
||||
async def _run_process_doc():
|
||||
return await process_document(file_name)
|
||||
doc_result = asyncio.run(_run_process_doc())
|
||||
if doc_result is not None:
|
||||
content_list, _ = doc_result
|
||||
|
||||
else:
|
||||
if not os.path.exists(file_path):
|
||||
raise FileNotFoundError(f"文件未找到: {file_path}")
|
||||
|
||||
with open(file_path, "rb") as f:
|
||||
files = {"file": (os.path.basename(file_path), f, "application/pdf")}
|
||||
logger.info(file_path)
|
||||
response = requests.post(api_url, files=files, timeout=12000)
|
||||
|
||||
if response.status_code != 200:
|
||||
logger.error(f"❌ API 请求失败,状态码: {response.status_code}, 响应: {response.text}")
|
||||
raise Exception(f"PDF/MD 解析 API 失败: {response.status_code}")
|
||||
|
||||
result = response.json()
|
||||
data = result.get("data", {})
|
||||
content_list = data.get("content_list")
|
||||
|
||||
logger.info(f"[{task_id}] 开始文件 {file_name} 的实体和关系抽取")
|
||||
if check_cancelled_callback():
|
||||
raise InterruptedError("任务已取消")
|
||||
|
||||
async def run_extraction():
|
||||
return await get_md_node_relation(
|
||||
content_list=content_list,
|
||||
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 = device,
|
||||
device_type = device_type
|
||||
)
|
||||
|
||||
final_all_entity, final_all_relation = asyncio.run(run_extraction())
|
||||
merged_output_data = {"entities": final_all_entity, "relationships": final_all_relation}
|
||||
|
||||
# --- 处理 DOCX ---
|
||||
elif file_extension in ['.docx']:
|
||||
base_name = os.path.splitext(file_name)[0]
|
||||
pdf_filename = f"{base_name}.pdf"
|
||||
pdf_path = os.path.join(data_dir, pdf_filename)
|
||||
|
||||
logger.info(f"[{task_id}] 转换 DOCX -> PDF: {file_name}")
|
||||
|
||||
try:
|
||||
converter.convert(file_path, output_dir=data_dir)
|
||||
|
||||
if not os.path.exists(pdf_path):
|
||||
raise FileNotFoundError(f"PDF 转换失败: {pdf_path}")
|
||||
|
||||
# 与 PDF 分支保持一致:先尝试 process_document,失败则上传 API
|
||||
async def _run_process_doc_docx():
|
||||
return await process_document(pdf_filename)
|
||||
doc_result = asyncio.run(_run_process_doc_docx())
|
||||
if doc_result is not None:
|
||||
content_list, _ = doc_result
|
||||
else:
|
||||
with open(pdf_path, "rb") as f:
|
||||
files = {"file": (pdf_filename, f, "application/pdf")}
|
||||
logger.info(pdf_path)
|
||||
response = requests.post(api_url, files=files, timeout=12000)
|
||||
|
||||
if response.status_code != 200:
|
||||
logger.error(f"❌ API 请求失败,状态码: {response.status_code}, 响应: {response.text}")
|
||||
raise Exception(f"PDF 解析 API 失败: {response.status_code}")
|
||||
|
||||
result = response.json()
|
||||
data = result.get("data", {})
|
||||
content_list = data.get("content_list")
|
||||
|
||||
logger.info(f"[{task_id}] 开始文件 {file_name} 的实体和关系抽取")
|
||||
if check_cancelled_callback():
|
||||
raise InterruptedError("任务已取消")
|
||||
|
||||
async def run_extraction_docx():
|
||||
return await get_md_node_relation(
|
||||
content_list=content_list,
|
||||
input_pdf_path=pdf_path,
|
||||
Ontology=ontology,
|
||||
Ontology_DESCRIPTION=ontology_desc,
|
||||
Relationships=relationships,
|
||||
RELATIONSHIP_TYPES=relationship_types,
|
||||
NODE_TYPES=node_types,
|
||||
PREFIX_URL=prefix_url,
|
||||
device=device,
|
||||
device_type=device_type,
|
||||
)
|
||||
|
||||
final_all_entity, final_all_relation = asyncio.run(run_extraction_docx())
|
||||
merged_output_data = {"entities": final_all_entity, "relationships": final_all_relation}
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"[{task_id}] DOCX 处理失败: {str(e)}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise
|
||||
finally:
|
||||
pass
|
||||
|
||||
# --- 处理 XLSX ---
|
||||
elif file_extension in ['.xlsx']:
|
||||
base_name = os.path.splitext(file_name)[0]
|
||||
pdf_filename = f"{base_name}.pdf"
|
||||
logger.info(f"[{task_id}] 转换 xlsx -> PDF (仅生成索引引用): {file_name}")
|
||||
# 注意:原代码中这里也调用了 converter,但后续逻辑主要依赖 Excel 解析
|
||||
try:
|
||||
converter.convert(file_path, output_dir=os.path.join(data_dir, "uploads_v1"))
|
||||
except Exception as e:
|
||||
logger.warning(f"Excel 转 PDF 失败,但不影响数据提取: {e}")
|
||||
|
||||
all_entities, all_relationships = get_excel_node_relation(file_path, sheet_name="Sheet1")
|
||||
|
||||
# 补充知识源信息
|
||||
for entity in all_entities:
|
||||
knowledge_source = []
|
||||
qiepian = entity["properties"].get("切片")
|
||||
url = f"{prefix_url}/files/uploads/{pdf_filename}"
|
||||
knowledge_source.append({
|
||||
"filename": file_name,
|
||||
"info": qiepian,
|
||||
"url": url
|
||||
})
|
||||
entity["properties"]["knowledge_source"] = knowledge_source
|
||||
|
||||
merged_output_data = {"entities": all_entities, "relationships": all_relationships}
|
||||
|
||||
else:
|
||||
raise ValueError(f"不支持的文件类型: {file_extension}")
|
||||
|
||||
# 统一取消检查
|
||||
if check_cancelled_callback():
|
||||
raise InterruptedError("任务已取消")
|
||||
|
||||
# 保存缓存 (可选,保持原有逻辑)
|
||||
try:
|
||||
safe_fname = file_name.replace(os.sep, "_").replace(" ", "_")
|
||||
cache_dir = os.path.join(data_dir, "kg_cache")
|
||||
os.makedirs(cache_dir, exist_ok=True)
|
||||
cache_path = os.path.join(cache_dir, f"merged_output_{task_id}_{safe_fname}.json")
|
||||
with open(cache_path, "w", encoding="utf-8") as cf:
|
||||
json.dump(merged_output_data, cf, ensure_ascii=False, indent=2)
|
||||
logger.info(f"[{task_id}] Saved merged_output_data cache: {cache_path}")
|
||||
except Exception as e:
|
||||
logger.warning(f"[{task_id}] Failed to save merged_output_data cache: {e}")
|
||||
|
||||
return merged_output_data["entities"], merged_output_data["relationships"]
|
||||
|
||||
def check_filename_ocr(filename: str) -> Optional[str]:
|
||||
"""
|
||||
根据文件名关键词判断是否需要 OCR。
|
||||
|
||||
Returns:
|
||||
"ocr_true" — 文件名包含维修/操作类关键词
|
||||
"ocr_false" — 文件名包含原理/接线/总图类关键词
|
||||
None — 未命中任何关键词
|
||||
"""
|
||||
#ocr_true_keywords = ["维修手册", "维修说明书", "操作使用手册", "图解目录","培训手册",""]
|
||||
# ocr_false_keywords = ["原理图", "接线图", "总图","随机文件目录","安装说明书","芯线表","链接图","备品备件及工具清单"]
|
||||
ocr_false_keywords = ["原理图", "接线图", "总图","连接图","线图","布置图","安装图","曲线图","标志图","综合图","结构图","基座图","型线图","展开图","编号图","剖面图","吊装图","总装图","曲线图","代号图","视界图","射角图","管路图","周界图","外形图","关系图","装配图","覆盖图","尺寸图","完工图","芯线表","明细表","接线表","编制表","值表册","线路表","矩阵表","汇总表","信息表","连接表","随机文件目录", "随机文件","备品备件","图解目录","电流图","接线图和芯线表","外形图","外形图、外部接线图、安装图","设备总图", "线路图"]
|
||||
|
||||
# for kw in ocr_true_keywords:
|
||||
# if kw in filename:
|
||||
# return "ocr_true"
|
||||
|
||||
for kw in ocr_false_keywords:
|
||||
if kw in filename:
|
||||
return "ocr_false"
|
||||
# else:
|
||||
# return "ocr_true"
|
||||
|
||||
#return None
|
||||
return "ocr_true"
|
||||
|
||||
|
||||
154
extract_util.py
Normal file
154
extract_util.py
Normal file
@ -0,0 +1,154 @@
|
||||
|
||||
import json
|
||||
import fitz # PyMuPDF
|
||||
import os
|
||||
import fitz
|
||||
from collections import Counter
|
||||
import uuid
|
||||
import re
|
||||
def safe_json_loads(s: str):
|
||||
try:
|
||||
return json.loads(s)
|
||||
except json.JSONDecodeError as e:
|
||||
# 尝试修复常见问题
|
||||
s = s.replace("\\n", "\\\n").replace("\r", "\\r").replace("\t", "\\t")
|
||||
# 或者直接清理非法字符
|
||||
s = re.sub(r'[\x00-\x1f\x7f]', '', s) # 移除控制字符
|
||||
try:
|
||||
return json.loads(s)
|
||||
except:
|
||||
raise e
|
||||
|
||||
def highlight_pdf_file(input_pdf_path, highlights, color=(1, 1, 0)):
|
||||
"""
|
||||
在单个 PDF 文件的多个位置添加高亮注释,并返回处理后路径及最常出现的 page_idx。
|
||||
|
||||
:param input_pdf_path: 输入 PDF 文件路径(字符串)
|
||||
:param highlights: 高亮区域列表,每个元素为 dict,格式:
|
||||
{"page_idx": int, "bbox": [x0, y0, x1, y1]}
|
||||
:param output_dir: 输出目录路径。若为 None,则使用输入文件所在目录。
|
||||
:param color: 高亮颜色,RGB 元组,范围 0~1,例如黄色 (1, 1, 0)
|
||||
:return: tuple (output_pdf_path: str, most_common_page: int or None)
|
||||
"""
|
||||
unique_id = uuid.uuid4()
|
||||
|
||||
if not highlights:
|
||||
print("⚠️ 无高亮区域,跳过处理。")
|
||||
return input_pdf_path, None,None
|
||||
|
||||
# 提取所有有效的 page_idx 用于统计
|
||||
valid_page_indices = []
|
||||
|
||||
# 打开 PDF
|
||||
doc = fitz.open(input_pdf_path)
|
||||
|
||||
for item in highlights:
|
||||
page_idx = item.get("page_idx")
|
||||
bbox = item.get("bbox")
|
||||
|
||||
# 基本校验
|
||||
if page_idx is None or bbox is None:
|
||||
print(f"⚠️ 跳过高亮项(缺少 page_idx 或 bbox): {item}")
|
||||
continue
|
||||
if not (0 <= page_idx < doc.page_count):
|
||||
print(f"⚠️ 页面索引 {page_idx} 超出范围(总页数: {doc.page_count}),跳过。")
|
||||
continue
|
||||
if len(bbox) != 4:
|
||||
print(f"⚠️ bbox 格式错误(需 [x0,y0,x1,y1]): {bbox},跳过。")
|
||||
continue
|
||||
|
||||
try:
|
||||
page = doc[page_idx]
|
||||
highlight_annot = page.add_highlight_annot(bbox)
|
||||
highlight_annot.set_colors(stroke=color)
|
||||
highlight_annot.update()
|
||||
valid_page_indices.append(page_idx) # 记录成功处理的 page_idx
|
||||
except Exception as e:
|
||||
print(f"❌ 在 page={page_idx}, bbox={bbox} 添加高亮失败: {e}")
|
||||
|
||||
# 确定输出路径
|
||||
filename = os.path.basename(input_pdf_path)
|
||||
base, ext = os.path.splitext(filename)
|
||||
output_dir = "/app/files/uploads_v1"
|
||||
|
||||
output_pdf_path = os.path.join(output_dir, f"{base}_{unique_id}{ext}")
|
||||
highlight_filename = f"{base}_{unique_id}{ext}"
|
||||
# 保存 PDF
|
||||
doc.save(output_pdf_path, garbage=3, deflate=True)
|
||||
doc.close()
|
||||
|
||||
print(f"✅ 高亮完成!已保存至: {output_pdf_path}")
|
||||
|
||||
# 统计最常出现的 page_idx
|
||||
most_common_page = None
|
||||
if valid_page_indices:
|
||||
counter = Counter(valid_page_indices)
|
||||
most_common_page, _ = counter.most_common(1)[0]
|
||||
|
||||
return highlight_filename, most_common_page,filename
|
||||
|
||||
|
||||
def get_content_bbox(data):
|
||||
content = ""
|
||||
highlight_list = []
|
||||
for ins in data:
|
||||
if "type" not in ins:
|
||||
continue
|
||||
|
||||
if ins["type"] == "text":
|
||||
text_content = ins.get("text", "")
|
||||
content += text_content + "\n"
|
||||
bbox = ins.get("bbox", [])
|
||||
page_idx = ins.get("page_idx", -1)
|
||||
highlight_list.append({
|
||||
"page_idx": page_idx,
|
||||
"bbox": bbox
|
||||
})
|
||||
|
||||
elif ins["type"] == "table":
|
||||
table_caption = ins.get("table_caption")
|
||||
table_body = ins.get("table_body")
|
||||
|
||||
# 处理 table_caption:可能是 list 或 str
|
||||
caption_str = ""
|
||||
if table_caption:
|
||||
if isinstance(table_caption, list):
|
||||
caption_str = "\n".join(str(c) for c in table_caption)
|
||||
else:
|
||||
caption_str = str(table_caption)
|
||||
|
||||
content += caption_str + "\n" + (table_body or "") + "\n"
|
||||
bbox = ins.get("bbox", [])
|
||||
page_idx = ins.get("page_idx", -1)
|
||||
highlight_list.append({
|
||||
"page_idx": page_idx,
|
||||
"bbox": bbox
|
||||
})
|
||||
|
||||
elif ins["type"] == "image":
|
||||
image_caption = ins.get("image_caption", "")
|
||||
img_path = ins.get("img_path", "")
|
||||
|
||||
# 关键修复:处理 image_caption 可能是 list 的情况
|
||||
if isinstance(image_caption, list):
|
||||
image_caption = "\n".join(str(item) for item in image_caption)
|
||||
elif not isinstance(image_caption, str):
|
||||
image_caption = str(image_caption)
|
||||
|
||||
if isinstance(img_path, list):
|
||||
img_path = "\n".join(str(p) for p in img_path)
|
||||
elif not isinstance(img_path, str):
|
||||
img_path = str(img_path)
|
||||
|
||||
if image_caption or img_path:
|
||||
content += image_caption + "\n" + img_path + "\n"
|
||||
bbox = ins.get("bbox", [])
|
||||
page_idx = ins.get("page_idx", -1)
|
||||
highlight_list.append({
|
||||
"page_idx": page_idx,
|
||||
"bbox": bbox
|
||||
})
|
||||
else:
|
||||
continue
|
||||
|
||||
return content, highlight_list
|
||||
4
fault_wenjian/timeout_20260331.txt
Normal file
4
fault_wenjian/timeout_20260331.txt
Normal file
@ -0,0 +1,4 @@
|
||||
test_超时文件.pdf
|
||||
test_超时文件.pdf
|
||||
file_B
|
||||
file_B
|
||||
3
fault_wenjian/timeout_20260401.txt
Normal file
3
fault_wenjian/timeout_20260401.txt
Normal file
@ -0,0 +1,3 @@
|
||||
test_超时文件.pdf
|
||||
122-06A0014-B01001_燃油发动马自达-维修手册.pdf
|
||||
122-06A0014-B01001_燃油发动马自达-维修手册.pdf
|
||||
1
fault_wenjian/timeout_20260610.txt
Normal file
1
fault_wenjian/timeout_20260610.txt
Normal file
@ -0,0 +1 @@
|
||||
163-06A0016-B01003_雷达-操作使用手册.pdf
|
||||
862
filename_proceess_and_kgquery.py
Normal file
862
filename_proceess_and_kgquery.py
Normal file
@ -0,0 +1,862 @@
|
||||
"""
|
||||
PDF 用户上传接口(FastAPI)
|
||||
- 处理文件名和其对应关系
|
||||
"""
|
||||
|
||||
from typing import Dict, Tuple
|
||||
import os
|
||||
import re
|
||||
from typing import Dict
|
||||
from fastapi import FastAPI, HTTPException
|
||||
import logging
|
||||
import json
|
||||
from modelsAPI.model_api import OpenaiAPI
|
||||
from openai import OpenAI
|
||||
from typing import Dict, List, Any
|
||||
|
||||
app = FastAPI(title="PDF Upload Service")
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s")
|
||||
logger = logging.getLogger("filename_neo4j_process")
|
||||
|
||||
|
||||
# =========================
|
||||
# 图方法
|
||||
# =========================
|
||||
|
||||
|
||||
def node_to_json(node):
|
||||
props = dict(node)
|
||||
# ks_raw = props.get("knowledge_source")
|
||||
# knowledge_source = json.loads(ks_raw) if ks_raw else []
|
||||
|
||||
ks_raw = props.get("knowledge_source")
|
||||
knowledge_source = []
|
||||
|
||||
if ks_raw and isinstance(ks_raw, str):
|
||||
try:
|
||||
ks_list = json.loads(ks_raw)
|
||||
if isinstance(ks_list, list):
|
||||
for ks in ks_list:
|
||||
if isinstance(ks, dict):
|
||||
ks.pop("info", None)
|
||||
knowledge_source.append(ks)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
# 解析失败或类型错误时,保持 knowledge_source 为空列表
|
||||
pass
|
||||
|
||||
name = props.get("名称")
|
||||
props.pop("embedding", None)
|
||||
props.pop("knowledge_source", None)
|
||||
props.pop("name", None) # 现在安全地移除
|
||||
props.pop("fulltext", None)
|
||||
if "last_updated" in props:
|
||||
props["最后更新时间"] = props.pop("last_updated")
|
||||
|
||||
if "created_at" in props:
|
||||
props["创建时间"] = props.pop("created_at")
|
||||
return {
|
||||
"id": node.element_id,
|
||||
"name": name, # 使用提前保存的值
|
||||
"labels": list(node.labels),
|
||||
"knowledge_source": knowledge_source,
|
||||
"properties": props,
|
||||
}
|
||||
|
||||
|
||||
def rel_to_json(rel):
|
||||
props = dict(rel)
|
||||
props.pop("embedding", None)
|
||||
props.pop("fact_timeline", None)
|
||||
return {"id": rel.element_id, "label": rel.type, "source": rel.start_node.element_id, "target": rel.end_node.element_id, "properties": props}
|
||||
|
||||
|
||||
def is_simple_path(path):
|
||||
"""
|
||||
判断是否为无环路径(simple path)
|
||||
"""
|
||||
seen = set()
|
||||
for n in path.nodes:
|
||||
nid = n.element_id
|
||||
if nid in seen:
|
||||
return False
|
||||
seen.add(nid)
|
||||
return True
|
||||
|
||||
def build_tree(nodes: List[Dict], links: List[Dict]) -> List[Dict]:
|
||||
"""
|
||||
按照固定层级构建树:舰船(根) -> 舰艇 -> 系统 -> 子系统
|
||||
不再依赖 links 的连线,直接根据 labels 归类,彻底避免环状数据导致的死循环。
|
||||
默认只展开第一艘舰艇的系统/子系统,其它舰艇 children 置空(前端可点击按需加载)。
|
||||
"""
|
||||
|
||||
# 1. 辅助函数:判断标签是否包含关键字
|
||||
def has_label(node, keyword):
|
||||
labels = node.get("labels", [])
|
||||
if isinstance(labels, list):
|
||||
return any(keyword in lbl for lbl in labels)
|
||||
return False
|
||||
|
||||
# 2. 辅助函数:构建标准的树节点格式
|
||||
def make_node(node_data):
|
||||
return {
|
||||
"id": node_data["id"],
|
||||
"name": node_data.get("name", ""),
|
||||
"labels": node_data.get("labels", []),
|
||||
|
||||
"properties": node_data.get("properties", {}),
|
||||
"color": node_data.get("color"),
|
||||
"size": node_data.get("size"),
|
||||
"children": []
|
||||
}
|
||||
|
||||
# 3. 将节点按层级分类存储到字典中,方便快速查找
|
||||
fleet_roots = {} # 第零层:舰船/舰号(父根节点)
|
||||
ships = {} # 第一层:舰艇
|
||||
systems = {} # 第二层:系统(不含子系统)
|
||||
sub_systems = {} # 第三层:子系统
|
||||
|
||||
for node in nodes:
|
||||
# 注意:判断逻辑要互斥,防止"子系统"被误判为"系统"
|
||||
if has_label(node, "舰船") or has_label(node, "舰号"):
|
||||
fleet_roots[node["id"]] = make_node(node)
|
||||
elif has_label(node, "舰艇"):
|
||||
ships[node["id"]] = make_node(node)
|
||||
elif has_label(node, "子系统"):
|
||||
sub_systems[node["id"]] = make_node(node)
|
||||
elif has_label(node, "系统"):
|
||||
systems[node["id"]] = make_node(node)
|
||||
else:
|
||||
pass
|
||||
|
||||
# 4. 组装系统 -> 子系统、舰艇 -> 系统 关系(先把完整子树挂好)
|
||||
for link in links:
|
||||
if link.get("label") != "包含":
|
||||
continue
|
||||
|
||||
src_id = link["source"]
|
||||
tgt_id = link["target"]
|
||||
|
||||
if src_id in ships and tgt_id in systems:
|
||||
ships[src_id]["children"].append(systems[tgt_id])
|
||||
elif src_id in systems and tgt_id in sub_systems:
|
||||
systems[src_id]["children"].append(sub_systems[tgt_id])
|
||||
elif src_id in ships and tgt_id in sub_systems:
|
||||
ships[src_id]["children"].append(sub_systems[tgt_id])
|
||||
|
||||
# 5. 所有舰艇 children 全部保留,仅用 collapsed 标记默认折叠状态。
|
||||
# 前端读取 collapsed=False 的舰艇默认展开,其它折叠;
|
||||
# 用户点击折叠的舰艇时,前端切换 collapsed 即可展示子树。
|
||||
# 只保留名称包含 "122" 的舰艇(隐藏 JZ贱 等其它分支)
|
||||
#sorted_ships = sorted(
|
||||
# (s for s in ships.values() if "122" in s.get("name", "")),
|
||||
# key=lambda s: s.get("name", ""),
|
||||
#)
|
||||
sorted_ships = sorted(ships.values(), key=lambda s: s.get("name", ""))
|
||||
for idx, ship in enumerate(sorted_ships):
|
||||
ship["collapsed"] = (idx != 0)
|
||||
|
||||
# 6. 挂载到舰船根节点;如果数据里没有舰船根,就创建一个虚拟根
|
||||
if fleet_roots:
|
||||
# 取第一个舰船根节点作为父(一般场景下只有一个)
|
||||
root = next(iter(fleet_roots.values()))
|
||||
root["children"] = sorted_ships
|
||||
return [root]
|
||||
else:
|
||||
virtual_root = {
|
||||
"id": "__fleet_root__",
|
||||
"name": "舰号",
|
||||
"labels": ["舰船"],
|
||||
"properties": {},
|
||||
"color": None,
|
||||
"size": None,
|
||||
"children": sorted_ships,
|
||||
}
|
||||
return [virtual_root]
|
||||
def build_tree1(nodes: List[Dict], links: List[Dict]) -> List[Dict]:
|
||||
"""
|
||||
从 nodes + links 构建树结构。
|
||||
- 根节点:label 包含 '舰艇' 的节点
|
||||
- 父子关系:links 中 label == '包含' 的边 (source -> target)
|
||||
- 支持多棵树(多个舰艇根节点)
|
||||
"""
|
||||
|
||||
# 1. 建立 id -> node 的映射
|
||||
node_map: Dict[str, Dict] = {n["id"]: n for n in nodes}
|
||||
|
||||
# 2. 建立 父id -> [子id列表] 的邻接表(只处理"包含"关系)
|
||||
children_map: Dict[str, List[str]] = {n["id"]: [] for n in nodes}
|
||||
parent_map: Dict[str, str] = {} # 子id -> 父id,用于识别根节点
|
||||
|
||||
for link in links:
|
||||
if link.get("label") == "包含":
|
||||
src = link["source"]
|
||||
tgt = link["target"]
|
||||
if src in children_map:
|
||||
children_map[src].append(tgt)
|
||||
if tgt not in parent_map:
|
||||
parent_map[tgt] = src
|
||||
|
||||
# 3. 递归构建子树
|
||||
def build_subtree(node_id: str, ancestors: set = None) -> Dict[str, Any]:
|
||||
if ancestors is None:
|
||||
ancestors = set()
|
||||
if node_id in ancestors:
|
||||
return None
|
||||
ancestors.add(node_id)
|
||||
node = node_map[node_id]
|
||||
subtree = {"id": node_id, "name": node.get("name", ""), "labels": node.get("labels", []), "properties": node.get("properties", {}), "color": node.get("color"), "size": node.get("size"), "children": []}
|
||||
current_ancestors = ancestors | {node_id}
|
||||
for child_id in children_map.get(node_id, []):
|
||||
child = build_subtree(child_id, current_ancestors)
|
||||
if child is not None:
|
||||
subtree["children"].append(child)
|
||||
return subtree
|
||||
|
||||
# def build_subtree(node_id: str) -> Dict[str, Any]:
|
||||
# node = node_map[node_id]
|
||||
# subtree = {
|
||||
# "id": node_id,
|
||||
# "name": node.get("name", ""),
|
||||
# "labels": node.get("labels", []),
|
||||
# "properties": node.get("properties", {}),
|
||||
# "color": node.get("color"),
|
||||
# "size": node.get("size"),
|
||||
# "children": []
|
||||
# }
|
||||
# for child_id in children_map.get(node_id, []):
|
||||
# subtree["children"].append(build_subtree(child_id))
|
||||
# return subtree
|
||||
|
||||
# 4. 找出所有根节点:label 包含 '舰艇',且没有父节点
|
||||
roots = [n["id"] for n in nodes if "舰艇" in n.get("labels", []) and n["id"] not in parent_map]
|
||||
|
||||
# 5. 为每棵树构建结构,并打上树编号
|
||||
forest = []
|
||||
for i, root_id in enumerate(roots):
|
||||
tree = build_subtree(root_id)
|
||||
tree["tree_index"] = i # 区分不同树
|
||||
forest.append(tree)
|
||||
|
||||
return forest
|
||||
|
||||
|
||||
def _filter_searchable(graph: dict) -> dict:
|
||||
"""剥掉节点的 Searchable 标签;若剥完后 labels 为空则整体移除该节点及其关联边"""
|
||||
for n in graph["nodes"]:
|
||||
n["labels"] = [l for l in n.get("labels", []) if l != "Searchable"]
|
||||
valid_ids = {n["id"] for n in graph["nodes"] if n["labels"]}
|
||||
return {
|
||||
"nodes": [n for n in graph["nodes"] if n["id"] in valid_ids],
|
||||
"links": [
|
||||
l for l in graph["links"]
|
||||
if l["source"] in valid_ids and l["target"] in valid_ids
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def build_graph_from_paths(paths):
|
||||
"""
|
||||
核心公共方法
|
||||
从 Neo4j 路径结果构建 graphData
|
||||
"""
|
||||
node_map: Dict[str, dict] = {}
|
||||
link_map: Dict[str, dict] = {}
|
||||
|
||||
def add_edge(rel):
|
||||
s = rel.start_node
|
||||
t = rel.end_node
|
||||
|
||||
if s.element_id not in node_map:
|
||||
node_map[s.element_id] = node_to_json(s)
|
||||
|
||||
if t.element_id not in node_map:
|
||||
node_map[t.element_id] = node_to_json(t)
|
||||
|
||||
if rel.element_id not in link_map:
|
||||
link_map[rel.element_id] = rel_to_json(rel)
|
||||
|
||||
for record in paths:
|
||||
p = record["p"]
|
||||
|
||||
if not is_simple_path(p):
|
||||
continue
|
||||
|
||||
for r in p.relationships:
|
||||
add_edge(r)
|
||||
|
||||
return _filter_searchable({"nodes": list(node_map.values()), "links": list(link_map.values())})
|
||||
|
||||
|
||||
# def fetch_graph_sample(driver,limit: int = 200):
|
||||
# """提取采样逻辑为可复用函数"""
|
||||
# with driver.session() as session:
|
||||
# result = session.run("""
|
||||
# MATCH (n)
|
||||
# WITH n
|
||||
# ORDER BY elementId(n) ASC
|
||||
# LIMIT $limit
|
||||
# MATCH (n)-[rel]->(m)
|
||||
# RETURN n, rel, m
|
||||
# """, limit=limit)
|
||||
|
||||
# node_map, link_map = {}, {}
|
||||
# for rec in result:
|
||||
# n, r, m = rec["n"], rec["rel"], rec["m"]
|
||||
# node_map[n.element_id] = node_to_json(n)
|
||||
# node_map[m.element_id] = node_to_json(m)
|
||||
# link_map[r.element_id] = rel_to_json(r)
|
||||
|
||||
# return {
|
||||
# "nodes": list(node_map.values()),
|
||||
# "links": list(link_map.values())
|
||||
# }
|
||||
|
||||
# def fetch_graph_sample(driver, limit: int = 200):
|
||||
# """
|
||||
# 获取采样数据:
|
||||
# 1. 选取前 $limit 个种子节点(Level 0 或 1,取决于业务定义,这里定义种子为 Level 0)。
|
||||
# 2. 对每个节点向下探索最多 3 跳。
|
||||
# 3. 每个节点对象中增加 level 字段,表示其距离种子节点的深度。
|
||||
# """
|
||||
# with driver.session() as session:
|
||||
# # Cypher 逻辑:
|
||||
# # - 选出种子节点 n
|
||||
# # - 匹配 1-3 跳路径
|
||||
# # - 返回种子节点 n,路径 path,以及该路径的长度(即层级)
|
||||
# result = session.run("""
|
||||
# MATCH (n)
|
||||
# WITH n ORDER BY elementId(n) ASC LIMIT $limit
|
||||
# OPTIONAL MATCH path = (n)-[*1..3]-(m)
|
||||
# UNWIND nodes(path) AS x
|
||||
# WITH n, path, m, COUNT(DISTINCT x) AS nodeCount
|
||||
# WHERE path IS NULL OR nodeCount = LENGTH(path) + 1
|
||||
# RETURN n, path, LENGTH(path) AS depth
|
||||
# """, limit=limit)
|
||||
|
||||
# node_map, link_map = {}, {}
|
||||
|
||||
# for rec in result:
|
||||
# seed_node = rec["n"]
|
||||
# path = rec["path"]
|
||||
# depth = rec["depth"] or 0 # 如果没有路径,深度为 0
|
||||
|
||||
# # 1. 处理种子节点 (Level 0)
|
||||
# if seed_node.element_id not in node_map:
|
||||
# node_data = node_to_json(seed_node)
|
||||
# node_data["level"] = 0 # 种子节点设为 0 层
|
||||
# node_map[seed_node.element_id] = node_data
|
||||
|
||||
# # 2. 如果存在路径,解析路径中的所有节点和关系
|
||||
# if path:
|
||||
# # 路径中的节点处理
|
||||
# # path.nodes 包含了从起始到终点的所有节点
|
||||
# for i, node in enumerate(path.nodes):
|
||||
# if node.element_id not in node_map:
|
||||
# node_data = node_to_json(node)
|
||||
# # 层级即为该节点在当前路径中的索引
|
||||
# node_data["level"] = i
|
||||
# node_map[node.element_id] = node_data
|
||||
# else:
|
||||
# # 如果节点已存在,保留最小的 level (即最靠近种子的距离)
|
||||
# node_map[node.element_id]["level"] = min(node_map[node.element_id].get("level", 3), i)
|
||||
|
||||
# # 路径中的关系处理
|
||||
# for rel in path.relationships:
|
||||
# if rel.element_id not in link_map:
|
||||
# link_map[rel.element_id] = rel_to_json(rel)
|
||||
|
||||
# return {
|
||||
# "nodes": list(node_map.values()),
|
||||
# "links": list(link_map.values())
|
||||
# }
|
||||
|
||||
|
||||
# def fetch_graph_sample(driver, limit=200, max_depth=3):
|
||||
# with driver.session() as session:
|
||||
# result = session.run("""
|
||||
# MATCH (c)
|
||||
# WITH c ORDER BY elementId(c) ASC LIMIT $limit
|
||||
|
||||
# MATCH path = (c)-[*0..3]-(n)
|
||||
# WITH n, MIN(length(path)) AS dist
|
||||
# OPTIONAL MATCH (n)-[r]-(m)
|
||||
|
||||
# RETURN n, dist, r, m
|
||||
# """, limit=limit)
|
||||
|
||||
# node_map, link_map = {}, {}
|
||||
|
||||
# for rec in result:
|
||||
# n = rec["n"]
|
||||
# dist = rec["dist"]
|
||||
# r = rec["r"]
|
||||
# m = rec["m"]
|
||||
|
||||
# # ---- 节点 ----
|
||||
# nid = n.element_id
|
||||
# if nid not in node_map:
|
||||
# node_data = node_to_json(n)
|
||||
# node_data["level"] = min(dist, 3) if dist is not None else 0
|
||||
# node_map[nid] = node_data
|
||||
|
||||
# # ---- 边 ----
|
||||
# if r and m:
|
||||
# mid = m.element_id
|
||||
# if mid not in node_map:
|
||||
# node_map[mid] = node_to_json(m)
|
||||
|
||||
# link_map.setdefault(r.element_id, rel_to_json(r))
|
||||
|
||||
# return {
|
||||
# "nodes": list(node_map.values()),
|
||||
# "links": list(link_map.values())
|
||||
# }
|
||||
|
||||
def fetch_graph_sample(driver):
|
||||
"""
|
||||
查询所有舰艇及其关联的系统和子系统
|
||||
路径结构:(舰艇) -> (系统) -> (子系统)
|
||||
"""
|
||||
|
||||
with driver.session() as session:
|
||||
# =========================
|
||||
# 统一查询:舰艇 -> 系统 -> 子系统
|
||||
# =========================
|
||||
# 逻辑说明:
|
||||
# 1. MATCH (ship:舰艇):选中所有舰艇
|
||||
# 2. -[*1..3]->(target):查找深度为 1 到 3 的路径
|
||||
# - 深度 1:舰艇 -> 系统
|
||||
# - 深度 2:舰艇 -> 系统 -> 子系统
|
||||
# 3. WHERE target:系统 OR target:子系统:确保终点是我们关心的节点类型
|
||||
# (防止查询出其他无关的深层节点)
|
||||
|
||||
query = """
|
||||
MATCH path = (ship:舰艇)-[*1..3]->(target)
|
||||
WHERE target:系统 OR target:子系统
|
||||
RETURN path
|
||||
"""
|
||||
|
||||
result = session.run(query)
|
||||
|
||||
# =========================
|
||||
# 解析 paths
|
||||
# =========================
|
||||
all_nodes = {}
|
||||
all_links = {}
|
||||
|
||||
for record in result:
|
||||
path = record["path"]
|
||||
|
||||
# 解析节点
|
||||
for node in path.nodes:
|
||||
# 兼容 ID 获取
|
||||
node_key = node.element_id if hasattr(node, 'element_id') else node.id
|
||||
if node_key not in all_nodes:
|
||||
all_nodes[node_key] = node_to_json(node)
|
||||
|
||||
# 解析关系
|
||||
for rel in path.relationships:
|
||||
rel_key = rel.element_id if hasattr(rel, 'element_id') else rel.id
|
||||
if rel_key not in all_links:
|
||||
all_links[rel_key] = rel_to_json(rel)
|
||||
|
||||
return _filter_searchable({
|
||||
"nodes": list(all_nodes.values()),
|
||||
"links": list(all_links.values())
|
||||
})
|
||||
def fetch_graph_sample1(driver):
|
||||
"""
|
||||
查找路径:
|
||||
1. 优先:舰艇 -> 子系统
|
||||
2. fallback:舰艇 -> 系统
|
||||
|
||||
注意:根据数据库实际结构,使用 Label (:舰艇, :子系统) 进行匹配,
|
||||
而不是 WHERE category = '...'
|
||||
"""
|
||||
|
||||
# 定义采样数量
|
||||
SAMPLE_LIMIT = 400
|
||||
|
||||
with driver.session() as session:
|
||||
# =========================
|
||||
# 第一阶段:舰艇 -> 子系统
|
||||
# =========================
|
||||
# 修改点:直接使用 (:舰艇) 和 (:子系统) 标签匹配,去掉 WHERE 子句
|
||||
result = session.run(
|
||||
"""
|
||||
MATCH (ship:舰艇)
|
||||
WITH ship LIMIT $limit
|
||||
|
||||
MATCH path = (ship)-[*1..3]->(sub:子系统)
|
||||
RETURN path
|
||||
""",
|
||||
limit=SAMPLE_LIMIT,
|
||||
)
|
||||
|
||||
paths = [record["path"] for record in result]
|
||||
|
||||
# =========================
|
||||
# fallback:舰艇 -> 系统
|
||||
# =========================
|
||||
if not paths:
|
||||
# 修改点:直接使用 (:舰艇) 和 (:系统) 标签匹配
|
||||
result = session.run(
|
||||
"""
|
||||
MATCH (ship:舰艇)
|
||||
WITH ship LIMIT $limit
|
||||
|
||||
MATCH path = (ship)-[*1..2]->(sys:系统)
|
||||
RETURN path
|
||||
""",
|
||||
limit=SAMPLE_LIMIT,
|
||||
)
|
||||
|
||||
paths = [record["path"] for record in result]
|
||||
|
||||
# =========================
|
||||
# 解析 paths
|
||||
# =========================
|
||||
all_nodes = {}
|
||||
all_links = {}
|
||||
|
||||
for path in paths:
|
||||
# 节点
|
||||
for node in path.nodes:
|
||||
# 兼容 Neo4j 5.x (element_id) 和 4.x (id)
|
||||
node_key = node.element_id if hasattr(node, "element_id") else node.id
|
||||
if node_key not in all_nodes:
|
||||
all_nodes[node_key] = node_to_json(node)
|
||||
|
||||
# 关系
|
||||
for rel in path.relationships:
|
||||
rel_key = rel.element_id if hasattr(rel, "element_id") else rel.id
|
||||
if rel_key not in all_links:
|
||||
all_links[rel_key] = rel_to_json(rel)
|
||||
|
||||
return {"nodes": list(all_nodes.values()), "links": list(all_links.values())}
|
||||
|
||||
|
||||
# =============================
|
||||
# 文件名处理方法
|
||||
# =============================
|
||||
|
||||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
TREE_JSON_PATH = os.path.join(BASE_DIR, "tree_data.json")
|
||||
|
||||
|
||||
def _load_tree_data() -> Dict:
|
||||
# 注意:实际运行时请确保 result.json 存在,此处仅为代码结构展示
|
||||
if not os.path.exists(TREE_JSON_PATH):
|
||||
return {}
|
||||
with open(TREE_JSON_PATH, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def _resolve_ship_root(data: Dict, ship_code: str) -> Dict:
|
||||
raw = data.get(ship_code)
|
||||
if not isinstance(raw, dict):
|
||||
return {}
|
||||
|
||||
# 结构 1:直接是节点 {"code","name","children"}
|
||||
if {"code", "name", "children"}.issubset(raw.keys()):
|
||||
return raw
|
||||
|
||||
# 结构 2:包装层 {"101": {...root node...}}
|
||||
if ship_code in raw and isinstance(raw[ship_code], dict):
|
||||
node = raw[ship_code]
|
||||
if {"code", "name", "children"}.issubset(node.keys()):
|
||||
return node
|
||||
|
||||
# 兜底:取第一个 node-like 值
|
||||
for value in raw.values():
|
||||
if isinstance(value, dict) and {"code", "name", "children"}.issubset(value.keys()):
|
||||
return value
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
def _longest_common_substring_len(a: str, b: str) -> int:
|
||||
if not a or not b:
|
||||
return 0
|
||||
# 优化:确保 a 是较短的字符串以节省空间(可选),原逻辑保持不动也没问题
|
||||
dp = [0] * (len(b) + 1)
|
||||
best = 0
|
||||
for i in range(1, len(a) + 1):
|
||||
prev = 0
|
||||
for j in range(1, len(b) + 1):
|
||||
temp = dp[j]
|
||||
if a[i - 1] == b[j - 1]:
|
||||
dp[j] = prev + 1
|
||||
if dp[j] > best:
|
||||
best = dp[j]
|
||||
else:
|
||||
dp[j] = 0
|
||||
prev = temp
|
||||
return best
|
||||
|
||||
|
||||
def _match_level4_tool(filename_text: str, device_node: Dict) -> Tuple[str, str]:
|
||||
"""
|
||||
匹配 Level 4 零件。
|
||||
规则:如果文件名与 Level 3 (设备名) 的最大共同字符串长度 >= 与任何 Level 4 (零件名) 的长度,
|
||||
则只匹配到 Level 3,返回空。
|
||||
只有当 Level 4 的匹配度严格高于 Level 3 时,才返回 Level 4 信息。
|
||||
"""
|
||||
children = device_node.get("children") or {}
|
||||
|
||||
# 1. 计算文件名与当前设备 (Level 3) 名称的匹配度
|
||||
device_name = str(device_node.get("name") or "")
|
||||
level3_score = _longest_common_substring_len(filename_text, device_name)
|
||||
|
||||
best_name = ""
|
||||
best_code = ""
|
||||
best_score = 0 # 记录 Level 4 中的最佳得分
|
||||
|
||||
# 2. 遍历 Level 4 子节点
|
||||
for child in children.values():
|
||||
tool_name = str(child.get("name") or "")
|
||||
tool_code = str(child.get("code") or "")
|
||||
|
||||
score = _longest_common_substring_len(filename_text, tool_name)
|
||||
|
||||
# 关键修改:只有当 Level 4 的得分严格大于 Level 3 的得分时,才视为有效匹配
|
||||
# 并且要比当前找到的其他 Level 4 更好
|
||||
if score > level3_score and score > best_score:
|
||||
best_score = score
|
||||
best_name = tool_name
|
||||
best_code = tool_code
|
||||
|
||||
# 如果没有找到比 Level 3 匹配度更高的 Level 4,则返回空
|
||||
if best_score <= level3_score:
|
||||
return "", ""
|
||||
|
||||
return best_name, best_code
|
||||
|
||||
|
||||
def process_filename(filename: str) -> Dict[str, str]:
|
||||
name, _ = os.path.splitext(os.path.basename(filename))
|
||||
parts = re.split(r"[_-]", name)
|
||||
|
||||
xian_number = ""
|
||||
level_1_system_name = ""
|
||||
level_2_system_name = ""
|
||||
device_name = ""
|
||||
tool_name = ""
|
||||
tool_code = ""
|
||||
|
||||
if not parts:
|
||||
return {}
|
||||
|
||||
xian_number_code = parts[0].strip()
|
||||
level1_code = ""
|
||||
level2_code = ""
|
||||
device_code = ""
|
||||
|
||||
if len(parts) > 1:
|
||||
part = parts[1].strip()
|
||||
# 防止索引越界
|
||||
if len(part) < 2:
|
||||
return {}
|
||||
|
||||
level1_code = part[:2]
|
||||
level2_code = part[2:4] if len(part) > 2 else ""
|
||||
device_code = part[4:] if len(part) > 4 else ""
|
||||
# 若 parts[2] 存在,则用 parts[2] 作为 device_code(支持 122-06A0014-B01001_发动机 格式)
|
||||
if len(parts) > 2:
|
||||
device_code_alt = parts[2].strip()
|
||||
if device_code_alt:
|
||||
device_code = device_code_alt
|
||||
|
||||
data = _load_tree_data()
|
||||
root = _resolve_ship_root(data, xian_number_code)
|
||||
if not root:
|
||||
return {}
|
||||
|
||||
xian_number = root.get("name", "")
|
||||
|
||||
# 一级
|
||||
s1 = root.get("children", {}).get(level1_code)
|
||||
if s1:
|
||||
level_1_system_name = s1.get("name", "")
|
||||
|
||||
# 二级
|
||||
s2 = s1.get("children", {}).get(level2_code)
|
||||
if s2:
|
||||
level_2_system_name = s2.get("name", "")
|
||||
|
||||
# 三级设备
|
||||
device = s2.get("children", {}).get(device_code)
|
||||
if device:
|
||||
device_name = device.get("name", "")
|
||||
# 若存在 level4,匹配与文件名最大相同字符串的零件
|
||||
# 内部已包含逻辑:如果设备名匹配度更高,则不返回零件
|
||||
tool_name, tool_code = _match_level4_tool(name, device)
|
||||
|
||||
return {
|
||||
"Xian_Number": xian_number,
|
||||
"Xian_Number_code": xian_number_code,
|
||||
"Level_1_System_Name": level_1_system_name,
|
||||
"level1_code": level1_code,
|
||||
"Level_2_System_Name": level_2_system_name,
|
||||
"level2_code": level2_code,
|
||||
"Device_Name": device_name,
|
||||
"device_code": device_code,
|
||||
"tool_name": tool_name,
|
||||
"tool_code": tool_code,
|
||||
}
|
||||
|
||||
|
||||
def get_entity(filename):
|
||||
result = process_filename(filename)
|
||||
entities = []
|
||||
relationships = []
|
||||
low_level = ""
|
||||
|
||||
if result:
|
||||
xian_number = result.get("Xian_Number")
|
||||
xian_number_code = result.get("Xian_Number_code")
|
||||
level_1_system_name = result.get("Level_1_System_Name")
|
||||
level1_code = result.get("level1_code")
|
||||
level_2_system_name = result.get("Level_2_System_Name")
|
||||
level2_code = result.get("level2_code")
|
||||
device_name = result.get("Device_Name")
|
||||
device_code = result.get("device_code")
|
||||
tool_name = result.get("tool_name")
|
||||
tool_code = result.get("tool_code")
|
||||
|
||||
if xian_number:
|
||||
entities.append(
|
||||
{
|
||||
"type": "舰艇",
|
||||
"properties": {
|
||||
"名称": xian_number,
|
||||
"舷号": xian_number_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = xian_number
|
||||
|
||||
if level_1_system_name:
|
||||
entities.append(
|
||||
{
|
||||
"type": "系统",
|
||||
"properties": {
|
||||
"名称": level_1_system_name,
|
||||
"系统编码": level1_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = level_1_system_name
|
||||
|
||||
if xian_number and level_1_system_name:
|
||||
relationships.append(
|
||||
{
|
||||
"type": "包含",
|
||||
"from_entity": xian_number,
|
||||
"to_entity": level_1_system_name,
|
||||
}
|
||||
)
|
||||
|
||||
if level_2_system_name:
|
||||
entities.append(
|
||||
{
|
||||
"type": "子系统",
|
||||
"properties": {
|
||||
"名称": level_2_system_name,
|
||||
"子系统编码": level2_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = level_2_system_name
|
||||
|
||||
if level_1_system_name and level_2_system_name:
|
||||
relationships.append(
|
||||
{
|
||||
"type": "包含",
|
||||
"from_entity": level_1_system_name,
|
||||
"to_entity": level_2_system_name,
|
||||
}
|
||||
)
|
||||
|
||||
if device_name:
|
||||
entities.append(
|
||||
{
|
||||
"type": "设备",
|
||||
"properties": {
|
||||
"名称": device_name,
|
||||
"设备编码": device_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = device_name
|
||||
|
||||
if level_2_system_name and device_name:
|
||||
relationships.append(
|
||||
{
|
||||
"type": "包含",
|
||||
"from_entity": level_2_system_name,
|
||||
"to_entity": device_name,
|
||||
}
|
||||
)
|
||||
|
||||
if tool_name:
|
||||
entities.append(
|
||||
{
|
||||
"type": "零件",
|
||||
"properties": {
|
||||
"名称": tool_name,
|
||||
"零件编码": tool_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = tool_name
|
||||
|
||||
if device_name and tool_name:
|
||||
relationships.append(
|
||||
{
|
||||
"type": "包含",
|
||||
"from_entity": device_name,
|
||||
"to_entity": tool_name,
|
||||
}
|
||||
)
|
||||
|
||||
final_result = {
|
||||
"entities": entities,
|
||||
"relationships": relationships,
|
||||
"low_level": low_level,
|
||||
}
|
||||
|
||||
return final_result
|
||||
|
||||
|
||||
# 示例测试逻辑(非必须,仅供验证)
|
||||
if __name__ == "__main__":
|
||||
# 阀控蓄电池脉冲式快速充电装置
|
||||
test_list = [
|
||||
"163-06A0015-B01001_发动机-维修手册.pdf",
|
||||
|
||||
]
|
||||
for i in test_list:
|
||||
result = get_entity(i)
|
||||
lower_entity = result.get("low_level")
|
||||
print(result)
|
||||
print(lower_entity)
|
||||
#from neo4j import GraphDatabase
|
||||
|
||||
#NEO4J_URI = "bolt://192.168.0.111:7687"
|
||||
#NEO4J_USER = os.getenv("NEO4J_USER", "neo4j")
|
||||
#NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD","zdht123@") # 不设默认值,强制要求提供
|
||||
#NEO4J_DATABASE = os.getenv("NEO4J_DATABASE", "neo4j")
|
||||
#driver = GraphDatabase.driver(
|
||||
# NEO4J_URI,
|
||||
# auth=(NEO4J_USER, NEO4J_PASSWORD),
|
||||
# database=NEO4J_DATABASE
|
||||
#)
|
||||
#res = fetch_graph_sample(driver)
|
||||
#print(res)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
821
filename_proceess_and_kgquery_0417.py
Normal file
821
filename_proceess_and_kgquery_0417.py
Normal file
@ -0,0 +1,821 @@
|
||||
"""
|
||||
PDF 用户上传接口(FastAPI)
|
||||
- 处理文件名和其对应关系
|
||||
"""
|
||||
|
||||
from typing import Dict, Tuple
|
||||
import os
|
||||
import re
|
||||
from typing import Dict
|
||||
from fastapi import FastAPI, HTTPException
|
||||
import logging
|
||||
import json
|
||||
from modelsAPI.model_api import OpenaiAPI
|
||||
from openai import OpenAI
|
||||
from typing import Dict, List, Any
|
||||
|
||||
app = FastAPI(title="PDF Upload Service")
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s")
|
||||
logger = logging.getLogger("filename_neo4j_process")
|
||||
|
||||
|
||||
# =========================
|
||||
# 图方法
|
||||
# =========================
|
||||
|
||||
|
||||
def node_to_json(node):
|
||||
props = dict(node)
|
||||
# ks_raw = props.get("knowledge_source")
|
||||
# knowledge_source = json.loads(ks_raw) if ks_raw else []
|
||||
|
||||
ks_raw = props.get("knowledge_source")
|
||||
knowledge_source = []
|
||||
|
||||
if ks_raw and isinstance(ks_raw, str):
|
||||
try:
|
||||
ks_list = json.loads(ks_raw)
|
||||
if isinstance(ks_list, list):
|
||||
for ks in ks_list:
|
||||
if isinstance(ks, dict):
|
||||
ks.pop("info", None)
|
||||
knowledge_source.append(ks)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
# 解析失败或类型错误时,保持 knowledge_source 为空列表
|
||||
pass
|
||||
|
||||
name = props.get("名称")
|
||||
props.pop("embedding", None)
|
||||
props.pop("knowledge_source", None)
|
||||
props.pop("name", None) # 现在安全地移除
|
||||
props.pop("fulltext", None)
|
||||
if "last_updated" in props:
|
||||
props["最后更新时间"] = props.pop("last_updated")
|
||||
|
||||
if "created_at" in props:
|
||||
props["创建时间"] = props.pop("created_at")
|
||||
return {
|
||||
"id": node.element_id,
|
||||
"name": name, # 使用提前保存的值
|
||||
"labels": list(node.labels),
|
||||
"knowledge_source": knowledge_source,
|
||||
"properties": props,
|
||||
}
|
||||
|
||||
|
||||
def rel_to_json(rel):
|
||||
props = dict(rel)
|
||||
props.pop("embedding", None)
|
||||
props.pop("fact_timeline", None)
|
||||
return {"id": rel.element_id, "label": rel.type, "source": rel.start_node.element_id, "target": rel.end_node.element_id, "properties": props}
|
||||
|
||||
|
||||
def is_simple_path(path):
|
||||
"""
|
||||
判断是否为无环路径(simple path)
|
||||
"""
|
||||
seen = set()
|
||||
for n in path.nodes:
|
||||
nid = n.element_id
|
||||
if nid in seen:
|
||||
return False
|
||||
seen.add(nid)
|
||||
return True
|
||||
|
||||
def build_tree(nodes: List[Dict], links: List[Dict]) -> List[Dict]:
|
||||
"""
|
||||
按照固定层级构建树:舰艇 -> 系统 -> 子系统
|
||||
不再依赖 links 的连线,直接根据 labels 归类,彻底避免环状数据导致的死循环。
|
||||
"""
|
||||
|
||||
# 1. 辅助函数:判断标签是否包含关键字
|
||||
def has_label(node, keyword):
|
||||
labels = node.get("labels", [])
|
||||
if isinstance(labels, list):
|
||||
return any(keyword in lbl for lbl in labels)
|
||||
return False
|
||||
|
||||
# 2. 辅助函数:构建标准的树节点格式
|
||||
def make_node(node_data):
|
||||
return {
|
||||
"id": node_data["id"],
|
||||
"name": node_data.get("name", ""),
|
||||
"labels": node_data.get("labels", []),
|
||||
|
||||
"properties": node_data.get("properties", {}),
|
||||
"color": node_data.get("color"),
|
||||
"size": node_data.get("size"),
|
||||
"children": []
|
||||
}
|
||||
|
||||
# 3. 将节点按层级分类存储到字典中,方便快速查找
|
||||
# key: id, value: node_dict_with_children
|
||||
ships = {} # 第一层:舰艇
|
||||
systems = {} # 第二层:系统(不含子系统)
|
||||
sub_systems = {} # 第三层:子系统
|
||||
|
||||
for node in nodes:
|
||||
# 注意:判断逻辑要互斥,防止“子系统”被误判为“系统”
|
||||
if has_label(node, "舰艇"):
|
||||
ships[node["id"]] = make_node(node)
|
||||
elif has_label(node, "子系统"):
|
||||
sub_systems[node["id"]] = make_node(node)
|
||||
elif has_label(node, "系统"):
|
||||
systems[node["id"]] = make_node(node)
|
||||
else:
|
||||
# 如果有其他无关节点,可以在这里处理或忽略
|
||||
pass
|
||||
|
||||
# 4. 开始组装:利用 links 建立关系(此时 links 只充当“连接器”,不决定层级)
|
||||
# 因为我们已经确定了层级,所以不会出现循环引用
|
||||
|
||||
for link in links:
|
||||
if link.get("label") != "包含":
|
||||
continue
|
||||
|
||||
src_id = link["source"]
|
||||
tgt_id = link["target"]
|
||||
|
||||
# 逻辑 A: 舰艇 -> 系统
|
||||
if src_id in ships and tgt_id in systems:
|
||||
ships[src_id]["children"].append(systems[tgt_id])
|
||||
|
||||
# 逻辑 B: 系统 -> 子系统
|
||||
elif src_id in systems and tgt_id in sub_systems:
|
||||
systems[src_id]["children"].append(sub_systems[tgt_id])
|
||||
|
||||
# 逻辑 C: 舰艇 -> 子系统 (防止有些数据直接连了子系统,虽然不规范但也兼容一下)
|
||||
elif src_id in ships and tgt_id in sub_systems:
|
||||
ships[src_id]["children"].append(sub_systems[tgt_id])
|
||||
|
||||
# 5. 返回所有舰艇作为根节点列表
|
||||
return list(ships.values())
|
||||
def build_tree1(nodes: List[Dict], links: List[Dict]) -> List[Dict]:
|
||||
"""
|
||||
从 nodes + links 构建树结构。
|
||||
- 根节点:label 包含 '舰艇' 的节点
|
||||
- 父子关系:links 中 label == '包含' 的边 (source -> target)
|
||||
- 支持多棵树(多个舰艇根节点)
|
||||
"""
|
||||
|
||||
# 1. 建立 id -> node 的映射
|
||||
node_map: Dict[str, Dict] = {n["id"]: n for n in nodes}
|
||||
|
||||
# 2. 建立 父id -> [子id列表] 的邻接表(只处理"包含"关系)
|
||||
children_map: Dict[str, List[str]] = {n["id"]: [] for n in nodes}
|
||||
parent_map: Dict[str, str] = {} # 子id -> 父id,用于识别根节点
|
||||
|
||||
for link in links:
|
||||
if link.get("label") == "包含":
|
||||
src = link["source"]
|
||||
tgt = link["target"]
|
||||
if src in children_map:
|
||||
children_map[src].append(tgt)
|
||||
if tgt not in parent_map:
|
||||
parent_map[tgt] = src
|
||||
|
||||
# 3. 递归构建子树
|
||||
def build_subtree(node_id: str, ancestors: set = None) -> Dict[str, Any]:
|
||||
if ancestors is None:
|
||||
ancestors = set()
|
||||
if node_id in ancestors:
|
||||
return None
|
||||
ancestors.add(node_id)
|
||||
node = node_map[node_id]
|
||||
subtree = {"id": node_id, "name": node.get("name", ""), "labels": node.get("labels", []), "properties": node.get("properties", {}), "color": node.get("color"), "size": node.get("size"), "children": []}
|
||||
current_ancestors = ancestors | {node_id}
|
||||
for child_id in children_map.get(node_id, []):
|
||||
child = build_subtree(child_id, current_ancestors)
|
||||
if child is not None:
|
||||
subtree["children"].append(child)
|
||||
return subtree
|
||||
|
||||
# def build_subtree(node_id: str) -> Dict[str, Any]:
|
||||
# node = node_map[node_id]
|
||||
# subtree = {
|
||||
# "id": node_id,
|
||||
# "name": node.get("name", ""),
|
||||
# "labels": node.get("labels", []),
|
||||
# "properties": node.get("properties", {}),
|
||||
# "color": node.get("color"),
|
||||
# "size": node.get("size"),
|
||||
# "children": []
|
||||
# }
|
||||
# for child_id in children_map.get(node_id, []):
|
||||
# subtree["children"].append(build_subtree(child_id))
|
||||
# return subtree
|
||||
|
||||
# 4. 找出所有根节点:label 包含 '舰艇',且没有父节点
|
||||
roots = [n["id"] for n in nodes if "舰艇" in n.get("labels", []) and n["id"] not in parent_map]
|
||||
|
||||
# 5. 为每棵树构建结构,并打上树编号
|
||||
forest = []
|
||||
for i, root_id in enumerate(roots):
|
||||
tree = build_subtree(root_id)
|
||||
tree["tree_index"] = i # 区分不同树
|
||||
forest.append(tree)
|
||||
|
||||
return forest
|
||||
|
||||
|
||||
def build_graph_from_paths(paths):
|
||||
"""
|
||||
核心公共方法
|
||||
从 Neo4j 路径结果构建 graphData
|
||||
"""
|
||||
node_map: Dict[str, dict] = {}
|
||||
link_map: Dict[str, dict] = {}
|
||||
|
||||
def add_edge(rel):
|
||||
s = rel.start_node
|
||||
t = rel.end_node
|
||||
|
||||
if s.element_id not in node_map:
|
||||
node_map[s.element_id] = node_to_json(s)
|
||||
|
||||
if t.element_id not in node_map:
|
||||
node_map[t.element_id] = node_to_json(t)
|
||||
|
||||
if rel.element_id not in link_map:
|
||||
link_map[rel.element_id] = rel_to_json(rel)
|
||||
|
||||
for record in paths:
|
||||
p = record["p"]
|
||||
|
||||
if not is_simple_path(p):
|
||||
continue
|
||||
|
||||
for r in p.relationships:
|
||||
add_edge(r)
|
||||
|
||||
return {"nodes": list(node_map.values()), "links": list(link_map.values())}
|
||||
|
||||
|
||||
# def fetch_graph_sample(driver,limit: int = 200):
|
||||
# """提取采样逻辑为可复用函数"""
|
||||
# with driver.session() as session:
|
||||
# result = session.run("""
|
||||
# MATCH (n)
|
||||
# WITH n
|
||||
# ORDER BY elementId(n) ASC
|
||||
# LIMIT $limit
|
||||
# MATCH (n)-[rel]->(m)
|
||||
# RETURN n, rel, m
|
||||
# """, limit=limit)
|
||||
|
||||
# node_map, link_map = {}, {}
|
||||
# for rec in result:
|
||||
# n, r, m = rec["n"], rec["rel"], rec["m"]
|
||||
# node_map[n.element_id] = node_to_json(n)
|
||||
# node_map[m.element_id] = node_to_json(m)
|
||||
# link_map[r.element_id] = rel_to_json(r)
|
||||
|
||||
# return {
|
||||
# "nodes": list(node_map.values()),
|
||||
# "links": list(link_map.values())
|
||||
# }
|
||||
|
||||
# def fetch_graph_sample(driver, limit: int = 200):
|
||||
# """
|
||||
# 获取采样数据:
|
||||
# 1. 选取前 $limit 个种子节点(Level 0 或 1,取决于业务定义,这里定义种子为 Level 0)。
|
||||
# 2. 对每个节点向下探索最多 3 跳。
|
||||
# 3. 每个节点对象中增加 level 字段,表示其距离种子节点的深度。
|
||||
# """
|
||||
# with driver.session() as session:
|
||||
# # Cypher 逻辑:
|
||||
# # - 选出种子节点 n
|
||||
# # - 匹配 1-3 跳路径
|
||||
# # - 返回种子节点 n,路径 path,以及该路径的长度(即层级)
|
||||
# result = session.run("""
|
||||
# MATCH (n)
|
||||
# WITH n ORDER BY elementId(n) ASC LIMIT $limit
|
||||
# OPTIONAL MATCH path = (n)-[*1..3]-(m)
|
||||
# UNWIND nodes(path) AS x
|
||||
# WITH n, path, m, COUNT(DISTINCT x) AS nodeCount
|
||||
# WHERE path IS NULL OR nodeCount = LENGTH(path) + 1
|
||||
# RETURN n, path, LENGTH(path) AS depth
|
||||
# """, limit=limit)
|
||||
|
||||
# node_map, link_map = {}, {}
|
||||
|
||||
# for rec in result:
|
||||
# seed_node = rec["n"]
|
||||
# path = rec["path"]
|
||||
# depth = rec["depth"] or 0 # 如果没有路径,深度为 0
|
||||
|
||||
# # 1. 处理种子节点 (Level 0)
|
||||
# if seed_node.element_id not in node_map:
|
||||
# node_data = node_to_json(seed_node)
|
||||
# node_data["level"] = 0 # 种子节点设为 0 层
|
||||
# node_map[seed_node.element_id] = node_data
|
||||
|
||||
# # 2. 如果存在路径,解析路径中的所有节点和关系
|
||||
# if path:
|
||||
# # 路径中的节点处理
|
||||
# # path.nodes 包含了从起始到终点的所有节点
|
||||
# for i, node in enumerate(path.nodes):
|
||||
# if node.element_id not in node_map:
|
||||
# node_data = node_to_json(node)
|
||||
# # 层级即为该节点在当前路径中的索引
|
||||
# node_data["level"] = i
|
||||
# node_map[node.element_id] = node_data
|
||||
# else:
|
||||
# # 如果节点已存在,保留最小的 level (即最靠近种子的距离)
|
||||
# node_map[node.element_id]["level"] = min(node_map[node.element_id].get("level", 3), i)
|
||||
|
||||
# # 路径中的关系处理
|
||||
# for rel in path.relationships:
|
||||
# if rel.element_id not in link_map:
|
||||
# link_map[rel.element_id] = rel_to_json(rel)
|
||||
|
||||
# return {
|
||||
# "nodes": list(node_map.values()),
|
||||
# "links": list(link_map.values())
|
||||
# }
|
||||
|
||||
|
||||
# def fetch_graph_sample(driver, limit=200, max_depth=3):
|
||||
# with driver.session() as session:
|
||||
# result = session.run("""
|
||||
# MATCH (c)
|
||||
# WITH c ORDER BY elementId(c) ASC LIMIT $limit
|
||||
|
||||
# MATCH path = (c)-[*0..3]-(n)
|
||||
# WITH n, MIN(length(path)) AS dist
|
||||
# OPTIONAL MATCH (n)-[r]-(m)
|
||||
|
||||
# RETURN n, dist, r, m
|
||||
# """, limit=limit)
|
||||
|
||||
# node_map, link_map = {}, {}
|
||||
|
||||
# for rec in result:
|
||||
# n = rec["n"]
|
||||
# dist = rec["dist"]
|
||||
# r = rec["r"]
|
||||
# m = rec["m"]
|
||||
|
||||
# # ---- 节点 ----
|
||||
# nid = n.element_id
|
||||
# if nid not in node_map:
|
||||
# node_data = node_to_json(n)
|
||||
# node_data["level"] = min(dist, 3) if dist is not None else 0
|
||||
# node_map[nid] = node_data
|
||||
|
||||
# # ---- 边 ----
|
||||
# if r and m:
|
||||
# mid = m.element_id
|
||||
# if mid not in node_map:
|
||||
# node_map[mid] = node_to_json(m)
|
||||
|
||||
# link_map.setdefault(r.element_id, rel_to_json(r))
|
||||
|
||||
# return {
|
||||
# "nodes": list(node_map.values()),
|
||||
# "links": list(link_map.values())
|
||||
# }
|
||||
|
||||
def fetch_graph_sample(driver):
|
||||
"""
|
||||
查询所有舰艇及其关联的系统和子系统
|
||||
路径结构:(舰艇) -> (系统) -> (子系统)
|
||||
"""
|
||||
|
||||
with driver.session() as session:
|
||||
# =========================
|
||||
# 统一查询:舰艇 -> 系统 -> 子系统
|
||||
# =========================
|
||||
# 逻辑说明:
|
||||
# 1. MATCH (ship:舰艇):选中所有舰艇
|
||||
# 2. -[*1..3]->(target):查找深度为 1 到 3 的路径
|
||||
# - 深度 1:舰艇 -> 系统
|
||||
# - 深度 2:舰艇 -> 系统 -> 子系统
|
||||
# 3. WHERE target:系统 OR target:子系统:确保终点是我们关心的节点类型
|
||||
# (防止查询出其他无关的深层节点)
|
||||
|
||||
query = """
|
||||
MATCH path = (ship:舰艇)-[*1..3]->(target)
|
||||
WHERE target:系统 OR target:子系统
|
||||
RETURN path
|
||||
"""
|
||||
|
||||
result = session.run(query)
|
||||
|
||||
# =========================
|
||||
# 解析 paths
|
||||
# =========================
|
||||
all_nodes = {}
|
||||
all_links = {}
|
||||
|
||||
for record in result:
|
||||
path = record["path"]
|
||||
|
||||
# 解析节点
|
||||
for node in path.nodes:
|
||||
# 兼容 ID 获取
|
||||
node_key = node.element_id if hasattr(node, 'element_id') else node.id
|
||||
if node_key not in all_nodes:
|
||||
all_nodes[node_key] = node_to_json(node)
|
||||
|
||||
# 解析关系
|
||||
for rel in path.relationships:
|
||||
rel_key = rel.element_id if hasattr(rel, 'element_id') else rel.id
|
||||
if rel_key not in all_links:
|
||||
all_links[rel_key] = rel_to_json(rel)
|
||||
|
||||
return {
|
||||
"nodes": list(all_nodes.values()),
|
||||
"links": list(all_links.values())
|
||||
}
|
||||
def fetch_graph_sample1(driver):
|
||||
"""
|
||||
查找路径:
|
||||
1. 优先:舰艇 -> 子系统
|
||||
2. fallback:舰艇 -> 系统
|
||||
|
||||
注意:根据数据库实际结构,使用 Label (:舰艇, :子系统) 进行匹配,
|
||||
而不是 WHERE category = '...'
|
||||
"""
|
||||
|
||||
# 定义采样数量
|
||||
SAMPLE_LIMIT = 400
|
||||
|
||||
with driver.session() as session:
|
||||
# =========================
|
||||
# 第一阶段:舰艇 -> 子系统
|
||||
# =========================
|
||||
# 修改点:直接使用 (:舰艇) 和 (:子系统) 标签匹配,去掉 WHERE 子句
|
||||
result = session.run(
|
||||
"""
|
||||
MATCH (ship:舰艇)
|
||||
WITH ship LIMIT $limit
|
||||
|
||||
MATCH path = (ship)-[*1..3]->(sub:子系统)
|
||||
RETURN path
|
||||
""",
|
||||
limit=SAMPLE_LIMIT,
|
||||
)
|
||||
|
||||
paths = [record["path"] for record in result]
|
||||
|
||||
# =========================
|
||||
# fallback:舰艇 -> 系统
|
||||
# =========================
|
||||
if not paths:
|
||||
# 修改点:直接使用 (:舰艇) 和 (:系统) 标签匹配
|
||||
result = session.run(
|
||||
"""
|
||||
MATCH (ship:舰艇)
|
||||
WITH ship LIMIT $limit
|
||||
|
||||
MATCH path = (ship)-[*1..2]->(sys:系统)
|
||||
RETURN path
|
||||
""",
|
||||
limit=SAMPLE_LIMIT,
|
||||
)
|
||||
|
||||
paths = [record["path"] for record in result]
|
||||
|
||||
# =========================
|
||||
# 解析 paths
|
||||
# =========================
|
||||
all_nodes = {}
|
||||
all_links = {}
|
||||
|
||||
for path in paths:
|
||||
# 节点
|
||||
for node in path.nodes:
|
||||
# 兼容 Neo4j 5.x (element_id) 和 4.x (id)
|
||||
node_key = node.element_id if hasattr(node, "element_id") else node.id
|
||||
if node_key not in all_nodes:
|
||||
all_nodes[node_key] = node_to_json(node)
|
||||
|
||||
# 关系
|
||||
for rel in path.relationships:
|
||||
rel_key = rel.element_id if hasattr(rel, "element_id") else rel.id
|
||||
if rel_key not in all_links:
|
||||
all_links[rel_key] = rel_to_json(rel)
|
||||
|
||||
return {"nodes": list(all_nodes.values()), "links": list(all_links.values())}
|
||||
|
||||
|
||||
# =============================
|
||||
# 文件名处理方法
|
||||
# =============================
|
||||
|
||||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
TREE_JSON_PATH = os.path.join(BASE_DIR, "tree_data.json")
|
||||
|
||||
|
||||
def _load_tree_data() -> Dict:
|
||||
# 注意:实际运行时请确保 result.json 存在,此处仅为代码结构展示
|
||||
if not os.path.exists(TREE_JSON_PATH):
|
||||
return {}
|
||||
with open(TREE_JSON_PATH, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def _resolve_ship_root(data: Dict, ship_code: str) -> Dict:
|
||||
raw = data.get(ship_code)
|
||||
if not isinstance(raw, dict):
|
||||
return {}
|
||||
|
||||
# 结构 1:直接是节点 {"code","name","children"}
|
||||
if {"code", "name", "children"}.issubset(raw.keys()):
|
||||
return raw
|
||||
|
||||
# 结构 2:包装层 {"101": {...root node...}}
|
||||
if ship_code in raw and isinstance(raw[ship_code], dict):
|
||||
node = raw[ship_code]
|
||||
if {"code", "name", "children"}.issubset(node.keys()):
|
||||
return node
|
||||
|
||||
# 兜底:取第一个 node-like 值
|
||||
for value in raw.values():
|
||||
if isinstance(value, dict) and {"code", "name", "children"}.issubset(value.keys()):
|
||||
return value
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
def _longest_common_substring_len(a: str, b: str) -> int:
|
||||
if not a or not b:
|
||||
return 0
|
||||
# 优化:确保 a 是较短的字符串以节省空间(可选),原逻辑保持不动也没问题
|
||||
dp = [0] * (len(b) + 1)
|
||||
best = 0
|
||||
for i in range(1, len(a) + 1):
|
||||
prev = 0
|
||||
for j in range(1, len(b) + 1):
|
||||
temp = dp[j]
|
||||
if a[i - 1] == b[j - 1]:
|
||||
dp[j] = prev + 1
|
||||
if dp[j] > best:
|
||||
best = dp[j]
|
||||
else:
|
||||
dp[j] = 0
|
||||
prev = temp
|
||||
return best
|
||||
|
||||
|
||||
def _match_level4_tool(filename_text: str, device_node: Dict) -> Tuple[str, str]:
|
||||
"""
|
||||
匹配 Level 4 零件。
|
||||
规则:如果文件名与 Level 3 (设备名) 的最大共同字符串长度 >= 与任何 Level 4 (零件名) 的长度,
|
||||
则只匹配到 Level 3,返回空。
|
||||
只有当 Level 4 的匹配度严格高于 Level 3 时,才返回 Level 4 信息。
|
||||
"""
|
||||
children = device_node.get("children") or {}
|
||||
|
||||
# 1. 计算文件名与当前设备 (Level 3) 名称的匹配度
|
||||
device_name = str(device_node.get("name") or "")
|
||||
level3_score = _longest_common_substring_len(filename_text, device_name)
|
||||
|
||||
best_name = ""
|
||||
best_code = ""
|
||||
best_score = 0 # 记录 Level 4 中的最佳得分
|
||||
|
||||
# 2. 遍历 Level 4 子节点
|
||||
for child in children.values():
|
||||
tool_name = str(child.get("name") or "")
|
||||
tool_code = str(child.get("code") or "")
|
||||
|
||||
score = _longest_common_substring_len(filename_text, tool_name)
|
||||
|
||||
# 关键修改:只有当 Level 4 的得分严格大于 Level 3 的得分时,才视为有效匹配
|
||||
# 并且要比当前找到的其他 Level 4 更好
|
||||
if score > level3_score and score > best_score:
|
||||
best_score = score
|
||||
best_name = tool_name
|
||||
best_code = tool_code
|
||||
|
||||
# 如果没有找到比 Level 3 匹配度更高的 Level 4,则返回空
|
||||
if best_score <= level3_score:
|
||||
return "", ""
|
||||
|
||||
return best_name, best_code
|
||||
|
||||
|
||||
def process_filename(filename: str) -> Dict[str, str]:
|
||||
name, _ = os.path.splitext(os.path.basename(filename))
|
||||
parts = re.split(r"[_-]", name)
|
||||
|
||||
xian_number = ""
|
||||
level_1_system_name = ""
|
||||
level_2_system_name = ""
|
||||
device_name = ""
|
||||
tool_name = ""
|
||||
tool_code = ""
|
||||
|
||||
if not parts:
|
||||
return {}
|
||||
|
||||
xian_number_code = parts[0].strip()
|
||||
level1_code = ""
|
||||
level2_code = ""
|
||||
device_code = ""
|
||||
|
||||
if len(parts) > 1:
|
||||
part = parts[1].strip()
|
||||
# 防止索引越界
|
||||
if len(part) < 2:
|
||||
return {}
|
||||
|
||||
level1_code = part[:2]
|
||||
level2_code = part[2:4] if len(part) > 2 else ""
|
||||
device_code = part[4:] if len(part) > 4 else ""
|
||||
# 若 parts[2] 存在,则用 parts[2] 作为 device_code(支持 122-06A0014-B01001_发动机 格式)
|
||||
if len(parts) > 2:
|
||||
device_code_alt = parts[2].strip()
|
||||
if device_code_alt:
|
||||
device_code = device_code_alt
|
||||
|
||||
data = _load_tree_data()
|
||||
root = _resolve_ship_root(data, xian_number_code)
|
||||
if not root:
|
||||
return {}
|
||||
|
||||
xian_number = root.get("name", "")
|
||||
|
||||
# 一级
|
||||
s1 = root.get("children", {}).get(level1_code)
|
||||
if s1:
|
||||
level_1_system_name = s1.get("name", "")
|
||||
|
||||
# 二级
|
||||
s2 = s1.get("children", {}).get(level2_code)
|
||||
if s2:
|
||||
level_2_system_name = s2.get("name", "")
|
||||
|
||||
# 三级设备
|
||||
device = s2.get("children", {}).get(device_code)
|
||||
if device:
|
||||
device_name = device.get("name", "")
|
||||
# 若存在 level4,匹配与文件名最大相同字符串的零件
|
||||
# 内部已包含逻辑:如果设备名匹配度更高,则不返回零件
|
||||
tool_name, tool_code = _match_level4_tool(name, device)
|
||||
|
||||
return {
|
||||
"Xian_Number": xian_number,
|
||||
"Xian_Number_code": xian_number_code,
|
||||
"Level_1_System_Name": level_1_system_name,
|
||||
"level1_code": level1_code,
|
||||
"Level_2_System_Name": level_2_system_name,
|
||||
"level2_code": level2_code,
|
||||
"Device_Name": device_name,
|
||||
"device_code": device_code,
|
||||
"tool_name": tool_name,
|
||||
"tool_code": tool_code,
|
||||
}
|
||||
|
||||
|
||||
def get_entity(filename):
|
||||
result = process_filename(filename)
|
||||
entities = []
|
||||
relationships = []
|
||||
low_level = ""
|
||||
|
||||
if result:
|
||||
xian_number = result.get("Xian_Number")
|
||||
xian_number_code = result.get("Xian_Number_code")
|
||||
level_1_system_name = result.get("Level_1_System_Name")
|
||||
level1_code = result.get("level1_code")
|
||||
level_2_system_name = result.get("Level_2_System_Name")
|
||||
level2_code = result.get("level2_code")
|
||||
device_name = result.get("Device_Name")
|
||||
device_code = result.get("device_code")
|
||||
tool_name = result.get("tool_name")
|
||||
tool_code = result.get("tool_code")
|
||||
|
||||
if xian_number:
|
||||
entities.append(
|
||||
{
|
||||
"type": "舰艇",
|
||||
"properties": {
|
||||
"名称": xian_number,
|
||||
"舷号": xian_number_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = xian_number
|
||||
|
||||
if level_1_system_name:
|
||||
entities.append(
|
||||
{
|
||||
"type": "系统",
|
||||
"properties": {
|
||||
"名称": level_1_system_name,
|
||||
"系统编码": level1_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = level_1_system_name
|
||||
|
||||
if xian_number and level_1_system_name:
|
||||
relationships.append(
|
||||
{
|
||||
"type": "包含",
|
||||
"from_entity": xian_number,
|
||||
"to_entity": level_1_system_name,
|
||||
}
|
||||
)
|
||||
|
||||
if level_2_system_name:
|
||||
entities.append(
|
||||
{
|
||||
"type": "子系统",
|
||||
"properties": {
|
||||
"名称": level_2_system_name,
|
||||
"子系统编码": level2_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = level_2_system_name
|
||||
|
||||
if level_1_system_name and level_2_system_name:
|
||||
relationships.append(
|
||||
{
|
||||
"type": "包含",
|
||||
"from_entity": level_1_system_name,
|
||||
"to_entity": level_2_system_name,
|
||||
}
|
||||
)
|
||||
|
||||
if device_name:
|
||||
entities.append(
|
||||
{
|
||||
"type": "设备",
|
||||
"properties": {
|
||||
"名称": device_name,
|
||||
"设备编码": device_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = device_name
|
||||
|
||||
if level_2_system_name and device_name:
|
||||
relationships.append(
|
||||
{
|
||||
"type": "包含",
|
||||
"from_entity": level_2_system_name,
|
||||
"to_entity": device_name,
|
||||
}
|
||||
)
|
||||
|
||||
if tool_name:
|
||||
entities.append(
|
||||
{
|
||||
"type": "零件",
|
||||
"properties": {
|
||||
"名称": tool_name,
|
||||
"零件编码": tool_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = tool_name
|
||||
|
||||
if device_name and tool_name:
|
||||
relationships.append(
|
||||
{
|
||||
"type": "包含",
|
||||
"from_entity": device_name,
|
||||
"to_entity": tool_name,
|
||||
}
|
||||
)
|
||||
|
||||
final_result = {
|
||||
"entities": entities,
|
||||
"relationships": relationships,
|
||||
"low_level": low_level,
|
||||
}
|
||||
|
||||
return final_result
|
||||
|
||||
|
||||
# 示例测试逻辑(非必须,仅供验证)
|
||||
if __name__ == "__main__":
|
||||
# 阀控蓄电池脉冲式快速充电装置
|
||||
# test_list = [
|
||||
# "163-06A0015-B01001_发动机-维修手册.pdf",
|
||||
|
||||
# ]
|
||||
# for i in test_list:
|
||||
# result = get_entity(i)
|
||||
# print(result)
|
||||
from neo4j import GraphDatabase
|
||||
|
||||
NEO4J_URI = "bolt://192.168.0.111:7687"
|
||||
NEO4J_USER = os.getenv("NEO4J_USER", "neo4j")
|
||||
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD","zdht123@") # 不设默认值,强制要求提供
|
||||
NEO4J_DATABASE = os.getenv("NEO4J_DATABASE", "neo4j")
|
||||
driver = GraphDatabase.driver(
|
||||
NEO4J_URI,
|
||||
auth=(NEO4J_USER, NEO4J_PASSWORD),
|
||||
database=NEO4J_DATABASE
|
||||
)
|
||||
res = fetch_graph_sample(driver)
|
||||
print(res)
|
||||
|
||||
|
||||
|
||||
815
filename_proceess_and_kgquery_hdf.py
Normal file
815
filename_proceess_and_kgquery_hdf.py
Normal file
@ -0,0 +1,815 @@
|
||||
"""
|
||||
PDF 用户上传接口(FastAPI)
|
||||
- 处理文件名和其对应关系
|
||||
"""
|
||||
|
||||
from typing import Dict, Tuple
|
||||
import os
|
||||
import re
|
||||
from typing import Dict
|
||||
from fastapi import FastAPI, HTTPException
|
||||
import logging
|
||||
import json
|
||||
from modelsAPI.model_api import OpenaiAPI
|
||||
from openai import OpenAI
|
||||
from typing import Dict, List, Any
|
||||
|
||||
app = FastAPI(title="PDF Upload Service")
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s")
|
||||
logger = logging.getLogger("filename_neo4j_process")
|
||||
|
||||
|
||||
# =========================
|
||||
# 图方法
|
||||
# =========================
|
||||
|
||||
|
||||
def node_to_json(node):
|
||||
props = dict(node)
|
||||
# ks_raw = props.get("knowledge_source")
|
||||
# knowledge_source = json.loads(ks_raw) if ks_raw else []
|
||||
|
||||
ks_raw = props.get("knowledge_source")
|
||||
knowledge_source = []
|
||||
|
||||
if ks_raw and isinstance(ks_raw, str):
|
||||
try:
|
||||
ks_list = json.loads(ks_raw)
|
||||
if isinstance(ks_list, list):
|
||||
for ks in ks_list:
|
||||
if isinstance(ks, dict):
|
||||
ks.pop("info", None)
|
||||
knowledge_source.append(ks)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
# 解析失败或类型错误时,保持 knowledge_source 为空列表
|
||||
pass
|
||||
|
||||
name = props.get("名称")
|
||||
props.pop("embedding", None)
|
||||
props.pop("knowledge_source", None)
|
||||
props.pop("name", None) # 现在安全地移除
|
||||
props.pop("fulltext", None)
|
||||
if "last_updated" in props:
|
||||
props["最后更新时间"] = props.pop("last_updated")
|
||||
|
||||
if "created_at" in props:
|
||||
props["创建时间"] = props.pop("created_at")
|
||||
return {
|
||||
"id": node.element_id,
|
||||
"name": name, # 使用提前保存的值
|
||||
"labels": list(node.labels),
|
||||
"knowledge_source": knowledge_source,
|
||||
"properties": props,
|
||||
}
|
||||
|
||||
|
||||
def rel_to_json(rel):
|
||||
props = dict(rel)
|
||||
props.pop("embedding", None)
|
||||
props.pop("fact_timeline", None)
|
||||
return {"id": rel.element_id, "label": rel.type, "source": rel.start_node.element_id, "target": rel.end_node.element_id, "properties": props}
|
||||
|
||||
|
||||
def is_simple_path(path):
|
||||
"""
|
||||
判断是否为无环路径(simple path)
|
||||
"""
|
||||
seen = set()
|
||||
for n in path.nodes:
|
||||
nid = n.element_id
|
||||
if nid in seen:
|
||||
return False
|
||||
seen.add(nid)
|
||||
return True
|
||||
|
||||
def build_tree(nodes: List[Dict], links: List[Dict]) -> List[Dict]:
|
||||
"""
|
||||
按照固定层级构建树:舰艇 -> 系统 -> 子系统
|
||||
不再依赖 links 的连线,直接根据 labels 归类,彻底避免环状数据导致的死循环。
|
||||
"""
|
||||
|
||||
# 1. 辅助函数:判断标签是否包含关键字
|
||||
def has_label(node, keyword):
|
||||
labels = node.get("labels", [])
|
||||
if isinstance(labels, list):
|
||||
return any(keyword in lbl for lbl in labels)
|
||||
return False
|
||||
|
||||
# 2. 辅助函数:构建标准的树节点格式
|
||||
def make_node(node_data):
|
||||
return {
|
||||
"id": node_data["id"],
|
||||
"name": node_data.get("name", ""),
|
||||
"labels": node_data.get("labels", []),
|
||||
|
||||
"properties": node_data.get("properties", {}),
|
||||
"color": node_data.get("color"),
|
||||
"size": node_data.get("size"),
|
||||
"children": []
|
||||
}
|
||||
|
||||
# 3. 将节点按层级分类存储到字典中,方便快速查找
|
||||
# key: id, value: node_dict_with_children
|
||||
ships = {} # 第一层:舰艇
|
||||
systems = {} # 第二层:系统(不含子系统)
|
||||
sub_systems = {} # 第三层:子系统
|
||||
|
||||
for node in nodes:
|
||||
# 注意:判断逻辑要互斥,防止“子系统”被误判为“系统”
|
||||
if has_label(node, "舰艇"):
|
||||
ships[node["id"]] = make_node(node)
|
||||
elif has_label(node, "子系统"):
|
||||
sub_systems[node["id"]] = make_node(node)
|
||||
elif has_label(node, "系统"):
|
||||
systems[node["id"]] = make_node(node)
|
||||
else:
|
||||
# 如果有其他无关节点,可以在这里处理或忽略
|
||||
pass
|
||||
|
||||
# 4. 开始组装:利用 links 建立关系(此时 links 只充当“连接器”,不决定层级)
|
||||
# 因为我们已经确定了层级,所以不会出现循环引用
|
||||
|
||||
for link in links:
|
||||
if link.get("label") != "包含":
|
||||
continue
|
||||
|
||||
src_id = link["source"]
|
||||
tgt_id = link["target"]
|
||||
|
||||
# 逻辑 A: 舰艇 -> 系统
|
||||
if src_id in ships and tgt_id in systems:
|
||||
ships[src_id]["children"].append(systems[tgt_id])
|
||||
|
||||
# 逻辑 B: 系统 -> 子系统
|
||||
elif src_id in systems and tgt_id in sub_systems:
|
||||
systems[src_id]["children"].append(sub_systems[tgt_id])
|
||||
|
||||
# 逻辑 C: 舰艇 -> 子系统 (防止有些数据直接连了子系统,虽然不规范但也兼容一下)
|
||||
elif src_id in ships and tgt_id in sub_systems:
|
||||
ships[src_id]["children"].append(sub_systems[tgt_id])
|
||||
|
||||
# 5. 返回所有舰艇作为根节点列表
|
||||
return list(ships.values())
|
||||
def build_tree1(nodes: List[Dict], links: List[Dict]) -> List[Dict]:
|
||||
"""
|
||||
从 nodes + links 构建树结构。
|
||||
- 根节点:label 包含 '舰艇' 的节点
|
||||
- 父子关系:links 中 label == '包含' 的边 (source -> target)
|
||||
- 支持多棵树(多个舰艇根节点)
|
||||
"""
|
||||
|
||||
# 1. 建立 id -> node 的映射
|
||||
node_map: Dict[str, Dict] = {n["id"]: n for n in nodes}
|
||||
|
||||
# 2. 建立 父id -> [子id列表] 的邻接表(只处理"包含"关系)
|
||||
children_map: Dict[str, List[str]] = {n["id"]: [] for n in nodes}
|
||||
parent_map: Dict[str, str] = {} # 子id -> 父id,用于识别根节点
|
||||
|
||||
for link in links:
|
||||
if link.get("label") == "包含":
|
||||
src = link["source"]
|
||||
tgt = link["target"]
|
||||
if src in children_map:
|
||||
children_map[src].append(tgt)
|
||||
if tgt not in parent_map:
|
||||
parent_map[tgt] = src
|
||||
|
||||
# 3. 递归构建子树
|
||||
def build_subtree(node_id: str, ancestors: set = None) -> Dict[str, Any]:
|
||||
if ancestors is None:
|
||||
ancestors = set()
|
||||
if node_id in ancestors:
|
||||
return None
|
||||
ancestors.add(node_id)
|
||||
node = node_map[node_id]
|
||||
subtree = {"id": node_id, "name": node.get("name", ""), "labels": node.get("labels", []), "properties": node.get("properties", {}), "color": node.get("color"), "size": node.get("size"), "children": []}
|
||||
current_ancestors = ancestors | {node_id}
|
||||
for child_id in children_map.get(node_id, []):
|
||||
child = build_subtree(child_id, current_ancestors)
|
||||
if child is not None:
|
||||
subtree["children"].append(child)
|
||||
return subtree
|
||||
|
||||
# def build_subtree(node_id: str) -> Dict[str, Any]:
|
||||
# node = node_map[node_id]
|
||||
# subtree = {
|
||||
# "id": node_id,
|
||||
# "name": node.get("name", ""),
|
||||
# "labels": node.get("labels", []),
|
||||
# "properties": node.get("properties", {}),
|
||||
# "color": node.get("color"),
|
||||
# "size": node.get("size"),
|
||||
# "children": []
|
||||
# }
|
||||
# for child_id in children_map.get(node_id, []):
|
||||
# subtree["children"].append(build_subtree(child_id))
|
||||
# return subtree
|
||||
|
||||
# 4. 找出所有根节点:label 包含 '舰艇',且没有父节点
|
||||
roots = [n["id"] for n in nodes if "舰艇" in n.get("labels", []) and n["id"] not in parent_map]
|
||||
|
||||
# 5. 为每棵树构建结构,并打上树编号
|
||||
forest = []
|
||||
for i, root_id in enumerate(roots):
|
||||
tree = build_subtree(root_id)
|
||||
tree["tree_index"] = i # 区分不同树
|
||||
forest.append(tree)
|
||||
|
||||
return forest
|
||||
|
||||
|
||||
def build_graph_from_paths(paths):
|
||||
"""
|
||||
核心公共方法
|
||||
从 Neo4j 路径结果构建 graphData
|
||||
"""
|
||||
node_map: Dict[str, dict] = {}
|
||||
link_map: Dict[str, dict] = {}
|
||||
|
||||
def add_edge(rel):
|
||||
s = rel.start_node
|
||||
t = rel.end_node
|
||||
|
||||
if s.element_id not in node_map:
|
||||
node_map[s.element_id] = node_to_json(s)
|
||||
|
||||
if t.element_id not in node_map:
|
||||
node_map[t.element_id] = node_to_json(t)
|
||||
|
||||
if rel.element_id not in link_map:
|
||||
link_map[rel.element_id] = rel_to_json(rel)
|
||||
|
||||
for record in paths:
|
||||
p = record["p"]
|
||||
|
||||
if not is_simple_path(p):
|
||||
continue
|
||||
|
||||
for r in p.relationships:
|
||||
add_edge(r)
|
||||
|
||||
return {"nodes": list(node_map.values()), "links": list(link_map.values())}
|
||||
|
||||
|
||||
# def fetch_graph_sample(driver,limit: int = 200):
|
||||
# """提取采样逻辑为可复用函数"""
|
||||
# with driver.session() as session:
|
||||
# result = session.run("""
|
||||
# MATCH (n)
|
||||
# WITH n
|
||||
# ORDER BY elementId(n) ASC
|
||||
# LIMIT $limit
|
||||
# MATCH (n)-[rel]->(m)
|
||||
# RETURN n, rel, m
|
||||
# """, limit=limit)
|
||||
|
||||
# node_map, link_map = {}, {}
|
||||
# for rec in result:
|
||||
# n, r, m = rec["n"], rec["rel"], rec["m"]
|
||||
# node_map[n.element_id] = node_to_json(n)
|
||||
# node_map[m.element_id] = node_to_json(m)
|
||||
# link_map[r.element_id] = rel_to_json(r)
|
||||
|
||||
# return {
|
||||
# "nodes": list(node_map.values()),
|
||||
# "links": list(link_map.values())
|
||||
# }
|
||||
|
||||
# def fetch_graph_sample(driver, limit: int = 200):
|
||||
# """
|
||||
# 获取采样数据:
|
||||
# 1. 选取前 $limit 个种子节点(Level 0 或 1,取决于业务定义,这里定义种子为 Level 0)。
|
||||
# 2. 对每个节点向下探索最多 3 跳。
|
||||
# 3. 每个节点对象中增加 level 字段,表示其距离种子节点的深度。
|
||||
# """
|
||||
# with driver.session() as session:
|
||||
# # Cypher 逻辑:
|
||||
# # - 选出种子节点 n
|
||||
# # - 匹配 1-3 跳路径
|
||||
# # - 返回种子节点 n,路径 path,以及该路径的长度(即层级)
|
||||
# result = session.run("""
|
||||
# MATCH (n)
|
||||
# WITH n ORDER BY elementId(n) ASC LIMIT $limit
|
||||
# OPTIONAL MATCH path = (n)-[*1..3]-(m)
|
||||
# UNWIND nodes(path) AS x
|
||||
# WITH n, path, m, COUNT(DISTINCT x) AS nodeCount
|
||||
# WHERE path IS NULL OR nodeCount = LENGTH(path) + 1
|
||||
# RETURN n, path, LENGTH(path) AS depth
|
||||
# """, limit=limit)
|
||||
|
||||
# node_map, link_map = {}, {}
|
||||
|
||||
# for rec in result:
|
||||
# seed_node = rec["n"]
|
||||
# path = rec["path"]
|
||||
# depth = rec["depth"] or 0 # 如果没有路径,深度为 0
|
||||
|
||||
# # 1. 处理种子节点 (Level 0)
|
||||
# if seed_node.element_id not in node_map:
|
||||
# node_data = node_to_json(seed_node)
|
||||
# node_data["level"] = 0 # 种子节点设为 0 层
|
||||
# node_map[seed_node.element_id] = node_data
|
||||
|
||||
# # 2. 如果存在路径,解析路径中的所有节点和关系
|
||||
# if path:
|
||||
# # 路径中的节点处理
|
||||
# # path.nodes 包含了从起始到终点的所有节点
|
||||
# for i, node in enumerate(path.nodes):
|
||||
# if node.element_id not in node_map:
|
||||
# node_data = node_to_json(node)
|
||||
# # 层级即为该节点在当前路径中的索引
|
||||
# node_data["level"] = i
|
||||
# node_map[node.element_id] = node_data
|
||||
# else:
|
||||
# # 如果节点已存在,保留最小的 level (即最靠近种子的距离)
|
||||
# node_map[node.element_id]["level"] = min(node_map[node.element_id].get("level", 3), i)
|
||||
|
||||
# # 路径中的关系处理
|
||||
# for rel in path.relationships:
|
||||
# if rel.element_id not in link_map:
|
||||
# link_map[rel.element_id] = rel_to_json(rel)
|
||||
|
||||
# return {
|
||||
# "nodes": list(node_map.values()),
|
||||
# "links": list(link_map.values())
|
||||
# }
|
||||
|
||||
|
||||
# def fetch_graph_sample(driver, limit=200, max_depth=3):
|
||||
# with driver.session() as session:
|
||||
# result = session.run("""
|
||||
# MATCH (c)
|
||||
# WITH c ORDER BY elementId(c) ASC LIMIT $limit
|
||||
|
||||
# MATCH path = (c)-[*0..3]-(n)
|
||||
# WITH n, MIN(length(path)) AS dist
|
||||
# OPTIONAL MATCH (n)-[r]-(m)
|
||||
|
||||
# RETURN n, dist, r, m
|
||||
# """, limit=limit)
|
||||
|
||||
# node_map, link_map = {}, {}
|
||||
|
||||
# for rec in result:
|
||||
# n = rec["n"]
|
||||
# dist = rec["dist"]
|
||||
# r = rec["r"]
|
||||
# m = rec["m"]
|
||||
|
||||
# # ---- 节点 ----
|
||||
# nid = n.element_id
|
||||
# if nid not in node_map:
|
||||
# node_data = node_to_json(n)
|
||||
# node_data["level"] = min(dist, 3) if dist is not None else 0
|
||||
# node_map[nid] = node_data
|
||||
|
||||
# # ---- 边 ----
|
||||
# if r and m:
|
||||
# mid = m.element_id
|
||||
# if mid not in node_map:
|
||||
# node_map[mid] = node_to_json(m)
|
||||
|
||||
# link_map.setdefault(r.element_id, rel_to_json(r))
|
||||
|
||||
# return {
|
||||
# "nodes": list(node_map.values()),
|
||||
# "links": list(link_map.values())
|
||||
# }
|
||||
|
||||
def fetch_graph_sample(driver):
|
||||
"""
|
||||
查询所有舰艇及其关联的系统和子系统
|
||||
路径结构:(舰艇) -> (系统) -> (子系统)
|
||||
"""
|
||||
|
||||
with driver.session() as session:
|
||||
# =========================
|
||||
# 统一查询:舰艇 -> 系统 -> 子系统
|
||||
# =========================
|
||||
# 逻辑说明:
|
||||
# 1. MATCH (ship:舰艇):选中所有舰艇
|
||||
# 2. -[*1..3]->(target):查找深度为 1 到 3 的路径
|
||||
# - 深度 1:舰艇 -> 系统
|
||||
# - 深度 2:舰艇 -> 系统 -> 子系统
|
||||
# 3. WHERE target:系统 OR target:子系统:确保终点是我们关心的节点类型
|
||||
# (防止查询出其他无关的深层节点)
|
||||
|
||||
query = """
|
||||
MATCH path = (ship:舰艇)-[*1..3]->(target)
|
||||
WHERE target:系统 OR target:子系统
|
||||
RETURN path
|
||||
"""
|
||||
|
||||
result = session.run(query)
|
||||
|
||||
# =========================
|
||||
# 解析 paths
|
||||
# =========================
|
||||
all_nodes = {}
|
||||
all_links = {}
|
||||
|
||||
for record in result:
|
||||
path = record["path"]
|
||||
|
||||
# 解析节点
|
||||
for node in path.nodes:
|
||||
# 兼容 ID 获取
|
||||
node_key = node.element_id if hasattr(node, 'element_id') else node.id
|
||||
if node_key not in all_nodes:
|
||||
all_nodes[node_key] = node_to_json(node)
|
||||
|
||||
# 解析关系
|
||||
for rel in path.relationships:
|
||||
rel_key = rel.element_id if hasattr(rel, 'element_id') else rel.id
|
||||
if rel_key not in all_links:
|
||||
all_links[rel_key] = rel_to_json(rel)
|
||||
|
||||
return {
|
||||
"nodes": list(all_nodes.values()),
|
||||
"links": list(all_links.values())
|
||||
}
|
||||
def fetch_graph_sample1(driver):
|
||||
"""
|
||||
查找路径:
|
||||
1. 优先:舰艇 -> 子系统
|
||||
2. fallback:舰艇 -> 系统
|
||||
|
||||
注意:根据数据库实际结构,使用 Label (:舰艇, :子系统) 进行匹配,
|
||||
而不是 WHERE category = '...'
|
||||
"""
|
||||
|
||||
# 定义采样数量
|
||||
SAMPLE_LIMIT = 400
|
||||
|
||||
with driver.session() as session:
|
||||
# =========================
|
||||
# 第一阶段:舰艇 -> 子系统
|
||||
# =========================
|
||||
# 修改点:直接使用 (:舰艇) 和 (:子系统) 标签匹配,去掉 WHERE 子句
|
||||
result = session.run(
|
||||
"""
|
||||
MATCH (ship:舰艇)
|
||||
WITH ship LIMIT $limit
|
||||
|
||||
MATCH path = (ship)-[*1..3]->(sub:子系统)
|
||||
RETURN path
|
||||
""",
|
||||
limit=SAMPLE_LIMIT,
|
||||
)
|
||||
|
||||
paths = [record["path"] for record in result]
|
||||
|
||||
# =========================
|
||||
# fallback:舰艇 -> 系统
|
||||
# =========================
|
||||
if not paths:
|
||||
# 修改点:直接使用 (:舰艇) 和 (:系统) 标签匹配
|
||||
result = session.run(
|
||||
"""
|
||||
MATCH (ship:舰艇)
|
||||
WITH ship LIMIT $limit
|
||||
|
||||
MATCH path = (ship)-[*1..2]->(sys:系统)
|
||||
RETURN path
|
||||
""",
|
||||
limit=SAMPLE_LIMIT,
|
||||
)
|
||||
|
||||
paths = [record["path"] for record in result]
|
||||
|
||||
# =========================
|
||||
# 解析 paths
|
||||
# =========================
|
||||
all_nodes = {}
|
||||
all_links = {}
|
||||
|
||||
for path in paths:
|
||||
# 节点
|
||||
for node in path.nodes:
|
||||
# 兼容 Neo4j 5.x (element_id) 和 4.x (id)
|
||||
node_key = node.element_id if hasattr(node, "element_id") else node.id
|
||||
if node_key not in all_nodes:
|
||||
all_nodes[node_key] = node_to_json(node)
|
||||
|
||||
# 关系
|
||||
for rel in path.relationships:
|
||||
rel_key = rel.element_id if hasattr(rel, "element_id") else rel.id
|
||||
if rel_key not in all_links:
|
||||
all_links[rel_key] = rel_to_json(rel)
|
||||
|
||||
return {"nodes": list(all_nodes.values()), "links": list(all_links.values())}
|
||||
|
||||
|
||||
# =============================
|
||||
# 文件名处理方法
|
||||
# =============================
|
||||
|
||||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||||
TREE_JSON_PATH = os.path.join(BASE_DIR, "tree_data.json")
|
||||
|
||||
|
||||
def _load_tree_data() -> Dict:
|
||||
# 注意:实际运行时请确保 result.json 存在,此处仅为代码结构展示
|
||||
if not os.path.exists(TREE_JSON_PATH):
|
||||
return {}
|
||||
with open(TREE_JSON_PATH, "r", encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def _resolve_ship_root(data: Dict, ship_code: str) -> Dict:
|
||||
raw = data.get(ship_code)
|
||||
if not isinstance(raw, dict):
|
||||
return {}
|
||||
|
||||
# 结构 1:直接是节点 {"code","name","children"}
|
||||
if {"code", "name", "children"}.issubset(raw.keys()):
|
||||
return raw
|
||||
|
||||
# 结构 2:包装层 {"101": {...root node...}}
|
||||
if ship_code in raw and isinstance(raw[ship_code], dict):
|
||||
node = raw[ship_code]
|
||||
if {"code", "name", "children"}.issubset(node.keys()):
|
||||
return node
|
||||
|
||||
# 兜底:取第一个 node-like 值
|
||||
for value in raw.values():
|
||||
if isinstance(value, dict) and {"code", "name", "children"}.issubset(value.keys()):
|
||||
return value
|
||||
|
||||
return {}
|
||||
|
||||
|
||||
def _longest_common_substring_len(a: str, b: str) -> int:
|
||||
if not a or not b:
|
||||
return 0
|
||||
# 优化:确保 a 是较短的字符串以节省空间(可选),原逻辑保持不动也没问题
|
||||
dp = [0] * (len(b) + 1)
|
||||
best = 0
|
||||
for i in range(1, len(a) + 1):
|
||||
prev = 0
|
||||
for j in range(1, len(b) + 1):
|
||||
temp = dp[j]
|
||||
if a[i - 1] == b[j - 1]:
|
||||
dp[j] = prev + 1
|
||||
if dp[j] > best:
|
||||
best = dp[j]
|
||||
else:
|
||||
dp[j] = 0
|
||||
prev = temp
|
||||
return best
|
||||
|
||||
|
||||
def _match_level4_tool(filename_text: str, device_node: Dict) -> Tuple[str, str]:
|
||||
"""
|
||||
匹配 Level 4 零件。
|
||||
规则:如果文件名与 Level 3 (设备名) 的最大共同字符串长度 >= 与任何 Level 4 (零件名) 的长度,
|
||||
则只匹配到 Level 3,返回空。
|
||||
只有当 Level 4 的匹配度严格高于 Level 3 时,才返回 Level 4 信息。
|
||||
"""
|
||||
children = device_node.get("children") or {}
|
||||
|
||||
# 1. 计算文件名与当前设备 (Level 3) 名称的匹配度
|
||||
device_name = str(device_node.get("name") or "")
|
||||
level3_score = _longest_common_substring_len(filename_text, device_name)
|
||||
|
||||
best_name = ""
|
||||
best_code = ""
|
||||
best_score = 0 # 记录 Level 4 中的最佳得分
|
||||
|
||||
# 2. 遍历 Level 4 子节点
|
||||
for child in children.values():
|
||||
tool_name = str(child.get("name") or "")
|
||||
tool_code = str(child.get("code") or "")
|
||||
|
||||
score = _longest_common_substring_len(filename_text, tool_name)
|
||||
|
||||
# 关键修改:只有当 Level 4 的得分严格大于 Level 3 的得分时,才视为有效匹配
|
||||
# 并且要比当前找到的其他 Level 4 更好
|
||||
if score > level3_score and score > best_score:
|
||||
best_score = score
|
||||
best_name = tool_name
|
||||
best_code = tool_code
|
||||
|
||||
# 如果没有找到比 Level 3 匹配度更高的 Level 4,则返回空
|
||||
if best_score <= level3_score:
|
||||
return "", ""
|
||||
|
||||
return best_name, best_code
|
||||
|
||||
|
||||
def process_filename(filename: str) -> Dict[str, str]:
|
||||
name, _ = os.path.splitext(os.path.basename(filename))
|
||||
parts = re.split(r"[_-]", name)
|
||||
|
||||
xian_number = ""
|
||||
level_1_system_name = ""
|
||||
level_2_system_name = ""
|
||||
device_name = ""
|
||||
tool_name = ""
|
||||
tool_code = ""
|
||||
|
||||
if not parts:
|
||||
return {}
|
||||
|
||||
xian_number_code = parts[0].strip()
|
||||
level1_code = ""
|
||||
level2_code = ""
|
||||
device_code = ""
|
||||
|
||||
if len(parts) > 1:
|
||||
part = parts[1].strip()
|
||||
# 防止索引越界
|
||||
if len(part) < 2:
|
||||
return {}
|
||||
|
||||
level1_code = part[:2]
|
||||
level2_code = part[2:4] if len(part) > 2 else ""
|
||||
device_code = part[4:] if len(part) > 4 else ""
|
||||
|
||||
data = _load_tree_data()
|
||||
root = _resolve_ship_root(data, xian_number_code)
|
||||
if not root:
|
||||
return {}
|
||||
|
||||
xian_number = root.get("name", "")
|
||||
|
||||
# 一级
|
||||
s1 = root.get("children", {}).get(level1_code)
|
||||
if s1:
|
||||
level_1_system_name = s1.get("name", "")
|
||||
|
||||
# 二级
|
||||
s2 = s1.get("children", {}).get(level2_code)
|
||||
if s2:
|
||||
level_2_system_name = s2.get("name", "")
|
||||
|
||||
# 三级设备
|
||||
device = s2.get("children", {}).get(device_code)
|
||||
if device:
|
||||
device_name = device.get("name", "")
|
||||
# 若存在 level4,匹配与文件名最大相同字符串的零件
|
||||
# 内部已包含逻辑:如果设备名匹配度更高,则不返回零件
|
||||
tool_name, tool_code = _match_level4_tool(name, device)
|
||||
|
||||
return {
|
||||
"Xian_Number": xian_number,
|
||||
"Xian_Number_code": xian_number_code,
|
||||
"Level_1_System_Name": level_1_system_name,
|
||||
"level1_code": level1_code,
|
||||
"Level_2_System_Name": level_2_system_name,
|
||||
"level2_code": level2_code,
|
||||
"Device_Name": device_name,
|
||||
"device_code": device_code,
|
||||
"tool_name": tool_name,
|
||||
"tool_code": tool_code,
|
||||
}
|
||||
|
||||
|
||||
def get_entity(filename):
|
||||
result = process_filename(filename)
|
||||
entities = []
|
||||
relationships = []
|
||||
low_level = ""
|
||||
|
||||
if result:
|
||||
xian_number = result.get("Xian_Number")
|
||||
xian_number_code = result.get("Xian_Number_code")
|
||||
level_1_system_name = result.get("Level_1_System_Name")
|
||||
level1_code = result.get("level1_code")
|
||||
level_2_system_name = result.get("Level_2_System_Name")
|
||||
level2_code = result.get("level2_code")
|
||||
device_name = result.get("Device_Name")
|
||||
device_code = result.get("device_code")
|
||||
tool_name = result.get("tool_name")
|
||||
tool_code = result.get("tool_code")
|
||||
|
||||
if xian_number:
|
||||
entities.append(
|
||||
{
|
||||
"type": "舰艇",
|
||||
"properties": {
|
||||
"名称": xian_number,
|
||||
"舷号": xian_number_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = xian_number
|
||||
|
||||
if level_1_system_name:
|
||||
entities.append(
|
||||
{
|
||||
"type": "系统",
|
||||
"properties": {
|
||||
"名称": level_1_system_name,
|
||||
"系统编码": level1_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = level_1_system_name
|
||||
|
||||
if xian_number and level_1_system_name:
|
||||
relationships.append(
|
||||
{
|
||||
"type": "包含",
|
||||
"from_entity": xian_number,
|
||||
"to_entity": level_1_system_name,
|
||||
}
|
||||
)
|
||||
|
||||
if level_2_system_name:
|
||||
entities.append(
|
||||
{
|
||||
"type": "子系统",
|
||||
"properties": {
|
||||
"名称": level_2_system_name,
|
||||
"子系统编码": level2_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = level_2_system_name
|
||||
|
||||
if level_1_system_name and level_2_system_name:
|
||||
relationships.append(
|
||||
{
|
||||
"type": "包含",
|
||||
"from_entity": level_1_system_name,
|
||||
"to_entity": level_2_system_name,
|
||||
}
|
||||
)
|
||||
|
||||
if device_name:
|
||||
entities.append(
|
||||
{
|
||||
"type": "设备",
|
||||
"properties": {
|
||||
"名称": device_name,
|
||||
"设备编码": device_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = device_name
|
||||
|
||||
if level_2_system_name and device_name:
|
||||
relationships.append(
|
||||
{
|
||||
"type": "包含",
|
||||
"from_entity": level_2_system_name,
|
||||
"to_entity": device_name,
|
||||
}
|
||||
)
|
||||
|
||||
if tool_name:
|
||||
entities.append(
|
||||
{
|
||||
"type": "零件",
|
||||
"properties": {
|
||||
"名称": tool_name,
|
||||
"零件编码": tool_code,
|
||||
},
|
||||
}
|
||||
)
|
||||
low_level = tool_name
|
||||
|
||||
if device_name and tool_name:
|
||||
relationships.append(
|
||||
{
|
||||
"type": "包含",
|
||||
"from_entity": device_name,
|
||||
"to_entity": tool_name,
|
||||
}
|
||||
)
|
||||
|
||||
final_result = {
|
||||
"entities": entities,
|
||||
"relationships": relationships,
|
||||
"low_level": low_level,
|
||||
}
|
||||
|
||||
return final_result
|
||||
|
||||
|
||||
# 示例测试逻辑(非必须,仅供验证)
|
||||
if __name__ == "__main__":
|
||||
# 阀控蓄电池脉冲式快速充电装置
|
||||
# test_list = [
|
||||
# "163-06A0015-B01001_发动机-维修手册.pdf",
|
||||
|
||||
# ]
|
||||
# for i in test_list:
|
||||
# result = get_entity(i)
|
||||
# print(result)
|
||||
from neo4j import GraphDatabase
|
||||
|
||||
NEO4J_URI = "bolt://192.168.0.111:7687"
|
||||
NEO4J_USER = os.getenv("NEO4J_USER", "neo4j")
|
||||
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD","zdht123@") # 不设默认值,强制要求提供
|
||||
NEO4J_DATABASE = os.getenv("NEO4J_DATABASE", "neo4j")
|
||||
driver = GraphDatabase.driver(
|
||||
NEO4J_URI,
|
||||
auth=(NEO4J_USER, NEO4J_PASSWORD),
|
||||
database=NEO4J_DATABASE
|
||||
)
|
||||
res = fetch_graph_sample(driver)
|
||||
print(res)
|
||||
|
||||
|
||||
168
fileparse_util.py
Normal file
168
fileparse_util.py
Normal file
@ -0,0 +1,168 @@
|
||||
import os
|
||||
import json
|
||||
import base64
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Optional, Tuple, Dict
|
||||
import aiofiles # pip install aiofiles
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def find_and_read_content_list(
|
||||
directory: str,
|
||||
original_filename: str,
|
||||
encoding: str = 'utf-8'
|
||||
) -> Tuple[Optional[list], Optional[str]]:
|
||||
"""根据原始文件名查找并读取对应的 _content_list.json 文件。"""
|
||||
file_name, _ = os.path.splitext(original_filename)
|
||||
target_filename = f"{file_name}_content_list.json"
|
||||
logger.info(f"正在查找文件: {target_filename}")
|
||||
|
||||
# os.walk 没有原生 async 版本,放到线程里执行避免阻塞事件循环
|
||||
def _walk_for_file():
|
||||
for root, _, files in os.walk(directory):
|
||||
if target_filename in files:
|
||||
return os.path.join(root, target_filename)
|
||||
return None
|
||||
|
||||
file_path = await asyncio.to_thread(_walk_for_file)
|
||||
if file_path is None:
|
||||
return None, None
|
||||
|
||||
try:
|
||||
async with aiofiles.open(file_path, 'r', encoding=encoding) as f:
|
||||
text = await f.read()
|
||||
# JSON 解析是 CPU 操作,大文件可考虑放线程
|
||||
return json.loads(text), file_path
|
||||
except json.JSONDecodeError:
|
||||
logger.error(f"文件 {file_path} 不是有效的 JSON 格式。")
|
||||
return None, None
|
||||
except Exception as e:
|
||||
logger.error(f"读取文件 {file_path} 时发生错误: {e}")
|
||||
return None, None
|
||||
|
||||
|
||||
def extract_all_jpg_filenames(doc_data: list) -> set:
|
||||
"""从文档数据中提取所有 .jpg 图片的纯文件名。纯 CPU 操作,不需要 async。"""
|
||||
result_set = set()
|
||||
for item in doc_data:
|
||||
if 'img_path' in item and item['img_path']:
|
||||
img_path = item['img_path']
|
||||
if img_path.endswith('.jpg'):
|
||||
result_set.add(os.path.basename(img_path))
|
||||
return result_set
|
||||
|
||||
|
||||
async def _encode_one_image(
|
||||
img_file: str,
|
||||
local_image_dir: str,
|
||||
mime_map: dict,
|
||||
) -> Tuple[str, Optional[str]]:
|
||||
"""编码单张图片,返回 (文件名, data_uri 或 None)。"""
|
||||
img_path = os.path.join(local_image_dir, img_file)
|
||||
if not os.path.isfile(img_path):
|
||||
logger.warning(f"图片不存在: {img_path}")
|
||||
return img_file, None
|
||||
|
||||
try:
|
||||
async with aiofiles.open(img_path, "rb") as f:
|
||||
img_bytes = await f.read()
|
||||
# base64 编码是 CPU 操作,小文件直接做即可;大文件可放线程
|
||||
img_base64 = base64.b64encode(img_bytes).decode('utf-8')
|
||||
ext = img_file.lower().rsplit('.', 1)[-1]
|
||||
mime_type = f"image/{mime_map.get(ext, 'png')}"
|
||||
return img_file, f"data:{mime_type};base64,{img_base64}"
|
||||
except Exception as e:
|
||||
logger.error(f"编码图片失败 {img_file}: {str(e)}")
|
||||
return img_file, None
|
||||
|
||||
|
||||
async def encode_images_to_base64(
|
||||
local_image_dir: str,
|
||||
referenced_images: set,
|
||||
concurrency: int = 32,
|
||||
) -> Dict[str, str]:
|
||||
"""并发将引用的图片编码为 base64。"""
|
||||
if not os.path.exists(local_image_dir) or not referenced_images:
|
||||
return {}
|
||||
|
||||
mime_map = {'jpg': 'jpeg', 'jpeg': 'jpeg', 'png': 'png', 'gif': 'gif', 'webp': 'webp'}
|
||||
|
||||
# 用 Semaphore 限制并发数,避免一次性打开几百个文件句柄
|
||||
sem = asyncio.Semaphore(concurrency)
|
||||
|
||||
async def _bounded(img_file: str):
|
||||
async with sem:
|
||||
return await _encode_one_image(img_file, local_image_dir, mime_map)
|
||||
|
||||
results = await asyncio.gather(*(_bounded(f) for f in referenced_images))
|
||||
return {name: data for name, data in results if data is not None}
|
||||
|
||||
|
||||
async def process_document(
|
||||
input_file_name: str,
|
||||
search_dir: Optional[str] = "/app/mineru_output",
|
||||
local_image_dir: Optional[str] = None,
|
||||
concurrency: int = 32,
|
||||
) -> Optional[Tuple[list, Dict[str, str]]]:
|
||||
"""
|
||||
异步处理文档:查找 content_list.json,提取图片引用,并发编码为 base64。
|
||||
|
||||
Args:
|
||||
input_file_name: 原始文件名(如 "xxx.pdf")
|
||||
search_dir: 搜索 _content_list.json 的根目录
|
||||
local_image_dir: 图片所在目录,None 则自动推断
|
||||
concurrency: 图片编码的并发数上限,默认 32
|
||||
|
||||
Returns:
|
||||
成功: (content_list, images_dict)
|
||||
失败: None
|
||||
"""
|
||||
content, content_path = await find_and_read_content_list(search_dir, input_file_name)
|
||||
if content is None:
|
||||
logger.error("❌ 未找到指定的 _content_list.json 文件。")
|
||||
return None
|
||||
|
||||
logger.info(f"✅ 找到文件: {content_path}")
|
||||
logger.info(f"内容类型: {type(content).__name__}, 长度: {len(str(content))}")
|
||||
|
||||
if local_image_dir is None:
|
||||
local_image_dir = os.path.join(os.path.dirname(content_path), "images")
|
||||
logger.info(f"图片目录: {local_image_dir}")
|
||||
|
||||
referenced_images = extract_all_jpg_filenames(content)
|
||||
logger.info(f"引用图片数量: {len(referenced_images)}")
|
||||
|
||||
images_dict = await encode_images_to_base64(
|
||||
local_image_dir, referenced_images, concurrency=concurrency
|
||||
)
|
||||
logger.info(f"成功编码 {len(images_dict)} 张图片")
|
||||
|
||||
return content, images_dict
|
||||
|
||||
|
||||
async def main():
|
||||
result = await process_document(
|
||||
input_file_name="163-06A0014-B01001_发动机-维修手册.pdf",
|
||||
search_dir="/app/mineru_output",
|
||||
concurrency=32,
|
||||
)
|
||||
|
||||
if result is None:
|
||||
return 1
|
||||
|
||||
content, images_dict = result
|
||||
print(f"\n=== 处理完成 ===")
|
||||
print(f"文档段落数: {len(content)}")
|
||||
print(f"图片数量: {len(images_dict)}")
|
||||
if images_dict:
|
||||
first_key = next(iter(images_dict))
|
||||
print(f"示例图片 [{first_key}]: {images_dict[first_key][:80]}...")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
exit_code = asyncio.run(main())
|
||||
exit(exit_code)
|
||||
0
flow/__init__.py
Normal file
0
flow/__init__.py
Normal file
15
flow/main.py
Normal file
15
flow/main.py
Normal file
@ -0,0 +1,15 @@
|
||||
from database.milvus_api import Milvus_Database
|
||||
from modelsAPI.model_api import get_embeddings
|
||||
|
||||
async def autorag_flow(query,milvus_db,collection_name,top_k=5):
|
||||
query_vector = get_embeddings(list(query))
|
||||
search_results = milvus_db.search(collection_name=collection_name,query_vector=query_vector, top_k=top_k)
|
||||
|
||||
search_prompt="""你作为一个由中科院的资深文档问答专家,请你从以下文本信息中选取与用户问题相关的内容作为参考:
|
||||
-----
|
||||
文本信息:{context}
|
||||
-----
|
||||
请你仔细思考后用中文回答用户的问题,如果文本中没有关于问题的答案,则直接回答 文档中没有相关内容
|
||||
用户提问:{query}
|
||||
你的回答:"""
|
||||
|
||||
BIN
font/DejaVuMathTeXGyre.ttf
Normal file
BIN
font/DejaVuMathTeXGyre.ttf
Normal file
Binary file not shown.
BIN
font/DejaVuSans-Bold.ttf
Normal file
BIN
font/DejaVuSans-Bold.ttf
Normal file
Binary file not shown.
BIN
font/DejaVuSans-BoldOblique.ttf
Normal file
BIN
font/DejaVuSans-BoldOblique.ttf
Normal file
Binary file not shown.
BIN
font/DejaVuSans-ExtraLight.ttf
Normal file
BIN
font/DejaVuSans-ExtraLight.ttf
Normal file
Binary file not shown.
BIN
font/DejaVuSans-Oblique.ttf
Normal file
BIN
font/DejaVuSans-Oblique.ttf
Normal file
Binary file not shown.
Some files were not shown because too many files have changed in this diff Show More
Loading…
x
Reference in New Issue
Block a user