RAG 知识库功能实现

// lib/actions/knowledge.ts
'use server'
 
export async function uploadDocument(knowledgeBaseId: string, formData: FormData) {
  const { session, membership } = await requirePermission('ai:knowledge:write')
 
  const file = formData.get('file') as File
  if (!file) throw new Error('No file provided')
 
  // 验证文件类型
  const allowedTypes = ['application/pdf', 'text/plain', 'text/markdown',
    'application/vnd.openxmlformats-officedocument.wordprocessingml.document']
  if (!allowedTypes.includes(file.type)) {
    throw new Error('Unsupported file type')
  }
 
  // 上传到 S3
  const key = `${membership.tenantId}/knowledge/${knowledgeBaseId}/${nanoid()}/${file.name}`
  const uploadUrl = await createUploadUrl(key, file.type)
  // 客户端直传（或在这里转发——小文件可以服务端中转）
 
  // 写入数据库
  const [doc] = await db.insert(documents).values({
    knowledgeBaseId,
    name: file.name,
    fileKey: key,
    fileSize: file.size,
    mimeType: file.type,
    status: 'pending',
  }).returning()
 
  // 触发异步处理
  await inngest.send({
    name: 'knowledge/document.process',
    data: { documentId: doc.id, knowledgeBaseId },
  })
 
  revalidatePath(`/knowledge/${knowledgeBaseId}`)
  return doc
}

// lib/ai/rag/parse.ts
 
export async function parseDocument(fileKey: string, mimeType: string): Promise<string> {
  const fileBuffer = await downloadFromS3(fileKey)
 
  switch (mimeType) {
    case 'text/plain':
    case 'text/markdown':
      return fileBuffer.toString('utf-8')
 
    case 'application/pdf':
      return parsePDF(fileBuffer)
 
    case 'application/vnd.openxmlformats-officedocument.wordprocessingml.document':
      return parseDocx(fileBuffer)
 
    default:
      throw new Error(`Unsupported mime type: ${mimeType}`)
  }
}
 
// PDF 解析
async function parsePDF(buffer: Buffer): Promise<string> {
  const { default: pdfParse } = await import('pdf-parse')
  const data = await pdfParse(buffer)
  return data.text
}
 
// Word 解析
async function parseDocx(buffer: Buffer): Promise<string> {
  const mammoth = await import('mammoth')
  const result = await mammoth.extractRawText({ buffer })
  return result.value
}

// lib/ai/rag/chunk.ts
 
interface ChunkOptions {
  chunkSize: number      // 每块目标大小（字符数）
  chunkOverlap: number   // 重叠区域
  separators?: string[]  // 分割优先级
}
 
export function splitIntoChunks(
  text: string,
  options: ChunkOptions = { chunkSize: 1000, chunkOverlap: 100 }
): string[] {
  const { chunkSize, chunkOverlap } = options
  const separators = options.separators || ['\n\n', '\n', '. ', ' ', '']
 
  return recursiveSplit(text, separators, chunkSize, chunkOverlap)
}
 
function recursiveSplit(
  text: string,
  separators: string[],
  chunkSize: number,
  overlap: number
): string[] {
  if (text.length <= chunkSize) return [text]
 
  // 找到能分割文本的分隔符
  let separator = ''
  for (const sep of separators) {
    if (text.includes(sep)) {
      separator = sep
      break
    }
  }
 
  // 按分隔符拆分
  const splits = separator ? text.split(separator) : [text]
  const chunks: string[] = []
  let currentChunk = ''
 
  for (const split of splits) {
    const piece = separator ? split + separator : split
 
    if ((currentChunk + piece).length > chunkSize && currentChunk.length > 0) {
      chunks.push(currentChunk.trim())
      // 重叠：保留上一块末尾的内容
      currentChunk = currentChunk.slice(-overlap) + piece
    } else {
      currentChunk += piece
    }
  }
 
  if (currentChunk.trim()) chunks.push(currentChunk.trim())
 
  return chunks
}

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

// lib/inngest/functions/process-document.ts
import { inngest } from '../client'
 
