Skip to content

Conversation

jayshrivastava0
Copy link

This commit addresses issue #128 by modifying the __init__ method of the ChatDatabricks class to allow a custom client to be passed in via the constructor.

Previously, any client passed in the constructor would be overwritten by the default client. With this change:

  • If a client is provided in kwargs, that client instance is used.
  • Otherwise, the client is initialized using get_deployment_client(self.target_uri).

A new unit test, test_chat_model_with_custom_client, has been added to integrations/langchain/tests/unit_tests/test_chat_models.py to verify this functionality. The test ensures that the provided custom client is correctly assigned and its methods are called when the ChatDatabricks instance makes predictions.

google-labs-jules bot and others added 2 commits June 14, 2025 02:06
This commit addresses issue databricks#128 by modifying the `__init__` method of the `ChatDatabricks` class to allow a custom client to be passed in via the constructor.

Previously, any client passed in the constructor would be overwritten by the default client. With this change:
- If a `client` is provided in `kwargs`, that client instance is used.
- Otherwise, the client is initialized using `get_deployment_client(self.target_uri)`.

A new unit test, `test_chat_model_with_custom_client`, has been added to `integrations/langchain/tests/unit_tests/test_chat_models.py` to verify this functionality. The test ensures that the provided custom client is correctly assigned and its methods are called when the `ChatDatabricks` instance makes predictions.
…ks-128

Fix: Allow custom client to be passed to ChatDatabricks

This commit addresses issue databricks#128 by modifying the `__init__` method of the `ChatDatabricks` class to allow a custom client to be passed in via the constructor.

Previously, any client passed in the constructor would be overwritten by the default client. With this change:
- If a `client` is provided in `kwargs`, that client instance is used.
- Otherwise, the client is initialized using `get_deployment_client(self.target_uri)`.

A new unit test, `test_chat_model_with_custom_client`, has been added to `integrations/langchain/tests/unit_tests/test_chat_models.py` to verify this functionality. The test ensures that the provided custom client is correctly assigned and its methods are called when the `ChatDatabricks` instance makes predictions.
@aravind-segu
Copy link
Contributor

@jayshrivastava0 what is the usecase for passing in the client directly? What kind of clients are you passing in here?

@jayshrivastava0
Copy link
Author

Hey @aravind-segu, great question. We enabled this to give developers more flexibility.

The two main use cases are:

  1. Testing: It allows you to inject a mock client for unit tests, so you can verify your logic without making actual network calls.
  2. Custom Configuration: You can pass in a client that's already configured with specific enterprise needs, like custom authentication, proxy settings, or advanced retry logic.

The client just needs to be an object that behaves like MLflow's BaseDeploymentClient.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants