Supabase向量搜索 #

向量搜索特点
├── 基于语义相似度搜索
├── 支持自然语言查询
├── 支持多模态搜索
├── 适合AI应用
└── 使用pgvector扩展

-- 启用pgvector扩展
CREATE EXTENSION IF NOT EXISTS vector;

-- 查看已安装扩展
SELECT * FROM pg_extension WHERE extname = 'vector';

-- 创建带向量列的表
CREATE TABLE documents (
    id BIGSERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding vector(1536),  -- OpenAI ada-002 维度
    metadata JSONB DEFAULT '{}',
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- 创建向量索引
CREATE INDEX idx_documents_embedding 
ON documents 
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);

// Edge Function: generate-embedding
import { createClient } from 'https://esm.sh/@supabase/supabase-js@2'

Deno.serve(async (req) => {
  const { content } = await req.json()

  // 调用OpenAI API
  const response = await fetch('https://api.openai.com/v1/embeddings', {
    method: 'POST',
    headers: {
      'Authorization': `Bearer ${Deno.env.get('OPENAI_API_KEY')}`,
      'Content-Type': 'application/json',
    },
    body: JSON.stringify({
      model: 'text-embedding-ada-002',
      input: content,
    }),
  })

  const data = await response.json()
  const embedding = data.data[0].embedding

  // 存储到数据库
  const supabase = createClient(
    Deno.env.get('SUPABASE_URL')!,
    Deno.env.get('SUPABASE_SERVICE_ROLE_KEY')!
  )

  const { error } = await supabase
    .from('documents')
    .insert({ content, embedding })

  return new Response(JSON.stringify({ embedding }), {
    headers: { 'Content-Type': 'application/json' },
  })
})

-- 创建触发器自动生成嵌入
-- 注意：需要在Edge Function中实现

-- 查找最相似的文档
SELECT 
    id,
    content,
    1 - (embedding <=> '[0.1, 0.2, ...]'::vector) as similarity
FROM documents
ORDER BY embedding <=> '[0.1, 0.2, ...]'::vector
LIMIT 10;

CREATE OR REPLACE FUNCTION search_documents(
    query_embedding vector(1536),
    match_threshold FLOAT DEFAULT 0.7,
    match_count INTEGER DEFAULT 10
)
RETURNS TABLE(
    id BIGINT,
    content TEXT,
    similarity FLOAT
) AS $$
BEGIN
    RETURN QUERY
    SELECT 
        documents.id,
        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;
END;
$$ LANGUAGE plpgsql;

async function searchSimilarDocuments(query: string) {
  // 1. 生成查询向量
  const embedding = await generateEmbedding(query)

  // 2. 搜索相似文档
  const { data, error } = await supabase.rpc('search_documents', {
    query_embedding: embedding,
    match_threshold: 0.7,
    match_count: 10,
  })

  return data
}

-- 创建IVFFlat索引
CREATE INDEX idx_documents_embedding 
ON documents 
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);

-- lists参数建议
-- 数据量 < 1000: lists = 10
-- 数据量 < 10000: lists = 50
-- 数据量 < 100000: lists = 100
-- 数据量 > 100000: lists = sqrt(行数)

-- 创建HNSW索引（更高性能）
CREATE INDEX idx_documents_embedding_hnsw
ON documents
USING hnsw (embedding vector_cosine_ops)
WITH (
    m = 16,        -- 连接数
    ef_construction = 64  -- 构建时搜索范围
);

索引选择建议
├── IVFFlat
│   ├── 构建快
│   ├── 内存占用小
│   └── 适合中等规模数据
│
└── HNSW
    ├── 查询快
    ├── 构建慢
    ├── 内存占用大
    └── 适合大规模数据

CREATE TABLE knowledge_base (
    id BIGSERIAL PRIMARY KEY,
    title TEXT NOT NULL,
    content TEXT NOT NULL,
    embedding vector(1536),
    source TEXT,
    metadata JSONB DEFAULT '{}',
    created_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE INDEX idx_knowledge_embedding 
ON knowledge_base 
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100);

