LangGraph技术底座

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

# 简化的执行流程
class GraphExecutor:
    def __init__(self, graph: nx.DiGraph):
        self.graph = graph
        self.state_schema = ...
    
    async def execute(self, initial_state):
        current_node = self.get_entry_point()
        state = initial_state
        
        while current_node != END:
            # 执行节点
            node_fn = self.graph.nodes[current_node]['function']
            state = await node_fn(state)
            
            # 决定下一个节点
            edges = self.graph.out_edges(current_node)
            current_node = self.route(state, edges)
        
        return state

1
2
3
4
5
6

# 条件边的实现
def add_conditional_edges(source, path_map, path_fn):
    for target, condition in path_map.items():
        def edge_condition(state):
            return path_fn(state) == condition
        graph.add_edge(source, target, condition=edge_condition)

 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

from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
import pydantic, asyncio

class FlowState(pydantic.BaseModel):
    history: list[str] = []
    user_input: str = ""

llm = ChatOpenAI(model="gpt-3.5-turbo", streaming=True)

# 1. 定义节点
async def answer_q(state: FlowState) -> FlowState:
    rsp = await llm.ainvoke(state.user_input)
    state.history.append(rsp.content)
    return state

def check_continue(state: FlowState) -> str:
    return END if state.user_input.lower() == "quit" else "ask"

# 2. 画图
g = StateGraph(FlowState)
g.add_node("ask", answer_q)
g.set_entry_point("ask")
g.add_conditional_edges("ask", check_continue)

graph_runnable = g.compile()

# 3. 调用
asyncio.run(
    graph_runnable.ainvoke({"user_input": "LangGraph 的技术底座？"})
)

 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

class StateGraph:
    def __init__(self, state_schema: Type[BaseModel]):
        self.nodes: Dict[str, NodeSpec] = {}
        self.edges: Dict[str, EdgeSpec] = {}
        self.channels: Dict[str, Channel] = {}
        self.state_schema = state_schema
        self.compiled_graph: Optional[CompiledGraph] = None
        
    def add_node(self, name: str, func: Callable, **kwargs):
        """节点注册机制"""
        self.nodes[name] = NodeSpec(
            name=name,
            func=func,
            input_channels=self._infer_input_channels(func),
            output_channels=self._infer_output_channels(func),
            metadata=kwargs
        )
        
    def add_edge(self, source: str, target: str):
        """边的添加与验证"""
        if source not in self.nodes or target not in self.nodes:
            raise GraphValidationError(f"Unknown node: {source} or {target}")
        
        self.edges[f"{source}->{target}"] = EdgeSpec(
            source=source,
            target=target,
            condition=None,
            edge_type=EdgeType.UNCONDITIONAL
        )
        
    def add_conditional_edges(self, source: str, condition: Callable, 
                            mapping: Dict[str, str]):
        """条件边实现"""
        for condition_result, target in mapping.items():
            self.edges[f"{source}->{target}@{condition_result}"] = EdgeSpec(
                source=source,
                target=target,
                condition=condition,
                condition_result=condition_result,
                edge_type=EdgeType.CONDITIONAL
            )

 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

class GraphCompiler:
    def compile(self, graph: StateGraph) -> CompiledGraph:
        """图编译的核心逻辑"""
        # 1. 拓扑排序验证
        self._validate_topology(graph)
        
        # 2. 状态通道构建
        channels = self._build_channels(graph)
        
        # 3. 执行计划生成
        execution_plan = self._generate_execution_plan(graph)
        
        # 4. 优化策略应用
        optimized_plan = self._optimize_execution_plan(execution_plan)
        
        return CompiledGraph(
            nodes=graph.nodes,
            edges=graph.edges,
            channels=channels,
            execution_plan=optimized_plan,
            state_schema=graph.state_schema
        )
        
    def _validate_topology(self, graph: StateGraph):
        """拓扑验证算法"""
        visited = set()
        rec_stack = set()
        
        def dfs(node):
            if node in rec_stack:
                raise CyclicGraphError(f"Cycle detected involving node: {node}")
            if node in visited:
                return
                
            visited.add(node)
            rec_stack.add(node)
            
            for edge in graph.edges.values():
                if edge.source == node:
                    dfs(edge.target)
                    
            rec_stack.remove(node)
        
        for node in graph.nodes:
            if node not in visited:
                dfs(node)

 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

