document_layout_node.py
4.61 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
"""
根据模板目录与多源报告,生成整本报告的标题/目录/主题设计。
"""
from __future__ import annotations
import json
from typing import Any, Dict, List
from loguru import logger
from ..core import TemplateSection
from ..prompts import (
SYSTEM_PROMPT_DOCUMENT_LAYOUT,
build_document_layout_prompt,
)
from ..utils.json_parser import RobustJSONParser, JSONParseError
from .base_node import BaseNode
class DocumentLayoutNode(BaseNode):
"""
负责生成全局标题、目录与Hero设计。
结合模板切片、报告摘要与论坛讨论,指导整本书的视觉与结构基调。
"""
def __init__(self, llm_client):
"""记录LLM客户端并设置节点名字,供BaseNode日志使用"""
super().__init__(llm_client, "DocumentLayoutNode")
# 初始化鲁棒JSON解析器,启用所有修复策略
self.json_parser = RobustJSONParser(
enable_json_repair=True,
enable_llm_repair=False, # 可以根据需要启用LLM修复
max_repair_attempts=3,
)
def run(
self,
sections: List[TemplateSection],
template_markdown: str,
reports: Dict[str, str],
forum_logs: str,
query: str,
template_overview: Dict[str, Any] | None = None,
) -> Dict[str, Any]:
"""
综合模板+多源内容,生成全书的标题、目录结构与主题色板。
参数:
sections: 模板切片后的章节列表。
template_markdown: 模板原文,用于LLM理解上下文。
reports: 三个引擎的内容映射。
forum_logs: 论坛讨论摘要。
query: 用户查询词。
template_overview: 预生成的模板概览,可复用以减少提示词长度。
返回:
dict: 包含 title/subtitle/toc/hero/themeTokens 等设计信息的字典。
"""
# 将模板原文、切片结构与多源报告一并喂给LLM,便于其理解层级与素材
payload = {
"query": query,
"template": {
"raw": template_markdown,
"sections": [section.to_dict() for section in sections],
},
"templateOverview": template_overview
or {
"title": sections[0].title if sections else "",
"chapters": [section.to_dict() for section in sections],
},
"reports": reports,
"forumLogs": forum_logs,
}
user_message = build_document_layout_prompt(payload)
response = self.llm_client.stream_invoke_to_string(
SYSTEM_PROMPT_DOCUMENT_LAYOUT,
user_message,
temperature=0.3,
top_p=0.9,
)
design = self._parse_response(response)
logger.info("文档标题/目录设计已生成")
return design
def _parse_response(self, raw: str) -> Dict[str, Any]:
"""
解析LLM返回的JSON文本,若失败则抛出友好错误。
使用鲁棒JSON解析器进行多重修复尝试:
1. 清理markdown标记和思考内容
2. 本地语法修复(括号平衡、逗号补全、控制字符转义等)
3. 使用json_repair库进行高级修复
4. 可选的LLM辅助修复
参数:
raw: LLM原始返回字符串,允许带```包裹、思考内容等。
返回:
dict: 结构化的设计稿。
异常:
ValueError: 当响应为空或JSON解析失败时抛出。
"""
try:
result = self.json_parser.parse(
raw,
context_name="文档设计",
expected_keys=["title", "toc", "hero"],
)
# 验证关键字段的类型
if not isinstance(result.get("title"), str):
logger.warning("文档设计缺少title字段或类型错误,使用默认值")
result.setdefault("title", "未命名报告")
if not isinstance(result.get("toc"), (list, dict)):
logger.warning("文档设计缺少toc字段或类型错误,使用空列表")
result.setdefault("toc", [])
if not isinstance(result.get("hero"), dict):
logger.warning("文档设计缺少hero字段或类型错误,使用空对象")
result.setdefault("hero", {})
return result
except JSONParseError as exc:
# 转换为原有的异常类型以保持向后兼容
raise ValueError(f"文档设计JSON解析失败: {exc}") from exc
__all__ = ["DocumentLayoutNode"]