68 lines
2.3 KiB
Python
68 lines
2.3 KiB
Python
import inspect
|
||
import time
|
||
from abc import ABC, abstractmethod
|
||
from contextlib import contextmanager
|
||
from typing import TypeVar
|
||
|
||
from langgraph.errors import GraphInterrupt
|
||
from loguru import logger
|
||
|
||
from app.config import Settings, settings
|
||
from app.processor.progress import node_progress
|
||
|
||
T = TypeVar("T")
|
||
|
||
|
||
class BaseNode(ABC):
|
||
name: str = "base_node"
|
||
|
||
def __init__(self, config: Settings | None = None):
|
||
self.config = config if config is not None else settings
|
||
self.logger = logger.bind(node=self.name)
|
||
|
||
@contextmanager
|
||
def _instrument(self, state):
|
||
"""统一记录日志、耗时和进度事件，异常保持原类型向上抛出"""
|
||
start_time = time.perf_counter()
|
||
self.logger.info("节点开始")
|
||
try:
|
||
with node_progress(self.name, state):
|
||
yield
|
||
except GraphInterrupt:
|
||
raise
|
||
except Exception:
|
||
self.logger.error("节点执行失败")
|
||
raise
|
||
else:
|
||
self.logger.info("节点完成")
|
||
finally:
|
||
self.logger.info(f"节点耗时 {time.perf_counter() - start_time:.3f} 秒")
|
||
|
||
def __call__(self, state: T) -> T:
|
||
# 异步节点直接调用时返回协程，日志与进度在 _acall 内等待完成后记录
|
||
if inspect.iscoroutinefunction(self.process):
|
||
return self._acall(state)
|
||
with self._instrument(state):
|
||
return self.process(state)
|
||
|
||
async def _acall(self, state: T) -> T:
|
||
"""异步节点入口，由图根据 process 是否为协程选择"""
|
||
with self._instrument(state):
|
||
return await self.process(state)
|
||
|
||
@abstractmethod
|
||
def process(self, state: T) -> T:
|
||
"""子类实现节点核心逻辑，返回状态更新"""
|
||
|
||
def fail_on_invalid_output(self, state: T, error: ValueError, reason: str) -> T:
|
||
"""模型输出未通过校验时收尾为失败结果，原因交给失败收尾由模型向用户说明"""
|
||
self.logger.warning(f"{reason}: {error}")
|
||
return {**state, "status": "failed", "failure_reason": f"{reason}：{error}"}
|
||
|
||
def log_step(self, step_name: str, message: str = ""):
|
||
"""记录步骤日志。"""
|
||
log_msg = f"[{step_name}]"
|
||
if message:
|
||
log_msg += f" {message}"
|
||
self.logger.info(log_msg)
|