RAG 及其与 LangChain、Bun、Ollama 和 Qdrant 的实践实现

1// RAG 系统中数据流的简化示例
2async function processQuery(query: string): Promise<string> {
3  // 1. 将查询转换为嵌入
4  const queryEmbedding = await embedder.embedQuery(query);
5  
6  // 2. 搜索相关文档
7  const relevantDocs = await vectorStore.similaritySearch(queryEmbedding);
8  
9  // 3. 使用上下文构建提示
10  const prompt = buildPromptWithContext(query, relevantDocs);
11  
12  // 4. 使用 LLM 生成响应
13  const response = await llm.generate(prompt);
14  
15  return response;
16}

1// src/config.ts
2export const CONFIG = {
3  // Ollama 配置
4  OLLAMA_BASE_URL: 'http://localhost:11434',
5  OLLAMA_MODEL_NAME: 'llama3',
6  
7  // Qdrant 配置
8  QDRANT_URL: 'http://localhost:6333',
9  QDRANT_COLLECTION_NAME: 'code_repository',
10  
11  // 索引配置
12  CHUNK_SIZE: 1000,
13  CHUNK_OVERLAP: 200,
14  
15  // 要索引的代码库路径
16  REPOSITORY_PATH: './repository',
17};

1// src/indexer.ts
2import { DirectoryLoader } from 'langchain/document_loaders/fs/directory';
3import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
4import { OllamaEmbeddings } from '@langchain/community/embeddings/ollama';
5import { QdrantClient } from 'qdrant-js';
6import { CONFIG } from './config';
7
8export async function indexRepository() {
9  // 1. 加载存储库文件
10  const loader = new DirectoryLoader(
11    CONFIG.REPOSITORY_PATH,
12    {
13      '.ts': (path) => new TextLoader(path),
14      '.js': (path) => new TextLoader(path),
15      '.tsx': (path) => new TextLoader(path),
16      '.jsx': (path) => new TextLoader(path),
17      '.md': (path) => new TextLoader(path),
18    }
19  );
20  
21  const docs = await loader.load();
22  console.log(`加载 ${docs.length} 个文件完成`);
23  
24  // 2. 分割成块
25  const textSplitter = new RecursiveCharacterTextSplitter({
26    chunkSize: CONFIG.CHUNK_SIZE,
27    chunkOverlap: CONFIG.CHUNK_OVERLAP,
28  });
29  
30  const chunks = await textSplitter.splitDocuments(docs);
31  console.log(`创建 ${chunks.length} 个块完成`);
32  
33  // 3. 生成嵌入
34  const embeddings = new OllamaEmbeddings({
35    baseUrl: CONFIG.OLLAMA_BASE_URL,
36    model: CONFIG.OLLAMA_MODEL_NAME,
37  });
38  
39  // 4. 存储在 Qdrant 中
40  const qdrant = new QdrantClient({ url: CONFIG.QDRANT_URL });
41  
42  // 如果集合不存在则创建
43  try {
44    await qdrant.createCollection(CONFIG.QDRANT_COLLECTION_NAME, {
45      vectors: {
46        size: 1536, // 嵌入向量的大小
47        distance: 'Cosine',
48      },
49    });
50  } catch (e) {
51    console.log('集合已存在或出错:', e.message);
52  }
53  
54  // 插入文档及其嵌入
55  for (let i = 0; i < chunks.length; i += 100) {
56    const batch = chunks.slice(i, i + 100);
57    const batchEmbeddings = await embeddings.embedDocuments(
58      batch.map(chunk => chunk.pageContent)
59    );
60    
61    await qdrant.upsert(CONFIG.QDRANT_COLLECTION_NAME, {
62      points: batchEmbeddings.map((vector, idx) => ({
63        id: `${i + idx}`,
64        vector,
65        payload: {
66          text: batch[idx].pageContent,
67          metadata: batch[idx].metadata,
68        },
69      })),
70    });
71    
72    console.log(`处理批次 ${i}-${i + batch.length} 完成`);
73  }
74  
75  console.log('索引成功完成！');
76}

1// src/rag.ts
2import { OllamaEmbeddings } from '@langchain/community/embeddings/ollama';
3import { QdrantClient } from 'qdrant-js';
4import { Ollama } from '@langchain/community/llms/ollama';
5import { CONFIG } from './config';
6
7export async function generateResponse(query: string): Promise<string> {
8  // 1. 初始化组件
9  const embeddings = new OllamaEmbeddings({
10    baseUrl: CONFIG.OLLAMA_BASE_URL,
11    model: CONFIG.OLLAMA_MODEL_NAME,
12  });
13  
14  const qdrant = new QdrantClient({ url: CONFIG.QDRANT_URL });
15  
16  const llm = new Ollama({
17    baseUrl: CONFIG.OLLAMA_BASE_URL,
18    model: CONFIG.OLLAMA_MODEL_NAME,
19  });
20  
21  // 2. 生成查询嵌入
22  const queryEmbedding = await embeddings.embedQuery(query);
23  
24  // 3. 搜索相关文档
25  const searchResults = await qdrant.search(CONFIG.QDRANT_COLLECTION_NAME, {
26    vector: queryEmbedding,
27    limit: 5,
28  });
29  
30  const relevantDocs = searchResults.map(result => result.payload.text).join('\n\n');
31  
32  // 4. 使用上下文构建提示
33  const prompt = `
34  您是一位专家编程助手，帮助理解项目的源代码。
35  使用以下上下文来回答用户的问题。
36  如果您在上下文中找不到信息，请明确说明。
37  
38  上下文:
39  ${relevantDocs}
40  
41  问题: ${query}
42  
43  回答:
44  `;
45  
46  // 5. 生成响应
47  const response = await llm.call(prompt);
48  
49  return response;
50}

1// src/index.ts
2import { indexRepository } from './indexer';
3import { generateResponse } from './rag';
4
5async function main() {
6  // 用于索引或查询的参数
7  const args = process.argv.slice(2);
8  
9  if (args[0] === 'index') {
10    console.log('开始索引...');
11    await indexRepository();
12  } else if (args[0] === 'query') {
13    const query = args.slice(1).join(' ');
14    if (!query) {
15      console.log('用法: bun run src/index.ts query "您关于代码的问题"');
16      process.exit(1);
17    }
18    
19    console.log(`问题: ${query}`);
20    console.log('正在生成响应...');
21    const response = await generateResponse(query);
22    console.log('\n回答:');
23    console.log(response);
24  } else {
25    console.log('用法: bun run src/index.ts [index|query "您的问题"]');
26  }
27}
28
29main().catch(console.error);

1// 更高级提示的示例
2const prompt = `
3# 角色
4您是一位专门从事源代码分析的专家编程助手。
5
6# 上下文
7${relevantDocs}
8
9# 说明
10- 仅根据提供的上下文回答用户的问题
11- 如果上下文不包含必要的信息，请明确指出
12- 在您的回答中引用相关的代码摘录
13- 以教学的方式解释代码
14- 如果发现任何问题，请提出改进建议
15
16# 问题
17${query}
18`;