class Channel(ABC):
    """状态通道抽象基类"""
    def __init__(self, value_type: Type, reducer: Optional[Callable] = None):
        self.value_type = value_type
        self.reducer = reducer or self._default_reducer
        
    @abstractmethod
    def update(self, old_value: Any, new_value: Any) -> Any:
        """状态更新逻辑"""
        pass
        
    def _default_reducer(self, old_value: Any, new_value: Any) -> Any:
        """默认的状态归约策略"""
        return new_value

class LastValueChannel(Channel):
    """最后值策略通道"""
    def update(self, old_value: Any, new_value: Any) -> Any:
        return new_value

class BinaryOperatorChannel(Channel):
    """二元操作通道"""
    def __init__(self, value_type: Type, operator: Callable):
        super().__init__(value_type)
        self.operator = operator
        
    def update(self, old_value: Any, new_value: Any) -> Any:
        if old_value is None:
            return new_value
        return self.operator(old_value, new_value)

class AppendChannel(BinaryOperatorChannel):
    """追加操作通道"""
    def __init__(self, value_type: Type):
        super().__init__(value_type, lambda a, b: a + [b] if isinstance(a, list) else [a, b])

 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

class StateManager:
    def __init__(self, channels: Dict[str, Channel], state_schema: Type[BaseModel]):
        self.channels = channels
        self.state_schema = state_schema
        self.current_state: Dict[str, Any] = {}
        self.state_history: List[Dict[str, Any]] = []
        
    def update_state(self, updates: Dict[str, Any]) -> Dict[str, Any]:
        """状态更新的核心逻辑"""
        new_state = self.current_state.copy()
        
        for key, new_value in updates.items():
            if key not in self.channels:
                raise StateError(f"Unknown state key: {key}")
                
            channel = self.channels[key]
            old_value = new_state.get(key)
            new_state[key] = channel.update(old_value, new_value)
            
        # 状态验证
        self._validate_state(new_state)
        
        # 历史记录
        self.state_history.append(self.current_state.copy())
        self.current_state = new_state
        
        return new_state
        
    def _validate_state(self, state: Dict[str, Any]):
        """状态一致性验证"""
        try:
            self.state_schema(**state)
        except ValidationError as e:
            raise StateValidationError(f"State validation failed: {e}")

 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

class NodeExecutor:
    def __init__(self, node_spec: NodeSpec, state_manager: StateManager):
        self.node_spec = node_spec
        self.state_manager = state_manager
        self.execution_context = ExecutionContext()
        
    async def execute(self, input_state: Dict[str, Any], 
                     config: ExecutionConfig) -> Dict[str, Any]:
        """节点执行的核心逻辑"""
        try:
            # 1. 前置处理
            processed_input = self._preprocess_input(input_state)
            
            # 2. 函数调用
            if asyncio.iscoroutinefunction(self.node_spec.func):
                result = await self.node_spec.func(processed_input)
            else:
                result = await asyncio.get_event_loop().run_in_executor(
                    None, self.node_spec.func, processed_input
                )
            
            # 3. 结果处理
            output_state = self._postprocess_output(result)
            
            # 4. 状态更新
            updated_state = self.state_manager.update_state(output_state)
            
            return updated_state
            
        except Exception as e:
            await self._handle_execution_error(e, input_state, config)
            raise
            
    def _preprocess_input(self, state: Dict[str, Any]) -> Any:
        """输入预处理"""
        if len(self.node_spec.input_channels) == 1:
            key = list(self.node_spec.input_channels.keys())[0]
            return state.get(key)
        else:
            return {key: state.get(key) for key in self.node_spec.input_channels}

 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

