Supabase Vector Search (pgvector)

-- Enable the vector extension
CREATE EXTENSION IF NOT EXISTS vector;

-- Create a table with a vector column
CREATE TABLE documents (
  id bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
  title text NOT NULL,
  content text NOT NULL,
  embedding vector(1536),  -- OpenAI text-embedding-3-small dimension
  metadata jsonb DEFAULT '{}',
  created_at timestamptz DEFAULT NOW()
);

-- Enable RLS
ALTER TABLE documents ENABLE ROW LEVEL SECURITY;

import { createClient } from '@supabase/supabase-js'
import OpenAI from 'openai'

const supabase = createClient(SUPABASE_URL, SUPABASE_KEY)
const openai = new OpenAI()

async function storeDocument(title: string, content: string) {
  // Generate embedding with OpenAI
  const embeddingResponse = await openai.embeddings.create({
    model: 'text-embedding-3-small',
    input: content,
  })

  const embedding = embeddingResponse.data[0].embedding

  // Store in Supabase
  const { data, error } = await supabase.from('documents').insert({
    title,
    content,
    embedding,
  })

  return data
}

-- Create a function for similarity search
CREATE OR REPLACE FUNCTION match_documents(
  query_embedding vector(1536),
  match_threshold float DEFAULT 0.78,
  match_count int DEFAULT 10
)
RETURNS TABLE (
  id bigint,
  title text,
  content text,
  similarity float
)
LANGUAGE sql STABLE
AS $$
  SELECT
    documents.id,
    documents.title,
    documents.content,
    1 - (documents.embedding <=> query_embedding) AS similarity
  FROM documents
  WHERE 1 - (documents.embedding <=> query_embedding) > match_threshold
  ORDER BY documents.embedding <=> query_embedding
  LIMIT match_count;
$$;

// Call similarity search from your app
async function searchDocuments(query: string) {
  // Generate embedding for the search query
  const embeddingResponse = await openai.embeddings.create({
    model: 'text-embedding-3-small',
    input: query,
  })

  const queryEmbedding = embeddingResponse.data[0].embedding

  // Search Supabase
  const { data: documents } = await supabase.rpc('match_documents', {
    query_embedding: queryEmbedding,
    match_threshold: 0.78,
    match_count: 5,
  })

  return documents
}

-- HNSW index (recommended for most use cases)
-- Faster queries, more memory during build
CREATE INDEX idx_documents_embedding ON documents
  USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 64);

-- IVFFlat index (better for very large datasets)
-- Must have data in the table before creating
CREATE INDEX idx_documents_embedding_ivf ON documents
  USING ivfflat (embedding vector_cosine_ops)
  WITH (lists = 100);  -- Roughly sqrt(total_rows)

-- Distance operators:
-- <=>  Cosine distance    (vector_cosine_ops)
-- <->  L2/Euclidean       (vector_l2_ops)
-- <#>  Inner product       (vector_ip_ops)

async function askQuestion(question: string) {
  // 1. Find relevant documents
  const documents = await searchDocuments(question)

  // 2. Build context from retrieved documents
  const context = documents
    .map(doc => doc.content)
    .join('\n\n')

  // 3. Ask the LLM with context
  const response = await openai.chat.completions.create({
    model: 'gpt-4o',
    messages: [
      {
        role: 'system',
        content: 'Answer based on the provided context. If the context does not contain relevant information, say so.',
      },
      {
        role: 'user',
        content: `Context:\n${context}\n\nQuestion: ${question}`,
      },
    ],
  })

  return response.choices[0].message.content
}