tests on imports fix

langchain-ai · Apr 23, 2024 · 24aae95 · 24aae95
1 parent 940dc42
commit 24aae95
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/libs/elasticsearch/langchain_elasticsearch/__init__.py b/libs/elasticsearch/langchain_elasticsearch/__init__.py
@@ -11,11 +11,11 @@
 
 __all__ = [
     "ApproxRetrievalStrategy",
+    "ElasticsearchCache",
     "ElasticsearchChatMessageHistory",
     "ElasticsearchEmbeddings",
     "ElasticsearchRetriever",
     "ElasticsearchStore",
     "ExactRetrievalStrategy",
     "SparseRetrievalStrategy",
-    "ElasticsearchCache"
 ]
diff --git a/libs/elasticsearch/langchain_elasticsearch/cache.py b/libs/elasticsearch/langchain_elasticsearch/cache.py
@@ -134,7 +134,7 @@ def lookup(self, prompt: str, llm_string: str) -> Optional[RETURN_VAL_TYPE]:
                 index=self._index_name,
                 body={
                     "query": {"term": {"_id": cache_key}},
-                    "sort": {"timestamp" : {"order" : "asc"}}
+                    "sort": {"timestamp": {"order": "asc"}}
                 },
                 source_includes=["llm_output"],
             )

diff --git a/libs/elasticsearch/tests/unit_tests/test_imports.py b/libs/elasticsearch/tests/unit_tests/test_imports.py
@@ -2,6 +2,7 @@
 
 EXPECTED_ALL = [
     "ApproxRetrievalStrategy",
+    "ElasticsearchCache",
     "ElasticsearchChatMessageHistory",
     "ElasticsearchEmbeddings",
     "ElasticsearchRetriever",