RAG 检索增强生成

// lib/rag/parser.ts
import pdf from 'pdf-parse'
 
export async function extractText(file: File): Promise<string> {
  const buffer = Buffer.from(await file.arrayBuffer())
 
  switch (file.type) {
    case 'application/pdf':
      const pdfData = await pdf(buffer)
      return pdfData.text
 
    case 'text/markdown':
    case 'text/plain':
      return buffer.toString('utf-8')
 
    case 'text/html':
      // 去除 HTML 标签，保留文本
      return buffer.toString('utf-8').replace(/<[^>]*>/g, ' ').replace(/\s+/g, ' ')
 
    default:
      throw new Error(`不支持的文件类型：${file.type}`)
  }
}

// lib/rag/chunker.ts
interface Chunk {
  content: string
  index: number
  metadata: {
    source: string
    page?: number
  }
}
 
export function chunkText(
  text: string,
  options: {
    chunkSize?: number      // 目标块大小（字符数）
    chunkOverlap?: number   // 重叠字符数
    source: string
  },
): Chunk[] {
  const { chunkSize = 1000, chunkOverlap = 200, source } = options
  const chunks: Chunk[] = []
 
  // 按段落分割
  const paragraphs = text.split(/\n\n+/)
  let currentChunk = ''
  let chunkIndex = 0
 
  for (const paragraph of paragraphs) {
    if (currentChunk.length + paragraph.length > chunkSize && currentChunk.length > 0) {
      chunks.push({
        content: currentChunk.trim(),
        index: chunkIndex++,
        metadata: { source },
      })
 
      // 保留重叠部分
      currentChunk = currentChunk.slice(-chunkOverlap) + '\n\n' + paragraph
    } else {
      currentChunk += (currentChunk ? '\n\n' : '') + paragraph
    }
  }
 
  // 最后一块
  if (currentChunk.trim()) {
    chunks.push({
      content: currentChunk.trim(),
      index: chunkIndex,
      metadata: { source },
    })
  }
 
  return chunks
}

// lib/rag/embeddings.ts
import { embed, embedMany } from 'ai'
import { openai } from '@ai-sdk/openai'
 
// 单个文本
export async function getEmbedding(text: string): Promise<number[]> {
  const { embedding } = await embed({
    model: openai.embedding('text-embedding-3-small'),
    value: text,
  })
  return embedding
}
 
// 批量生成（效率更高）
export async function getEmbeddings(texts: string[]): Promise<number[][]> {
  const { embeddings } = await embedMany({
    model: openai.embedding('text-embedding-3-small'),
    values: texts,
  })
  return embeddings
}

// src/db/schema.ts — 向量表
import { pgTable, uuid, text, timestamp, integer, vector } from 'drizzle-orm/pg-core'
 
export const documentChunks = pgTable('document_chunks', {
  id: uuid('id').defaultRandom().primaryKey(),
  documentId: uuid('document_id').references(() => documents.id, { onDelete: 'cascade' }).notNull(),
  content: text('content').notNull(),
  chunkIndex: integer('chunk_index').notNull(),
  embedding: vector('embedding', { dimensions: 1536 }),  // pgvector 类型
  metadata: text('metadata'),     // JSON 字符串
  createdAt: timestamp('created_at').defaultNow().notNull(),
})
 
export const documents = pgTable('documents', {
  id: uuid('id').defaultRandom().primaryKey(),
  name: text('name').notNull(),
  type: text('type').notNull(),   // 'pdf' | 'markdown' | 'url'
  workspaceId: uuid('workspace_id').references(() => workspaces.id).notNull(),
  chunkCount: integer('chunk_count').default(0),
  status: text('status', { enum: ['processing', 'ready', 'error'] }).default('processing'),
  createdAt: timestamp('created_at').defaultNow().notNull(),
})

-- migration
CREATE EXTENSION IF NOT EXISTS vector;
CREATE INDEX ON document_chunks USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

// lib/rag/ingest.ts
export async function ingestDocument(file: File, workspaceId: string) {
  // 1. 创建文档记录
  const [doc] = await db.insert(documents).values({
    name: file.name,
    type: file.type,
    workspaceId,
  }).returning()
 
  try {
    // 2. 提取文本
    const text = await extractText(file)
 
    // 3. 分块
    const chunks = chunkText(text, { source: file.name })
 
    // 4. 批量生成 Embedding
    const embeddings = await getEmbeddings(chunks.map((c) => c.content))
 
    // 5. 批量写入数据库
    await db.insert(documentChunks).values(
      chunks.map((chunk, i) => ({
        documentId: doc.id,
        content: chunk.content,
        chunkIndex: chunk.index,
        embedding: embeddings[i],
        metadata: JSON.stringify(chunk.metadata),
      })),
    )
 
    // 6. 更新状态
    await db.update(documents)
      .set({ status: 'ready', chunkCount: chunks.length })
      .where(eq(documents.id, doc.id))
 
    return doc
  } catch (error) {
    await db.update(documents)
      .set({ status: 'error' })
      .where(eq(documents.id, doc.id))
    throw error
  }
}

// lib/rag/retrieve.ts
import { sql, cosineDistance, desc } from 'drizzle-orm'
 