export const processDocument = inngest.createFunction(
  { id: 'process-document', retries: 3 },
  { event: 'knowledge/document.process' },
  async ({ event, step }) => {
    const { documentId, knowledgeBaseId } = event.data
 
    // 1. 获取文档信息
    const doc = await step.run('get-document', async () => {
      const [d] = await db.select().from(documents).where(eq(documents.id, documentId))
      return d
    })
 
    // 2. 更新状态为处理中
    await step.run('update-status-processing', async () => {
      await db.update(documents)
        .set({ status: 'processing' })
        .where(eq(documents.id, documentId))
    })
 
    // 3. 解析文档
    const text = await step.run('parse-document', async () => {
      return parseDocument(doc.fileKey!, doc.mimeType!)
    })
 
    // 4. 获取知识库配置
    const kb = await step.run('get-kb-config', async () => {
      const [k] = await db.select().from(knowledgeBases)
        .where(eq(knowledgeBases.id, knowledgeBaseId))
      return k
    })
 
    // 5. 分块
    const chunks = await step.run('chunk-text', async () => {
      return splitIntoChunks(text, {
        chunkSize: (kb.chunkSize || 512) * 4,  // tokens → 字符近似
        chunkOverlap: (kb.chunkOverlap || 50) * 4,
      })
    })
 
    // 6. 批量 Embedding（每批 100 条）
    const allEmbeddings: number[][] = []
    for (let i = 0; i < chunks.length; i += 100) {
      const batch = chunks.slice(i, i + 100)
      const embeddings = await step.run(`embed-batch-${i}`, async () => {
        return embedTexts(batch)
      })
      allEmbeddings.push(...embeddings)
    }
 
    // 7. 写入 document_chunks
    await step.run('insert-chunks', async () => {
      const values = chunks.map((content, index) => ({
        documentId,
        content,
        embedding: allEmbeddings[index],
        chunkIndex: index,
        metadata: { documentName: doc.name },
      }))
      await db.insert(documentChunks).values(values)
    })
 
    // 8. 更新文档状态
    await step.run('update-status-ready', async () => {
      await db.update(documents)
        .set({ status: 'ready', chunkCount: chunks.length })
        .where(eq(documents.id, documentId))
 
      await db.update(knowledgeBases)
        .set({
          documentCount: sql`${knowledgeBases.documentCount} + 1`,
        })
        .where(eq(knowledgeBases.id, knowledgeBaseId))
    })
 
    return { chunkCount: chunks.length }
  }
)

// lib/ai/rag/retrieve.ts
import { sql } from 'drizzle-orm'
 
interface RetrievalResult {
  content: string
  similarity: number
  documentName: string
  chunkIndex: number
  documentId: string
}
 
export async function retrieveRelevantChunks(
  query: string,
  knowledgeBaseId: string,
  options: { topK?: number; threshold?: number } = {}
): Promise<RetrievalResult[]> {
  const { topK = 5, threshold = 0.7 } = options
 
  // 1. 将查询转为向量
  const queryEmbedding = await embedText(query)
 
  // 2. pgvector 余弦相似度搜索
  const results = await db.execute(sql`
    SELECT
      dc.content,
      dc.chunk_index,
      dc.document_id,
      dc.metadata->>'documentName' as document_name,
      1 - (dc.embedding <=> ${JSON.stringify(queryEmbedding)}::vector) as similarity
    FROM document_chunks dc
    JOIN documents d ON d.id = dc.document_id
    WHERE d.knowledge_base_id = ${knowledgeBaseId}
      AND d.status = 'ready'
      AND 1 - (dc.embedding <=> ${JSON.stringify(queryEmbedding)}::vector) > ${threshold}
    ORDER BY dc.embedding <=> ${JSON.stringify(queryEmbedding)}::vector
    LIMIT ${topK}
  `)
 
  return results.rows as RetrievalResult[]
}

export async function hybridRetrieve(
  query: string,
  knowledgeBaseId: string,
  options: { topK?: number } = {}
): Promise<RetrievalResult[]> {
  const { topK = 5 } = options
  const queryEmbedding = await embedText(query)
 
  // 同时执行向量搜索和全文搜索
  const results = await db.execute(sql`
    WITH vector_results AS (
      SELECT dc.id, dc.content, dc.document_id, dc.chunk_index, dc.metadata,
        1 - (dc.embedding <=> ${JSON.stringify(queryEmbedding)}::vector) as vector_score
      FROM document_chunks dc
      JOIN documents d ON d.id = dc.document_id
      WHERE d.knowledge_base_id = ${knowledgeBaseId} AND d.status = 'ready'
      ORDER BY dc.embedding <=> ${JSON.stringify(queryEmbedding)}::vector
      LIMIT ${topK * 2}
    ),
    keyword_results AS (
      SELECT dc.id, dc.content, dc.document_id, dc.chunk_index, dc.metadata,
        ts_rank(to_tsvector('english', dc.content), plainto_tsquery('english', ${query})) as keyword_score
      FROM document_chunks dc
      JOIN documents d ON d.id = dc.document_id
      WHERE d.knowledge_base_id = ${knowledgeBaseId} AND d.status = 'ready'
        AND to_tsvector('english', dc.content) @@ plainto_tsquery('english', ${query})
      LIMIT ${topK * 2}
    )
    SELECT DISTINCT ON (id)
      content, document_id, chunk_index,
      metadata->>'documentName' as document_name,
      COALESCE(v.vector_score, 0) * 0.7 + COALESCE(k.keyword_score, 0) * 0.3 as combined_score
    FROM vector_results v
    FULL OUTER JOIN keyword_results k USING (id, content, document_id, chunk_index, metadata)
    ORDER BY id, combined_score DESC
    LIMIT ${topK}
  `)
 
  return results.rows as RetrievalResult[]
}

