fix(cloud): add generic types into agent data responses (#2102)

Co-authored-by: Alex Yang <himself65@outlook.com>

.changeset/wicked-pianos-share.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/cloud": patch
+---
+
+fix: add generic types into agent data responses

packages/cloud/src/beta/agent/client.ts

@@ -7,70 +7,18 @@ import {
   getAgentDataApiV1BetaAgentDataItemIdGet,
   searchAgentDataApiV1BetaAgentDataSearchPost,
   updateAgentDataApiV1BetaAgentDataItemIdPut,
-  type AggregateRequest as OriginalAggregateRequest,
-  type SearchRequest as OriginalSearchRequest,
+  type AgentData,
+  type AggregateGroup,
 } from "../../client";
 import type {
-  AgentData,
-  AggregateRequest,
-  PaginatedResponseAgentData,
-  PaginatedResponseAggregateGroup,
-  SearchRequest,
+  AggregateAgentDataOptions,
+  SearchAgentDataOptions,
+  TypedAgentData,
+  TypedAgentDataItems,
+  TypedAggregateGroup,
+  TypedAggregateGroupItems,
 } from "./types";
 
-// Utility functions to convert between camelCase and snake_case
-function camelToSnake(str: string): string {
-  return str.replace(/[A-Z]/g, (letter) => `_${letter.toLowerCase()}`);
-}
-
-function snakeToCamel(str: string): string {
-  return str.replace(/_([a-z])/g, (_, letter) => letter.toUpperCase());
-}
-
-function convertKeysToSnakeCase(obj: unknown): unknown {
-  if (Array.isArray(obj)) {
-    return obj.map(convertKeysToSnakeCase);
-  } else if (
-    obj !== null &&
-    typeof obj === "object" &&
-    obj.constructor === Object
-  ) {
-    const result: Record<string, unknown> = {};
-    for (const [key, value] of Object.entries(obj)) {
-      const snakeKey = camelToSnake(key);
-      result[snakeKey] = convertKeysToSnakeCase(value);
-    }
-    return result;
-  }
-  return obj;
-}
-
-function convertKeysToCamelCase<T>(obj: unknown): T {
-  if (Array.isArray(obj)) {
-    return obj.map((item) => convertKeysToCamelCase(item)) as T;
-  } else if (
-    obj !== null &&
-    typeof obj === "object" &&
-    obj.constructor === Object
-  ) {
-    const result: Record<string, unknown> = {};
-    for (const [key, value] of Object.entries(obj)) {
-      const camelKey = snakeToCamel(key);
-      result[camelKey] = convertKeysToCamelCase(value);
-    }
-    return result as T;
-  }
-  return obj as T;
-}
-
-type AgentClientOptions = {
-  apiKey?: string;
-  baseUrl?: string;
-  collection?: string;
-  agentUrlId?: string;
-  windowUrl?: string;
-};
-
 /**
  * Async client for agent data operations
  */
@@ -81,26 +29,18 @@ export class AgentClient<T = unknown> {
   private collection: string;
   private agentUrlId: string;
 
-  constructor(options: AgentClientOptions = {}) {
-    // Handle windowUrl to infer agentUrlId
-    let inferredAgentUrlId: string | undefined;
-    if (options.windowUrl && !options.agentUrlId) {
-      try {
-        const path = new URL(options.windowUrl).pathname;
-        // /deployments/<agent-url-id>/ui/ -> ["", "deployments", "<agent-url-id>", "ui"]
-        inferredAgentUrlId = path.split("/")[2];
-      } catch (error) {
-        console.warn(
-          "Failed to infer agent url id from window url, falling back to default",
-          error,
-        );
-      }
-    }
-
-    this.collection = options.collection || "default";
-    this.agentUrlId = options.agentUrlId || inferredAgentUrlId || "default";
-    const apiKey = options?.apiKey || getEnv("LLAMA_CLOUD_API_KEY");
-    this.baseUrl = options?.baseUrl || "https://api.cloud.llamaindex.ai/";
+  constructor({
+    apiKey = getEnv("LLAMA_CLOUD_API_KEY"),
+    baseUrl = "https://api.cloud.llamaindex.ai/",
+    collection = "default",
+    agentUrlId = "default",
+  }: {
+    apiKey?: string;
+    baseUrl?: string;
+    collection?: string;
+    agentUrlId?: string;
+  }) {
+    this.baseUrl = baseUrl;
 
     this.headers = {
       "X-SDK-Name": "llamaindex-ts",
@@ -113,29 +53,32 @@ export class AgentClient<T = unknown> {
         headers: this.headers,
       }),
     );
+
+    this.collection = collection;
+    this.agentUrlId = agentUrlId;
   }
 
   /**
    * Create new agent data
    */
-  async createItem(data: T): Promise<AgentData> {
+  async createItem(data: T): Promise<TypedAgentData<T>> {
     const response = await createAgentDataApiV1BetaAgentDataPost({
       throwOnError: true,
       body: {
-        collection: this.collection,
         agent_slug: this.agentUrlId,
-        data: convertKeysToSnakeCase(data) as Record<string, unknown>,
+        collection: this.collection,
+        data: data as Record<string, unknown>,
       },
       client: this.client,
     });
 
-    return convertKeysToCamelCase<AgentData>(response.data);
+    return this.transformResponse(response.data);
   }
 
   /**
    * Get agent data by ID
    */
-  async getItem(id: string): Promise<AgentData | null> {
+  async getItem(id: string): Promise<TypedAgentData<T> | null> {
     try {
       const response = await getAgentDataApiV1BetaAgentDataItemIdGet({
         throwOnError: true,
@@ -143,7 +86,7 @@ export class AgentClient<T = unknown> {
         client: this.client,
       });
 
-      return convertKeysToCamelCase<AgentData>(response.data);
+      return this.transformResponse(response.data);
     } catch (error) {
       if (
         error instanceof Error &&
@@ -159,17 +102,17 @@ export class AgentClient<T = unknown> {
   /**
    * Update agent data
    */
-  async updateItem(id: string, data: T): Promise<AgentData> {
+  async updateItem(id: string, data: T): Promise<TypedAgentData<T>> {
     const response = await updateAgentDataApiV1BetaAgentDataItemIdPut({
       throwOnError: true,
       path: { item_id: id },
       body: {
-        data: convertKeysToSnakeCase(data) as Record<string, unknown>,
+        data: data as Record<string, unknown>,
       },
       client: this.client,
     });
 
-    return convertKeysToCamelCase<AgentData>(response.data);
+    return this.transformResponse(response.data);
   }
 
   /**
@@ -184,54 +127,188 @@ export class AgentClient<T = unknown> {
   }
 
   /**
-   * Search agent data
+   * List agent data
    */
   async search(
-    options: Partial<SearchRequest> = {},
-  ): Promise<PaginatedResponseAgentData> {
-    const snakeCaseOptions = convertKeysToSnakeCase(
-      options,
-    ) as Partial<OriginalSearchRequest>;
+    options: SearchAgentDataOptions,
+  ): Promise<TypedAgentDataItems<T>> {
     const response = await searchAgentDataApiV1BetaAgentDataSearchPost({
       throwOnError: true,
       body: {
-        ...snakeCaseOptions,
         agent_slug: this.agentUrlId,
-        collection: this.collection,
+        ...(this.collection !== undefined && {
+          collection: this.collection,
+        }),
+        ...(options.filter !== undefined && { filter: options.filter }),
+        ...(options.orderBy !== undefined && { order_by: options.orderBy }),
+        ...(options.pageSize !== undefined && { page_size: options.pageSize }),
+        ...(options.offset !== undefined && { offset: options.offset }),
+        ...(options.includeTotal !== undefined && {
+          include_total: options.includeTotal,
+        }),
       },
       client: this.client,
     });
 
-    return convertKeysToCamelCase<PaginatedResponseAgentData>(response.data);
+    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;
   }
 
   /**
    * Aggregate agent data into groups
    */
   async aggregate(
-    options: Partial<AggregateRequest> = {},
-  ): Promise<PaginatedResponseAggregateGroup> {
-    const snakeCaseOptions = convertKeysToSnakeCase(
-      options,
-    ) as Partial<OriginalAggregateRequest>;
+    options: AggregateAgentDataOptions,
+  ): Promise<TypedAggregateGroupItems<T>> {
     const response = await aggregateAgentDataApiV1BetaAgentDataAggregatePost({
       throwOnError: true,
       body: {
-        ...snakeCaseOptions,
         agent_slug: this.agentUrlId,
-        collection: this.collection,
+        ...(this.collection !== undefined && {
+          collection: this.collection,
+        }),
+        ...(options.filter !== undefined && { filter: options.filter }),
+        ...(options.groupBy !== undefined && { group_by: options.groupBy }),
+        ...(options.count !== undefined && { count: options.count }),
+        ...(options.first !== undefined && { first: options.first }),
+        ...(options.orderBy !== undefined && { order_by: options.orderBy }),
+        ...(options.offset !== undefined && { offset: options.offset }),
+        ...(options.pageSize !== undefined && { page_size: options.pageSize }),
       },
       client: this.client,
     });
 
-    return convertKeysToCamelCase<PaginatedResponseAggregateGroup>(
-      response.data,
-    );
+    const result: TypedAggregateGroupItems<T> = {
+      items: response.data.items.map((item) =>
+        this.transformAggregateResponse(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;
+  }
+
+  /**
+   * Transform API response to typed data
+   */
+  private transformResponse(data: AgentData): TypedAgentData<T> {
+    const result: TypedAgentData<T> = {
+      id: data.id!,
+      agentUrlId: 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;
+  }
+
+  /**
+   * Transform API aggregate response to typed data
+   */
+  private transformAggregateResponse(
+    data: AggregateGroup,
+  ): TypedAggregateGroup<T> {
+    const result: TypedAggregateGroup<T> = {
+      groupKey: data.group_key,
+    };
+
+    if (data.count !== null && data.count !== undefined) {
+      result.count = data.count;
+    }
+
+    if (data.first_item !== null && data.first_item !== undefined) {
+      result.firstItem = data.first_item as T;
+    }
+
+    return result;
   }
 }
 
-export function createAgentDataClient<T = unknown>(
-  options: AgentClientOptions = {},
-): AgentClient<T> {
-  return new AgentClient<T>(options);
+export interface AgentDataClientOptions<T = unknown> {
+  /** API key for the client */
+  apiKey?: string;
+  /** Base URL for the client */
+  /** Base URL of the llama cloud api */
+  baseUrl?: string;
+  /** If running in an agent runtime, optionally provide the window url to infer the agent url id */
+  windowUrl?: string;
+  /** Agent URL ID for the client, if not provided, it will be inferred from the window url, or fall back to "default" */
+  agentUrlId?: string;
+  /** Collection name for the client, defaults to "default" */
+  collection?: string;
+}
+/**
+ * Create a new AsyncAgentDataClient instance
+ * @param options - The options for the client
+ * @returns A new AgentClient instance
+ */
+export function createAgentDataClient<T = unknown>({
+  apiKey,
+  baseUrl,
+  windowUrl,
+  agentUrlId,
+  collection = "default",
+}: {
+  apiKey?: string;
+  baseUrl?: string;
+  windowUrl?: string;
+  agentUrlId?: string;
+  collection?: string;
+} = {}): AgentClient<T> {
+  if (windowUrl && !agentUrlId) {
+    try {
+      const path = new URL(windowUrl).pathname;
+      // /deployments/<agent-url-id>/ui/ -> ["", "deployments", "<agent-url-id>", "ui"]
+      agentUrlId = path.split("/")[2];
+    } catch (error) {
+      console.warn(
+        "Failed to infer agent url id from window url, falling back to default",
+        error,
+      );
+    }
+  }
+
+  return new AgentClient({
+    ...(apiKey && { apiKey }),
+    ...(baseUrl && { baseUrl }),
+    ...(agentUrlId && { agentUrlId }),
+    collection,
+  });
 }

packages/cloud/src/beta/agent/index.ts

@@ -1,15 +1,16 @@
 export { AgentClient, createAgentDataClient } from "./client";
 
 export type {
-  AgentData,
-  AggregateGroup,
-  AggregateRequest,
+  AggregateAgentDataOptions,
+  ComparisonOperator,
   ExtractedData,
   FilterOperation,
-  PaginatedResponseAgentData,
-  PaginatedResponseAggregateGroup,
-  SearchRequest,
+  SearchAgentDataOptions,
   StatusType,
+  TypedAgentData,
+  TypedAgentDataItems,
+  TypedAggregateGroup,
+  TypedAggregateGroupItems,
 } from "./types";
 
 export { StatusType as StatusTypeEnum } from "./types";

packages/cloud/src/beta/agent/types.ts

@@ -1,40 +1,4 @@
-// Utility type to convert snake_case to camelCase
-type SnakeToCamelCase<S extends string> = S extends `${infer T}_${infer U}`
-  ? `${T}${Capitalize<SnakeToCamelCase<U>>}`
-  : S;
-
-// Utility type to recursively convert all keys in an object from snake_case to camelCase
-type KeysToCamelCase<T> =
-  T extends Array<infer U>
-    ? Array<KeysToCamelCase<U>>
-    : T extends object
-      ? {
-          [K in keyof T as SnakeToCamelCase<K & string>]: KeysToCamelCase<T[K]>;
-        }
-      : T;
-
-// Import original types from client
-import type {
-  AgentData as OriginalAgentData,
-  AggregateGroup as OriginalAggregateGroup,
-  AggregateRequest as OriginalAggregateRequest,
-  FilterOperation as OriginalFilterOperation,
-  PaginatedResponseAgentData as OriginalPaginatedResponseAgentData,
-  PaginatedResponseAggregateGroup as OriginalPaginatedResponseAggregateGroup,
-  SearchRequest as OriginalSearchRequest,
-} from "../../client";
-
-// Re-export types with camelCase conversion
-export type AgentData = KeysToCamelCase<OriginalAgentData>;
-export type PaginatedResponseAgentData =
-  KeysToCamelCase<OriginalPaginatedResponseAgentData>;
-export type PaginatedResponseAggregateGroup =
-  KeysToCamelCase<OriginalPaginatedResponseAggregateGroup>;
-export type SearchRequest = KeysToCamelCase<OriginalSearchRequest>;
-export type AggregateRequest = KeysToCamelCase<OriginalAggregateRequest>;
-export type FilterOperation = KeysToCamelCase<OriginalFilterOperation>;
-export type AggregateGroup = KeysToCamelCase<OriginalAggregateGroup>;
-
+import type { FilterOperation as RawFilterOperation } from "../../client/types.gen";
 /**
  * Status types for agent data processing
  */
@@ -47,24 +11,128 @@ export const StatusType = {
 
 export type StatusType = (typeof StatusType)[keyof typeof StatusType];
 
+export const ComparisonOperator = {
+  GT: "gt",
+  GTE: "gte",
+  LT: "lt",
+  LTE: "lte",
+  EQ: "eq",
+  INCLUDES: "includes",
+} as const;
+
+export type ComparisonOperator =
+  (typeof ComparisonOperator)[keyof typeof ComparisonOperator];
+
+/**
+ * Filter operation for searching/filtering agent data
+ */
+export type FilterOperation = RawFilterOperation;
+
 /**
  * Base extracted data interface
  */
 export interface ExtractedData<T = unknown> {
   /** The original data that was extracted from the document. For tracking changes. Should not be updated. */
-  originalData: T;
+  original_data: T;
   /** The latest state of the data. Will differ if data has been updated. */
   data?: T;
   /** The status of the extracted data. Prefer to use the StatusType values, but any string is allowed. */
   status: StatusType | string;
-  /** Confidence scores, if any, for each primitive field in the originalData data. */
+  /** Confidence scores, if any, for each primitive field in the original_data data. */
   confidence?: Record<string, unknown>;
   /** The ID of the file that was used to extract the data. */
-  fileId?: string;
+  file_id?: string;
   /** The name of the file that was used to extract the data. */
-  fileName?: string;
+  file_name?: string;
   /** The hash of the file that was used to extract the data. */
-  fileHash?: string;
+  file_hash?: string;
   /** Additional metadata about the extracted data, such as errors, tokens, etc. */
   metadata?: Record<string, unknown>;
 }
+
+/**
+ * TypedAgentData interface for typed agent data
+ */
+export interface TypedAgentData<T = unknown> {
+  /** The unique ID of the agent data record. */
+  id: string;
+  /** The ID of the agent that created the data. */
+  agentUrlId: string;
+  /** The collection of the agent data. */
+  collection?: string;
+  /** The data of the agent data. Usually an ExtractedData&lt;SomeOtherType&gt; */
+  data: T;
+  /** The date and time the data was created. */
+  createdAt: Date;
+  /** The date and time the data was last updated. */
+  updatedAt: Date;
+}
+
+/**
+ * Paginated response of typed agent data items
+ */
+export interface TypedAgentDataItems<T = unknown> {
+  items: TypedAgentData<T>[];
+  totalSize?: number;
+  nextPageToken?: string;
+}
+
+/**
+ * Options for listing agent data
+ */
+export interface SearchAgentDataOptions {
+  /** Filter options for the list. */
+  filter?: Record<string, FilterOperation>;
+  /** Order by options for the list. */
+  orderBy?: string;
+  /** Page size for the list. */
+  pageSize?: number;
+  /** Offset for the list. */
+  offset?: number;
+  /**
+   * Whether to include the total number of items in the response.
+   * Should use only for first request to build total pagination, and not subsequent requests.
+   */
+  includeTotal?: boolean;
+}
+
+/**
+ * Options for aggregating agent data
+ */
+export interface AggregateAgentDataOptions {
+  /** Filter options for the aggregation. */
+  filter?: Record<string, FilterOperation>;
+  /** Fields to group by. */
+  groupBy?: string[];
+  /** Whether to count the number of items in each group. */
+  count?: boolean;
+  /** Whether to return the first item in each group. */
+  first?: boolean;
+  /** Order by options for the aggregation. */
+  orderBy?: string;
+  /** Offset for the aggregation. */
+  offset?: number;
+  /** Page size for the aggregation. */
+  pageSize?: number;
+}
+
+/**
+ * Single aggregation group result
+ */
+export interface TypedAggregateGroup<T = unknown> {
+  /** The group key values */
+  groupKey: Record<string, unknown>;
+  /** Count of items in the group */
+  count?: number;
+  /** First item in the group */
+  firstItem?: T;
+}
+
+/**
+ * Paginated response of aggregated agent data
+ */
+export interface TypedAggregateGroupItems<T = unknown> {
+  items: TypedAggregateGroup<T>[];
+  totalSize?: number;
+  nextPageToken?: string;
+}