LangGraph从新手到老师傅 - 6 - Context上下文的使用

前言

在构建复杂的工作流系统时，我们经常需要在不改变核心状态的情况下传递一些额外信息，如用户身份、环境配置、请求元数据等。LangGraph提供了Context（上下文）功能，完美地解决了这个问题。本文将通过分析示例，深入讲解Context的概念、实现方式和应用场景，帮助你在实际项目中灵活运用这一强大功能。

Context基础概念

Context是LangGraph中的一个特殊机制，它允许我们在执行图时传递额外的信息，而不需要将这些信息作为状态的一部分。与状态不同，上下文通常包含以下特点：

不参与状态更新：上下文数据不会被节点函数修改或更新
贯穿整个执行过程：上下文在整个图执行过程中保持不变
辅助决策：上下文通常用于辅助节点函数做出决策或个性化处理
类型安全：可以通过类型标注确保上下文数据的类型安全

示例代码

python 复制代码

from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END

print("======= 示例4: 使用Context =======")

# 定义状态类型
class ContextState(TypedDict):
    message: str

# 定义上下文类型
class UserContext(TypedDict):
    user_name: str
    user_role: str

# 定义节点函数
def process_with_context(state: ContextState, runtime) -> dict:
    """使用上下文处理消息的节点"""
    # 访问上下文
    context = runtime.context
    user_name = context.get("user_name", "Guest")
    user_role = context.get("user_role", "User")
    
    processed_message = f"[{user_role}: {user_name}] {state['message']}"
    return {"message": processed_message}

# 创建StateGraph，指定上下文类型
context_graph = StateGraph(ContextState, context_schema=UserContext)

# 添加节点
context_graph.add_node("process", process_with_context)

# 添加边
context_graph.add_edge(START, "process")
context_graph.add_edge("process", END)

# 编译图
compiled_context_graph = context_graph.compile()

# 执行图，提供上下文
result = compiled_context_graph.invoke(
    {"message": "Hello, LangGraph!"}, 
    context={"user_name": "Alice", "user_role": "Admin"}
)
print(f"输入: {{'message': 'Hello, LangGraph!'}}")
print(f"上下文: {{'user_name': 'Alice', 'user_role': 'Admin'}}")
print(f"输出: {result}")

# 示例说明：
# 1. 这个示例展示了如何在StateGraph中使用上下文(context)传递额外信息
# 2. 通过定义context_schema参数，指定了上下文的类型结构
# 3. 在节点函数中，通过runtime.context访问上下文字段
# 4. 执行图时，通过context参数传入具体的上下文数据
# 5. 这种机制非常适合传递用户身份、环境配置等非状态但影响处理逻辑的信息

输出结果

css 复制代码

======= 示例4: 使用Context =======
输入: {'message': 'Hello, LangGraph!'}
上下文: {'user_name': 'Alice', 'user_role': 'Admin'}
输出: {'message': '[Admin: Alice] Hello, LangGraph!'}

代码解析：Context的实现与应用

1. 定义状态类型

python 复制代码

class ContextState(TypedDict):
    message: str

这个状态类型非常简单，只包含一个message字段，用于存储要处理的消息内容。与前面的示例不同，这里的状态类型不包含用户信息，因为这些信息将通过上下文传递。

2. 定义上下文类型

python 复制代码

class UserContext(TypedDict):
    user_name: str
    user_role: str

这里定义了一个UserContext类型，用于描述上下文数据的结构：

user_name：字符串类型，表示用户名
user_role：字符串类型，表示用户角色

使用TypedDict定义上下文类型可以提供类型提示和类型检查，确保上下文数据的正确性。

3. 定义使用上下文的节点函数

python 复制代码

def process_with_context(state: ContextState, runtime) -> dict:
    """使用上下文处理消息的节点"""
    # 访问上下文
    context = runtime.context
    user_name = context.get("user_name", "Guest")
    user_role = context.get("user_role", "User")
    
    processed_message = f"[{user_role}: {user_name}] {state['message']}"
    return {"message": processed_message}

这个节点函数与前面示例中的节点函数有一个重要区别：它接收两个参数：

state：当前状态
runtime：运行时对象，包含上下文信息

在函数内部，我们通过runtime.context访问上下文数据，并使用get方法获取特定字段的值（同时提供默认值以防字段不存在）。然后，我们使用上下文中的用户名和角色来处理消息，并返回更新后的状态。

4. 创建带上下文的StateGraph

python 复制代码

