AutoRAG 0.0.0__py3-none-any.whl

autorag/vectordb/milvus.py

@@ -0,0 +1,169 @@
+import logging
+from typing import Any, Dict, List, Tuple, Optional, Union
+
+from pymilvus import (
+	DataType,
+	FieldSchema,
+	CollectionSchema,
+	connections,
+	Collection,
+	MilvusException,
+)
+from pymilvus.orm import utility
+
+from autorag.utils.util import apply_recursive
+from autorag.vectordb import BaseVectorStore
+
+
+logger = logging.getLogger("AutoRAG")
+
+
+class Milvus(BaseVectorStore):
+	def __init__(
+		self,
+		embedding_model: Union[str, List[dict]],
+		collection_name: str,
+		embedding_batch: int = 100,
+		similarity_metric: str = "cosine",
+		index_type: str = "IVF_FLAT",
+		uri: str = "http://localhost:19530",
+		db_name: str = "",
+		token: str = "",
+		user: str = "",
+		password: str = "",
+		timeout: Optional[float] = None,
+		params: Dict[str, Any] = {},
+	):
+		super().__init__(embedding_model, similarity_metric, embedding_batch)
+
+		# Connect to Milvus server
+		connections.connect(
+			"default",
+			uri=uri,
+			token=token,
+			db_name=db_name,
+			user=user,
+			password=password,
+		)
+		self.collection_name = collection_name
+		self.timeout = timeout
+		self.params = params
+		self.index_type = index_type
+
+		# Set Collection
+		if not utility.has_collection(collection_name, timeout=timeout):
+			# Get the dimension of the embeddings
+			test_embedding_result: List[float] = self.embedding.get_query_embedding(
+				"test"
+			)
+			dimension = len(test_embedding_result)
+
+			pk = FieldSchema(
+				name="id",
+				dtype=DataType.VARCHAR,
+				max_length=128,
+				is_primary=True,
+				auto_id=False,
+			)
+			field = FieldSchema(
+				name="vector", dtype=DataType.FLOAT_VECTOR, dim=dimension
+			)
+			schema = CollectionSchema(fields=[pk, field])
+
+			self.collection = Collection(name=self.collection_name, schema=schema)
+			index_params = {
+				"metric_type": self.similarity_metric.upper(),
+				"index_type": self.index_type.upper(),
+				"params": self.params,
+			}
+			self.collection.create_index(
+				field_name="vector", index_params=index_params, timeout=self.timeout
+			)
+		else:
+			self.collection = Collection(name=self.collection_name)
+
+	async def add(self, ids: List[str], texts: List[str]):
+		texts = self.truncated_inputs(texts)
+		text_embeddings: List[
+			List[float]
+		] = await self.embedding.aget_text_embedding_batch(texts)
+		self.add_embedding(ids, text_embeddings)
+
+	def add_embedding(self, ids: List[str], embeddings: List[List[float]]):
+		data = list(
+			map(lambda _id, vector: {"id": _id, "vector": vector}, ids, embeddings)
+		)
+
+		# Insert data into the collection
+		res = self.collection.insert(data=data, timeout=self.timeout)
+		assert res.insert_count == len(ids), (
+			f"Insertion failed. Try to insert {len(ids)} but only {res['insert_count']} inserted."
+		)
+
+		self.collection.flush(timeout=self.timeout)
+
+	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)
+
+		self.collection.load(timeout=self.timeout)
+
+		# Perform similarity search
+		results = self.collection.search(
+			data=query_embeddings,
+			limit=top_k,
+			anns_field="vector",
+			param={"metric_type": self.similarity_metric.upper()},
+			timeout=self.timeout,
+			**kwargs,
+		)
+
+		# Extract IDs and distances
+		ids = [[str(hit.id) for hit in result] for result in results]
+		distances = [[hit.distance for hit in result] for result in results]
+
+		if self.similarity_metric in ["l2"]:
+			distances = apply_recursive(lambda x: -x, distances)
+
+		return ids, distances
+
+	async def fetch(self, ids: List[str]) -> List[List[float]]:
+		try:
+			self.collection.load(timeout=self.timeout)
+		except MilvusException as e:
+			logger.warning(f"Failed to load collection: {e}")
+			return [[]] * len(ids)
+		# Fetch vectors by IDs
+		results = self.collection.query(
+			expr=f"id in {ids}", output_fields=["id", "vector"], timeout=self.timeout
+		)
+		id_vector_dict = {str(result["id"]): result["vector"] for result in results}
+		result = [id_vector_dict[_id] for _id in ids]
+		return result
+
+	async def is_exist(self, ids: List[str]) -> List[bool]:
+		try:
+			self.collection.load(timeout=self.timeout)
+		except MilvusException:
+			return [False] * len(ids)
+		# Check the existence of IDs
+		results = self.collection.query(
+			expr=f"id in {ids}", output_fields=["id"], timeout=self.timeout
+		)
+		# Determine existence
+		existing_ids = {str(result["id"]) for result in results}
+		return [str(_id) in existing_ids for _id in ids]
+
+	async def delete(self, ids: List[str]):
+		# Delete entries by IDs
+		self.collection.delete(expr=f"id in {ids}", timeout=self.timeout)
+
+	def delete_collection(self):
+		# Delete the collection
+		self.collection.release(timeout=self.timeout)
+		self.collection.drop_index(timeout=self.timeout)
+		self.collection.drop(timeout=self.timeout)

