Merge pull request #7 from weni-ai/es-timeout-config

add elasticsearch timeout config
weni-ai · Jan 5, 2024 · 761eeed · 761eeed
2 parents c04d873 + 9fa0ad9
commit 761eeed
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/app/config.py b/app/config.py
@@ -24,3 +24,4 @@ def __init__(self):
  "HUGGINGFACE_API_TOKEN", "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"
  ),
  }
+ self.es_timeout = os.environ.get("ELASTICSEARCH_TIMEOUT", "30")
diff --git a/app/main.py b/app/main.py
@@ -1,3 +1,4 @@
+from elasticsearch import Elasticsearch
 from fastapi import FastAPI
 from langchain.embeddings import SagemakerEndpointEmbeddings, HuggingFaceHubEmbeddings
 from langchain.embeddings.base import Embeddings
@@ -44,6 +45,9 @@ def __init__(self, config: AppConfig):
  index_name=config.product_index_name,
  embedding=self.embeddings,
  )
+ self.vectorstore.client = Elasticsearch(
+ hosts=config.es_url, timeout=int(config.es_timeout)
+ )
  self.elasticStore = ElasticsearchVectorStoreIndex(self.vectorstore)
  self.products_indexer = ProductsIndexer(self.elasticStore)
  self.products_handler = ProductsHandler(self.products_indexer)