Bring back client attribute on ElasticsearchStore (#29)

langchain-ai · May 16, 2024 · 87b82f5 · 87b82f5
1 parent fd84e71
commit 87b82f5
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/libs/elasticsearch/langchain_elasticsearch/vectorstores.py b/libs/elasticsearch/langchain_elasticsearch/vectorstores.py
@@ -717,6 +717,7 @@ def __init__(
         )
 
         self.embedding = embedding
+        self.client = self._store.client
         self._embedding_service = embedding_service
         self.query_field = query_field
         self.vector_query_field = vector_query_field

diff --git a/libs/elasticsearch/tests/integration_tests/test_vectorstores.py b/libs/elasticsearch/tests/integration_tests/test_vectorstores.py
@@ -660,7 +660,7 @@ def assert_query(
             )
 
         # 2. check query result is okay
-        es_output = docsearch._store.client.search(
+        es_output = docsearch.client.search(
             index=index_name,
             query={
                 "bool": {

diff --git a/libs/elasticsearch/tests/unit_tests/test_vectorstores.py b/libs/elasticsearch/tests/unit_tests/test_vectorstores.py
@@ -228,7 +228,7 @@ def dummy_custom_query(query_body: dict, query: Optional[str]) -> Dict[str, Any]
         return {"dummy": "query"}
 
     def test_agent_header(self, store: ElasticsearchStore) -> None:
-        agent = store._store.client._headers["User-Agent"]
+        agent = store.client._headers["User-Agent"]
         assert (
             re.match(r"^langchain-py-vs/\d+\.\d+\.\d+(?:rc\d+)?$", agent) is not None
         ), f"The string '{agent}' does not match the expected pattern."