Usando LangChain com o Banco de Dados do Azure para PostgreSQL

O Banco de Dados do Azure para PostgreSQL integra-se perfeitamente com os principais pacotes de orquestração llm (modelo de linguagem grande), como LangChain, permitindo que os desenvolvedores aproveitem o poder das funcionalidades avançadas de IA em seus aplicativos. O LangChain pode simplificar o gerenciamento e o uso de LLMs, modelos de inserção e bancos de dados, tornando ainda mais fácil desenvolver aplicativos de IA generativos.

Este tutorial mostra como usar o banco de dados de vetor integrado do Banco de Dados do Azure para PostgreSQL para armazenar e gerenciar documentos em coleções com LangChain. Ele também mostra como criar índices e executar consultas de pesquisa de vetor usando algoritmos vizinhos mais próximos aproximados, como Distância cossina, L2 (distância euclidiana) e IP (produto interno) para localizar documentos próximos aos vetores de consulta.

Suporte ao vetor

O Banco de Dados do Azure para PostgreSQL – Servidor Flexível permite que você armazene e consulte milhões de inserções de vetor no PostgreSQL com eficiência e dimensione seus casos de uso de IA de POC (prova de conceito) para produção:

• Fornece uma interface SQL familiar para consultar inserções de vetor e dados relacionais.
• Melhora pgvector com uma pesquisa de similaridade 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, inserções de vetor e dados de série temporal em um único banco de dados.
• Usa o poder do ecossistema robusto do PostgreSQL e do Azure Cloud para recursos de nível empresarial, incluindo replicação e alta disponibilidade.

Autenticação

O Banco de Dados do Azure para PostgreSQL – Servidor Flexível dá suporte à autenticação baseada em senha, bem como ao Microsoft Entra (antigo Azure Active Directory). A autenticação Entra permite que você use a identidade Entra para autenticar-se no servidor PostgreSQL. A ID do Entra elimina a necessidade de gerenciar nomes de usuário e senhas separados para os usuários do banco de dados e permite que você use os mesmos mecanismos de segurança usados para outros serviços do Azure.

Este notebook está configurado para usar qualquer um dos métodos de autenticação. Você pode configurar se deseja ou não usar a autenticação do Entra posteriormente no notebook.

Configuração

O Banco de Dados do Azure para PostgreSQL usa o suporte do Postgres do LangChain, que é de software livre, para se conectar ao Banco de Dados do Azure para PostgreSQL. Primeiro baixe o 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 – Servidor Flexível

Consulte as instruções de habilitação para o Banco de Dados do Azure para PostgreSQL.

Credenciais

Você precisa dos detalhes da conexão do Banco de Dados do Azure para PostgreSQL e adicioná-los como variáveis de ambiente para executar este notebook.

Defina o USE_ENTRA_AUTH sinalizador para True se você quiser usar a autenticação do Microsoft Entra. Se estiver usando a autenticação do Entra, você só precisará fornecer o nome do host e do banco de dados. Se estiver usando a autenticação de senha, você 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 using a username and password, supply them here
    os.environ["DBUSER"] = "<username>"
    os.environ["DBPASSWORD"] = getpass.getpass("Database Password:")

Configurar Azure OpenAI Embeddings

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

Autenticação do Microsoft Entra

A célula a seguir contém funções que configuram LangChain para usar a autenticação do Entra. Ele fornece uma função get_token_and_username que recupera tokens para o serviço Bancos 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 JWT (Token Web Java), para extrair o nome de usuário usado para se conectar ao banco de dados.

A função create_postgres_engine cria um sqlalchemy Engine que define dinamicamente o nome de usuário e a senha com base no token buscado do TokenManager. Este Engine pode ser passado para o parâmetro connection do LangChain VectorStore PGVector.

Registro no Azure

Para fazer logon no Azure, verifique se você tem a CLI do Azure instalada. Você precisa executar o seguinte comando em seu terminal:

az login

Depois de fazer logon, o código a seguir buscará 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 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

Autenticação de senha

Se não estiver usando a autenticação do Entra, o get_connection_uri fornece um URI de conexão que recupera 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 connection URI
    # Use psycopg 3!
    db_uri = (
        f"postgresql+psycopg://{dbuser}:{password}@{dbhost}/{dbname}?sslmode={sslmode}"
    )
    return db_uri

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

Gerenciar repositório de vetores

Adicionar itens ao repositório de vetores

Adicionar documentos por ID sobrescreve 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"])

Repositório de vetores de consulta

Quando o repositório de vetores tiver sido criado e os documentos relevantes tiverem sido adicionados, você poderá consultar o repositório de vetores em sua cadeia ou agente.

Suporte à filtragem

O repositório de vetores dá suporte a 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 a (<=)
$gt	Maior que (>)
$gte	Maior ou igual a (>=)
$in	Especial cased (em)
$nin	Especial cased (não em)
$between	Especial Cased (entre)
$like	Texto (como)
$ilike	Texto (sem diferenciar maiúsculas de minúsculas como)
$and	Lógico (e)
$or	Lógico (ou)

Consultar diretamente

A execução de uma pesquisa de similaridade simples pode ser feita 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 ditado com vários campos, mas sem operadores, o nível superior será interpretado como um filtro AND lógico

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, poderá 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 diferentes pesquisas que você pode executar em um PGVector repositório de vetores, consulte a referência de API.

Consulta transformando-se em recuperador

Você também pode transformar o repositório de vetores em um recuperador 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 psycopg3. Atualizar suas cadeias de conexão de postgresql+psycopg2://... para postgresql+psycopg://langchain:langchain@...
• O esquema do repositório de inserção e da coleção foi alterado para fazer add_documents funcionar corretamente com as IDs especificadas pelo usuário.
• Agora é preciso passar um objeto de conexão explícito.

Atualmente, não há nenhum mecanismo que dê suporte à migração de dados fácil em alterações de esquema. Portanto, todas as alterações de esquema no repositório de vetores exigem que o usuário recrie as tabelas e leia os documentos.

Conteúdo relacionado

• Saiba mais sobre o suporte a LangChain PGVector
• GenAI Frameworks e Banco de Dados do Azure para PostgreSQL
• Agentes de IA no Banco de Dados do Azure para PostgreSQL
• Saiba mais sobre a integração do Serviço OpenAI do Azure
• IA generativa com servidor flexível do Banco de Dados do Azure para PostgreSQL.
• Habilite e use o pgvector no servidor flexível do Banco de Dados do Azure para PostgreSQL.