143 lines
4.7 KiB
Python
143 lines
4.7 KiB
Python
"""
|
||
文件概述生成工具
|
||
把 wiki_engine 的文档概述能力独立出来
|
||
输入：markdown 格式的文件内容
|
||
输出：文件概述（SUMMARY 行 + markdown 正文）
|
||
支持并发
|
||
"""
|
||
import asyncio
|
||
from typing import List, Optional
|
||
from config import LLM_CONFIG
|
||
from openai import AsyncOpenAI
|
||
|
||
|
||
# ==================== 配置 ====================
|
||
_MAX_CONCURRENCY = 8 # 最大并发数
|
||
|
||
_client: Optional[AsyncOpenAI] = None
|
||
_semaphore: Optional[asyncio.Semaphore] = None
|
||
|
||
|
||
def _get_client() -> AsyncOpenAI:
|
||
"""获取或创建 AsyncOpenAI 单例客户端"""
|
||
global _client
|
||
if _client is None:
|
||
_client = AsyncOpenAI(api_key=LLM_CONFIG['api_key'], base_url=LLM_CONFIG['base_url'])
|
||
return _client
|
||
|
||
|
||
def _get_semaphore() -> asyncio.Semaphore:
|
||
"""获取或创建并发信号量"""
|
||
global _semaphore
|
||
if _semaphore is None:
|
||
_semaphore = asyncio.Semaphore(_MAX_CONCURRENCY)
|
||
return _semaphore
|
||
|
||
|
||
# ==================== Prompt（与 wiki_builder.py 的 WIKI_SUMMARY_PROMPT 一致，仅去掉 extracted_slugs 相关部分） ====================
|
||
SUMMARY_PROMPT = """You are a wiki editor. Given the following document content, create a structured wiki summary page in Markdown format.
|
||
|
||
<document>
|
||
<content>
|
||
{content}
|
||
</content>
|
||
</document>
|
||
|
||
<instructions>
|
||
1. The FIRST line of your output MUST be: SUMMARY: {{one sentence, 15-40 words, describing what this document is about for wiki index listing}}
|
||
2. Create a concise but useful document-level summary.
|
||
3. Include the document's main subject, scope, important procedures, standards, systems, equipment, tables, fields.
|
||
4. Do NOT invent facts. Stay grounded in the document content.
|
||
5. Write in Chinese.
|
||
6. If the content is empty or has no substantive information, output exactly: "SUMMARY: No textual content was extractable from this document." followed by a brief note.
|
||
</instructions>
|
||
|
||
Output the SUMMARY line first, then the Markdown content. Do not include any other preamble."""
|
||
|
||
|
||
async def generate_summary(content: str) -> str:
|
||
"""
|
||
从 markdown 内容生成文件概述
|
||
|
||
Args:
|
||
content: markdown 格式的文件内容
|
||
|
||
Returns:
|
||
文件概述（SUMMARY 行 + markdown 正文）
|
||
"""
|
||
if not content or not content.strip():
|
||
return ""
|
||
|
||
prompt = SUMMARY_PROMPT.format(content=content)
|
||
client = _get_client()
|
||
|
||
async with _get_semaphore():
|
||
try:
|
||
response = await client.chat.completions.create(
|
||
model=LLM_CONFIG['model'],
|
||
messages=[
|
||
{"role": "system", "content": "You are a grounded wiki editor. Do not invent facts."},
|
||
{"role": "user", "content": prompt},
|
||
],
|
||
temperature=0.1,
|
||
max_tokens=LLM_CONFIG['max_tokens'],
|
||
stream=False,
|
||
extra_body={"chat_template_kwargs": {"enable_thinking": False}},
|
||
)
|
||
return (response.choices[0].message.content or "").strip()
|
||
except Exception as exc:
|
||
print(f"[generate_summary] 概述生成失败: {exc}")
|
||
return ""
|
||
|
||
|
||
async def generate_summaries(contents: List[str]) -> List[str]:
|
||
"""
|
||
并发生成多个文件的概述
|
||
|
||
Args:
|
||
contents: markdown 格式的文件内容列表
|
||
|
||
Returns:
|
||
文件概述列表，顺序与输入一致
|
||
"""
|
||
return await asyncio.gather(*(generate_summary(c) for c in contents))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# import sys
|
||
from pathlib import Path
|
||
|
||
# if len(sys.argv) > 1:
|
||
# with open(sys.argv[1], "r", encoding="utf-8") as f:
|
||
# test_content = f.read()
|
||
# else:
|
||
# test_content = """
|
||
# # 船舶动力系统维护规程
|
||
|
||
# ## 概述
|
||
# 本文档详细介绍了船舶动力系统的日常维护和故障处理流程。
|
||
|
||
# ## 发动机日常检查
|
||
# - 润滑油位检查：每日检查发动机润滑油位，保持在标尺正常范围
|
||
# - 冷却液位检查：确保冷却系统液位正常，无泄漏
|
||
# - 皮带张紧度：检查传动皮带张紧度，过松或过紧均需调整
|
||
|
||
# ## 冷却系统维护
|
||
# 定期清洗热交换器，检查水泵密封性，更换老化管路。
|
||
|
||
# ## 故障诊断流程
|
||
# 1. 现象观察：记录故障现象和发生条件
|
||
# 2. 数据采集：收集运行参数和报警信息
|
||
# 3. 原因分析：对照标准参数分析故障原因
|
||
# 4. 处理方案：制定维修方案并执行
|
||
|
||
# ## 安全注意事项
|
||
# 所有维护操作必须在停机状态下进行，操作人员需佩戴防护装备。
|
||
# """
|
||
md_path = Path(r"E:\ZKYNLP\Hjunproject\project0506\kgrag\船舶主机燃油泵自动控制系统故障树分析.md")
|
||
with open(md_path, "r", encoding="utf-8") as f:
|
||
test_content = f.read()
|
||
|
||
result = asyncio.run(generate_summary(test_content))
|
||
print(result)
|