diff --git a/osschat_insert.py b/osschat_insert.py index f7946f7..f34ae70 100644 --- a/osschat_insert.py +++ b/osschat_insert.py @@ -39,7 +39,7 @@ class OSSChatInsertConfig(BaseModel, extra=Extra.allow): milvus_password: Optional[str] = None # config for elasticsearch es_enable: Optional[bool] = True - es_connection_kwargs: Optional[dict] = {hosts=['https://127.0.0.1:9200'], basic_auth=('elastic', 'my_password')} + es_connection_kwargs: Optional[dict] = {'hosts': ['https://127.0.0.1:9200'], 'basic_auth': ('elastic', 'my_password')} _hf_models = ops.sentence_embedding.transformers().get_op().supported_model_names() @@ -105,12 +105,7 @@ def osschat_insert_pipe(config): p = p.map(('project_name', 'doc', 'sentence', 'embedding'), 'milvus_res', insert_milvus_op) if config.es_enable: - es_index_op = ops.elasticsearch.osschat_index(host=config.es_host, - port=config.es_port, - user=config.es_user, - password=config.es_password, - ca_certs=config.es_ca_certs, - ) + es_index_op = ops.elasticsearch.osschat_index(**config.es_connection_kwargs) p = ( p.map('sentence', 'es_sentence', lambda x: {'sentence': x}) .map(('project_name', 'es_sentence'), 'es_res', es_index_op)