From 99826cff438985cd2e24967e0b9ad66cb8cb90cf Mon Sep 17 00:00:00 2001
From: Thuc Pham <51660321+thucpn@users.noreply.github.com>
Date: Fri, 5 Jul 2024 20:43:26 +0700
Subject: [PATCH] fix: missing dispatch retrieve event on llamacloud retriever
(#1018)
---
.../src/cloud/LlamaCloudRetriever.ts | 19 +++++++++++++++++--
1 file changed, 17 insertions(+), 2 deletions(-)
diff --git a/packages/llamaindex/src/cloud/LlamaCloudRetriever.ts b/packages/llamaindex/src/cloud/LlamaCloudRetriever.ts
index 1ae722f2f..c437d6e3d 100644
--- a/packages/llamaindex/src/cloud/LlamaCloudRetriever.ts
+++ b/packages/llamaindex/src/cloud/LlamaCloudRetriever.ts
@@ -8,6 +8,7 @@ import type { NodeWithScore } from "@llamaindex/core/schema";
import { jsonToNode, ObjectType } from "@llamaindex/core/schema";
import type { BaseRetriever, RetrieveParams } from "../Retriever.js";
import { wrapEventCaller } from "../internal/context/EventCaller.js";
+import { getCallbackManager } from "../internal/settings/CallbackManager.js";
import { extractText } from "../llm/utils.js";
import type { ClientParams, CloudConstructorParams } from "./constants.js";
import { DEFAULT_PROJECT_NAME } from "./constants.js";
@@ -28,9 +29,14 @@ export class LlamaCloudRetriever implements BaseRetriever {
nodes: TextNodeWithScore[],
): NodeWithScore[] {
return nodes.map((node: TextNodeWithScore) => {
+ const textNode = jsonToNode(node.node, ObjectType.TEXT);
+ textNode.metadata = {
+ ...textNode.metadata,
+ ...node.node.extra_info, // append LlamaCloud extra_info to node metadata (file_name, pipeline_id, etc.)
+ };
return {
// Currently LlamaCloud only supports text nodes
- node: jsonToNode(node.node, ObjectType.TEXT),
+ node: textNode,
score: node.score,
};
});
@@ -83,6 +89,15 @@ export class LlamaCloudRetriever implements BaseRetriever {
},
});
- return this.resultNodesToNodeWithScore(results.retrieval_nodes);
+ const nodesWithScores = this.resultNodesToNodeWithScore(
+ results.retrieval_nodes,
+ );
+ getCallbackManager().dispatchEvent("retrieve-end", {
+ payload: {
+ query,
+ nodes: nodesWithScores,
+ },
+ });
+ return nodesWithScores;
}
}
--
GitLab