Refactoring vectordb naming convention in embedchain.config (#1469)

2024-07-08 19:01:17 -04:00
parent 1a5d0d236a
commit 83e8c97295
20 changed files with 124 additions and 124 deletions
--- a/embedchain/config/vector_db/elasticsearch.py
+++ b/embedchain/config/vector_db/elasticsearch.py
@@ -0,0 +1,56 @@
+import os
+from typing import Optional, Union
+
+from embedchain.config.vector_db.base import BaseVectorDbConfig
+from embedchain.helpers.json_serializable import register_deserializable
+
+
+@register_deserializable
+class ElasticsearchDBConfig(BaseVectorDbConfig):
+    def __init__(
+        self,
+        collection_name: Optional[str] = None,
+        dir: Optional[str] = None,
+        es_url: Union[str, list[str]] = None,
+        cloud_id: Optional[str] = None,
+        batch_size: Optional[int] = 100,
+        **ES_EXTRA_PARAMS: dict[str, any],
+    ):
+        """
+        Initializes a configuration class instance for an Elasticsearch client.
+
+        :param collection_name: Default name for the collection, defaults to None
+        :type collection_name: Optional[str], optional
+        :param dir: Path to the database directory, where the database is stored, defaults to None
+        :type dir: Optional[str], optional
+        :param es_url: elasticsearch url or list of nodes url to be used for connection, defaults to None
+        :type es_url: Union[str, list[str]], optional
+        :param cloud_id: cloud id of the elasticsearch cluster, defaults to None
+        :type cloud_id: Optional[str], optional
+        :param batch_size: Number of items to insert in one batch, defaults to 100
+        :type batch_size: Optional[int], optional
+        :param ES_EXTRA_PARAMS: extra params dict that can be passed to elasticsearch.
+        :type ES_EXTRA_PARAMS: dict[str, Any], optional
+        """
+        if es_url and cloud_id:
+            raise ValueError("Only one of `es_url` and `cloud_id` can be set.")
+        # self, es_url: Union[str, list[str]] = None, **ES_EXTRA_PARAMS: dict[str, any]):
+        self.ES_URL = es_url or os.environ.get("ELASTICSEARCH_URL")
+        self.CLOUD_ID = cloud_id or os.environ.get("ELASTICSEARCH_CLOUD_ID")
+        if not self.ES_URL and not self.CLOUD_ID:
+            raise AttributeError(
+                "Elasticsearch needs a URL or CLOUD_ID attribute, "
+                "this can either be passed to `ElasticsearchDBConfig` or as `ELASTICSEARCH_URL` or `ELASTICSEARCH_CLOUD_ID` in `.env`"  # noqa: E501
+            )
+        self.ES_EXTRA_PARAMS = ES_EXTRA_PARAMS
+        # Load API key from .env if it's not explicitly passed.
+        # Can only set one of 'api_key', 'basic_auth', and 'bearer_auth'
+        if (
+            not self.ES_EXTRA_PARAMS.get("api_key")
+            and not self.ES_EXTRA_PARAMS.get("basic_auth")
+            and not self.ES_EXTRA_PARAMS.get("bearer_auth")
+        ):
+            self.ES_EXTRA_PARAMS["api_key"] = os.environ.get("ELASTICSEARCH_API_KEY")
+
+        self.batch_size = batch_size
+        super().__init__(collection_name=collection_name, dir=dir)