Поделиться через


Использование LangChain с базой данных Azure для PostgreSQL

База данных Azure для PostgreSQL легко интегрируется с ведущими пакетами оркестрации большой языковой модели (LLM), такими как LangChain. Эта интеграция позволяет разработчикам использовать расширенные возможности искусственного интеллекта в своих приложениях. LangChain может упростить управление и использование LLM, моделей встраивания и баз данных, чтобы разработка приложений ИИ стала проще.

В этой статье показано, как использовать интегрированную базу данных векторов в Базе данных Azure для PostgreSQL для хранения документов и управления ими в коллекциях с помощью LangChain. В нем также показано, как создавать индексы и выполнять векторные запросы на поиск, используя алгоритмы ближайшего соседа, такие как косинусное расстояние, Евклидово расстояние (расстояние L2), а также внутренний продукт, для поиска документов, похожих на запросные вектора.

Поддержка векторов

Базу данных Azure для PostgreSQL можно использовать для эффективного хранения и выполнения запросов миллионов векторных представлений в PostgreSQL. Данный сервис поможет вам масштабировать варианты использования ИИ от доказательства концепции до промышленной эксплуатации. Он предлагает следующие преимущества:

  • Предоставляет знакомый интерфейс SQL для запросов векторных представлений и реляционных данных.
  • Ускоряет pgvector благодаря более быстрому и более точному поиску сходства по более чем 100 миллионам векторов с использованием алгоритма индексирования DiskANN.
  • Упрощает операции за счет интеграции реляционных метаданных, векторных встраиваний и данных временных рядов в одну базу данных.
  • Использует возможности надежной экосистемы PostgreSQL и облачной платформы Azure для функций корпоративного уровня, включая репликацию и высокий уровень доступности.

Аутентификация

База данных Azure для PostgreSQL поддерживает проверку подлинности на основе паролей и Microsoft Entra (прежнее название — Azure Active Directory).

Проверка подлинности Microsoft Entra позволяет использовать идентификатор Microsoft Entra для проверки подлинности на сервере PostgreSQL. Идентификатор Microsoft Entra устраняет необходимость управления отдельными именами пользователей и паролями для пользователей базы данных. Он позволяет использовать те же механизмы безопасности, которые используются для других служб Azure.

В этой статье можно использовать любой метод проверки подлинности.

Настройка

Azure Database для PostgreSQL использует поддержку LangChain Postgres с открытым исходным кодом для подключения к Azure Database для PostgreSQL. Сначала скачайте пакет партнера:

%pip install -qU langchain_postgres
%pip install -qU langchain-openai
%pip install -qU azure-identity

Включение pgvector в Базе данных Azure для PostgreSQL

См. раздел "Включение и использование pgvector" в Базе данных Azure для PostgreSQL.

Настройка учетных данных

Необходимо получить сведения о подключении базы данных Azure для PostgreSQL и добавить их в качестве переменных среды.

Задайте для флага USE_ENTRA_AUTH значение True, если вы хотите использовать аутентификацию Microsoft Entra. Если вы используете проверку подлинности Microsoft Entra, необходимо указать только имена узлов и баз данных. Если вы используете проверку подлинности паролей, необходимо также задать имя пользователя и пароль.

import getpass
import os

USE_ENTRA_AUTH = True

# Supply the connection details for the database
os.environ["DBHOST"] = "<server-name>"
os.environ["DBNAME"] = "<database-name>"
os.environ["SSLMODE"] = "require"

if not USE_ENTRA_AUTH:
    # If you're using a username and password, supply them here
    os.environ["DBUSER"] = "<username>"
    os.environ["DBPASSWORD"] = getpass.getpass("Database Password:")

Настройка внедрения Azure OpenAI

os.environ["AZURE_OPENAI_ENDPOINT"] = "<azure-openai-endpoint>"
os.environ["AZURE_OPENAI_API_KEY"] = getpass.getpass("Azure OpenAI API Key:")
AZURE_OPENAI_ENDPOINT = os.environ["AZURE_OPENAI_ENDPOINT"]
AZURE_OPENAI_API_KEY = os.environ["AZURE_OPENAI_API_KEY"]

from langchain_openai import AzureOpenAIEmbeddings

embeddings = AzureOpenAIEmbeddings(
    model="text-embedding-3-small",
    api_key=AZURE_OPENAI_API_KEY,
    azure_endpoint=AZURE_OPENAI_ENDPOINT,
    azure_deployment="text-embedding-3-small",
)

Инициализация

Используйте проверку подлинности Microsoft Entra

