AI RAG 前端集成实战：从向量检索到流式问答的全链路实现

1
// 工具：调用不同 Embedding API（伪代码）
2
async function embed(text: string, model: 'bge-large' | 'text-embedding-3-small' | 'm3e-base') {
3
  switch (model) {
4
    case 'bge-large':
5
      return await fetch('http://localhost:8000/embed', {
6
        method: 'POST',
7
        body: JSON.stringify({ text })
8
      }).then(r => r.json());
9

10
    case 'text-embedding-3-small':
11
      return await openai.embeddings.create({
12
        model: 'text-embedding-3-small',
13
        input: text
14
      });
15

16
    case 'm3e-base':
17
      return await fetch('http://localhost:8001/embed', {
18
        method: 'POST',
19
        body: JSON.stringify({ text })
20
      }).then(r => r.json());
21
  }
22
}

1
// 混合检索的简易实现
2
async function hybridSearch(query: string, k = 5) {
3
  // 1. 向量召回 Top 2K
4
  const vectorResults = await vectorDB.search(query, k * 4);
5

6
  // 2. BM25 关键词召回 Top 2K
7
  const bm25Results = await bm25Index.search(query, k * 4);
8

9
  // 3. 倒数排名融合 (Reciprocal Rank Fusion)
10
  const fused = reciprocalRankFusion(
11
    vectorResults,
12
    bm25Results,
13
    weights: [0.7, 0.3]  // 向量权重更高
14
  );
15

16
  return fused.slice(0, k);
17
}
18

19
function reciprocalRankFusion(...lists) {
20
  const scores = new Map<string, number>();
21
  for (const [weight, list] of lists) {
22
    list.forEach((item, rank) => {
23
      const rr = weight / (60 + rank + 1);  // k 常数 = 60
24
      scores.set(item.id, (scores.get(item.id) || 0) + rr);
25
    });
26
  }
27
  return [...scores.entries()]
28
    .sort((a, b) => b[1] - a[1])
29
    .map(([id]) => id);
30
}

1
// chunker.ts - 工业级 Markdown 分块器
2
import { unified } from 'unified';
3
import remarkParse from 'remark-parse';
4
import { visit } from 'unist-util-visit';
5
import mdast from 'mdast';
6

7
interface Chunk {
8
  content: string;
9
  heading: string;
10
  level: number;
11
  index: number;
12
  source: string;
13
  tokenCount: number;
14
}
15

16
const MAX_CHUNK_TOKENS = 500;  // 推荐 200~500
17
const MIN_CHUNK_TOKENS = 50;   // 太小的 chunk 没意义
18

19
export function chunkMarkdown(
20
  markdown: string,
21
  source: string
22
): Chunk[] {
23
  const tree = unified().use(remarkParse).parse(markdown);
24
  const chunks: Chunk[] = [];
25
  let currentChunk: string[] = [];
26
  let currentHeading = '';
27
  let currentLevel = 0;
28
  let chunkIndex = 0;
29

30
  function flushChunk() {
31
    const content = currentChunk.join('\n\n').trim();
32
    const tokens = estimateTokens(content);
33
    if (tokens >= MIN_CHUNK_TOKENS) {
34
      chunks.push({
35
        content,
36
        heading: currentHeading,
37
        level: currentLevel,
38
        index: chunkIndex++,
39
        source,
40
        tokenCount: tokens,
41
      });
42
    }
43
    currentChunk = [];
44
  }
45

46
  visit(tree, (node) => {
47
    // 遇到新的标题，刷新当前 chunk
48
    if (node.type === 'heading') {
49
      flushChunk();
50
      currentHeading = mdastToString(node);
51
      currentLevel = node.depth;
52
    } else {
53
      const text = mdastToString(node);
54
      if (text) {
55
        currentChunk.push(text);
56

57
        // 超过 max tokens 就切分
58
        if (estimateTokens(currentChunk.join('\n\n')) > MAX_CHUNK_TOKENS) {
59
          flushChunk();
60
        }
61
      }
62
    }
63
  });
64

65
  flushChunk();
66
  return chunks;
67
}
68

69
// 简单的 token 估算（实际项目用 gpt-tokenizer）
70
function estimateTokens(text: string): number {
71
  return Math.ceil(text.length / 1.5);  // 中文 1.5 字/token
72
}
73