// lib/ai/rag/query-rewrite.ts
import { generateText } from 'ai'
import { openai } from '@ai-sdk/openai'
 
export async function rewriteQuery(
  query: string,
  chatHistory?: string
): Promise<string> {
  const { text } = await generateText({
    model: openai('gpt-4o-mini'),
    system: `You are a search query optimizer. Given a user question (and optional chat history), 
rewrite it into a clear, self-contained search query. 
Rules:
- Resolve pronouns (it, this, that) using chat history
- Expand abbreviations
- Keep it concise (under 50 words)
- Return ONLY the rewritten query`,
    prompt: chatHistory
      ? `Chat history:\n${chatHistory}\n\nLatest question: ${query}`
      : query,
    maxTokens: 100,
  })
 
  return text.trim()
}

// app/api/chat/route.ts（RAG 增强版）
export async function POST(req: Request) {
  const session = await getSession()
  if (!session) return new Response('Unauthorized', { status: 401 })
 
  const tenantId = await getCurrentTenantId()
  const { messages: clientMessages, conversationId, knowledgeBaseId } = await req.json()
 
  const lastUserMsg = clientMessages[clientMessages.length - 1]
 
  let ragContext = ''
  let sources: RetrievalResult[] = []
 
  // 如果指定了知识库，执行 RAG 检索
  if (knowledgeBaseId && lastUserMsg.role === 'user') {
    const rewrittenQuery = await rewriteQuery(lastUserMsg.content)
    sources = await retrieveRelevantChunks(rewrittenQuery, knowledgeBaseId)
 
    if (sources.length > 0) {
      ragContext = sources
        .map((s, i) => `[Source ${i + 1}: ${s.documentName}]\n${s.content}`)
        .join('\n\n---\n\n')
    }
  }
 
  const systemPrompt = buildRAGSystemPrompt(ragContext)
 
  const result = streamText({
    model: openai('gpt-4o-mini'),
    system: systemPrompt,
    messages: clientMessages,
    onFinish: async ({ text, usage }) => {
      await saveAssistantMessage(conversationId, text, usage, {
        sources: sources.map(s => ({
          documentName: s.documentName,
          chunkIndex: s.chunkIndex,
          similarity: s.similarity,
        })),
      })
    },
  })
 
  // 将 sources 作为流数据附带返回
  return result.toDataStreamResponse({
    getErrorMessage: (error) => String(error),
  })
}
 
function buildRAGSystemPrompt(context: string): string {
  if (!context) {
    return 'You are a helpful assistant. If you cannot find relevant information, say so honestly.'
  }
 
  return `You are a helpful assistant that answers questions based on the provided documents.
 
IMPORTANT RULES:
- Base your answers ONLY on the provided context. Do not use outside knowledge.
- If the context doesn't contain relevant information, say "I couldn't find relevant information in the documents."
- When citing information, reference the source (e.g., "According to [Source 1]...").
- Be precise and factual. Do not speculate or infer beyond what the documents state.
 
CONTEXT:
${context}`
}

// 在流开始前，通过 data 附加 sources
import { createDataStream } from 'ai'
 
const dataStream = createDataStream({
  execute: async (writer) => {
    // 先发送 sources 数据
    writer.writeData({
      type: 'sources',
      sources: sources.map(s => ({
        documentName: s.documentName,
        similarity: Math.round(s.similarity * 100),
        preview: s.content.slice(0, 150) + '...',
      })),
    })
 
    // 然后开始流式生成
    const result = streamText({
      model: openai('gpt-4o-mini'),
      system: systemPrompt,
      messages: clientMessages,
    })
 
    result.mergeIntoDataStream(writer)
  },
})
 
