Состояние графа: схема, редьюсеры, каналы | Курс LangGraph урок 2

"""Пример 1 урока 2: три способа задать схему состояния.

Граф из одного узла над стендом mini_repo собран трижды: на TypedDict, на
dataclass со значением по умолчанию и на модели Pydantic с проверкой входа.
Печатается, чем три схемы отличаются при запуске.
"""

from dataclasses import dataclass, field
from pathlib import Path

from langgraph.graph import END, START, StateGraph
from pydantic import BaseModel, ValidationError
from typing_extensions import TypedDict

REPO = Path(__file__).parent / "mini_repo"


class DictState(TypedDict):
    root: str
    suffix: str
    files: list[str]


def collect_dict(state: DictState) -> dict:
    """Собирает имена файлов и приписывает к ним пометку из состояния."""
    names = sorted(path.name for path in Path(state["root"]).glob("*.py"))
    return {"files": [name + state["suffix"] for name in names]}


@dataclass
class DataState:
    root: str
    suffix: str = " [py]"
    files: list[str] = field(default_factory=list)


def collect_data(state: DataState) -> dict:
    """То же самое, но состояние приходит объектом dataclass."""
    names = sorted(path.name for path in Path(state.root).glob("*.py"))
    return {"files": [name + state.suffix for name in names]}


class ModelState(BaseModel):
    root: str
    suffix: str = " [py]"
    files: list[str] = []


def collect_model(state: ModelState) -> dict:
    """То же самое, но состояние приходит моделью Pydantic."""
    names = sorted(path.name for path in Path(state.root).glob("*.py"))
    return {"files": [name + state.suffix for name in names]}


def build(schema, node) -> object:
    """Собирает граф из одного узла на переданной схеме состояния."""
    builder = StateGraph(schema)
    builder.add_node("collect", node)
    builder.add_edge(START, "collect")
    builder.add_edge("collect", END)
    return builder.compile()


# 1. TypedDict: ключа suffix на входе нет, значения по умолчанию тоже нет.
try:
    build(DictState, collect_dict).invoke({"root": str(REPO)})
except KeyError as error:
    print("1. TypedDict без suffix:", type(error).__name__, error)

# 2. TypedDict: ключ передан явно.
result = build(DictState, collect_dict).invoke({"root": str(REPO), "suffix": " [py]"})
print("2. TypedDict с suffix:", result["files"][:2])

# 3. dataclass: тот же вход без suffix, значение берётся из схемы.
result = build(DataState, collect_data).invoke({"root": str(REPO)})
print("3. dataclass без suffix:", result["files"][:2])
print("   тип результата:", type(result).__name__)

# 4. Pydantic: неверный тип на входе.
try:
    build(ModelState, collect_model).invoke({"root": 42})
except ValidationError as error:
    print("4. Pydantic, неверный тип:", str(error).splitlines()[0])

# 5. Pydantic: верный вход и тип результата.
result = build(ModelState, collect_model).invoke({"root": str(REPO)})
print("5. Pydantic, верный вход:", result["files"][:2])
print("   тип результата:", type(result).__name__)

# Вывод:
# 1. TypedDict без suffix: KeyError 'suffix'
# 2. TypedDict с suffix: ['app.py [py]', 'config.py [py]']
# 3. dataclass без suffix: ['app.py [py]', 'config.py [py]']
#    тип результата: dict
# 4. Pydantic, неверный тип: 1 validation error for ModelState
# 5. Pydantic, верный вход: ['app.py [py]', 'config.py [py]']
#    тип результата: dict

new_value = reducer(left=current_state[key], right=node_update[key])

"""Пример 2 урока 2: редьюсер сводит значение ключа с правкой узла.

Два узла над стендом mini_repo пишут находки в один ключ состояния. Граф
собран трижды: без редьюсера, с редьюсером на цепочке узлов и с редьюсером на
параллельной паре. Редьюсер печатает свои аргументы. В конце показана
ошибка StateGraph на редьюсер неверной формы.
"""

from pathlib import Path
from typing import Annotated

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

REPO = Path(__file__).parent / "mini_repo"


def append_findings(left: list[str], right: list[str]) -> list[str]:
    """Складывает накопленное значение ключа (left) с правкой узла (right)."""
    print(f"   редьюсер: left={left}, right={right}")
    return left + right


class PlainState(TypedDict):
    root: str
    findings: list[str]


class ReducedState(TypedDict):
    root: str
    findings: Annotated[list[str], append_findings]


def count_lines(state) -> dict:
    """Считает строки во всех файлах стенда."""
    root = Path(state["root"])
    total = sum(
        len(path.read_text(encoding="utf-8").splitlines())
        for path in sorted(root.glob("*.py"))
    )
    return {"findings": [f"строк {total}"]}


