1150 lines
47 KiB
Python
1150 lines
47 KiB
Python
import sys
|
||
from pathlib import Path
|
||
import asyncio
|
||
import uuid
|
||
|
||
from docxcompose.composer import Composer
|
||
from docx import Document as Document_compose
|
||
from docxtpl import DocxTemplate
|
||
from typing import List, Optional
|
||
import json
|
||
import os
|
||
import re
|
||
from datetime import datetime, date
|
||
|
||
from fastapi import HTTPException
|
||
from md2docx.converter import MarkdownToDocx
|
||
from sqlalchemy import select, func, update, delete
|
||
from agents.ai_agents.task_hub.schemas.task_schema import ProjectTaskRecordCreate, TaskTypeEnum, TaskStatusEnum
|
||
from agents.ai_agents.task_hub.services.task_service import async_create_or_update_task_record_service
|
||
from agents.config.database import AsyncSessionLocal
|
||
from agents.interface.service.project_service import async_get_project_by_id_service
|
||
from agents.utils.document_parser import DocumentParser
|
||
|
||
from agents.ai_agents.document_agent.schemas.document_schema import DocumentGenerationRequest, DocumentCategoryEnum, \
|
||
DocumentSSEData,GenerateCheckRoleRequest
|
||
from agents.ai_agents.document_agent.services.document_service import DOCUMENT_TYPES
|
||
from agents.interface.schemas.project_file_link_schema import FileLinkSchema, CreateProjectFileLinkSchema
|
||
from agents.interface.service.project_file_link_service import async_create_file_link
|
||
from agents.models.base.project_model import FileEntity
|
||
from agents.utils.minio_client import minio_client
|
||
from agents.utils.openai_client import get_open_ai_client
|
||
from agents.ai_agents.document_agent.services.code_parser_service import CodeParserService
|
||
from agents.utils.run_command import run_command
|
||
from agents.ai_agents.prompt_template.models.prompt_template_model import ProjectTemplateEntity
|
||
from agents.ai_agents.prompt_template.services.project_template_custom_service import ProjectTemplateCustomService
|
||
from agents.ai_agents.test_verification_agent.services.task_service import TaskService
|
||
from agents.ai_agents.diagram_agent.utils.drawio_utils import replace_drawio_host
|
||
|
||
|
||
_STREAM_END = object()
|
||
|
||
|
||
def _next_stream_chunk(iterator):
|
||
try:
|
||
return next(iterator)
|
||
except StopIteration:
|
||
return _STREAM_END
|
||
|
||
|
||
async def _iterate_chat_completion_stream(openai_client, messages):
|
||
"""
|
||
The OpenAI wrapper exposes a blocking sync generator. Running next() in a
|
||
worker thread keeps concurrent SSE requests from blocking the event loop.
|
||
"""
|
||
iterator = openai_client.chat_completion_stream([dict(message) for message in messages])
|
||
while True:
|
||
chunk = await asyncio.to_thread(_next_stream_chunk, iterator)
|
||
if chunk is _STREAM_END:
|
||
break
|
||
yield chunk
|
||
|
||
|
||
class DocumentGenerationService:
|
||
"""文档生成服务类"""
|
||
|
||
@classmethod
|
||
def load_template_prompt(cls, file_name: str) -> str:
|
||
"""
|
||
加载系统提示词模板
|
||
|
||
:return: 系统提示词模板内容
|
||
"""
|
||
template_path = os.path.join(
|
||
os.path.dirname(os.path.dirname(__file__)),
|
||
"templates",
|
||
file_name
|
||
)
|
||
with open(template_path, 'r', encoding='utf-8') as f:
|
||
return f.read()
|
||
|
||
async def generate_stream_event_check_role(
|
||
request: GenerateCheckRoleRequest,
|
||
):
|
||
"""
|
||
根据项目文档模板id，
|
||
获取到模板信息，
|
||
调用ai生成校验规则文档
|
||
"""
|
||
|
||
|
||
# 调用OpenAI流式生成
|
||
full_content = ""
|
||
try:
|
||
project_template = await async_get_project_template_info(request.templateId)
|
||
project = await async_get_project_by_id_service(project_template.project_id)
|
||
content = project_template.content
|
||
name = project_template.name
|
||
file_name = name + "校验规则"
|
||
prompt = load_check_role_prompt(content, name)
|
||
messages = [
|
||
{"role": "system", "content": prompt},
|
||
{"role": "user", "content": ""}
|
||
]
|
||
|
||
start_event = DocumentSSEData(
|
||
eventType="start",
|
||
message=f"开始生成{file_name}",
|
||
documentName=None
|
||
)
|
||
yield f"data: {json.dumps(start_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
openai_client = get_open_ai_client(project.api_key, project.base_url, project.model_name)
|
||
async for chunk in _iterate_chat_completion_stream(openai_client, messages):
|
||
full_content += chunk
|
||
# 发送内容块事件
|
||
content_event = DocumentSSEData(
|
||
eventType="document_generating",
|
||
message=chunk,
|
||
documentName=file_name
|
||
)
|
||
yield f"data: {json.dumps(content_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
error_event = DocumentSSEData(
|
||
eventType="end",
|
||
message="校验规则生成成功",
|
||
documentName=file_name
|
||
)
|
||
yield f"data: {json.dumps(error_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
|
||
except Exception as e:
|
||
error_event = DocumentSSEData(
|
||
eventType="error",
|
||
message=f"文档生成失败， 异常信息：{str(e)}",
|
||
documentName=file_name
|
||
)
|
||
yield f"data: {json.dumps(error_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
|
||
|
||
async def generate_stream_event_generator(
|
||
request: DocumentGenerationRequest,
|
||
# db: Session
|
||
):
|
||
"""
|
||
SSE流式事件生成器
|
||
|
||
:param request: 文档生成请求
|
||
:param db: 数据库会话
|
||
:return: SSE事件流
|
||
"""
|
||
|
||
generation_request_id = uuid.uuid4().hex
|
||
project = await async_get_project_by_id_service(request.projectId)
|
||
if not project:
|
||
error_event = DocumentSSEData(
|
||
eventType="end",
|
||
message="项目不存在",
|
||
documentName=None,
|
||
documentIndex=0
|
||
)
|
||
yield f"data: {json.dumps(error_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
return
|
||
# 3. 创建任务记录
|
||
task_record_data = ProjectTaskRecordCreate(
|
||
project_id=request.projectId,
|
||
task_type=TaskTypeEnum.GENERATE_DOCUMENT,
|
||
task_status=TaskStatusEnum.RUNNING,
|
||
task_log="开始生成文档"
|
||
)
|
||
task_record_result = await async_create_or_update_task_record_service(task_record_data)
|
||
task_id = task_record_result.data.id
|
||
try:
|
||
# 发送开始事件
|
||
start_event = DocumentSSEData(
|
||
eventType="start",
|
||
# message=f"开始生成 {len(request.produceDocumentList)} 个文档",
|
||
message=f"开始生成 {len(request.projectTemplateIdList)} 个文档",
|
||
documentName=None,
|
||
documentIndex=0
|
||
)
|
||
yield f"data: {json.dumps(start_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
# 获取参考文档内容
|
||
reference_content = ""
|
||
sys_data = ""
|
||
if request.referToDocument is not None and len(request.referToDocument) > 0:
|
||
analysis_start_event = DocumentSSEData(
|
||
eventType="analysisDocumentStart",
|
||
message="开始分析参考文档",
|
||
documentName=None,
|
||
documentIndex=0
|
||
)
|
||
yield f"data: {json.dumps(analysis_start_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
# 使用非流式版本获取参考文档内容
|
||
reference_content = await async_handle_refer_document(request, project)
|
||
# 输出分析过程中的事件（模拟）
|
||
analysis_complete_event = DocumentSSEData(
|
||
eventType="analysisDocumentEnd",
|
||
message="参考文档分析完成",
|
||
documentName=None,
|
||
documentIndex=0
|
||
)
|
||
yield f"data: {json.dumps(analysis_complete_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
# 逐个生成文档
|
||
save_document_list: List[FileLinkSchema] = []
|
||
for index, project_template_id in enumerate(request.projectTemplateIdList, 1):
|
||
# 获取文档模板信息
|
||
project_template = await async_get_project_template_info(project_template_id)
|
||
doc_name = project_template.name
|
||
check_doc_name = doc_name + "校验报告"
|
||
project_check_role = project_template.check_role
|
||
project_template_content = project_template.content
|
||
custom_template_content = await ProjectTemplateCustomService.get_latest_custom_content(
|
||
request.projectId,
|
||
project_template_id
|
||
)
|
||
if custom_template_content is not None:
|
||
project_template_content = custom_template_content
|
||
category = project_template.category
|
||
# 发送文档开始事件
|
||
doc_start_event = DocumentSSEData(
|
||
eventType="document_start",
|
||
message=f"开始生成文档: {doc_name}",
|
||
documentName=doc_name,
|
||
documentIndex=index
|
||
)
|
||
yield f"data: {json.dumps(doc_start_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
# 获取编写提示词模板
|
||
system_prompt = load_prompt(request, sys_data, doc_name, reference_content, project_template_content)
|
||
|
||
messages = [
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": f"{request.userInput}"}
|
||
]
|
||
# 如果选择了仓库信息，可以将仓库先解析，之后再填充至仓库的解析结果
|
||
if request.repoInfo is not None:
|
||
if request.repoInfo.username and request.repoInfo.password:
|
||
# 在URL中插入用户名和密码
|
||
from urllib.parse import urlparse
|
||
parsed = urlparse(request.repoInfo.url)
|
||
auth_url = (f"{parsed.scheme}://{request.repoInfo.username}"
|
||
f":{request.repoInfo.password}@{parsed.netloc}{parsed.path}")
|
||
repo_url = auth_url
|
||
coed_parser_service = CodeParserService()
|
||
parse_result = coed_parser_service.parse_code_repository(
|
||
repo_url=repo_url, branch=request.repoInfo.branch, use_doxygen=True)
|
||
prompt = (coed_parser_service.build_architecture_prompt
|
||
(parse_result=parse_result, project_name=project.name))
|
||
messages.append({
|
||
"role": "user", "content": f"工程代码信息如下：{prompt}"
|
||
})
|
||
# 调用OpenAI流式生成
|
||
full_content = ""
|
||
try:
|
||
openai_client = get_open_ai_client(project.api_key, project.base_url, project.model_name)
|
||
async for chunk in _iterate_chat_completion_stream(openai_client, messages):
|
||
full_content += chunk
|
||
# 发送内容块事件
|
||
content_event = DocumentSSEData(
|
||
eventType="document_generating",
|
||
message=chunk,
|
||
documentName=doc_name,
|
||
documentIndex=index
|
||
)
|
||
yield f"data: {json.dumps(content_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
except Exception as e:
|
||
error_event = DocumentSSEData(
|
||
eventType="error",
|
||
message=f"文档生成失败: {doc_name}， 异常信息：{str(e)}",
|
||
documentName=doc_name,
|
||
documentIndex=index
|
||
)
|
||
update_task_data = ProjectTaskRecordCreate(
|
||
id=task_id,
|
||
project_id=request.projectId,
|
||
task_type=TaskTypeEnum.GENERATE_DOCUMENT,
|
||
task_status=TaskStatusEnum.ERROR,
|
||
task_log=f"文档生成失败: {doc_name}， 异常信息：{str(e)}"
|
||
)
|
||
await async_create_or_update_task_record_service(update_task_data)
|
||
yield f"data: {json.dumps(error_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
continue
|
||
|
||
# 开始将文档转换为md格式
|
||
convert_succeeded = False
|
||
convert_failed = False
|
||
for event in handle_conver_file(
|
||
doc_name,
|
||
full_content,
|
||
project,
|
||
save_document_list,
|
||
request,
|
||
category,
|
||
generation_request_id
|
||
):
|
||
if event.eventType == "convert_complete":
|
||
convert_succeeded = True
|
||
elif event.eventType == "convert_error":
|
||
convert_failed = True
|
||
yield f"data: {json.dumps(event.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
if convert_failed or not convert_succeeded:
|
||
error_message = f"文档转换失败: {doc_name}"
|
||
update_task_data = ProjectTaskRecordCreate(
|
||
id=task_id,
|
||
project_id=request.projectId,
|
||
task_type=TaskTypeEnum.GENERATE_DOCUMENT,
|
||
task_status=TaskStatusEnum.ERROR,
|
||
task_log=error_message
|
||
)
|
||
await async_create_or_update_task_record_service(update_task_data)
|
||
error_event = DocumentSSEData(
|
||
eventType="error",
|
||
message=error_message,
|
||
documentName=doc_name,
|
||
documentIndex=index
|
||
)
|
||
yield f"data: {json.dumps(error_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
continue
|
||
|
||
# 发送文档完成事件
|
||
doc_complete_event = DocumentSSEData(
|
||
eventType="document_complete",
|
||
message=f"文档生成完成: {doc_name}",
|
||
documentName=doc_name,
|
||
documentIndex=index
|
||
)
|
||
update_task_data = ProjectTaskRecordCreate(
|
||
id=task_id,
|
||
project_id=request.projectId,
|
||
task_type=TaskTypeEnum.GENERATE_DOCUMENT,
|
||
task_status=TaskStatusEnum.PENDING,
|
||
task_log=f"文档生成完成: {doc_name}"
|
||
)
|
||
await async_create_or_update_task_record_service(update_task_data)
|
||
yield f"data: {json.dumps(doc_complete_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
# 文档校验，生成校验报告
|
||
if request.check:
|
||
if not project_check_role:
|
||
check_start_error = DocumentSSEData(
|
||
eventType="check_start_error",
|
||
message=f"文档校验规则不存在: {doc_name}",
|
||
documentName=doc_name,
|
||
documentIndex=index
|
||
)
|
||
yield f"data: {json.dumps(check_start_error.to_dict(), ensure_ascii=False)}\n\n"
|
||
continue
|
||
|
||
check_start = DocumentSSEData(
|
||
eventType="check_start",
|
||
message=f"开始生成校验报告: {check_doc_name}",
|
||
documentName=doc_name,
|
||
documentIndex=index
|
||
)
|
||
yield f"data: {json.dumps(check_start.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
check_prompt = load_check_prompt(full_content, project_check_role, doc_name)
|
||
check_messages = [
|
||
{"role": "system", "content": check_prompt},
|
||
{"role": "user", "content": "请按照规则生成校验报告"}
|
||
]
|
||
check_full_content = ""
|
||
try:
|
||
openai_client = get_open_ai_client(project.api_key, project.base_url, project.model_name)
|
||
async for chunk in _iterate_chat_completion_stream(openai_client, check_messages):
|
||
check_full_content += chunk
|
||
# 发送内容块事件
|
||
content_event = DocumentSSEData(
|
||
eventType="document_checking",
|
||
message=chunk,
|
||
documentName=check_doc_name,
|
||
documentIndex=index
|
||
)
|
||
yield f"data: {json.dumps(content_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
except Exception as e:
|
||
error_event = DocumentSSEData(
|
||
eventType="error",
|
||
message=f"校验报告生成失败: {check_doc_name}， 异常信息：{str(e)}",
|
||
documentName=doc_name,
|
||
documentIndex=index
|
||
)
|
||
update_task_data = ProjectTaskRecordCreate(
|
||
id=task_id,
|
||
project_id=request.projectId,
|
||
task_type=TaskTypeEnum.GENERATE_DOCUMENT,
|
||
task_status=TaskStatusEnum.ERROR,
|
||
task_log=f"校验报告生成失败: {check_doc_name}， 异常信息：{str(e)}"
|
||
)
|
||
await async_create_or_update_task_record_service(update_task_data)
|
||
yield f"data: {json.dumps(error_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
continue
|
||
check_convert_succeeded = False
|
||
check_convert_failed = False
|
||
for event in handle_conver_file(
|
||
check_doc_name,
|
||
check_full_content,
|
||
project,
|
||
save_document_list,
|
||
request,
|
||
category,
|
||
generation_request_id
|
||
):
|
||
if event.eventType == "convert_complete":
|
||
check_convert_succeeded = True
|
||
elif event.eventType == "convert_error":
|
||
check_convert_failed = True
|
||
yield f"data: {json.dumps(event.to_dict(), ensure_ascii=False)}\n\n"
|
||
if check_convert_failed or not check_convert_succeeded:
|
||
error_message = f"校验报告转换失败: {check_doc_name}"
|
||
update_task_data = ProjectTaskRecordCreate(
|
||
id=task_id,
|
||
project_id=request.projectId,
|
||
task_type=TaskTypeEnum.GENERATE_DOCUMENT,
|
||
task_status=TaskStatusEnum.ERROR,
|
||
task_log=error_message
|
||
)
|
||
await async_create_or_update_task_record_service(update_task_data)
|
||
error_event = DocumentSSEData(
|
||
eventType="error",
|
||
message=error_message,
|
||
documentName=check_doc_name,
|
||
documentIndex=index
|
||
)
|
||
yield f"data: {json.dumps(error_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
continue
|
||
# 发送文档完成事件
|
||
doc_complete_event = DocumentSSEData(
|
||
eventType="check_document_complete",
|
||
message=f"校验报告生成完成: {check_doc_name}",
|
||
documentName=check_doc_name,
|
||
documentIndex=index
|
||
)
|
||
update_task_data = ProjectTaskRecordCreate(
|
||
id=task_id,
|
||
project_id=request.projectId,
|
||
task_type=TaskTypeEnum.GENERATE_DOCUMENT,
|
||
task_status=TaskStatusEnum.PENDING,
|
||
task_log=f"校验报告生成完成: {check_doc_name}"
|
||
)
|
||
await async_create_or_update_task_record_service(update_task_data)
|
||
yield f"data: {json.dumps(doc_complete_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
if save_document_list is not None and len(save_document_list) > 0:
|
||
try:
|
||
file_link_schema = CreateProjectFileLinkSchema(
|
||
project_id=request.projectId,
|
||
file_list=save_document_list,
|
||
# type=None
|
||
)
|
||
# 创建附件关联
|
||
links = await async_create_file_link(file_link_schema)
|
||
|
||
# 添加任务文件关联记录
|
||
if request.taskId and links:
|
||
file_ids = [link.file_id for link in links]
|
||
await TaskService.async_add_task_file_links(request.taskId, file_ids)
|
||
except Exception as e:
|
||
error_event = DocumentSSEData(
|
||
eventType="error",
|
||
message=f"文件上传失败: {e}",
|
||
documentName=None,
|
||
documentIndex=0
|
||
)
|
||
update_task_data = ProjectTaskRecordCreate(
|
||
id=task_id,
|
||
project_id=request.projectId,
|
||
task_type=TaskTypeEnum.GENERATE_DOCUMENT,
|
||
task_status=TaskStatusEnum.ERROR,
|
||
task_log=f"文件上传失败: {e}"
|
||
)
|
||
await async_create_or_update_task_record_service(update_task_data)
|
||
yield f"data: {json.dumps(error_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
# 发送结束事件
|
||
end_event = DocumentSSEData(
|
||
eventType="end",
|
||
message="",
|
||
documentName=None,
|
||
documentIndex=len(request.produceDocumentList)
|
||
)
|
||
update_task_data = ProjectTaskRecordCreate(
|
||
id=task_id,
|
||
project_id=request.projectId,
|
||
task_type=TaskTypeEnum.GENERATE_DOCUMENT,
|
||
task_status=TaskStatusEnum.COMPLETED,
|
||
task_log=f"文档生成成功"
|
||
)
|
||
await async_create_or_update_task_record_service(update_task_data)
|
||
yield f"data: {json.dumps(end_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
except Exception as e:
|
||
error_event = DocumentSSEData(
|
||
eventType="error",
|
||
message=f"生成过程出错: {str(e)}",
|
||
documentName=None,
|
||
documentIndex=0
|
||
)
|
||
yield f"data: {json.dumps(error_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
end_event = DocumentSSEData(
|
||
eventType="end",
|
||
message="",
|
||
documentName=None,
|
||
documentIndex=len(request.produceDocumentList)
|
||
)
|
||
yield f"data: {json.dumps(end_event.to_dict(), ensure_ascii=False)}\n\n"
|
||
|
||
|
||
async def async_get_project_template_info(project_template_id) -> ProjectTemplateEntity:
|
||
async with AsyncSessionLocal() as db:
|
||
result = await db.execute(
|
||
select(ProjectTemplateEntity).where(ProjectTemplateEntity.id == project_template_id)
|
||
)
|
||
|
||
project_template = result.scalar_one_or_none()
|
||
if not project_template:
|
||
raise HTTPException(status_code=404, detail=f"项目模板不存在: {project_template_id}")
|
||
if not project_template.content:
|
||
raise HTTPException(status_code=404, detail=f"项目模板内容不存在: {project_template.name}")
|
||
return project_template
|
||
|
||
|
||
def load_prompt(request, sys_data, doc_name, reference_content, project_template_content) -> str:
|
||
document_template_prompt = ""
|
||
# project_template_content 为空就读内置文件
|
||
if project_template_content is None or project_template_content == '':
|
||
document_template_prompt = DocumentGenerationService.load_template_prompt(file_name=doc_name + ".md")
|
||
else:
|
||
document_template_prompt = project_template_content
|
||
# 构建消息
|
||
system_prompt = DocumentGenerationService.load_template_prompt(file_name="document_system.md")
|
||
system_prompt = system_prompt.replace("{SYSTEM_DATA}", sys_data)
|
||
system_prompt = system_prompt.replace("{USER_SUPPLEMENT_QUANTITY}", request.userInput)
|
||
system_prompt = system_prompt.replace("{WRITE_TEMPLATE}", document_template_prompt)
|
||
system_prompt = system_prompt.replace("{REFERENCE_CONTENT}", reference_content)
|
||
system_prompt = system_prompt.replace("{USER_EXTRA_CONSTRAINT}", "无额外约束")
|
||
system_prompt = system_prompt.replace("{OTHER_WRITE_CONSTRAINT}", f"国军标选项: {request.GJBSelectOption}")
|
||
system_prompt = system_prompt.replace("{CURRENT_DATE}", f'{date.today().strftime("%Y年%m月%d日")}')
|
||
return system_prompt
|
||
|
||
|
||
def load_check_prompt(document_content, check_role_connect, doc_name) -> str:
|
||
"""构造校验文档的提示词信息"""
|
||
check_prompt = DocumentGenerationService.load_template_prompt(file_name="document_check.md")
|
||
check_prompt = check_prompt.replace("{AI_GENERATED_DOCUMENT}", document_content)
|
||
check_prompt = check_prompt.replace("{CHECK_DOCUMENT}", check_role_connect)
|
||
check_prompt = check_prompt.replace("{WRITE_TEMPLATE}", doc_name)
|
||
return check_prompt
|
||
|
||
def load_check_role_prompt(document_content, document_name) -> str:
|
||
"""获取生成校验规则文档的提示词信息"""
|
||
check_prompt = DocumentGenerationService.load_template_prompt(file_name="generate_check_role.md")
|
||
check_prompt = check_prompt.replace("{TEMPLATE_DOCUMENT}", document_content)
|
||
check_prompt = check_prompt.replace("{DOCUMENT_NAME}", document_name)
|
||
return check_prompt
|
||
|
||
|
||
async def handle_refer_document(request, project) -> str:
|
||
"""处理参考文档，返回分析后的内容（不流式输出）"""
|
||
content_map = {}
|
||
openai_client = get_open_ai_client(project.api_key, project.base_url, project.model_name)
|
||
async with AsyncSessionLocal() as db:
|
||
result = await db.execute(
|
||
select(FileEntity).where(FileEntity.id.in_(request.referToDocument))
|
||
)
|
||
file_list = result.scalars().all()
|
||
|
||
for doc in file_list:
|
||
if doc:
|
||
# 从Minio下载并解析文档
|
||
try:
|
||
text_content = DocumentParser.parse_all_document(doc, openai_client=openai_client)
|
||
content_map[doc.name] = text_content
|
||
except Exception as e:
|
||
raise Exception("文档分析失败", str(e))
|
||
return json.dumps(content_map, ensure_ascii=False)
|
||
# 组装提示词
|
||
# remote_prompt = DocumentGenerationService.load_template_prompt(file_name="remote_analysis.md")
|
||
# remote_prompt = remote_prompt.replace("{ANALYSIS_CONTENT}", json.dumps(content_map))
|
||
#
|
||
# messages = [
|
||
# {"role": "system", "content": remote_prompt},
|
||
# {"role": "user", "content": "请你帮我分析以下内容"}
|
||
# ]
|
||
# try:
|
||
# # 使用非流式方式获取完整内容
|
||
# full_content = openai_client.chat_completion(messages)
|
||
# return full_content
|
||
# except Exception as e:
|
||
# error_msg = f"生成文档失败: {str(e)}"
|
||
# raise Exception("文档分析失败：", error_msg)
|
||
|
||
|
||
async def async_handle_refer_document(request, project) -> str:
|
||
content_map = {}
|
||
openai_client = get_open_ai_client(project.api_key, project.base_url, project.model_name)
|
||
async with AsyncSessionLocal() as db:
|
||
result = await db.execute(
|
||
select(FileEntity).where(FileEntity.id.in_(request.referToDocument))
|
||
)
|
||
file_list = result.scalars().all()
|
||
|
||
for doc in file_list:
|
||
if doc:
|
||
try:
|
||
text_content = DocumentParser.parse_all_document(doc, openai_client=openai_client)
|
||
content_map[doc.name] = text_content
|
||
except Exception as e:
|
||
raise Exception("文档分析失败", str(e))
|
||
return json.dumps(content_map, ensure_ascii=False)
|
||
|
||
|
||
# draw.io CLI 路径解析
|
||
def _resolve_drawio_cli() -> Optional[str]:
|
||
"""解析 Draw.io 可执行文件路径（不含 xvfb-run）。"""
|
||
import logging
|
||
import platform
|
||
import shutil
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 1. 环境变量 DRAWIO_CLI_PATH
|
||
env_path = os.getenv("DRAWIO_CLI_PATH")
|
||
if env_path and os.path.isfile(env_path):
|
||
logger.info(f"使用环境变量 DRAWIO_CLI_PATH: {env_path}")
|
||
return env_path
|
||
|
||
# 2. 项目内默认路径（与 article_agent 保持一致）
|
||
default = os.path.normpath(os.path.join(
|
||
os.path.dirname(os.path.abspath(__file__)),
|
||
"..", "..", "..", "..", "..", "soft", "drawio", "draw.io", "draw.io.exe",
|
||
))
|
||
if os.path.isfile(default):
|
||
logger.info(f"使用项目默认路径: {default}")
|
||
return default
|
||
|
||
# 3. tool 目录下的 drawio.exe（兼容旧配置）
|
||
tool_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "tool/drawio.exe")
|
||
if os.path.isfile(tool_path):
|
||
logger.info(f"使用 tool 目录: {tool_path}")
|
||
return tool_path
|
||
|
||
# 4. PATH 查找
|
||
for name in ("draw.io.exe", "drawio", "draw.io"):
|
||
found = shutil.which(name)
|
||
if found:
|
||
logger.info(f"使用系统 PATH: {found}")
|
||
return found
|
||
|
||
logger.warning("Draw.io 未找到，请安装或设置 DRAWIO_CLI_PATH 环境变量")
|
||
return None
|
||
|
||
|
||
DRAWIO_EXE_PATH = _resolve_drawio_cli()
|
||
|
||
|
||
def extract_drawio_blocks(content: str) -> list:
|
||
"""
|
||
从Markdown内容中提取draw.io XML代码块
|
||
|
||
:param content: Markdown内容
|
||
:return: draw.io XML列表
|
||
"""
|
||
import re
|
||
pattern = r'```(?:drawio|xml)\n(.*?)\n```'
|
||
matches = re.findall(pattern, content, re.DOTALL)
|
||
# 替换每个 XML 块中的 host
|
||
replaced_matches = [replace_drawio_host(match) for match in matches]
|
||
return replaced_matches
|
||
|
||
|
||
def convert_drawio_to_png(drawio_xml: str, output_path: str, output_dir) -> bool:
|
||
"""
|
||
使用draw.io将XML转换为PNG
|
||
|
||
:param output_dir: 输出目录
|
||
:param drawio_xml: draw.io XML内容
|
||
:param output_path: 输出PNG文件路径
|
||
:return: 是否转换成功
|
||
"""
|
||
import tempfile
|
||
import logging
|
||
import re
|
||
import platform
|
||
import shutil
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 优化 drawio XML 中的画布大小，确保适配内容
|
||
page_width_match = re.search(r'pageWidth="(\d+)"', drawio_xml)
|
||
page_height_match = re.search(r'pageHeight="(\d+)"', drawio_xml)
|
||
|
||
if page_width_match and page_height_match:
|
||
page_width = int(page_width_match.group(1))
|
||
page_height = int(page_height_match.group(1))
|
||
# 如果画布过大（超过2000），调整为更合适的尺寸
|
||
if page_width > 1200 or page_height > 1000:
|
||
new_width = min(800, page_width)
|
||
new_height = min(600, page_height)
|
||
drawio_xml = drawio_xml.replace(
|
||
f'pageWidth="{page_width}"', f'pageWidth="{new_width}"'
|
||
).replace(
|
||
f'pageHeight="{page_height}"', f'pageHeight="{new_height}"'
|
||
)
|
||
|
||
with tempfile.NamedTemporaryFile(mode='w', suffix='.drawio', delete=False, encoding='utf-8') as f:
|
||
f.write(drawio_xml)
|
||
xml_file = f.name
|
||
|
||
try:
|
||
# 检查draw.io命令是否存在
|
||
if os.path.exists(DRAWIO_EXE_PATH):
|
||
drawio_cmd = DRAWIO_EXE_PATH
|
||
logger.info(f"使用内置draw.io: {DRAWIO_EXE_PATH}")
|
||
else:
|
||
drawio_cmd = shutil.which("draw.io") or shutil.which("drawio")
|
||
if not drawio_cmd:
|
||
logger.error("draw.io命令未找到，请确认已安装并添加到PATH")
|
||
return False
|
||
logger.info(f"使用系统draw.io: {drawio_cmd}")
|
||
|
||
if platform.system() == "Linux":
|
||
drawio_cmd = f"xvfb-run -a {drawio_cmd}"
|
||
|
||
command = (
|
||
f'{drawio_cmd} -x -f png '
|
||
f'-o "{output_path}" '
|
||
f'-b 0 '
|
||
f'--crop '
|
||
f'"{xml_file}"'
|
||
)
|
||
logger.info(f"执行draw.io转换命令: {command}")
|
||
run_command(command, output_dir)
|
||
|
||
# 检查输出文件是否生成
|
||
if os.path.exists(output_path):
|
||
file_size = os.path.getsize(output_path)
|
||
logger.info(f"draw.io转换成功，输出文件: {output_path}, 大小: {file_size} bytes")
|
||
return True
|
||
else:
|
||
logger.error(f"draw.io转换完成但输出文件不存在: {output_path}")
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"draw.io转换异常: {str(e)}", exc_info=True)
|
||
return False
|
||
finally:
|
||
if os.path.exists(xml_file):
|
||
os.unlink(xml_file)
|
||
|
||
|
||
def process_drawio_images(content: str, output_dir: str):
|
||
"""
|
||
处理Markdown中的draw.io代码块，转换为PNG图片引用
|
||
|
||
:param content: Markdown内容
|
||
:param output_dir: 输出目录
|
||
:return: 处理后的Markdown内容，draw.io块被替换为图片引用
|
||
"""
|
||
import re
|
||
import logging
|
||
logger = logging.getLogger(__name__)
|
||
|
||
pattern = r'```(?:drawio|xml)\n(.*?)\n```'
|
||
matches = re.findall(pattern, content, re.DOTALL)
|
||
total_images = len(matches)
|
||
|
||
if total_images == 0:
|
||
yield DocumentSSEData(
|
||
eventType="convert_processing",
|
||
message="无draw.io图片需要处理",
|
||
documentName=None,
|
||
documentIndex=0
|
||
)
|
||
yield content
|
||
return
|
||
|
||
yield DocumentSSEData(
|
||
eventType="convert_processing",
|
||
message=f"开始处理 {total_images} 个draw.io图片",
|
||
documentName=None,
|
||
documentIndex=0
|
||
)
|
||
|
||
parts = []
|
||
last_end = 0
|
||
image_index = 0
|
||
|
||
for match in re.finditer(pattern, content, re.DOTALL):
|
||
parts.append(content[last_end:match.start()])
|
||
last_end = match.end()
|
||
drawio_xml = match.group(1).strip()
|
||
|
||
# 替换 draw.io XML 中的 host
|
||
drawio_xml = replace_drawio_host(drawio_xml)
|
||
|
||
image_index += 1
|
||
|
||
yield DocumentSSEData(
|
||
eventType="convert_processing",
|
||
message=f"处理图片 {image_index}/{total_images}",
|
||
documentName=None,
|
||
documentIndex=0
|
||
)
|
||
|
||
png_filename = f"diagram_{image_index}.png"
|
||
png_path = os.path.join(output_dir, png_filename)
|
||
|
||
xml_filename = f"diagram_{image_index}.drawio"
|
||
xml_path = os.path.join(output_dir, xml_filename)
|
||
|
||
with open(xml_path, 'w', encoding='utf-8') as f:
|
||
f.write(drawio_xml)
|
||
|
||
success = convert_drawio_to_png(drawio_xml, png_path, output_dir)
|
||
if success:
|
||
logger.info(f"draw.io图片生成成功: {png_path}")
|
||
# 使用相对路径而不是绝对路径，避免转换问题
|
||
parts.append(f"")
|
||
yield DocumentSSEData(
|
||
eventType="convert_processing",
|
||
message=f"图片 {image_index}/{total_images} 生成完成",
|
||
documentName=None,
|
||
documentIndex=0
|
||
)
|
||
else:
|
||
logger.warning(f"draw.io图片生成失败: {png_filename}，继续处理")
|
||
# 即使转换失败，也保留占位符，避免破坏文档结构
|
||
parts.append(f"")
|
||
# 注意：这里不发送convert_error事件，因为图片转换失败不应影响整体转换结果
|
||
yield DocumentSSEData(
|
||
eventType="convert_processing",
|
||
message=f"图片 {image_index}/{total_images} 转换失败（已保留占位符）",
|
||
documentName=None,
|
||
documentIndex=0
|
||
)
|
||
|
||
parts.append(content[last_end:])
|
||
processed_content = "".join(parts)
|
||
|
||
yield DocumentSSEData(
|
||
eventType="convert_processing",
|
||
message=f"draw.io图片处理完成，共处理 {total_images} 张图片",
|
||
documentName=None,
|
||
documentIndex=0
|
||
)
|
||
|
||
yield processed_content
|
||
|
||
|
||
def sanitize_generated_markdown(content: str) -> str:
|
||
"""
|
||
Clean model output before Markdown->Word conversion.
|
||
|
||
Current fixes:
|
||
- drop template artifact lines like ``\3.1 接口标识和接口图``
|
||
- de-duplicate immediately repeated heading text caused by prompt examples
|
||
"""
|
||
cleaned_lines = []
|
||
previous_non_empty = ""
|
||
|
||
for raw_line in content.splitlines():
|
||
line = raw_line.rstrip()
|
||
stripped = line.strip()
|
||
|
||
if re.match(r"^\\\d+(?:\.\d+)*\s+", stripped):
|
||
normalized = stripped[1:].strip()
|
||
if normalized == previous_non_empty:
|
||
continue
|
||
# These are template residue lines, not valid final content.
|
||
continue
|
||
|
||
if stripped and stripped == previous_non_empty:
|
||
continue
|
||
|
||
cleaned_lines.append(line)
|
||
if stripped:
|
||
previous_non_empty = stripped
|
||
|
||
return "\n".join(cleaned_lines)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
base_path = os.path.dirname(os.path.abspath(sys.argv[0]))
|
||
if not os.path.exists(f"{base_path}/workspace"):
|
||
os.makedirs(f"{base_path}/workspace")
|
||
import uuid
|
||
current_temp_path = f"{base_path}/workspace/{uuid.uuid4()}"
|
||
os.makedirs(current_temp_path)
|
||
template_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "templates/template.docx")
|
||
# mermaid_filter_lua_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "tool/mermaid-filter.lua")
|
||
# 使用库来进行文档转换
|
||
|
||
converter = MarkdownToDocx(template_file, font_name="宋体")
|
||
docx_file_path = f"{current_temp_path}/aaa.docx"
|
||
with open(
|
||
"C:/workspace/ai/DocumentGenerateAgent/agents/ai_agents/workspace/1.md",
|
||
'r',
|
||
encoding="UTF-8") as f:
|
||
txt = f.read()
|
||
|
||
converter.convert(txt, docx_file_path)
|
||
# 转换图表的题注。
|
||
from agents.ai_agents.document_agent.utils import docx_util
|
||
|
||
docx_util.fix_table_borders(docx_file_path)
|
||
docx_util.add_captions(docx_file_path)
|
||
|
||
|
||
# for event1 in process_drawio_images(txt, "D:/workspace/agent/DocumentGenerateAgent/agents/workspace/test"):
|
||
# if isinstance(event1, str):
|
||
# processed_content1 = event1
|
||
# else:
|
||
# print(f'{event1}')
|
||
# convert_drawio_to_png(txt,
|
||
# "D:/workspace/agent/DocumentGenerateAgent/agents/workspace/bcb79d47-4b33-44fd-9aa6-7334ee714f57/2.png",
|
||
# "D:/workspace/agent/DocumentGenerateAgent/agents/workspace/bcb79d47-4b33-44fd-9aa6-7334ee714f57")
|
||
# with open("D:/workspace/agent/DocumentGenerateAgent/agents/ai_agents/document_agent/templates/txt.txt", 'r',
|
||
# encoding="UTF-8") as f:
|
||
# txt = f.read()
|
||
# process_drawio_images(txt,
|
||
# "D:/workspace/agent/DocumentGenerateAgent/agents/workspace/3ce4660d-e4f6-4fd3-a5d4-29fb0dd25a22/output")
|
||
|
||
|
||
# 处理基于模板转换的文件
|
||
def convert_template_docx(docx_file_path, project, doc_name, current_temp_path) -> str:
|
||
context_data = {
|
||
"projectName": project.name,
|
||
"name": doc_name,
|
||
}
|
||
# 加载模板
|
||
template_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "templates/base_template.docx")
|
||
template = DocxTemplate(template_file)
|
||
# 执行填充，生成一个填充好的中间文档
|
||
template.render(context_data)
|
||
# 保存中间文档，此文档包含了填充后的所有内容
|
||
filled_file = f'{current_temp_path}/filled_temp.docx'
|
||
template.save(filled_file)
|
||
|
||
# --- 第二阶段：将填充好的文档与其他文档合并 ---
|
||
# 创建一个基于填充后文档的合并器
|
||
base_doc = Document_compose(filled_file)
|
||
composer = Composer(base_doc)
|
||
|
||
# 添加其他需要合并的文档
|
||
other_files = [docx_file_path]
|
||
for file in other_files:
|
||
if os.path.exists(file):
|
||
doc_to_append = Document_compose(file)
|
||
composer.append(doc_to_append)
|
||
|
||
# 保存最终合并的文档
|
||
save_path = f'{current_temp_path}/{project.name}-{doc_name}.docx'
|
||
composer.save(save_path)
|
||
return save_path
|
||
|
||
|
||
# 处理文档转换
|
||
def handle_conver_file(doc_name, full_content, project, save_document, request, category, generation_request_id: str):
|
||
import logging
|
||
logger = logging.getLogger(__name__)
|
||
|
||
try:
|
||
base_path = os.path.dirname(os.path.abspath(sys.argv[0]))
|
||
logger.info(f"base_path: {base_path}")
|
||
if not os.path.exists(f"{base_path}/workspace"):
|
||
os.makedirs(f"{base_path}/workspace")
|
||
import uuid
|
||
current_temp_path = f"{base_path}/workspace/{generation_request_id}_{uuid.uuid4().hex}"
|
||
os.makedirs(current_temp_path)
|
||
|
||
yield DocumentSSEData(
|
||
eventType="convert_start",
|
||
message=f"开始转换文档: {doc_name}",
|
||
documentName=doc_name,
|
||
documentIndex=0
|
||
)
|
||
|
||
processed_content = None
|
||
for event in process_drawio_images(full_content, current_temp_path):
|
||
if isinstance(event, str):
|
||
processed_content = event
|
||
else:
|
||
yield event
|
||
|
||
if processed_content is None:
|
||
processed_content = full_content
|
||
processed_content = sanitize_generated_markdown(processed_content)
|
||
|
||
yield DocumentSSEData(
|
||
eventType="convert_processing",
|
||
message=f"Markdown文件生成完成",
|
||
documentName=doc_name,
|
||
documentIndex=0
|
||
)
|
||
|
||
md_file_path = f"{current_temp_path}/{doc_name}.md"
|
||
with open(md_file_path, mode='w', encoding="utf-8") as file:
|
||
file.write(processed_content)
|
||
|
||
template_file = os.path.join(os.path.dirname(os.path.dirname(__file__)), "templates/template.docx")
|
||
# mermaid_filter_lua_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "tool/mermaid-filter.lua")
|
||
# 使用库来进行文档转换
|
||
|
||
# 验证模板文件是否存在
|
||
if not os.path.exists(template_file):
|
||
logger.error(f"模板文件不存在: {template_file}")
|
||
yield DocumentSSEData(
|
||
eventType="convert_error",
|
||
message=f"模板文件不存在: {template_file}",
|
||
documentName=doc_name,
|
||
documentIndex=0
|
||
)
|
||
return
|
||
|
||
logger.info(f"开始Markdown转Word转换，文档名: {doc_name}, 模板路径: {template_file}")
|
||
logger.info(f"Markdown内容长度: {len(processed_content)} 字符")
|
||
|
||
converter = MarkdownToDocx(template_file, font_name="宋体")
|
||
docx_file_path = f"{current_temp_path}/{doc_name}.docx"
|
||
|
||
try:
|
||
converter.convert(processed_content, docx_file_path)
|
||
logger.info(f"Markdown转Word成功，输出文件: {docx_file_path}")
|
||
except Exception as convert_ex:
|
||
logger.error(f"Markdown转Word失败: {str(convert_ex)}", exc_info=True)
|
||
# 保存原始Markdown以便调试
|
||
debug_md_path = f"{current_temp_path}/{doc_name}_debug.md"
|
||
with open(debug_md_path, 'w', encoding='utf-8') as debug_f:
|
||
debug_f.write(processed_content)
|
||
logger.info(f"已保存调试用的Markdown文件: {debug_md_path}")
|
||
raise convert_ex
|
||
# 转换图表的题注。
|
||
from agents.ai_agents.document_agent.utils import docx_util
|
||
|
||
try:
|
||
logger.info(f"开始修复表格边框: {docx_file_path}")
|
||
docx_util.fix_table_borders(docx_file_path)
|
||
logger.info(f"开始添加题注: {docx_file_path}")
|
||
docx_util.add_captions(docx_file_path)
|
||
logger.info(f"开始修复表格单元格样式: {docx_file_path}")
|
||
docx_util.fix_table_cell_styles(docx_file_path)
|
||
logger.info(f"文档后处理完成: {docx_file_path}")
|
||
except Exception as postprocess_ex:
|
||
logger.error(f"文档后处理失败: {str(postprocess_ex)}", exc_info=True)
|
||
raise postprocess_ex
|
||
|
||
yield DocumentSSEData(
|
||
eventType="convert_processing",
|
||
message=f"Word文档生成完成",
|
||
documentName=doc_name,
|
||
documentIndex=0
|
||
)
|
||
|
||
# pandoc_path = "pandoc"
|
||
# sys_platform = platform.system()
|
||
#
|
||
# if sys_platform == "Linux":
|
||
# pandoc_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "tool/pandoc")
|
||
# elif sys_platform == "Windows":
|
||
# pandoc_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "tool/pandoc.exe")
|
||
# # 执行命令行
|
||
# try:
|
||
# command = (f'{pandoc_path} "{md_file_path}" -o '
|
||
# f'"{current_temp_path}/{doc_name}.docx" --standalone '
|
||
# f'--reference-doc "{template_file}"')
|
||
# logging.info(f"command: {command}")
|
||
# run_command(command, work=current_temp_path)
|
||
# except Exception as ex:
|
||
# logger.error(f"文档转换失败 ({doc_name})，返回码: {ex}")
|
||
# # 检查输出文件是否存在
|
||
# docx_file_path = f"{current_temp_path}/{doc_name}.docx"
|
||
if not os.path.exists(docx_file_path):
|
||
yield DocumentSSEData(
|
||
eventType="convert_error",
|
||
message=f"Word文档生成失败",
|
||
documentName=doc_name,
|
||
documentIndex=0
|
||
)
|
||
return
|
||
# 基于模板再次生成详细的文档
|
||
convert_docx_path = convert_template_docx(docx_file_path, project, doc_name, current_temp_path)
|
||
docx_util.postprocess_merged_document(convert_docx_path, project_name=project.name, doc_name=doc_name)
|
||
docx_util.enable_update_fields_on_open(convert_docx_path)
|
||
|
||
yield DocumentSSEData(
|
||
eventType="convert_processing",
|
||
message=f"模板文档生成完成，开始上传",
|
||
documentName=doc_name,
|
||
documentIndex=0
|
||
)
|
||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
|
||
object_name = f'{doc_name}_{timestamp}.docx'
|
||
if request.version != '' and request.version is not None:
|
||
object_name = f'{doc_name}_{request.version}_{timestamp}.docx'
|
||
minio_url = minio_client.upload_file_from_path(
|
||
file_path=convert_docx_path,
|
||
object_name=object_name,
|
||
content_type="application/octet-stream"
|
||
)
|
||
|
||
# 确定文档分类
|
||
# category = determine_document_category(doc_name)
|
||
|
||
# 保存到数据库
|
||
# show_name 包含时间戳（不包含UUID），格式如：运行方案说明_20260625_112409_671606.docx
|
||
timestamp_without_microseconds = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]
|
||
file_name = f'{doc_name}_{timestamp_without_microseconds}.docx'
|
||
if request.version != '' and request.version is not None:
|
||
file_name = f'{doc_name}_{request.version}_{timestamp_without_microseconds}.docx'
|
||
generated_doc = FileLinkSchema(
|
||
object_name=object_name,
|
||
url=minio_url,
|
||
file_name=file_name,
|
||
type=category
|
||
)
|
||
save_document.append(generated_doc)
|
||
logger.info(f"文档转换并上传成功: {doc_name}")
|
||
|
||
yield DocumentSSEData(
|
||
eventType="convert_complete",
|
||
message=f"文档转换并上传成功: {doc_name}",
|
||
documentName=doc_name,
|
||
documentIndex=0
|
||
)
|
||
except Exception as ex:
|
||
logger.error(f"文档转换失败 ({doc_name}): {str(ex)}", exc_info=True)
|
||
yield DocumentSSEData(
|
||
eventType="convert_error",
|
||
message=f"文档转换失败: {str(ex)}",
|
||
documentName=doc_name,
|
||
documentIndex=0
|
||
)
|
||
return
|
||
|
||
|
||
def determine_document_category(doc_name: str) -> str:
|
||
"""
|
||
根据文档名称确定分类
|
||
|
||
:param doc_name: 文档名称
|
||
:return: 文档分类
|
||
"""
|
||
for category, doc_types in DOCUMENT_TYPES.items():
|
||
if doc_name in doc_types:
|
||
return DocumentCategoryEnum(category).name
|
||
|
||
# 默认返回需求文档类
|
||
return DocumentCategoryEnum.REQUIREMENT.name
|