Busca Vetorial em Ação: Personalização com Embeddings de IA

Este tutorial usa Supabase + JS/TS, então ter alguma experiência básica com essas ferramentas vai ajudar.

Grande parte do que vou explicar aqui, eu desenvolvi no Dissolutus.

Para começar, precisamos ter algo que os usuários possam buscar. No meu caso, são templates, para que eles possam criar seus próprios posts.

Fica mais ou menos assim:

asicamente, cada um dos meus templates precisa de uma descrição.

E agora começa a nossa jornada com IA!

Você pode pensar em um embedding como a forma que um modelo de IA entende a diferença entre, por exemplo, uma caneta e um carro.

Um embedding nada mais é do que um vetor, ou seja, uma longa lista de números. Esses números representam o significado ou a essência de um texto. Quanto mais próximos dois vetores estão, mais relacionados são seus significados.

Por exemplo, se eu pesquisar por “promoção de comida”, templates com descrições como “pizza” ou “hambúrguer” terão uma semelhança maior — porque estão relacionados a comida. Já algo como “trabalhe conosco” ficaria bem distante.

O que precisamos agora é transformar nosso termo de busca, como "promoção de comida", em um vetor, e então procurar na nossa tabela por vetores semelhantes para encontrar os templates que mais fazem sentido.

Pra isso, vamos instalar a dependência:

pnpm i @huggingface/transformers

E agora é só adicionar esse snippet que ele faz a mágica pra gente:

import { pipeline } from '@huggingface/transformers';

const pipe = await pipeline('feature-extraction', 'Supabase/gte-small');

const output = await pipe(searchText, {
  pooling: 'mean',
  normalize: true
 });

const embedding = Array.from(output.data);

Algo muito importante: você precisa usar o mesmo modelo tanto para o termo de busca quanto para sua tabela. Caso contrário, não vai funcionar.

Aqui, vou usar o modelo Supabase/gte-small, que retorna um vetor com 384 dimensões. Vetores com menos dimensões executam consultas mais rápidas e consomem menos memória RAM. Quer aprender mais? Dá uma olhada neste artigo criado pela equipe do Supabase.

Estamos gerando o embedding da busca, e isso é ótimo.
Mas agora, vamos para o nosso banco de dados.

Precisamos adicionar ao Postgres a extensão pgvector.

Você pode fazer isso usando o painel do Supabase ou rodando no terminal do Postgres (ou em um arquivo de migração):

CREATE EXTENSION IF NOT EXISTS vector;

E para executar requisições externas a partir do Postgres, também precisamos rodar

CREATE EXTENSION IF NOT EXISTS pg_net;

A tabela templates é onde vou armazenar meu embedding.
Então eu executo:

ALTER TABLE templates ADD COLUMN embedding vector(384);

Quando o usuário admin cria um novo template, precisamos de um trigger, que chama uma função, e essa função chama outra função (sim, é isso mesmo!) que gera nosso embedding.

Nossa função é:

create or replace function embed_server()
returns trigger
language plpgsql
as $$
declare
  content_column text = TG_ARGV[0];
  embedding_column text = TG_ARGV[1];
  batch_size int = case when array_length(TG_ARGV, 1) >= 3 then TG_ARGV[2]::int else 5 end;
  timeout_milliseconds int = case when array_length(TG_ARGV, 1) >= 4 then TG_ARGV[3]::int else 5 * 60 * 1000 end;
  batch_count int = ceiling((select count(*) from inserted) / batch_size::float);
begin
  -- Loop through each batch and invoke an edge function to handle the embedding generation
  for i in 0 .. (batch_count-1) loop
  perform
    net.http_post(
      url := supabase_url() || '/functions/v1/embed_server',
      headers := jsonb_build_object(
        'Content-Type', 'application/json',
        'Authorization', current_setting('request.headers')::json->>'authorization'
      ),
      body := jsonb_build_object(
        'ids', (select json_agg(ds.id) from (select id from inserted limit batch_size offset i*batch_size) ds),
        'table', TG_TABLE_NAME,
        'contentColumn', content_column,
        'embeddingColumn', embedding_column
      ),
      timeout_milliseconds := timeout_milliseconds
    );
  end loop;

  return null;
end;
$$;

O que esse cara aí faz nada mais é do que pegar os IDs que foram modificados (criados ou atualizados) e chamar a função que gera o embedding para a tabela e colunas especificadas (a coluna que vai receber o embedding e a coluna que servirá de base para gerar o embedding).

Você pergunta: "e esse supabase_url?"
Ele funciona como uma espécie de “env”, onde podemos armazenar informações sensíveis (secrets).

Se quiser aprender mais sobre o Vault, dá uma olhada neste link.

Podemos criar um novo secret assim:

