LangChain для бывалых — память и агенты. часть 2

Добро пожаловать во вторую часть статьи о фреймворке LangChain. В предыдущей части мы научились добавлять память в диалоги с LLM, а также создали простого агента, который подключался к поисковой системе Google. Некоторые запросы вызвали беспокойство модераторов на Хабре и они попросили добавить немного цензуры отредактировать статью. Но мы знаем, что рукописи не горят, поэтому с оригинальным фрагментом можно познакомиться здесь

Как уже отмечено в предыдущей части, агенты являются одним из самых мощных инструментов фреймворка LangChain. Они позволяют подключаться к внешним ресурсам, что значительно расширяет возможности языковой модели.

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

Cоздаем базу данных

Чтобы подключить агента к базе данных, необходимо сначала создать саму базу данных. Представим, что наш заказчик бота — небольшой магазин, специализирующийся на продаже футболок с принтами. Допустим, его база имеет следующую структуру:

Таблица tshirts:

• Поля:

  • id: Уникальный идентификатор футболки

  • type: Тип футболки (например, «Type 1», «Type 2» и т.д.)

  • color: Цвет футболки

  • size: Размер футболки

  • quantity: Количество футболок в наличии

  • print_id: Внешний ключ на таблицу «prints»

  • price: Цена футболки с учетом атрибутов

Таблица prints:

• Поля:

  • id: Уникальный идентификатор принта

  • name: Название принта

  • image_url: URL изображения принта

Таблица clients:

• Поля:

  • id: Уникальный идентификатор клиента

  • name: Имя клиента

  • address: Адрес клиента

  • contact_info: Контактная информация клиента

Таблица orders :

• Поля:

  • id: Уникальный идентификатор заказа

  • order_date: Дата заказа

  • client_id: Внешний ключ на таблицу «clients»

  • price: Общая стоимость заказа

Структура не идеальна, но лучше пока ее не усложнять.

Теперь можно воплотить идею в коде. Я буду использовать SQLAlchemy. С помощью SQLAlchemy можно создавать объекты, которые представляют таблицы в базе данных, и работать с ними через python. Она упрощает взаимодействие с базами данных и помогает избежать ошибок, связанных с написанием SQL‑запросов.

Возможно, у вас возникнет вопрос: зачем заморачиваться, если можно просто скачать готовую базу данных и протестировать ее. Действительно, это более простой способ. Но за простоту мы заплатим свободой и гибкостью. К тому же инструменты LangChain работают с алхимией, поэтому будет полезно попробовать этот фреймворк в работе.

Код для создания базы:

from sqlalchemy import Column, Integer, String, ForeignKey, Date, Numeric, Table   from sqlalchemy.orm import relationship   from sqlalchemy.orm import declarative_base      from sqlalchemy import create_engine      Base = declarative_base()         class TShirt(Base):       __tablename__ = 'tshirts'       id = Column(Integer, primary_key=True)       type = Column(String)  # Тип футболки (например, "Type 1", "Type 2" и т.д.)       color = Column(String)        size = Column(String)       quantity = Column(Integer)  # Количество футболок в наличии       print_id = Column(Integer, ForeignKey('prints.id'))  # Внешний ключ на таблицу "Prints"       print = relationship("Print")  # Отношение между футболками и принтами       price = Column(Numeric(10, 2))  # Цена футболки с учетом атрибутов          def calculate_tshirt_price(self):           base_price = 500.0           type_multiplier = 1.2           color_multiplier = 1.3           size_multiplier = 1           if self.size in ['XL', 'XXL']:               size_multiplier = 1.5           price = base_price * type_multiplier * color_multiplier * size_multiplier           self.price = round(price, 2)         class Print(Base):       __tablename__ = 'prints'       id = Column(Integer, primary_key=True)       name = Column(String)  # Название принта       image_url = Column(String)  # URL изображения принта         class Client(Base):       __tablename__ = 'clients'       id = Column(Integer, primary_key=True)       name = Column(String)  # Имя клиента       address = Column(String)  # Адрес клиента       contact_info = Column(String)  # Контактная информация клиента       orders = relationship("Order", back_populates="client")  # Отношение         order_tshirt_table = Table('order_tshirt', Base.metadata,                              Column('order_id', Integer, ForeignKey('orders.id'), primary_key=True),                              Column('tshirt_id', Integer, ForeignKey('tshirts.id'), primary_key=True)                              )         class Order(Base):       __tablename__ = 'orders'       id = Column(Integer, primary_key=True)       order_date = Column(Date)  # Дата заказа       client_id = Column(Integer, ForeignKey('clients.id'))  # Внешний ключ на таблицу "Clients"       client = relationship("Client", back_populates="orders")  # Отношение между заказами и клиентами       tshirts = relationship("TShirt", secondary=order_tshirt_table)       quantity = Column(Integer)       price = Column(Numeric(10, 2))  # Общая стоимость заказа          def calculate_order_price(self):           total_price = 0.0           self.quantity = len(self.tshirts)           for tshirt in self.tshirts:               total_price += tshirt.price           self.price = round(total_price, 2) * self.quantity         def create_tables(engine):       Base.metadata.create_all(engine)       print('Tables created successfully!')         if __name__ == '__main__':       # Создаем экземпляр движка базы данных       engine = create_engine('sqlite:///tshirt_store.db')       print('Сreate engine')          # Создаем таблицы       create_tables(engine)

