wehub-resource-sync
79 lines
2.4 KiB
Python
79 lines
2.4 KiB
Python
# Test integration with the `databricks-langchain` package.
|
|
from typing import Generator
|
|
from unittest import mock
|
|
|
|
import langchain
|
|
import pytest
|
|
from langchain_core.output_parsers import StrOutputParser
|
|
from langchain_core.prompts import PromptTemplate
|
|
from openai.types.chat.chat_completion import ChatCompletion
|
|
from packaging.version import Version
|
|
|
|
import mlflow
|
|
|
|
_MOCK_CHAT_RESPONSE = {
|
|
"id": "chatcmpl_id",
|
|
"object": "chat.completion",
|
|
"created": 1721875529,
|
|
"model": "meta-llama-3.1-70b-instruct-072424",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": "What is MLflow?",
|
|
},
|
|
"finish_reason": "stop",
|
|
"logprobs": None,
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 30, "completion_tokens": 36, "total_tokens": 66},
|
|
}
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def mock_client(monkeypatch) -> Generator:
|
|
# In databricks-langchain <= 0.7.0, ChatDatabricks uses MLflow deployment client
|
|
deploy_client = mock.MagicMock()
|
|
deploy_client.predict.return_value = _MOCK_CHAT_RESPONSE
|
|
# For newer version, ChatDatabricks uses workspace OpenAI client
|
|
openai_client = mock.MagicMock()
|
|
openai_client.chat.completions.create.return_value = ChatCompletion.model_validate(
|
|
_MOCK_CHAT_RESPONSE
|
|
)
|
|
|
|
with (
|
|
mock.patch("mlflow.deployments.get_deploy_client", return_value=deploy_client),
|
|
mock.patch(
|
|
"databricks_langchain.chat_models.get_openai_client", return_value=openai_client
|
|
),
|
|
):
|
|
yield
|
|
|
|
|
|
@pytest.fixture
|
|
def model_path(tmp_path):
|
|
return tmp_path / "model"
|
|
|
|
|
|
# TODO: Remove this once databricks-langchain supports v1
|
|
@pytest.mark.skipif(
|
|
Version(langchain.__version__).major >= 1,
|
|
reason="databricks-langchain does not support v1 yet",
|
|
)
|
|
def test_save_and_load_chat_databricks(model_path):
|
|
from databricks_langchain import ChatDatabricks
|
|
|
|
llm = ChatDatabricks(endpoint="databricks-meta-llama-3-70b-instruct")
|
|
prompt = PromptTemplate.from_template("What is {product}?")
|
|
chain = prompt | llm | StrOutputParser()
|
|
|
|
mlflow.langchain.save_model(chain, path=model_path)
|
|
|
|
loaded_model = mlflow.langchain.load_model(model_path)
|
|
assert loaded_model == chain
|
|
|
|
loaded_pyfunc_model = mlflow.pyfunc.load_model(model_path)
|
|
prediction = loaded_pyfunc_model.predict([{"product": "MLflow"}])
|
|
assert prediction == ["What is MLflow?"]
|