В следующих разделах содержатся функции, которые настраивают LangChain для использования проверки подлинности Microsoft Entra. Функция get_token_and_username извлекает токены для службы Azure Database для PostgreSQL, используя DefaultAzureCredential из библиотеки azure.identity. Это гарантирует, что обработчик SQLAlchemy имеет действительный маркер, с помощью которого необходимо создать новые подключения. Он также анализирует маркер, который является веб-маркером JSON (JWT), чтобы извлечь имя пользователя, используемое для подключения к базе данных.

Функция create_postgres_engine создает подсистему SQLAlchemy, которая динамически задает имя пользователя и пароль на основе маркера, полученного из диспетчера маркеров. Этот движок можно передать в параметр connection хранилища векторов PGVector LangChain.

Вход в Azure

Чтобы войти в Azure, убедитесь, что у вас установлен Azure CLI . Выполните следующую команду в терминале:

az login

После входа в систему следующий код извлекает маркер:

import base64
import json
from functools import lru_cache

from azure.identity import DefaultAzureCredential
from sqlalchemy import create_engine, event
from sqlalchemy.engine.url import URL


@lru_cache(maxsize=1)
def get_credential():
    """Memoized function to create the Azure credential, which caches tokens."""
    return DefaultAzureCredential()


def decode_jwt(token):
    """Decode the JWT payload to extract claims."""
    payload = token.split(".")[1]
    padding = "=" * (4 - len(payload) % 4)
    decoded_payload = base64.urlsafe_b64decode(payload + padding)
    return json.loads(decoded_payload)


def get_token_and_username():
    """Fetches a token and returns the username and token."""
    # Fetch a new token and extract the username
    token = get_credential().get_token(
        "https://ossrdbms-aad.database.windows.net/.default"
    )
    claims = decode_jwt(token.token)
    username = claims.get("upn")
    if not username:
        raise ValueError("Could not extract username from token. Have you logged in?")

    return username, token.token


def create_postgres_engine():
    db_url = URL.create(
        drivername="postgresql+psycopg",
        username="",  # This will be replaced dynamically
        password="",  # This will be replaced dynamically
        host=os.environ["DBHOST"],
        port=os.environ.get("DBPORT", 5432),
        database=os.environ["DBNAME"],
    )

    # Create a SQLAlchemy engine
    engine = create_engine(db_url, echo=True)

    # Listen for the connection event to inject dynamic credentials
    @event.listens_for(engine, "do_connect")
    def provide_dynamic_credentials(dialect, conn_rec, cargs, cparams):
        # Fetch the dynamic username and token
        username, token = get_token_and_username()

        # Override the connection parameters
        cparams["user"] = username
        cparams["password"] = token

    return engine

Использование проверки подлинности паролей

Если вы не используете проверку подлинности Microsoft Entra, get_connection_uri предоставляет URI подключения, который извлекает имя пользователя и пароль из переменных окружения.

import urllib.parse


def get_connection_uri():
    # Read URI parameters from the environment
    dbhost = os.environ["DBHOST"]
    dbname = os.environ["DBNAME"]
    dbuser = urllib.parse.quote(os.environ["DBUSER"])
    password = os.environ["DBPASSWORD"]
    sslmode = os.environ["SSLMODE"]

    # Construct the connection URI
    # Use Psycopg 3!
    db_uri = (
        f"postgresql+psycopg://{dbuser}:{password}@{dbhost}/{dbname}?sslmode={sslmode}"
    )
    return db_uri

Создание хранилища векторов

from langchain_core.documents import Document
from langchain_postgres import PGVector
from langchain_postgres.vectorstores import PGVector

collection_name = "my_docs"

# The connection is either a SQLAlchemy engine or a connection URI
connection = create_postgres_engine() if USE_ENTRA_AUTH else get_connection_uri()

vector_store = PGVector(
    embeddings=embeddings,
    collection_name=collection_name,
    connection=connection,
    use_jsonb=True,
)

Управление хранилищем векторов

Добавление элементов в векторное хранилище

Добавление документов по идентификатору перезаписывает все существующие документы, соответствующие идентификатору.

docs = [
    Document(
        page_content="there are cats in the pond",
        metadata={"id": 1, "location": "pond", "topic": "animals"},
    ),
    Document(
        page_content="ducks are also found in the pond",
        metadata={"id": 2, "location": "pond", "topic": "animals"},
    ),
    Document(
        page_content="fresh apples are available at the market",
        metadata={"id": 3, "location": "market", "topic": "food"},
    ),
    Document(
        page_content="the market also sells fresh oranges",
        metadata={"id": 4, "location": "market", "topic": "food"},
    ),
    Document(
        page_content="the new art exhibit is fascinating",
        metadata={"id": 5, "location": "museum", "topic": "art"},
    ),
    Document(
        page_content="a sculpture exhibit is also at the museum",
        metadata={"id": 6, "location": "museum", "topic": "art"},
    ),
    Document(
        page_content="a new coffee shop opened on Main Street",
        metadata={"id": 7, "location": "Main Street", "topic": "food"},
    ),
    Document(
        page_content="the book club meets at the library",
        metadata={"id": 8, "location": "library", "topic": "reading"},
    ),
    Document(
        page_content="the library hosts a weekly story time for kids",
        metadata={"id": 9, "location": "library", "topic": "reading"},
    ),
    Document(
        page_content="a cooking class for beginners is offered at the community center",
        metadata={"id": 10, "location": "community center", "topic": "classes"},
    ),
]

