Fix langchain tests (mlflow#10495)

Signed-off-by: Daniel Lok <[email protected]>
daniellok-db · Nov 24, 2023 · bf5eafb · bf5eafb
1 parent ddfeea6
commit bf5eafb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/mlflow/langchain/__init__.py b/mlflow/langchain/__init__.py
@@ -738,6 +738,7 @@ def predict(
         :return: Model predictions.
         """
         import langchain
+        from langchain.schema.retriever import BaseRetriever
 
         from mlflow.openai.utils import TEST_CONTENT, TEST_INTERMEDIATE_STEPS, TEST_SOURCE_DOCUMENTS
 
@@ -748,7 +749,7 @@ def predict(
             (
                 langchain.chains.llm.LLMChain,
                 langchain.chains.RetrievalQA,
-                langchain.schema.retriever.BaseRetriever,
+                BaseRetriever,
             ),
         ):
             mockContent = TEST_CONTENT

diff --git a/tests/langchain/test_langchain_model_export.py b/tests/langchain/test_langchain_model_export.py
@@ -485,7 +485,9 @@ def embed_query(self, text: str) -> List[float]:
 
 
 def assert_equal_retrievers(retriever, expected_retreiver):
-    assert isinstance(retriever, langchain.schema.retriever.BaseRetriever)
+    from langchain.schema.retriever import BaseRetriever
+
+    assert isinstance(retriever, BaseRetriever)
     assert isinstance(retriever, type(expected_retreiver))
     assert isinstance(retriever.vectorstore, type(expected_retreiver.vectorstore))
     assert retriever.tags == expected_retreiver.tags