AutoRAG 0.3.22__tar.gz → 0.3.23__tar.gz

{autorag-0.3.22 → autorag-0.3.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: AutoRAG
-Version: 0.3.22
+Version: 0.3.23
 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>
@@ -217,8 +217,8 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: <3.13,>=3.10
-Requires-Dist: aiohttp>=3.12.13
-Requires-Dist: banks>=2.1.2
+Requires-Dist: aiohttp<3.14,>=3.12.13
+Requires-Dist: banks>=2.4.2
 Requires-Dist: chromadb>=1.0.0
 Requires-Dist: click>=8.2.1
 Requires-Dist: cohere>=5.18.0
@@ -228,19 +228,19 @@ Requires-Dist: emoji>=2.14.1
 Requires-Dist: evaluate>=0.4.4
 Requires-Dist: fastapi>=0.115.13
 Requires-Dist: fastparquet>=2024.11.0
-Requires-Dist: gradio>=5.34.2
-Requires-Dist: grpcio-health-checking<1.68.0,>=1.66.2
-Requires-Dist: grpcio-status<1.68.0,>=1.66.2
-Requires-Dist: grpcio-tools<1.68.0,>=1.66.2
-Requires-Dist: grpcio<1.68.0,>=1.66.2
+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: ipykernel>=6.29.5
 Requires-Dist: ipywidgets-bokeh>=1.6.0
 Requires-Dist: ipywidgets>=8.1.7
 Requires-Dist: langchain-community>=0.4.1
-Requires-Dist: langchain-core>=1.2.23
-Requires-Dist: langchain-text-splitters>=1.1.1
-Requires-Dist: langchain>=1.2.14
-Requires-Dist: langsmith>=0.7.23
+Requires-Dist: langchain-core>=1.3.3
+Requires-Dist: langchain-text-splitters>=1.1.2
+Requires-Dist: langchain>=1.3.9
+Requires-Dist: langsmith>=0.8.18
 Requires-Dist: llama-index-core>=0.12.42
 Requires-Dist: llama-index-embeddings-ollama>=0.6.0
 Requires-Dist: llama-index-embeddings-openai-like>=0.1.1
@@ -252,14 +252,14 @@ 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: mixedbread-ai>=2.2.6
-Requires-Dist: numpy==1.26.4
+Requires-Dist: numpy>=2.2.6
 Requires-Dist: openai>=2.8.0
 Requires-Dist: pandas>=2.2.3
 Requires-Dist: panel>=1.7.1
 Requires-Dist: pinecone[grpc]
-Requires-Dist: pyarrow>=20.0.0
+Requires-Dist: pyarrow>=23.0.1
 Requires-Dist: pydantic>=2.9.2
-Requires-Dist: pymilvus>=2.6.0b0
+Requires-Dist: pymilvus<3,>=2.6.0b0
 Requires-Dist: pyngrok>=7.2.11
 Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: qdrant-client>=1.12.1
@@ -270,7 +270,7 @@ Requires-Dist: rouge-score>=0.1.2
 Requires-Dist: sacrebleu>=2.5.1
 Requires-Dist: scikit-learn>=1.7.0
 Requires-Dist: seaborn>=0.13.2
-Requires-Dist: streamlit>=1.46.0
+Requires-Dist: streamlit>=1.54.0
 Requires-Dist: tiktoken>=0.9.0
 Requires-Dist: tokenlog>=0.0.3
 Requires-Dist: tqdm>=4.67.1
@@ -292,7 +292,7 @@ Requires-Dist: pdfminer-six>=20250506; extra == 'all'
 Requires-Dist: pdfplumber>=0.11.7; extra == 'all'
 Requires-Dist: peft>=0.15.2; extra == 'all'
 Requires-Dist: pymupdf>=1.26.1; extra == 'all'
-Requires-Dist: pypdf2==2.12.1; extra == 'all'
+Requires-Dist: pypdf2>=3.0.1; extra == 'all'
 Requires-Dist: sentence-transformers>=4.1.0; extra == 'all'
 Requires-Dist: sentencepiece>=0.2.0; extra == 'all'
 Requires-Dist: sudachidict-core; extra == 'all'
@@ -327,11 +327,14 @@ Requires-Dist: pdf2image>=1.17.0; extra == 'parse'
 Requires-Dist: pdfminer-six>=20250506; extra == 'parse'
 Requires-Dist: pdfplumber>=0.11.7; extra == 'parse'
 Requires-Dist: pymupdf>=1.26.1; extra == 'parse'
-Requires-Dist: pypdf2==2.12.1; extra == 'parse'
+Requires-Dist: pypdf2>=3.0.1; extra == 'parse'
 Requires-Dist: unstructured[pdf]>=0.17.2; extra == 'parse'
 Description-Content-Type: text/markdown
 
-# AutoRAG
+# AutoRAG (Legacy)
+
+> [!NOTE]
+> This is the original Python-based AutoRAG — the RAG AutoML tool for automatically finding an optimal RAG pipeline for your data. It now lives in the `legacy/` directory of the AutoRAG monorepo and is in **maintenance mode**: it continues to receive bug fixes, dependency updates, and PyPI releases (`pip install AutoRAG`), but new feature development happens in [AutoRAG 2.0](../README.md) at the repository root.
 
 RAG AutoML tool for automatically finding an optimal RAG pipeline for your data.
 

{autorag-0.3.22 → autorag-0.3.23}/README.md

@@ -1,4 +1,7 @@
-# AutoRAG
+# AutoRAG (Legacy)
+
+> [!NOTE]
+> This is the original Python-based AutoRAG — the RAG AutoML tool for automatically finding an optimal RAG pipeline for your data. It now lives in the `legacy/` directory of the AutoRAG monorepo and is in **maintenance mode**: it continues to receive bug fixes, dependency updates, and PyPI releases (`pip install AutoRAG`), but new feature development happens in [AutoRAG 2.0](../README.md) at the repository root.
 
 RAG AutoML tool for automatically finding an optimal RAG pipeline for your data.
 

{autorag-0.3.22 → autorag-0.3.23}/autorag/data/parse/table_hybrid_parse.py

@@ -3,7 +3,7 @@ import tempfile
 from glob import glob
 from typing import List, Tuple, Dict
 
-from PyPDF2 import PdfFileReader, PdfFileWriter
+from PyPDF2 import PdfReader, PdfWriter
 import pdfplumber
 
 from autorag.support import get_support_modules
@@ -76,8 +76,8 @@ def save_page_by_table(data_path: str, text_dir: str, table_dir: str) -> Dict[st
 	file_name = os.path.basename(data_path).split(".pdf")[0]
 
 	with open(data_path, "rb") as input_data:
-		pdf_reader = PdfFileReader(input_data)
-		num_pages = pdf_reader.getNumPages()
+		pdf_reader = PdfReader(input_data)
+		num_pages = len(pdf_reader.pages)
 
 		path_map_dict = {}
 		for page_num in range(num_pages):
@@ -100,9 +100,9 @@ def _get_output_path(
 		return os.path.join(directory, f"{file_name}_page_{page_num + 1}.pdf")
 
 
-def _save_single_page(pdf_reader: PdfFileReader, page_num: int, output_pdf_path: str):
-	pdf_writer = PdfFileWriter()
-	pdf_writer.addPage(pdf_reader.getPage(page_num))
+def _save_single_page(pdf_reader: PdfReader, page_num: int, output_pdf_path: str):
+	pdf_writer = PdfWriter()
+	pdf_writer.add_page(pdf_reader.pages[page_num])
 
 	with open(output_pdf_path, "wb") as output_file:
 		pdf_writer.write(output_file)

{autorag-0.3.22 → autorag-0.3.23}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "AutoRAG"
-version = "v0.3.22"
+version = "0.3.23"
 authors = [
     { name = "Marker-Inc", email = "vkehfdl1@gmail.com" }
 ]
@@ -23,14 +23,14 @@ classifiers = [
 urls = { Homepage = "https://github.com/Marker-Inc-Korea/AutoRAG" }
 dependencies = [
     "pydantic>=2.9.2",
-    "numpy==1.26.4",  # temporal not using numpy 2.0.0
+    "numpy>=2.2.6",
     "pandas>=2.2.3",
     "tqdm>=4.67.1",
     "tiktoken>=0.9.0",  # for counting token
     "openai>=2.8.0",
     "rank_bm25>=0.2.2",  # for bm25 retrieval
     "pyyaml>=6.0.2",  # for yaml file
-    "pyarrow>=20.0.0",  # for pandas with parquet
+    "pyarrow>=23.0.1",  # for pandas with parquet
     "fastparquet>=2024.11.0",  # for pandas with parquet
     "sacrebleu>=2.5.1",  # for bleu score
     "evaluate>=0.4.4",  # for meteor and other scores
@@ -39,21 +39,20 @@ dependencies = [
     "click>=8.2.1",  # for cli
     "cohere>=5.18.0",  # for cohere services
     "tokenlog>=0.0.3",  # for token logging
-    "aiohttp>=3.12.13",  # for async http requests
+    "aiohttp>=3.12.13,<3.14",  # aioresponses (dev) incompatible with aiohttp 3.14+ (stream_writer)
     "voyageai>=0.3.2",  # for voyageai reranker
     "mixedbread-ai>=2.2.6",  # for mixedbread-ai reranker
     "scikit-learn>=1.7.0",
     "emoji>=2.14.1",
     "fastapi>=0.115.13",
-    "banks>=2.1.2",
-    "grpcio>=1.66.2,<1.68.0",
-    "grpcio-health-checking>=1.66.2,<1.68.0",
-    "grpcio-status>=1.66.2,<1.68.0",
-    "grpcio-tools>=1.66.2,<1.68.0",
+    "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",
     "datasets>=3.5.1",
-    "pyarrow>=20.0.0",
     # Vector DB
-    "pymilvus>=2.6.0b0",  # for using milvus vectordb
+    "pymilvus>=2.6.0b0,<3",  # 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
@@ -74,14 +73,14 @@ dependencies = [
     "llama-index-retrievers-bm25>=0.5.2",
     "llama-index-llms-bedrock>=0.3.8",
     # WebUI
-    "streamlit>=1.46.0",
-    "gradio>=5.34.2",
+    "streamlit>=1.54.0",
+    "gradio>=6.15.0",
     # LangChain
-    "langchain>=1.2.14",
-    "langchain-core>=1.2.23",
+    "langchain>=1.3.9",
+    "langchain-core>=1.3.3",
     "langchain-community>=0.4.1",
-    "langchain-text-splitters>=1.1.1",
-    "langsmith>=0.7.23",
+    "langchain-text-splitters>=1.1.2",
+    "langsmith>=0.8.18",
     # autorag dashboard
     "panel>=1.7.1",
     "seaborn>=0.13.2",
@@ -118,7 +117,7 @@ dev = [
     "pytest-xdist",
     "pytest-asyncio",
     "pytest-mock",
-    "aioresponses",
+    "aioresponses>=0.7.9",
     "asyncstdlib",
     "vllm-mock>=0.0.3",
     "ty",
@@ -127,7 +126,7 @@ dev = [
 
 [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"]
+parse = ["PyMuPDF>=1.26.1", "pdfminer.six>=20250506", "pdfplumber>=0.11.7", "jq>=1.9.1", "unstructured[pdf]>=0.17.2", "PyPDF2>=3.0.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", "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",

{autorag-0.3.22 → autorag-0.3.23}/tests/autorag/nodes/generator/test_vllm.py

@@ -1,10 +1,12 @@
 import pytest
 import pandas as pd
 
-vllm_mock = pytest.importorskip(
-	"vllm_mock", reason="vllm_mock not compatible with installed vllm"
-)
-LLM = vllm_mock.LLM
+pytest.importorskip("vllm", reason="vllm not installed")
+try:
+	import vllm_mock
+	from vllm_mock import LLM
+except Exception as exc:  # noqa: BLE001 - skip when mock is incompatible with installed vllm
+	pytest.skip(f"vllm_mock not compatible with installed vllm: {exc}", allow_module_level=True)
 
 from autorag.nodes.generator import Vllm  # noqa: E402
 from tests.autorag.nodes.generator.test_generator_base import (  # noqa: E402
@@ -18,7 +20,6 @@ from tests.autorag.nodes.generator.test_generator_base import (  # noqa: E402
 	check_generated_log_probs_chat,
 )
 
-
 def test_vllm(mocker):
 	mock_class = mocker.patch("vllm.LLM")
 	mock_class.return_value = LLM(model="mock-model")

{autorag-0.3.22 → autorag-0.3.23}/tests/autorag/nodes/hybridretrieval/test_run_hybrid_retrieval.py

@@ -171,10 +171,11 @@ def test_run_hybrid_retrieval_node(node_line_dir):
     assert len(summary_df) == 2
     assert summary_df["filename"][0] == "0.parquet"
     assert summary_df["filename"][1] == "1.parquet"
-    assert summary_df["retrieval_f1"][0] == vectordb_top_k_df["retrieval_f1"].mean()
-    assert (
-        summary_df["retrieval_recall"][0]
-        == vectordb_top_k_df["retrieval_recall"].mean()
+    assert summary_df["retrieval_f1"][0] == pytest.approx(
+        float(vectordb_top_k_df["retrieval_f1"].mean())
+    )
+    assert summary_df["retrieval_recall"][0] == pytest.approx(
+        float(vectordb_top_k_df["retrieval_recall"].mean())
     )
     assert summary_df["module_name"][0] == "HybridCC"
     assert summary_df["module_name"][1] == "HybridRRF"

{autorag-0.3.22 → autorag-0.3.23}/tests/autorag/nodes/lexicalretrieval/test_run_lexical_retrieval.py

@@ -122,8 +122,12 @@ def test_run_lexical_retrieval_node(node_line_dir):
     }
     assert len(summary_df) == 1
     assert summary_df["filename"][0] == "0.parquet"
-    assert summary_df["retrieval_f1"][0] == bm25_top_k_df["retrieval_f1"].mean()
-    assert summary_df["retrieval_recall"][0] == bm25_top_k_df["retrieval_recall"].mean()
+    assert summary_df["retrieval_f1"][0] == pytest.approx(
+        float(bm25_top_k_df["retrieval_f1"].mean())
+    )
+    assert summary_df["retrieval_recall"][0] == pytest.approx(
+        float(bm25_top_k_df["retrieval_recall"].mean())
+    )
     assert summary_df["module_name"][0] == "BM25"
     assert summary_df["module_params"][0] == {"top_k": 4, "bm25_tokenizer": "gpt2"}
     assert summary_df["execution_time"][0] > 0

{autorag-0.3.22 → autorag-0.3.23}/tests/autorag/nodes/passagefilter/test_passage_filter_run.py

@@ -78,12 +78,11 @@ def test_run_passage_filter_node(node_line_dir):
     }
     assert len(summary_df) == 1
     assert summary_df["filename"][0] == "0.parquet"
-    assert (
-        summary_df["passage_filter_retrieval_f1"][0] == result_df["retrieval_f1"].mean()
+    assert summary_df["passage_filter_retrieval_f1"][0] == pytest.approx(
+        float(result_df["retrieval_f1"].mean())
     )
-    assert (
-        summary_df["passage_filter_retrieval_recall"][0]
-        == result_df["retrieval_recall"].mean()
+    assert summary_df["passage_filter_retrieval_recall"][0] == pytest.approx(
+        float(result_df["retrieval_recall"].mean())
     )
     assert summary_df["module_name"][0] == "SimilarityThresholdCutoff"
     assert summary_df["module_params"][0] == {"threshold": 0.87}

{autorag-0.3.22 → autorag-0.3.23}/tests/autorag/nodes/passagereranker/test_passage_reranker_run.py

@@ -136,13 +136,11 @@ def test_run_passage_reranker_node(node_line_dir):
     }
     assert len(summary_df) == 1
     assert summary_df["filename"][0] == "0.parquet"
-    assert (
-        summary_df["passage_reranker_retrieval_f1"][0]
-        == single_result_df["retrieval_f1"].mean()
+    assert summary_df["passage_reranker_retrieval_f1"][0] == pytest.approx(
+        float(single_result_df["retrieval_f1"].mean())
     )
-    assert (
-        summary_df["passage_reranker_retrieval_recall"][0]
-        == single_result_df["retrieval_recall"].mean()
+    assert summary_df["passage_reranker_retrieval_recall"][0] == pytest.approx(
+        float(single_result_df["retrieval_recall"].mean())
     )
     assert summary_df["module_name"][0] == "MonoT5"
     assert summary_df["module_params"][0] == {

{autorag-0.3.22 → autorag-0.3.23}/tests/autorag/nodes/semanticretrieval/test_run_semantic_retrieval.py

@@ -121,10 +121,11 @@ def test_run_semantic_retrieval_node(node_line_dir):
     }
     assert len(summary_df) == 1
     assert summary_df["filename"][0] == "0.parquet"
-    assert summary_df["retrieval_f1"][0] == vectordb_top_k_df["retrieval_f1"].mean()
-    assert (
-        summary_df["retrieval_recall"][0]
-        == vectordb_top_k_df["retrieval_recall"].mean()
+    assert summary_df["retrieval_f1"][0] == pytest.approx(
+        float(vectordb_top_k_df["retrieval_f1"].mean())
+    )
+    assert summary_df["retrieval_recall"][0] == pytest.approx(
+        float(vectordb_top_k_df["retrieval_recall"].mean())
     )
     assert summary_df["module_name"][0] == "VectorDB"
     assert summary_df["module_params"][0] == {"top_k": 4, "vectordb": "test_mock"}

{autorag-0.3.22 → autorag-0.3.23}/tests/autorag/test_cli.py

@@ -1,4 +1,5 @@
 import os
+import sys
 import pathlib
 import subprocess
 import tempfile
@@ -12,13 +13,14 @@ from tests.delete_tests import is_github_action
 
 root_dir = pathlib.PurePath(os.path.dirname(os.path.realpath(__file__))).parent
 resource_dir = os.path.join(root_dir, "resources")
+AUTORAG_CLI = [sys.executable, "-m", "autorag.cli"]
 
 
 def test_evaluator_cli():
     with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as project_dir:
         result = subprocess.run(
             [
-                "autorag",
+                *AUTORAG_CLI,
                 "evaluate",
                 "--config",
                 os.path.join(resource_dir, "simple_mock.yaml"),
@@ -52,7 +54,7 @@ def test_evaluator_cli():
 def test_validator_cli():
     result = subprocess.run(
         [
-            "autorag",
+            *AUTORAG_CLI,
             "validate",
             "--config",
             os.path.join(resource_dir, "simple_mock.yaml"),
@@ -93,7 +95,7 @@ def test_extract_best_config_cli():
         output_path = os.path.join(project_dir, "best.yaml")
         subprocess.run(
             [
-                "autorag",
+                *AUTORAG_CLI,
                 "extract_best_config",
                 "--trial_path",
                 trial_path,
@@ -109,7 +111,7 @@ def test_restart_evaluate():
         original_path = os.path.join(resource_dir, "result_project")
         copytree(original_path, project_dir, dirs_exist_ok=True)
         trial_path = os.path.join(project_dir, "1")
-        subprocess.run(["autorag", "restart_evaluate", "--trial_path", trial_path])
+        subprocess.run([*AUTORAG_CLI, "restart_evaluate", "--trial_path", trial_path])
         assert os.path.exists(os.path.join(trial_path, "summary.csv"))
 
 
@@ -121,5 +123,5 @@ def test_restart_evaluate_leads_start_evaluate():
         original_path = os.path.join(resource_dir, "result_project")
         copytree(original_path, project_dir, dirs_exist_ok=True)
         trial_path = os.path.join(project_dir, "3")
-        subprocess.run(["autorag", "restart_evaluate", "--trial_path", trial_path])
+        subprocess.run([*AUTORAG_CLI, "restart_evaluate", "--trial_path", trial_path])
         assert os.path.exists(os.path.join(project_dir, "4", "summary.csv"))

autorag-0.3.23/tests/autorag/vectordb/test_couchbase.py

@@ -0,0 +1,77 @@
+import asyncio
+import os
+
+import pytest
+
+from autorag.vectordb.couchbase import Couchbase
+from tests.delete_tests import is_github_action
+
+
+pytestmark = pytest.mark.skipif(
+	is_github_action(),
+	reason="This test needs couchbase connection string, username, and password",
+)
+
+
+@pytest.fixture
+def couchbase_instance():
+	couchbase = Couchbase(
+		embedding_model="mock",
+		bucket_name="autorag",
+		scope_name="autorag",
+		collection_name="autorag",
+		index_name="autorag_search",
+		connection_string=os.environ["COUCHBASE_CONNECTION_STRING"],
+		username=os.environ["COUCHBASE_USERNAME"],
+		password=os.environ["COUCHBASE_PASSWORD"],
+	)
+	yield couchbase
+
+
+@pytest.mark.asyncio
+async def test_add_and_query_documents(couchbase_instance):
+	# Add documents
+	ids = ["doc1", "doc2"]
+	texts = ["This is a test document.", "This is another test document."]
+	await couchbase_instance.add(ids, texts)
+
+	await asyncio.sleep(1)
+
+	# Query documents
+	queries = ["test document"]
+	contents, scores = await couchbase_instance.query(queries, top_k=2)
+
+	assert len(contents) == 1
+	assert len(scores) == 1
+	assert len(contents[0]) == 2
+	assert len(scores[0]) == 2
+	assert scores[0][0] > scores[0][1]
+
+	embeddings = await couchbase_instance.fetch([ids[0]])
+	assert len(embeddings) == 1
+	assert len(embeddings[0]) == 768
+
+	exist = await couchbase_instance.is_exist([ids[0], "doc3"])
+	assert len(exist) == 2
+	assert exist[0] is True
+	assert exist[1] is False
+
+
+@pytest.mark.asyncio
+async def test_delete_documents(couchbase_instance):
+	# Add documents
+	ids = ["doc1", "doc2"]
+	texts = ["This is a test document.", "This is another test document."]
+	await couchbase_instance.add(ids, texts)
+
+	await asyncio.sleep(1)
+
+	# Delete documents
+	await couchbase_instance.delete([ids[0]])
+
+	# Query documents to ensure they are deleted
+	queries = ["test document"]
+	contents, scores = await couchbase_instance.query(queries, top_k=2)
+
+	assert len(contents[0]) == 1
+	assert len(scores[0]) == 1

autorag-0.3.23/tests/autorag/vectordb/test_milvus.py

@@ -0,0 +1,75 @@
+import asyncio
+import os
+
+import pytest
+
+from autorag.vectordb.milvus import Milvus
+from tests.delete_tests import is_github_action
+
+
+pytestmark = pytest.mark.skipif(
+	is_github_action(),
+	reason="This test needs milvus uri and token which is confidential.",
+)
+
+
+@pytest.fixture
+def milvus_instance():
+	milvus = Milvus(
+		uri=os.environ["MILVUS_URI"],
+		token=os.environ["MILVUS_TOKEN"],
+		embedding_model="mock",
+		collection_name="test_collection",
+		similarity_metric="ip",
+	)
+	yield milvus
+	milvus.delete_collection()
+
+
+@pytest.mark.asyncio
+async def test_add_and_query_documents(milvus_instance):
+	# Add documents
+	ids = ["doc1", "doc2"]
+	texts = ["This is a test document.", "This is another test document."]
+	await milvus_instance.add(ids, texts)
+
+	await asyncio.sleep(1)
+
+	# Query documents
+	queries = ["test document"]
+	contents, scores = await milvus_instance.query(queries, top_k=2)
+
+	assert len(contents) == 1
+	assert len(scores) == 1
+	assert len(contents[0]) == 2
+	assert len(scores[0]) == 2
+	assert scores[0][0] > scores[0][1]
+
+	embeddings = await milvus_instance.fetch([ids[0]])
+	assert len(embeddings) == 1
+	assert len(embeddings[0]) == 768
+
+	exist = await milvus_instance.is_exist([ids[0], "doc3"])
+	assert len(exist) == 2
+	assert exist[0] is True
+	assert exist[1] is False
+
+
+@pytest.mark.asyncio
+async def test_delete_documents(milvus_instance):
+	# Add documents
+	ids = ["doc1", "doc2"]
+	texts = ["This is a test document.", "This is another test document."]
+	await milvus_instance.add(ids, texts)
+
+	await asyncio.sleep(1)
+
+	# Delete documents
+	await milvus_instance.delete([ids[0]])
+
+	# Query documents to ensure they are deleted
+	queries = ["test document"]
+	contents, scores = await milvus_instance.query(queries, top_k=2)
+
+	assert len(contents[0]) == 1
+	assert len(scores[0]) == 1

autorag-0.3.23/tests/autorag/vectordb/test_pinecone.py

@@ -0,0 +1,74 @@
+import asyncio
+import os
+
+import pytest
+
+from autorag.vectordb.pinecone import Pinecone
+from tests.delete_tests import is_github_action
+
+
+pytestmark = pytest.mark.skipif(
+	is_github_action(),
+	reason="This test needs pinecone api key",
+)
+
+
+@pytest.fixture
+def pinecone_instance():
+	pinecone = Pinecone(
+		embedding_model="mock",
+		index_name="nodonggeon",
+		dimension=768,  # mock embedding model has 768 dimensions
+		api_key=os.environ["PINECONE_API_KEY"],
+	)
+	yield pinecone
+	pinecone.delete_index()
+
+
+@pytest.mark.asyncio
+async def test_add_and_query_documents(pinecone_instance):
+	# Add documents
+	ids = ["doc1", "doc2"]
+	texts = ["This is a test document.", "This is another test document."]
+	await pinecone_instance.add(ids, texts)
+
+	await asyncio.sleep(10)
+
+	# Query documents
+	queries = ["test document"]
+	contents, scores = await pinecone_instance.query(queries, top_k=2)
+
+	assert len(contents) == 1
+	assert len(scores) == 1
+	assert len(contents[0]) == 2
+	assert len(scores[0]) == 2
+	assert scores[0][0] > scores[0][1]
+
+	embeddings = await pinecone_instance.fetch([ids[0]])
+	assert len(embeddings) == 1
+	assert len(embeddings[0]) == 768
+
+	exist = await pinecone_instance.is_exist([ids[0], "doc3"])
+	assert len(exist) == 2
+	assert exist[0] is True
+	assert exist[1] is False
+
+
+@pytest.mark.asyncio
+async def test_delete_documents(pinecone_instance):
+	# Add documents
+	ids = ["doc1", "doc2"]
+	texts = ["This is a test document.", "This is another test document."]
+	await pinecone_instance.add(ids, texts)
+
+	await asyncio.sleep(10)
+
+	# Delete documents
+	await pinecone_instance.delete([ids[0]])
+
+	# Query documents to ensure they are deleted
+	queries = ["test document"]
+	contents, scores = await pinecone_instance.query(queries, top_k=2)
+
+	assert len(contents[0]) == 1
+	assert len(scores[0]) == 1