document_layout_node.py
2.74 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
"""
根据模板目录与多源报告,生成整本报告的标题/目录/主题设计。
"""
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 .base_node import BaseNode
class DocumentLayoutNode(BaseNode):
"""
负责生成全局标题、目录与Hero设计。
结合模板切片、报告摘要与论坛讨论,指导整本书的视觉与结构基调。
"""
def __init__(self, llm_client):
"""记录LLM客户端并设置节点名字,供BaseNode日志使用"""
super().__init__(llm_client, "DocumentLayoutNode")
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]:
"""综合模板+多源内容,生成全书的标题、目录结构与主题色板"""
# 将模板原文、切片结构与多源报告一并喂给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文本,若失败则抛出友好错误"""
cleaned = raw.strip()
if cleaned.startswith("```json"):
cleaned = cleaned[7:]
if cleaned.startswith("```"):
cleaned = cleaned[3:]
if cleaned.endswith("```"):
cleaned = cleaned[:-3]
cleaned = cleaned.strip()
if not cleaned:
raise ValueError("文档设计LLM返回空内容")
try:
return json.loads(cleaned)
except json.JSONDecodeError as exc:
raise ValueError(f"文档设计JSON解析失败: {exc}") from exc
__all__ = ["DocumentLayoutNode"]