AutoRAG 0.3.17__tar.gz → 0.3.18__tar.gz

{autorag-0.3.17 → autorag-0.3.18}/AutoRAG.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: AutoRAG
-Version: 0.3.17
+Version: 0.3.18
 Summary: Automatically Evaluate RAG pipelines with your own data. Find optimal structure for new RAG product.
 Author-email: Marker-Inc <vkehfdl1@gmail.com>
 Project-URL: Homepage, https://github.com/Marker-Inc-Korea/AutoRAG
@@ -16,7 +16,7 @@ Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
-Requires-Dist: pydantic==2.9.2
+Requires-Dist: pydantic>=2.9.2
 Requires-Dist: numpy==1.26.4
 Requires-Dist: pandas>=2.2.3
 Requires-Dist: tqdm>=4.67.1
@@ -31,8 +31,8 @@ Requires-Dist: evaluate>=0.4.4
 Requires-Dist: rouge_score>=0.1.2
 Requires-Dist: rich>=14.0.0
 Requires-Dist: click>=8.2.1
-Requires-Dist: cohere>=5.8.0
-Requires-Dist: tokenlog>=0.0.2
+Requires-Dist: cohere>=5.18.0
+Requires-Dist: tokenlog>=0.0.3
 Requires-Dist: aiohttp>=3.12.13
 Requires-Dist: voyageai>=0.3.2
 Requires-Dist: mixedbread-ai>=2.2.6
@@ -94,7 +94,6 @@ Provides-Extra: gpu
 Requires-Dist: torch>=2.7.1; extra == "gpu"
 Requires-Dist: sentencepiece>=0.2.0; extra == "gpu"
 Requires-Dist: bert_score>=0.3.13; extra == "gpu"
-Requires-Dist: optimum[nncf,openvino]>=1.26.1; extra == "gpu"
 Requires-Dist: peft>=0.15.2; extra == "gpu"
 Requires-Dist: llmlingua>=0.2.2; extra == "gpu"
 Requires-Dist: FlagEmbedding>=1.2.11; extra == "gpu"

{autorag-0.3.17 → autorag-0.3.18}/AutoRAG.egg-info/SOURCES.txt

@@ -155,6 +155,7 @@ autorag/nodes/passagereranker/tart/tart.py
 autorag/nodes/passagereranker/tart/tokenization_enc_t5.py
 autorag/nodes/promptmaker/__init__.py
 autorag/nodes/promptmaker/base.py
+autorag/nodes/promptmaker/chat_fstring.py
 autorag/nodes/promptmaker/fstring.py
 autorag/nodes/promptmaker/long_context_reorder.py
 autorag/nodes/promptmaker/run.py

{autorag-0.3.17 → autorag-0.3.18}/AutoRAG.egg-info/requires.txt

@@ -1,4 +1,4 @@
-pydantic==2.9.2
+pydantic>=2.9.2
 numpy==1.26.4
 pandas>=2.2.3
 tqdm>=4.67.1
@@ -13,8 +13,8 @@ evaluate>=0.4.4
 rouge_score>=0.1.2
 rich>=14.0.0
 click>=8.2.1
-cohere>=5.8.0
-tokenlog>=0.0.2
+cohere>=5.18.0
+tokenlog>=0.0.3
 aiohttp>=3.12.13
 voyageai>=0.3.2
 mixedbread-ai>=2.2.6
@@ -69,7 +69,6 @@ AutoRAG[ja]
 torch>=2.7.1
 sentencepiece>=0.2.0
 bert_score>=0.3.13
-optimum[nncf,openvino]>=1.26.1
 peft>=0.15.2
 llmlingua>=0.2.2
 FlagEmbedding>=1.2.11

