diff --git a/.changeset/moody-pears-lick.md b/.changeset/moody-pears-lick.md new file mode 100644 index 0000000000..6297bd36ab --- /dev/null +++ b/.changeset/moody-pears-lick.md @@ -0,0 +1,5 @@ +--- +"llamaindex": patch +--- + +feat: support organization id in llamacloud index diff --git a/packages/llamaindex/src/cloud/LlamaCloudIndex.ts b/packages/llamaindex/src/cloud/LlamaCloudIndex.ts index 3353e6b84b..ebd25a9433 100644 --- a/packages/llamaindex/src/cloud/LlamaCloudIndex.ts +++ b/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 { 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 { + 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", }, diff --git a/packages/llamaindex/src/cloud/LlamaCloudRetriever.ts b/packages/llamaindex/src/cloud/LlamaCloudRetriever.ts index 0f63fdf56c..b75f48f90d 100644 --- a/packages/llamaindex/src/cloud/LlamaCloudRetriever.ts +++ b/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 { const pipelines = await PipelinesService.searchPipelinesApiV1PipelinesGet({ - projectName: this.projectName, + projectId: await getProjectId(this.projectName, this.organizationId), pipelineName: this.pipelineName, }); diff --git a/packages/llamaindex/src/cloud/constants.ts b/packages/llamaindex/src/cloud/constants.ts index 2a3e868985..7d75e13613 100644 --- a/packages/llamaindex/src/cloud/constants.ts +++ b/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; diff --git a/packages/llamaindex/src/cloud/utils.ts b/packages/llamaindex/src/cloud/utils.ts index e030fcd43e..b31ae9cf1c 100644 --- a/packages/llamaindex/src/cloud/utils.ts +++ b/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 { + 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; +}