From 8b66cf434122b80c123f7324febff143c2fed651 Mon Sep 17 00:00:00 2001
From: Thuc Pham <51660321+thucpn@users.noreply.github.com>
Date: Thu, 15 Aug 2024 13:51:48 +0700
Subject: [PATCH] feat: support organization id in llamacloud index (#1123)
Co-authored-by: Marcus Schiesser <mail@marcusschiesser.de>
---
.changeset/moody-pears-lick.md | 5 ++++
.../llamaindex/src/cloud/LlamaCloudIndex.ts | 23 ++++++++++----
.../src/cloud/LlamaCloudRetriever.ts | 8 +++--
packages/llamaindex/src/cloud/constants.ts | 1 +
packages/llamaindex/src/cloud/utils.ts | 30 ++++++++++++++++++-
5 files changed, 58 insertions(+), 9 deletions(-)
create mode 100644 .changeset/moody-pears-lick.md
diff --git a/.changeset/moody-pears-lick.md b/.changeset/moody-pears-lick.md
new file mode 100644
index 000000000..6297bd36a
--- /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 3353e6b84..ebd25a943 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<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",
},
diff --git a/packages/llamaindex/src/cloud/LlamaCloudRetriever.ts b/packages/llamaindex/src/cloud/LlamaCloudRetriever.ts
index 0f63fdf56..b75f48f90 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<NodeWithScore[]> {
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 2a3e86898..7d75e1361 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 e030fcd43..b31ae9cf1 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<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;
+}
--
GitLab