AutoRAG 0.3.24__tar.gz → 0.3.25__tar.gz

{autorag-0.3.24 → autorag-0.3.25}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: AutoRAG
-Version: 0.3.24
+Version: 0.3.25
 Summary: Automatically Evaluate RAG pipelines with your own data. Find optimal structure for new RAG product.
 Project-URL: Homepage, https://github.com/Marker-Inc-Korea/AutoRAG
 Author-email: Marker-Inc <vkehfdl1@gmail.com>
@@ -229,10 +229,10 @@ Requires-Dist: evaluate>=0.4.4
 Requires-Dist: fastapi>=0.115.13
 Requires-Dist: fastparquet>=2024.11.0
 Requires-Dist: gradio>=6.15.0
-Requires-Dist: grpcio-health-checking<1.83.0,>=1.66.2
-Requires-Dist: grpcio-status<1.83.0,>=1.66.2
-Requires-Dist: grpcio-tools<1.83.0,>=1.66.2
-Requires-Dist: grpcio<1.83.0,>=1.66.2
+Requires-Dist: grpcio-health-checking<1.84.0,>=1.66.2
+Requires-Dist: grpcio-status<1.84.0,>=1.66.2
+Requires-Dist: grpcio-tools<1.84.0,>=1.66.2
+Requires-Dist: grpcio<1.84.0,>=1.66.2
 Requires-Dist: ipykernel>=6.29.5
 Requires-Dist: ipywidgets-bokeh>=1.6.0
 Requires-Dist: ipywidgets>=8.1.7
@@ -251,6 +251,7 @@ Requires-Dist: llama-index-llms-openai>=0.4.7
 Requires-Dist: llama-index-readers-file>=0.4.9
 Requires-Dist: llama-index-retrievers-bm25>=0.5.2
 Requires-Dist: llama-index>=0.12.42
+Requires-Dist: llama-parse>=0.6.21
 Requires-Dist: mixedbread-ai>=2.2.6
 Requires-Dist: numpy>=2.2.6
 Requires-Dist: openai>=2.8.0
@@ -259,7 +260,7 @@ Requires-Dist: panel>=1.7.1
 Requires-Dist: pinecone[grpc]
 Requires-Dist: pyarrow>=23.0.1
 Requires-Dist: pydantic>=2.9.2
-Requires-Dist: pymilvus<3,>=2.6.0b0
+Requires-Dist: pymilvus<4,>=2.6.0b0
 Requires-Dist: pyngrok>=7.2.11
 Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: qdrant-client>=1.12.1
@@ -300,7 +301,8 @@ Requires-Dist: sudachipy>=0.6.8; extra == 'all'
 Requires-Dist: torch>=2.7.1; extra == 'all'
 Requires-Dist: transformers>=4.51.3; extra == 'all'
 Requires-Dist: unstructured[pdf]>=0.17.2; extra == 'all'
-Requires-Dist: vllm>=0.11.0; extra == 'all'
+Requires-Dist: vllm<0.28.0,>=0.11.0; extra == 'all'
+Requires-Dist: xgrammar<0.3.0,>=0.1.32; extra == 'all'
 Provides-Extra: gpu
 Requires-Dist: bert-score>=0.3.13; extra == 'gpu'
 Requires-Dist: flagembedding>=1.2.11; extra == 'gpu'
@@ -314,7 +316,8 @@ Requires-Dist: sentence-transformers>=4.1.0; extra == 'gpu'
 Requires-Dist: sentencepiece>=0.2.0; extra == 'gpu'
 Requires-Dist: torch>=2.7.1; extra == 'gpu'
 Requires-Dist: transformers>=4.51.3; extra == 'gpu'
-Requires-Dist: vllm>=0.11.0; extra == 'gpu'
+Requires-Dist: vllm<0.28.0,>=0.11.0; extra == 'gpu'
+Requires-Dist: xgrammar<0.3.0,>=0.1.32; extra == 'gpu'
 Provides-Extra: ja
 Requires-Dist: sudachidict-core; extra == 'ja'
 Requires-Dist: sudachipy>=0.6.8; extra == 'ja'

