Добро пожаловать во вторую часть статьи о фреймворке 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. Напомню три источника, три составные части агента:
Языковая модель
Инструмент
Тело агента
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 в телеграм.