autorag/vectordb/pinecone.py

@@ -0,0 +1,121 @@
+import logging
+
+from pinecone.grpc import PineconeGRPC as Pinecone_client
+from pinecone import ServerlessSpec
+
+from typing import List, Optional, Tuple, Union
+
+from autorag.utils.util import make_batch, apply_recursive
+from autorag.vectordb import BaseVectorStore
+
+logger = logging.getLogger("AutoRAG")
+
+
+class Pinecone(BaseVectorStore):
+	def __init__(
+		self,
+		embedding_model: Union[str, List[dict]],
+		index_name: str,
+		embedding_batch: int = 100,
+		dimension: int = 1536,
+		similarity_metric: str = "cosine",  # "cosine", "dotproduct", "euclidean"
+		cloud: Optional[str] = "aws",
+		region: Optional[str] = "us-east-1",
+		api_key: Optional[str] = None,
+		deletion_protection: Optional[str] = "disabled",  # "enabled" or "disabled"
+		namespace: Optional[str] = "default",
+		ingest_batch: int = 200,
+	):
+		super().__init__(embedding_model, similarity_metric, embedding_batch)
+
+		self.index_name = index_name
+		self.namespace = namespace
+		self.ingest_batch = ingest_batch
+
+		self.client = Pinecone_client(api_key=api_key)
+
+		if similarity_metric == "ip":
+			similarity_metric = "dotproduct"
+		elif similarity_metric == "l2":
+			similarity_metric = "euclidean"
+
+		if not self.client.has_index(index_name):
+			self.client.create_index(
+				name=index_name,
+				dimension=dimension,
+				metric=similarity_metric,
+				spec=ServerlessSpec(
+					cloud=cloud,
+					region=region,
+				),
+				deletion_protection=deletion_protection,
+			)
+		self.index = self.client.Index(index_name)
+
+	async def add(self, ids: List[str], texts: List[str]):
+		texts = self.truncated_inputs(texts)
+		text_embeddings: List[
+			List[float]
+		] = await self.embedding.aget_text_embedding_batch(texts)
+		self.add_embedding(ids, text_embeddings)
+
+	def add_embedding(self, ids: List[str], embeddings: List[List[float]]):
+		vector_tuples = list(zip(ids, embeddings))
+		batch_vectors = make_batch(vector_tuples, self.ingest_batch)
+
+		async_res = [
+			self.index.upsert(
+				vectors=batch_vector_tuples,
+				namespace=self.namespace,
+				async_req=True,
+			)
+			for batch_vector_tuples in batch_vectors
+		]
+		# Wait for the async requests to finish
+		[async_result.result() for async_result in async_res]
+
+	async def fetch(self, ids: List[str]) -> List[List[float]]:
+		results = self.index.fetch(ids=ids, namespace=self.namespace)
+		id_vector_dict = {
+			str(key): val["values"] for key, val in results["vectors"].items()
+		}
+		result = [id_vector_dict[_id] for _id in ids]
+		return result
+
+	async def is_exist(self, ids: List[str]) -> List[bool]:
+		fetched_result = self.index.fetch(ids=ids, namespace=self.namespace)
+		existed_ids = list(map(str, fetched_result.get("vectors", {}).keys()))
+		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)
+
+		ids, scores = [], []
+		for query_embedding in query_embeddings:
+			response = self.index.query(
+				vector=query_embedding,
+				top_k=top_k,
+				include_values=True,
+				namespace=self.namespace,
+			)
+
+			ids.append([o.id for o in response.matches])
+			scores.append([o.score for o in response.matches])
+
+		if self.similarity_metric in ["l2"]:
+			scores = apply_recursive(lambda x: -x, scores)
+
+		return ids, scores
+
+	async def delete(self, ids: List[str]):
+		# Delete entries by IDs
+		self.index.delete(ids=ids, namespace=self.namespace)
+
+	def delete_index(self):
+		# Delete the index
+		self.client.delete_index(self.index_name)

