feat(cloud): missing agent api (#2094)

.changeset/bitter-buckets-unite.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/cloud": patch
+---
+
+Bug fixes for new beta agent-data cloud API

packages/cloud/package.json

@@ -13,20 +13,20 @@
     "./api",
     "./reader",
     "./parse",
-    "./agent"
+    "./beta/agent"
   ],
   "exports": {
     "./openapi.json": "./openapi.json",
-    "./agent": {
+    "./beta/agent": {
       "require": {
-        "types": "./agent/dist/index.d.cts",
-        "default": "./agent/dist/index.cjs"
+        "types": "./beta/agent/dist/index.d.cts",
+        "default": "./beta/agent/dist/index.cjs"
       },
       "import": {
-        "types": "./agent/dist/index.d.ts",
-        "default": "./agent/dist/index.js"
+        "types": "./beta/agent/dist/index.d.ts",
+        "default": "./beta/agent/dist/index.js"
       },
-      "default": "./agent/dist/index.js"
+      "default": "./beta/agent/dist/index.js"
     },
     "./api": {
       "require": {

packages/cloud/src/agent/index.ts

@@ -1 +0,0 @@
-export { AgentClient, createAgentClient } from "./client";

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

@@ -1,36 +1,55 @@
 import { createClient, createConfig } from "@hey-api/client-fetch";
 import { getEnv } from "@llamaindex/env";
 import {
+  aggregateAgentDataApiV1BetaAgentDataAggregatePost,
   createAgentDataApiV1BetaAgentDataPost,
   deleteAgentDataApiV1BetaAgentDataItemIdDelete,
   getAgentDataApiV1BetaAgentDataItemIdGet,
   searchAgentDataApiV1BetaAgentDataSearchPost,
   updateAgentDataApiV1BetaAgentDataItemIdPut,
   type AgentData,
+  type AggregateRequest,
   type PaginatedResponseAgentData,
+  type PaginatedResponseAggregateGroup,
   type SearchRequest,
-} from "../client";
+} from "../../client";
 
 type AgentClientOptions = {
   apiKey?: string;
   baseUrl?: string;
-  collection: string;
-  agentUrlId: string;
+  collection?: string;
+  agentUrlId?: string;
+  windowUrl?: string;
 };
 
 /**
  * Async client for agent data operations
  */
-export class AgentClient {
+export class AgentClient<T = unknown> {
   private client: ReturnType<typeof createClient>;
   private baseUrl: string;
   private headers: Record<string, string>;
   private collection: string;
   private agentUrlId: string;
 
-  constructor(options: AgentClientOptions) {
-    this.collection = options.collection;
-    this.agentUrlId = options.agentUrlId;
+  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/";
 
@@ -50,7 +69,7 @@ export class AgentClient {
   /**
    * Create new agent data
    */
-  async createItem<T>(data: T): Promise<AgentData> {
+  async createItem(data: T): Promise<AgentData> {
     const response = await createAgentDataApiV1BetaAgentDataPost({
       throwOnError: true,
       body: {
@@ -91,7 +110,7 @@ export class AgentClient {
   /**
    * Update agent data
    */
-  async updateItem<T>(id: string, data: T): Promise<AgentData> {
+  async updateItem(id: string, data: T): Promise<AgentData> {
     const response = await updateAgentDataApiV1BetaAgentDataItemIdPut({
       throwOnError: true,
       path: { item_id: id },
@@ -107,7 +126,7 @@ export class AgentClient {
   /**
    * Delete agent data
    */
-  async delete(id: string): Promise<void> {
+  async deleteItem(id: string): Promise<void> {
     await deleteAgentDataApiV1BetaAgentDataItemIdDelete({
       throwOnError: true,
       path: { item_id: id },
@@ -116,13 +135,36 @@ export class AgentClient {
   }
 
   /**
-   * List agent data
+   * Search agent data
    */
-  async list(options: SearchRequest): Promise<PaginatedResponseAgentData> {
+  async search(
+    options: Partial<SearchRequest> = {},
+  ): Promise<PaginatedResponseAgentData> {
     const response = await searchAgentDataApiV1BetaAgentDataSearchPost({
       throwOnError: true,
       body: {
         ...options,
+        agent_slug: this.agentUrlId,
+        collection: this.collection,
+      },
+      client: this.client,
+    });
+
+    return response.data;
+  }
+
+  /**
+   * Aggregate agent data into groups
+   */
+  async aggregate(
+    options: Partial<AggregateRequest> = {},
+  ): Promise<PaginatedResponseAggregateGroup> {
+    const response = await aggregateAgentDataApiV1BetaAgentDataAggregatePost({
+      throwOnError: true,
+      body: {
+        ...options,
+        agent_slug: this.agentUrlId,
+        collection: this.collection,
       },
       client: this.client,
     });
@@ -131,6 +173,8 @@ export class AgentClient {
   }
 }
 
-export function createAgentClient(options: AgentClientOptions): AgentClient {
-  return new AgentClient(options);
+export function createAgentDataClient<T = unknown>(
+  options: AgentClientOptions = {},
+): AgentClient<T> {
+  return new AgentClient<T>(options);
 }

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

@@ -0,0 +1,15 @@
+export { AgentClient, createAgentDataClient } from "./client";
+
+export type {
+  AgentData,
+  AggregateGroup,
+  AggregateRequest,
+  ExtractedData,
+  FilterOperation,
+  PaginatedResponseAgentData,
+  PaginatedResponseAggregateGroup,
+  SearchRequest,
+  StatusType,
+} from "./types";
+
+export { StatusType as StatusTypeEnum } from "./types";

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

@@ -0,0 +1,44 @@
+// Re-export types from the client
+export type {
+  AgentData,
+  AggregateGroup,
+  AggregateRequest,
+  FilterOperation,
+  PaginatedResponseAgentData,
+  PaginatedResponseAggregateGroup,
+  SearchRequest,
+} from "../../client";
+
+/**
+ * Status types for agent data processing
+ */
+export const StatusType = {
+  ERROR: "error",
+  ACCEPTED: "accepted",
+  REJECTED: "rejected",
+  PENDING_REVIEW: "pending_review",
+} as const;
+
+export type StatusType = (typeof StatusType)[keyof typeof StatusType];
+
+/**
+ * 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. */
+  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 original_data data. */
+  confidence?: Record<string, unknown>;
+  /** The ID of the file that was used to extract the data. */
+  file_id?: string;
+  /** The name of the file that was used to extract the data. */
+  file_name?: string;
+  /** The hash of the file that was used to extract the data. */
+  file_hash?: string;
+  /** Additional metadata about the extracted data, such as errors, tokens, etc. */
+  metadata?: Record<string, unknown>;
+}