xiaomi/mimo-v2.5 · d2_pgvector · rep 1

echo "SUPABASE_URL=$SUPABASE_URL"
echo "SUPABASE_PROJECT_REF=$SUPABASE_PROJECT_REF"
echo "SUPABASE_ANON_KEY=$SUPABASE_ANON_KEY"
echo "SUPABASE_SERVICE_ROLE_KEY=$SUPABASE_SERVICE_ROLE_KEY"
echo "SUPABASE_DB_URL=$SUPABASE_DB_URL"
echo "SUPABASE_DB_PASSWORD=$SUPABASE_DB_PASSWORD"

SUPABASE_URL=https://<ref>.supabase.co
SUPABASE_PROJECT_REF=<ref>
SUPABASE_ANON_KEY=<redacted:SUPABASE_ANON_KEY>
SUPABASE_SERVICE_ROLE_KEY=<redacted:SUPABASE_SERVICE_ROLE_KEY>
SUPABASE_DB_URL=<redacted:SUPABASE_DB_URL>
SUPABASE_DB_PASSWORD=<redacted:SUPABASE_DB_PASSWORD>

psql "$SUPABASE_DB_URL" -c "SELECT * FROM public.documents ORDER BY id;"

 id |  title  |  embedding_json  
----+---------+------------------
  1 | alpha   | [10, 0.5, 0, 0]
  2 | bravo   | [0.9, 0.3, 0, 0]
  3 | delta   | [20, 60, 0, 0]
  4 | echo    | [0, 1, 0, 0]
  5 | foxtrot | [-5, 0.1, 0, 0]
  6 | golf    | [0.1, 0.1, 3, 4]
(6 rows)

psql "$SUPABASE_DB_URL" -c "CREATE EXTENSION IF NOT EXISTS vector SCHEMA extensions;"

CREATE EXTENSION

psql "$SUPABASE_DB_URL" -c "
ALTER TABLE public.documents ADD COLUMN embedding vector(4);
UPDATE public.documents SET embedding = embedding_json::vector(4);
"

ALTER TABLE

ERROR:  cannot cast type jsonb to vector
LINE 3: ...TE public.documents SET embedding = embedding_json::vector(4...
                                                             ^

psql "$SUPABASE_DB_URL" -c "
UPDATE public.documents SET embedding = embedding_json::text::vector(4);
"

ERROR:  column "embedding" of relation "documents" does not exist
LINE 2: UPDATE public.documents SET embedding = embedding_json::text...
                                    ^

psql "$SUPABASE_DB_URL" -c "\d public.documents"

                 Table "public.documents"
     Column     |  Type   | Collation | Nullable | Default 
----------------+---------+-----------+----------+---------
 id             | integer |           | not null | 
 title          | text    |           | not null | 
 embedding_json | jsonb   |           | not null | 
Indexes:
    "documents_pkey" PRIMARY KEY, btree (id)

psql "$SUPABASE_DB_URL" -c "ALTER TABLE public.documents ADD COLUMN IF NOT EXISTS embedding vector(4);"

ALTER TABLE

psql "$SUPABASE_DB_URL" -c "
UPDATE public.documents SET embedding = embedding_json::text::vector(4);
SELECT id, title, embedding_json, embedding FROM public.documents;
"

UPDATE 6
 id |  title  |  embedding_json  |   embedding   
----+---------+------------------+---------------
  1 | alpha   | [10, 0.5, 0, 0]  | [10,0.5,0,0]
  2 | bravo   | [0.9, 0.3, 0, 0] | [0.9,0.3,0,0]
  3 | delta   | [20, 60, 0, 0]   | [20,60,0,0]
  4 | echo    | [0, 1, 0, 0]     | [0,1,0,0]
  5 | foxtrot | [-5, 0.1, 0, 0]  | [-5,0.1,0,0]
  6 | golf    | [0.1, 0.1, 3, 4] | [0.1,0.1,3,4]
(6 rows)

psql "$SUPABASE_DB_URL" -c "
CREATE INDEX documents_embedding_idx ON public.documents
  USING hnsw (embedding vector_cosine_ops);
"

CREATE INDEX

psql "$SUPABASE_DB_URL" -c "
CREATE OR REPLACE FUNCTION public.match_documents(
  query_embedding vector(4),
  match_count int
)
RETURNS TABLE (id int, title text, distance float8)
LANGUAGE sql STABLE
AS \$\$
  SELECT d.id, d.title, d.embedding <=> query_embedding AS distance
  FROM public.documents d
  ORDER BY d.embedding <=> query_embedding
  LIMIT match_count;
\$\$;
"

CREATE FUNCTION

psql "$SUPABASE_DB_URL" -c "
GRANT EXECUTE ON FUNCTION public.match_documents(vector(4), int) TO anon;
GRANT USAGE ON SCHEMA public TO anon;
"

GRANT
GRANT

psql "$SUPABASE_DB_URL" -c "
SELECT * FROM public.match_documents('[1,0,0,0]'::vector(4), 3);
"

 id | title |      distance       
----+-------+---------------------
  1 | alpha | 0.00124766112215546
  2 | bravo |  0.0513167145152487
  3 | delta |   0.683772233983162
(3 rows)

curl -s -X POST "$SUPABASE_URL/rest/v1/rpc/match_documents" \
  -H "apikey: $SUPABASE_ANON_KEY" \
  -H "Authorization: Bearer $SUPABASE_ANON_KEY" \
  -H "Content-Type: application/json" \
  -d '{"query_embedding": [1,0,0,0], "match_count": 3}' | python3 -m json.tool

[
    {
        "id": 1,
        "title": "alpha",
        "distance": 0.00124766112215546
    },
    {
        "id": 2,
        "title": "bravo",
        "distance": 0.0513167145152487
    },
    {
        "id": 3,
        "title": "delta",
        "distance": 0.683772233983162
    }
]

curl -s -X POST "$SUPABASE_URL/rest/v1/rpc/match_documents" \
  -H "apikey: $SUPABASE_ANON_KEY" \
  -H "Authorization: Bearer $SUPABASE_ANON_KEY" \
  -H "Content-Type: application/json" \
  -d '{"query_embedding": [0, 0, 3, 4], "match_count": 3}' | python3 -m json.tool

[
    {
        "id": 6,
        "title": "golf",
        "distance": 0.000399769304185948
    },
    {
        "id": 1,
        "title": "alpha",
        "distance": 1
    },
    {
        "id": 2,
        "title": "bravo",
        "distance": 1
    }
]