From 1eb141a45f794af987f171442820703e447e3e53 Mon Sep 17 00:00:00 2001 From: seehi <6580@pm.me> Date: Tue, 26 Mar 2024 16:40:10 +0800 Subject: [PATCH] add object ranker --- metagpt/rag/rankers/object_ranker.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/metagpt/rag/rankers/object_ranker.py b/metagpt/rag/rankers/object_ranker.py index fe45f9395..b8456803f 100644 --- a/metagpt/rag/rankers/object_ranker.py +++ b/metagpt/rag/rankers/object_ranker.py @@ -38,12 +38,10 @@ class ObjectSortPostprocessor(BaseNodePostprocessor): return [] self._check_metadata(nodes[0].node) + sort_key = lambda node: json.loads(node.node.metadata["obj_json"])[self.field_name] return self._get_sort_func()(self.top_n, nodes, key=sort_key) - def _get_sort_func(self): - return heapq.nlargest if self.order == "desc" else heapq.nsmallest - def _check_metadata(self, node: ObjectNode): try: obj_dict = json.loads(node.metadata.get("obj_json")) @@ -52,3 +50,6 @@ class ObjectSortPostprocessor(BaseNodePostprocessor): if self.field_name not in obj_dict: raise ValueError(f"Field '{self.field_name}' not found in object: {obj_dict}") + + def _get_sort_func(self): + return heapq.nlargest if self.order == "desc" else heapq.nsmallest