vector_store.add_documents(docs, ids=[doc.metadata["id"] for doc in docs])

Обновление элементов в хранилище векторов

docs = [
    Document(
        page_content="Updated - cooking class for beginners is offered at the community center",
        metadata={"id": 10, "location": "community center", "topic": "classes"},
    )
]
vector_store.add_documents(docs, ids=[doc.metadata["id"] for doc in docs])

Удаление элементов из хранилища векторов

vector_store.delete(ids=["3"])

Запросы к хранилищу векторов

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

Поддержка фильтрации

Хранилище векторов поддерживает набор фильтров, которые можно применить к полям метаданных документов:

Оператор Значение или категория
$eq Равенство (==)
$ne Неравенство (!=)
$lt Меньше (<)
$lte Меньше или равно (<=)
$gt Больше (>)
$gte Больше или равно (>=)
$in Особый регистр (в)
$nin Особая обработка (не включено)
$between Особая обработка регистров (между)
$like Текст (например)
$ilike Текст (без учета регистра)
$and Логические (и)
$or Логический (или)

Прямой запрос

Вы можете выполнить простой поиск сходства следующим образом:

results = vector_store.similarity_search(
    "kitty", k=10, filter={"id": {"$in": [1, 5, 2, 9]}}
)
for doc in results:
    print(f"* {doc.page_content} [{doc.metadata}]")
    * there are cats in the pond [{'id': 1, 'topic': 'animals', 'location': 'pond'}]
    * ducks are also found in the pond [{'id': 2, 'topic': 'animals', 'location': 'pond'}]
    * the new art exhibit is fascinating [{'id': 5, 'topic': 'art', 'location': 'museum'}]
    * the library hosts a weekly story time for kids [{'id': 9, 'topic': 'reading', 'location': 'library'}]

Если вы предоставляете словарь с несколькими полями, но без операторов, верхний уровень интерпретируется как логический AND фильтр:

vector_store.similarity_search(
    "ducks",
    k=10,
    filter={"id": {"$in": [1, 5, 2, 9]}, "location": {"$in": ["pond", "market"]}},
)
[Document(id='2', metadata={'id': 2, 'topic': 'animals', 'location': 'pond'}, page_content='ducks are also found in the pond'),
 Document(id='1', metadata={'id': 1, 'topic': 'animals', 'location': 'pond'}, page_content='there are cats in the pond')]
vector_store.similarity_search(
    "ducks",
    k=10,
    filter={
        "$and": [
            {"id": {"$in": [1, 5, 2, 9]}},
            {"location": {"$in": ["pond", "market"]}},
        ]
    },
)
[Document(id='2', metadata={'id': 2, 'topic': 'animals', 'location': 'pond'}, page_content='ducks are also found in the pond'),
 Document(id='1', metadata={'id': 1, 'topic': 'animals', 'location': 'pond'}, page_content='there are cats in the pond')]

Если вы хотите выполнить поиск сходства и получить соответствующие оценки, можно выполнить следующее:

results = vector_store.similarity_search_with_score(query="cats", k=1)
for doc, score in results:
    print(f"* [SIM={score:3f}] {doc.page_content} [{doc.metadata}]")
* [SIM=0.528338] there are cats in the pond [{'id': 1, 'topic': 'animals', 'location': 'pond'}]

Полный список поисковых запросов, которые можно выполнить в PGVector хранилище векторов, см. в справочнике по API.

Преобразование в ретривер

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

retriever = vector_store.as_retriever(search_type="mmr", search_kwargs={"k": 1})
retriever.invoke("kitty")
[Document(id='1', metadata={'id': 1, 'topic': 'animals', 'location': 'pond'}, page_content='there are cats in the pond')]

Текущие ограничения

  • langchain_postgres работает только с Psycopg 3 (psycopg3). Обновите строки подключения из postgresql+psycopg2://... в postgresql+psycopg://langchain:langchain@....
  • Схема внедренного хранилища и коллекции изменилась, чтобы обеспечить add_documents правильную работу с указанными пользователем идентификаторами.
  • Теперь необходимо передать явный объект подключения.
  • В настоящее время нет механизма , поддерживающего простую миграцию данных при изменениях схемы. Любые изменения схемы в хранилище векторов требуют повторного создания таблиц и повторного добавления документов.