LangGraph. Делаем State для LLM

Часто многие LLM требуют историю сообщений, чтобы лучше понимать контекст.

Так как передача истории сообщений довольно частая ситуация, то рассмотрим на примере:

Как использовать State для определения схемы графа
Как использовать reducers для управления обработкой обновлений состояния.

Сперва установим библиотеки

pip install -U langgraph langchain

Пример простого графа

Состоянием (State) в LangGraph может быть TypedDict, Pydantic model или dataclass. Ниже мы будем использовать TypedDict.

Определим State

from langchain_core.messages import AnyMessage
from typing_extensions import TypedDict


class State(TypedDict):
    messages: list[AnyMessage]
    extra_field: int

В состоянии будем хранить историю сообщений.

Определим структуру графа

Давайте построим пример графа с одним узлом (node). Наш node - это просто функция Python, которая считывает состояние нашего графа и вносит в него изменения. Первым аргументом этой функции всегда будет состояние:

from langchain_core.messages import AIMessage


def node(state: State):
    messages = state["messages"]
    new_message = AIMessage("Hello!")
    return {"messages": messages + [new_message], "extra_field": 10}

Этот узел (node) просто добавляет сообщение в наш список сообщений и заполняет дополнительное поле. Как видим, мы храним всю историю сообщений и дополняем новым сообщением историю.

Давайте далее определим простой граф, содержащий этот узел. Мы используем StateGraph для определения графа, который работает с этим состоянием. Затем мы используем add_node для заполнения нашего графа.

from langgraph.graph import StateGraph

graph_builder = StateGraph(State)
graph_builder.add_node(node)
graph_builder.set_entry_point("node")
graph = graph_builder.compile()

LangGraph предоставляет встроенные утилиты для визуализации вашего графа. Давайте рассмотрим наш граф.

from IPython.display import Image, display

display(Image(graph.get_graph().draw_mermaid_png()))

Использование графа

from langchain_core.messages import HumanMessage

result = graph.invoke({"messages": [HumanMessage("Hi")]})
print(result)

{'messages': [
HumanMessage(content='Hi', additional_kwargs={}, response_metadata={}),
AIMessage(content='Hello!', additional_kwargs={}, response_metadata={})
], 'extra_field': 10}

Для удобства мы часто проверяем содержимое объектов сообщений с помощью pretty-print:

for message in result["messages"]:
    message.pretty_print()

========================== Human Message ============================

Hi
=========================== Ai Message ==============================

Hello!

Обновление состояния графа с помощью редьюсеров (reducers)

Каждый ключ в состоянии может иметь свою собственную независимую функцию-reducer, которая управляет тем, как применяются обновления с узлов. Если функция-reducer явно не указана, то предполагается, что все обновления ключа должны переопределять ее.

Для схем состояний TypedDict мы можем определить reducer, снабдив соответствующее поле состояния аннотацией с помощью функции reducer.

В предыдущем примере наш узел обновил ключ "сообщения" в состоянии, добавив к нему сообщение. Ниже мы добавим к этому ключу reducer, чтобы обновления добавлялись автоматически:

from typing_extensions import Annotated


def add(left: list[AnyMessage], right: list[AnyMessage]):
    return left + right


class State(TypedDict):
    messages: Annotated[list[AnyMessage], add]
    extra_field: int

Теперь наш узел можно упростить:

def node(state: State):
    new_message = AIMessage("Hello!")
    return {"messages": [new_message], "extra_field": 10}

Теперь наша объявленная функция add будет сама делать конкатенацию списков.

from langgraph.graph import START


graph_builder = StateGraph(State)
graph = graph_builder.add_node(node).add_edge(START, "node").compile()

result = graph.invoke({"messages": [HumanMessage("Hi")]})

for message in result["messages"]:
    message.pretty_print()

========================== Human Message ============================

Hi
=========================== Ai Message ==============================

Hello!

MessagesState

LangGraph имеет встроенный reducer add_messages, который обрабатывает эти соображения:

from langgraph.graph.message import add_messages


class State(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    extra_field: int


def node(state: State):
    new_message = AIMessage("Hello!")
    return {"messages": [new_message], "extra_field": 10}


graph_builder = StateGraph(State)
graph = graph_builder.add_node(node).set_entry_point("node").compile()

input_message = {"role": "user", "content": "Hi"}

result = graph.invoke({"messages": [input_message]})

for message in result["messages"]:
    message.pretty_print()

========================== Human Message ============================

Hi
=========================== Ai Message ==============================

Hello!

Это универсальное представление состояния для приложений, использующих модели чата. Для удобства в LangGraph встроен MessagesState, так что мы можем написать:

from langgraph.graph import MessagesState


class State(MessagesState):
    extra_field: int

Таким образом, данный MessageState сам делает структуру:

class State(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

А мы лишь дальше от него наследуемся и можем переиспользовать функционал. Чтобы вручную самим не объявлять.

✨Основной Telegram-канал со всеми обновлениями✨

Вам также может понравиться:

✔️ Обучение Python с нуля за 9 месяцев 🔝

#ai #ии #langchain #langgraph #llm #openai