diff --git a/src/vanna/opensearch/opensearch_vector.py b/src/vanna/opensearch/opensearch_vector.py index c49e2bd7..7fab1ecd 100644 --- a/src/vanna/opensearch/opensearch_vector.py +++ b/src/vanna/opensearch/opensearch_vector.py @@ -155,6 +155,11 @@ def __init__(self, config=None): else: max_retries = 10 + if config is not None and "es_http_compress" in config: + es_http_compress = config["es_http_compress"] + else: + es_http_compress = False + print("OpenSearch_VectorStore initialized with es_urls: ", es_urls, " host: ", host, " port: ", port, " ssl: ", ssl, " verify_certs: ", verify_certs, " timeout: ", timeout, " max_retries: ", max_retries) @@ -162,7 +167,7 @@ def __init__(self, config=None): # Initialize the OpenSearch client by passing a list of URLs self.client = OpenSearch( hosts=[es_urls], - http_compress=True, + http_compress=es_http_compress, use_ssl=ssl, verify_certs=verify_certs, timeout=timeout, @@ -175,7 +180,7 @@ def __init__(self, config=None): # Initialize the OpenSearch client by passing a host and port self.client = OpenSearch( hosts=[{'host': host, 'port': port}], - http_compress=True, + http_compress=es_http_compress, use_ssl=ssl, verify_certs=verify_certs, timeout=timeout, @@ -267,6 +272,7 @@ def get_related_ddl(self, question: str, **kwargs) -> List[str]: } } } + print(query) response = self.client.search(index=self.ddl_index, body=query, **kwargs) return [hit['_source']['ddl'] for hit in response['hits']['hits']] @@ -279,6 +285,7 @@ def get_related_documentation(self, question: str, **kwargs) -> List[str]: } } } + print(query) response = self.client.search(index=self.document_index, body=query, **kwargs) @@ -292,6 +299,7 @@ def get_similar_question_sql(self, question: str, **kwargs) -> List[str]: } } } + print(query) response = self.client.search(index=self.question_sql_index, body=query, **kwargs) @@ -307,6 +315,7 @@ def get_training_data(self, **kwargs) -> pd.DataFrame: body={"query": {"match_all": {}}}, size=1000 ) + print(query) # records = [hit['_source'] for hit in response['hits']['hits']] for hit in response['hits']['hits']: data.append(