my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import os
#2	from typing import Optional, Union
#3
#4	from embedchain.config.vector_db.base import BaseVectorDbConfig
#5	from embedchain.helpers.json_serializable import register_deserializable
#6
#7
#8	@register_deserializable
#9	class ElasticsearchDBConfig(BaseVectorDbConfig):
#10	def __init__(
#11	self,
#12	collection_name: Optional[str] = None,
#13	dir: Optional[str] = None,
#14	es_url: Union[str, list[str]] = None,
#15	cloud_id: Optional[str] = None,
#16	batch_size: Optional[int] = 100,
#17	**ES_EXTRA_PARAMS: dict[str, any],
#18	):
#19	"""
#20	Initializes a configuration class instance for an Elasticsearch client.
#21
#22	:param collection_name: Default name for the collection, defaults to None
#23	:type collection_name: Optional[str], optional
#24	:param dir: Path to the database directory, where the database is stored, defaults to None
#25	:type dir: Optional[str], optional
#26	:param es_url: elasticsearch url or list of nodes url to be used for connection, defaults to None
#27	:type es_url: Union[str, list[str]], optional
#28	:param cloud_id: cloud id of the elasticsearch cluster, defaults to None
#29	:type cloud_id: Optional[str], optional
#30	:param batch_size: Number of items to insert in one batch, defaults to 100
#31	:type batch_size: Optional[int], optional
#32	:param ES_EXTRA_PARAMS: extra params dict that can be passed to elasticsearch.
#33	:type ES_EXTRA_PARAMS: dict[str, Any], optional
#34	"""
#35	if es_url and cloud_id:
#36	raise ValueError("Only one of `es_url` and `cloud_id` can be set.")
#37	# self, es_url: Union[str, list[str]] = None, **ES_EXTRA_PARAMS: dict[str, any]):
#38	self.ES_URL = es_url or os.environ.get("ELASTICSEARCH_URL")
#39	self.CLOUD_ID = cloud_id or os.environ.get("ELASTICSEARCH_CLOUD_ID")
#40	if not self.ES_URL and not self.CLOUD_ID:
#41	raise AttributeError(
#42	"Elasticsearch needs a URL or CLOUD_ID attribute, "
#43	"this can either be passed to `ElasticsearchDBConfig` or as `ELASTICSEARCH_URL` or `ELASTICSEARCH_CLOUD_ID` in `.env`" # noqa: E501
#44	)
#45	self.ES_EXTRA_PARAMS = ES_EXTRA_PARAMS
#46	# Load API key from .env if it's not explicitly passed.
#47	# Can only set one of 'api_key', 'basic_auth', and 'bearer_auth'
#48	if (
#49	not self.ES_EXTRA_PARAMS.get("api_key")
#50	and not self.ES_EXTRA_PARAMS.get("basic_auth")
#51	and not self.ES_EXTRA_PARAMS.get("bearer_auth")
#52	):
#53	self.ES_EXTRA_PARAMS["api_key"] = os.environ.get("ELASTICSEARCH_API_KEY")
#54
#55	self.batch_size = batch_size
#56	super().__init__(collection_name=collection_name, dir=dir)
#57

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public