class GraphExecutor:
    def __init__(self, compiled_graph: CompiledGraph):
        self.graph = compiled_graph
        self.checkpoint_manager = CheckpointManager()
        self.execution_tracer = ExecutionTracer()
        
    async def invoke(self, input_data: Dict[str, Any], 
                    config: Optional[ExecutionConfig] = None) -> Dict[str, Any]:
        """图执行的主入口"""
        config = config or ExecutionConfig()
        
        # 初始化状态
        state_manager = StateManager(
            self.graph.channels, 
            self.graph.state_schema
        )
        state_manager.update_state(input_data)
        
        # 创建执行上下文
        context = ExecutionContext(
            graph=self.graph,
            state_manager=state_manager,
            config=config,
            trace_id=uuid.uuid4()
        )
        
        try:
            # 执行图
            final_state = await self._execute_graph(context)
            return final_state
            
        except Exception as e:
            await self._handle_graph_error(e, context)
            raise
            
    async def _execute_graph(self, context: ExecutionContext) -> Dict[str, Any]:
        """图执行的核心算法"""
        current_nodes = {self.graph.entry_point}
        visited_nodes = set()
        
        while current_nodes:
            # 并发执行当前层的所有节点
            tasks = []
            for node_name in current_nodes:
                if node_name not in visited_nodes:
                    node_executor = NodeExecutor(
                        self.graph.nodes[node_name],
                        context.state_manager
                    )
                    tasks.append(
                        self._execute_node_with_trace(
                            node_executor, context, node_name
                        )
                    )
                    visited_nodes.add(node_name)
            
            if tasks:
                await asyncio.gather(*tasks)
            
            # 计算下一批要执行的节点
            next_nodes = set()
            for node_name in current_nodes:
                next_nodes.update(
                    self._get_next_nodes(node_name, context.state_manager.current_state)
                )
            
            current_nodes = next_nodes - visited_nodes
            
        return context.state_manager.current_state
        
    def _get_next_nodes(self, current_node: str, state: Dict[str, Any]) -> Set[str]:
        """计算下一步要执行的节点"""
        next_nodes = set()
        
        for edge in self.graph.edges.values():
            if edge.source == current_node:
                if edge.edge_type == EdgeType.UNCONDITIONAL:
                    next_nodes.add(edge.target)
                elif edge.edge_type == EdgeType.CONDITIONAL:
                    condition_result = edge.condition(state)
                    if condition_result == edge.condition_result:
                        next_nodes.add(edge.target)
                        
        return next_nodes

 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

class CheckpointManager:
    def __init__(self, storage_backend: StorageBackend):
        self.storage = storage_backend
        self.checkpoint_strategy = CheckpointStrategy.EVERY_NODE
        
    async def save_checkpoint(self, execution_id: str, node_name: str, 
                            state: Dict[str, Any], metadata: Dict[str, Any]):
        """保存检查点"""
        checkpoint = Checkpoint(
            execution_id=execution_id,
            node_name=node_name,
            state=state,
            metadata=metadata,
            timestamp=datetime.utcnow(),
            checkpoint_id=uuid.uuid4()
        )
        
        await self.storage.save(checkpoint)
        
    async def restore_checkpoint(self, execution_id: str, 
                               checkpoint_id: Optional[str] = None) -> Checkpoint:
        """恢复检查点"""
        if checkpoint_id:
            return await self.storage.load_by_id(checkpoint_id)
        else:
            return await self.storage.load_latest(execution_id)

class StorageBackend(ABC):
    @abstractmethod
    async def save(self, checkpoint: Checkpoint):
        pass
        
    @abstractmethod
    async def load_by_id(self, checkpoint_id: str) -> Checkpoint:
        pass
        
    @abstractmethod
    async def load_latest(self, execution_id: str) -> Checkpoint:
        pass

class RedisStorageBackend(StorageBackend):
    def __init__(self, redis_client):
        self.redis = redis_client
        
    async def save(self, checkpoint: Checkpoint):
        key = f"checkpoint:{checkpoint.execution_id}:{checkpoint.checkpoint_id}"
        data = checkpoint.to_dict()
        await self.redis.setex(key, 3600, json.dumps(data))
        
        # 维护最新检查点索引
        latest_key = f"latest_checkpoint:{checkpoint.execution_id}"
        await self.redis.setex(latest_key, 3600, checkpoint.checkpoint_id)

 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

