fix several errors

Author: vkehfdl1

autorag/autorag/nodes/passageaugmenter/base.py

@@ -11,7 +11,7 @@
 	validate_corpus_dataset,
 	cast_corpus_dataset,
 )
-from autorag.utils.cast import cast_retrieve_infos
+from autorag.utils.cast import cast_retrieved_ids
 from autorag.utils.util import select_top_k
 
 logger = logging.getLogger("AutoRAG")
@@ -42,8 +42,7 @@ def cast_to_run(self, previous_result: pd.DataFrame, *args, **kwargs):
 		validate_qa_dataset(previous_result)
 
 		# find ids columns
-		retrieve_infos = cast_retrieve_infos(previous_result)
-		return retrieve_infos["retrieved_ids"]
+		return cast_retrieved_ids(previous_result)
 
 	@staticmethod
 	def sort_by_scores(

autorag/autorag/nodes/passagecompressor/base.py

@@ -8,7 +8,7 @@
 from autorag import generator_models
 from autorag.schema import BaseModule
 from autorag.utils import result_to_dataframe
-from autorag.utils.cast import cast_retrieve_infos
+from autorag.utils.cast import cast_retrieved_contents
 
 logger = logging.getLogger("AutoRAG")
 
@@ -34,9 +34,8 @@ def cast_to_run(self, previous_result: pd.DataFrame, *args, **kwargs):
 		assert len(previous_result) > 0, "previous_result must have at least one row."
 
 		queries = previous_result["query"].tolist()
-		retrieve_infos = cast_retrieve_infos(previous_result)
 
-		return queries, retrieve_infos["retrieved_contents"]
+		return queries, cast_retrieved_contents(previous_result)
 
 
 class LlamaIndexCompressor(BasePassageCompressor, metaclass=abc.ABCMeta):

autorag/autorag/nodes/promptmaker/base.py

@@ -6,7 +6,7 @@
 import pandas as pd
 
 from autorag.schema.base import BaseModule
-from autorag.utils.cast import cast_retrieve_infos
+from autorag.utils.cast import cast_retrieved_contents
 
 logger = logging.getLogger("AutoRAG")
 
@@ -28,6 +28,5 @@ def cast_to_run(self, previous_result: pd.DataFrame, *args, **kwargs):
 		), "previous_result must have query column."
 
 		query = previous_result["query"].tolist()
-		retrieve_infos = cast_retrieve_infos(previous_result)
 		prompt = kwargs.pop("prompt")
-		return query, retrieve_infos["retrieved_contents"], prompt
+		return query, cast_retrieved_contents(previous_result), prompt