{autorag-0.3.24 → autorag-0.3.25}/autorag/evaluator.py

@@ -151,6 +151,9 @@ class Evaluator:
 		)
 		yaml_dict = load_yaml_config(yaml_path)
 		vectordb = yaml_dict.get("vectordb", [])
+		for vectordb_config in vectordb:
+			if "path" in vectordb_config:
+				vectordb_config["path"] = os.path.abspath(vectordb_config["path"])
 
 		vectordb_config_path = os.path.join(
 			self.project_dir, "resources", "vectordb.yaml"

autorag-0.3.25/autorag/nodes/generator/chat.py

@@ -0,0 +1,65 @@
+"""Helpers for truncating chat prompts without changing their message roles."""
+
+from typing import Callable, Dict, List, Sequence, TypeVar, Union
+
+
+Token = TypeVar("Token")
+Encode = Callable[[str], Sequence[Token]]
+Decode = Callable[[Sequence[Token]], str]
+Prompt = Union[str, List[Dict]]
+
+
+def messages_to_string(messages: List[Dict[str, str]]) -> str:
+	"""Render messages only for token counting, never for model input."""
+	formatted_parts = [
+		f"<|im_start|>{message['role']}\n{message['content']}<|im_end|>"
+		for message in messages
+	]
+	formatted_parts.append("<|im_start|>assistant")
+	return "\n".join(formatted_parts)
+
+
+def count_message_tokens(messages: List[Dict], encode: Encode) -> int:
+	"""Count the tokenized representation of a chat prompt."""
+	return len(encode(messages_to_string(messages)))
+
+
+def truncate_messages_by_token(
+	messages: List[Dict], encode: Encode, decode: Decode, max_token_size: int
+) -> List[Dict]:
+	"""Fit a chat prompt in a token budget while preserving message structure."""
+	truncated_messages = [dict(message) for message in messages]
+	if count_message_tokens(truncated_messages, encode) <= max_token_size:
+		return truncated_messages
+
+	result = [{**message, "content": ""} for message in truncated_messages]
+	if count_message_tokens(result, encode) > max_token_size:
+		return result
+
+	for index, message in enumerate(truncated_messages):
+		content_tokens = encode(message.get("content", ""))
+		low, high = 0, len(content_tokens)
+		best_content = ""
+		while low <= high:
+			mid = (low + high) // 2
+			candidate = [dict(item) for item in result]
+			candidate[index]["content"] = decode(content_tokens[:mid])
+			if count_message_tokens(candidate, encode) <= max_token_size:
+				best_content = candidate[index]["content"]
+				low = mid + 1
+			else:
+				high = mid - 1
+
+		result[index]["content"] = best_content
+
+	return result
+
+
+def truncate_prompt_by_token(
+	prompt: Prompt, encode: Encode, decode: Decode, max_token_size: int
+) -> Prompt:
+	"""Truncate text or chat prompts while retaining a chat prompt's roles."""
+	if isinstance(prompt, list):
+		return truncate_messages_by_token(prompt, encode, decode, max_token_size)
+	tokens = encode(prompt)
+	return decode(tokens[:max_token_size])

{autorag-0.3.24 → autorag-0.3.25}/autorag/nodes/generator/llama_index_llm.py

@@ -55,7 +55,12 @@ class LlamaIndexLLM(BaseGenerator):
 						"`model` or `model_name` parameter must be provided for using huggingfacellm."
 					)
 			kwargs["tokenizer_name"] = kwargs["model_name"]
-		self.llm_instance: BaseLLM = llm_class(**pop_params(llm_class.__init__, kwargs))
+		llm_params = pop_params(llm_class.__init__, kwargs)
+		model_fields = getattr(llm_class, "model_fields", {})
+		for key in list(kwargs.keys()):
+			if key in model_fields:
+				llm_params[key] = kwargs.pop(key)
+		self.llm_instance: BaseLLM = llm_class(**llm_params)
 
 	def __del__(self):
 		super().__del__()

{autorag-0.3.24 → autorag-0.3.25}/autorag/nodes/generator/minimax_llm.py

