Fix issues in batch fallback strategies

Author: ianrob

lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_llm_proposition_extractor_sync.py

@@ -63,20 +63,29 @@ def _get_json(self, node, llm, inference_parameters):
         }
 
     def _run_non_batch_extractor(self, nodes):
+        
+        all_nodes = [node for node in nodes]
 
         extractor = LLMPropositionExtractor(
             prompt_template=self.prompt_template, 
             source_metadata_field=self.source_metadata_field
         )
-
-        return extractor.extract(nodes)
+        
+        extracted = extractor.extract(all_nodes)
+        
+        results = [{n.node_id: e[PROPOSITIONS_KEY]} for (n, e) in zip(all_nodes, extracted)]
+        
+        return results
 
     def _update_node(self, node:TextNode, node_metadata_map):
         if node.node_id in node_metadata_map:
-            raw_response = node_metadata_map[node.node_id]
-            propositions = raw_response.split('\n')
-            propositions_model = Propositions(propositions=[p for p in propositions if p])
-            node.metadata[PROPOSITIONS_KEY] = propositions_model.model_dump()['propositions']                
+            proposition_data = node_metadata_map[node.node_id]
+            if isinstance(proposition_data, list):
+                node.metadata[PROPOSITIONS_KEY] = proposition_data
+            else:
+                propositions = proposition_data.split('\n')
+                propositions_model = Propositions(propositions=[p for p in propositions if p])
+                node.metadata[PROPOSITIONS_KEY] = propositions_model.model_dump()['propositions']                
         else:
             node.metadata[PROPOSITIONS_KEY] = []
         return node

lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/batch_topic_extractor_sync.py

@@ -76,6 +76,8 @@ def _get_json(self, node, llm, inference_parameters):
         }
 
     def _run_non_batch_extractor(self, nodes):
+        
+        all_nodes = [node for node in nodes]
 
         extractor = TopicExtractor( 
             prompt_template=self.prompt_template, 
@@ -84,14 +86,20 @@ def _run_non_batch_extractor(self, nodes):
             topic_provider=self.topic_provider
         )
 
-        return extractor.extract(nodes)
+        extracted = extractor.extract(all_nodes)
+        
+        results = [{n.id_: e[TOPICS_KEY]} for (n, e) in zip(all_nodes, extracted)]
+        
+        return results
 
     def _update_node(self, node:TextNode, node_metadata_map):
         if node.node_id in node_metadata_map:
-            raw_response = node_metadata_map[node.node_id]
-            (topics, _) = parse_extracted_topics(raw_response)
-            node.metadata[TOPICS_KEY] = topics.model_dump()             
+            topic_data = node_metadata_map[node.node_id]
+            if isinstance(topic_data, dict):
+                node.metadata[TOPICS_KEY] = topic_data
+            else:
+                (topics, _) = parse_extracted_topics(topic_data)
+                node.metadata[TOPICS_KEY] = topics.model_dump()             
         else:
             node.metadata[TOPICS_KEY] = []
-        return node
-    
+        return node

lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/llm_proposition_extractor.py

@@ -133,7 +133,7 @@ async def _extract_propositions_for_nodes(self, nodes):
             jobs, 
             show_progress=self.show_progress, 
             workers=self.num_workers, 
-            desc=f'Extracting propositions [nodes: {len(nodes)}, num_workers: {self.num_workers}]'
+            desc=f'Extracting propositions [nodes: {len(jobs)}, num_workers: {self.num_workers}]'
         )
 
     async def _extract_propositions_for_node(self, node):

lexical-graph/src/graphrag_toolkit/lexical_graph/indexing/extract/topic_extractor.py

@@ -125,7 +125,7 @@ async def _extract_for_nodes(self, nodes):
             jobs, 
             show_progress=self.show_progress, 
             workers=self.num_workers, 
-            desc=f'Extracting topics [nodes: {len(nodes)}, num_workers: {self.num_workers}]'
+            desc=f'Extracting topics [nodes: {len(jobs)}, num_workers: {self.num_workers}]'
         )
 
     def _get_metadata_or_default(self, metadata, key, default):