RAG 실전 가이드: TypeScript와 AI SDK로 검색 증강 생성 구축하기

# Next.js 프로젝트에 필요한 패키지 설치
pnpm add ai @ai-sdk/react drizzle-orm @neondatabase/serverless
pnpm add -D drizzle-kit

// src/db/schema.ts
import { pgTable, text, varchar, vector, index, serial, timestamp } from 'drizzle-orm/pg-core'

// 문서 원본 테이블
export const documents = pgTable('documents', {
  id: serial('id').primaryKey(),
  title: varchar('title', { length: 256 }).notNull(),
  content: text('content').notNull(),
  source: varchar('source', { length: 512 }),
  createdAt: timestamp('created_at').defaultNow(),
})

// 청크 + 임베딩 테이블
export const chunks = pgTable(
  'chunks',
  {
    id: serial('id').primaryKey(),
    documentId: serial('document_id').references(() => documents.id),
    content: text('content').notNull(),
    embedding: vector('embedding', { dimensions: 1536 }).notNull(),
    metadata: text('metadata'), // JSON 문자열로 추가 정보 저장
  },
  (table) => [
    // HNSW 인덱스로 벡터 검색 성능 최적화
    index('chunks_embedding_idx').using(
      'hnsw',
      table.embedding.op('vector_cosine_ops')
    ),
  ]
)

// src/lib/chunking.ts

interface Chunk {
  content: string
  metadata: {
    startIndex: number
    endIndex: number
    chunkIndex: number
  }
}

const CHUNK_SIZE = 512    // 토큰 기준 400~512 추천
const CHUNK_OVERLAP = 64  // 10~20% 오버랩

const SEPARATORS = ['\n\n', '\n', '. ', ' ', '']

export function splitIntoChunks(text: string): Chunk[] {
  const chunks: Chunk[] = []

  function recursiveSplit(text: string, separatorIndex: number): string[] {
    if (text.length <= CHUNK_SIZE) return [text]

    const separator = SEPARATORS[separatorIndex] ?? ''
    const parts = text.split(separator)
    const result: string[] = []
    let current = ''

    for (const part of parts) {
      const candidate = current ? current + separator + part : part

      if (candidate.length > CHUNK_SIZE && current) {
        result.push(current)
        // 오버랩: 이전 청크의 마지막 부분을 다음 청크 시작에 포함
        const overlapText = current.slice(-CHUNK_OVERLAP)
        current = overlapText + separator + part
      } else {
        current = candidate
      }
    }

    if (current) result.push(current)
    return result
  }

  const parts = recursiveSplit(text, 0)

  parts.forEach((content, index) => {
    chunks.push({
      content: content.trim(),
      metadata: {
        startIndex: text.indexOf(content),
        endIndex: text.indexOf(content) + content.length,
        chunkIndex: index,
      },
    })
  })

  return chunks
}

// src/lib/indexing.ts
import { embedMany } from 'ai'
import { db } from '@/db'
import { chunks, documents } from '@/db/schema'
import { splitIntoChunks } from './chunking'

export async function indexDocument(title: string, content: string, source?: string) {
  // 1. 문서 원본 저장
  const [doc] = await db
    .insert(documents)
    .values({ title, content, source })
    .returning({ id: documents.id })

  // 2. 청킹
  const textChunks = splitIntoChunks(content)

  // 3. 배치 임베딩 생성
  const { embeddings } = await embedMany({
    model: 'openai/text-embedding-3-small', // AI Gateway 경유
    values: textChunks.map((chunk) => chunk.content),
  })

  // 4. 벡터 DB에 저장
  const chunkRecords = textChunks.map((chunk, i) => ({
    documentId: doc.id,
    content: chunk.content,
    embedding: embeddings[i],
    metadata: JSON.stringify(chunk.metadata),
  }))

  await db.insert(chunks).values(chunkRecords)

  return {
    documentId: doc.id,
    chunksCreated: chunkRecords.length,
  }
}

// src/lib/retrieval.ts
import { embed } from 'ai'
import { db } from '@/db'
import { chunks } from '@/db/schema'
import { cosineDistance, desc, sql, gt } from 'drizzle-orm'

export async function searchSimilarChunks(query: string, topK = 5) {
  // 1. 질문을 벡터로 변환
  const { embedding: queryEmbedding } = await embed({
    model: 'openai/text-embedding-3-small',
    value: query,
  })

  // 2. 코사인 유사도 기반 검색
  const similarity = sql<number>`1 - (${cosineDistance(chunks.embedding, queryEmbedding)})`

  const results = await db
    .select({
      id: chunks.id,
      content: chunks.content,
      similarity,
    })
    .from(chunks)
    .where(gt(similarity, 0.3)) // 유사도 0.3 이상만
    .orderBy(desc(similarity))
    .limit(topK)

  return results
}