return new Response(dataStream.stream, {
  headers: { 'Content-Type': 'text/event-stream' },
})

const { messages, data } = useChat({ api: '/api/chat' })
 
// data 中包含 sources 信息
const sources = data?.find(d => d.type === 'sources')?.sources

// lib/actions/knowledge.ts
 
export async function createKnowledgeBase(data: { name: string; description?: string }) {
  const { membership } = await requirePermission('ai:knowledge:write')
 
  const [kb] = await db.insert(knowledgeBases).values({
    tenantId: membership.tenantId,
    name: data.name,
    description: data.description,
  }).returning()
 
  return kb
}
 
export async function deleteKnowledgeBase(knowledgeBaseId: string) {
  const { membership } = await requirePermission('ai:knowledge:write')
 
  // 验证归属
  const [kb] = await db.select().from(knowledgeBases)
    .where(and(
      eq(knowledgeBases.id, knowledgeBaseId),
      eq(knowledgeBases.tenantId, membership.tenantId),
    ))
  if (!kb) throw new Error('Knowledge base not found')
 
  // 级联删除：documents → document_chunks
  await db.delete(knowledgeBases).where(eq(knowledgeBases.id, knowledgeBaseId))
 
  // 清理 S3 文件（异步）
  await inngest.send({
    name: 'knowledge/cleanup-files',
    data: { tenantId: membership.tenantId, knowledgeBaseId },
  })
}

// components/document-status.tsx
const STATUS_MAP = {
  pending: { label: 'Queued', color: 'bg-yellow-100 text-yellow-700' },
  processing: { label: 'Processing', color: 'bg-blue-100 text-blue-700' },
  ready: { label: 'Ready', color: 'bg-green-100 text-green-700' },
  failed: { label: 'Failed', color: 'bg-red-100 text-red-700' },
}
 
export function DocumentStatus({ status }: { status: string }) {
  const config = STATUS_MAP[status as keyof typeof STATUS_MAP]
  return (
    <span className={`inline-flex rounded-full px-2 py-0.5 text-xs font-medium ${config.color}`}>
      {config.label}
    </span>
  )
}

export async function retryDocumentProcessing(documentId: string) {
  const { membership } = await requirePermission('ai:knowledge:write')
 
  // 清除旧的 chunks
  await db.delete(documentChunks).where(eq(documentChunks.documentId, documentId))
 
  // 重置状态
  await db.update(documents)
    .set({ status: 'pending', errorMessage: null, chunkCount: 0 })
    .where(eq(documents.id, documentId))
 
  // 重新触发处理
  const [doc] = await db.select().from(documents).where(eq(documents.id, documentId))
  await inngest.send({
    name: 'knowledge/document.process',
    data: { documentId, knowledgeBaseId: doc.knowledgeBaseId },
  })
}

// scripts/eval-rag.ts
interface TestCase {
  question: string
  expectedDocumentName: string
  expectedKeywords: string[]
}
 
async function evaluateRAG(knowledgeBaseId: string, testCases: TestCase[]) {
  let hits = 0
 
  for (const tc of testCases) {
    const results = await retrieveRelevantChunks(tc.question, knowledgeBaseId)
 
    const hasCorrectDoc = results.some(r => r.documentName === tc.expectedDocumentName)
    const hasKeywords = tc.expectedKeywords.some(kw =>
      results.some(r => r.content.toLowerCase().includes(kw.toLowerCase()))
    )
 
    if (hasCorrectDoc && hasKeywords) hits++
 
    console.log(`${hasCorrectDoc && hasKeywords ? '✅' : '❌'} ${tc.question}`)
  }
 
  console.log(`\nAccuracy: ${hits}/${testCases.length} (${(hits/testCases.length*100).toFixed(1)}%)`)
}

export async function retrieveWithContext(
  query: string,
  knowledgeBaseId: string
): Promise<RetrievalResult[]> {
  const results = await retrieveRelevantChunks(query, knowledgeBaseId, { topK: 3 })
 
  // 对每个结果，获取前后各一个块
  const expanded: RetrievalResult[] = []
  for (const result of results) {
    const neighbors = await db.select()
      .from(documentChunks)
      .where(and(
        eq(documentChunks.documentId, result.documentId),
        sql`${documentChunks.chunkIndex} BETWEEN ${result.chunkIndex - 1} AND ${result.chunkIndex + 1}`,
      ))
      .orderBy(asc(documentChunks.chunkIndex))
 
    const combined = neighbors.map(n => n.content).join('\n')
    expanded.push({ ...result, content: combined })
  }
 
  return expanded
}