from collections import Counter, defaultdict from unittest import mock import langchain import pytest from databricks.vector_search.client import VectorSearchIndex from packaging.version import Version from mlflow.langchain.databricks_dependencies import ( _detect_databricks_dependencies, _extract_databricks_dependencies_from_chat_model, _extract_databricks_dependencies_from_llm, _extract_databricks_dependencies_from_retriever, _extract_dependency_list_from_lc_model, ) from mlflow.models.resources import ( DatabricksFunction, DatabricksServingEndpoint, DatabricksSQLWarehouse, DatabricksVectorSearchIndex, ) # TODO: Remove this once databricks-langchain supports v1 if Version(langchain.__version__).major >= 1: pytest.skip("databricks-langchain does not support v1 yet", allow_module_level=True) class MockDatabricksServingEndpointClient: def __init__( self, host: str, api_token: str, endpoint_name: str, databricks_uri: str, task: str, ): self.host = host self.api_token = api_token self.endpoint_name = endpoint_name self.databricks_uri = databricks_uri self.task = task def _is_partner_package_installed(): try: import databricks_langchain # noqa: F401 return True except ImportError: return False def remove_langchain_community(monkeypatch): # Simulate the environment where langchain_community is not installed original_import = __import__ def mock_import(name, *args, **kwargs): if name.startswith("langchain_community"): raise ImportError("No module named 'langchain_community'") return original_import(name, *args, **kwargs) monkeypatch.setattr("builtins.__import__", mock_import) def test_parsing_dependency_from_databricks_llm(monkeypatch: pytest.MonkeyPatch): from langchain_community.llms import Databricks from mlflow.langchain.utils.logging import IS_PICKLE_SERIALIZATION_RESTRICTED monkeypatch.setattr( "langchain_community.llms.databricks._DatabricksServingEndpointClient", MockDatabricksServingEndpointClient, ) monkeypatch.setenv("DATABRICKS_HOST", "my-default-host") monkeypatch.setenv("DATABRICKS_TOKEN", "my-default-token") llm_kwargs = {"endpoint_name": "databricks-mixtral-8x7b-instruct"} if IS_PICKLE_SERIALIZATION_RESTRICTED: llm_kwargs["allow_dangerous_deserialization"] = True llm = Databricks(**llm_kwargs) resources = list(_extract_databricks_dependencies_from_llm(llm)) assert resources == [ DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct") ] class MockVectorSearchIndex(VectorSearchIndex): def __init__(self, endpoint_name, index_name, has_embedding_endpoint=False) -> None: self.endpoint_name = endpoint_name self.name = index_name self.has_embedding_endpoint = has_embedding_endpoint def describe(self): if self.has_embedding_endpoint: return { "name": self.name, "endpoint_name": self.endpoint_name, "primary_key": "id", "index_type": "DELTA_SYNC", "delta_sync_index_spec": { "source_table": "ml.schema.databricks_documentation", "embedding_source_columns": [ {"name": "content", "embedding_model_endpoint_name": "embedding-model"} ], "pipeline_type": "TRIGGERED", "pipeline_id": "79a76fcc-67ad-4ac6-8d8e-20f7d485ffa6", }, "status": { "detailed_state": "OFFLINE_FAILED", "message": "Index creation failed.", "indexed_row_count": 0, "failed_status": {"error_message": ""}, "ready": False, "index_url": "e2-dogfood.staging.cloud.databricks.com/rest_of_url", }, "creator": "first.last@databricks.com", } else: return { "name": self.name, "endpoint_name": self.endpoint_name, "primary_key": "id", "index_type": "DELTA_SYNC", "delta_sync_index_spec": { "source_table": "ml.schema.databricks_documentation", "embedding_vector_columns": [], "pipeline_type": "TRIGGERED", "pipeline_id": "fbbd5bf1-2b9b-4a7e-8c8d-c0f6cc1030de", }, "status": { "detailed_state": "ONLINE", "message": "Index is currently online", "indexed_row_count": 17183, "ready": True, "index_url": "e2-dogfood.staging.cloud.databricks.com/rest_of_url", }, "creator": "first.last@databricks.com", } def get_vector_search( endpoint_name: str, index_name: str, has_embedding_endpoint=False, **kwargs, ): index = MockVectorSearchIndex(endpoint_name, index_name, has_embedding_endpoint) from databricks_langchain import DatabricksVectorSearch with mock.patch("databricks.vector_search.client.VectorSearchClient") as mock_client: mock_client().get_index.return_value = index return DatabricksVectorSearch( endpoint=endpoint_name, index_name=index_name, **kwargs, ) def test_parsing_dependency_from_databricks_retriever(monkeypatch): from databricks_langchain import ChatDatabricks, DatabricksEmbeddings remove_langchain_community(monkeypatch) with pytest.raises(ImportError, match="No module named 'langchain_community"): from langchain_community.embeddings import DatabricksEmbeddings embedding_model = DatabricksEmbeddings(endpoint="databricks-bge-large-en") # vs_index_1 is a direct access index vectorstore_1 = get_vector_search( endpoint_name="vs_endpoint", index_name="mlflow.rag.vs_index_1", text_column="content", embedding=embedding_model, ) retriever_1 = vectorstore_1.as_retriever() # vs_index_2 has builtin embedding endpoint "embedding-model" vectorstore_2 = get_vector_search( endpoint_name="vs_endpoint", index_name="mlflow.rag.vs_index_2", has_embedding_endpoint=True, ) retriever_2 = vectorstore_2.as_retriever() llm = ChatDatabricks(endpoint="databricks-llama-2-70b-chat", temperature=0) assert list(_extract_databricks_dependencies_from_retriever(retriever_1)) == [ DatabricksVectorSearchIndex(index_name="mlflow.rag.vs_index_1"), DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"), ] assert list(_extract_databricks_dependencies_from_retriever(retriever_2)) == [ DatabricksVectorSearchIndex(index_name="mlflow.rag.vs_index_2"), DatabricksServingEndpoint(endpoint_name="embedding-model"), ] try: from langchain.retrievers import ( ContextualCompressionRetriever, EnsembleRetriever, TimeWeightedVectorStoreRetriever, ) from langchain.retrievers.document_compressors import LLMChainExtractor from langchain.retrievers.multi_query import MultiQueryRetriever except ImportError: from langchain_classic.retrievers import ( ContextualCompressionRetriever, EnsembleRetriever, TimeWeightedVectorStoreRetriever, ) from langchain_classic.retrievers.document_compressors import LLMChainExtractor from langchain_classic.retrievers.multi_query import MultiQueryRetriever multi_query_retriever = MultiQueryRetriever.from_llm(retriever=retriever_1, llm=llm) assert list(_extract_databricks_dependencies_from_retriever(multi_query_retriever)) == [ DatabricksVectorSearchIndex(index_name="mlflow.rag.vs_index_1"), DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"), ] compressor = LLMChainExtractor.from_llm(llm) compression_retriever = ContextualCompressionRetriever( base_compressor=compressor, base_retriever=retriever_1 ) assert list(_extract_databricks_dependencies_from_retriever(compression_retriever)) == [ DatabricksVectorSearchIndex(index_name="mlflow.rag.vs_index_1"), DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"), ] ensemble_retriever = EnsembleRetriever( retrievers=[retriever_1, retriever_2], weights=[0.5, 0.5] ) assert list(_extract_databricks_dependencies_from_retriever(ensemble_retriever)) == [ DatabricksVectorSearchIndex(index_name="mlflow.rag.vs_index_1"), DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"), DatabricksVectorSearchIndex(index_name="mlflow.rag.vs_index_2"), DatabricksServingEndpoint(endpoint_name="embedding-model"), ] time_weighted_retriever = TimeWeightedVectorStoreRetriever( vectorstore=vectorstore_1, decay_rate=0.0000000000000000000000001, k=1 ) assert list(_extract_databricks_dependencies_from_retriever(time_weighted_retriever)) == [ DatabricksVectorSearchIndex(index_name="mlflow.rag.vs_index_1"), DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"), ] def test_parsing_dependency_from_retriever_with_embedding_endpoint_in_index(): vectorstore = get_vector_search( endpoint_name="dbdemos_vs_endpoint", index_name="mlflow.rag.vs_index", has_embedding_endpoint=True, ) retriever = vectorstore.as_retriever() resources = list(_extract_databricks_dependencies_from_retriever(retriever)) assert resources == [ DatabricksVectorSearchIndex(index_name="mlflow.rag.vs_index"), DatabricksServingEndpoint(endpoint_name="embedding-model"), ] def test_parsing_dependency_from_agent(monkeypatch: pytest.MonkeyPatch): from databricks.sdk.service.catalog import FunctionInfo from databricks_langchain import ChatDatabricks from langchain.agents import initialize_agent try: from langchain_community.tools.databricks import UCFunctionToolkit except Exception: return # When get is called return a function def mock_function_get(self, function_name): components = function_name.split(".") # Initialize agent used below requires functions to take in exactly one parameter param_dict = { "parameters": [ { "name": "param", "parameter_type": "PARAM", "position": 0, "type_json": '{"name":"param","type":"string","nullable":true,"metadata":{}}', "type_name": "STRING", "type_precision": 0, "type_scale": 0, "type_text": "string", } ] } # Add the catalog, schema and name to the function Info followed by the parameter return FunctionInfo.from_dict({ "catalog_name": components[0], "schema_name": components[1], "name": components[2], "input_params": param_dict, }) monkeypatch.setenv("DATABRICKS_HOST", "my-default-host") monkeypatch.setenv("DATABRICKS_TOKEN", "my-default-token") monkeypatch.setattr("databricks.sdk.service.catalog.FunctionsAPI.get", mock_function_get) toolkit = UCFunctionToolkit(warehouse_id="testId1").include("rag.test.test_function") llm = ChatDatabricks(endpoint="databricks-llama-2-70b-chat", temperature=0) agent = initialize_agent( toolkit.get_tools(), llm, verbose=True, ) resources = sorted(_extract_dependency_list_from_lc_model(agent), key=lambda x: x.name) assert resources == [ DatabricksServingEndpoint(endpoint_name="databricks-llama-2-70b-chat"), DatabricksFunction(function_name="rag.test.test_function"), DatabricksSQLWarehouse(warehouse_id="testId1"), ] def test_parsing_multiple_dependency_from_agent(monkeypatch): from databricks.sdk.service.catalog import FunctionInfo from databricks_langchain import ChatDatabricks from langchain.agents import initialize_agent from langchain.tools.retriever import create_retriever_tool remove_langchain_community(monkeypatch) def mock_function_get(self, function_name): components = function_name.split(".") param_dict = { "parameters": [ { "name": "param", "parameter_type": "PARAM", "position": 0, "type_json": '{"name":"param","type":"string","nullable":true,"metadata":{}}', "type_name": "STRING", "type_precision": 0, "type_scale": 0, "type_text": "string", } ] } return FunctionInfo.from_dict({ "catalog_name": components[0], "schema_name": components[1], "name": components[2], "input_params": param_dict, }) # In addition to above now handle the case where a '*' is passed in and list all the functions def mock_function_list(self, catalog_name, schema_name): assert catalog_name == "rag" assert schema_name == "test" return [ FunctionInfo(full_name="rag.test.test_function"), FunctionInfo(full_name="rag.test.test_function_2"), FunctionInfo(full_name="rag.test.test_function_3"), ] monkeypatch.setenv("DATABRICKS_HOST", "my-default-host") monkeypatch.setenv("DATABRICKS_TOKEN", "my-default-token") monkeypatch.setattr("databricks.sdk.service.catalog.FunctionsAPI.get", mock_function_get) monkeypatch.setattr("databricks.sdk.service.catalog.FunctionsAPI.list", mock_function_list) include_uc_function_tools = False try: from langchain_community.tools.databricks import UCFunctionToolkit include_uc_function_tools = True except Exception: include_uc_function_tools = False uc_function_tools = ( (UCFunctionToolkit(warehouse_id="testId1").include("rag.test.*").get_tools()) if include_uc_function_tools else [] ) chat_model = ChatDatabricks(endpoint="databricks-llama-2-70b-chat", max_tokens=500) vectorstore = get_vector_search( endpoint_name="dbdemos_vs_endpoint", index_name="mlflow.rag.vs_index", has_embedding_endpoint=True, ) retriever = vectorstore.as_retriever() retriever_tool = create_retriever_tool(retriever, "vs_index_name", "vs_index_desc") agent = initialize_agent( uc_function_tools + [retriever_tool], chat_model, verbose=True, ) resources = list(_extract_dependency_list_from_lc_model(agent)) # Ensure all resources are added in expected = [ DatabricksVectorSearchIndex(index_name="mlflow.rag.vs_index"), DatabricksServingEndpoint(endpoint_name="embedding-model"), DatabricksServingEndpoint(endpoint_name="databricks-llama-2-70b-chat"), ] if include_uc_function_tools: expected = [ DatabricksServingEndpoint(endpoint_name="databricks-llama-2-70b-chat"), DatabricksFunction(function_name="rag.test.test_function"), DatabricksFunction(function_name="rag.test.test_function_2"), DatabricksFunction(function_name="rag.test.test_function_3"), DatabricksVectorSearchIndex(index_name="mlflow.rag.vs_index"), DatabricksServingEndpoint(endpoint_name="embedding-model"), DatabricksSQLWarehouse(warehouse_id="testId1"), ] def build_resource_map(resources): resource_map = defaultdict(list) for resource in resources: resource_type = resource.type.value resource_name = resource.to_dict()[resource_type][0]["name"] resource_map[resource_type].append(resource_name) return dict(resource_map) # Build maps for resources and expected resources resource_maps = build_resource_map(resources) expected_maps = build_resource_map(expected) assert len(resource_maps) == len(expected_maps) for resource_type in resource_maps: assert Counter(resource_maps[resource_type]) == Counter( expected_maps.get(resource_type, []) ) def test_parsing_dependency_from_databricks_chat(monkeypatch): from databricks_langchain import ChatDatabricks # in databricks-langchain > 0.7.0, ChatDatabricks instantiates # workspace client in __init__ which requires Databricks creds monkeypatch.setenv("DATABRICKS_HOST", "my-default-host") monkeypatch.setenv("DATABRICKS_TOKEN", "my-default-token") remove_langchain_community(monkeypatch) with pytest.raises(ImportError, match="No module named 'langchain_community"): from langchain_community.chat_models import ChatDatabricks chat_model = ChatDatabricks(endpoint="databricks-llama-2-70b-chat", max_tokens=500) resources = list(_extract_databricks_dependencies_from_chat_model(chat_model)) assert resources == [DatabricksServingEndpoint(endpoint_name="databricks-llama-2-70b-chat")] def test_parsing_dependency_from_databricks(monkeypatch): from databricks_langchain import ChatDatabricks # in databricks-langchain > 0.7.0, ChatDatabricks instantiates # workspace client in __init__ which requires Databricks creds monkeypatch.setenv("DATABRICKS_HOST", "my-default-host") monkeypatch.setenv("DATABRICKS_TOKEN", "my-default-token") remove_langchain_community(monkeypatch) with pytest.raises(ImportError, match="No module named 'langchain_community"): from langchain_community.chat_models import ChatDatabricks vectorstore = get_vector_search( endpoint_name="dbdemos_vs_endpoint", index_name="mlflow.rag.vs_index", has_embedding_endpoint=True, ) retriever = vectorstore.as_retriever() llm = ChatDatabricks(endpoint="databricks-llama-2-70b-chat", max_tokens=500) llm2 = ChatDatabricks(endpoint="databricks-llama-2-70b-chat", max_tokens=500) model = retriever | llm | llm2 resources = _detect_databricks_dependencies(model) assert resources == [ DatabricksVectorSearchIndex(index_name="mlflow.rag.vs_index"), DatabricksServingEndpoint(endpoint_name="embedding-model"), DatabricksServingEndpoint(endpoint_name="databricks-llama-2-70b-chat"), ] def test_parsing_unitycatalog_tool_as_dependency(monkeypatch: pytest.MonkeyPatch): from databricks.sdk.service.catalog import FunctionInfo from databricks_langchain import ChatDatabricks from langchain.agents import initialize_agent from unitycatalog.ai.core.databricks import DatabricksFunctionClient from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit # When get is called return a function def mock_function_get(self, function_name): components = function_name.split(".") # Initialize agent used below requires functions to take in exactly one parameter param_dict = { "parameters": [ { "name": "param", "parameter_type": "PARAM", "position": 0, "type_json": '{"name":"param","type":"string","nullable":true,"metadata":{}}', "type_name": "STRING", "type_precision": 0, "type_scale": 0, "type_text": "string", } ] } # Add the catalog, schema and name to the function Info followed by the parameter return FunctionInfo.from_dict({ "catalog_name": components[0], "schema_name": components[1], "name": components[2], "input_params": param_dict, }) monkeypatch.setenv("DATABRICKS_HOST", "my-default-host") monkeypatch.setenv("DATABRICKS_TOKEN", "my-default-token") monkeypatch.setattr("databricks.sdk.service.catalog.FunctionsAPI.get", mock_function_get) # TODO: remove this mock after unitycatalog-ai release a new version to avoid setting # spark session during initialization with mock.patch("unitycatalog.ai.core.databricks.DatabricksFunctionClient.set_spark_session"): client = DatabricksFunctionClient() toolkit = UCFunctionToolkit(function_names=["rag.test.test_function"], client=client) llm = ChatDatabricks(endpoint="databricks-llama-2-70b-chat", temperature=0) agent = initialize_agent( toolkit.tools, llm, verbose=True, ) resources = sorted(_extract_dependency_list_from_lc_model(agent), key=lambda x: x.name) assert resources == [ DatabricksServingEndpoint(endpoint_name="databricks-llama-2-70b-chat"), DatabricksFunction(function_name="rag.test.test_function"), ]