// src/app/api/chat/route.ts
import { streamText, convertToModelMessages, tool } from 'ai'
import { z } from 'zod'
import { searchSimilarChunks } from '@/lib/retrieval'

export async function POST(req: Request) {
  const { messages } = await req.json()

  const result = streamText({
    model: 'anthropic/claude-sonnet-4.6',
    system: `You are a helpful assistant that answers questions based on the provided context.
If the context doesn't contain relevant information, say so honestly.
Always cite which document your answer is based on.`,
    messages: convertToModelMessages(messages),
    tools: {
      searchKnowledge: tool({
        description: 'Search the knowledge base for relevant information',
        inputSchema: z.object({
          query: z.string().describe('The search query to find relevant documents'),
        }),
        execute: async ({ query }) => {
          const results = await searchSimilarChunks(query, 5)
          return results.map((r) => ({
            content: r.content,
            relevance: Math.round(r.similarity * 100) + '%',
          }))
        },
      }),
    },
    stopWhen: stepCountIs(3),
  })

  return result.toUIMessageStreamResponse()
}

// src/lib/agentic-rag.ts
import { streamText, tool } from 'ai'
import { z } from 'zod'
import { searchSimilarChunks } from './retrieval'

// Agentic RAG: LLM이 검색 전략을 스스로 결정
export function createAgenticRAG() {
  return {
    // 1. 검색 도구 — 필요할 때만 호출
    searchKnowledge: tool({
      description:
        'Search the knowledge base. Use this when the user asks about specific topics, policies, or technical details that require factual information.',
      inputSchema: z.object({
        query: z.string().describe('Search query - be specific and focused'),
      }),
      execute: async ({ query }) => {
        const results = await searchSimilarChunks(query, 5)
        return results
      },
    }),

    // 2. 쿼리 재작성 도구 — 검색 결과가 부족할 때
    rewriteQuery: tool({
      description:
        'Rewrite a search query to get better results. Use when initial search results are not relevant enough.',
      inputSchema: z.object({
        originalQuery: z.string(),
        reason: z.string().describe('Why the original query needs rewriting'),
        rewrittenQuery: z.string().describe('The improved query'),
      }),
      execute: async ({ rewrittenQuery }) => {
        const results = await searchSimilarChunks(rewrittenQuery, 5)
        return results
      },
    }),

    // 3. 출처 인용 도구 — 답변의 근거 명시
    citeSource: tool({
      description: 'Cite a specific source that supports your answer',
      inputSchema: z.object({
        content: z.string().describe('The relevant excerpt from the source'),
        relevance: z.string().describe('How this source supports the answer'),
      }),
      execute: async ({ content, relevance }) => {
        return { cited: true, content, relevance }
      },
    }),
  }
}

// src/lib/hybrid-search.ts
import { embed } from 'ai'
import { db } from '@/db'
import { chunks } from '@/db/schema'
import { cosineDistance, desc, sql, or, ilike } from 'drizzle-orm'

export async function hybridSearch(query: string, topK = 5) {
  const { embedding } = await embed({
    model: 'openai/text-embedding-3-small',
    value: query,
  })

  // 시맨틱 검색 점수
  const semanticScore = sql<number>`1 - (${cosineDistance(chunks.embedding, embedding)})`

  // 키워드 검색 점수 (PostgreSQL ts_rank)
  const keywordScore = sql<number>`
    ts_rank(
      to_tsvector('simple', ${chunks.content}),
      plainto_tsquery('simple', ${query})
    )
  `

  // 하이브리드 점수: 시맨틱 70% + 키워드 30% (가중치 조절 가능)
  const hybridScore = sql<number>`(${semanticScore} * 0.7 + ${keywordScore} * 0.3)`

  const results = await db
    .select({
      id: chunks.id,
      content: chunks.content,
      semanticScore,
      keywordScore,
      hybridScore,
    })
    .from(chunks)
    .orderBy(desc(hybridScore))
    .limit(topK)

  return results
}

import { createHash } from 'crypto'

function contentHash(text: string): string {
  return createHash('sha256').update(text).digest('hex')
}

// 이미 임베딩된 청크인지 해시로 확인
const existing = await db
  .select()
  .from(chunks)
  .where(eq(chunks.contentHash, contentHash(newContent)))
  .limit(1)

if (existing.length > 0) {
  // 이미 인덱싱됨, 스킵
  return existing[0]
}