select vault.create_secret(
  '<api-url>', -- altere aqui!
  'supabase_url'
);

E daí podemos criar a função que busca esse secret:

create function supabase_url()
returns text
language plpgsql
security definer
as $$
declare
  secret_value text;
begin
  select decrypted_secret into secret_value from vault.decrypted_secrets where name = 'supabase_url';
  return secret_value;
end;
$$;

Agora podemos criar nosso trigger, que será acionado tanto em updates quanto em inserts.

-- add  trigger to embedd the action with the function embed
create trigger embed_templates_insert
  after insert on public.templates
  referencing new table as inserted
  for each statement
  execute procedure embed_server(content, embedding);

-- when a update is made, the trigger will be executed
create trigger embed_templates_update 
  after update on public.templates
  referencing new table as inserted
  for each statement
  execute procedure embed_server(content, embedding);

Agora vamos para nossa Edge Function.
Não sabe o que isso significa? Dá uma olhada neste artigo.

Como você pode ver, o nome é "embed_server", então precisamos de um arquivo com esse mesmo nome.

No seu import_map, você deve usar:

E o nosso arquivo vai ficar assim:

import { createClient } from "@supabase/supabase-js";
const model = new Supabase.ai.Session("gte-small");

const supabaseUrl = Deno.env.get("SUPABASE_URL");
const supabaseRoleKey = Deno.env.get("SUPABASE_SERVICE_ROLE_KEY");

Deno.serve(async (req) => {
  if (!supabaseUrl || !supabaseRoleKey) {
    return new Response(
      JSON.stringify({
        error: "Missing environment variables.",
      }),
      {
        status: 500,
        headers: { "Content-Type": "application/json" },
      }
    );
  }

  const { ids, table, contentColumn, embeddingColumn } = await req.json();

  if (!ids || !table || !contentColumn || !embeddingColumn) {
    return new Response(
      JSON.stringify({
        error: "Missing required parameters.",
      }),
      {
        status: 400,
        headers: { "Content-Type": "application/json" },
      }
    );
  }

  const supabase = createClient(supabaseUrl, supabaseRoleKey);

  const { data: rows, error: selectError } = await supabase
    .from(table)
    .select(`id, ${contentColumn}`)
    .in("id", ids);

  if (selectError) {
    return new Response(JSON.stringify({ error: selectError }), {
      status: 500,
      headers: { "Content-Type": "application/json" },
    });
  }

  for (const row of rows) {
    const { id, [contentColumn]: content } = row;

    if (!content) {
      console.error(`No content available in column '${contentColumn}'`);
      continue;
    }

    if (!id) {
      console.error(`No id available in column 'id'`);
      continue;
    }

    const output = (await model.run(content, {
      mean_pool: true,
      normalize: true,
    })) as number[];

    const embedding = JSON.stringify(output);

    if (!embedding) {
      console.error(`No embedding generated for id ${id}`);
      continue;
    }

    const { error } = await supabase
      .from(table)
      .update({
        [embeddingColumn]: embedding,
      })
      .eq("id", id);

    if (error) {
      console.error(
        `Failed to save embedding on '${table}' table with id ${id}`
      );

      console.error(error);
    }

    console.log(
      `Generated embedding ${JSON.stringify({
        table,
        id,
        contentColumn,
        embeddingColumn,
      })}`
    );
  }

  return new Response(null, {
    status: 204,
    headers: { "Content-Type": "application/json" },
  });
});

A ideia base dessa função e dos helpers eu peguei deste repositório.

Estamos quase lá!

Agora, quando o usuário fizer uma busca, precisamos de uma última função no Postgres que retorna uma verificação de similaridade.

Se parece com isso:

create or replace function match_table_embeddings(
  table_name text,
  embedding vector(384),
  match_threshold float
)
returns table (id uuid, similarity float)
language plpgsql
as $$
#variable_conflict use_variable
declare
  query text;
begin
  query := format(
    'select table_name.id, embedding <#> $1 as similarity
     from %I as table_name
     where embedding <#> $1 < -$2
     order by similarity;',
    table_name
  );

  -- Execute the query dynamically
  return query execute query using embedding, match_threshold;
end;
$$;

Na sua função de busca, tendo o embedding do usuário, você pode fazer algo assim:

const { data, error } = await supabase
  .rpc('match_table_embeddings', {
    table_name: 'creativehub_actions',
    embedding,
    match_threshold: 0.8
  })
  .select('id')
  .limit(LIMIT);

Sobre o match_threshold:
Quanto mais próximo de 1, mais ele retorna similaridades exatas.
Quanto mais próximo de 0, bem... ele pode começar a delirar.

Mas ainda tem mais sobre isso.

O meu artigo original está em inglês, neste link aqui

Se quiser, pode se conectar comigo no linkedin