autorag/vectordb/qdrant.py

@@ -0,0 +1,155 @@
+import logging
+
+from qdrant_client import QdrantClient
+from qdrant_client.models import (
+	Distance,
+	VectorParams,
+	PointStruct,
+	PointIdsList,
+	HasIdCondition,
+	Filter,
+	SearchRequest,
+)
+
+from typing import List, Tuple, Union
+
+from autorag.vectordb import BaseVectorStore
+
+logger = logging.getLogger("AutoRAG")
+
+
+class Qdrant(BaseVectorStore):
+	def __init__(
+		self,
+		embedding_model: Union[str, List[dict]],
+		collection_name: str,
+		embedding_batch: int = 100,
+		similarity_metric: str = "cosine",
+		client_type: str = "docker",
+		url: str = "http://localhost:6333",
+		host: str = "",
+		api_key: str = "",
+		dimension: int = 1536,
+		ingest_batch: int = 64,
+		parallel: int = 1,
+		max_retries: int = 3,
+	):
+		super().__init__(embedding_model, similarity_metric, embedding_batch)
+
+		self.collection_name = collection_name
+		self.ingest_batch = ingest_batch
+		self.parallel = parallel
+		self.max_retries = max_retries
+
+		if similarity_metric == "cosine":
+			distance = Distance.COSINE
+		elif similarity_metric == "ip":
+			distance = Distance.DOT
+		elif similarity_metric == "l2":
+			distance = Distance.EUCLID
+		else:
+			raise ValueError(
+				f"similarity_metric {similarity_metric} is not supported\n"
+				"supported similarity metrics are: cosine, ip, l2"
+			)
+
+		if client_type == "docker":
+			self.client = QdrantClient(
+				url=url,
+			)
+		elif client_type == "cloud":
+			self.client = QdrantClient(
+				host=host,
+				api_key=api_key,
+			)
+		else:
+			raise ValueError(
+				f"client_type {client_type} is not supported\n"
+				"supported client types are: docker, cloud"
+			)
+
+		if not self.client.collection_exists(collection_name):
+			self.client.create_collection(
+				collection_name,
+				vectors_config=VectorParams(
+					size=dimension,
+					distance=distance,
+				),
+			)
+		self.collection = self.client.get_collection(collection_name)
+
+	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)
+		self.add_embedding(ids, text_embeddings)
+
+	def add_embedding(self, ids: List[str], embeddings: List[List[float]]):
+		points = list(
+			map(lambda x: PointStruct(id=x[0], vector=x[1]), zip(ids, embeddings))
+		)
+
+		self.client.upload_points(
+			collection_name=self.collection_name,
+			points=points,
+			batch_size=self.ingest_batch,
+			parallel=self.parallel,
+			max_retries=self.max_retries,
+			wait=True,
+		)
+
+	async def fetch(self, ids: List[str]) -> List[List[float]]:
+		# Fetch vectors by IDs
+		fetched_results = self.client.retrieve(
+			collection_name=self.collection_name,
+			ids=ids,
+			with_vectors=True,
+		)
+		return list(map(lambda x: x.vector, fetched_results))
+
+	async def is_exist(self, ids: List[str]) -> List[bool]:
+		existed_result = self.client.scroll(
+			collection_name=self.collection_name,
+			scroll_filter=Filter(
+				must=[
+					HasIdCondition(has_id=ids),
+				],
+			),
+		)
+		# existed_result is tuple. So we use existed_result[0] to get list of Record
+		existed_ids = list(map(lambda x: x.id, existed_result[0]))
+		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)
+
+		search_queries = list(
+			map(
+				lambda x: SearchRequest(vector=x, limit=top_k, with_vector=True),
+				query_embeddings,
+			)
+		)
+
+		search_result = self.client.search_batch(
+			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]
+
+		return ids, scores
+
+	async def delete(self, ids: List[str]):
+		self.client.delete(
+			collection_name=self.collection_name,
+			points_selector=PointIdsList(points=ids),
+		)
+
+	def delete_collection(self):
+		# Delete the collection
+		self.client.delete_collection(self.collection_name)