Заполняем базу

Хорошо, давайте заполним базу данных футболками. Вот пример кода, который поможет нам насемплировать любое количество футболок с использованием SQLAlchemy:

import random   from random import randint      from create_database import TShirt, Print   from sqlalchemy import create_engine   from sqlalchemy.orm import sessionmaker         def random_color():       colors = ['Red', 'Blue', 'Green', 'Yellow', 'Black', 'White']       return random.choice(colors)         def random_size():       sizes = ['S', 'M', 'L', 'XL', 'XXL']       return random.choice(sizes)         def create_tshirt_samples(session, num_samples):       tshirt_samples = []       print_samples = []       for _ in range(num_samples):           tshirt = TShirt(               type=f'Type {randint(1, 5)}',               color=random_color(),               size=random_size(),               quantity=randint(1, num_samples // 5),               print_id=0,  # Заглушка для внешнего ключа, будет обновлено ниже               price=0.0           )           tshirt.calculate_tshirt_price()  # Вычисление цены футболки           tshirt_samples.append(tshirt)              print = Print(               name=f'Print_{randint(1, num_samples * 2)}',               image_url=f'https://example.com/print_{randint(1, 10)}.jpg'           )           print_samples.append(print)          session.add_all(print_samples)       session.flush()  # Получение сгенерированных ID принтов          for i, tshirt in enumerate(tshirt_samples):           tshirt.print_id = print_samples[i].id          session.add_all(tshirt_samples)       session.commit()         if __name__ == '__main__':       # Создаем экземпляр движка базы данных       engine = create_engine('sqlite:///tshirt_store.db')       print('Сreate engine')          Session = sessionmaker(bind=engine)       session = Session()          create_tshirt_samples(session, 50)       print('Shirts successfully added into DB')          session.close()

Если ничего не упадет, то у вас должна появиться табличка с футболками:

Дальше нужно сформировать заказы. Попробуем сделать это согласно определенной функциональной зависимости: например периодической с наличием небольшого тренда.

import random from datetime import datetime, timedelta from random import randint from create_database import Order, Client, TShirt from sqlalchemy.orm import sessionmaker from sqlalchemy import create_engine, func  from math import sin, pi   def create_random_client(n):     # Создание случайного клиента с использованием случайно сгенерированных имени, адреса и контактной информации     name = f'Client_{randint(1, n)}'     address = f'Address_{randint(1, n)}'     contact_info = f'Contact_{randint(1, n)}'     return Client(name=name, address=address, contact_info=contact_info)   def get_random_client(session, r_id=1000, client_proba=0.5):     # Получение случайного клиента из базы данных     client = session.query(Client).order_by(func.random()).first()     if client is None or client_proba < 0.5:         # Создание случайного клиента, если база данных клиентов пуста или вероятность меньше 0.5         client = create_random_client(r_id)         session.add(client)         session.commit()     return client   def get_random_tshirts(session, num_tshirts):     # Получение случайного количества футболок из базы данных     tshirts = session.query(TShirt).order_by(func.random()).limit(num_tshirts).all()     return tshirts   def create_orders(         session,         start_date,         end_date,         order_function, ):     current_date = start_date     while current_date <= end_date:         # Получение количества заказов на текущую дату с использованием функции order_function         num_orders = order_function(current_date)         for _ in range(num_orders):             # Получение случайного клиента             client = get_random_client(session, client_proba=random.random())             # Создание экземпляра класса Order с указанной датой и клиентом             order = Order(order_date=current_date, client=client)             order.client = client             # Получение случайных футболок             tshirts = get_random_tshirts(session, randint(1, num_orders))             order.tshirts = tshirts              # Вычисление общей цены заказа             order.calculate_order_price()              session.add(order)          current_date += timedelta(days=1)      session.commit()   def order_function(date):       day_of_year = date.timetuple().tm_yday       amplitude = 10  # Амплитуда количества заказов       period = 7  # Период синусоиды в днях          angle = 2 * pi * (day_of_year % period) / period       num_orders = amplitude * (sin(angle) + 1)          return int(num_orders * day_of_year / 365)   if __name__ == '__main__':     # Создаем экземпляр движка базы данных     engine = create_engine('sqlite:///tshirt_store.db')     print('Сreate engine')      Session = sessionmaker(bind=engine)     session = Session()      start_date = datetime(2023, 5, 1)     end_date = datetime(2023, 6, 30)      create_orders(session, start_date, end_date, order_function)    print('Orders successfully added into DB')    session.close()

Должна появится похожая на эту таблица:

Теперь у нас есть фундамент для работы агента.

Создаем агента

Для удобства дальнейшего воспроизведения, код я буду писать в сниппетах, а результаты выводить в виде скринов ячеек jupyter notebook. Напомню три источника, три составные части агента:

1. Языковая модель

2. Инструмент

3. Тело агента

import os  from langchain.sql_database import SQLDatabase from langchain.chains import SQLDatabaseChain  from langchain.chat_models import ChatOpenAI from langchain.agents import Tool from langchain.agents import initialize_agent  # Инициализируем языковую модель llm = ChatOpenAI( openai_api_key=open_ai_api_key, model_name='gpt-3.5-turbo', temperature=0 ) # Создаем движок для работы с БД engine = create_engine('sqlite:///tshirt_store.db')  # Регистрируем движок в langchain db = SQLDatabase(engine)  # Создаем цепочку для работы с базой sql_chain = SQLDatabaseChain( llm=llm, database=db, verbose=True )  # Создаем на основе цепочки инструмент db_tool = Tool(     name='tshirt store db',     func=sql_chain.run,     description='Use for extaract data from database'  )  # Ну и наконец создаем самого агента(пока без использования памяти) db_agent = initialize_agent(     agent='zero-shot-react-description',      tools=[db_tool],      llm=llm,     verbose=True,     max_iterations=5, )

Пора начинать задавать вопросы нашему эксперту:

Неплохо, попробуем задать вопрос посложнее:

Видно, что цепочка рассуждений верна. Агент даже правильно понял, как связать различные таблицы, чтобы получить верный ответ.

На самом деле, запрос к нашей базе можно сделать и без агента, а напрямую через цепочку sql_chain:

Но вот уже со вторым запросом цепочка не справилась:

Посмотрим, как выглядит шаблон промпта нашего агента:

А так выглядит промпт для цепочки:

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

Кажется, что агент работает идеально, но сломать его оказалось довольно просто:

Что же делать? Можно попытаться модифицировать промпт. Однако, есть и другой путь — воспользоваться более мощным инструментом. В LangChain есть SQLDatabaseToolkit, под капотом этого швейцарского ножа сразу 4 инструмента:

• QuerySQLDataBaseTool

• InfoSQLDatabaseTool

• ListSQLDatabaseTool

• QueryCheckerTool

Посмотрим, смогут ли они нам помочь. Но сначала нужно инициализировать нового агента:

from langchain.agents import create_sql_agent from langchain.agents.agent_toolkits import SQLDatabaseToolkit  toolkit = SQLDatabaseToolkit(db=db, llm=llm)  db_agent_toolkit = create_sql_agent(     llm=llm,     toolkit=toolkit,     verbose=True )

Задаем тот же вопрос:

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

Мы научились извлекать информацию из базы данных. Почему бы нам не научить агента строить модели на основе этих данных? Изначально я планировал написать свой собственный инструмент для этой задачи, но оказалось, что уже существует готовая реализация в виде PythonREPLTool. Ну что ж, давайте добавим и этот инструмент в наш арсенал.

from langchain.agents.agent_toolkits import create_python_agent from langchain.tools.python.tool import PythonREPLTool  agent_model = create_python_agent(     llm=llm,     tool=PythonREPLTool(),     verbose=True )

Попробуем спрогнозировать заказы на футболки:

Ну что, кожаные, впечатлены? Время эльфов людей прошло, настало время агентов. На это й оптимистичной ноте с агентами можно пока закончить. В заключительной части мы попробуем написать собственные инструменты для работы с агентами. Спасибо за внимание!

Пишу про AI и NLP в телеграм.