AutoRAG 0.0.0__py3-none-any.whl

autorag/validator.py

@@ -0,0 +1,98 @@
+import itertools
+import logging
+import os
+import tempfile
+
+import pandas as pd
+
+from autorag.evaluator import Evaluator
+from autorag.utils import (
+	cast_qa_dataset,
+	cast_corpus_dataset,
+	validate_qa_from_corpus_dataset,
+)
+
+logger = logging.getLogger("AutoRAG")
+
+
+class Validator:
+	def __init__(self, qa_data_path: str, corpus_data_path: str):
+		"""
+		Initialize a Validator object.
+
+		:param qa_data_path: The path to the QA dataset.
+		    Must be parquet file.
+		:param corpus_data_path: The path to the corpus dataset.
+		    Must be parquet file.
+		"""
+		# validate data paths
+		if not os.path.exists(qa_data_path):
+			raise ValueError(f"QA data path {qa_data_path} does not exist.")
+		if not os.path.exists(corpus_data_path):
+			raise ValueError(f"Corpus data path {corpus_data_path} does not exist.")
+		if not qa_data_path.endswith(".parquet"):
+			raise ValueError(f"QA data path {qa_data_path} is not a parquet file.")
+		if not corpus_data_path.endswith(".parquet"):
+			raise ValueError(
+				f"Corpus data path {corpus_data_path} is not a parquet file."
+			)
+		self.qa_data = pd.read_parquet(qa_data_path, engine="pyarrow")
+		self.corpus_data = pd.read_parquet(corpus_data_path, engine="pyarrow")
+		self.qa_data = cast_qa_dataset(self.qa_data)
+		self.corpus_data = cast_corpus_dataset(self.corpus_data)
+
+	def validate(self, yaml_path: str, qa_cnt: int = 5, random_state: int = 42):
+		# Determine the sample size and log a warning if qa_cnt is larger than available records
+		available_records = len(self.qa_data)
+		safe_sample_size = min(qa_cnt, available_records)  # 먼저 safe_sample_size 계산
+
+		if safe_sample_size < qa_cnt:
+			logger.warning(
+				f"Minimal Requested sample size ({qa_cnt}) is larger than available records ({available_records}). "
+				f"Sampling will be limited to {safe_sample_size} records. "
+			)
+
+		# safe sample QA data
+		sample_qa_df = self.qa_data.sample(
+			n=safe_sample_size, random_state=random_state
+		)
+		sample_qa_df.reset_index(drop=True, inplace=True)
+
+		# get doc_id
+		temp_qa_df = sample_qa_df.copy(deep=True)
+		flatten_retrieval_gts = (
+			temp_qa_df["retrieval_gt"]
+			.apply(lambda x: list(itertools.chain.from_iterable(x)))
+			.tolist()
+		)
+		target_doc_ids = list(itertools.chain.from_iterable(flatten_retrieval_gts))
+
+		# make sample corpus data
+		sample_corpus_df = self.corpus_data.loc[
+			self.corpus_data["doc_id"].isin(target_doc_ids)
+		]
+		sample_corpus_df.reset_index(drop=True, inplace=True)
+
+		validate_qa_from_corpus_dataset(sample_qa_df, sample_corpus_df)
+
+		# start Evaluate at temp project directory
+		with (
+			tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as qa_path,
+			tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as corpus_path,
+			tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_project_dir,
+		):
+			sample_qa_df.to_parquet(qa_path.name, index=False)
+			sample_corpus_df.to_parquet(corpus_path.name, index=False)
+
+			evaluator = Evaluator(
+				qa_data_path=qa_path.name,
+				corpus_data_path=corpus_path.name,
+				project_dir=temp_project_dir,
+			)
+			evaluator.start_trial(yaml_path, skip_validation=True)
+			qa_path.close()
+			corpus_path.close()
+			os.unlink(qa_path.name)
+			os.unlink(corpus_path.name)
+
+		logger.info("Validation complete.")

autorag/vectordb/__init__.py