@@ -8,6 +8,11 @@ from openai import AsyncOpenAI
 from tiktoken import Encoding
 
 from autorag.nodes.generator.base import BaseGenerator
+from autorag.nodes.generator.chat import (
+	count_message_tokens as count_message_tokens_with_encoder,
+	messages_to_string as render_messages_to_string,
+	truncate_prompt_by_token,
+)
 from autorag.utils.util import (
 	get_event_loop,
 	process_batch,
@@ -186,63 +191,34 @@ class MiniMaxLLM(BaseGenerator):
 def truncate_by_token(
 	prompt: Union[str, List[Dict]], tokenizer: Encoding, max_token_size: int
 ) -> Union[str, List[Dict]]:
-	if isinstance(prompt, list):
-		return truncate_messages_by_token(prompt, tokenizer, max_token_size)
-	tokens = tokenizer.encode(prompt, allowed_special="all")
-	return tokenizer.decode(tokens[:max_token_size])
+	return truncate_prompt_by_token(
+		prompt,
+		lambda text: tokenizer.encode(text, allowed_special="all"),
+		tokenizer.decode,
+		max_token_size,
+	)
 
 
 def truncate_messages_by_token(
 	messages: List[Dict], tokenizer: Encoding, max_token_size: int
 ) -> List[Dict]:
-	truncated_messages = [dict(message) for message in messages]
-	if count_message_tokens(truncated_messages, tokenizer) <= max_token_size:
-		return truncated_messages
-
-	result: List[Dict] = []
-	for message in truncated_messages:
-		candidate = result + [message]
-		if count_message_tokens(candidate, tokenizer) <= max_token_size:
-			result.append(message)
-			continue
-
-		content_tokens = tokenizer.encode(
-			message.get("content", ""), allowed_special="all"
-		)
-		low, high = 0, len(content_tokens)
-		best_content = ""
-		while low <= high:
-			mid = (low + high) // 2
-			truncated_message = {
-				**message,
-				"content": tokenizer.decode(content_tokens[:mid]),
-			}
-			candidate = result + [truncated_message]
-			if count_message_tokens(candidate, tokenizer) <= max_token_size:
-				best_content = truncated_message["content"]
-				low = mid + 1
-			else:
-				high = mid - 1
-
-		if best_content or not result:
-			result.append({**message, "content": best_content})
-		break
-
-	return result
+	return truncate_prompt_by_token(
+		messages,
+		lambda text: tokenizer.encode(text, allowed_special="all"),
+		tokenizer.decode,
+		max_token_size,
+	)
 
 
 def count_message_tokens(messages: List[Dict], tokenizer: Encoding) -> int:
-	return len(tokenizer.encode(messages_to_string(messages), allowed_special="all"))
+	return count_message_tokens_with_encoder(
+		messages, lambda text: tokenizer.encode(text, allowed_special="all")
+	)
 
 
 def messages_to_string(messages: List[Dict[str, str]]) -> str:
-	"""Convert chat messages to string format for token counting."""
-	formatted_parts = [
-		f"<|im_start|>{message['role']}\n{message['content']}<|im_end|>"
-		for message in messages
-	]
-	formatted_parts.append("<|im_start|>assistant")
-	return "\n".join(formatted_parts)
+	"""Convert chat messages to a string for token counting compatibility."""
+	return render_messages_to_string(messages)
 
 
 def strip_think_tags(text: str) -> str:

{autorag-0.3.24 → autorag-0.3.25}/autorag/nodes/generator/openai_llm.py

@@ -8,6 +8,7 @@ from openai import AsyncOpenAI
 from tiktoken import Encoding
 
 from autorag.nodes.generator.base import BaseGenerator
+from autorag.nodes.generator.chat import messages_to_string, truncate_prompt_by_token
 from autorag.utils.util import (
 	get_event_loop,
 	process_batch,
@@ -333,16 +334,9 @@ class OpenAILLM(BaseGenerator):
 		if not self.llm.startswith("gpt-5"):
 			raise ValueError("get_result_gpt_5 is only for gpt-5 models.")
 		messages = parse_prompt(prompt)
-		instruction = "\n\n".join(
-			[msg["content"] for msg in messages if msg["role"] == "system"]
-		)
-		user_input = "\n\n".join(
-			[msg["content"] for msg in messages if msg["role"] == "user"]
-		)
 		response = await self.client.responses.create(
 			model=self.llm,
-			instructions=instruction,
-			input=user_input,
+			input=messages,
 			**kwargs,
 		)
 		answer: str = response.output_text
@@ -354,21 +348,17 @@ class OpenAILLM(BaseGenerator):
 def truncate_by_token(
 	prompt: Union[str, List[Dict]], tokenizer: Encoding, max_token_size: int
 ):
-	if isinstance(prompt, list):
-		prompt = tiktoken_messages_to_string(prompt)
-	tokens = tokenizer.encode(prompt, allowed_special="all")
-	return tokenizer.decode(tokens[:max_token_size])
+	return truncate_prompt_by_token(
+		prompt,
+		lambda text: tokenizer.encode(text, allowed_special="all"),
+		tokenizer.decode,
+		max_token_size,
+	)
 
 
 def tiktoken_messages_to_string(messages: List[Dict[str, str]]) -> str:
-	"""Convert chat messages to string format for accurate token counting"""
-	formatted_parts = [
-		f"<|im_start|>{message['role']}\n{message['content']}<|im_end|>"
-		for message in messages
-	]
-	formatted_parts.append("<|im_start|>assistant")
-	full_string = "\n".join(formatted_parts)
-	return full_string
+	"""Convert chat messages to a string for token counting compatibility."""
+	return messages_to_string(messages)
 
 
 def parse_prompt(prompt: Union[str, List[Dict]]) -> List[Dict]:

{autorag-0.3.24 → autorag-0.3.25}/autorag/nodes/generator/vllm_api.py

@@ -7,6 +7,7 @@ import requests
 from asyncio import to_thread
 
 from autorag.nodes.generator.base import BaseGenerator
+from autorag.nodes.generator.chat import truncate_prompt_by_token
 from autorag.utils.util import get_event_loop, process_batch, result_to_dataframe
 
 logger = logging.getLogger("AutoRAG")
@@ -84,15 +85,18 @@ class VllmAPI(BaseGenerator):
 		logprob_result = list(map(lambda x: x[2], results))
 		return answer_result, token_result, logprob_result
 
-	def truncate_by_token(self, prompt: Union[str, List[Dict]]) -> str:
+	def truncate_by_token(
+		self, prompt: Union[str, List[Dict]]
+	) -> Union[str, List[Dict]]:
 		"""
 		Function to truncate prompts to fit within the maximum token limit.
 		"""
-		if not isinstance(prompt, str):
-			content_list = [msg["content"] for msg in prompt]
-			prompt = "\n".join(content_list)
-		tokens = self.encoding_for_model(prompt)["tokens"]  # Simple tokenization
-		return self.decoding_for_model(tokens[: self.max_model_len])["prompt"]
+		return truncate_prompt_by_token(
+			prompt,
+			lambda text: self.encoding_for_model(text)["tokens"],
+			lambda tokens: self.decoding_for_model(list(tokens))["prompt"],
+			self.max_model_len,
+		)
 
 	def call_vllm_api(self, prompt: Union[str, List[Dict]], **kwargs) -> dict:
 		"""

{autorag-0.3.24 → autorag-0.3.25}/autorag/schema/node.py

@@ -22,6 +22,12 @@ class Node:
 	run_node: Callable = field(init=False)
 
 	def __post_init__(self):
+		if self.node_type == "hybrid_retrieval" and isinstance(
+			self.node_params.get("top_k"), list
+		):
+			raise ValueError(
+				"hybrid_retrieval top_k must be an integer, not a list."
+			)
 		self.run_node = get_support_nodes(self.node_type)
 		if self.run_node is None:
 			raise ValueError(f"Node type {self.node_type} is not supported.")

{autorag-0.3.24 → autorag-0.3.25}/autorag/vectordb/qdrant.py

@@ -8,7 +8,7 @@ from qdrant_client.models import (
 	PointIdsList,
 	HasIdCondition,
 	Filter,
-	SearchRequest,
+	QueryRequest,
 )
 
 from typing import List, Tuple, Union
@@ -148,18 +148,18 @@ class Qdrant(BaseVectorStore):
 
 		search_queries = list(
 			map(
-				lambda x: SearchRequest(vector=x, limit=top_k, with_vector=True),
+				lambda x: QueryRequest(query=x, limit=top_k, with_vector=True),
 				query_embeddings,
 			)
 		)
 
-		search_result = self.client.search_batch(
+		search_result = self.client.query_batch_points(
 			collection_name=self.collection_name, requests=search_queries
 		)
 
 		# Extract IDs and distances
-		ids = [[str(hit.id) for hit in result] for result in search_result]
-		scores = [[hit.score for hit in result] for result in search_result]
+		ids = [[str(hit.id) for hit in result.points] for result in search_result]
+		scores = [[hit.score for hit in result.points] for result in search_result]
 
 		return ids, scores
 

{autorag-0.3.24 → autorag-0.3.25}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "AutoRAG"
-version = "0.3.24"
+version = "0.3.25"
 authors = [
     { name = "Marker-Inc", email = "vkehfdl1@gmail.com" }
 ]
@@ -46,13 +46,13 @@ dependencies = [
     "emoji>=2.14.1",
     "fastapi>=0.115.13",
     "banks>=2.4.2",
-    "grpcio>=1.66.2,<1.83.0",
-    "grpcio-health-checking>=1.66.2,<1.83.0",
-    "grpcio-status>=1.66.2,<1.83.0",
-    "grpcio-tools>=1.66.2,<1.83.0",
+    "grpcio>=1.66.2,<1.84.0",
+    "grpcio-health-checking>=1.66.2,<1.84.0",
+    "grpcio-status>=1.66.2,<1.84.0",
+    "grpcio-tools>=1.66.2,<1.84.0",
     "datasets>=3.5.1",
     # Vector DB
-    "pymilvus>=2.6.0b0,<3",  # for using milvus vectordb (3.x drops milvus-lite)
+    "pymilvus>=2.6.0b0,<4",  # for using milvus vectordb (3.x drops milvus-lite)
     "chromadb>=1.0.0",  # for chroma vectordb
     "weaviate-client>=4.15.2",  # for weaviate vectordb
     "pinecone[grpc]",  # for pinecone vectordb
@@ -72,6 +72,7 @@ dependencies = [
     "llama-index-llms-openai-like>=0.4.0",
     "llama-index-retrievers-bm25>=0.5.2",
     "llama-index-llms-bedrock>=0.3.8",
+    "llama-parse>=0.6.21",  # llamaparse parser; no longer transitively provided by llama-index
     # WebUI
     "streamlit>=1.54.0",
     "gradio>=6.15.0",
@@ -119,7 +120,7 @@ dev = [
     "pytest-mock",
     "aioresponses>=0.7.9",
     "asyncstdlib",
-    "vllm-mock>=0.0.3",
+    "vllm-mock>=0.0.3; sys_platform == 'linux'",
     "ty",
     "deptry",
 ]
@@ -130,7 +131,8 @@ parse = ["PyMuPDF>=1.26.1", "pdfminer.six>=20250506", "pdfplumber>=0.11.7", "jq>
 ja = ["sudachipy>=0.6.8", "sudachidict_core"]
 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", "vllm>=0.11.0"]
+    "llama-index-llms-huggingface>=0.5.0", "onnxruntime>=1.22.0", "xgrammar>=0.1.32,<0.3.0",
+    "vllm>=0.11.0,<0.28.0"]
 all = ["AutoRAG[gpu]", "AutoRAG[ko]", "AutoRAG[parse]", "AutoRAG[ja]"]
 
 [project.scripts]

{autorag-0.3.24 → autorag-0.3.25}/sample_dataset/README.md

@@ -2,11 +2,11 @@
 
 The sample_dataset folder does not includes a `qa.parquet`, `corpus.parquet` file that is significantly large and cannot be uploaded directly to Git due to size limitations.
 
-To prepare and use datasets available in the sample_dataset folder, specifically `triviaqa`, `hotpotqa`, `msmarco` and `eli5`, you can follow the outlined methods below.
+To prepare and use datasets available in the sample_dataset folder, specifically `triviaqa`, `hotpotqa`, `msmarco`, `eli5` and `clapnq`, you can follow the outlined methods below.
 
 ## Usage
 
- The example provided uses `triviaqa`, but the same approach applies to `msmarco`, `eli5` and `hotpotqa`.
+ The example provided uses `triviaqa`, but the same approach applies to `msmarco`, `eli5`, `hotpotqa` and `clapnq`.
 
 ### 1. Run with a specified save path
 To execute the Python script from the terminal and save the dataset to a specified path, use the command:
@@ -22,4 +22,16 @@ If you run the script without the `--save_path` argument, the dataset will be sa
 ```bash
 python ./sample_dataset/triviaqa/load_triviaqa_dataset.py
 ```
+
+ClapNQ's official retrieval corpus and answerable train/development question
+files are downloaded from the `PrimeQA/clapnq` repository. The corpus includes
+the complete retrieval collection rather than only gold passages. Questions
+without a gold passage are omitted because they cannot produce AutoRAG
+`retrieval_gt` entries.
+
+To create the ClapNQ sample directly:
+
+```bash
+python ./sample_dataset/clapnq/load_clapnq_dataset.py --save_path /path/to/save/dataset
+```
 This behavior allows for a straightforward execution without needing to specify a path, making it convenient for quick tests or when working directly within the target directory.

autorag-0.3.25/sample_dataset/clapnq/load_clapnq_dataset.py

@@ -0,0 +1,86 @@
+import pathlib
+from collections.abc import Iterable
+from typing import Any
+
+import click
+import pandas as pd
+from datasets import load_dataset
+
+
+PASSAGES_URL = (
+	"https://huggingface.co/datasets/PrimeQA/clapnq_passages/resolve/main/passages.tsv"
+)
+RETRIEVAL_URL = "https://raw.githubusercontent.com/primeqa/clapnq/main/retrieval"
+TRAIN_QUESTIONS_URL = f"{RETRIEVAL_URL}/train/question_train_answerable.tsv"
+TEST_QUESTIONS_URL = f"{RETRIEVAL_URL}/dev/question_dev_answerable.tsv"
+
+
+def _load_tsv(url: str):
+	return load_dataset("csv", data_files={"split": url}, delimiter="\t", split="split")
+
+
+def _records_to_dataframes(
+	passages: Iterable[dict[str, Any]], questions: Iterable[dict[str, Any]]
+) -> tuple[pd.DataFrame, pd.DataFrame]:
+	corpus = pd.DataFrame(
+		{
+			"doc_id": passage["id"],
+			"contents": passage["text"],
+			"metadata": {"title": passage.get("title", "")},
+		}
+		for passage in passages
+	)
+	corpus_ids = set(corpus["doc_id"])
+	qa_rows = []
+	for question in questions:
+		doc_id = question.get("doc-id-list")
+		answers = [
+			answer.strip() for answer in (question.get("answers") or "").split("::")
+		]
+		answers = list(dict.fromkeys(answer for answer in answers if answer))
+		if not doc_id or doc_id not in corpus_ids or not answers:
+			continue
+		qa_rows.append(
+			{
+				"qid": str(question["id"]),
+				"query": question["question"].strip(),
+				"retrieval_gt": [[doc_id]],
+				"generation_gt": answers,
+			}
+		)
+
+	return corpus, pd.DataFrame(qa_rows)
+
+
+@click.command()
+@click.option(
+	"--save_path",
+	type=click.Path(file_okay=False, dir_okay=True, path_type=pathlib.Path),
+	default=pathlib.Path(__file__).parent,
+	help="Path to save sample ClapNQ dataset.",
+)
+def load_clapnq_dataset(save_path: pathlib.Path):
+	"""Download the official ClapNQ retrieval corpus and QA splits."""
+	paths = {
+		"corpus": save_path / "corpus.parquet",
+		"train": save_path / "qa_train.parquet",
+		"test": save_path / "qa_test.parquet",
+	}
+	if any(path.exists() for path in paths.values()):
+		raised = next(path for path in paths.values() if path.exists())
+		raise ValueError(f"{raised.name} already exists")
+
+	passages = _load_tsv(PASSAGES_URL)
+	train_questions = _load_tsv(TRAIN_QUESTIONS_URL)
+	test_questions = _load_tsv(TEST_QUESTIONS_URL)
+	corpus, train_qa = _records_to_dataframes(passages, train_questions)
+	_, test_qa = _records_to_dataframes(passages, test_questions)
+
+	save_path.mkdir(parents=True, exist_ok=True)
+	corpus.to_parquet(paths["corpus"], index=False)
+	train_qa.to_parquet(paths["train"], index=False)
+	test_qa.to_parquet(paths["test"], index=False)
+
+
+if __name__ == "__main__":
+	load_clapnq_dataset()

{autorag-0.3.24 → autorag-0.3.25}/tests/autorag/nodes/generator/test_llama_index_llm.py

@@ -23,12 +23,43 @@ from tests.mock import MockLLM
 logger = logging.getLogger("AutoRAG")
 
 
+class PydanticCapturingLLM(BaseModel):
+    model: str
+    is_chat_model: bool
+    api_base: str
+    api_key: str
+
+    @classmethod
+    def class_name(cls):
+        return "OpenAILike"
+
+
 @pytest.fixture
 def llama_index_llm_instance():
     generator_models["mock"] = MockLLM
     return LlamaIndexLLM(project_dir=".", llm="mock", temperature=0.5, top_p=0.9)
 
 
+def test_llama_index_llm_forwards_pydantic_model_fields():
+    generator_models["pydantic_capturing"] = PydanticCapturingLLM
+    try:
+        instance = LlamaIndexLLM(
+            project_dir=".",
+            llm="pydantic_capturing",
+            model="qwen-plus",
+            is_chat_model=True,
+            api_base="https://example.test/v1",
+            api_key="secret",
+        )
+    finally:
+        generator_models.pop("pydantic_capturing")
+
+    assert instance.llm_instance.model == "qwen-plus"
+    assert instance.llm_instance.is_chat_model is True
+    assert instance.llm_instance.api_base == "https://example.test/v1"
+    assert instance.llm_instance.api_key == "secret"
+
+
 def test_llama_index_llm(llama_index_llm_instance):
     answers, tokens, log_probs = llama_index_llm_instance._pure(prompts)
     check_generated_texts(answers)
@@ -110,8 +141,8 @@ Hello, my name is John Doe. My phone number is 1234567890. I am 30 years old. I
 
 
 @pytest.mark.skipif(
-    is_github_action(),
-    reason="Skipping this test on GitHub Actions because it uses the real OpenAI API.",
+    is_github_action() or not os.getenv("OPENAI_API_KEY"),
+    reason="Skipping this test because it uses the real OpenAI API.",
 )
 @pytest.mark.asyncio()
 async def test_llama_index_llm_astream():

{autorag-0.3.24 → autorag-0.3.25}/tests/autorag/nodes/generator/test_minimax.py

@@ -426,6 +426,32 @@ class TestTruncateByToken:
 		assert isinstance(result, list)
 		assert result == messages
 
+	def test_truncate_messages_preserves_roles(self):
+		class CharacterTokenizer:
+			def encode(self, text, allowed_special="none"):
+				return list(text)
+
+			def decode(self, tokens):
+				return "".join(tokens)
+
+		messages = [
+			{"role": "system", "content": "Follow the conversation."},
+			{"role": "user", "content": "What is the capital of France?"},
+			{"role": "assistant", "content": "Paris."},
+			{"role": "user", "content": "And Germany?"},
+		]
+
+		result = truncate_by_token(messages, CharacterTokenizer(), 10_000)
+
+		assert result == messages
+		assert result is not messages
+		assert [message["role"] for message in result] == [
+			"system",
+			"user",
+			"assistant",
+			"user",
+		]
+
 
 class TestThinkTagHelpers:
 	def test_strip_think_tags(self):