Context engineering: чем управлять в контексте агента | Курс LangChain урок 14

"""Общая сборка модели для примеров урока 14.

Тот же модуль, что course_model.py уроков 4-13, без изменений: у build_model есть
необязательный первый аргумент с именем модели, а рядом лежит gateway_kwargs()
для примеров, которые собирают модель сами.

Файл .env берётся тот же, что в уроке 0. Положите его рядом с этой папкой или
выше по дереву: load_dotenv() ищет файл начиная с папки этого модуля и поднимается
вверх.
"""

import os

from dotenv import load_dotenv
from langchain.chat_models import init_chat_model

load_dotenv()


def gateway_kwargs():
    """Возвращает аргументы доступа к провайдеру: имя, адрес, ключ.

    Нужны примерам, которые вызывают init_chat_model сами. У настраиваемой модели
    имени модели при создании нет, поэтому build_model ей не подходит.
    """
    base_url = os.getenv("MODEL_BASE_URL")

    if base_url:
        return {
            "model_provider": "openai",
            "base_url": base_url,
            "api_key": os.environ["OPENAI_API_KEY"],
        }

    return {}


def build_model(model_name=None, **kwargs):
    """Собирает модель курса.

    model_name без значения означает модель из переменной MODEL_NAME. Явное имя
    нужно примерам, где моделей в приложении больше одной.

    Все именованные аргументы уходят в init_chat_model как есть: temperature,
    max_tokens, timeout, max_retries, rate_limiter, profile и прочее из раздела
    Parameters.
    """
    model_name = model_name or os.environ["MODEL_NAME"]
    access = gateway_kwargs()

    if access:
        # Путь для любого адреса, совместимого с OpenAI Chat Completions API.
        return init_chat_model(model=model_name, **access, **kwargs)

    # Путь напрямую к провайдеру.
    return init_chat_model(model_name, **kwargs)

"""Пример 1 урока 14: из чего собран запрос к модели и сколько весит каждая часть.

Middleware печатает строку на каждый вызов модели: сколько сообщений уходит,
во что оценены системное сообщение и сообщения каждого типа, сколько отправлено
инструментов и какой вход насчитал провайдер.

Задача выбрана так, чтобы агент вызвал инструменты несколько раз, и в таблице
стало видно, какая именно часть запроса растёт от вызова к вызову.
"""

from collections import Counter
from itertools import count
from typing import Callable

from langchain.agents import create_agent
from langchain.agents.middleware import ModelRequest, ModelResponse, wrap_model_call
from langchain_core.messages.utils import count_tokens_approximately

from course_model import build_model

ORDERS = {
    "Z-771": {"city": "Казань", "status": "укомплектован"},
    "Z-772": {"city": "Омск", "status": "ждёт оплаты"},
}

DELIVERY = {"Казань": 2, "Омск": 5}

RULES = (
    "Регламент приёмки груза на складе города {city}.\n"
    "1. Машина встаёт под разгрузку по электронной очереди, талон берут на "
    "въезде, время в талоне сдвигать нельзя.\n"
    "2. Паллеты принимают по счёту мест и по весу брутто, расхождение больше "
    "двух процентов фиксируют актом в тот же день.\n"
    "3. Упаковку вскрывают только при видимом повреждении обрешётки, вскрытие "
    "снимают на камеру и прикладывают файл к акту.\n"
    "4. Температурный груз замеряют дважды: в момент открытия дверей и через "
    "пятнадцать минут после начала разгрузки.\n"
    "5. Брак раскладывают по трём коробам: бой, подмочка, пересорт. Короб "
    "опечатывают и передают на карантинную зону до конца смены.\n"
    "6. Документы возвращают водителю с отметкой о приёмке, копию подшивают в "
    "журнал склада и дублируют в учётную систему до 20:00 того же дня.\n"
    "7. Расхождение по количеству больше пяти мест останавливает приёмку: "
    "вызывают старшего смены и представителя перевозчика.\n"
)


def find_order(order_id: str) -> str:
    """Находит заказ по номеру вида Z-771 и возвращает город и статус."""
    order = ORDERS.get(order_id)

    if order is None:
        return f"Заказ {order_id} не найден."

    return f"Заказ {order_id}: город {order['city']}, статус {order['status']}."


def delivery_days(city: str) -> str:
    """Сообщает срок доставки в днях до указанного города."""
    days = DELIVERY.get(city)

    if days is None:
        return f"Город {city} не обслуживается."

    return f"Доставка до города {city} занимает {days} дн."


