|
|
@@ -6,8 +6,6 @@
|
|
|
|
|
|
import os
|
|
|
from typing import Dict, Any, List, Optional
|
|
|
-from langchain_core.prompts import ChatPromptTemplate
|
|
|
-
|
|
|
from src.datasets.parser.core.base import BaseNode, BaseState
|
|
|
from src.datasets.parser.core.registry import register_node
|
|
|
from src.conf.settings import model_settings
|
|
|
@@ -25,7 +23,7 @@ class SummaryNode(BaseNode):
|
|
|
对图像解析结果进行总结提要,并将结果保存到文件。
|
|
|
"""
|
|
|
|
|
|
- def __init__(self, dimension_id: int, model_name: Optional[str] = None):
|
|
|
+ def __init__(self, dimension_id: int = 0, model_name: Optional[str] = None):
|
|
|
"""
|
|
|
初始化总结节点
|
|
|
|
|
|
@@ -36,8 +34,7 @@ class SummaryNode(BaseNode):
|
|
|
self.dimension_id = dimension_id
|
|
|
self.model_name = model_name or model_settings.chat_model_name
|
|
|
# 总结提示模板
|
|
|
- self.summary_prompt = ChatPromptTemplate.from_template(
|
|
|
- """
|
|
|
+ self.summary_prompt = """
|
|
|
你是一位专业的内容总结专家,擅长从长篇文本中提取核心内容并生成简洁的总结。
|
|
|
|
|
|
请对以下内容进行总结,要求:
|
|
|
@@ -51,7 +48,6 @@ class SummaryNode(BaseNode):
|
|
|
|
|
|
总结:
|
|
|
"""
|
|
|
- )
|
|
|
|
|
|
@property
|
|
|
def name(self) -> str:
|
|
|
@@ -69,11 +65,11 @@ class SummaryNode(BaseNode):
|
|
|
"""
|
|
|
try:
|
|
|
# 构建提示
|
|
|
- messages = self.summary_prompt.format_messages(content=content)
|
|
|
+ messages = self.summary_prompt.format(content=content)
|
|
|
# 调用模型生成总结
|
|
|
chat_model = QWenVLParser(self.model_name)
|
|
|
- response = chat_model.invoke(messages)
|
|
|
- return response.content
|
|
|
+ response = chat_model.chat(prompt=messages)
|
|
|
+ return response
|
|
|
except Exception as e:
|
|
|
logger.error(f"总结内容时出错: {str(e)}")
|
|
|
return f"总结失败: {str(e)}"
|
|
|
@@ -90,7 +86,7 @@ class SummaryNode(BaseNode):
|
|
|
"""
|
|
|
# 获取解析结果
|
|
|
parsed_results = getattr(state, 'parsed_results', [])
|
|
|
-
|
|
|
+ original_filename = getattr(state, 'original_filename', '')
|
|
|
if not parsed_results:
|
|
|
logger.warning("没有解析结果可总结")
|
|
|
return {
|
|
|
@@ -128,7 +124,7 @@ class SummaryNode(BaseNode):
|
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
|
|
|
|
# 保存总结到文件
|
|
|
- file_path = os.path.join(temp_dir, f"{state.original_filename}_{self.dimension_id}.md")
|
|
|
+ file_path = os.path.join(temp_dir, f"{original_filename}_{self.dimension_id}.md")
|
|
|
try:
|
|
|
with open(file_path, 'w', encoding='utf-8') as f:
|
|
|
f.write(summary)
|