def count_imports(state) -> dict:
    """Считает строки импорта во всех файлах стенда."""
    root = Path(state["root"])
    total = 0
    for path in sorted(root.glob("*.py")):
        for line in path.read_text(encoding="utf-8").splitlines():
            if line.startswith("import ") or line.startswith("from "):
                total += 1
    return {"findings": [f"импортов {total}"]}


def chain(schema) -> object:
    """Два узла цепочкой: count_lines, затем count_imports."""
    builder = StateGraph(schema)
    builder.add_node("count_lines", count_lines)
    builder.add_node("count_imports", count_imports)
    builder.add_edge(START, "count_lines")
    builder.add_edge("count_lines", "count_imports")
    builder.add_edge("count_imports", END)
    return builder.compile()


def pair(schema) -> object:
    """Те же два узла параллельно: обе правки приходят в одном шаге."""
    builder = StateGraph(schema)
    builder.add_node("count_lines", count_lines)
    builder.add_node("count_imports", count_imports)
    builder.add_edge(START, "count_lines")
    builder.add_edge(START, "count_imports")
    builder.add_edge("count_lines", END)
    builder.add_edge("count_imports", END)
    return builder.compile()


# 1. Цепочка без редьюсера: вторая правка затирает первую.
print("1. цепочка, ключ без редьюсера")
result = chain(PlainState).invoke({"root": str(REPO), "findings": []})
print("   итог:", result["findings"])

# 2. Та же цепочка с редьюсером: правки копятся.
print("2. цепочка, ключ с редьюсером")
result = chain(ReducedState).invoke({"root": str(REPO)})
print("   итог:", result["findings"])

# 3. Параллельная пара с редьюсером: обе правки в одном шаге.
print("3. параллельная пара, ключ с редьюсером")
result = pair(ReducedState).invoke({"root": str(REPO)})
print("   итог:", result["findings"])


# 4. Редьюсер неверной формы: один аргумент вместо двух.
def only_one(right: list[str]) -> list[str]:
    return right


class BrokenState(TypedDict):
    findings: Annotated[list[str], only_one]


try:
    StateGraph(BrokenState)
except ValueError as error:
    print("4. редьюсер с одним аргументом:", error)

# Вывод:
# 1. цепочка, ключ без редьюсера
#    итог: ['импортов 3']
# 2. цепочка, ключ с редьюсером
#    редьюсер: left=[], right=['строк 44']
#    редьюсер: left=['строк 44'], right=['импортов 3']
#    итог: ['строк 44', 'импортов 3']
# 3. параллельная пара, ключ с редьюсером
#    редьюсер: left=[], right=['импортов 3']
#    редьюсер: left=['импортов 3'], right=['строк 44']
#    итог: ['импортов 3', 'строк 44']
# 4. редьюсер с одним аргументом: Invalid reducer signature. Expected (a, b) -> c. Got (right: list[str]) -> list[str]
#
# В пункте 3 правки шага применены по алфавиту имён узлов,
# поэтому count_imports идёт раньше count_lines.

"""Пример 3 урока 2: почему списку сообщений нужен add_messages.

Три узла складывают переписку про стенд mini_repo: вопрос, черновик ответа и
его исправленная версия с тем же идентификатором. Граф собран трижды: ключ без
редьюсера, ключ с operator.add и ключ с add_messages. Печатается, что каждый вариант
ключа делает с исправленным сообщением и с краткой формой записи.
"""

import operator
from pathlib import Path
from typing import Annotated

from langchain.messages import AIMessage, AnyMessage
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict

REPO = Path(__file__).parent / "mini_repo"


class PlainState(TypedDict):
    messages: list[AnyMessage]


class AddState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]


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


def ask(state) -> dict:
    """Кладёт вопрос пользователя краткой формой записи."""
    return {"messages": [{"role": "user", "content": "Сколько файлов в репозитории?"}]}


def draft(state) -> dict:
    """Кладёт черновик ответа с явным идентификатором."""
    return {"messages": [AIMessage(content="Файлов несколько.", id="answer")]}


def fix(state) -> dict:
    """Возвращает исправленный ответ под тем же идентификатором."""
    count = len(sorted(REPO.glob("*.py")))
    return {"messages": [AIMessage(content=f"Файлов {count}.", id="answer")]}


def build(schema) -> object:
    """Собирает цепочку из трёх узлов на переданной схеме."""
    builder = StateGraph(schema)
    builder.add_node("ask", ask)
    builder.add_node("draft", draft)
    builder.add_node("fix", fix)
    builder.add_edge(START, "ask")
    builder.add_edge("ask", "draft")
    builder.add_edge("draft", "fix")
    builder.add_edge("fix", END)
    return builder.compile()