def acceptance_rules(city: str) -> str:
    """Возвращает полный регламент приёмки груза для склада города."""
    return RULES.format(city=city)


step = count(1)
sent_tools = []


@wrap_model_call
def breakdown(
    request: ModelRequest,
    handler: Callable[[ModelRequest], ModelResponse],
) -> ModelResponse:
    """Считает состав запроса до вызова модели, печатает его с расходом после ответа."""
    sent_tools[:] = [tool.name for tool in request.tools]
    prompt = (
        count_tokens_approximately([request.system_message])
        if request.system_message
        else 0
    )

    by_type = Counter()
    for message in request.messages:
        by_type[message.type] += count_tokens_approximately([message])

    estimate = prompt + sum(by_type.values())
    response = handler(request)
    # Расход возвращает не каждый провайдер, поэтому пустой словарь как запасной путь.
    usage = response.result[0].usage_metadata or {}

    print(
        f"{next(step):>5}  {len(request.messages):>9}  {prompt:>6}  "
        f"{by_type['human']:>6}  {by_type['ai']:>4}  {by_type['tool']:>6}  "
        f"{len(request.tools):>6}  {estimate:>6}  {usage.get('input_tokens', 0):>5}"
    )

    return response


agent = create_agent(
    model=build_model(temperature=0, max_tokens=1024),
    tools=[find_order, delivery_days, acceptance_rules],
    system_prompt=(
        "Вы помощник логистического отдела. Отвечайте по-русски. Номер заказа "
        "всегда проверяйте инструментом, по памяти не отвечайте. Если в ответе "
        "нужен срок доставки, берите его инструментом для города из заказа. "
        "Правила приёмки пересказывайте своими словами, целиком регламент не "
        "цитируйте."
    ),
    middleware=[breakdown],
)

print(
    f"{'вызов':>5}  {'сообщений':>9}  {'промпт':>6}  {'human':>6}  {'ai':>4}  "
    f"{'tool':>6}  {'инстр':>6}  {'оценка':>6}  {'вход':>5}"
)

result = agent.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": (
                    "Заказ Z-771: когда он приедет и что учесть при приёмке?"
                ),
            }
        ]
    }
)

print()
print("ИНСТРУМЕНТЫ В ЗАПРОСЕ:", ", ".join(sent_tools))
print("ОТВЕТ:", result["messages"][-1].text.replace("\n", " "))

# Вывод:
# вызов  сообщений  промпт   human    ai    tool   инстр  оценка   вход
#     1          1      76      18     0       0       3      94    507
#     2          3      76      18    45      26       3     165    596
#     3          6      76      18   129     291       3     514   1031
#
# ИНСТРУМЕНТЫ В ЗАПРОСЕ: find_order, delivery_days, acceptance_rules
# ОТВЕТ:  Вот вся информация по заказу Z-771:  **📦 Статус и доставка:** - Заказ
# **Z-771** направлен в **Казань**, статус — **укомплектован**. - Срок доставки до
# Казани — **2 дня**.  **📋 Что учесть при приёмке (основные моменты):**  1.
# **Электронная очередь** — машина встаёт под разгрузку по талону, который берут на
# въезде. Время в талоне сдвигать нельзя. ... 7. **Важно:** если расхождение по
# количеству превышает **5 мест** — приёмка останавливается, вызывают старшего смены
# и представителя перевозчика.
# [вывод подрезан: между пунктами 1 и 7 идут ещё пять пунктов регламента (паллеты,
# вскрытие упаковки, температурный груз, брак, документы) и завершающая фраза
# "Если нужны уточнения — обращайтесь!"]

"""Пример 2 урока 14: инструкция собирается под собеседника и под ход разговора.

Один и тот же агент отвечает на один и тот же вопрос трижды. Меняются два
источника: конфигурация запуска (кто спрашивает) и состояние разговора (сколько
сообщений уже накопилось). Промпт печатается перед каждым ответом, чтобы было
видно, что именно ушло в модель.
"""

from dataclasses import dataclass

from langchain.agents import create_agent
from langchain.agents.middleware import ModelRequest, dynamic_prompt

from course_model import build_model

QUESTION = "Где сейчас заказ Z-771 и когда он приедет?"

