feat: support organization id in llamacloud index (#1123)

Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>

.changeset/moody-pears-lick.md

@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+feat: support organization id in llamacloud index

packages/llamaindex/src/cloud/LlamaCloudIndex.ts

@@ -8,7 +8,7 @@ import type { CloudRetrieveParams } from "./LlamaCloudRetriever.js";
 import { LlamaCloudRetriever } from "./LlamaCloudRetriever.js";
 import { getPipelineCreate } from "./config.js";
 import type { CloudConstructorParams } from "./constants.js";
-import { getAppBaseUrl, initService } from "./utils.js";
+import { getAppBaseUrl, getProjectId, initService } from "./utils.js";
 
 import { PipelinesService, ProjectsService } from "@llamaindex/cloud/api";
 import { SentenceSplitter } from "@llamaindex/core/node-parser";
@@ -132,18 +132,28 @@ export class LlamaCloudIndex {
     await this.waitForPipelineIngestion(verbose, raiseOnError);
   }
 
-  private async getPipelineId(
-    name: string,
-    projectName: string,
+  public async getPipelineId(
+    name?: string,
+    projectName?: string,
   ): Promise<string> {
     const pipelines = await PipelinesService.searchPipelinesApiV1PipelinesGet({
-      projectName,
-      pipelineName: name,
+      projectId: await this.getProjectId(projectName),
+      pipelineName: name ?? this.params.name,
     });
 
     return pipelines[0].id;
   }
 
+  public async getProjectId(
+    projectName?: string,
+    organizationId?: string,
+  ): Promise<string> {
+    return await getProjectId(
+      projectName ?? this.params.projectName,
+      organizationId ?? this.params.organizationId,
+    );
+  }
+
   static async fromDocuments(
     params: {
       documents: Document[];
@@ -168,6 +178,7 @@ export class LlamaCloudIndex {
     });
 
     const project = await ProjectsService.upsertProjectApiV1ProjectsPut({
+      organizationId: params.organizationId,
       requestBody: {
         name: params.projectName ?? "default",
       },

packages/llamaindex/src/cloud/LlamaCloudRetriever.ts

@@ -11,7 +11,7 @@ import { extractText, wrapEventCaller } from "@llamaindex/core/utils";
 import type { BaseRetriever, RetrieveParams } from "../Retriever.js";
 import type { ClientParams, CloudConstructorParams } from "./constants.js";
 import { DEFAULT_PROJECT_NAME } from "./constants.js";
-import { initService } from "./utils.js";
+import { getProjectId, initService } from "./utils.js";
 
 export type CloudRetrieveParams = Omit<
   RetrievalParams,
@@ -21,6 +21,7 @@ export type CloudRetrieveParams = Omit<
 export class LlamaCloudRetriever implements BaseRetriever {
   clientParams: ClientParams;
   retrieveParams: CloudRetrieveParams;
+  organizationId?: string;
   projectName: string = DEFAULT_PROJECT_NAME;
   pipelineName: string;
 
@@ -49,6 +50,9 @@ export class LlamaCloudRetriever implements BaseRetriever {
     if (params.projectName) {
       this.projectName = params.projectName;
     }
+    if (params.organizationId) {
+      this.organizationId = params.organizationId;
+    }
   }
 
   @wrapEventCaller
@@ -57,7 +61,7 @@ export class LlamaCloudRetriever implements BaseRetriever {
     preFilters,
   }: RetrieveParams): Promise<NodeWithScore[]> {
     const pipelines = await PipelinesService.searchPipelinesApiV1PipelinesGet({
-      projectName: this.projectName,
+      projectId: await getProjectId(this.projectName, this.organizationId),
       pipelineName: this.pipelineName,
     });
 

packages/llamaindex/src/cloud/constants.ts

@@ -8,5 +8,6 @@ export type ClientParams = { apiKey?: string; baseUrl?: string };
 export type CloudConstructorParams = {
   name: string;
   projectName: string;
+  organizationId?: string;
   serviceContext?: ServiceContext;
 } & ClientParams;

packages/llamaindex/src/cloud/utils.ts

@@ -1,4 +1,4 @@
-import { OpenAPI } from "@llamaindex/cloud/api";
+import { OpenAPI, ProjectsService } from "@llamaindex/cloud/api";
 import { getEnv } from "@llamaindex/env";
 import type { ClientParams } from "./constants.js";
 import { DEFAULT_BASE_URL } from "./constants.js";
@@ -20,3 +20,31 @@ export function initService({ apiKey, baseUrl }: ClientParams = {}) {
     );
   }
 }
+
+export async function getProjectId(
+  projectName: string,
+  organizationId?: string,
+): Promise<string> {
+  const projects = await ProjectsService.listProjectsApiV1ProjectsGet({
+    projectName: projectName,
+    organizationId: organizationId,
+  });
+
+  if (projects.length === 0) {
+    throw new Error(
+      `Unknown project name ${projectName}. Please confirm a managed project with this name exists.`,
+    );
+  } else if (projects.length > 1) {
+    throw new Error(
+      `Multiple projects found with name ${projectName}. Please specify organization_id.`,
+    );
+  }
+
+  const project = projects[0];
+
+  if (!project.id) {
+    throw new Error(`No project found with name ${projectName}`);
+  }
+
+  return project.id;
+}