CREATE OR REPLACE FUNCTION rag_search(
    query_embedding vector(1536),
    match_count INTEGER DEFAULT 5
)
RETURNS TABLE(
    id BIGINT,
    title TEXT,
    content TEXT,
    source TEXT,
    similarity FLOAT
) AS $$
BEGIN
    RETURN QUERY
    SELECT 
        knowledge_base.id,
        knowledge_base.title,
        knowledge_base.content,
        knowledge_base.source,
        1 - (knowledge_base.embedding <=> query_embedding) as similarity
    FROM knowledge_base
    ORDER BY knowledge_base.embedding <=> query_embedding
    LIMIT match_count;
END;
$$ LANGUAGE plpgsql;

// Edge Function: rag-chat
import { createClient } from 'https://esm.sh/@supabase/supabase-js@2'

Deno.serve(async (req) => {
  const { question } = await req.json()

  // 1. 生成问题向量
  const embeddingResponse = await fetch('https://api.openai.com/v1/embeddings', {
    method: 'POST',
    headers: {
      'Authorization': `Bearer ${Deno.env.get('OPENAI_API_KEY')}`,
      'Content-Type': 'application/json',
    },
    body: JSON.stringify({
      model: 'text-embedding-ada-002',
      input: question,
    }),
  })
  const embeddingData = await embeddingResponse.json()
  const queryEmbedding = embeddingData.data[0].embedding

  // 2. 搜索相关知识
  const supabase = createClient(
    Deno.env.get('SUPABASE_URL')!,
    Deno.env.get('SUPABASE_SERVICE_ROLE_KEY')!
  )

  const { data: contexts } = await supabase.rpc('rag_search', {
    query_embedding: queryEmbedding,
    match_count: 5,
  })

  // 3. 构建提示词
  const context = contexts?.map((c: any) => c.content).join('\n\n')

  // 4. 调用LLM生成回答
  const chatResponse = await fetch('https://api.openai.com/v1/chat/completions', {
    method: 'POST',
    headers: {
      'Authorization': `Bearer ${Deno.env.get('OPENAI_API_KEY')}`,
      'Content-Type': 'application/json',
    },
    body: JSON.stringify({
      model: 'gpt-4-turbo-preview',
      messages: [
        {
          role: 'system',
          content: 'Based on the following context, answer the question.',
        },
        {
          role: 'user',
          content: `Context:\n${context}\n\nQuestion: ${question}`,
        },
      ],
    }),
  })

  const chatData = await chatResponse.json()
  const answer = chatData.choices[0].message.content

  return new Response(JSON.stringify({
    answer,
    sources: contexts?.map((c: any) => ({ id: c.id, title: c.title })),
  }), {
    headers: { 'Content-Type': 'application/json' },
  })
})

-- 向量加法
SELECT '[1,2,3]'::vector + '[4,5,6]'::vector;

-- 向量减法
SELECT '[4,5,6]'::vector - '[1,2,3]'::vector;

-- 向量乘法
SELECT '[1,2,3]'::vector * 2;

-- 内积
SELECT '[1,2,3]'::vector <#> '[4,5,6]'::vector;

-- 余弦距离
SELECT '[1,2,3]'::vector <=> '[4,5,6]'::vector;

-- L2距离
SELECT '[1,2,3]'::vector <-> '[4,5,6]'::vector;

-- 获取向量维度
SELECT vector_dims('[1,2,3]'::vector);

-- 获取向量范数
SELECT vector_norm('[1,2,3]'::vector);

// 文本分块
function chunkText(text: string, maxTokens: number = 500): string[] {
  const sentences = text.split(/[.!?]+/)
  const chunks: string[] = []
  let currentChunk = ''

  for (const sentence of sentences) {
    if ((currentChunk + sentence).length > maxTokens) {
      if (currentChunk) chunks.push(currentChunk.trim())
      currentChunk = sentence
    } else {
      currentChunk += sentence + '.'
    }
  }

  if (currentChunk) chunks.push(currentChunk.trim())

  return chunks
}

向量搜索优化建议
├── 使用合适的索引
├── 限制返回数量
├── 设置相似度阈值
├── 批量插入向量
└── 定期维护索引