74
function mdastToString(node: any): string {
75
  if (node.value) return node.value;
76
  if (node.children) return node.children.map(mdastToString).join('');
77
  return '';
78
}

1
// 客户端：EventSource 或 fetch + ReadableStream
2
async function* streamChat(prompt: string, history: Message[]) {
3
  const response = await fetch('/api/rag/chat', {
4
    method: 'POST',
5
    headers: {
6
      'Content-Type': 'application/json',
7
      'Accept': 'text/event-stream',  // 关键
8
    },
9
    body: JSON.stringify({ prompt, history }),
10
  });
11

12
  if (!response.body) throw new Error('No response body');
13

14
  const reader = response.body.getReader();
15
  const decoder = new TextDecoder();
16

17
  while (true) {
18
    const { done, value } = await reader.read();
19
    if (done) break;
20

21
    const chunk = decoder.decode(value);
22
    const lines = chunk.split('\n').filter(l => l.startsWith('data:'));
23

24
    for (const line of lines) {
25
      const data = line.slice(5).trim();
26
      if (data === '[DONE]') return;
27
      try {
28
        yield JSON.parse(data);
29
      } catch {}
30
    }
31
  }
32
}

1
// api/rag/chat.ts (Next.js API Route)
2
import { OpenAIStream, StreamingTextResponse } from 'ai';
3
import { retrieveContext } from '@/lib/rag';
4
import { checkRateLimit } from '@/lib/rate-limit';
5

6
export const runtime = 'edge';  // Edge Runtime，低延迟
7

8
export async function POST(req: Request) {
9
  const { prompt, history = [] } = await req.json();
10

11
  // 1. 限流（防止滥用和超支）
12
  const ip = req.headers.get('x-forwarded-for') || 'unknown';
13
  if (!await checkRateLimit(ip)) {
14
    return new Response('Too Many Requests', { status: 429 });
15
  }
16

17
  // 2. 检索上下文（通常 100-300ms）
18
  const t0 = performance.now();
19
  const context = await retrieveContext(prompt, { topK: 5 });
20
  const retrievalMs = performance.now() - t0;
21

22
  // 3. 构造 Prompt
23
  const systemPrompt = `你是一个专业的助手。请基于以下参考资料回答用户问题。
24
如果参考资料中没有答案，请明确告知"根据现有资料无法回答"。
25
回答时请在引用的事实后用 [1][2] 这样的标记标注来源编号。`;
26

27
  const userPrompt = `参考资料：
28
${context.map((c, i) => `[${i + 1}] ${c.content}\n来源: ${c.source}`).join('\n\n')}
29

30
用户问题：${prompt}
31

32
请用中文回答。`;
33

34
  // 4. 调用 LLM（流式）
35
  const response = await openai.chat.completions.create({
36
    model: 'gpt-4o',
37
    stream: true,
38
    temperature: 0.3,
39
    messages: [
40
      { role: 'system', content: systemPrompt },
41
      ...history,
42
      { role: 'user', content: userPrompt },
43
    ],
44
  });
45

46
  // 5. 包装成 SSE（Vercel AI SDK 帮你做了）
47
  const stream = OpenAIStream(response, {
48
    onFinal: async (completion) => {
49
      // 记录日志、做埋点
50
      await logChat({
51
        prompt,
52
        completion,
53
        retrievalMs,
54
        totalSources: context.length,
55
        sources: context.map(c => c.source),
56
      });
57
    },
58
  });
59

60
  return new StreamingTextResponse(stream, {
61
    headers: {
62
      'X-Retrieval-Ms': retrievalMs.toString(),  // 自定义 header，前端可以读
63
    },
64
  });
65
}

1
// 1. 检索时返回的 chunk 要带位置信息
2
interface RetrievedChunk {
3
  id: string;
4
  content: string;
5
  source: string;
6
  sourceUrl: string;
7
  startOffset: number;  // 在原文中的起始位置
8
  endOffset: number;
9
  score: number;
10
}
11

12
// 2. 让 LLM 输出带引用的答案
13
const prompt = `...
14
回答时，请在引用的事实后用 <<<1>>>、<<<2>>> 这样的标记标注来源编号。
15
示例：员工试用期辞退需要看是否符合录用条件<<<1>>>。
16
...
17
`;
18

