如何用LlamaIndex加载大CSV文件适配OpenAI API查询？

问题背景 #

用create-llama搭建的应用，SimpleDirectoryReader().loadData可正常处理100多页PDF，但加载含17566个token的大型CSV时，触发以下错误：

BadRequestError: 400 This model's maximum context length is 8192 tokens, however you requested 17566 tokens (17566 in your prompt; 0 for the completion). Please reduce your prompt; or completion length.
at Function.generate (C:\Code\canoe\node_modules\openai\src\error.ts:70:14)
at OpenAI.makeStatusError (C:\Code\canoe\node_modules\openai\src\core.ts:383:21)
at OpenAI.makeRequest (C:\Code\canoe\node_modules\openai\src\core.ts:446:24)
at process.processTicksAndRejections (node:internal/process/task_queues:95:5)
at async OpenAIEmbedding.getOpenAIEmbedding (C:\Code\canoe\node_modules\llamaindex\dist\cjs\embeddings\OpenAIEmbedding.js:100:26)
at async OpenAIEmbedding.getTextEmbeddings (C:\Code\canoe\node_modules\llamaindex\dist\cjs\embeddings\OpenAIEmbedding.js:111:16)
at async batchEmbeddings (C:\Code\canoe\node_modules\llamaindex\dist\cjs\embeddings\types.js:61:32)
at async OpenAIEmbedding.getTextEmbeddingsBatch (C:\Code\canoe\node_modules\llamaindex\dist\cjs\embeddings\types.js:43:16)
at async OpenAIEmbedding.transform (C:\Code\canoe\node_modules\llamaindex\dist\cjs\embeddings\types.js:47:28)
at async VectorStoreIndex.getNodeEmbeddingResults (C:\Code\canoe\node_modules\llamaindex\dist\cjs\indices\vectorStore\index.js:487:9)
at async VectorStoreIndex.insertNodes (C:\Code\canoe\node_modules\llamaindex\dist\cjs\indices\vectorStore\index.js:572:17)
at async VectorStoreIndex.buildIndexFromNodes (C:\Code\canoe\node_modules\llamaindex\dist\cjs\indices\vectorStore\index.js:497:9)
at async VectorStoreIndex.init (C:\Code\canoe\node_modules\llamaindex\dist\cjs\indices\vectorStore\index.js:445:13)
at async VectorStoreIndex.fromDocuments (C:\Code\canoe\node_modules\llamaindex\dist\cjs\indices\vectorStore\index.js:523:16)
at <anonymous> (C:\Code\canoe\app\api\chat\engine\generate.ts:27:5)
at getRuntime (C:\Code\canoe\app\api\chat\engine\generate.ts:14:3) {

核心失败代码：

async function generateDatasource() {
  console.log(`Generating storage context...`);
  // Split documents, create embeddings and store them in the storage context
  const ms = await getRuntime(async () => {
    const storageContext = await storageContextFromDefaults({
      persistDir: STORAGE_CACHE_DIR,
    });
    const documents = await getDocuments();
    await VectorStoreIndex.fromDocuments(documents, {
      storageContext,
    });
  });
  console.log(`Storage context successfully generated in ${ms / 1000}s.`);
}

.env配置：

MODEL=gpt-4-turbo 
EMBEDDING_MODEL=text-embedding-3-large
EMBEDDING_DIM=1024

解决方案 #

1. 创建无错误的文档存储、索引存储和向量存储 #

问题根源是默认读取方式将整个CSV文件作为单个Document处理，导致生成embedding时token数超限制。需通过文档拆分解决：

方式一：使用文本拆分器自动拆分

修改generateDatasource函数，在创建索引时指定文本拆分规则，确保每个chunk的token数低于embedding模型限制（text-embedding-3-large最大支持8192 token，建议留余量设为4000）：

import { RecursiveCharacterTextSplitter } from "llamaindex";

async function generateDatasource() {
  console.log(`Generating storage context...`);
  const ms = await getRuntime(async () => {
    const storageContext = await storageContextFromDefaults({
      persistDir: STORAGE_CACHE_DIR,
    });
    const documents = await getDocuments();
    // 添加文本拆分器，拆分大型CSV为小chunk
    await VectorStoreIndex.fromDocuments(documents, {
      storageContext,
      transformations: [
        new RecursiveCharacterTextSplitter({
          chunkSize: 4000, // 每个chunk的token数
          chunkOverlap: 200, // 重叠部分，保持上下文连贯性
        }),
      ],
    });
  });
  console.log(`Storage context successfully generated in ${ms / 1000}s.`);
}

方式二：使用CSVReader针对性读取

如果CSV是结构化表格，用CSVReader替代默认读取器，可按行或自定义规则拆分文档：

import { CSVReader, RecursiveCharacterTextSplitter } from "llamaindex";

async function getDocuments() {
  // 初始化CSV读取器，支持配置表头、分隔符等
  const csvReader = new CSVReader({
    headerRow: true,
    separator: ",",
  });
  // 读取CSV文件，返回Document数组
  const documents = await csvReader.loadData("./path/to/large.csv");
  // 配合拆分器进一步拆分
  const splitter = new RecursiveCharacterTextSplitter({
    chunkSize: 4000,
    chunkOverlap: 200,
  });
  return splitter.splitDocuments(documents);
}

2. 在token限制内查询CSV数据 #

针对大型CSV的查询，需优化召回逻辑，避免超出LLM的token限制：

方式一：优化向量查询引擎

构建查询引擎时，控制召回的chunk数量，并设置压缩模式减少token使用：

// 创建索引后构建查询引擎
const index = await VectorStoreIndex.fromDocuments(splitDocuments, { storageContext });
const queryEngine = index.asQueryEngine({
  similarityTopK: 5, // 召回前5个最相关的chunk，根据需求调整
  responseMode: "compact", // 压缩召回内容，减少输入token数
});

// 执行查询
const response = await queryEngine.query("请统计某列的平均值");
console.log(response.response);

方式二：使用PandasQueryEngine处理结构化数据

如果CSV是标准表格，用PandasQueryEngine直接基于结构化数据查询，效率更高且无需拆分文本：

import { PandasQueryEngine, OpenAI } from "llamaindex";
import pandas from "pandas-js";

// 将CSV转为DataFrame
const df = pandas.read_csv("./path/to/large.csv");

// 初始化Pandas查询引擎
const queryEngine = new PandasQueryEngine({
  pandasDF: df,
  llm: new OpenAI({ model: "gpt-4-turbo" }),
});

// 执行SQL-like查询
const response = await queryEngine.query("计算所有行中'销售额'列的总和");
console.log(response.response);

内容的提问来源于stack exchange，提问作者TheFastCat