def show(title: str, schema) -> None:
    """Печатает итоговую переписку одной строкой на сообщение."""
    print(title)
    try:
        result = build(schema).invoke({"messages": []})
    except Exception as error:
        print("   упало:", type(error).__name__, error)
        return
    for message in result["messages"]:
        kind = type(message).__name__
        text = message.content if hasattr(message, "content") else message
        print(f"   {kind}(id={getattr(message, 'id', None)}): {text}")


show("1. ключ без редьюсера", PlainState)
show("2. ключ с operator.add", AddState)
show("3. ключ с add_messages", MessagesGraphState)

# Вывод:
# 1. ключ без редьюсера
#    AIMessage(id=answer): Файлов 4.
# 2. ключ с operator.add
#    dict(id=None): {'role': 'user', 'content': 'Сколько файлов в репозитории?'}
#    AIMessage(id=answer): Файлов несколько.
#    AIMessage(id=answer): Файлов 4.
# 3. ключ с add_messages
#    HumanMessage(id=d36b75fa-1ef2-418a-823a-8b47ce57379b): Сколько файлов в репозитории?
#    AIMessage(id=answer): Файлов 4.
#
# Идентификатор HumanMessage случайный, при своём запуске он будет другим.

from langgraph.graph import MessagesState

class State(MessagesState):
    documents: list[str]

"""Пример 4 урока 2: обход редьюсера типом Overwrite.

Ключ errors копит сообщения об ошибках обхода стенда. Показаны три попытки
очистить его перед новым заходом: пустым списком, типом Overwrite и его же
записью словарём. В конце показаны две правки Overwrite в одном шаге.
"""

import operator
from typing import Annotated

from langgraph.errors import InvalidUpdateError
from langgraph.graph import END, START, StateGraph
from langgraph.types import Overwrite
from typing_extensions import TypedDict


class State(TypedDict):
    errors: Annotated[list[str], operator.add]


def record_error(state: State) -> dict:
    """Записывает находку в общий буфер ошибок."""
    return {"errors": ["в utils.py функции без докстрингов"]}


def clear_with_empty(state: State) -> dict:
    """Пробует очистить буфер пустым списком."""
    return {"errors": []}


def clear_with_overwrite(state: State) -> dict:
    """Очищает буфер в обход редьюсера."""
    return {"errors": Overwrite([])}


def clear_with_dict(state: State) -> dict:
    """То же самое записью словарём."""
    return {"errors": {"__overwrite__": []}}


def chain(second) -> object:
    """Собирает цепочку: запись ошибки, затем переданный узел очистки."""
    builder = StateGraph(State)
    builder.add_node("record_error", record_error)
    builder.add_node("clear", second)
    builder.add_edge(START, "record_error")
    builder.add_edge("record_error", "clear")
    builder.add_edge("clear", END)
    return builder.compile()


print("1. очистка пустым списком:", chain(clear_with_empty).invoke({"errors": []}))
print("2. очистка типом Overwrite:", chain(clear_with_overwrite).invoke({"errors": []}))
print("3. очистка словарём:", chain(clear_with_dict).invoke({"errors": []}))

# 4. Два узла одного шага пишут Overwrite в один ключ.
builder = StateGraph(State)
builder.add_node("clear_a", clear_with_overwrite)
builder.add_node("clear_b", clear_with_dict)
builder.add_edge(START, "clear_a")
builder.add_edge(START, "clear_b")
builder.add_edge("clear_a", END)
builder.add_edge("clear_b", END)

try:
    builder.compile().invoke({"errors": ["старая находка"]})
except InvalidUpdateError as error:
    print("4. два Overwrite в одном шаге:", str(error).splitlines()[0])

# Вывод:
# 1. очистка пустым списком: {'errors': ['в utils.py функции без докстрингов']}
# 2. очистка типом Overwrite: {'errors': []}
# 3. очистка словарём: {'errors': []}
# 4. два Overwrite в одном шаге: Can receive only one Overwrite value per super-step.

"""Пример 5 урока 2: раздельные схемы входа и выхода, приватный ключ и контекст.

Граф строит карточку стенда mini_repo. На входе принимается один ключ, наружу
отдаётся тоже один, а промежуточные размеры файлов идут приватным ключом от
первого узла ко второму. Настройки запуска передаются контекстом.
"""

from dataclasses import dataclass
from pathlib import Path

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

REPO = Path(__file__).parent / "mini_repo"


class InputState(TypedDict):
    root: str


class OutputState(TypedDict):
    card: str


class OverallState(TypedDict):
    root: str
    card: str
    files: list[str]


class PrivateState(TypedDict):
    sizes: dict[str, int]


