[Improvements] allow setting up the elasticsearch cloud instance (#997)

Co-authored-by: Deven Patel <deven298@yahoo.com>
This commit is contained in:
Deven Patel
2023-12-07 14:17:59 -08:00
committed by GitHub
parent d62a23edf6
commit 0ea8ab228c
4 changed files with 49 additions and 21 deletions

View File

@@ -57,7 +57,7 @@ class BaseLlmConfig(BaseConfig):
def __init__(
self,
number_documents: int = 1,
number_documents: int = 3,
template: Optional[Template] = None,
model: Optional[str] = None,
temperature: float = 0,

View File

@@ -12,6 +12,7 @@ class ElasticsearchDBConfig(BaseVectorDbConfig):
collection_name: Optional[str] = None,
dir: Optional[str] = None,
es_url: Union[str, List[str]] = None,
cloud_id: Optional[str] = None,
**ES_EXTRA_PARAMS: Dict[str, any],
):
"""
@@ -26,12 +27,15 @@ class ElasticsearchDBConfig(BaseVectorDbConfig):
: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")
if not self.ES_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 attribute, "
"this can either be passed to `ElasticsearchDBConfig` or as `ELASTICSEARCH_URL` in `.env`"
"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.
@@ -40,7 +44,6 @@ class ElasticsearchDBConfig(BaseVectorDbConfig):
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")
and not self.ES_EXTRA_PARAMS.get("http_auth")
):
self.ES_EXTRA_PARAMS["api_key"] = os.environ.get("ELASTICSEARCH_API_KEY")
super().__init__(collection_name=collection_name, dir=dir)