class StreamingExecutor:
    def __init__(self, graph_executor: GraphExecutor):
        self.graph_executor = graph_executor
        
    async def stream(self, input_data: Dict[str, Any], 
                    config: Optional[ExecutionConfig] = None):
        """流式执行图"""
        async for event in self._stream_execution(input_data, config):
            yield event
            
    async def _stream_execution(self, input_data: Dict[str, Any], 
                              config: ExecutionConfig):
        """流式执行的核心逻辑"""
        # 创建执行上下文
        context = self._create_streaming_context(input_data, config)
        
        current_nodes = {self.graph_executor.graph.entry_point}
        visited_nodes = set()
        
        while current_nodes:
            for node_name in current_nodes:
                if node_name not in visited_nodes:
                    # 发送节点开始事件
                    yield StreamEvent(
                        event_type=EventType.NODE_START,
                        node_name=node_name,
                        state=context.state_manager.current_state.copy(),
                        timestamp=datetime.utcnow()
                    )
                    
                    # 执行节点
                    node_executor = NodeExecutor(
                        self.graph_executor.graph.nodes[node_name],
                        context.state_manager
                    )
                    
                    updated_state = await node_executor.execute(
                        context.state_manager.current_state, 
                        config
                    )
                    
                    # 发送节点完成事件
                    yield StreamEvent(
                        event_type=EventType.NODE_COMPLETE,
                        node_name=node_name,
                        state=updated_state,
                        timestamp=datetime.utcnow()
                    )
                    
                    visited_nodes.add(node_name)
            
            # 计算下一批节点
            next_nodes = set()
            for node_name in current_nodes:
                next_nodes.update(
                    self.graph_executor._get_next_nodes(
                        node_name, 
                        context.state_manager.current_state
                    )
                )
            
            current_nodes = next_nodes - visited_nodes

 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

class LoopController:
    def __init__(self, max_iterations: int = 100):
        self.max_iterations = max_iterations
        self.iteration_count = defaultdict(int)
        
    def should_continue(self, loop_id: str, condition_func: Callable, 
                      state: Dict[str, Any]) -> bool:
        """循环控制逻辑"""
        self.iteration_count[loop_id] += 1
        
        if self.iteration_count[loop_id] > self.max_iterations:
            raise LoopLimitExceededError(
                f"Loop {loop_id} exceeded maximum iterations: {self.max_iterations}"
            )
            
        return condition_func(state)

class ConditionalBranch:
    def __init__(self, condition: Callable, branches: Dict[Any, str]):
        self.condition = condition
        self.branches = branches
        
    def evaluate(self, state: Dict[str, Any]) -> str:
        """条件分支评估"""
        result = self.condition(state)
        
        if result in self.branches:
            return self.branches[result]
        elif 'default' in self.branches:
            return self.branches['default']
        else:
            raise BranchEvaluationError(
                f"No branch found for condition result: {result}"
            )

 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

class ErrorHandler:
    def __init__(self, retry_policy: RetryPolicy, 
                 fallback_strategy: FallbackStrategy):
        self.retry_policy = retry_policy
        self.fallback_strategy = fallback_strategy
        
    async def handle_error(self, error: Exception, context: ExecutionContext) -> bool:
        """错误处理的核心逻辑"""
        if self.retry_policy.should_retry(error, context):
            await self._execute_retry(context)
            return True
        else:
            await self._execute_fallback(error, context)
            return False
            
    async def _execute_retry(self, context: ExecutionContext):
        """重试逻辑"""
        delay = self.retry_policy.get_delay(context.retry_count)
        await asyncio.sleep(delay)
        context.retry_count += 1
        
    async def _execute_fallback(self, error: Exception, context: ExecutionContext):
        """回退策略执行"""
        if self.fallback_strategy.has_fallback():
            fallback_result = await self.fallback_strategy.execute(error, context)
            context.state_manager.update_state(fallback_result)
        else:
            raise error

class ExponentialBackoffRetryPolicy:
    def __init__(self, max_retries: int = 3, base_delay: float = 1.0):
        self.max_retries = max_retries
        self.base_delay = base_delay
        
    def should_retry(self, error: Exception, context: ExecutionContext) -> bool:
        return (context.retry_count < self.max_retries and 
                isinstance(error, RetryableError))
                
    def get_delay(self, retry_count: int) -> float:
        return self.base_delay * (2 ** retry_count)

 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

