Compare commits

...

3 Commits

Author SHA1 Message Date
Alex Yang dcb6c0dcea fix: api key 2025-03-12 17:56:39 +00:00
Alex Yang 4924c9e050 fix: package.json 2025-03-12 17:47:43 +00:00
Alex Yang 92d32a10bc feat: composite retriever 2025-03-12 17:46:19 +00:00
4 changed files with 101 additions and 0 deletions
+7
View File
@@ -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",
+22
View File
@@ -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}`,
},
});
}
+68
View File
@@ -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;
}
+4
View File
@@ -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();