autorag/vectordb/weaviate.py

@@ -0,0 +1,184 @@
+import logging
+
+import weaviate
+from weaviate.classes.init import Auth
+from weaviate.classes.config import Property, DataType
+import weaviate.classes as wvc
+from weaviate.classes.query import MetadataQuery
+
+from typing import List, Optional, Tuple, Union
+
+from autorag.vectordb import BaseVectorStore
+
+logger = logging.getLogger("AutoRAG")
+
+
+class Weaviate(BaseVectorStore):
+	def __init__(
+		self,
+		embedding_model: Union[str, List[dict]],
+		collection_name: str,
+		embedding_batch: int = 100,
+		similarity_metric: str = "cosine",
+		client_type: str = "docker",
+		host: str = "localhost",
+		port: int = 8080,
+		grpc_port: int = 50051,
+		url: Optional[str] = None,
+		api_key: Optional[str] = None,
+		text_key: str = "content",
+	):
+		super().__init__(embedding_model, similarity_metric, embedding_batch)
+
+		self.text_key = text_key
+
+		if client_type == "docker":
+			self.client = weaviate.connect_to_local(
+				host=host,
+				port=port,
+				grpc_port=grpc_port,
+			)
+		elif client_type == "cloud":
+			self.client = weaviate.connect_to_weaviate_cloud(
+				cluster_url=url,
+				auth_credentials=Auth.api_key(api_key),
+			)
+		else:
+			raise ValueError(
+				f"client_type {client_type} is not supported\n"
+				"supported client types are: docker, cloud"
+			)
+		if similarity_metric == "cosine":
+			distance_metric = wvc.config.VectorDistances.COSINE
+		elif similarity_metric == "ip":
+			distance_metric = wvc.config.VectorDistances.DOT
+		elif similarity_metric == "l2":
+			distance_metric = wvc.config.VectorDistances.L2_SQUARED
+		else:
+			raise ValueError(
+				f"similarity_metric {similarity_metric} is not supported\n"
+				"supported similarity metrics are: cosine, ip, l2"
+			)
+
+		if not self.client.collections.exists(collection_name):
+			self.client.collections.create(
+				collection_name,
+				properties=[
+					Property(
+						name="content", data_type=DataType.TEXT, skip_vectorization=True
+					),
+				],
+				vectorizer_config=wvc.config.Configure.Vectorizer.none(),
+				vector_index_config=wvc.config.Configure.VectorIndex.hnsw(  # hnsw, flat, dynamic,
+					distance_metric=distance_metric
+				),
+			)
+		self.collection = self.client.collections.get(collection_name)
+		self.collection_name = collection_name
+
+	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)
+
+		with self.client.batch.dynamic() as batch:
+			for i, text in enumerate(texts):
+				data_properties = {self.text_key: text}
+
+				batch.add_object(
+					collection=self.collection_name,
+					properties=data_properties,
+					uuid=ids[i],
+					vector=text_embeddings[i],
+				)
+
+		failed_objs = self.client.batch.failed_objects
+		for obj in failed_objs:
+			err_message = (
+				f"Failed to add object: {obj.original_uuid}\nReason: {obj.message}"
+			)
+
+			logger.error(err_message)
+
+	def add_embedding(self, ids: List[str], embeddings: List[List[float]]):
+		with self.client.batch.dynamic() as batch:
+			for i in range(len(ids)):
+				batch.add_object(
+					collection=self.collection_name,
+					uuid=ids[i],
+					vector=embeddings[i],
+				)
+
+		failed_objs = self.client.batch.failed_objects
+		for obj in failed_objs:
+			err_message = (
+				f"Failed to add object: {obj.original_uuid}\nReason: {obj.message}"
+			)
+
+			logger.error(err_message)
+
+	async def fetch(self, ids: List[str]) -> List[List[float]]:
+		# Fetch vectors by IDs
+		results = self.collection.query.fetch_objects(
+			filters=wvc.query.Filter.by_property("_id").contains_any(ids),
+			include_vector=True,
+		)
+		id_vector_dict = {
+			str(object.uuid): object.vector["default"] for object in results.objects
+		}
+		result = [id_vector_dict[_id] for _id in ids]
+		return result
+
+	async def is_exist(self, ids: List[str]) -> List[bool]:
+		fetched_result = self.collection.query.fetch_objects(
+			filters=wvc.query.Filter.by_property("_id").contains_any(ids),
+		)
+		existed_ids = [str(result.uuid) for result in fetched_result.objects]
+		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)
+
+		ids, scores = [], []
+		for query_embedding in query_embeddings:
+			response = self.collection.query.near_vector(
+				near_vector=query_embedding,
+				limit=top_k,
+				return_metadata=MetadataQuery(distance=True),
+			)
+
+			ids.append([o.uuid for o in response.objects])
+			scores.append(
+				[
+					distance_to_score(o.metadata.distance, self.similarity_metric)
+					for o in response.objects
+				]
+			)
+
+		return ids, scores
+
+	async def delete(self, ids: List[str]):
+		filter = wvc.query.Filter.by_id().contains_any(ids)
+		self.collection.data.delete_many(where=filter)
+
+	def delete_collection(self):
+		# Delete the collection
+		self.client.collections.delete(self.collection_name)
+
+
+def distance_to_score(distance: float, similarity_metric) -> float:
+	if similarity_metric == "cosine":
+		return 1 - distance
+	elif similarity_metric == "ip":
+		return -distance
+	elif similarity_metric == "l2":
+		return -distance
+	else:
+		raise ValueError(
+			f"similarity_metric {similarity_metric} is not supported\n"
+			"supported similarity metrics are: cosine, ip, l2"
+		)