@@ -0,0 +1,75 @@
+import os
+from typing import List
+
+from autorag.support import dynamically_find_function
+from autorag.utils.util import load_yaml_config
+from autorag.vectordb.base import BaseVectorStore
+
+
+def get_support_vectordb(vectordb_name: str):
+	support_vectordb = {
+		"chroma": ("autorag.vectordb.chroma", "Chroma"),
+		"Chroma": ("autorag.vectordb.chroma", "Chroma"),
+		"milvus": ("autorag.vectordb.milvus", "Milvus"),
+		"Milvus": ("autorag.vectordb.milvus", "Milvus"),
+		"weaviate": ("autorag.vectordb.weaviate", "Weaviate"),
+		"Weaviate": ("autorag.vectordb.weaviate", "Weaviate"),
+		"pinecone": ("autorag.vectordb.pinecone", "Pinecone"),
+		"Pinecone": ("autorag.vectordb.pinecone", "Pinecone"),
+		"couchbase": ("autorag.vectordb.couchbase", "Couchbase"),
+		"Couchbase": ("autorag.vectordb.couchbase", "Couchbase"),
+		"qdrant": ("autorag.vectordb.qdrant", "Qdrant"),
+		"Qdrant": ("autorag.vectordb.qdrant", "Qdrant"),
+	}
+	return dynamically_find_function(vectordb_name, support_vectordb)
+
+
+def load_vectordb(vectordb_name: str, **kwargs):
+	vectordb = get_support_vectordb(vectordb_name)
+	return vectordb(**kwargs)
+
+
+def load_vectordb_from_yaml(yaml_path: str, vectordb_name: str, project_dir: str):
+	config_dict = load_yaml_config(yaml_path)
+	vectordb_list = config_dict.get("vectordb", [])
+	if len(vectordb_list) == 0 or vectordb_name == "default":
+		chroma_path = os.path.join(project_dir, "resources", "chroma")
+		return load_vectordb(
+			"chroma",
+			client_type="persistent",
+			embedding_model="openai",
+			collection_name="openai",
+			path=chroma_path,
+		)
+
+	target_dict = list(filter(lambda x: x["name"] == vectordb_name, vectordb_list))
+	target_dict[0].pop("name")  # delete a name key
+	target_vectordb_name = target_dict[0].pop("db_type")
+	target_vectordb_params = target_dict[0]
+	return load_vectordb(target_vectordb_name, **target_vectordb_params)
+
+
+def load_all_vectordb_from_yaml(
+	yaml_path: str, project_dir: str
+) -> List[BaseVectorStore]:
+	config_dict = load_yaml_config(yaml_path)
+	vectordb_list = config_dict.get("vectordb", [])
+	if len(vectordb_list) == 0:
+		chroma_path = os.path.join(project_dir, "resources", "chroma")
+		return [
+			load_vectordb(
+				"chroma",
+				client_type="persistent",
+				embedding_model="openai",
+				collection_name="openai",
+				path=chroma_path,
+			)
+		]
+
+	result_vectordbs = []
+	for vectordb_dict in vectordb_list:
+		_ = vectordb_dict.pop("name")
+		vectordb_type = vectordb_dict.pop("db_type")
+		vectordb = load_vectordb(vectordb_type, **vectordb_dict)
+		result_vectordbs.append(vectordb)
+	return result_vectordbs

autorag/vectordb/base.py

@@ -0,0 +1,73 @@
+from abc import abstractmethod
+from typing import List, Tuple, Union
+
+from llama_index.embeddings.openai import OpenAIEmbedding
+
+from autorag.utils.util import openai_truncate_by_token
+from autorag.embedding.base import EmbeddingModel
+
+
+class BaseVectorStore:
+	support_similarity_metrics = ["l2", "ip", "cosine"]
+
+	def __init__(
+		self,
+		embedding_model: Union[str, List[dict]],
+		similarity_metric: str = "cosine",
+		embedding_batch: int = 100,
+	):
+		self.embedding = EmbeddingModel.load(embedding_model)()
+		self.embedding_batch = embedding_batch
+		self.embedding.embed_batch_size = embedding_batch
+		assert similarity_metric in self.support_similarity_metrics, (
+			f"search method {similarity_metric} is not supported"
+		)
+		self.similarity_metric = similarity_metric
+
+	@abstractmethod
+	async def add(
+		self,
+		ids: List[str],
+		texts: List[str],
+	):
+		pass
+
+	@abstractmethod
+	def add_embedding(self, ids: List[str], embeddings: List[List[float]]):
+		"""
+		Add the embeddings to the Vector DB.
+		"""
+		pass
+
+	@abstractmethod
+	async def query(
+		self, queries: List[str], top_k: int, **kwargs
+	) -> Tuple[List[List[str]], List[List[float]]]:
+		pass
+
+	@abstractmethod
+	async def fetch(self, ids: List[str]) -> List[List[float]]:
+		"""
+		Fetch the embeddings of the ids.
+		"""
+		pass
+
+	@abstractmethod
+	async def is_exist(self, ids: List[str]) -> List[bool]:
+		"""
+		Check if the ids exist in the Vector DB.
+		"""
+		pass
+
+	@abstractmethod
+	async def delete(self, ids: List[str]):
+		pass
+
+	def truncated_inputs(self, inputs: List[str]) -> List[str]:
+		if isinstance(self.embedding, OpenAIEmbedding):
+			openai_embedding_limit = 8000
+			results = openai_truncate_by_token(
+				inputs, openai_embedding_limit, self.embedding.model_name
+			)
+			return results
+		return inputs