export async function searchSimilarChunks(
  query: string,
  workspaceId: string,
  options: { topK?: number; threshold?: number } = {},
) {
  const { topK = 5, threshold = 0.3 } = options
 
  // 1. 把用户问题向量化
  const queryEmbedding = await getEmbedding(query)
 
  // 2. 余弦相似度搜索
  const similarity = sql<number>`1 - (${cosineDistance(documentChunks.embedding, queryEmbedding)})`
 
  const results = await db
    .select({
      content: documentChunks.content,
      similarity,
      documentName: documents.name,
      chunkIndex: documentChunks.chunkIndex,
    })
    .from(documentChunks)
    .innerJoin(documents, eq(documentChunks.documentId, documents.id))
    .where(
      and(
        eq(documents.workspaceId, workspaceId),
        eq(documents.status, 'ready'),
        sql`${similarity} > ${threshold}`,  // 过滤低相关性结果
      ),
    )
    .orderBy(desc(similarity))
    .limit(topK)
 
  return results
}

// lib/rag/prompt.ts
export function buildRAGPrompt(
  chunks: { content: string; documentName: string }[],
): string {
  if (chunks.length === 0) {
    return '知识库中没有找到相关内容。请基于你的通用知识回答，并提醒用户这不是基于知识库的回答。'
  }
 
  const context = chunks
    .map((c, i) => `[来源：${c.documentName}]\n${c.content}`)
    .join('\n\n---\n\n')
 
  return `你是一个知识库问答助手。请严格基于以下参考资料回答用户的问题。
 
规则：
1. 只使用参考资料中的信息，不要编造内容
2. 如果参考资料不足以回答问题，明确说明"知识库中未找到相关信息"
3. 引用信息时标注来源文档名称
4. 回答要简洁、准确
 
参考资料：
${context}`
}

// app/api/chat/rag/route.ts
import { streamText } from 'ai'
import { openai } from '@ai-sdk/openai'
import { searchSimilarChunks } from '@/lib/rag/retrieve'
import { buildRAGPrompt } from '@/lib/rag/prompt'
 
export async function POST(request: Request) {
  const { messages, chatId } = await request.json()
  const { workspace } = await requireWorkspaceAccess()
 
  // 用最后一条用户消息做检索
  const lastUserMessage = messages.findLast((m: any) => m.role === 'user')
 
  // 检索相关文档
  const chunks = await searchSimilarChunks(lastUserMessage.content, workspace.id, {
    topK: 5,
    threshold: 0.3,
  })
 
  const result = streamText({
    model: openai('gpt-4o'),
    system: buildRAGPrompt(chunks),
    messages,
    onFinish: async ({ text, usage }) => {
      // 保存消息和引用来源
      await db.insert(dbMessages).values({
        chatId,
        role: 'assistant',
        content: text,
        metadata: JSON.stringify({ sources: chunks.map((c) => c.documentName) }),
      })
    },
  })
 
  return result.toDataStreamResponse()
}

async function rewriteQuery(query: string): Promise<string> {
  const { text } = await generateText({
    model: openai('gpt-4o-mini'),
    prompt: `将以下用户问题改写为适合知识库检索的关键词查询（保持简短，提取核心概念）：\n\n${query}`,
    maxTokens: 100,
  })
  return text
}

// 同时做向量搜索和全文搜索，合并结果
async function hybridSearch(query: string, workspaceId: string) {
  const [vectorResults, textResults] = await Promise.all([
    searchSimilarChunks(query, workspaceId, { topK: 5 }),
    db.select()
      .from(documentChunks)
      .innerJoin(documents, eq(documentChunks.documentId, documents.id))
      .where(
        and(
          eq(documents.workspaceId, workspaceId),
          sql`to_tsvector('chinese', ${documentChunks.content}) @@ plainto_tsquery('chinese', ${query})`,
        ),
      )
      .limit(3),
  ])
 
  // 去重合并
  const seen = new Set<string>()
  const merged = []
  for (const result of [...vectorResults, ...textResults]) {
    if (!seen.has(result.content)) {
      seen.add(result.content)
      merged.push(result)
    }
  }
  return merged.slice(0, 8)
}

import { CohereClient } from 'cohere-ai'
 
const cohere = new CohereClient({ token: process.env.COHERE_API_KEY })
 
async function rerankResults(query: string, chunks: { content: string }[]) {
  const response = await cohere.rerank({
    query,
    documents: chunks.map((c) => ({ text: c.content })),
    topN: 5,
    model: 'rerank-v3.5',
  })
 
  return response.results.map((r) => chunks[r.index])
}

// app/(dashboard)/knowledge/page.tsx
export default async function KnowledgePage() {
  const { workspace } = await requireWorkspaceAccess()
 
  const docs = await db.query.documents.findMany({
    where: eq(documents.workspaceId, workspace.id),
    orderBy: desc(documents.createdAt),
  })
 
  return (
    <div className="p-6">
      <h1 className="text-xl font-bold mb-4">知识库</h1>
      <UploadForm workspaceId={workspace.id} />
      <div className="mt-6 space-y-2">
        {docs.map((doc) => (
          <div key={doc.id} className="flex items-center justify-between p-3 border rounded-lg">
            <div>
              <p className="font-medium">{doc.name}</p>
              <p className="text-sm text-gray-500">
                {doc.chunkCount} 个分块 · {doc.status === 'ready' ? '✅ 就绪' : doc.status === 'processing' ? '⏳ 处理中' : '❌ 失败'}
              </p>
            </div>
            <form action={deleteDocument.bind(null, doc.id)}>
              <button className="text-red-500 text-sm">删除</button>
            </form>
          </div>
        ))}
      </div>
    </div>
  )
}