HISTORY = [
    {"role": "user", "content": "Добрый день, нужна помощь по заказам."},
    {"role": "assistant", "content": "Здравствуйте. Назовите номер заказа."},
    {"role": "user", "content": "Сначала общий вопрос: вы работаете в выходные?"},
    {"role": "assistant", "content": "Склад принимает и выдаёт груз по субботам до 15:00."},
    {"role": "user", "content": "Понял, спасибо."},
    {"role": "assistant", "content": "Рад помочь. Что-то ещё?"},
]


@dataclass
class Session:
    """Конфигурация запуска: кто пришёл и из какого города."""

    role: str
    city: str


def find_order(order_id: str) -> str:
    """Находит заказ по номеру вида Z-771 и возвращает город, статус и срок."""
    if order_id != "Z-771":
        return f"Заказ {order_id} не найден."

    return (
        "Заказ Z-771: город Казань, статус укомплектован, ячейка хранения B-14, "
        "перевозчик СДЭК, срок доставки 2 дн."
    )


@dynamic_prompt
def prompt_for_session(request: ModelRequest) -> str:
    """Собирает системный промпт из конфигурации запуска и длины разговора."""
    session = request.runtime.context

    parts = ["Вы помощник службы доставки. Отвечайте по-русски."]

    if session.role == "оператор":
        parts.append(
            "Собеседник, это оператор склада: называйте ячейку хранения и "
            "перевозчика."
        )
    else:
        parts.append(
            "Собеседник, это клиент: внутренние коды ячеек и имя перевозчика "
            "не называйте."
        )

    parts.append(f"Город собеседника: {session.city}.")

    if len(request.messages) > 4:
        parts.append("Разговор длинный, отвечайте одной фразой.")

    prompt = " ".join(parts)
    print("ПРОМПТ:", prompt)

    return prompt


agent = create_agent(
    model=build_model(temperature=0, max_tokens=384),
    tools=[find_order],
    middleware=[prompt_for_session],
    context_schema=Session,
)

RUNS = [
    ("клиент, начало разговора", [], Session(role="клиент", city="Казань")),
    ("оператор, начало разговора", [], Session(role="оператор", city="Казань")),
    ("оператор, разговор идёт давно", HISTORY, Session(role="оператор", city="Казань")),
]

for label, history, session in RUNS:
    print()
    print("ЗАПУСК:", label)
    result = agent.invoke(
        {"messages": [*history, {"role": "user", "content": QUESTION}]},
        context=session,
    )
    print("ОТВЕТ: ", result["messages"][-1].text.replace("\n", " "))

# Вывод:
# (полный вывод длиннее десяти строк, второй запуск, "оператор, начало
# разговора", опущен: он устроен так же, как первый)
# ЗАПУСК: клиент, начало разговора
# ПРОМПТ: Вы помощник службы доставки. Отвечайте по-русски. Собеседник, это клиент: внутренние коды ячеек и имя перевозчика не называйте. Город собеседника: Казань.
# (строка ПРОМПТ печатается на каждый вызов модели, здесь их два: тем же текстом)
# ОТВЕТ:   Ваш заказ Z-771 уже укомплектован и находится в Казани. Доставка займёт 2 дня.
#
# ЗАПУСК: оператор, разговор идёт давно
# ПРОМПТ: Вы помощник службы доставки. Отвечайте по-русски. Собеседник, это оператор склада: называйте ячейку хранения и перевозчика. Город собеседника: Казань. Разговор длинный, отвечайте одной фразой.
# (та же строка печатается ещё раз тем же текстом)
# ОТВЕТ:  Заказ Z-771 в Казани, ячейка B-14, перевозчик СДЭК, доставка через 2 дня.

"""Пример 3 урока 14: набор инструментов под права собеседника и цена этого набора.

Агент один, инструментов восемь, задача одна и та же. Middleware оставляет в
запросе только те инструменты, которые разрешены роли из конфигурации запуска.
Наблюдатель печатает, сколько инструментов ушло в модель и во что провайдер
оценил вход.
"""

from dataclasses import dataclass
from typing import Callable

from langchain.agents import create_agent
from langchain.agents.middleware import ModelRequest, ModelResponse, wrap_model_call

from course_model import build_model

ORDERS = [
    {"id": "Z-771", "city": "Казань", "status": "укомплектован"},
    {"id": "Z-773", "city": "Казань", "status": "застрял на складе"},
    {"id": "Z-774", "city": "Казань", "status": "застрял на складе"},
    {"id": "Z-775", "city": "Омск", "status": "в пути"},
]


