pgvector with custom filter

429 views Asked by At

hi how can I add a custom filter to a sql function in a supabase database that uses pgvector to store embeddings.

create table documents (
  id bigserial primary key,
  content text,
  embedding vector(1536),
custom_id text // Custom id for each company
);

Something like this:

create or replace function match_documents (
  query_embedding vector(1536),
  match_threshold float,
  match_count int,
  custom_id text
)
returns table (
  id bigint,
  content text,
  similarity float
)
language sql stable
as $$
  select
    documents.id,
    documents.content,
    1 - (documents.embedding <=> query_embedding) as similarity
  from documents
  where documents.custom_id == custom_id
  order by similarity desc
  limit match_count;
$$;

The problem with this function is that if I ask for information about x custom_id, it responds with data from all the documents, it does not respect the where. I need to filter the documents by custom_id before checking the embedding similarity.

Thanks!!

1

There are 1 answers

0
dshukertjr On BEST ANSWER

I see two problems. First, Postgres uses =, and not == for equality checks. Second, since there is a column named custom_id and a parameter of the function also named custom_id, you need to disambiguate these two. You can do so by adding function name and a dot, in this case match_documents., before the custom_id in the filter.

create or replace function match_documents (
  query_embedding vector(1536),
  match_threshold float,
  match_count int,
  custom_id text
)
returns table (
  id bigint,
  content text,
  similarity float
)
language sql stable
as $$
  select
    documents.id,
    documents.content,
    1 - (documents.embedding <=> query_embedding) as similarity
  from documents
  where documents.custom_id = match_documents.custom_id
  order by similarity desc
  limit match_count;
$$;