Update unit test

databricks · Jan 16, 2025 · 4cdcaea · 4cdcaea
1 parent a213f4c
commit 4cdcaea
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py b/integrations/langchain/tests/unit_tests/test_vector_search_retriever_tool.py
@@ -5,6 +5,7 @@
 from databricks_ai_bridge.test_utils.vector_search import (  # noqa: F401
     ALL_INDEX_NAMES,
     DELTA_SYNC_INDEX,
+    INPUT_TEXTS,
     mock_vs_client,
     mock_workspace_client,
 )
@@ -15,6 +16,7 @@
 from databricks_langchain import ChatDatabricks, VectorSearchRetrieverTool
 from tests.utils.chat_models import llm, mock_client  # noqa: F401
 from tests.utils.vector_search import EMBEDDING_MODEL
+import json
 
 
 def init_vector_search_tool(
@@ -115,3 +117,7 @@ def test_vs_tool_tracing(index_name: str, tool_name: Optional[str]) -> None:
     trace = mlflow.get_last_active_trace()
     spans = trace.search_spans(name=tool_name or index_name, span_type=SpanType.RETRIEVER)
     assert len(spans) == 1
+    inputs = json.loads(trace.to_dict()["data"]["spans"][0]["attributes"]["mlflow.spanInputs"])
+    assert inputs["query"] == "Databricks Agent Framework"
+    outputs = json.loads(trace.to_dict()["data"]["spans"][0]["attributes"]["mlflow.spanOutputs"])
+    assert [d["page_content"] in INPUT_TEXTS for d in outputs]