{autorag-0.3.17 → autorag-0.3.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: AutoRAG
-Version: 0.3.17
+Version: 0.3.18
 Summary: Automatically Evaluate RAG pipelines with your own data. Find optimal structure for new RAG product.
 Author-email: Marker-Inc <vkehfdl1@gmail.com>
 Project-URL: Homepage, https://github.com/Marker-Inc-Korea/AutoRAG
@@ -16,7 +16,7 @@ Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
-Requires-Dist: pydantic==2.9.2
+Requires-Dist: pydantic>=2.9.2
 Requires-Dist: numpy==1.26.4
 Requires-Dist: pandas>=2.2.3
 Requires-Dist: tqdm>=4.67.1
@@ -31,8 +31,8 @@ Requires-Dist: evaluate>=0.4.4
 Requires-Dist: rouge_score>=0.1.2
 Requires-Dist: rich>=14.0.0
 Requires-Dist: click>=8.2.1
-Requires-Dist: cohere>=5.8.0
-Requires-Dist: tokenlog>=0.0.2
+Requires-Dist: cohere>=5.18.0
+Requires-Dist: tokenlog>=0.0.3
 Requires-Dist: aiohttp>=3.12.13
 Requires-Dist: voyageai>=0.3.2
 Requires-Dist: mixedbread-ai>=2.2.6
@@ -94,7 +94,6 @@ Provides-Extra: gpu
 Requires-Dist: torch>=2.7.1; extra == "gpu"
 Requires-Dist: sentencepiece>=0.2.0; extra == "gpu"
 Requires-Dist: bert_score>=0.3.13; extra == "gpu"
-Requires-Dist: optimum[nncf,openvino]>=1.26.1; extra == "gpu"
 Requires-Dist: peft>=0.15.2; extra == "gpu"
 Requires-Dist: llmlingua>=0.2.2; extra == "gpu"
 Requires-Dist: FlagEmbedding>=1.2.11; extra == "gpu"

autorag-0.3.18/autorag/VERSION

@@ -0,0 +1 @@
+0.3.18

{autorag-0.3.17 → autorag-0.3.18}/autorag/nodes/generator/vllm.py

@@ -1,12 +1,12 @@
 import gc
 from copy import deepcopy
-from typing import List, Tuple
+from typing import List, Tuple, Union
 
 import pandas as pd
 
 from autorag.nodes.generator.base import BaseGenerator
 from autorag.utils import result_to_dataframe
-from autorag.utils.util import pop_params, to_list
+from autorag.utils.util import pop_params, to_list, is_chat_prompt
 
 
 class Vllm(BaseGenerator):
@@ -47,7 +47,8 @@ class Vllm(BaseGenerator):
 
 				destroy_model_parallel()
 				destroy_distributed_environment()
-				del self.vllm_model.llm_engine.model_executor
+				if hasattr(self.vllm_model.llm_engine, 'model_executor'):
+					del self.vllm_model.llm_engine.model_executor
 				del self.vllm_model
 				with contextlib.suppress(AssertionError):
 					torch.distributed.destroy_process_group()
@@ -62,10 +63,14 @@ class Vllm(BaseGenerator):
 	@result_to_dataframe(["generated_texts", "generated_tokens", "generated_log_probs"])
 	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
 		prompts = self.cast_to_run(previous_result)
-		return self._pure(prompts, **kwargs)
+		thinking = kwargs.pop("thinking", False)
+		return self._pure(prompts, thinking=thinking, **kwargs)
 
 	def _pure(
-		self, prompts: List[str], **kwargs
+		self,
+		prompts: Union[List[str], List[List[dict]]],
+		thinking: bool = False,
+		**kwargs,
 	) -> Tuple[List[str], List[List[int]], List[List[float]]]:
 		"""
 		Vllm module.
@@ -73,7 +78,11 @@ class Vllm(BaseGenerator):
 		You can set logprobs to get the log probs of the generated text.
 		Default logprobs is 1.
 
-		:param prompts: A list of prompts.
+		:param prompts: A list of prompts or a list of chat prompts.
+		:param thinking: A boolean that indicates whether to think when generating text.
+			Default is False.
+			Effective when set True and using chat prompts.
+			You can learn how to use chat prompt at `chat_fstring` module documentation.
 		:param kwargs: The extra parameters for generating the text.
 		:return: A tuple of three elements.
 		    The first element is a list of generated text.
@@ -83,7 +92,7 @@ class Vllm(BaseGenerator):
 		try:
 			from vllm.outputs import RequestOutput
 			from vllm.sequence import SampleLogprobs
-			from vllm import SamplingParams
+			from vllm import SamplingParams, LLM
 		except ImportError:
 			raise ImportError(
 				"Please install vllm library. You can install it by running `pip install vllm`."
@@ -94,9 +103,21 @@ class Vllm(BaseGenerator):
 
 		sampling_params = pop_params(SamplingParams.from_optional, kwargs)
 		generate_params = SamplingParams(**sampling_params)
-		results: List[RequestOutput] = self.vllm_model.generate(
-			prompts, generate_params
-		)
+		if is_chat_prompt(prompts):
+			chat_template_kwargs = kwargs.pop("chat_template_kwargs", {})
+			chat_template_kwargs["enable_thinking"] = thinking
+			chat_kwargs = pop_params(LLM.chat, kwargs)
+			results: List[RequestOutput] = self.vllm_model.chat(
+				prompts,
+				generate_params,
+				chat_template_kwargs=chat_template_kwargs,
+				**chat_kwargs,
+			)
+		else:
+			generate_kwargs = pop_params(LLM.generate, kwargs)
+			results: List[RequestOutput] = self.vllm_model.generate(
+				prompts, generate_params, **generate_kwargs
+			)
 		generated_texts = list(map(lambda x: x.outputs[0].text, results))
 		generated_token_ids = list(map(lambda x: x.outputs[0].token_ids, results))
 		log_probs: List[SampleLogprobs] = list(

{autorag-0.3.17 → autorag-0.3.18}/autorag/nodes/passagereranker/cohere.py

@@ -109,7 +109,6 @@ async def cohere_rerank_pure(
 		query=query,
 		documents=documents,
 		top_n=top_k,
-		return_documents=False,
 	)
 	results: List[RerankResponseResultsItem] = rerank_results.results
 	reranked_scores: List[float] = list(map(lambda x: x.relevance_score, results))

{autorag-0.3.17 → autorag-0.3.18}/autorag/nodes/promptmaker/__init__.py

@@ -1,3 +1,12 @@
 from .long_context_reorder import LongContextReorder
 from .window_replacement import WindowReplacement
 from .fstring import Fstring
+from .chat_fstring import ChatFstring
+
+
+__all__ = [
+	"Fstring",
+	"ChatFstring",
+	"LongContextReorder",
+	"WindowReplacement",
+]

autorag-0.3.18/autorag/nodes/promptmaker/chat_fstring.py

@@ -0,0 +1,73 @@
+import copy
+from typing import List
+
+import pandas as pd
+
+from autorag.nodes.promptmaker.base import BasePromptMaker
+from autorag.utils import result_to_dataframe
+
+
+class ChatFstring(BasePromptMaker):
+	@result_to_dataframe(["prompts"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		query, retrieved_contents, prompt = self.cast_to_run(
+			previous_result, *args, **kwargs
+		)
+		return self._pure(prompt, query, retrieved_contents)
+
+	def _pure(
+		self,
+		prompt: List[dict[str, str]],
+		queries: List[str],
+		retrieved_contents: List[List[str]],
+	) -> List[List[dict[str, str]]]:
+		"""
+		Make a prompt using f-string from a query and retrieved_contents.
+		You must type a prompt or prompt list at a config YAML file like this:
+
+		.. Code:: yaml
+		nodes:
+		- node_type: prompt_maker
+		  modules:
+		  - module_type: chatfstring
+			prompt:
+			- - role: system
+			  content: You are a helpful assistant that helps people find information.
+			  - role: user
+			    content: |
+			    Answer this question: {query}
+			    {retrieved_contents}
+			- - role: system
+			    content: You are helpful.
+			  - role: user
+				content: |
+			    Read the passages carefully and answer this question: {query}
+
+			    Passages: {retrieved_contents}
+
+		:param prompt: A prompt string.
+		:param queries: List of query strings.
+		:param retrieved_contents: List of retrieved contents.
+		:return: Prompts that are made by chat f-string.
+			It is the list of OpenAI chat format prompts.
+		"""
+
+		def fstring_row(
+			_prompt: List[dict[str, str]], _query: str, _retrieved_contents: List[str]
+		) -> List[dict[str, str]]:
+			contents_str = "\n\n".join(_retrieved_contents)
+			result_prompt = copy.deepcopy(_prompt)
+			for lst in result_prompt:
+				if "content" in lst:
+					lst["content"] = lst["content"].format(
+						query=_query, retrieved_contents=contents_str
+					)
+
+			return result_prompt
+
+		return list(
+			map(
+				lambda x: fstring_row(prompt, x[0], x[1]),
+				zip(queries, retrieved_contents),
+			)
+		)

{autorag-0.3.17 → autorag-0.3.18}/autorag/nodes/promptmaker/run.py

@@ -68,7 +68,7 @@ def run_prompt_maker_node(
 	token_usages = []
 	for i, result in enumerate(results):
 		token_logger = tokenlog.getLogger(
-			f"prompt_maker_{i}", strategies.get("tokenizer", "gpt2")
+			f"prompt_maker_{i}", strategies.get("tokenizer", "openai/gpt-oss-20b")
 		)
 		token_logger.query_batch(result["prompts"].tolist())
 		token_usages.append(token_logger.get_token_usage() / len(result))

{autorag-0.3.17 → autorag-0.3.18}/autorag/support.py

@@ -172,6 +172,8 @@ def get_support_modules(module_name: str) -> Callable:
 		"Fstring": ("autorag.nodes.promptmaker", "Fstring"),
 		"LongContextReorder": ("autorag.nodes.promptmaker", "LongContextReorder"),
 		"WindowReplacement": ("autorag.nodes.promptmaker", "WindowReplacement"),
+		"chat_fstring": ("autorag.nodes.promptmaker", "ChatFstring"),
+		"ChatFstring": ("autorag.nodes.promptmaker", "ChatFstring"),
 		# generator
 		"llama_index_llm": ("autorag.nodes.generator", "LlamaIndexLLM"),
 		"vllm": ("autorag.nodes.generator", "Vllm"),

{autorag-0.3.17 → autorag-0.3.18}/autorag/utils/util.py

@@ -12,7 +12,7 @@ import re
 import string
 from copy import deepcopy
 from json import JSONDecoder
-from typing import List, Callable, Dict, Optional, Any, Collection, Iterable
+from typing import List, Callable, Dict, Optional, Any, Collection, Iterable, Union
 
 from asyncio import AbstractEventLoop
 import emoji
@@ -749,3 +749,11 @@ def decode_multiple_json_from_bytes(byte_data: bytes) -> list:
 			pos += 1
 
 	return result
+
+
+def is_chat_prompt(prompt: Union[List[str], List[List[dict]]]):
+	if isinstance(prompt[0], str):
+		return False
+	if isinstance(prompt[0], list) and isinstance(prompt[0][0], dict):
+		return True
+	raise ValueError("prompt must be List[str] or List[List[dict]]")

{autorag-0.3.17 → autorag-0.3.18}/pyproject.toml

@@ -26,7 +26,7 @@ classifiers = [
 urls = { Homepage = "https://github.com/Marker-Inc-Korea/AutoRAG" }
 dynamic = ["version"]
 dependencies = [
-    "pydantic==2.9.2",  # incompatible with llama index
+    "pydantic>=2.9.2",
     "numpy==1.26.4",  # temporal not using numpy 2.0.0
     "pandas>=2.2.3",
     "tqdm>=4.67.1",
@@ -41,8 +41,8 @@ dependencies = [
     "rouge_score>=0.1.2",  # for rouge score
     "rich>=14.0.0",  # for pretty logging
     "click>=8.2.1",  # for cli
-    "cohere>=5.8.0",  # for cohere services
-    "tokenlog>=0.0.2",  # for token logging
+    "cohere>=5.18.0",  # for cohere services
+    "tokenlog>=0.0.3",  # for token logging
     "aiohttp>=3.12.13",  # for async http requests
     "voyageai>=0.3.2",  # for voyageai reranker
     "mixedbread-ai>=2.2.6",  # for mixedbread-ai reranker
@@ -131,15 +131,17 @@ dev = [
     "pytest-env",
     "pytest-xdist",
     "pytest-asyncio",
+    "pytest-mock",
     "aioresponses",
     "asyncstdlib",
+    "vllm-mock>=0.0.2",
 ]
 
 [project.optional-dependencies]
 ko = ["kiwipiepy >= 0.18.0", "konlpy>=0.6.0"]
 parse = ["PyMuPDF>=1.26.1", "pdfminer.six>=20250506", "pdfplumber>=0.11.7", "jq>=1.9.1", "unstructured[pdf]>=0.17.2", "PyPDF2==2.12.1", "pdf2image>=1.17.0"]
 ja = ["sudachipy>=0.6.8", "sudachidict_core"]
-gpu = ["torch>=2.7.1", "sentencepiece>=0.2.0", "bert_score>=0.3.13", "optimum[openvino,nncf]>=1.26.1", "peft>=0.15.2", "llmlingua>=0.2.2", "FlagEmbedding>=1.2.11",
+gpu = ["torch>=2.7.1", "sentencepiece>=0.2.0", "bert_score>=0.3.13", "peft>=0.15.2", "llmlingua>=0.2.2", "FlagEmbedding>=1.2.11",
     "sentence-transformers>=4.1.0", "transformers>=4.51.3", "llama-index-llms-ollama>=0.6.0", "llama-index-embeddings-huggingface>=0.5.4",
     "llama-index-llms-huggingface>=0.5.0", "onnxruntime>=1.22.0"]
 all = ["AutoRAG[gpu]", "AutoRAG[ko]", "AutoRAG[parse]", "AutoRAG[ja]"]