improve handling of warnings

Author: blink1073

libs/langchain-mongodb/langchain_mongodb/retrievers/full_text_search.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Annotated, Any, Dict, List, Optional, Union
 
 from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun
@@ -46,7 +47,13 @@ def _get_relevant_documents(
         Returns:
             List of relevant documents
         """
-        default_k = self.k if self.k is not None else self.top_k
+        is_top_k_set = False
+        with warnings.catch_warnings():
+            # Ignore warning raised by checking the value of top_k.
+            warnings.simplefilter("ignore", DeprecationWarning)
+            if self.top_k is not None:
+                is_top_k_set = True
+        default_k = self.k if not is_top_k_set else self.top_k
         pipeline = text_search_stage(  # type: ignore
             query=query,
             search_field=self.search_field,

libs/langchain-mongodb/langchain_mongodb/retrievers/hybrid_search.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Annotated, Any, Dict, List, Optional
 
 from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun
@@ -83,7 +84,13 @@ def _get_relevant_documents(
         pipeline: List[Any] = []
 
         # Get the appropriate value for k.
-        default_k = self.top_k if self.top_k is not None else self.k
+        is_top_k_set = False
+        with warnings.catch_warnings():
+            # Ignore warnings raised by base class.
+            warnings.simplefilter("ignore", DeprecationWarning)
+            if self.top_k is not None:
+                is_top_k_set = True
+        default_k = self.k if not is_top_k_set else self.top_k
         k = kwargs.get("k", default_k)
 
         # First we build up the aggregation pipeline,

libs/langchain-mongodb/tests/integration_tests/conftest.py

@@ -1,4 +1,5 @@
 import os
+import warnings
 from typing import Generator, List
 
 import pytest
@@ -16,7 +17,10 @@
 def technical_report_pages() -> List[Document]:
     """Returns a Document for each of the 100 pages of a GPT-4 Technical Report"""
     loader = PyPDFLoader("https://arxiv.org/pdf/2303.08774.pdf")
-    pages = loader.load()
+    with warnings.catch_warnings():
+        # Ignore warnings raised by base class.
+        warnings.simplefilter("ignore", ResourceWarning)
+        pages = loader.load()
     return pages
 
 

libs/langchain-mongodb/tests/integration_tests/test_cache.py

@@ -157,6 +157,7 @@ def test_mongodb_cache(
         _execute_test(prompt, llm, response)
     finally:
         get_llm_cache().clear()
+        get_llm_cache().close()
 
 
 @pytest.mark.parametrize(
@@ -208,3 +209,4 @@ def test_mongodb_atlas_cache_matrix(
         generations=llm_generations, llm_output={}
     )
     get_llm_cache().clear()
+    get_llm_cache().close()

libs/langchain-mongodb/tests/integration_tests/test_chat_message_histories.py

@@ -1,4 +1,5 @@
 import json
+import warnings
 
 from langchain.memory import ConversationBufferMemory  # type: ignore[import-not-found]
 from langchain_core.messages import message_to_dict
@@ -19,9 +20,12 @@ def test_memory_with_message_store() -> None:
         database_name=DB_NAME,
         collection_name=COLLECTION,
     )
-    memory = ConversationBufferMemory(
-        memory_key="baz", chat_memory=message_history, return_messages=True
-    )
+    with warnings.catch_warnings():
+        # Ignore warnings raised by base class.
+        warnings.simplefilter("ignore", DeprecationWarning)
+        memory = ConversationBufferMemory(
+            memory_key="baz", chat_memory=message_history, return_messages=True
+        )
 
     # add some messages
     memory.chat_memory.add_ai_message("This is me, the AI")
@@ -38,3 +42,4 @@ def test_memory_with_message_store() -> None:
     memory.chat_memory.clear()
 
     assert memory.chat_memory.messages == []
+    memory.chat_memory.close()

libs/langchain-mongodb/tests/integration_tests/test_parent_document.py

@@ -74,3 +74,4 @@ def test_1clxn_retriever(
     assert len(responses) == 3
     assert all("GPT-4" in doc.page_content for doc in responses)
     assert {4, 5, 29} == set(doc.metadata["page"] for doc in responses)
+    client.close()

libs/langchain-mongodb/tests/integration_tests/test_retrievers.py

@@ -194,12 +194,14 @@ def test_hybrid_retriever_deprecated_top_k(
     )
 
     query1 = "When did I visit France?"
-    results = retriever.invoke(query1)
+    with pytest.warns(DeprecationWarning):
+        results = retriever.invoke(query1)
     assert len(results) == 3
     assert "Paris" in results[0].page_content
 
     query2 = "When was the last time I visited new orleans?"
-    results = retriever.invoke(query2)
+    with pytest.warns(DeprecationWarning):
+        results = retriever.invoke(query2)
     assert "New Orleans" in results[0].page_content
 
 

libs/langchain-mongodb/tests/integration_tests/test_retrievers_multi_field.py

@@ -207,12 +207,14 @@ def test_hybrid_retriever_deprecated_top_k(
     )
 
     query1 = "When did I visit France?"
-    results = retriever.invoke(query1)
+    with pytest.warns(DeprecationWarning):
+        results = retriever.invoke(query1)
     assert len(results) == 3
     assert "Paris" in results[0].page_content
 
     query2 = "When was the last time I visited new orleans?"
-    results = retriever.invoke(query2)
+    with pytest.warns(DeprecationWarning):
+        results = retriever.invoke(query2)
     assert "New Orleans" in results[0].page_content
 
 

libs/langchain-mongodb/tests/integration_tests/test_retrievers_standard.py

@@ -44,16 +44,16 @@ def setup_test() -> tuple[Collection, MongoDBAtlasVectorSearch]:
         text_key=PAGE_CONTENT_FIELD,
         auto_index_timeout=TIMEOUT,
     )
-
-    if coll.count_documents({}) == 0:
-        vs.add_documents(
-            [
-                Document(page_content="In 2023, I visited Paris"),
-                Document(page_content="In 2022, I visited New York"),
-                Document(page_content="In 2021, I visited New Orleans"),
-                Document(page_content="Sandwiches are beautiful. Sandwiches are fine."),
-            ]
-        )
+    coll.delete_many({})
+
+    vs.add_documents(
+        [
+            Document(page_content="In 2023, I visited Paris"),
+            Document(page_content="In 2022, I visited New York"),
+            Document(page_content="In 2021, I visited New Orleans"),
+            Document(page_content="Sandwiches are beautiful. Sandwiches are fine."),
+        ]
+    )
 
     # Set up the search index if needed.
     if not any([ix["name"] == SEARCH_INDEX_NAME for ix in coll.list_search_indexes()]):
@@ -68,16 +68,23 @@ def setup_test() -> tuple[Collection, MongoDBAtlasVectorSearch]:
 
 
 class TestMongoDBAtlasFullTextSearchRetriever(RetrieversIntegrationTests):
+    @classmethod
+    def setup_class(cls):
+        cls._coll, _ = setup_test()
+
+    @classmethod
+    def teardown_class(cls):
+        cls._coll.database.client.close()
+
     @property
     def retriever_constructor(self) -> Type[MongoDBAtlasFullTextSearchRetriever]:
         """Get a retriever for integration tests."""
         return MongoDBAtlasFullTextSearchRetriever
 
     @property
     def retriever_constructor_params(self) -> dict:
-        coll, _ = setup_test()
         return {
-            "collection": coll,
+            "collection": self._coll,
             "search_index_name": SEARCH_INDEX_NAME,
             "search_field": PAGE_CONTENT_FIELD,
         }
@@ -91,17 +98,24 @@ def retriever_query_example(self) -> str:
 
 
 class TestMongoDBAtlasHybridSearchRetriever(RetrieversIntegrationTests):
+    @classmethod
+    def setup_class(cls):
+        cls._coll, cls._vs = setup_test()
+
+    @classmethod
+    def teardown_class(cls):
+        cls._coll.database.client.close()
+
     @property
     def retriever_constructor(self) -> Type[MongoDBAtlasHybridSearchRetriever]:
         """Get a retriever for integration tests."""
         return MongoDBAtlasHybridSearchRetriever
 
     @property
     def retriever_constructor_params(self) -> dict:
-        coll, vs = setup_test()
         return {
-            "vectorstore": vs,
-            "collection": coll,
+            "vectorstore": self._vs,
+            "collection": self._coll,
             "search_index_name": SEARCH_INDEX_NAME,
             "search_field": PAGE_CONTENT_FIELD,
         }

libs/langchain-mongodb/tests/integration_tests/test_tools.py

@@ -12,55 +12,87 @@
 
 
 class TestQueryMongoDBDatabaseToolIntegration(ToolsIntegrationTests):
+    @classmethod
+    def setup_class(cls):
+        cls._db = create_database()
+
+    @classmethod
+    def teardown_class(cls):
+        cls._db.close()
+
     @property
     def tool_constructor(self) -> Type[QueryMongoDBDatabaseTool]:
         return QueryMongoDBDatabaseTool
 
     @property
     def tool_constructor_params(self) -> dict:
-        return dict(db=create_database())
+        return dict(db=self._db)
 
     @property
     def tool_invoke_params_example(self) -> dict:
         return dict(query='db.test.aggregate([{"$match": {}}])')
 
 
 class TestInfoMongoDBDatabaseToolIntegration(ToolsIntegrationTests):
+    @classmethod
+    def setup_class(cls):
+        cls._db = create_database()
+
+    @classmethod
+    def teardown_class(cls):
+        cls._db.close()
+
     @property
     def tool_constructor(self) -> Type[InfoMongoDBDatabaseTool]:
         return InfoMongoDBDatabaseTool
 
     @property
     def tool_constructor_params(self) -> dict:
-        return dict(db=create_database())
+        return dict(db=self._db)
 
     @property
     def tool_invoke_params_example(self) -> dict:
         return dict(collection_names="test")
 
 
 class TestListMongoDBDatabaseToolIntegration(ToolsIntegrationTests):
+    @classmethod
+    def setup_class(cls):
+        cls._db = create_database()
+
+    @classmethod
+    def teardown_class(cls):
+        cls._db.close()
+
     @property
     def tool_constructor(self) -> Type[ListMongoDBDatabaseTool]:
         return ListMongoDBDatabaseTool
 
     @property
     def tool_constructor_params(self) -> dict:
-        return dict(db=create_database())
+        return dict(db=self._db)
 
     @property
     def tool_invoke_params_example(self) -> dict:
         return dict()
 
 
 class TestQueryMongoDBCheckerToolIntegration(ToolsIntegrationTests):
+    @classmethod
+    def setup_class(cls):
+        cls._db = create_database()
+
+    @classmethod
+    def teardown_class(cls):
+        cls._db.close()
+
     @property
     def tool_constructor(self) -> Type[QueryMongoDBCheckerTool]:
         return QueryMongoDBCheckerTool
 
     @property
     def tool_constructor_params(self) -> dict:
-        return dict(db=create_database(), llm=create_llm())
+        return dict(db=self._db, llm=create_llm())
 
     @property
     def tool_invoke_params_example(self) -> dict: