# -*- coding: utf-8 -*-
|
# @file: doc_processor.py
|
# @author: lyg
|
# @date: 2025-5-13
|
# @version:
|
# @description: 处理文档，拆分文档，将拆分后的章节保存到数据库中。
|
from langchain_core.messages import HumanMessage
|
from langchain_core.output_parsers import JsonOutputParser
|
from langchain_core.prompts import ChatPromptTemplate, HumanMessagePromptTemplate
|
|
from knowledgebase.db.doc_db_models import TEntity
|
from knowledgebase.doc.docx_split import DocSplit
|
import asyncio
|
from knowledgebase.db.doc_db_helper import doc_dbh
|
from knowledgebase.doc.entity_helper import get_entity_helper
|
from knowledgebase.doc.entity_recognition import EntityRecognition
|
import os.path
|
|
from knowledgebase.doc.models import DocInfo, ParagraphInfo, DocType
|
from knowledgebase.llm import llm
|
from knowledgebase.log import Log
|
from knowledgebase import utils
|
|
|
class DocProcessor:
|
def __init__(self, docx_file: str):
|
"""
|
文档处理
|
:param docx_file: 要处理的文档
|
"""
|
Log.info(f'开始处理文档：{docx_file}')
|
self.docx_file = docx_file
|
self.doc_type = self.get_doc_type()
|
self.doc_split = DocSplit(docx_file, self.doc_type)
|
self.entity_recognition = EntityRecognition(self.doc_type)
|
self.doc_id = 0
|
|
def get_doc_type(self):
|
entity_helper = get_entity_helper()
|
Log.info(f'识别文档类型：{self.docx_file}')
|
rules = '；\n'.join([f'- {it}：{entity_helper.doc_prompt_map[it]}' for it in entity_helper.doc_prompt_map.keys()])
|
msg = HumanMessage(f'''
|
# 指令
|
请从下面的文件名中识别文档类型，如果识别失败不要输出任何字符。
|
文件名：{os.path.basename(self.docx_file)}
|
# 识别规则
|
{rules}
|
# 示例
|
遥测大纲
|
''')
|
resp = llm.invoke([msg])
|
Log.info(f'识别结果：{resp.content}')
|
return resp.content
|
|
async def get_tc_info(self, paragraph: ParagraphInfo):
|
if self.doc_type not in [DocType.tc_format]:
|
return ''
|
prompt = HumanMessagePromptTemplate.from_template('''
|
# 指令
|
请从下面的文本中识别指令信息，如果识别失败不要输出任何字符。
|
指令信息包括：指令名称。
|
# 识别规则
|
- 文本内容为遥控指令数据域或遥控指令应用数据的定义描述。
|
# 约束
|
- 如果文本内容是目录则不要输出任何字符；
|
- 指令名称在章节标题中，提取指令名称要和文本中的严格一致；
|
- 如果没有识别到指令信息不要输出任何字符；
|
- 识别失败，不要输出任何内容，包括解释性文本；
|
- 输出json格式。
|
# 示例 - 识别到指令
|
{{
|
"name": "xxx"
|
}}
|
# 示例 - 未识别到指令
|
""
|
# 文本内容：
|
{text}
|
''')
|
chain = prompt.prompt | llm | JsonOutputParser()
|
resp = await chain.ainvoke({"text": paragraph.full_text})
|
import json
|
# Log.info(f'>>>>>>指令识别：\n{paragraph.full_text}')
|
# Log.info(f'<<<<<<指令：{json.dumps(resp, ensure_ascii=False)}')
|
return resp
|
|
async def get_tm_pkt_info(self, paragraph: ParagraphInfo):
|
if self.doc_type not in [DocType.tm_outline, DocType.tm_pkt_design]:
|
return ''
|
prompt = HumanMessagePromptTemplate.from_template('''
|
# 指令
|
识别遥测包信息，请从下面的文本中识别遥测包信息，如果识别失败不要输出任何字符。
|
识别规则：章节标题中包含包名称和代号，章节内容为表格，表格中包括包头定义和包参数定义。
|
提取的遥测包信息包括：包名称，包代号。
|
# 约束
|
- 如果文本内容是目录则不要输出任何字符；
|
- 文本描述的内容是单个遥测包，如果有多个遥测包则不要输出任何字符；
|
- 文本结构通常是：包名称、代号和APID(应用过程标识)在开头（应用过程标识也有可能在表格中），后面紧接着是包头和参数定义表；
|
- 如果没有识别到遥测包信息不要输出任何字符；
|
- 识别失败，不要输出任何内容，包括解释性文本；
|
- 输出json格式。
|
# 符合要求的文本结构1
|
1.1.1 code xxx包(APID=0x123)
|
```json
|
表格内容
|
```
|
# 符合要求的文本结构2
|
1.1.1 code xxx包
|
```json
|
表格内容
|
应用过程标识
|
...
|
```
|
# 示例 - 识别到数据包
|
{{
|
"name": "xxx包",
|
"code": "TMS001"
|
}}
|
# 示例 - 未识别到数据包
|
""
|
# 文本内容：
|
{text}
|
''')
|
chain = prompt.prompt | llm | JsonOutputParser()
|
resp = await chain.ainvoke({"text": paragraph.full_text})
|
return resp
|
|
async def get_chapter_refs(self, paragraph: ParagraphInfo, toc: [str]) -> [str]:
|
if self.doc_type not in [DocType.tc_format]:
|
return ''
|
toc_text = '\n'.join(toc)
|
prompt = HumanMessagePromptTemplate.from_template(f'''
|
# 角色
|
你是一名资深的软件工程师。
|
# 指令
|
帮助我完成对文本中引用关系的抽取，判断当前文本中是否包含了引用信息，例如包含以下关键字：“详见1.1”、“见1.1”、“具体见1.1”、“见附录”等。
|
如果包含引用，将引用与“目录内容”中的目录条目进行匹配。
|
将匹配到的目录条目输出，输出格式为json格式。
|
# 约束
|
- 是否包含引用的判断条件中必须包含引用相关的描述，例如：“详见1.1”、“见1.1”、“具体见1.1”、“见附录”等；
|
- 注意不要自己引用自己；
|
- 仅提取目录内容中包含的条目，如果目录内容不包含则不提取；
|
- 如果仅靠标题号码无法确定目录条目的，根据文本内容匹配对应的目录条目；
|
- 输出的内容必须是目录中的条目；
|
- 输出json格式，不要输出任何json以外的字符。
|
# 输出案例
|
["1.1 xxx"]
|
# 目录内容：
|
{toc_text}
|
# 文本内容：
|
{{text}}
|
''')
|
chain = prompt.prompt | llm | JsonOutputParser()
|
resp = await chain.ainvoke({"text": paragraph.full_text})
|
return resp
|
|
async def gen_chapter_entities(self, paragraph: ParagraphInfo, paragraphs: [ParagraphInfo], toc: [str]):
|
# 获取章节实体词
|
entity_names_task = self.entity_recognition.run(paragraph.full_text)
|
# 获取指令信息
|
cmd_task = self.get_tc_info(paragraph)
|
# 获取遥测包信息
|
pkt_task = self.get_tm_pkt_info(paragraph)
|
# 获取文档引用
|
refs_task = self.get_chapter_refs(paragraph, toc)
|
entity_names, cmd, pkt, chapter_refs = await asyncio.gather(entity_names_task, cmd_task, pkt_task, refs_task)
|
|
Log.info(f'章节{paragraph.title_num}实体词：{entity_names}')
|
Log.info(f'章节{paragraph.title_num}引用：{chapter_refs}')
|
if entity_names:
|
paragraph.entities = doc_dbh.get_entities_by_names(entity_names)
|
|
if pkt:
|
entity = TEntity(name=pkt['code'], type='遥测包配置', prompts='', doc_type='')
|
e = doc_dbh.get_entity(entity)
|
if e:
|
entity.id = e.id
|
else:
|
doc_dbh.add_entity(entity)
|
Log.info(f"新增Entity：{entity.name}，id：{entity.id}")
|
paragraph.entities.append(entity)
|
|
if cmd:
|
entity = TEntity(name=cmd['name'], type='指令格式配置', prompts='', doc_type='')
|
e = doc_dbh.get_entity(entity)
|
if e:
|
entity.id = e.id
|
else:
|
doc_dbh.add_entity(entity)
|
Log.info(f"新增Entity：{entity.name}，id：{entity.id}")
|
paragraph.entities.append(entity)
|
# 获取引用信息
|
if chapter_refs:
|
for ref in chapter_refs:
|
_p = next(filter(lambda p: ref == p.title, self.doc_split.paragraphs), None)
|
if _p:
|
if paragraph != _p:
|
paragraph.refs.append(_p)
|
|
def process(self):
|
self.doc_split.split()
|
# 分批并发处理，每批10个
|
tasks = []
|
toc = []
|
for p in self.doc_split.paragraphs:
|
if p.title_level:
|
toc.append(p.title)
|
for paragraph in self.doc_split.paragraphs:
|
tasks.append(self.gen_chapter_entities(paragraph, self.doc_split.paragraphs, toc))
|
|
async def run():
|
await asyncio.gather(*tasks)
|
|
asyncio.run(run())
|
# 保存到数据库
|
self.save_to_db()
|
|
def save_to_db(self):
|
"""
|
保存段落和段落实体词关系到数据库。
|
"""
|
Log.info('保存段落和段落实体词关系到数据库...')
|
with open(self.docx_file, 'rb') as f:
|
file_bytes = f.read()
|
md5 = utils.generate_bytes_md5(file_bytes)
|
doc = DocInfo(os.path.basename(self.docx_file), md5, self.doc_type, self.doc_split.paragraph_tree)
|
self.doc_id = doc_dbh.add_doc(doc)
|
for paragraph in doc.paragraphs:
|
doc_dbh.add_paragraph(self.doc_id, None, paragraph)
|
for paragraph in self.doc_split.paragraphs:
|
for ref_paragraph in paragraph.refs:
|
doc_dbh.add_paragraph_ref_link(paragraph.id, ref_paragraph.id)
|
Log.info(f"{paragraph.title} 引用了-> {ref_paragraph.title}")
|
Log.info('保存段落和段落实体词关系到数据库完成')
|