autorag/vectordb/chroma.py

@@ -0,0 +1,118 @@
+from typing import List, Optional, Dict, Tuple, Union
+
+from chromadb import (
+	EphemeralClient,
+	PersistentClient,
+	DEFAULT_TENANT,
+	DEFAULT_DATABASE,
+	CloudClient,
+	AsyncHttpClient,
+)
+from chromadb.api.models.AsyncCollection import AsyncCollection
+from chromadb.api.types import QueryResult
+
+from autorag.utils.util import apply_recursive
+from autorag.vectordb.base import BaseVectorStore
+
+
+class Chroma(BaseVectorStore):
+	def __init__(
+		self,
+		embedding_model: Union[str, List[dict]],
+		collection_name: str,
+		embedding_batch: int = 100,
+		client_type: str = "persistent",
+		similarity_metric: str = "cosine",
+		path: str = None,
+		host: str = "localhost",
+		port: int = 8000,
+		ssl: bool = False,
+		headers: Optional[Dict[str, str]] = None,
+		api_key: Optional[str] = None,
+		tenant: str = DEFAULT_TENANT,
+		database: str = DEFAULT_DATABASE,
+	):
+		super().__init__(embedding_model, similarity_metric, embedding_batch)
+		if client_type == "ephemeral":
+			self.client = EphemeralClient(tenant=tenant, database=database)
+		elif client_type == "persistent":
+			assert path is not None, "path must be provided for persistent client"
+			self.client = PersistentClient(path=path, tenant=tenant, database=database)
+		elif client_type == "http":
+			self.client = AsyncHttpClient(
+				host=host,
+				port=port,
+				ssl=ssl,
+				headers=headers,
+				tenant=tenant,
+				database=database,
+			)
+		elif client_type == "cloud":
+			self.client = CloudClient(
+				tenant=tenant,
+				database=database,
+				api_key=api_key,
+			)
+		else:
+			raise ValueError(
+				f"client_type {client_type} is not supported\n"
+				"supported client types are: ephemeral, persistent, http, cloud"
+			)
+
+		self.collection = self.client.get_or_create_collection(
+			name=collection_name,
+			metadata={"hnsw:space": similarity_metric},
+		)
+
+	async def add(self, ids: List[str], texts: List[str]):
+		texts = self.truncated_inputs(texts)
+		text_embeddings = await self.embedding.aget_text_embedding_batch(texts)
+		if isinstance(self.collection, AsyncCollection):
+			await self.collection.add(ids=ids, embeddings=text_embeddings)
+		else:
+			self.collection.add(ids=ids, embeddings=text_embeddings)
+
+	def add_embedding(self, ids: List[str], embeddings: List[List[float]]):
+		self.collection.add(ids=ids, embeddings=embeddings)
+
+	async def fetch(self, ids: List[str]) -> List[List[float]]:
+		if isinstance(self.collection, AsyncCollection):
+			fetch_result = await self.collection.get(ids, include=["embeddings"])
+		else:
+			fetch_result = self.collection.get(ids, include=["embeddings"])
+		fetch_embeddings = fetch_result["embeddings"]
+		return fetch_embeddings
+
+	async def is_exist(self, ids: List[str]) -> List[bool]:
+		if isinstance(self.collection, AsyncCollection):
+			fetched_result = await self.collection.get(ids, include=[])
+		else:
+			fetched_result = self.collection.get(ids, include=[])
+		existed_ids = fetched_result["ids"]
+		return list(map(lambda x: x in existed_ids, ids))
+
+	async def query(
+		self, queries: List[str], top_k: int, **kwargs
+	) -> Tuple[List[List[str]], List[List[float]]]:
+		queries = self.truncated_inputs(queries)
+		query_embeddings: List[
+			List[float]
+		] = await self.embedding.aget_text_embedding_batch(queries)
+		if isinstance(self.collection, AsyncCollection):
+			query_result: QueryResult = await self.collection.query(
+				query_embeddings=query_embeddings, n_results=top_k
+			)
+		else:
+			query_result: QueryResult = self.collection.query(
+				query_embeddings=query_embeddings, n_results=top_k
+			)
+		ids = query_result["ids"]
+		scores = query_result["distances"]
+		scores = apply_recursive(lambda x: 1 - x, scores)
+		return ids, scores
+
+	async def delete(self, ids: List[str]):
+		if isinstance(self.collection, AsyncCollection):
+			await self.collection.delete(ids)
+		else:
+			self.collection.delete(ids)