19
// 3. 客户端解析引用标记并渲染
20
function renderAnswerWithCitations(text: string, sources: RetrievedChunk[]) {
21
  const parts = text.split(/(<<<\d+>>>)/g);
22

23
  return parts.map((part, i) => {
24
    const match = part.match(/<<<(\d+)>>>/);
25
    if (match) {
26
      const sourceIdx = parseInt(match[1]) - 1;
27
      const source = sources[sourceIdx];
28
      if (!source) return null;
29
      return (
30
        <sup key={i} className="citation">
31
          [{sourceIdx + 1}]
32
          <div className="citation-popup">
33
            <div className="source-title">{source.source}</div>
34
            <div className="source-content">{source.content.slice(0, 200)}...</div>
35
            <a href={source.sourceUrl} target="_blank">查看原文 →</a>
36
          </div>
37
        </sup>
38
      );
39
    }
40
    return <span key={i}>{part}</span>;
41
  });
42
}

1
// ❌ 错误：一次处理所有文档
2
const allDocs = await loadAllDocs();  // 10 万篇
3
const chunks = chunkAll(allDocs);  // 500 万 chunks
4
const vectors = await embedAll(chunks);  // 烧光 API 额度
5
await qdrant.upsert(vectors);  // OOM
6

7
// ✅ 正确：分批 + 异步队列
8
import { Queue } from 'bullmq';
9

10
const embeddingQueue = new Queue('embed-docs', {
11
  connection: { host: 'localhost', port: 6379 }
12
});
13

14
for (const doc of docs) {
15
  await embeddingQueue.add('embed', { docId: doc.id }, {
16
    attempts: 3,
17
    backoff: { type: 'exponential', delay: 1000 },
18
  });
19
}

1
// ❌ 错误：每次都重新 Embedding
2
async function getContext(prompt: string) {
3
  const vec = await embed(prompt);  // 每次都花钱
4
  return await qdrant.search(COLLECTION, vec);
5
}
6

7
// ✅ 正确：语义哈希 + LRU 缓存
8
import { LRUCache } from 'lru-cache';
9
import crypto from 'crypto';
10

11
const cache = new LRUCache<string, number[]>({ max: 1000 });
12

13
function semanticHash(text: string): string {
14
  // 归一化 + 哈希，让"你好"和"你好！" 命中同一缓存
15
  const normalized = text.toLowerCase().replace(/[^\w\u4e00-\u9fa5]/g, '');
16
  return crypto.createHash('md5').update(normalized).digest('hex');
17
}
18

19
async function getContext(prompt: string) {
20
  const hash = semanticHash(prompt);
21
  let vec = cache.get(hash);
22
  if (!vec) {
23
    vec = await embed(prompt);
24
    cache.set(hash, vec);
25
  }
26
  return await qdrant.search(COLLECTION, vec);
27
}

1
// 用 LLM-as-a-Judge 做自动评估
2
async function evaluateRAG(question: string, answer: string, contexts: RetrievedChunk[]) {
3
  const evalPrompt = `你是一个严格的 RAG 评估员。
4
请评估以下回答的质量：
5

6
问题：${question}
7
回答：${answer}
8
参考资料：
9
${contexts.map((c, i) => `[${i+1}] ${c.content}`).join('\n')}
10

11
请从三个维度打分（0-10）：
12
1. 准确性：回答是否符合参考资料
13
2. 完整性：是否充分回答了问题
14
3. 引用质量：是否正确标注了引用
15

16
返回 JSON：
17
{ "accuracy": 0, "completeness": 0, "citation": 0, "reason": "..." }`;
18

19
  const result = await openai.chat.completions.create({
20
    model: 'gpt-4o',
21
    response_format: { type: 'json_object' },
22
    messages: [{ role: 'user', content: evalPrompt }],
23
  });
24

25
  return JSON.parse(result.choices[0].message.content);
26
}

1
// ✅ 必要的：点赞/点踩按钮 + 反馈原因
2
async function submitFeedback(chatId: string, rating: 'up' | 'down', reason?: string) {
3
  await db.feedback.insert({
4
    chatId,
5
    rating,
6
    reason,
7
    timestamp: Date.now(),
8
    // 同时记录：用户问了什么、检索到了什么、LLM 回了什么
9
    // 这些数据是后续优化 RAG 的金矿
10
  });
11

12
  // 点踩时自动触发人工审核或专家回答
13
  if (rating === 'down') {
14
    await notifyHumanReview(chatId);
15
  }
16
}