@dataclass
class Access:
    """Конфигурация запуска: под какой ролью работает собеседник."""

    role: str


def read_orders(city: str) -> str:
    """Возвращает список заказов по городу вместе со статусом каждого."""
    rows = [o for o in ORDERS if o["city"] == city]

    if not rows:
        return f"Заказов по городу {city} нет."

    return "; ".join(f"{o['id']} - {o['status']}" for o in rows)


def read_stock(sku: str) -> str:
    """Возвращает остаток товара по артикулу на всех складах."""
    return f"Артикул {sku}: Казань 7 шт., Омск 0 шт."


def read_routes(city: str) -> str:
    """Возвращает список маршрутов, которые проходят через город."""
    return f"Через город {city} идут маршруты M-1 и M-4."


def read_notes(order_id: str) -> str:
    """Возвращает комментарии сотрудников к заказу."""
    return f"Комментариев к заказу {order_id} нет."


def write_order(order_id: str, status: str) -> str:
    """Меняет статус заказа на указанный."""
    return f"Заказ {order_id} переведён в статус {status}."


def cancel_order(order_id: str) -> str:
    """Отменяет заказ по номеру."""
    return f"Заказ {order_id} отменён."


def refund_payment(order_id: str) -> str:
    """Возвращает деньги за заказ на карту плательщика."""
    return f"По заказу {order_id} оформлен возврат."


def delete_order(order_id: str) -> str:
    """Удаляет заказ из системы без возможности восстановления."""
    return f"Заказ {order_id} удалён."


ALL_TOOLS = [
    read_orders,
    read_stock,
    read_routes,
    read_notes,
    write_order,
    cancel_order,
    refund_payment,
    delete_order,
]


@wrap_model_call
def tools_by_role(
    request: ModelRequest,
    handler: Callable[[ModelRequest], ModelResponse],
) -> ModelResponse:
    """Оставляет в запросе инструменты, разрешённые роли, и печатает цену набора."""
    role = request.runtime.context.role

    if role != "оператор":
        request = request.override(
            tools=[t for t in request.tools if t.name.startswith("read_")]
        )

    response = handler(request)
    # Расход возвращает не каждый провайдер, поэтому пустой словарь как запасной путь.
    usage = response.result[0].usage_metadata or {}

    print(
        f"{role:<10}  {len(request.tools):>6}  {usage.get('input_tokens', 0):>5}  "
        f"{', '.join(t.name for t in request.tools)}"
    )

    return response


agent = create_agent(
    model=build_model(temperature=0, max_tokens=384),
    tools=ALL_TOOLS,
    system_prompt=(
        "Вы помощник логистического отдела. Отвечайте по-русски, одной фразой. "
        "Данные берите инструментами, по памяти не отвечайте."
    ),
    middleware=[tools_by_role],
    context_schema=Access,
)

QUESTION = "Сколько заказов застряло на складе в Казани?"

print(f"{'роль':<10}  {'инстр':>6}  {'вход':>5}  инструменты в запросе")

for role in ("оператор", "наблюдатель"):
    result = agent.invoke(
        {"messages": [{"role": "user", "content": QUESTION}]},
        context=Access(role=role),
    )
    answer = result["messages"][-1].text.replace("\n", " ")
    print(f"{role:<10}  ОТВЕТ: {answer}")

# Вывод:
# роль         инстр   вход  инструменты в запросе
# оператор         8    739  read_orders, read_stock, read_routes, read_notes, write_order, cancel_order, refund_payment, delete_order
# оператор         8    852  read_orders, read_stock, read_routes, read_notes, write_order, cancel_order, refund_payment, delete_order
# оператор    ОТВЕТ:  В Казани на складе застряло 2 заказа (Z-773 и Z-774).
# наблюдатель       4    508  read_orders, read_stock, read_routes, read_notes
# наблюдатель       4    600  read_orders, read_stock, read_routes, read_notes
# наблюдатель  ОТВЕТ:  В Казани на складе застряло 2 заказа (Z-773 и Z-774).

"""Пример 4 урока 14: что инструмент отдаёт модели и что он кладёт в состояние.

Два агента решают одну задачу по одному и тому же отчёту. Первый получает отчёт
целиком текстом инструмента, и этот текст уходит в каждый следующий запрос.
Второй получает выжимку, а полный текст уходит в состояние, откуда его достаёт
второй инструмент по ключевому слову.

Печатается число вызовов модели и сумма входных токенов за весь запуск.
"""