autorag/vectordb/couchbase.py

@@ -0,0 +1,239 @@
+import logging
+
+from datetime import timedelta
+
+from couchbase.auth import PasswordAuthenticator
+from couchbase.cluster import Cluster
+from couchbase.options import ClusterOptions
+
+from typing import List, Tuple, Optional, Union
+
+from autorag.utils.util import make_batch
+from autorag.vectordb import BaseVectorStore
+
+logger = logging.getLogger("AutoRAG")
+
+
+class Couchbase(BaseVectorStore):
+	def __init__(
+		self,
+		embedding_model: Union[str, List[dict]],
+		bucket_name: str,
+		scope_name: str,
+		collection_name: str,
+		index_name: str,
+		embedding_batch: int = 100,
+		connection_string: str = "",
+		username: str = "",
+		password: str = "",
+		ingest_batch: int = 100,
+		text_key: Optional[str] = "text",
+		embedding_key: Optional[str] = "embedding",
+		scoped_index: bool = True,
+	):
+		super().__init__(
+			embedding_model=embedding_model,
+			similarity_metric="ip",
+			embedding_batch=embedding_batch,
+		)
+
+		self.index_name = index_name
+		self.bucket_name = bucket_name
+		self.scope_name = scope_name
+		self.collection_name = collection_name
+		self.scoped_index = scoped_index
+		self.text_key = text_key
+		self.embedding_key = embedding_key
+		self.ingest_batch = ingest_batch
+
+		auth = PasswordAuthenticator(username, password)
+		self.cluster = Cluster(connection_string, ClusterOptions(auth))
+
+		# Wait until the cluster is ready for use.
+		self.cluster.wait_until_ready(timedelta(seconds=5))
+
+		# Check if the bucket exists
+		if not self._check_bucket_exists():
+			raise ValueError(
+				f"Bucket {self.bucket_name} does not exist. "
+				" Please create the bucket before searching."
+			)
+
+		try:
+			self.bucket = self.cluster.bucket(self.bucket_name)
+			self.scope = self.bucket.scope(self.scope_name)
+			self.collection = self.scope.collection(self.collection_name)
+		except Exception as e:
+			raise ValueError(
+				"Error connecting to couchbase. "
+				"Please check the connection and credentials."
+			) from e
+
+		# Check if the index exists. Throws ValueError if it doesn't
+		try:
+			self._check_index_exists()
+		except Exception:
+			raise
+
+		# Reinitialize to ensure a consistent state
+		self.bucket = self.cluster.bucket(self.bucket_name)
+		self.scope = self.bucket.scope(self.scope_name)
+		self.collection = self.scope.collection(self.collection_name)
+
+	async def add(self, ids: List[str], texts: List[str]):
+		from couchbase.exceptions import DocumentExistsException
+
+		texts = self.truncated_inputs(texts)
+		text_embeddings: List[
+			List[float]
+		] = await self.embedding.aget_text_embedding_batch(texts)
+
+		documents_to_insert = []
+		for _id, text, embedding in zip(ids, texts, text_embeddings):
+			doc = {
+				self.text_key: text,
+				self.embedding_key: embedding,
+			}
+			documents_to_insert.append({_id: doc})
+
+		batch_documents_to_insert = make_batch(documents_to_insert, self.ingest_batch)
+
+		for batch in batch_documents_to_insert:
+			insert_batch = {}
+			for doc in batch:
+				insert_batch.update(doc)
+			try:
+				self.collection.upsert_multi(insert_batch)
+			except DocumentExistsException as e:
+				logger.debug(f"Document already exists: {e}")
+
+	def add_embedding(self, ids: List[str], embeddings: List[List[float]]):
+		from couchbase.exceptions import DocumentExistsException
+
+		documents_to_insert = []
+		for _id, embedding in zip(ids, embeddings):
+			doc = {
+				self.embedding_key: embedding,
+			}
+			documents_to_insert.append({_id: doc})
+
+		batch_documents_to_insert = make_batch(documents_to_insert, self.ingest_batch)
+
+		for batch in batch_documents_to_insert:
+			insert_batch = {}
+			for doc in batch:
+				insert_batch.update(doc)
+			try:
+				self.collection.upsert_multi(insert_batch)
+			except DocumentExistsException as e:
+				logger.debug(f"Document already exists: {e}")
+
+	async def fetch(self, ids: List[str]) -> List[List[float]]:
+		# Fetch vectors by IDs
+		fetched_result = self.collection.get_multi(ids)
+		fetched_vectors = {
+			k: v.value[f"{self.embedding_key}"]
+			for k, v in fetched_result.results.items()
+		}
+		return list(map(lambda x: fetched_vectors[x], ids))
+
+	async def is_exist(self, ids: List[str]) -> List[bool]:
+		existed_result = self.collection.exists_multi(ids)
+		existed_ids = {k: v.exists for k, v in existed_result.results.items()}
+		return list(map(lambda x: existed_ids[x], ids))
+
+	async def query(
+		self, queries: List[str], top_k: int, **kwargs
+	) -> Tuple[List[List[str]], List[List[float]]]:
+		import couchbase.search as search
+		from couchbase.options import SearchOptions
+		from couchbase.vector_search import VectorQuery, VectorSearch
+
+		queries = self.truncated_inputs(queries)
+		query_embeddings: List[
+			List[float]
+		] = await self.embedding.aget_text_embedding_batch(queries)
+
+		ids, scores = [], []
+		for query_embedding in query_embeddings:
+			# Create Search Request
+			search_req = search.SearchRequest.create(
+				VectorSearch.from_vector_query(
+					VectorQuery(
+						self.embedding_key,
+						query_embedding,
+						top_k,
+					)
+				)
+			)
+
+			# Search
+			if self.scoped_index:
+				search_iter = self.scope.search(
+					self.index_name,
+					search_req,
+					SearchOptions(limit=top_k),
+				)
+
+			else:
+				search_iter = self.cluster.search(
+					self.index_name,
+					search_req,
+					SearchOptions(limit=top_k),
+				)
+
+			# Parse the search results
+			# search_iter.rows() can only be iterated once.
+			id_list, score_list = [], []
+			for result in search_iter.rows():
+				id_list.append(result.id)
+				score_list.append(result.score)
+
+			ids.append(id_list)
+			scores.append(score_list)
+
+		return ids, scores
+
+	async def delete(self, ids: List[str]):
+		self.collection.remove_multi(ids)
+
+	def _check_bucket_exists(self) -> bool:
+		"""Check if the bucket exists in the linked Couchbase cluster.
+
+		Returns:
+		    True if the bucket exists
+		"""
+		bucket_manager = self.cluster.buckets()
+		try:
+			bucket_manager.get_bucket(self.bucket_name)
+			return True
+		except Exception as e:
+			logger.debug("Error checking if bucket exists:", e)
+			return False
+
+	def _check_index_exists(self) -> bool:
+		"""Check if the Search index exists in the linked Couchbase cluster
+		Returns:
+		    bool: True if the index exists, False otherwise.
+		    Raises a ValueError if the index does not exist.
+		"""
+		if self.scoped_index:
+			all_indexes = [
+				index.name for index in self.scope.search_indexes().get_all_indexes()
+			]
+			if self.index_name not in all_indexes:
+				raise ValueError(
+					f"Index {self.index_name} does not exist. "
+					" Please create the index before searching."
+				)
+		else:
+			all_indexes = [
+				index.name for index in self.cluster.search_indexes().get_all_indexes()
+			]
+			if self.index_name not in all_indexes:
+				raise ValueError(
+					f"Index {self.index_name} does not exist. "
+					" Please create the index before searching."
+				)
+
+		return True