-
Notifications
You must be signed in to change notification settings - Fork 1.9k
Updated the DatabricksRM class to use Databricks service principals with REST API path #8327
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
- Updated the base class to include the optional parameters - Updated logic on the _query_via_databricks_sdk to use the SP credentials if they exist otherwise will fallback to PAT.
…ity' into dbx_service_principal_functionality # Conflicts: # dspy/retrieve/databricks_rm.py
Updated print statements for which auth method is invoked. Updated docstrings. Created helper function: _get_oauth_token Updated logic for databricks token to use oauth token with SP when using the REST API with client secret and client id
…ity' into dbx_service_principal_functionality # Conflicts: # dspy/retrieve/databricks_rm.py
Updated print statements for which auth method is invoked. Updated docstrings. Created helper function: _get_oauth_token Updated logic for databricks token to use oauth token with SP when using the REST API with client secret and client id
Updated print statements for which auth method is invoked. Updated docstrings. Created helper function: _get_oauth_token Updated logic for databricks token to use oauth token with SP when using the REST API with client secret and client id
Updated print statements for which auth method is invoked. Updated docstrings. Created helper function: _get_oauth_token Updated logic for databricks token to use oauth token with SP when using the REST API with client secret and client id
…ity' into dbx_service_principal_functionality
Updated print statements for which auth method is invoked. Updated docstrings. Created helper function: _get_oauth_token Updated logic for databricks token to use oauth token with SP when using the REST API with client secret and client id
…ity' into dbx_service_principal_functionality
…ality Updated the DatabricksRM class to use Databricks service principals.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the PR!
I feel this is going too complex with supporting all Databricks auth methods. We should move this implementation into dastabricks-ai-bridge IMO. I will discuss with the team and make a decision.
Pardon me for the extra PR. This was supposed to be included with my other PR to add the same SP support for both methods that were already present (SDK and REST API). It's the same mechanism, but inadvertently developed for the SDK before the API when they should've been added together if that makes sense. Thank you for your thoughts on this! |
This PR is to add service principal support with the Databricks REST API path for the databricks retriever in DSPy. This is mainly the path for work outside of Databricks environments (i.e VS Code with Databricks Connect) where the Databricks SDK is not automatically installed or if the REST API is preferred.
Note: I synced the fork with the main branch which added a number of commits.
Screenshots of REST API with SP tests: