AutoRAG 0.2.16__tar.gz → 0.2.18__tar.gz

{autorag-0.2.16 → autorag-0.2.18}/AutoRAG.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.2.16
+Version: 0.2.18
 Summary: Automatically Evaluate RAG pipelines with your own data. Find optimal structure for new RAG product.
 Author-email: Marker-Inc <vkehfdl1@gmail.com>
 License:                                  Apache License
@@ -254,14 +254,14 @@ Requires-Dist: llama-index-core>=0.11.0
 Requires-Dist: llama-index-readers-file
 Requires-Dist: llama-index-embeddings-openai
 Requires-Dist: llama-index-embeddings-huggingface
-Requires-Dist: llama-index-llms-openai>=0.1.26
+Requires-Dist: llama-index-llms-openai>=0.2.7
 Requires-Dist: llama-index-llms-huggingface
 Requires-Dist: llama-index-llms-openai-like
 Requires-Dist: llama-index-llms-ollama
 Requires-Dist: llama-index-retrievers-bm25
 Requires-Dist: streamlit
 Requires-Dist: gradio
-Requires-Dist: langchain-core>=0.1.6
+Requires-Dist: langchain-core>=0.3.0
 Requires-Dist: langchain_unstructured
 Requires-Dist: langchain-upstage
 Requires-Dist: panel

{autorag-0.2.16 → autorag-0.2.18}/AutoRAG.egg-info/SOURCES.txt

@@ -34,6 +34,7 @@ autorag/validator.py
 autorag/web.py
 autorag/data/__init__.py
 autorag/data/beta/__init__.py
+autorag/data/beta/extract_evidence.py
 autorag/data/beta/sample.py
 autorag/data/beta/schema.py
 autorag/data/beta/filter/__init__.py
@@ -200,10 +201,12 @@ docs/source/_static/roadmap/cycle.png
 docs/source/_static/roadmap/merger.png
 docs/source/_static/roadmap/node_line_modular.png
 docs/source/_static/roadmap/policy.png
+docs/source/api_spec/autorag.data.beta.filter.rst
 docs/source/api_spec/autorag.data.beta.generation_gt.rst
 docs/source/api_spec/autorag.data.beta.query.rst
 docs/source/api_spec/autorag.data.beta.rst
 docs/source/api_spec/autorag.data.beta.schema.rst
+docs/source/api_spec/autorag.data.chunk.rst
 docs/source/api_spec/autorag.data.corpus.rst
 docs/source/api_spec/autorag.data.parse.rst
 docs/source/api_spec/autorag.data.qacreation.rst
@@ -229,9 +232,19 @@ docs/source/data_creation/data_format.md
 docs/source/data_creation/ragas.md
 docs/source/data_creation/tutorial.md
 docs/source/data_creation/beta/data_creation.md
-docs/source/data_creation/beta/filter.md
-docs/source/data_creation/beta/query_gen.md
 docs/source/data_creation/beta/tutorial.md
+docs/source/data_creation/beta/chunk/chunk.md
+docs/source/data_creation/beta/chunk/langchain_chunk.md
+docs/source/data_creation/beta/chunk/llama_index_chunk.md
+docs/source/data_creation/beta/parse/clova.md
+docs/source/data_creation/beta/parse/langchain_parse.md
+docs/source/data_creation/beta/parse/llama_parse.md
+docs/source/data_creation/beta/parse/parse.md
+docs/source/data_creation/beta/parse/table_hybrid_parse.md
+docs/source/data_creation/beta/qa_creation/answer_gen.md
+docs/source/data_creation/beta/qa_creation/filter.md
+docs/source/data_creation/beta/qa_creation/qa_creation.md
+docs/source/data_creation/beta/qa_creation/query_gen.md
 docs/source/deploy/api_endpoint.md
 docs/source/deploy/web.md
 docs/source/evaluate_metrics/generation.md

{autorag-0.2.16 → autorag-0.2.18}/AutoRAG.egg-info/requires.txt

@@ -32,14 +32,14 @@ llama-index-core>=0.11.0
 llama-index-readers-file
 llama-index-embeddings-openai
 llama-index-embeddings-huggingface
-llama-index-llms-openai>=0.1.26
+llama-index-llms-openai>=0.2.7
 llama-index-llms-huggingface
 llama-index-llms-openai-like
 llama-index-llms-ollama
 llama-index-retrievers-bm25
 streamlit
 gradio
-langchain-core>=0.1.6
+langchain-core>=0.3.0
 langchain_unstructured
 langchain-upstage
 panel

{autorag-0.2.16 → autorag-0.2.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.2.16
+Version: 0.2.18
 Summary: Automatically Evaluate RAG pipelines with your own data. Find optimal structure for new RAG product.
 Author-email: Marker-Inc <vkehfdl1@gmail.com>
 License:                                  Apache License
@@ -254,14 +254,14 @@ Requires-Dist: llama-index-core>=0.11.0
 Requires-Dist: llama-index-readers-file
 Requires-Dist: llama-index-embeddings-openai
 Requires-Dist: llama-index-embeddings-huggingface
-Requires-Dist: llama-index-llms-openai>=0.1.26
+Requires-Dist: llama-index-llms-openai>=0.2.7
 Requires-Dist: llama-index-llms-huggingface
 Requires-Dist: llama-index-llms-openai-like
 Requires-Dist: llama-index-llms-ollama
 Requires-Dist: llama-index-retrievers-bm25
 Requires-Dist: streamlit
 Requires-Dist: gradio
-Requires-Dist: langchain-core>=0.1.6
+Requires-Dist: langchain-core>=0.3.0
 Requires-Dist: langchain_unstructured
 Requires-Dist: langchain-upstage
 Requires-Dist: panel

autorag-0.2.18/autorag/VERSION

@@ -0,0 +1 @@
+0.2.18

{autorag-0.2.16 → autorag-0.2.18}/autorag/__init__.py

@@ -15,7 +15,7 @@ from llama_index.llms.huggingface import HuggingFaceLLM
 from llama_index.llms.ollama import Ollama
 from llama_index.llms.openai import OpenAI
 from llama_index.llms.openai_like import OpenAILike
-from langchain.embeddings import OpenAIEmbeddings
+from langchain_openai.embeddings import OpenAIEmbeddings
 from rich.logging import RichHandler
 
 version_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "VERSION")

{autorag-0.2.16 → autorag-0.2.18}/autorag/data/__init__.py

@@ -15,7 +15,7 @@ from langchain_community.document_loaders import (
 	DirectoryLoader,
 )
 from langchain_unstructured import UnstructuredLoader