from langchain.agents import AgentState, create_agent
from langchain.messages import ToolMessage
from langchain.tools import ToolRuntime, tool
from langgraph.types import Command

from course_model import build_model

REPORT = (
    "Отчёт по складу Казань за неделю.\n"
    "Приёмка: принято 412 паллет, из них 38 по электронной очереди с опозданием.\n"
    "Отгрузка: отгружено 389 паллет, средняя загрузка машины 84 процента.\n"
    "Брак: выявлено 17 мест, из них бой 9, подмочка 5, пересорт 3.\n"
    "Брак по перевозчикам: 11 мест у перевозчика М-1, 6 мест у перевозчика М-4.\n"
    "Возвраты: оформлено 23 возврата, средний срок обработки 1,8 дня.\n"
    "Инвентаризация: пересчитаны зоны A и B, расхождение 0,4 процента.\n"
    "Персонал: средняя численность смены 14 человек, две смены неполные.\n"
    "Техника: один погрузчик в ремонте с вторника, замена не выделена.\n"
    "Температурная зона: два срыва режима, оба короче пятнадцати минут.\n"
    "Документы: в учётную систему с опозданием попали 6 актов приёмки.\n"
    "Планы: на следующую неделю заявлено 430 паллет приёмки и 400 отгрузки.\n"
)

QUESTION = "Откройте отчёт по складу Казань и скажите, что в нём написано про брак."

PROMPT = (
    "Вы помощник руководителя склада. Отвечайте по-русски, одной-двумя фразами. "
    "Данные берите инструментами, по памяти не отвечайте."
)


class DeskState(AgentState):
    """Состояние агента: к messages добавлено место для загруженных документов."""

    documents: dict[str, str]


def open_report(name: str) -> str:
    """Возвращает полный текст отчёта по складу указанного города."""
    if name != "Казань":
        return f"Отчёта по складу {name} нет."

    return REPORT


@tool
def load_report(name: str, runtime: ToolRuntime[None, DeskState]) -> Command:
    """Загружает отчёт по складу города в рабочую память и отдаёт его состав."""
    if name != "Казань":
        return Command(
            update={
                "messages": [
                    ToolMessage(
                        content=f"Отчёта по складу {name} нет.",
                        tool_call_id=runtime.tool_call_id,
                    )
                ]
            }
        )

    lines = REPORT.splitlines()
    topics = ", ".join(line.split(":")[0] for line in lines[1:])
    digest = (
        f"Отчёт по складу {name} загружен: {len(lines)} строк, "
        f"{len(REPORT)} знаков. Разделы: {topics}. Чтобы прочитать нужный "
        f"раздел, вызовите find_in_report с ключевым словом."
    )

    return Command(
        update={
            "documents": {**runtime.state.get("documents", {}), name: REPORT},
            "messages": [
                ToolMessage(content=digest, tool_call_id=runtime.tool_call_id)
            ],
        }
    )


@tool
def find_in_report(
    name: str, keyword: str, runtime: ToolRuntime[None, DeskState]
) -> str:
    """Возвращает строки загруженного отчёта, в которых встречается слово."""
    text = runtime.state.get("documents", {}).get(name)

    if text is None:
        return f"Отчёт по складу {name} ещё не загружен."

    found = [line for line in text.splitlines() if keyword.lower() in line.lower()]

    if not found:
        return f"Слово '{keyword}' в отчёте по складу {name} не встречается."

    return "\n".join(found)


def spent(messages) -> int:
    """Складывает входные токены по всем ответам модели в истории запуска."""
    total = 0

    for message in messages:
        # Расход возвращает не каждый провайдер, поэтому пустой словарь как запасной путь.
        usage = getattr(message, "usage_metadata", None) or {}
        total += usage.get("input_tokens", 0)

    return total


def model_calls(messages) -> int:
    """Считает ответы модели в истории запуска."""
    return sum(1 for message in messages if message.type == "ai")


direct = create_agent(
    model=build_model(temperature=0, max_tokens=384),
    tools=[open_report],
    system_prompt=PROMPT,
)

filtered = create_agent(
    model=build_model(temperature=0, max_tokens=384),
    tools=[load_report, find_in_report],
    system_prompt=PROMPT,
    state_schema=DeskState,
)