# 创建StateGraph，指定上下文类型
context_graph = StateGraph(ContextState, context_schema=UserContext)

创建StateGraph时，我们通过context_schema参数指定了上下文的类型。这一步是可选的，但强烈推荐，因为它可以提供类型提示和类型检查。

5. 添加节点和边

python 复制代码

# 添加节点
context_graph.add_node("process", process_with_context)

# 添加边
context_graph.add_edge(START, "process")
context_graph.add_edge("process", END)

这部分代码与前面的示例类似，添加了一个处理节点并定义了从START到处理节点再到END的边。

6. 编译和执行图，并提供上下文

python 复制代码

# 编译图
compiled_context_graph = context_graph.compile()

# 执行图，提供上下文
result = compiled_context_graph.invoke(
    {"message": "Hello, LangGraph!"}, 
    context={"user_name": "Alice", "user_role": "Admin"}
)
print(f"输入: {{'message': 'Hello, LangGraph!'}}")
print(f"上下文: {{'user_name': 'Alice', 'user_role': 'Admin'}}")
print(f"输出: {result}")

编译图后，我们使用invoke方法执行图。与前面的示例不同，这里我们传递了两个参数：

第一个参数是初始状态：{"message": "Hello, LangGraph!"}
第二个参数是上下文数据：context={"user_name": "Alice", "user_role": "Admin"}

执行结果显示，消息被成功处理，并包含了上下文中的用户信息。

执行流程分析

让我们详细分析一下整个图的执行流程：

初始化 ：invoke()方法接收初始状态{"message": "Hello, LangGraph!"}和上下文{"user_name": "Alice", "user_role": "Admin"}
执行process节点 ：从START开始，执行process_with_context节点
访问上下文 ：在节点函数内部，通过runtime.context访问上下文数据
处理消息 ：使用上下文中的用户名和角色处理消息，生成格式为[role: name] message的新消息
更新状态 ：返回更新后的状态{"message": "[Admin: Alice] Hello, LangGraph!"}
结束：从process节点连接到END节点，执行结束并返回最终状态

为什么使用Context？

Context在以下场景中特别有用：

用户身份信息：传递用户身份、角色等信息，而不需要将其作为状态的一部分
环境配置：传递环境特定的配置信息，如API密钥、数据库连接等
请求元数据：传递请求ID、时间戳等元数据
上下文感知处理：根据上下文信息定制处理逻辑
避免状态膨胀：将不参与状态更新的信息从状态中分离出来

优化

虽然这个示例很好地展示了Context的基础用法，但还有一些可以改进的地方：

使用数据类增强类型安全 ：可以使用Python的@dataclass装饰器定义上下文类型
python 复制代码

from dataclasses import dataclass

@dataclass class UserContext: user_name: str user_role: str = "User" # 提供默认值

在执行图时使用

context = UserContext(user_name="Alice", user_role="Admin") result = compiled_context_graph.invoke({"message": "Hello, LangGraph!"}, context=context)

python 复制代码

2. **添加上下文验证**：在节点函数中添加对上下文数据的验证
```python
 def process_with_context(state: ContextState, runtime) -> dict:
     context = runtime.context

     # 验证上下文数据
     if not isinstance(context.get("user_name"), str):
         raise ValueError("user_name must be a string")
     if context.get("user_role") not in ["Admin", "User", "Guest"]:
         raise ValueError("Invalid user_role")

     user_name = context.get("user_name", "Guest")
     user_role = context.get("user_role", "User")

     processed_message = f"[{user_role}: {user_name}] {state['message']}"
     return {"message": processed_message}

Context的实际应用场景

Context在实际应用中有广泛的用途：

多租户系统：在多租户系统中传递租户信息
用户会话：在用户会话期间保持用户状态和首选项
API集成：传递API密钥、认证令牌等敏感信息
日志记录：传递请求ID、跟踪信息等用于日志记录
环境切换：在开发、测试和生产环境之间切换配置
A/B测试：在A/B测试中传递实验分组信息

总结

通过本文的学习，我们了解了LangGraph中Context（上下文）的概念、实现方式和应用场景。Context是一个强大的机制，它允许我们在执行图时传递额外的信息，而不需要将这些信息作为状态的一部分。在实际应用中，你可以根据业务需求传递各种类型的上下文信息，构建更加灵活和强大的工作流系统。

通过合理使用Context可以：

保持状态的简洁性和专注性
提高代码的可维护性和可测试性
实现更加灵活和动态的工作流逻辑
避免状态膨胀和不必要的复杂性