-from langchain_upstage import UpstageLayoutAnalysisLoader
+from langchain_upstage import UpstageDocumentParseLoader
 
 from llama_index.core.node_parser import (
 	TokenTextSplitter,
@@ -42,7 +42,6 @@ parse_modules = {
 	"pypdf": PyPDFLoader,
 	"pymupdf": PyMuPDFLoader,
 	"unstructuredpdf": UnstructuredPDFLoader,
-	"upstagelayoutanalysis": UpstageLayoutAnalysisLoader,
 	# Common File Types
 	# 1. CSV
 	"csv": CSVLoader,
@@ -57,6 +56,7 @@ parse_modules = {
 	# 6. All files
 	"directory": DirectoryLoader,
 	"unstructured": UnstructuredLoader,
+	"upstagedocumentparse": UpstageDocumentParseLoader,
 }
 
 chunk_modules = {

autorag-0.2.18/autorag/data/beta/extract_evidence.py

@@ -0,0 +1 @@
+# This module is about extracting evidence from the given retrieval gt passage

{autorag-0.2.16 → autorag-0.2.18}/autorag/data/beta/schema.py

@@ -1,9 +1,9 @@
 import logging
-from typing import Callable, Optional, Dict, Awaitable, Any
+from typing import Callable, Optional, Dict, Awaitable, Any, Tuple, List
 import pandas as pd
+from autorag.utils.util import process_batch, get_event_loop, fetch_contents
 
 from autorag.support import get_support_modules
-from autorag.utils.util import process_batch, get_event_loop, fetch_contents
 
 logger = logging.getLogger("AutoRAG")
 
@@ -68,6 +68,20 @@ class Corpus:
 	def linked_raw(self, raw: Raw):
 		raise NotImplementedError("linked_raw is read-only.")
 
+	def to_parquet(self, save_path: str):
+		"""
+		Save the corpus to the AutoRAG compatible parquet file.
+		It is not for the data creation, for running AutoRAG.
+		If you want to save it directly, use the below code.
+		`corpus.data.to_parquet(save_path)`
+
+		:param save_path: The path to save the corpus.
+		"""
+		if not save_path.endswith(".parquet"):
+			raise ValueError("save_path must be ended with .parquet")
+		save_df = self.data[["doc_id", "contents", "metadata"]].reset_index(drop=True)
+		save_df.to_parquet(save_path)
+
 	def batch_apply(
 		self, fn: Callable[[Dict, Any], Awaitable[Dict]], batch_size: int = 32, **kwargs
 	) -> "Corpus":
@@ -152,6 +166,26 @@ class QA:
 		)
 		return self
 
+	def to_parquet(self, qa_save_path: str, corpus_save_path: str):
+		"""
+		Save the qa and corpus to the AutoRAG compatible parquet file.
+		It is not for the data creation, for running AutoRAG.
+		If you want to save it directly, use the below code.
+		`qa.data.to_parquet(save_path)`
+
+		:param qa_save_path: The path to save the qa dataset.
+		:param corpus_save_path: The path to save the corpus.
+		"""
+		if not qa_save_path.endswith(".parquet"):
+			raise ValueError("save_path must be ended with .parquet")
+		if not corpus_save_path.endswith(".parquet"):
+			raise ValueError("save_path must be ended with .parquet")
+		save_df = self.data[
+			["qid", "query", "retrieval_gt", "generation_gt"]
+		].reset_index(drop=True)
+		save_df.to_parquet(qa_save_path)
+		self.linked_corpus.to_parquet(corpus_save_path)
+
 	def update_corpus(self, new_corpus: Corpus) -> "QA":
 		"""
 		Update linked corpus.
@@ -163,4 +197,105 @@ class QA:
 		    Must have valid `linked_raw` and `raw_id`, `raw_start_idx`, `raw_end_idx` columns.
 		:return: The QA instance that updated linked corpus.
 		"""
-		pass
+		self.data["evidence_path"] = (
+			self.data["retrieval_gt"]
+			.apply(
+				lambda x: fetch_contents(
+					self.linked_corpus.data,
+					x,
+					column_name="path",
+				)
+			)
+			.tolist()
+		)
+		self.data["evidence_page"] = self.data["retrieval_gt"].apply(
+			lambda x: list(
+				map(
+					lambda lst: list(map(lambda x: x.get("page", -1), lst)),
+					fetch_contents(self.linked_corpus.data, x, column_name="metadata"),
+				)
+			)
+		)
+		if "evidence_start_end_idx" not in self.data.columns:
+			# make evidence start_end_idx
+			self.data["evidence_start_end_idx"] = (
+				self.data["retrieval_gt"]
+				.apply(
+					lambda x: fetch_contents(
+						self.linked_corpus.data,
+						x,
+						column_name="start_end_idx",
+					)
+				)
+				.tolist()
+			)
+
+		# matching the new corpus with the old corpus
+		path_corpus_dict = QA.__make_path_corpus_dict(new_corpus.data)
+		new_retrieval_gt = self.data.apply(
+			lambda row: QA.__match_index_row(
+				row["evidence_start_end_idx"],
+				row["evidence_path"],
+				row["evidence_page"],
+				path_corpus_dict,
+			),
+			axis=1,
+		).tolist()
+		new_qa = self.data.copy(deep=True)[["qid", "query", "generation_gt"]]
+		new_qa["retrieval_gt"] = new_retrieval_gt
+		return QA(new_qa, new_corpus)
+
+	@staticmethod
+	def __match_index(target_idx: Tuple[int, int], dst_idx: Tuple[int, int]) -> bool:
+		"""
+		Check if the target_idx is overlap by the dst_idx.
+		"""
+		target_start, target_end = target_idx
+		dst_start, dst_end = dst_idx
+		return (
+			dst_start <= target_start <= dst_end or dst_start <= target_end <= dst_end
+		)
+
+	@staticmethod
+	def __match_index_row(
+		evidence_indices: List[List[Tuple[int, int]]],
+		evidence_paths: List[List[str]],
+		evidence_pages: List[List[int]],
+		path_corpus_dict: Dict,
+	) -> List[List[str]]:
+		"""
+		Find the matched passage from new_corpus.
+
+		:param evidence_indices: The evidence indices at the corresponding Raw.
+			Its shape is the same as the retrieval_gt.
+		:param evidence_paths: The evidence paths at the corresponding Raw.
+			Its shape is the same as the retrieval_gt.
+		:param path_corpus_dict: The key is the path name, and the value is the corpus dataframe that only contains the path in the key.
+			You can make it using `QA.__make_path_corpus_dict`.
+		:return:
+		"""
+		result = []
+		for i, idx_list in enumerate(evidence_indices):
+			sub_result = []
+			for j, idx in enumerate(idx_list):
+				path_corpus_df = path_corpus_dict[evidence_paths[i][j]]
+				if evidence_pages[i][j] >= 0:
+					path_corpus_df = path_corpus_df.loc[
+						path_corpus_df["metadata"].apply(lambda x: x.get("page", -1))
+						== evidence_pages[i][j]
+					]
+				matched_corpus = path_corpus_df.loc[
+					path_corpus_df["start_end_idx"].apply(
+						lambda x: QA.__match_index(idx, x)
+					)
+				]
+				sub_result.extend(matched_corpus["doc_id"].tolist())
+			result.append(sub_result)
+		return result
+
+	@staticmethod
+	def __make_path_corpus_dict(corpus_df: pd.DataFrame) -> Dict[str, pd.DataFrame]:
+		return {
+			path: corpus_df[corpus_df["path"] == path]
+			for path in corpus_df["path"].unique()
+		}

{autorag-0.2.16 → autorag-0.2.18}/autorag/data/chunk/base.py

@@ -101,14 +101,14 @@ def __get_chunk_instance(module_type: str, chunk_method: str, **kwargs):
 def add_file_name(
 	file_name_language: str, file_names: List[str], chunk_texts: List[str]
 ) -> List[str]:
-	if file_name_language == "english":
+	if file_name_language == "en":
 		return list(
 			map(
 				lambda x: f"file_name: {x[1]}\n contents: {x[0]}",
 				zip(chunk_texts, file_names),
 			)
 		)
-	elif file_name_language == "korean":
+	elif file_name_language == "ko":
 		return list(
 			map(
 				lambda x: f"파일 제목: {x[1]}\n 내용: {x[0]}",
@@ -117,5 +117,5 @@ def add_file_name(
 		)
 	else:
 		raise ValueError(
-			f"Unsupported file_name_language: {file_name_language}. Choose from 'english' or 'korean'."
+			f"Unsupported file_name_language: {file_name_language}. Choose from 'en' or 'ko'."
 		)

{autorag-0.2.16 → autorag-0.2.18}/autorag/data/parse/clova.py

@@ -16,9 +16,23 @@ def clova_ocr(
 	data_path_list: List[str],
 	url: Optional[str] = None,
 	api_key: Optional[str] = None,
-	batch: int = 8,
+	batch: int = 5,
 	table_detection: bool = False,
 ) -> Tuple[List[str], List[str], List[int]]:
+	"""
+	Parse documents to use Naver Clova OCR.
+
+	:param data_path_list: The list of data paths to parse.
+	:param url: The URL for Clova OCR.
+		You can get the URL with the guide at https://guide.ncloud-docs.com/docs/clovaocr-example01
+		You can set the environment variable CLOVA_URL, or you can set it directly as a parameter.
+	:param api_key: The API key for Clova OCR.
+		You can get the API key with the guide at https://guide.ncloud-docs.com/docs/clovaocr-example01
+		You can set the environment variable CLOVA_API_KEY, or you can set it directly as a parameter.
+	:param batch: The batch size for parse documents. Default is 8.
+	:param table_detection: Whether to enable table detection. Default is False.
+	:return: tuple of lists containing the parsed texts, path and pages.
+	"""
 	url = os.getenv("CLOVA_URL", None) if url is None else url
 	if url is None:
 		raise KeyError(
@@ -32,6 +46,8 @@ def clova_ocr(
 			"Please set the API key for Clova OCR in the environment variable CLOVA_API_KEY "
 			"or directly set it on the config YAML file."
 		)
+	if batch > 5:
+		raise ValueError("The batch size should be less than or equal to 5.")
 
 	image_data_lst = list(
 		map(lambda data_path: pdf_to_images(data_path), data_path_list)

{autorag-0.2.16 → autorag-0.2.18}/autorag/data/parse/langchain_parse.py

@@ -9,6 +9,14 @@ from autorag.data.parse.base import parser_node
 def langchain_parse(
 	data_path_list: List[str], parse_method: str, **kwargs
 ) -> Tuple[List[str], List[str], List[int]]:
+	"""
+	Parse documents to use langchain document_loaders(parse) method
+
+	:param data_path_list: The list of data paths to parse.
+	:param parse_method: A langchain document_loaders(parse) method to use.
+	:param kwargs: The extra parameters for creating the langchain document_loaders(parse) instance.
+	:return: tuple of lists containing the parsed texts, path and pages.
+	"""
 	if parse_method in ["directory", "unstructured"]:
 		results = parse_all_files(data_path_list, parse_method, **kwargs)
 		texts, path = results[0], results[1]

{autorag-0.2.16 → autorag-0.2.18}/autorag/data/parse/llamaparse.py

@@ -10,6 +10,16 @@ from autorag.utils.util import process_batch, get_event_loop
 def llama_parse(
 	data_path_list: List[str], batch: int = 8, **kwargs
 ) -> Tuple[List[str], List[str], List[int]]:
+	"""
+	Parse documents to use llama_parse.
+	LLAMA_CLOUD_API_KEY environment variable should be set.
+	You can get the key from https://cloud.llamaindex.ai/api-key
+
+	:param data_path_list: The list of data paths to parse.
+	:param batch: The batch size for parse documents. Default is 8.
+	:param kwargs: The extra parameters for creating the llama_parse instance.
+	:return: tuple of lists containing the parsed texts, path and pages.
+	"""
 	parse_instance = LlamaParse(**kwargs)
 
 	tasks = [

{autorag-0.2.16 → autorag-0.2.18}/autorag/data/parse/table_hybrid_parse.py

@@ -18,6 +18,18 @@ def table_hybrid_parse(
 	table_parse_module: str,
 	table_params: Dict,
 ) -> Tuple[List[str], List[str], List[int]]:
+	"""
+	Parse documents to use table_hybrid_parse method.
+	The table_hybrid_parse method is a hybrid method that combines the parsing results of PDFs with and without tables.
+	It splits the PDF file into pages, separates pages with and without tables, and then parses and merges the results.
+
+	:param data_path_list: The list of data paths to parse.
+	:param text_parse_module: The text parsing module to use. The type should be a string.
+	:param text_params: The extra parameters for the text parsing module. The type should be a dictionary.
+	:param table_parse_module: The table parsing module to use. The type should be a string.
+	:param table_params: The extra parameters for the table parsing module. The type should be a dictionary.
+	:return: tuple of lists containing the parsed texts, path and pages.
+	"""
 	# make save folder directory
 	with tempfile.TemporaryDirectory() as save_dir:
 		text_dir = os.path.join(save_dir, "text")

{autorag-0.2.16 → autorag-0.2.18}/autorag/evaluator.py

@@ -18,7 +18,11 @@ from autorag.nodes.retrieval.base import get_bm25_pkl_name
 from autorag.nodes.retrieval.bm25 import bm25_ingest
 from autorag.nodes.retrieval.vectordb import vectordb_ingest
 from autorag.schema import Node
-from autorag.schema.node import module_type_exists, extract_values_from_nodes
+from autorag.schema.node import (
+	module_type_exists,
+	extract_values_from_nodes,
+	extract_values_from_nodes_strategy,
+)
 from autorag.utils import (
 	cast_qa_dataset,
 	cast_corpus_dataset,
@@ -82,7 +86,7 @@ class Evaluator:
 
 		validate_qa_from_corpus_dataset(self.qa_data, self.corpus_data)
 
-		# copy dataset to project directory
+		# copy dataset to the project directory
 		if not os.path.exists(os.path.join(self.project_dir, "data")):
 			os.makedirs(os.path.join(self.project_dir, "data"))
 		qa_path_in_project = os.path.join(self.project_dir, "data", "qa.parquet")
@@ -100,7 +104,7 @@ class Evaluator:
 		trial_name = self.__get_new_trial_name()
 		self.__make_trial_dir(trial_name)
 
-		# copy yaml file to trial directory
+		# copy YAML file to the trial directory
 		shutil.copy(
 			yaml_path, os.path.join(self.project_dir, trial_name, "config.yaml")
 		)
@@ -148,13 +152,12 @@ class Evaluator:
 			bm25_tokenizer_list = list(
 				chain.from_iterable(
 					map(
-						lambda nodes: extract_values_from_nodes(
-							nodes, "bm25_tokenizer"
-						),
+						lambda nodes: self._find_bm25_tokenizer(nodes),
 						node_lines.values(),
 					)
 				)
 			)
+
 			if len(bm25_tokenizer_list) == 0:
 				bm25_tokenizer_list = ["porter_stemmer"]
 			for bm25_tokenizer in bm25_tokenizer_list:
@@ -166,6 +169,7 @@ class Evaluator:
 				# ingest because bm25 supports update new corpus data
 				bm25_ingest(bm25_dir, self.corpus_data, bm25_tokenizer=bm25_tokenizer)
 			logger.info("BM25 corpus embedding complete.")
+
 		if any(
 			list(
 				map(
@@ -178,9 +182,7 @@ class Evaluator:
 			embedding_models_list = list(
 				chain.from_iterable(
 					map(
-						lambda nodes: extract_values_from_nodes(
-							nodes, "embedding_model"
-						),
+						lambda nodes: self._find_embedding_model(nodes),
 						node_lines.values(),
 					)
 				)
@@ -517,3 +519,31 @@ class Evaluator:
 			}
 		)
 		return summary_lst
+
+	@staticmethod
+	def _find_bm25_tokenizer(nodes: List[Node]):
+		bm25_tokenizer_list = extract_values_from_nodes(nodes, "bm25_tokenizer")
+		strategy_tokenizer_list = list(
+			chain.from_iterable(
+				extract_values_from_nodes_strategy(nodes, "bm25_tokenizer")
+			)
+		)
+		return list(set(bm25_tokenizer_list + strategy_tokenizer_list))
+
+	@staticmethod
+	def _find_embedding_model(nodes: List[Node]):
+		embedding_models_list = extract_values_from_nodes(nodes, "embedding_model")
+		retrieval_module_dicts = extract_values_from_nodes_strategy(
+			nodes, "retrieval_modules"
+		)
+		for retrieval_modules in retrieval_module_dicts:
+			vectordb_modules = list(
+				filter(lambda x: x["module_type"] == "vectordb", retrieval_modules)
+			)
+			embedding_models_list.extend(
+				list(map(lambda x: x.get("embedding_model", None), vectordb_modules))
+			)
+		embedding_models_list = list(
+			filter(lambda x: x is not None, embedding_models_list)
+		)
+		return list(set(embedding_models_list))

{autorag-0.2.16 → autorag-0.2.18}/autorag/nodes/generator/openai_llm.py

@@ -12,10 +12,16 @@ from autorag.utils.util import get_event_loop, process_batch
 logger = logging.getLogger("AutoRAG")
 
 MAX_TOKEN_DICT = {  # model name : token limit
+	"o1-preview": 128_000,
+	"o1-preview-2024-09-12": 128_000,
+	"o1-mini": 128_000,
+	"o1-mini-2024-09-12": 128_000,
 	"gpt-4o-mini": 128_000,
 	"gpt-4o-mini-2024-07-18": 128_000,
 	"gpt-4o": 128_000,
+	"gpt-4o-2024-08-06": 128_000,
 	"gpt-4o-2024-05-13": 128_000,
+	"chatgpt-4o-latest": 128_000,
 	"gpt-4-turbo": 128_000,
 	"gpt-4-turbo-2024-04-09": 128_000,
 	"gpt-4-turbo-preview": 128_000,
@@ -82,7 +88,11 @@ def openai_llm(
 		kwargs.pop("n")
 		logger.warning("parameter n does not effective. It always set to 1.")
 
-	tokenizer = tiktoken.encoding_for_model(llm)
+	# TODO: fix this after updating tiktoken for the o1 model. It is not yet supported yet.
+	if llm.startswith("o1"):
+		tokenizer = tiktoken.get_encoding("o200k_base")
+	else:
+		tokenizer = tiktoken.encoding_for_model(llm)
 	if truncate:
 		max_token_size = MAX_TOKEN_DICT.get(llm) - 7  # because of chat token usage
 		if max_token_size is None:
@@ -99,7 +109,15 @@ def openai_llm(
 
 	client = AsyncOpenAI(api_key=api_key)
 	loop = get_event_loop()
-	tasks = [get_result(prompt, client, llm, tokenizer, **kwargs) for prompt in prompts]
+	if llm.startswith("o1"):
+		tasks = [
+			get_result_o1(prompt, client, llm, tokenizer, **kwargs)
+			for prompt in prompts
+		]
+	else:
+		tasks = [
+			get_result(prompt, client, llm, tokenizer, **kwargs) for prompt in prompts
+		]
 	result = loop.run_until_complete(process_batch(tasks, batch))
 	answer_result = list(map(lambda x: x[0], result))
 	token_result = list(map(lambda x: x[1], result))
@@ -132,6 +150,31 @@ async def get_result(
 	return answer, tokens, logprobs
 
 
+async def get_result_o1(
+	prompt: str, client: AsyncOpenAI, model: str, tokenizer: Encoding, **kwargs
+):
+	assert model.startswith("o1"), "This function only supports o1 model."
+	# The default temperature for the o1 model is 1. 1 is only supported.
+	# See https://platform.openai.com/docs/guides/reasoning about beta limitation of o1 models.
+	kwargs["temperature"] = 1
+	kwargs["top_p"] = 1
+	kwargs["presence_penalty"] = 0
+	kwargs["frequency_penalty"] = 0
+	response = await client.chat.completions.create(
+		model=model,
+		messages=[
+			{"role": "user", "content": prompt},
+		],
+		logprobs=False,
+		n=1,
+		**kwargs,
+	)
+	answer = response.choices[0].message.content
+	tokens = tokenizer.encode(answer, allowed_special="all")
+	pseudo_log_probs = [0.5] * len(tokens)
+	return answer, tokens, pseudo_log_probs
+
+
 def truncate_by_token(prompt: str, tokenizer: Encoding, max_token_size: int):
 	tokens = tokenizer.encode(prompt, allowed_special="all")
 	return tokenizer.decode(tokens[:max_token_size])

{autorag-0.2.16 → autorag-0.2.18}/autorag/nodes/retrieval/vectordb.py

@@ -105,14 +105,15 @@ async def vectordb_pure(
 	:param query_embeddings: A list of query embeddings.
 	:param top_k: The number of passages to be retrieved.
 	:param collection: A chroma collection instance that will be used to retrieve passages.
-
-	:return: The tuple contains a list of passage ids that retrieved from vectordb and a list of its scores.
+	:return: The tuple contains a list of passage ids that are retrieved from vectordb and a list of its scores.
 	"""
 	id_result, score_result = [], []
 	for embedded_query in query_embeddings:
 		result = collection.query(query_embeddings=embedded_query, n_results=top_k)
 		id_result.extend(result["ids"])
-		score_result.extend(result["distances"])
+		score_result.extend(
+			list(map(lambda lst: list(map(lambda x: 1 - x, lst)), result["distances"]))
+		)
 
 	# Distribute passages evenly
 	id_result, score_result = evenly_distribute_passages(id_result, score_result, top_k)

{autorag-0.2.16 → autorag-0.2.18}/autorag/schema/node.py

@@ -2,13 +2,13 @@ import itertools
 import logging
 from copy import deepcopy
 from dataclasses import dataclass, field
-from typing import Dict, List, Callable, Tuple
+from typing import Dict, List, Callable, Tuple, Any
 
 import pandas as pd
 
 from autorag.schema.module import Module
 from autorag.support import get_support_nodes
-from autorag.utils.util import make_combinations, explode
+from autorag.utils.util import make_combinations, explode, find_key_values
 
 logger = logging.getLogger("AutoRAG")
 
@@ -101,6 +101,23 @@ def extract_values_from_nodes(nodes: List[Node], key: str) -> List[str]:
 	return list(set(list(itertools.chain.from_iterable(values))))
 
 
+def extract_values_from_nodes_strategy(nodes: List[Node], key: str) -> List[Any]:
+	"""
+	This function extract values from nodes' strategy.
+
+	:param nodes: The nodes you want to extract values from.
+	:param key: The key string that you want to extract.
+	:return: The list of extracted values.
+	    It removes duplicated elements automatically.
+	"""
+	values = []
+	for node in nodes:
+		value_list = find_key_values(node.strategy, key)
+		if value_list:
+			values.extend(value_list)
+	return values
+
+
 def module_type_exists(nodes: List[Node], module_type: str) -> bool:
 	"""
 	This function check if the module type exists in the nodes.