print(f"{'агент':<22}  {'вызовов':>7}  {'вход всего':>10}")

for label, agent in (("отчёт целиком", direct), ("выжимка и состояние", filtered)):
    result = agent.invoke({"messages": [{"role": "user", "content": QUESTION}]})
    messages = result["messages"]
    answer = messages[-1].text.replace("\n", " ")
    print(f"{label:<22}  {model_calls(messages):>7}  {spent(messages):>10}")
    print(f"{'':<22}  ОТВЕТ: {answer}")

print()
print("В СОСТОЯНИИ ВТОРОГО АГЕНТА ОСТАЛОСЬ:")
for name, text in result.get("documents", {}).items():
    print(f"  {name}: {len(text)} знаков")

# Вывод:
# агент                   вызовов  вход всего
# отчёт целиком                 2        1052
#                         ОТВЕТ:  В отчёте по складу Казань указано, что выявлено 17 мест брака, из них: бой — 9, подмочка — 5, пересорт — 3. Также отмечено, что по перевозчикам: 11 мест брака у перевозчика М-1 и 6 мест у перевозчика М-4.
# выжимка и состояние           3        1727
#                         ОТВЕТ: В отчёте по складу Казань указано, что выявлено 17 мест брака: из них 9 — бой, 5 — подмочка, 3 — пересорт. Также отмечено, что 11 мест брака приходится на перевозчика М-1 и 6 мест — на перевозчика М-4.
#
# В СОСТОЯНИИ ВТОРОГО АГЕНТА ОСТАЛОСЬ:
#   Казань: 785 знаков

"""Пример 5 урока 14: встроенное редактирование контекста и его временность.

Три вопроса в одном thread, каждый с тяжёлым результатом инструмента. Middleware
ContextEditingMiddleware заменяет старые результаты заглушкой, когда история
превышает порог.

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

from itertools import count
from typing import Callable

from langchain.agents import create_agent
from langchain.agents.middleware import (
    ClearToolUsesEdit,
    ContextEditingMiddleware,
    ModelRequest,
    ModelResponse,
    wrap_model_call,
)
from langchain_core.messages.utils import count_tokens_approximately
from langgraph.checkpoint.memory import InMemorySaver

from course_model import build_model

REPORTS = {
    "Казань": "бой 9, подмочка 5, пересорт 3",
    "Омск": "бой 2, подмочка 1, пересорт 14",
    "Пермь": "бой 4, подмочка 12, пересорт 2",
}

FILLER = (
    "Приёмка идёт по электронной очереди, талон берут на въезде. Паллеты "
    "принимают по счёту мест и по весу брутто. Упаковку вскрывают при видимом "
    "повреждении обрешётки и снимают вскрытие на камеру. Температурный груз "
    "замеряют дважды. Документы возвращают водителю с отметкой о приёмке, копию "
    "подшивают в журнал склада и дублируют в учётную систему до 20:00 того же "
    "дня. Расхождение больше пяти мест останавливает приёмку до прихода "
    "старшего смены."
)


def warehouse_report(city: str) -> str:
    """Возвращает полный недельный отчёт по складу города."""
    defects = REPORTS.get(city)

    if defects is None:
        return f"Отчёта по складу {city} нет."

    return f"Отчёт по складу {city}. Регламент: {FILLER} Брак за неделю: {defects}."


step = count(1)


@wrap_model_call
def observer(
    request: ModelRequest,
    handler: Callable[[ModelRequest], ModelResponse],
) -> ModelResponse:
    """Печатает то, что дошло до модели после работы редактора контекста."""
    tool_messages = [m for m in request.messages if m.type == "tool"]
    cleared = sum(
        1
        for m in tool_messages
        if m.response_metadata.get("context_editing", {}).get("cleared")
    )
    shapes = ", ".join(f"{len(m.text):>4}" for m in tool_messages)

    print(
        f"{next(step):>5}  {len(request.messages):>9}  {len(tool_messages):>4}  "
        f"{cleared:>8}  {count_tokens_approximately(request.messages):>7}  {shapes}"
    )

    return handler(request)


agent = create_agent(
    model=build_model(temperature=0, max_tokens=384),
    tools=[warehouse_report],
    system_prompt=(
        "Вы помощник руководителя склада. Отвечайте по-русски, одной фразой. "
        "Цифры берите инструментом, по памяти не отвечайте."
    ),
    middleware=[
        ContextEditingMiddleware(
            edits=[ClearToolUsesEdit(trigger=300, keep=1)],
        ),
        observer,
    ],
    checkpointer=InMemorySaver(),
)

CONFIG = {"configurable": {"thread_id": "edit-1"}}

QUESTIONS = [
    "Сколько брака за неделю на складе Казань?",
    "А на складе Омск?",
    "А на складе Пермь?",
]

print(
    f"{'вызов':>5}  {'сообщений':>9}  {'tool':>4}  {'очищено':>8}  "
    f"{'оценка':>7}  длины результатов инструмента"
)

for question in QUESTIONS:
    result = agent.invoke(
        {"messages": [{"role": "user", "content": question}]},
        CONFIG,
    )

print()
print("ПОСЛЕДНИЙ ОТВЕТ:", result["messages"][-1].text.replace("\n", " "))

print()
print("В THREAD ПОСЛЕ ТРЁХ ВОПРОСОВ:")
for message in result["messages"]:
    if message.type == "tool":
        print(f"  tool  {len(message.text):>4} знаков  {message.text[:38]}")

# Вывод:
# вызов  сообщений  tool   очищено   оценка  длины результатов инструмента
#     1          1     0         0       15
# (вызовы 2 и 3 опущены: очистка ещё не сработала, оценка выросла до 232)
#     4          7     2         1      285     9,  523
#     5          9     2         1      320     9,  523
#     6         11     3         2      373     9,    9,  524
#
# ПОСЛЕДНИЙ ОТВЕТ:  За неделю на складе Пермь зафиксирован брак: бой — 4, подмочка — 12, пересорт — 2.
#
# В THREAD ПОСЛЕ ТРЁХ ВОПРОСОВ:
#   tool   524 знаков  Отчёт по складу Казань. Регламент: При
#   tool   523 знаков  Отчёт по складу Омск. Регламент: Приём
#   tool   524 знаков  Отчёт по складу Пермь. Регламент: Приё

"""Пример 6 урока 14: голый агент против собранного на одной и той же задаче.