class ConcurrencyController:
    def __init__(self, max_concurrent_nodes: int = 10):
        self.semaphore = asyncio.Semaphore(max_concurrent_nodes)
        self.node_dependencies = {}
        
    async def execute_nodes_concurrently(self, nodes: List[str], 
                                       context: ExecutionContext):
        """并发执行节点"""
        # 计算依赖关系
        dependency_graph = self._build_dependency_graph(nodes, context.graph)
        
        # 按层级并发执行
        execution_levels = self._topological_sort(dependency_graph)
        
        for level in execution_levels:
            tasks = []
            for node_name in level:
                task = self._execute_node_with_semaphore(node_name, context)
                tasks.append(task)
            
            if tasks:
                await asyncio.gather(*tasks)
                
    async def _execute_node_with_semaphore(self, node_name: str, 
                                         context: ExecutionContext):
        """带信号量控制的节点执行"""
        async with self.semaphore:
            node_executor = NodeExecutor(
                context.graph.nodes[node_name],
                context.state_manager
            )
            return await node_executor.execute(
                context.state_manager.current_state,
                context.config
            )

 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

class MemoryManager:
    def __init__(self, max_memory_mb: int = 1024):
        self.max_memory_bytes = max_memory_mb * 1024 * 1024
        self.state_cache = {}
        self.cache_policy = LRUCachePolicy(max_size=1000)
        
    def optimize_state_storage(self, state: Dict[str, Any]) -> Dict[str, Any]:
        """状态存储优化"""
        # 大对象压缩
        optimized_state = {}
        for key, value in state.items():
            if self._is_large_object(value):
                optimized_state[key] = self._compress_object(value)
            else:
                optimized_state[key] = value
                
        return optimized_state
        
    def _is_large_object(self, obj: Any) -> bool:
        """判断是否为大对象"""
        return sys.getsizeof(obj) > 1024 * 1024  # 1MB
        
    def _compress_object(self, obj: Any) -> CompressedObject:
        """对象压缩"""
        serialized = pickle.dumps(obj)
        compressed = gzip.compress(serialized)
        return CompressedObject(data=compressed, original_type=type(obj))

 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

class DistributedTracer:
    def __init__(self, tracer_provider):
        self.tracer = tracer_provider.get_tracer("langgraph")
        
    async def trace_graph_execution(self, graph_name: str, 
                                  execution_func: Callable):
        """图执行追踪"""
        with self.tracer.start_as_current_span(
            f"graph_execution_{graph_name}"
        ) as span:
            span.set_attribute("graph.name", graph_name)
            span.set_attribute("graph.version", "1.0")
            
            try:
                result = await execution_func()
                span.set_attribute("execution.status", "success")
                return result
            except Exception as e:
                span.set_attribute("execution.status", "error")
                span.set_attribute("error.message", str(e))
                span.record_exception(e)
                raise
                
    async def trace_node_execution(self, node_name: str, 
                                 execution_func: Callable):
        """节点执行追踪"""
        with self.tracer.start_as_current_span(
            f"node_execution_{node_name}"
        ) as span:
            span.set_attribute("node.name", node_name)
            
            start_time = time.time()
            try:
                result = await execution_func()
                execution_time = time.time() - start_time
                span.set_attribute("node.execution_time", execution_time)
                return result
            except Exception as e:
                span.record_exception(e)
                raise

 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

class PerformanceMonitor:
    def __init__(self):
        self.metrics_collector = MetricsCollector()
        self.performance_history = deque(maxlen=1000)
        
    async def monitor_execution(self, context: ExecutionContext):
        """执行性能监控"""
        metrics = ExecutionMetrics(
            execution_id=context.execution_id,
            start_time=datetime.utcnow(),
            memory_usage=self._get_memory_usage(),
            cpu_usage=self._get_cpu_usage()
        )
        
        try:
            yield metrics
        finally:
            metrics.end_time = datetime.utcnow()
            metrics.duration = (metrics.end_time - metrics.start_time).total_seconds()
            
            self.performance_history.append(metrics)
            await self.metrics_collector.record_metrics(metrics)
            
    def _get_memory_usage(self) -> float:
        """获取内存使用量"""
        process = psutil.Process()
        return process.memory_info().rss / 1024 / 1024  # MB
        
    def _get_cpu_usage(self) -> float:
        """获取CPU使用率"""
        return psutil.cpu_percent(interval=0.1)