autorag/web.py

@@ -0,0 +1,81 @@
+from typing import Optional
+
+import click
+import streamlit as st
+
+from autorag.deploy import Runner
+
+
+def get_runner(
+	yaml_path: Optional[str], project_dir: Optional[str], trial_path: Optional[str]
+):
+	if not yaml_path and not trial_path:
+		raise ValueError("yaml_path or trial_path must be given.")
+	elif yaml_path and trial_path:
+		raise ValueError("yaml_path and trial_path cannot be given at the same time.")
+	elif yaml_path:
+		return Runner.from_yaml(yaml_path, project_dir=project_dir)
+	elif trial_path:
+		return Runner.from_trial_folder(trial_path)
+
+
+def set_initial_state():
+	if "messages" not in st.session_state:
+		st.session_state["messages"] = [
+			{
+				"role": "assistant",
+				"content": "Welcome !",
+			}
+		]
+
+
+def set_page_config():
+	st.set_page_config(
+		page_title="AutoRAG",
+		page_icon="🤖",
+		layout="wide",
+		initial_sidebar_state="expanded",
+		menu_items={
+			"Get help": "https://github.com/Marker-Inc-Korea/AutoRAG/discussions",
+			"Report a bug": "https://github.com/Marker-Inc-Korea/AutoRAG/issues",
+		},
+	)
+
+
+def set_page_header():
+	st.header("📚 AutoRAG", anchor=False)
+	st.caption("Input a question and get an answer from the given documents. ")
+
+
+def chat_box(runner: Runner):
+	if query := st.chat_input("How can I help?"):
+		# Add the user input to messages state
+		st.session_state["messages"].append({"role": "user", "content": query})
+		with st.chat_message("user"):
+			st.markdown(query)
+
+		# Generate llama-index stream with user input
+		with st.chat_message("assistant"):
+			with st.spinner("Processing..."):
+				response = st.write(runner.run(query))
+
+		# Add the final response to messages state
+		st.session_state["messages"].append({"role": "assistant", "content": response})
+
+
+@click.command()
+@click.option("--yaml_path", type=str, help="Path to the YAML file.")
+@click.option("--project_dir", type=str, help="Path to the project directory.")
+@click.option("--trial_path", type=str, help="Path to the trial directory.")
+def run_web_server(
+	yaml_path: Optional[str], project_dir: Optional[str], trial_path: Optional[str]
+):
+	runner = get_runner(yaml_path, project_dir, trial_path)
+	set_initial_state()
+	set_page_config()
+	set_page_header()
+	chat_box(runner)
+
+
+if __name__ == "__main__":
+	run_web_server()