语义搜索

1
create extension if not exists vector with schema extensions;
2
3
create table embeddings (
4
  id bigint primary key generated always as identity,
5
  content text not null,
6
  embedding extensions.vector (384)
7
);
8
alter table embeddings enable row level security;
9
10
create index on embeddings using hnsw (embedding vector_ip_ops);

1
import { withSupabase } from 'npm:@supabase/server@^1'
2
3
const model = new Supabase.ai.Session('gte-small')
4
5
// Triggered by a Database Webhook, which authenticates with a secret key.
6
// Deploy with `verify_jwt = false`.
7
export default {
8
  fetch: withSupabase({ auth: 'secret' }, async (req, ctx) => {
9
    const payload: WebhookPayload = await req.json()
10
    const { content, id } = payload.record
11
12
    // Generate embedding.
13
    const embedding = await model.run(content, {
14
      mean_pool: true,
15
      normalize: true,
16
    })
17
18
    // Store in database.
19
    const { error } = await ctx.supabaseAdmin
20
      .from('embeddings')
21
      .update({ embedding: JSON.stringify(embedding) })
22
      .eq('id', id)
23
    if (error) console.warn(error.message)
24
25
    return Response.json({ ok: true })
26
  }),
27
}

1
-- Matches document sections using vector similarity search on embeddings
2
--
3
-- Returns a setof embeddings so that we can use PostgREST resource embeddings (joins with other tables)
4
-- Additional filtering like limits can be chained to this function call
5
create or replace function query_embeddings(embedding extensions.vector(384), match_threshold float)
6
returns setof embeddings
7
language plpgsql
8
as $$
9
begin
10
  return query
11
  select *
12
  from embeddings
13
14
  -- The inner product is negative, so we negate match_threshold
15
  where embeddings.embedding <#> embedding < -match_threshold
16
17
  -- Our embeddings are normalized to length 1, so cosine similarity
18
  -- and inner product will produce the same query results.
19
  -- Using inner product which can be computed faster.
20
  --
21
  -- For the different distance functions, see https://github.com/pgvector/pgvector
22
  order by embeddings.embedding <#> embedding;
23
end;
24
$$;

1
import { withSupabase } from 'npm:@supabase/server@^1'
2
3
const model = new Supabase.ai.Session('gte-small')
4
5
export default {
6
  fetch: withSupabase({ auth: 'user' }, async (req, ctx) => {
7
    const { search } = await req.json()
8
    if (!search) return Response.json({ error: 'Please provide a search param!' }, { status: 400 })
9
    // Generate embedding for search term.
10
    const embedding = await model.run(search, {
11
      mean_pool: true,
12
      normalize: true,
13
    })
14
15
    // Query embeddings.
16
    const { data: result, error } = await ctx.supabase
17
      .rpc('query_embeddings', {
18
        embedding,
19
        match_threshold: 0.8,
20
      })
21
      .select('content')
22
      .limit(3)
23
    if (error) {
24
      return Response.json({ error: error.message }, { status: 500 })
25
    }
26
27
    return Response.json({ search, result })
28
  }),
29
}