Partilhar via

Usar LangChain com o Banco de Dados do Azure para PostgreSQL

O Banco de Dados do Azure para PostgreSQL integra-se perfeitamente aos principais pacotes de orquestração de modelo de linguagem grande (LLM), como LangChain. Essa integração permite que os desenvolvedores usem recursos avançados de IA em seus aplicativos. LangChain pode simplificar o gerenciamento e o uso de LLMs, modelos de incorporação e bancos de dados para que os aplicativos de IA generativa sejam mais fáceis de desenvolver.

Este artigo mostra como usar o banco de dados vetorial integrado no Banco de Dados do Azure para PostgreSQL para armazenar e gerenciar documentos em coleções com LangChain. O texto também mostra como criar índices e executar consultas de pesquisa vetorial usando algoritmos de vizinhança mais próxima, como a distância do cosseno, a distância L2 (distância euclidiana) e o produto interno para localizar documentos que estão próximos aos vetores de consulta.

Suporte vetorial

Você pode usar o Banco de Dados do Azure para PostgreSQL para armazenar e consultar com eficiência milhões de incorporações vetoriais no PostgreSQL. O serviço pode ajudá-lo a escalar seus casos de uso de IA desde a prova de conceito até a produção. Oferece os seguintes benefícios:

  • Fornece uma interface SQL familiar para consultar incorporações vetoriais e dados relacionais.
  • Aumenta pgvector com uma pesquisa de semelhança mais rápida e precisa em mais de 100 milhões de vetores usando o algoritmo de indexação DiskANN.
  • Simplifica as operações integrando metadados relacionais, incorporações vetoriais e dados de séries cronológicas em um único banco de dados.
  • Usa o poder do ecossistema robusto do PostgreSQL e da plataforma de nuvem do Azure para recursos de nível empresarial, incluindo replicação e alta disponibilidade.

Autenticação

O Banco de Dados do Azure para PostgreSQL dá suporte à autenticação baseada em senha e ao Microsoft Entra (anteriormente Azure Ative Directory).

A autenticação do Microsoft Entra permite que você use o Microsoft Entra ID para autenticar no seu servidor PostgreSQL. O Microsoft Entra ID elimina a necessidade de gerenciar nomes de usuário e senhas separados para os usuários do banco de dados. Ele permite que você use os mesmos mecanismos de segurança que você usa para outros serviços do Azure.

Neste artigo, você pode usar qualquer um dos métodos de autenticação.

Configuração

O Banco de Dados do Azure para PostgreSQL usa o suporte LangChain Postgres de código aberto para se conectar ao Banco de Dados do Azure para PostgreSQL. Primeiro, faça o download do pacote de parceiros:

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

Habilitar pgvector no Banco de Dados do Azure para PostgreSQL

Consulte Habilitar e usar pgvector no Banco de Dados do Azure para PostgreSQL.

Configurar credenciais

Você precisa obter seus detalhes de conexão do Banco de Dados do Azure para PostgreSQL e adicioná-los como variáveis de ambiente.

Defina o sinalizador USE_ENTRA_AUTH para True se quiser usar a autenticação do Microsoft Entra. Se você estiver usando a autenticação do Microsoft Entra, precisará fornecer apenas os nomes de host e banco de dados. Se você estiver usando a autenticação de senha, também precisará definir o nome de usuário e a senha.

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:")

Configurar incorporações do 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",
)

Inicialização

Utilize a autenticação do Microsoft Entra

As seções a seguir contêm funções que configuram LangChain para usar a autenticação do Microsoft Entra. A função get_token_and_username recupera tokens para o serviço Banco de Dados do Azure para PostgreSQL usando DefaultAzureCredential da biblioteca azure.identity. Ele garante que o mecanismo SQLAlchemy tenha um token válido com o qual criar novas conexões. Ele também analisa o token, que é um JSON Web Token (JWT), para extrair o nome de usuário usado para se conectar ao banco de dados.

A create_postgres_engine função cria um mecanismo SQLAlchemy que define dinamicamente o nome de usuário e a senha com base no token obtido do gerenciador de tokens. Este motor pode ser passado para o parâmetro connection do armazenamento vetorial PGVector LangChain.

Iniciar sessão no Azure

Para entrar no Azure, verifique se você tem a CLI do Azure instalada. Execute o seguinte comando no seu terminal:

az login

Depois de entrar, o código a seguir busca o token:

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

Usar autenticação de senha

Se você não estiver usando a autenticação do Microsoft Entra, get_connection_uri forneça um URI de conexão que extraia o nome de usuário e a senha das variáveis de ambiente:

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

Criar o repositório de vetores

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,
)

Gestão da loja de vetores

Adicionar itens ao repositório de vetores

Adicionar documentos por ID substitui quaisquer documentos existentes que correspondam a esse ID.

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])

Atualizar itens no repositório de vetores

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])

Excluir itens do repositório de vetores

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

Consultas ao repositório de vetores

Depois de criar seu repositório de vetores e adicionar os documentos relevantes, você pode consultar o repositório de vetores em sua cadeia ou agente.

Suporte de filtragem

O armazenamento vetorial suporta um conjunto de filtros que podem ser aplicados nos campos de metadados dos documentos:

Operador Significado/Categoria
$eq Igualdade (==)
$ne Desigualdade (!=)
$lt Menor que (<)
$lte Menor ou igual (<=)
$gt Maior que (>)
$gte Maior ou igual (>=)
$in Tratado especialmente (no âmbito de)
$nin Exceção especial (não incluído)
$between Caso especial (entre)
$like Texto (gosto)
$ilike Texto (não diferencia maiúsculas de minúsculas)
$and Lógica (e)
$or Lógico (ou)

Consulta direta

Você pode realizar uma pesquisa de semelhança simples da seguinte maneira:

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'}]

Se você fornecer um dicionário com vários campos, mas sem operadores, o nível superior será interpretado como um filtro lógico 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')]

Se você quiser executar uma pesquisa de similaridade e receber as pontuações correspondentes, você pode executar:

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'}]

Para obter uma lista completa das pesquisas que você pode executar em um PGVector repositório de vetores, consulte a referência da API.

Transformação em retriever

Você também pode transformar o armazenamento vetorial em um retriever para facilitar o uso em suas cadeias:

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')]

Limitações atuais

  • langchain_postgres funciona apenas com Psycopg 3 (psycopg3). Atualize suas cadeias de conexão de postgresql+psycopg2://... para postgresql+psycopg://langchain:langchain@....
  • O esquema do armazenamento e da coleção de incorporação foi alterado para funcionar add_documents corretamente com IDs especificadas pelo usuário.
  • Você precisa passar um objeto de conexão explícito agora.
  • Atualmente, não há nenhum mecanismo que suporte a migração fácil de dados em alterações de esquema. Qualquer alteração de esquema no repositório de vetores exige que você recrie as tabelas e adicione os documentos novamente.

Conteúdo relacionado