pg_auto_embeddings — считаем эмбеддинги для текста прямо в Postgres, без экстеншенов10.01.2025 14:31

b2fbf4b2c666934e97bff421807ffa58.png

У вас есть PostgreSQL база, где хранится множество текстовых данных. Вы хотите использовать векторные представления (embeddings), к примеру, от OpenAI, чтобы построить систему рекомендаций, улучшенный поиск или реализовать RAG для работы с LLM. Но при этом ставить расширения (extensions) не хочется, а может, и вовсе нельзя — например, в облачных Managed PostgreSQL зачастую нет нужных прав.

pg_auto_embeddings — это лёгкое open-source решение, которое позволяет вычислять embeddings через модели OpenAI непосредственно в PostgreSQL без установки сторонних расширений. Оно использует механизм Foreign Data Wrappers (FDW) «под капотом» для того, чтобы делать запросы в OpenAI API, и работает синхронно и атомарно. В этой статье разберём, чем pg_auto_embeddings может помочь, как его установить (спойлер: очень легко) и в чём ключевые особенности проекта.

Что такое pg_auto_embeddings и какую задачу решает?

pg_auto_embeddings — это open-source-проект под лицензией MIT, который решает ключевую проблему: «Как вычислять текстовые векторные представления (embeddings) прямо из PostgreSQL без лишних телодвижений и без специальных расширений?»

Основные идеи:

1. Вызов через SQL: вы пишете простую функцию `pgae_embedding ('какой-то текст')` и этого достаточно. Функцию можно использовать в триггерах, и таким образом автоматически сохранять эмбеддинги для текстовых колонок.

2. Гибкие настройки: вы можете использовать публичные модели (например, OpenAI) или поднять on-premise-прокси, если нужно больше контроля — например, добавлять свои лимиты, закрытый доступ и т.п.

Благодаря этому, pg_auto_embeddings отлично подходит, когда нужно быстро «прокинуть» вычисление эмбеддингов в существующую БД и не заморачиваться с установкой сторонних бинарных расширений. Удобно для RAG-систем и прочих задач, где embeddings — это базовый функционал.

Ключевые возможности

  • Без расширений: Никакого дополнительного ПО в PostgreSQL ставить не нужно — только выполнить SQL-файл (да-да, в третий раз пишу, но это правда важно :).

  • Две варианта развёртывания:

    • Упрощённая установка: выполните один SQL-скрипт и готово.

    • On-Premise (через Docker): поднимите свой прокси-сервер, который обрабатывает запросы к API для embedding. Тоже поднимается одной `docker run` командой.

  • Поддержка OpenAI Embeddings: на данный момент из коробки работают модели OpenAI (text-embedding-3-small/large и некоторые другие).

  • Автоматический апдейт вектора при вставке или обновлении текстовых данных: можно «повесить» auto-embedding на столбец таблицы и не тратить время на написание триггера.

  • Удаление «за собой»: при необходимости вы можете полностью удалить pg_auto_embeddings и все его объекты одной функцией.

Шаг 1. Установка

  1. Возьмите файл simple/pgae_simple_install.sql и выполните его в вашей базе данных.

  2. Инициализируйте pg_auto_embeddings: 

    CALL pgae_init('openai-text-embedding-3-small', 'ВАШ_OPENAI_API_КЛЮЧ');

    Вот и всё :)

Шаг 2. Использование

Чтобы получить вектор (массив `double precision[]`): 

SELECT pgae_embedding('ваш текст');

Если у вас установлен pgvector и хочется формат vector, то: 

SELECT pgae_embedding_vec('some text');

Автоматический подсчет и запись эмбеддинга. Допустим, у вас есть таблица posts со столбцом title, и вы хотите автоматом сохранять эмбеддинги для заголовков в title_embedding: 

SELECT pgae_create_auto_embedding(
  'public', 'posts', 'title', 'title_embedding'
);

Вуаля.

Если вы вдруг решили снести pg_auto_embeddings, полностью убрав все его функции и объекты из вашей БД, просто выполните: 

SELECT pgae_self_destroy();

[Опционально] On-premise вариант (через Docker)

Если у вас есть ограничения на внешние запросы или вы хотите поднять собственный прокси-сервер для пущей безопасности, то это легко сделать.

Для этого нужно поднять 2 сервиса: Postgres базу, которая выступает в роли прокси между FDW и Node.js, и сам Node.js, который будет выполнять запросы к API моделей.

pg_auto_embeddings предоставляет docker-образ, в котором оба этих компонента уже развернуты и настроены для использования.

Пример docker-compose.yml: 

services:
    server:
        image: elkornacio/pg_auto_embeddings:latest
        environment:
            - PG_HOST=localhost # Хост, где расположена проксирующий Postgres
            - PG_PORT=5432 # Порт от него
            - PG_USERNAME=root_user # root-юзер от него
            - PG_PASSWORD=root_pass # пароль
            - DATABASE_SYNC=true
            - SERVER_HOST=localhost # Хост, который слушает Node.js-прокси
            - SERVER_PORT=3000 # Порт, который он слушает
            - SELF_URL=http://localhost:3000 # Главное! URL, по которому проксирующий Postgres обращается к Node.js прокси
        ports:
            # Порт проксирующего Postgres должен быть открыт - к нему будет подключаться ваша managed БД
            - 5432:5432

Далее, вместо pgae_init потребуется использовать pgae_init_onprem: 

CALL pgae_init_onprem(
  'your.host.com', '5432', -- хост и порт вашего проксирующего Postgres
  'openai-text-embedding-3-small', 'sk-...' -- тип модели и ключ API
);

Использование такое же, как и в «простом» варианте: 

SELECT pgae_embedding('ваш текст');

Как это работает внутри?

pg_auto_embeddings под капотом использует трюк на базе Foreign Data Wrappers (FDW):

  1. При установке создаётся «прокси-таблица» embeddings_* в локальной базе.

  2. Когда вы вызываете `SELECT pgae_embedding ('some text')`, под капотом происходит UPDATE в этой «прокси-таблице» с передачей текста.

  3. FDW перенаправляет запрос на удалённую таблицу (в Docker-прокси или публичный сервер).

  4. На удалённом сервере срабатывает триггер, который вызывает внутреннюю функцию pgae_embedding_internal().

  5. Эта функция делает HTTP-запрос к Node.js-прокси.

  6. Node.js-сервер обращается к OpenAI API (или к чему-то ещё, если у вас другой провайдер) и получает вектор.

  7. Вектор возвращается обратно в удалённую БД, дальше в локальную БД и в конечном итоге — в ваш SELECT-запрос.

Заключение

pg_auto_embeddings отлично подходит, когда нужно быстро, а главное — без сложных установок, получить возможность вычислять векторные представления. Он идеально подойдёт для тех, кто хочет «подружить» свою БД с LLM или сделать продвинутый полнотекстовый поиск, не выходя за рамки SQL.

Проект активен и открыт для предложений и PR-ов. А ещё для звездочек:) Если у вас возникли вопросы, смело пишите в Issues на GitHub.

Спасибо за внимание и удачных экспериментов :) Если вдруг в тексте нашли ошибки — напишите мне в ЛС, мигом поправлю.

© Habrahabr.ru