@dataclass
class RunContext:
    label: str = "обзор"
    max_files: int = 2


def scan(state: InputState, runtime: Runtime[RunContext]) -> PrivateState:
    """Считает строки в первых файлах стенда, число файлов задаёт контекст запуска."""
    print("   узел scan получил runtime.context =", runtime.context)
    context = runtime.context or RunContext()
    paths = sorted(Path(state["root"]).glob("*.py"))[: context.max_files]
    sizes = {
        path.name: len(path.read_text(encoding="utf-8").splitlines())
        for path in paths
    }
    return {"sizes": sizes}


def make_card(state: PrivateState, runtime: Runtime[RunContext]) -> OverallState:
    """Собирает карточку из приватных размеров и подписи из контекста."""
    context = runtime.context or RunContext()
    parts = [f"{name}: {lines}" for name, lines in state["sizes"].items()]
    return {
        "card": f"{context.label}, " + "; ".join(parts),
        "files": list(state["sizes"]),
    }


builder = StateGraph(
    OverallState,
    context_schema=RunContext,
    input_schema=InputState,
    output_schema=OutputState,
)
builder.add_node("scan", scan)
builder.add_node("make_card", make_card)
builder.add_edge(START, "scan")
builder.add_edge("scan", "make_card")
builder.add_edge("make_card", END)

graph = builder.compile()

print("1. без контекста вовсе:", graph.invoke({"root": str(REPO)}))
print("2. пустой контекст:", graph.invoke({"root": str(REPO)}, context={}))
print(
    "3. свой контекст:",
    graph.invoke({"root": str(REPO)}, context={"label": "полный обход", "max_files": 4}),
)

print("4. снимки состояния потоком:")
for chunk in graph.stream({"root": str(REPO)}, stream_mode="values", version="v2"):
    print("   ключи снимка:", sorted(chunk["data"]))

# Вывод:
#    узел scan получил runtime.context = None
# 1. без контекста вовсе: {'card': 'обзор, app.py: 13; config.py: 6'}
#    узел scan получил runtime.context = RunContext(label='обзор', max_files=2)
# 2. пустой контекст: {'card': 'обзор, app.py: 13; config.py: 6'}
#    узел scan получил runtime.context = RunContext(label='полный обход', max_files=4)
# 3. свой контекст: {'card': 'полный обход, app.py: 13; config.py: 6; db.py: 15; utils.py: 10'}
# 4. снимки состояния потоком:
#    ключи снимка: ['root']
#    узел scan получил runtime.context = None
#    ключи снимка: ['root', 'sizes']
#    ключи снимка: ['card', 'files', 'root', 'sizes']

"""Пример 6 урока 2: каналы, в которые превращается схема состояния.

Схема из четырёх ключей собрана в граф, и у собранного графа печатается тип
канала для каждого ключа. Дальше два узла цепочки пишут в ключ с редьюсером
и в ключ с Topic. В ключ messages не пишет никто, он показывает начальное
значение канала с редьюсером.
"""

import operator
from typing import Annotated

from langchain.messages import AnyMessage
from langgraph.channels import Topic
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages
from typing_extensions import TypedDict


class State(TypedDict):
    root: str
    findings: Annotated[list[str], operator.add]
    messages: Annotated[list[AnyMessage], add_messages]
    events: Annotated[list[str], Topic(str, accumulate=True)]


def first(state: State) -> dict:
    """Пишет по одной записи в ключ с редьюсером и в ключ с Topic."""
    return {"findings": ["первая находка"], "events": ["узел first"]}


def second(state: State) -> dict:
    """То же самое из второго узла цепочки."""
    return {"findings": ["вторая находка"], "events": ["узел second"]}


builder = StateGraph(State)
builder.add_node("first", first)
builder.add_node("second", second)
builder.add_edge(START, "first")
builder.add_edge("first", "second")
builder.add_edge("second", END)

graph = builder.compile()

print("ТИПЫ КАНАЛОВ:")
for key in ("root", "findings", "messages", "events"):
    channel = graph.channels[key]
    print(f"   {key}: {type(channel).__name__}")

print()
result = graph.invoke({"root": "mini_repo"})
print("НА ВЫХОДЕ:")
for key in ("root", "findings", "messages", "events"):
    print(f"   {key}: {result.get(key)}")

# Вывод:
# ТИПЫ КАНАЛОВ:
#    root: LastValue
#    findings: BinaryOperatorAggregate
#    messages: BinaryOperatorAggregate
#    events: Topic
#
# НА ВЫХОДЕ:
#    root: mini_repo
#    findings: ['первая находка', 'вторая находка']
#    messages: []
#    events: ['узел first', 'узел second']