Building RAG Systems with Supabase pgvector

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

-- Create documents table with embedding column
CREATE TABLE documents (
  id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
  content TEXT NOT NULL,
  metadata JSONB,
  embedding vector(1536), -- OpenAI ada-002 dimensions
  created_at TIMESTAMPTZ DEFAULT NOW(),
  updated_at TIMESTAMPTZ DEFAULT NOW()
);

-- Create HNSW index for fast similarity search
CREATE INDEX documents_embedding_idx ON documents 
USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- Create function for similarity search
CREATE OR REPLACE FUNCTION search_documents(
  query_embedding vector(1536),
  match_count INT DEFAULT 5,
  filter JSONB DEFAULT '{}'::JSONB
)
RETURNS TABLE(
  id UUID,
  content TEXT,
  metadata JSONB,
  similarity FLOAT
)
LANGUAGE plpgsql
AS $$
BEGIN
  RETURN QUERY
  SELECT
    d.id,
    d.content,
    d.metadata,
    1 - (d.embedding <=> query_embedding) AS similarity
  FROM documents d
  WHERE d.metadata @> filter
  ORDER BY d.embedding <=> query_embedding
  LIMIT match_count;
END;
$$;

// TypeScript implementation with Supabase client
import { createClient } from '@supabase/supabase-js';
import OpenAI from 'openai';

const supabase = createClient(
  process.env.SUPABASE_URL!,
  process.env.SUPABASE_ANON_KEY!
);

const openai = new OpenAI({
  apiKey: process.env.OPENAI_API_KEY!
});

async function generateEmbedding(text: string): Promise<number[]> {
  const response = await openai.embeddings.create({
    model: 'text-embedding-ada-002',
    input: text
  });
  return response.data[0].embedding;
}

async function indexDocument(
  content: string, 
  metadata: Record<string, any> = {}
) {
  // Chunk long content for better retrieval
  const chunks = chunkText(content, 500, 100); // 500 tokens, 100 overlap
  
  for (const chunk of chunks) {
    const embedding = await generateEmbedding(chunk);
    
    const { error } = await supabase
      .from('documents')
      .insert({
        content: chunk,
        metadata: {
          ...metadata,
          source: 'manual_upload',
          chunk_index: chunks.indexOf(chunk),
          total_chunks: chunks.length
        },
        embedding
      });
      
    if (error) throw error;
  }
}

function chunkText(
  text: string, 
  maxTokens: number = 500, 
  overlap: number = 100
): string[] {
  // Simple word-based chunking - use tiktoken for production
  const words = text.split(' ');
  const chunks: string[] = [];
  
  for (let i = 0; i < words.length; i += (maxTokens - overlap)) {
    const chunk = words.slice(i, i + maxTokens).join(' ');
    if (chunk.length > 50) chunks.push(chunk);
  }
  
  return chunks;
}

// Advanced RAG retrieval with hybrid search
interface RAGQuery {
  question: string;
  filters?: Record<string, any>;
  topK?: number;
  threshold?: number;
}

async function retrieveContext(query: RAGQuery): Promise<string> {
  const { question, filters = {}, topK = 5, threshold = 0.7 } = query;
  
  // Generate embedding for the query
  const queryEmbedding = await generateEmbedding(question);
  
  // Perform similarity search
  const { data: documents, error } = await supabase
    .rpc('search_documents', {
      query_embedding: queryEmbedding,
      match_count: topK * 2, // Get more for re-ranking
      filter: filters
    });
    
  if (error) throw error;
  
  // Re-rank using cross-encoder for better accuracy
  const reranked = await rerankDocuments(question, documents);
  
  // Filter by similarity threshold
  const relevant = reranked
    .filter(doc => doc.similarity >= threshold)
    .slice(0, topK);
    
  // Combine chunks intelligently
  const context = combineChunks(relevant);
  
  return context;
}

async function rerankDocuments(
  query: string, 
  documents: any[]
): Promise<any[]> {
  // Use a cross-encoder model for re-ranking
  // This significantly improves retrieval quality
  const scores = await Promise.all(
    documents.map(async doc => {
      const score = await crossEncoderScore(query, doc.content);
      return { ...doc, rerank_score: score };
    })
  );
  
  return scores.sort((a, b) => b.rerank_score - a.rerank_score);
}

