From 22a5b53efbc6a4eb3220561b31b895c747c270c3 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Sun, 19 May 2024 06:50:48 -0700 Subject: [PATCH] [MLflow] Fix langchain tests (#12054) Signed-off-by: Sunish Sheth --- tests/langchain/test_langchain_model_export.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/langchain/test_langchain_model_export.py b/tests/langchain/test_langchain_model_export.py index d863775cb172a..d8e37550f47eb 100644 --- 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",