Compare commits

...

4 Commits

Author SHA1 Message Date
Clelia (Astra) Bertelli 6195948ae1 chore: implement suggestions 2025-08-07 15:21:17 +02:00
Clelia (Astra) Bertelli 93e22a9b4f fix: loop logic to fix test 🙈 2025-08-07 14:21:47 +02:00
Clelia (Astra) Bertelli c7a944a9d1 feat: add tests 2025-08-07 14:07:51 +02:00
Clelia (Astra) Bertelli 8ec1add002 feat: add parse and getTables methods to LlamaParseReader 2025-08-07 13:54:53 +02:00
4 changed files with 182 additions and 6 deletions
+81 -6
View File
@@ -14,7 +14,8 @@ import {
getJobTextResultApiV1ParsingJobJobIdResultTextGet,
uploadFileApiV1ParsingUploadPost,
} from "./api";
import { sleep } from "./utils";
import { sleep, getSavePath } from "./utils";
import type { ParseResult } from "./type";
export type Language = ParserLanguages;
export type ResultType = "text" | "markdown" | "json";
@@ -594,15 +595,15 @@ export class LlamaParseReader extends FileReader {
}
/**
* Loads data from a file and returns an array of JSON objects.
* Loads data from a file and returns an array of ParseResult objects.
* To be used with resultType "json".
*
* @param filePathOrContent - The file path or the file content as a Uint8Array.
* @returns A Promise that resolves to an array of JSON objects.
* @returns A Promise that resolves to an array of ParseResult objects, encapsulating the JSON array resulting from the parsing within the pages attribute.
*/
async loadJson(
filePathOrContent: string | Uint8Array,
): Promise<Record<string, any>[]> {
): Promise<ParseResult[]> {
let jobId;
const isFilePath =
typeof filePathOrContent === "string" &&
@@ -628,7 +629,7 @@ export class LlamaParseReader extends FileReader {
const resultJson = await this.getJobResult(jobId, "json");
resultJson.job_id = jobId;
resultJson.file_path = isFilePath ? filePathOrContent : undefined;
return [resultJson];
return [resultJson] as ParseResult[];
} catch (e) {
console.error(`Error while parsing the file under job id ${jobId}`, e);
if (this.ignoreErrors) {
@@ -639,6 +640,80 @@ export class LlamaParseReader extends FileReader {
}
}
/**
* Loads data from a file or file bytes (or an array of those) and returns an array of ParseResult objects.
*
* @param filePathOrContent - The file path or the file content as a Uint8Array, or an array of one of those two types.
* @returns A Promise that resolves to an array of ParseResult objects, encapsulating the JSON array resulting from the parsing within the pages attribute.
*/
async parse(
filePathOrContent: string | Uint8Array | string[] | Uint8Array[],
): Promise<ParseResult[]> {
const jsonResults: Record<string, any>[][] = [];
if (!Array.isArray(filePathOrContent)) {
const jsonResult = await this.loadJson(filePathOrContent);
jsonResults.push(jsonResult);
} else {
for (let i = 0; i < filePathOrContent.length; i++) {
console.log(
`Processing file ${i + 1} of ${filePathOrContent.length}...`,
);
const jsonResult = await this.loadJson(
filePathOrContent[i] as string | Uint8Array,
);
jsonResults.push(jsonResult);
}
}
const parseResults: ParseResult[] = [];
for (const jsonResult of jsonResults) {
for (const result of jsonResult) {
const parseResult = {
pages: result.pages,
job_metadata: result.job_metadata,
job_id: result.job_id,
file_path: result?.file_path ?? "",
is_completed: true,
} as ParseResult;
parseResults.push(parseResult);
}
}
return parseResults;
}
/**
* Downloads and saves tables from a given array of ParseResult to a specified download path.
*
* @param jsonResults - The array of ParseResult containing table information.
* @param downloadPath - The path where the downloaded tables will be saved as CSV files.
* @returns A Promise that resolves to an array of strings representing the paths to the tables.
*/
async getTables(
jsonResults: ParseResult[],
downloadPath: string,
): Promise<string[]> {
const tables: string[] = [];
for (const result of jsonResults) {
for (const page of result.pages) {
if ("items" in page && Array.isArray(page.items)) {
for (let i = 0; i < page.items.length; i++) {
if (
"type" in page.items[i] &&
page.items[i].type === "table" &&
"csv" in page.items[i] &&
typeof page.items[i].csv === "string" &&
page.items[i].csv != ""
) {
const savePath = getSavePath(downloadPath, i);
await fs.writeFile(savePath, page.items[i].csv);
tables.push(savePath);
}
}
}
}
}
return tables;
}
/**
* Downloads and saves images from a given JSON result to a specified download path.
* Currently only supports resultType "json".
@@ -648,7 +723,7 @@ export class LlamaParseReader extends FileReader {
* @returns A Promise that resolves to an array of image objects.
*/
async getImages(
jsonResult: Record<string, any>[],
jsonResult: ParseResult[],
downloadPath: string,
): Promise<Record<string, any>[]> {
try {
+10
View File
@@ -49,3 +49,13 @@ export type ExtractResult = {
| null;
};
};
export type ParseResult = {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
pages: Record<string, any>[];
// eslint-disable-next-line @typescript-eslint/no-explicit-any
job_metadata: Record<string, any>;
job_id: string;
is_completed: boolean;
file_path: string;
};
+18
View File
@@ -6,6 +6,8 @@ import {
import { DEFAULT_BASE_URL } from "@llamaindex/core/global";
import { getEnv } from "@llamaindex/env";
import type { ClientParams } from "./type.js";
import fs from "fs";
import path from "path";
function getBaseUrl(baseUrl?: string): string {
return baseUrl ?? getEnv("LLAMA_CLOUD_BASE_URL") ?? DEFAULT_BASE_URL;
@@ -99,3 +101,19 @@ export async function getPipelineId(
export async function sleep(ms: number): Promise<void> {
return new Promise((resolve) => setTimeout(resolve, ms));
}
export function getSavePath(downloadPath: string, i: number): string {
const now = new Date();
const formatted =
now.toISOString().replace(/[-:T]/g, "_").replace(/\..+/, "") +
"_" +
now.getMilliseconds().toString().padStart(3, "0");
let savePath = path.join(downloadPath, `table_${formatted}.csv`);
if (fs.existsSync(savePath)) {
savePath = savePath.replace(".csv", "_") + i.toString() + ".csv";
}
return savePath;
}
@@ -5,6 +5,7 @@ import { LlamaExtract, LlamaExtractAgent } from "../src/LlamaExtract.js";
import { Document } from "@llamaindex/core/schema";
import { fs } from "@llamaindex/env";
import { ExtractConfig } from "../src/api.js";
import { ParseResult } from "../src/type.js";
// Integration tests that require actual API keys and files
describe("Integration Tests", () => {
@@ -284,6 +285,78 @@ describe("Integration Tests", () => {
},
60000,
);
it.skipIf(skipIfNoApiKey)(
"should parse a file and return a ParseResult array",
async () => {
const parseReader = new LlamaParseReader({
apiKey: process.env.LLAMA_CLOUD_API_KEY!,
verbose: false,
});
const testContent = "Test document for JSON parsing";
const testFilePath = "test-json.txt";
await fs.writeFile(testFilePath, new TextEncoder().encode(testContent));
try {
const result = await parseReader.parse(testFilePath);
expect(result).toBeDefined();
expect(Array.isArray(result)).toBe(true);
expect(result.length).toBeGreaterThan(0);
expect(result[0]).toHaveProperty("job_id");
expect(result[0]).toHaveProperty("job_metadata");
expect(result[0]).toHaveProperty("file_path");
expect(result[0]).toHaveProperty("pages");
await fs.unlink(testFilePath);
} catch (error) {
try {
await fs.unlink(testFilePath);
} catch {}
throw error;
}
},
60000,
);
it.skipIf(skipIfNoApiKey)(
"should extract tables correctly from a JSON result",
async () => {
const parseReader = new LlamaParseReader({
apiKey: process.env.LLAMA_CLOUD_API_KEY!,
verbose: false,
});
const pseudoJsonResult = [
{
pages: [
{
items: [
{
type: "table",
csv: "Name,Age,Height (cm)\nAnna,12,140\nBob,22,175\nClaire,33,173\nDenis,44,185\n",
},
],
},
],
job_id: "jobId",
job_metadata: { job_id: "jobId" },
file_path: "table.csv",
is_completed: true,
},
] as ParseResult[];
const tmpdir = await fs.mkdtemp("tables");
const result = await parseReader.getTables(pseudoJsonResult, tmpdir);
expect(result).toBeDefined();
expect(Array.isArray(result)).toBe(true);
expect(result.length).toBeGreaterThan(0);
expect(typeof result[0] === "string").toBe(true);
},
60000,
);
});
describe("LlamaCloudIndex Integration", () => {