function combineChunks(documents: any[]): string {
  // Group by source and maintain order
  const grouped = documents.reduce((acc, doc) => {
    const source = doc.metadata.source;
    if (!acc[source]) acc[source] = [];
    acc[source].push(doc);
    return acc;
  }, {} as Record<string, any[]>);
  
  // Intelligently combine chunks
  let context = '';
  for (const [source, docs] of Object.entries(grouped)) {
    const sorted = docs.sort((a, b) => 
      a.metadata.chunk_index - b.metadata.chunk_index
    );
    
    context += `\n\nSource: ${source}\n`;
    context += sorted.map(d => d.content).join('\n');
  }
  
  return context.trim();
}

-- Performance monitoring query
WITH search_stats AS (
  SELECT 
    date_trunc('hour', created_at) as hour,
    COUNT(*) as searches,
    AVG(similarity) as avg_similarity,
    PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY similarity) as p95_similarity
  FROM search_logs
  WHERE created_at > NOW() - INTERVAL '24 hours'
  GROUP BY hour
)
SELECT 
  hour,
  searches,
  ROUND(avg_similarity::numeric, 3) as avg_similarity,
  ROUND(p95_similarity::numeric, 3) as p95_similarity,
  CASE 
    WHEN avg_similarity < 0.7 THEN 'Poor'
    WHEN avg_similarity < 0.8 THEN 'Fair'
    WHEN avg_similarity < 0.9 THEN 'Good'
    ELSE 'Excellent'
  END as quality_rating
FROM search_stats
ORDER BY hour DESC;

// Production-ready RAG API endpoint
import { NextRequest, NextResponse } from 'next/server';
import { z } from 'zod';
import { RateLimiter } from '@/lib/rate-limiter';
import { metrics } from '@/lib/monitoring';

const querySchema = z.object({
  question: z.string().min(1).max(500),
  context: z.enum(['general', 'technical', 'legal']).optional(),
  language: z.enum(['en', 'no', 'sv', 'da']).default('en')
});

const rateLimiter = new RateLimiter({
  tokensPerInterval: 100,
  interval: 'minute'
});

export async function POST(req: NextRequest) {
  const startTime = Date.now();
  
  try {
    // Rate limiting
    const identifier = req.headers.get('x-api-key') || req.ip;
    if (!rateLimiter.check(identifier)) {
      return NextResponse.json(
        { error: 'Rate limit exceeded' },
        { status: 429 }
      );
    }
    
    // Validate input
    const body = await req.json();
    const { question, context, language } = querySchema.parse(body);
    
    // Log request
    await logRequest({ question, context, language, identifier });
    
    // Retrieve context with caching
    const cacheKey = `rag:${hashQuery(question)}:${context}:${language}`;
    let retrievedContext = await cache.get(cacheKey);
    
    if (!retrievedContext) {
      retrievedContext = await retrieveContext({
        question,
        filters: { context, language },
        topK: 5,
        threshold: 0.75
      });
      
      // Cache for 1 hour
      await cache.set(cacheKey, retrievedContext, 3600);
    }
    
    // Generate response
    const response = await generateResponse(question, retrievedContext);
    
    // Track metrics
    const duration = Date.now() - startTime;
    metrics.recordLatency('rag_query', duration);
    metrics.incrementCounter('rag_queries_total');
    
    // Return response with metadata
    return NextResponse.json({
      answer: response.answer,
      sources: response.sources,
      confidence: response.confidence,
      metadata: {
        processingTime: duration,
        contextLength: retrievedContext.length,
        model: 'gpt-4-turbo',
        timestamp: new Date().toISOString()
      }
    });
    
  } catch (error) {
    metrics.incrementCounter('rag_errors_total');
    console.error('RAG query error:', error);
    
    if (error instanceof z.ZodError) {
      return NextResponse.json(
        { error: 'Invalid input', details: error.errors },
        { status: 400 }
      );
    }
    
    return NextResponse.json(
      { error: 'Internal server error' },
      { status: 500 }
    );
  }
}