feat(cloud): composite retriever

packages/cloud/package.json

@@ -37,6 +37,13 @@
       },
       "default": "./reader/dist/index.js"
     },
+    "./retriever": {
+      "import": {
+        "types": "./reader/dist/index.d.ts",
+        "default": "./reader/dist/index.js"
+      },
+      "default": "./reader/dist/index.js"
+    },
     ".": {
       "require": {
         "types": "./reader/dist/index.d.cts",

packages/cloud/src/api.ts

@@ -1 +1,23 @@
+import { getEnv } from "@llamaindex/env";
+import { client } from "./client";
+
 export * from "./client";
+
+export type ClientOptions = {
+  baseUrl?: string;
+  apiKey?: string;
+};
+
+export function initClient(options: ClientOptions) {
+  const baseUrl = options.baseUrl ?? "https://api.cloud.llamaindex.ai";
+  const apiKey = options.apiKey ?? getEnv("LLAMA_CLOUD_API_KEY");
+  if (!apiKey) {
+    throw new Error("Cannot find LlamaCloud API key");
+  }
+  client.setConfig({
+    baseUrl,
+    headers: {
+      Authorization: `Bearer ${apiKey}`,
+    },
+  });
+}

packages/cloud/src/retriever.ts

@@ -0,0 +1,68 @@
+import type { QueryType } from "@llamaindex/core/query-engine";
+import type { Retriever } from "@llamaindex/core/retriever";
+import { Document, type NodeWithScore } from "@llamaindex/core/schema";
+import { extractText } from "@llamaindex/core/utils";
+import {
+  listRetrieversApiV1RetrieversGet,
+  retrieveApiV1RetrieversRetrieverIdRetrievePost,
+  type CompositeRetrievalParams,
+} from "./api";
+
+type RetrieversParams =
+  | {
+      id: string;
+    }
+  | {
+      name: string;
+    };
+
+export function compositeRetriever(params: RetrieversParams) {
+  return {
+    retrieve: async (
+      query: QueryType,
+      retrievalParams?: Omit<CompositeRetrievalParams, "query">,
+    ): Promise<NodeWithScore[]> => {
+      let retrieverId: string | null = null;
+      if ("name" in params) {
+        const result = await listRetrieversApiV1RetrieversGet({
+          query: {
+            name: params.name,
+          },
+          throwOnError: true,
+        });
+        const retriever = result.data.find(
+          (retriever) => retriever.name === params.name,
+        );
+        if (retriever) {
+          retrieverId = retriever.id;
+        }
+      } else {
+        retrieverId = params.id;
+      }
+      if (!retrieverId) {
+        throw new Error("Retriever not found");
+      }
+      const result = await retrieveApiV1RetrieversRetrieverIdRetrievePost({
+        path: {
+          retriever_id: retrieverId,
+        },
+        body: {
+          ...retrievalParams,
+          query: extractText(query),
+        },
+        throwOnError: true,
+      });
+      return (
+        result.data.nodes?.map(({ node, score }) => ({
+          node: new Document({
+            text: node.text,
+            metadata: node.metadata,
+            startCharIdx: node.start_char_idx ?? undefined,
+            endCharIdx: node.end_char_idx ?? undefined,
+          }),
+          score: score ?? undefined,
+        })) ?? []
+      );
+    },
+  } satisfies Retriever;
+}

packages/core/src/retriever/index.ts

@@ -21,6 +21,10 @@ export type RetrieveEndEvent = {
   nodes: NodeWithScore[];
 };
 
+export interface Retriever {
+  retrieve: (params: QueryType) => Promise<NodeWithScore[]>;
+}
+
 export abstract class BaseRetriever extends PromptMixin {
   objectMap: Map<string, unknown> = new Map();