Три вопроса подряд, у каждого агента свой thread, один набор инструментов, один
провайдер. Разница только в харнессе: у второго агента промпт собирается под роль,
в запрос уходят только разрешённые роли инструменты, а старые результаты
инструментов заменяются заглушкой по порогу.

Печатается число вызовов модели и сумма входных токенов за весь разговор.
"""

from dataclasses import dataclass
from typing import Callable

from langchain.agents import create_agent
from langchain.agents.middleware import (
    ClearToolUsesEdit,
    ContextEditingMiddleware,
    ModelRequest,
    ModelResponse,
    dynamic_prompt,
    wrap_model_call,
)
from langgraph.checkpoint.memory import InMemorySaver

from course_model import build_model

REPORTS = {
    "Казань": "бой 9, подмочка 5, пересорт 3",
    "Омск": "бой 2, подмочка 1, пересорт 14",
    "Пермь": "бой 4, подмочка 12, пересорт 2",
}

FILLER = (
    "Приёмка идёт по электронной очереди, талон берут на въезде. Паллеты "
    "принимают по счёту мест и по весу брутто. Упаковку вскрывают при видимом "
    "повреждении обрешётки и снимают вскрытие на камеру. Температурный груз "
    "замеряют дважды. Документы возвращают водителю с отметкой о приёмке, копию "
    "подшивают в журнал склада и дублируют в учётную систему до 20:00 того же "
    "дня. Расхождение больше пяти мест останавливает приёмку до прихода "
    "старшего смены."
)

LONG_PROMPT = (
    "Вы помощник логистического отдела. Отвечайте по-русски, одной фразой. "
    "Данные берите инструментами, по памяти не отвечайте. "
    "Инструмент read_report отдаёт недельный отчёт по складу города вместе с "
    "регламентом приёмки. Инструмент read_stock отдаёт остаток по артикулу. "
    "Инструмент read_routes отдаёт маршруты через город. Инструмент "
    "write_order меняет статус заказа, его разрешено звать только оператору "
    "склада, и только после того, как оператор назвал номер заказа целиком. "
    "Инструмент cancel_order отменяет заказ, его разрешено звать только "
    "оператору склада и только после подтверждения суммы возврата. "
    "Наблюдателю запрещено менять данные: если наблюдатель просит отменить "
    "заказ или сменить статус, откажите и предложите обратиться к оператору."
)


@dataclass
class Access:
    """Конфигурация запуска: под какой ролью работает собеседник."""

    role: str


def read_report(city: str) -> str:
    """Возвращает полный недельный отчёт по складу города."""
    defects = REPORTS.get(city)

    if defects is None:
        return f"Отчёта по складу {city} нет."

    return f"Отчёт по складу {city}. Регламент: {FILLER} Брак за неделю: {defects}."


def read_stock(sku: str) -> str:
    """Возвращает остаток товара по артикулу на всех складах."""
    return f"Артикул {sku}: Казань 7 шт., Омск 0 шт."


def read_routes(city: str) -> str:
    """Возвращает список маршрутов, которые проходят через город."""
    return f"Через город {city} идут маршруты M-1 и M-4."


def write_order(order_id: str, status: str) -> str:
    """Меняет статус заказа на указанный."""
    return f"Заказ {order_id} переведён в статус {status}."


def cancel_order(order_id: str) -> str:
    """Отменяет заказ по номеру."""
    return f"Заказ {order_id} отменён."


ALL_TOOLS = [read_report, read_stock, read_routes, write_order, cancel_order]


@dynamic_prompt
def prompt_for_role(request: ModelRequest) -> str:
    """Собирает короткую инструкцию под роль собеседника."""
    base = (
        "Вы помощник логистического отдела. Отвечайте по-русски, одной фразой. "
        "Данные берите инструментами, по памяти не отвечайте."
    )

    if request.runtime.context.role == "оператор":
        return base + " Менять статус и отменять заказ разрешено."

    return base + " Менять данные вам нельзя, доступно только чтение."


@wrap_model_call
def tools_by_role(
    request: ModelRequest,
    handler: Callable[[ModelRequest], ModelResponse],
) -> ModelResponse:
    """Оставляет в запросе инструменты, разрешённые роли собеседника."""
    if request.runtime.context.role != "оператор":
        request = request.override(
            tools=[t for t in request.tools if t.name.startswith("read_")]
        )

    return handler(request)


def spent(messages) -> int:
    """Складывает входные токены по всем ответам модели в истории thread."""
    total = 0

    for message in messages:
        # Расход возвращает не каждый провайдер, поэтому пустой словарь как запасной путь.
        usage = getattr(message, "usage_metadata", None) or {}
        total += usage.get("input_tokens", 0)

    return total


bare = create_agent(
    model=build_model(temperature=0, max_tokens=384),
    tools=ALL_TOOLS,
    system_prompt=LONG_PROMPT,
    context_schema=Access,
    checkpointer=InMemorySaver(),
)

tuned = create_agent(
    model=build_model(temperature=0, max_tokens=384),
    tools=ALL_TOOLS,
    context_schema=Access,
    middleware=[
        prompt_for_role,
        tools_by_role,
        ContextEditingMiddleware(
            edits=[ClearToolUsesEdit(trigger=300, keep=1)],
        ),
    ],
    checkpointer=InMemorySaver(),
)

QUESTIONS = [
    "Сколько брака за неделю на складе Казань?",
    "А на складе Омск?",
    "А на складе Пермь?",
]

print(f"{'агент':<10}  {'вызовов':>7}  {'вход всего':>10}  последний ответ")

for label, agent in (("голый", bare), ("собранный", tuned)):
    config = {"configurable": {"thread_id": f"together-{label}"}}

    for question in QUESTIONS:
        result = agent.invoke(
            {"messages": [{"role": "user", "content": question}]},
            config,
            context=Access(role="наблюдатель"),
        )

    messages = result["messages"]
    calls = sum(1 for message in messages if message.type == "ai")
    answer = messages[-1].text.replace("\n", " ")
    print(f"{label:<10}  {calls:>7}  {spent(messages):>10}  {answer}")

# Вывод:
# агент       вызовов  вход всего  последний ответ
# голый             6        6891  Брак за неделю на складе Пермь: бой 4, подмочка 12, пересорт 2.
# собранный         6        4471  Брак за неделю на складе Пермь: бой 4, подмочка 12, пересорт 2.

@wrap_model_call
def wrong(request, handler):
    request.tools = [t for t in request.tools if t.name.startswith("read_")]
    return handler(request)

@wrap_model_call
def right(request, handler):
    tools = [t for t in request.tools if t.name.startswith("read_")]
    return handler(request.override(tools=tools))

ContextEditingMiddleware(edits=[ClearToolUsesEdit(trigger=300, keep=10)])