Add Comments

ReportEngine/nodes/template_selection_node.py

@@ -1,6 +1,8 @@
 """
-模板选择节点
-根据查询内容和可用模板选择最合适的报告模板
+模板选择节点。
+
+综合用户查询、三引擎报告、论坛日志与本地模板库，
+调用LLM挑选最合适的报告骨架。
 """
 
 import os
@@ -13,7 +15,12 @@ from ..prompts import SYSTEM_PROMPT_TEMPLATE_SELECTION
 
 
 class TemplateSelectionNode(BaseNode):
-    """模板选择处理节点"""
+    """
+    模板选择处理节点。
+
+    负责准备模板候选列表、构建提示词、解析LLM返回结果，
+    并在失败时回退到内置模板。
+    """
     
     def __init__(self, llm_client, template_dir: str = "ReportEngine/report_template"):
         """
@@ -28,7 +35,7 @@ class TemplateSelectionNode(BaseNode):
         
     def run(self, input_data: Dict[str, Any], **kwargs) -> Dict[str, Any]:
         """
-        执行模板选择
+        执行模板选择。
         
         Args:
             input_data: 包含查询和报告内容的字典
@@ -37,7 +44,7 @@ class TemplateSelectionNode(BaseNode):
                 - forum_logs: 论坛日志内容
                 
         Returns:
-            选择的模板信息
+            选择的模板信息，包含名称、内容与选择理由
         """
         logger.info("开始模板选择...")
         
@@ -67,7 +74,12 @@ class TemplateSelectionNode(BaseNode):
     
     def _llm_template_selection(self, query: str, reports: List[Any], forum_logs: str, 
                               available_templates: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
-        """使用LLM进行模板选择"""
+        """
+        使用LLM进行模板选择。
+
+        构造模板列表与报告摘要 → 调用LLM → 解析JSON →
+        验证模板是否存在并返回标准结构。
+        """
         logger.info("尝试使用LLM进行模板选择...")
         
         # 构建模板列表
@@ -150,7 +162,11 @@ class TemplateSelectionNode(BaseNode):
             return self._extract_template_from_text(response, available_templates)
     
     def _clean_llm_response(self, response: str) -> str:
-        """清理LLM响应"""
+        """
+        清理LLM响应。
+
+        去掉 ```json``` 包裹以及前后空白，方便 `json.loads`。
+        """
         # 移除可能的markdown代码块标记
         if '```json' in response:
             response = response.split('```json')[1].split('```')[0]
@@ -163,7 +179,11 @@ class TemplateSelectionNode(BaseNode):
         return response
     
     def _extract_template_from_text(self, response: str, available_templates: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
-        """从文本响应中提取模板信息"""
+        """
+        从文本响应中提取模板信息。
+
+        当LLM未输出合法JSON时，尝试匹配模板名称关键字做降级。
+        """
         logger.info("尝试从文本响应中提取模板信息")
         
         # 查找响应中是否包含模板名称
@@ -186,7 +206,11 @@ class TemplateSelectionNode(BaseNode):
         return None
     
     def _get_available_templates(self) -> List[Dict[str, Any]]:
-        """获取可用的模板列表"""
+        """
+        获取可用的模板列表。
+
+        枚举模板目录下的 `.md` 文件并读取内容与描述字段。
+        """
         templates = []
         
         if not os.path.exists(self.template_dir):
@@ -216,7 +240,7 @@ class TemplateSelectionNode(BaseNode):
         return templates
     
     def _extract_template_description(self, template_name: str) -> str:
-        """根据模板名称生成描述"""
+        """根据模板名称生成描述，方便LLM理解模板定位。"""
         if '企业品牌' in template_name:
             return "适用于企业品牌声誉和形象分析"
         elif '市场竞争' in template_name:
@@ -235,7 +259,7 @@ class TemplateSelectionNode(BaseNode):
 
     
     def _get_fallback_template(self) -> Dict[str, Any]:
-        """获取备用默认模板（空模板，让LLM自行发挥）"""
+        """获取备用默认模板（空模板，让LLM自行发挥）。"""
         logger.info("未找到合适模板，使用空模板让LLM自行发挥")
         
         return {