feat: add support for llamaextract (untested)

packages/llamaindex/package.json

@@ -28,7 +28,8 @@
     "@types/lodash": "^4.17.7",
     "@types/node": "^24.0.13",
     "lodash": "^4.17.21",
-    "magic-bytes.js": "^1.10.0"
+    "magic-bytes.js": "^1.10.0",
+    "zod": "^3.25.76"
   },
   "devDependencies": {
     "@types/node": "^22.9.0",

packages/llamaindex/src/cloud/LlamaExtract/index.ts

@@ -0,0 +1,87 @@
+import * as z from "zod/v4";
+import type { CreateAgentRequest, ExtractionResult } from "./interfaces.js";
+import { extractDataFromFile } from "./utils.js";
+
+class LlamaExtract {
+  apiKey: string;
+
+  constructor(apiKey: string) {
+    this.apiKey = apiKey;
+  }
+
+  getAgent(agentId: string): LlamaExtractAgent {
+    return new LlamaExtractAgent(this.apiKey, agentId);
+  }
+
+  async createAgent(
+    agentName: string,
+    dataSchema: z.ZodType,
+    returnId: boolean = true,
+  ): Promise<string | LlamaExtractAgent | null> {
+    const jsonSchema = z.toJSONSchema(dataSchema) as CreateAgentRequest;
+    jsonSchema.name = agentName;
+    const response = await fetch(
+      "https://api.cloud.llamaindex.ai/api/v1/extraction/extraction-agents",
+      {
+        method: "POST",
+        headers: {
+          Authorization: `Bearer ${this.apiKey}`,
+          accept: "application/json",
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify(jsonSchema),
+      },
+    );
+
+    if (!response.ok) {
+      console.log(
+        `Failed to create agent: ${response.status} ${response.statusText}`,
+      );
+      return null;
+    }
+
+    const jsonResponse = await response.json();
+
+    if (returnId) {
+      return jsonResponse["id"] as string;
+    } else {
+      return new LlamaExtractAgent(this.apiKey, jsonResponse["id"] as string);
+    }
+  }
+}
+
+class LlamaExtractAgent {
+  apiKey: string;
+  agentId: string;
+
+  constructor(apiKey: string, agentId: string) {
+    this.apiKey = apiKey;
+    this.agentId = agentId;
+  }
+
+  async extract(
+    filePath: string,
+    fileName?: string,
+    pollInterval: number = 2000,
+    maxRetries: number = 150,
+  ): Promise<ExtractionResult | null> {
+    try {
+      return extractDataFromFile(
+        this.apiKey,
+        this.agentId,
+        filePath,
+        fileName,
+        pollInterval,
+        maxRetries,
+      );
+    } catch (error) {
+      console.log(
+        "An error has occurred while extracting data from your file",
+        error,
+      );
+      return null;
+    }
+  }
+}
+
+export { LlamaExtract, LlamaExtractAgent };

packages/llamaindex/src/cloud/LlamaExtract/interfaces.ts

@@ -0,0 +1,31 @@
+export interface UploadResponse {
+  id: string;
+  name: string;
+  status: string;
+}
+
+export interface ExtractionJobResponse {
+  id: string;
+  status: string;
+  extraction_agent_id: string;
+  file_id: string;
+}
+
+export interface JobStatusResponse {
+  id: string;
+  status: "PENDING" | "RUNNING" | "SUCCESS" | "FAILED";
+  extraction_agent_id: string;
+  file_id: string;
+  created_at: string;
+  updated_at: string;
+}
+
+export interface ExtractionResult {
+  // The actual structure will depend on your extraction agent configuration
+  [key: string]: unknown;
+}
+
+export interface CreateAgentRequest {
+  name: string;
+  [key: string]: unknown; // for the JSON schema properties
+}

packages/llamaindex/src/cloud/LlamaExtract/utils.ts

@@ -0,0 +1,233 @@
+import fs from "fs";
+import path from "path";
+import type {
+  ExtractionJobResponse,
+  ExtractionResult,
+  JobStatusResponse,
+  UploadResponse,
+} from "./interfaces.js";
+
+/**
+ * Extracts data from a file using LlamaIndex Cloud
+ * @param apiKey - Your LlamaIndex Cloud API key
+ * @param agentId - The ID of your extraction agent
+ * @param filePath - Path to the file to extract data from
+ * @param fileName - Name of the file (optional, will be inferred from path if not provided)
+ * @param pollInterval - How often to poll for job completion in milliseconds (default: 2000)
+ * @param maxRetries - Maximum number of polling attempts (default: 150, ~5 minutes)
+ * @param returnAsJson - Whether to return the result as a JSON string (default: false)
+ * @returns Promise that resolves to the extracted data (object or JSON string)
+ */
+async function extractDataFromFile(
+  apiKey: string,
+  agentId: string,
+  filePath: string,
+  fileName?: string,
+  pollInterval: number = 2000,
+  maxRetries: number = 150,
+): Promise<ExtractionResult> {
+  // Step 1: Upload the file
+  console.log("Uploading file...");
+  const fileId = await uploadFile(apiKey, filePath, fileName);
+  console.log(`File uploaded with ID: ${fileId}`);
+
+  // Step 2: Run extraction job
+  console.log("Starting extraction job...");
+  const jobId = await runExtractionJob(apiKey, agentId, fileId);
+  console.log(`Extraction job started with ID: ${jobId}`);
+
+  // Step 3: Poll for job completion
+  console.log("Polling for job completion...");
+  await pollForJobCompletion(apiKey, jobId, pollInterval, maxRetries);
+  console.log("Job completed successfully!");
+
+  // Step 4: Get results
+  console.log("Retrieving extraction results...");
+  const results = await getExtractionResults(apiKey, jobId);
+  console.log("Extraction completed successfully!");
+
+  return results;
+}
+
+/**
+ * Uploads a file to LlamaIndex Cloud
+ */
+async function uploadFile(
+  apiKey: string,
+  filePath: string,
+  fileName?: string,
+): Promise<string> {
+  // Check if file exists
+  if (!fs.existsSync(filePath)) {
+    throw new Error(`File not found: ${filePath}`);
+  }
+
+  // Read the file
+  const fileBuffer = fs.readFileSync(filePath);
+  const finalFileName = fileName || path.basename(filePath);
+
+  // Create FormData
+  const formData = new FormData();
+
+  // Convert Node.js Buffer to Uint8Array (Blob-compatible)
+  const uint8Array = new Uint8Array(fileBuffer);
+
+  // Create a Blob using Uint8Array
+  const fileBlob = new Blob([uint8Array], {
+    type: getContentType(finalFileName),
+  });
+
+  formData.append("upload_file", fileBlob, finalFileName);
+  const response = await fetch("https://api.cloud.llamaindex.ai/api/v1/files", {
+    method: "POST",
+    headers: {
+      Authorization: `Bearer ${apiKey}`,
+      accept: "application/json",
+    },
+    body: formData,
+  });
+
+  if (!response.ok) {
+    throw new Error(
+      `File upload failed: ${response.status} ${response.statusText}`,
+    );
+  }
+
+  const result: UploadResponse = (await response.json()) as UploadResponse;
+  return result.id;
+}
+
+/**
+ * Helper function to determine content type based on file extension
+ */
+function getContentType(fileName: string): string {
+  const ext = fileName.toLowerCase().split(".").pop();
+  switch (ext) {
+    case "pdf":
+      return "application/pdf";
+    case "docx":
+      return "application/vnd.openxmlformats-officedocument.wordprocessingml.document";
+    case "doc":
+      return "application/msword";
+    case "txt":
+      return "text/plain";
+    case "html":
+      return "text/html";
+    case "md":
+      return "text/markdown";
+    default:
+      return "application/octet-stream";
+  }
+}
+
+/**
+ * Starts an extraction job
+ */
+async function runExtractionJob(
+  apiKey: string,
+  agentId: string,
+  fileId: string,
+): Promise<string> {
+  const response = await fetch(
+    "https://api.cloud.llamaindex.ai/api/v1/extraction/jobs",
+    {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${apiKey}`,
+        accept: "application/json",
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        extraction_agent_id: agentId,
+        file_id: fileId,
+      }),
+    },
+  );
+
+  if (!response.ok) {
+    throw new Error(
+      `Extraction job failed to start: ${response.status} ${response.statusText}`,
+    );
+  }
+
+  const result: ExtractionJobResponse =
+    (await response.json()) as ExtractionJobResponse;
+  return result.id;
+}
+
+/**
+ * Polls for job completion
+ */
+async function pollForJobCompletion(
+  apiKey: string,
+  jobId: string,
+  pollInterval: number,
+  maxRetries: number,
+): Promise<void> {
+  let retries = 0;
+
+  while (retries < maxRetries) {
+    const response = await fetch(
+      `https://api.cloud.llamaindex.ai/api/v1/extraction/jobs/${jobId}`,
+      {
+        method: "GET",
+        headers: {
+          Authorization: `Bearer ${apiKey}`,
+          accept: "application/json",
+        },
+      },
+    );
+
+    if (!response.ok) {
+      throw new Error(
+        `Failed to get job status: ${response.status} ${response.statusText}`,
+      );
+    }
+
+    const status: JobStatusResponse =
+      (await response.json()) as JobStatusResponse;
+
+    if (status.status === "SUCCESS") {
+      return; // Job completed successfully
+    }
+
+    if (status.status === "FAILED") {
+      throw new Error(`Extraction job failed`);
+    }
+
+    // Job is still PENDING or RUNNING, wait and retry
+    await new Promise((resolve) => setTimeout(resolve, pollInterval));
+    retries++;
+  }
+
+  throw new Error(`Job polling timed out after ${maxRetries} attempts`);
+}
+
+/**
+ * Retrieves the results of a completed extraction job
+ */
+async function getExtractionResults(
+  apiKey: string,
+  jobId: string,
+): Promise<ExtractionResult> {
+  const response = await fetch(
+    `https://api.cloud.llamaindex.ai/api/v1/extraction/jobs/${jobId}/result`,
+    {
+      method: "GET",
+      headers: {
+        Authorization: `Bearer ${apiKey}`,
+        accept: "application/json",
+      },
+    },
+  );
+
+  if (!response.ok) {
+    throw new Error(
+      `Failed to get extraction results: ${response.status} ${response.statusText}`,
+    );
+  }
+
+  return (await response.json()) as ExtractionResult;
+}
+
+export { extractDataFromFile };

packages/llamaindex/src/cloud/index.ts

@@ -4,4 +4,5 @@ export {
   LlamaCloudRetriever,
   type CloudRetrieveParams,
 } from "./LlamaCloudRetriever.js";
+export { LlamaExtract, LlamaExtractAgent } from "./LlamaExtract/index.js";
 export type { CloudConstructorParams } from "./type.js";

pnpm-lock.yaml

@@ -1099,6 +1099,9 @@ importers:
       magic-bytes.js:
         specifier: ^1.10.0
         version: 1.10.0
+      zod:
+        specifier: ^3.25.76
+        version: 3.25.76
     devDependencies:
       ajv:
         specifier: ^8.17.1