[MLflow] Fix langchain tests (#12054)

Signed-off-by: Sunish Sheth <sunishsheth2009@gmail.com>
mlflow · May 19, 2024 · 22a5b53 · 22a5b53
1 parent 9b36854
commit 22a5b53
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/langchain/test_langchain_model_export.py b/tests/langchain/test_langchain_model_export.py
@@ -2416,11 +2416,11 @@ def test_save_load_chain_as_code(chain_model_signature):
     )
     assert response["choices"][0]["message"]["content"] == "Databricks"
     trace = mlflow.get_trace(tracer._request_id)
-    assert trace.info.tags["vector_search_index"] == json.dumps(
+    assert trace.info.tags[DependenciesSchemasType.RETRIEVERS.value] == json.dumps(
         [
             {
                 "doc_uri": "doc-uri",
-                "name": "vector_search_index",
+                "name": "retriever",
                 "other_columns": ["column1", "column2"],
                 "primary_key": "primary-key",
                 "text_column": "text-column",