feat: init agent api on cloud sdk (#2069)

.changeset/cold-sloths-scream.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/cloud": patch
+---
+
+feat: init agent api on cloud sdk

.changeset/funny-donuts-approve.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/cloud": patch
+---
+
+feat: init agent api on cloud sdk

packages/cloud/package.json

@@ -15,6 +15,17 @@
   ],
   "exports": {
     "./openapi.json": "./openapi.json",
+    "./agent": {
+      "require": {
+        "types": "./agent/dist/index.d.cts",
+        "default": "./agent/dist/index.cjs"
+      },
+      "import": {
+        "types": "./agent/dist/index.d.ts",
+        "default": "./agent/dist/index.js"
+      },
+      "default": "./agent/dist/index.js"
+    },
     "./api": {
       "require": {
         "types": "./api/dist/index.d.cts",

packages/cloud/src/agent/client.ts

@@ -0,0 +1,291 @@
+import { createClient, createConfig } from "@hey-api/client-fetch";
+import { getEnv } from "@llamaindex/env";
+import pRetry from "p-retry";
+import {
+  createAgentDataApiV1BetaAgentDataPost,
+  deleteAgentDataApiV1BetaAgentDataItemIdDelete,
+  getAgentDataApiV1BetaAgentDataItemIdGet,
+  searchAgentDataApiV1BetaAgentDataSearchPost,
+  updateAgentDataApiV1BetaAgentDataItemIdPut,
+  type AgentData,
+} from "../client";
+import type {
+  CreateAgentDataOptions,
+  ExtractedData,
+  ExtractOptions,
+  ListAgentDataOptions,
+  TypedAgentData,
+  TypedAgentDataItems,
+  UpdateAgentDataOptions,
+} from "./types";
+
+/**
+ * Async client for agent data operations
+ */
+export class AgentClient {
+  private client: ReturnType<typeof createClient>;
+  private baseUrl: string;
+  private headers: Record<string, string>;
+
+  constructor(options?: { apiKey?: string; baseUrl?: string }) {
+    const apiKey = options?.apiKey || getEnv("LLAMA_CLOUD_API_KEY");
+    this.baseUrl = options?.baseUrl || "https://api.cloud.llamaindex.ai/";
+
+    this.headers = {
+      "X-SDK-Name": "llamaindex-ts",
+      ...(apiKey && { Authorization: `Bearer ${apiKey}` }),
+    };
+
+    this.client = createClient(
+      createConfig({
+        baseUrl: this.baseUrl,
+        headers: this.headers,
+      }),
+    );
+  }
+
+  /**
+   * Create new agent data
+   */
+  async create<T = unknown>(
+    options: CreateAgentDataOptions<T>,
+  ): Promise<TypedAgentData<T>> {
+    const response = await createAgentDataApiV1BetaAgentDataPost({
+      throwOnError: true,
+      body: {
+        agent_slug: options.agentSlug,
+        ...(options.collection !== undefined && {
+          collection: options.collection,
+        }),
+        data: options.data as Record<string, unknown>,
+      },
+      client: this.client,
+    });
+
+    return this.transformResponse(response.data);
+  }
+
+  /**
+   * Get agent data by ID
+   */
+  async get<T = unknown>(id: string): Promise<TypedAgentData<T> | null> {
+    try {
+      const response = await getAgentDataApiV1BetaAgentDataItemIdGet({
+        throwOnError: true,
+        path: { item_id: id },
+        client: this.client,
+      });
+
+      return this.transformResponse(response.data);
+    } catch (error) {
+      if (
+        error instanceof Error &&
+        "response" in error &&
+        (error as { response?: { status?: number } }).response?.status === 404
+      ) {
+        return null;
+      }
+      throw error;
+    }
+  }
+
+  /**
+   * Update agent data
+   */
+  async update<T = unknown>(
+    id: string,
+    options: UpdateAgentDataOptions<T>,
+  ): Promise<TypedAgentData<T>> {
+    const response = await updateAgentDataApiV1BetaAgentDataItemIdPut({
+      throwOnError: true,
+      path: { item_id: id },
+      body: {
+        data: options.data as Record<string, unknown>,
+      },
+      client: this.client,
+    });
+
+    return this.transformResponse(response.data);
+  }
+
+  /**
+   * Delete agent data
+   */
+  async delete(id: string): Promise<void> {
+    await deleteAgentDataApiV1BetaAgentDataItemIdDelete({
+      throwOnError: true,
+      path: { item_id: id },
+      client: this.client,
+    });
+  }
+
+  /**
+   * List agent data
+   */
+  async list<T = unknown>(
+    options: ListAgentDataOptions,
+  ): Promise<TypedAgentDataItems<T>> {
+    const response = await searchAgentDataApiV1BetaAgentDataSearchPost({
+      throwOnError: true,
+      body: {
+        agent_slug: options.agentSlug,
+        ...(options.collection !== undefined && {
+          collection: options.collection,
+        }),
+        ...(options.filter !== undefined && { filter: options.filter }),
+        ...(options.orderBy !== undefined && { order_by: options.orderBy }),
+        ...(options.pageSize !== undefined && { page_size: options.pageSize }),
+        ...(options.pageToken !== undefined && {
+          page_token: options.pageToken,
+        }),
+        ...(options.offset !== undefined && { offset: options.offset }),
+      },
+      client: this.client,
+    });
+
+    const result: TypedAgentDataItems<T> = {
+      items: response.data.items.map((item: AgentData) =>
+        this.transformResponse(item),
+      ),
+    };
+
+    if (
+      response.data.total_size !== null &&
+      response.data.total_size !== undefined
+    ) {
+      result.totalSize = response.data.total_size;
+    }
+
+    if (
+      response.data.next_page_token !== null &&
+      response.data.next_page_token !== undefined
+    ) {
+      result.nextPageToken = response.data.next_page_token;
+    }
+
+    return result;
+  }
+
+  /**
+   * Extract data from agent with retry logic
+   */
+  async extract<T = unknown>(
+    agentId: string,
+    input: unknown,
+    options?: ExtractOptions,
+  ): Promise<ExtractedData<T>> {
+    const extractOptions = {
+      retries: options?.retryCount || 3,
+      onFailedAttempt: (error: {
+        attemptNumber: number;
+        retriesLeft: number;
+      }) => {
+        console.log(
+          `Extraction attempt ${error.attemptNumber} failed. ${error.retriesLeft} retries left.`,
+        );
+      },
+      minTimeout: options?.retryDelay || 1000,
+      maxTimeout: options?.timeout || 30000,
+    };
+
+    return pRetry(async () => {
+      // Note: The extract endpoint might not be in the generated client yet
+      // Using the native fetch API for this endpoint
+      const response = await fetch(
+        `${this.baseUrl}/api/v1/beta/agent-data/${agentId}/extract`,
+        {
+          method: "POST",
+          body: JSON.stringify({ input }),
+          headers: {
+            "Content-Type": "application/json",
+            ...this.headers,
+          },
+        },
+      );
+
+      if (!response.ok) {
+        throw new Error(`Failed to extract data: ${response.statusText}`);
+      }
+
+      const extractedData = (await response.json()) as ExtractedData<T>;
+
+      // If status is still pending or in progress, poll for completion
+      if (
+        extractedData.status === "pending" ||
+        extractedData.status === "in_progress"
+      ) {
+        return this.pollExtraction<T>(extractedData.id, options);
+      }
+
+      return extractedData;
+    }, extractOptions);
+  }
+
+  /**
+   * Poll for extraction completion
+   */
+  private async pollExtraction<T = unknown>(
+    extractionId: string,
+    options?: ExtractOptions,
+  ): Promise<ExtractedData<T>> {
+    const pollInterval = 2000; // 2 seconds
+    const maxAttempts = Math.floor((options?.timeout || 30000) / pollInterval);
+
+    for (let i = 0; i < maxAttempts; i++) {
+      await new Promise((resolve) => setTimeout(resolve, pollInterval));
+
+      const response = await fetch(
+        `${this.baseUrl}/api/v1/extractions/${extractionId}`,
+        {
+          headers: this.headers,
+        },
+      );
+
+      if (!response.ok) {
+        throw new Error(
+          `Failed to get extraction status: ${response.statusText}`,
+        );
+      }
+
+      const extractedData = (await response.json()) as ExtractedData<T>;
+
+      if (
+        extractedData.status === "completed" ||
+        extractedData.status === "failed"
+      ) {
+        return extractedData;
+      }
+    }
+
+    throw new Error("Extraction timeout exceeded");
+  }
+
+  /**
+   * Transform API response to typed data
+   */
+  private transformResponse<T = unknown>(data: AgentData): TypedAgentData<T> {
+    const result: TypedAgentData<T> = {
+      id: data.id!,
+      agentSlug: data.agent_slug,
+      data: data.data as T,
+      createdAt: new Date(data.created_at!),
+      updatedAt: new Date(data.updated_at!),
+    };
+
+    if (data.collection !== undefined) {
+      result.collection = data.collection;
+    }
+
+    return result;
+  }
+}
+
+/**
+ * Create a new AsyncAgentDataClient instance
+ */
+export function createAgentDataClient(options?: {
+  apiKey?: string;
+  baseUrl?: string;
+}): AgentClient {
+  return new AgentClient(options);
+}

packages/cloud/src/agent/index.ts

@@ -0,0 +1,18 @@
+export { AgentClient, createAgentDataClient } from "./client";
+
+export type {
+  AgentDataT,
+  CreateAgentDataOptions,
+  ExtractOptions,
+  ExtractedData,
+  ExtractedT,
+  FilterOperation,
+  ListAgentDataOptions,
+  SortOptions,
+  StatusType,
+  TypedAgentData,
+  TypedAgentDataItems,
+  UpdateAgentDataOptions,
+} from "./types";
+
+export { StatusType as StatusTypeEnum } from "./types";

packages/cloud/src/agent/types.ts

@@ -0,0 +1,98 @@
+/**
+ * Status types for agent data processing
+ */
+export enum StatusType {
+  PENDING = "pending",
+  IN_PROGRESS = "in_progress",
+  COMPLETED = "completed",
+  FAILED = "failed",
+}
+
+/**
+ * Filter operation for searching/filtering agent data
+ */
+export interface FilterOperation {
+  [key: string]: unknown;
+}
+
+/**
+ * Base extracted data interface
+ */
+export interface ExtractedData<T = unknown> {
+  id: string;
+  status: StatusType;
+  data?: T;
+  error?: string;
+  createdAt: Date;
+  updatedAt: Date;
+}
+
+/**
+ * TypedAgentData interface for typed agent data
+ */
+export interface TypedAgentData<T = unknown> {
+  id: string;
+  agentSlug: string;
+  collection?: string;
+  data: T;
+  createdAt: Date;
+  updatedAt: Date;
+}
+
+/**
+ * Collection of typed agent data items
+ */
+export interface TypedAgentDataItems<T = unknown> {
+  items: TypedAgentData<T>[];
+  totalSize?: number;
+  nextPageToken?: string;
+}
+
+/**
+ * Options for creating agent data
+ */
+export interface CreateAgentDataOptions<T = unknown> {
+  agentSlug: string;
+  collection?: string;
+  data: T;
+}
+
+/**
+ * Options for updating agent data
+ */
+export interface UpdateAgentDataOptions<T = unknown> {
+  data: T;
+}
+
+/**
+ * Sort options for listing
+ */
+export interface SortOptions {
+  field: string;
+  order: "asc" | "desc";
+}
+
+/**
+ * Options for listing agent data
+ */
+export interface ListAgentDataOptions {
+  agentSlug: string;
+  collection?: string;
+  filter?: Record<string, FilterOperation>;
+  orderBy?: string;
+  pageSize?: number;
+  pageToken?: string;
+  offset?: number;
+}
+
+/**
+ * Options for extraction
+ */
+export interface ExtractOptions {
+  timeout?: number;
+  retryCount?: number;
+  retryDelay?: number;
+}
+
+export type ExtractedT<T> = T;
+export type AgentDataT<T> = T;

packages/cloud/src/api.ts

@@ -7,5 +7,21 @@ client.setConfig({
   },
 });
 
+export {
+  AgentClient,
+  createAgentDataClient,
+  StatusTypeEnum,
+  type AgentDataT,
+  type CreateAgentDataOptions,
+  type ExtractedData,
+  type ExtractedT,
+  type ExtractOptions,
+  type ListAgentDataOptions,
+  type SortOptions,
+  type StatusType,
+  type TypedAgentData,
+  type TypedAgentDataItems,
+  type UpdateAgentDataOptions,
+} from "./agent";
 export * from "./client";
 export { client };