Skip to content
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

[Feature] Introduce new Credential Strategies for Agents #882

Open
wants to merge 5 commits into
base: main
Choose a base branch
from

Conversation

aravind-segu
Copy link
Contributor

@aravind-segu aravind-segu commented Jan 31, 2025

What changes are proposed in this pull request?

This PR introduces two new credential strategies for Agents, (AgentEmbeddedCredentials, AgentUserCredentials).

Agents currently use the databricks.sdk in order to interact with databricks resources. However the authentication method for these resources is a little unique where we store the token for the authentication in a Credential File on the Kubernetes Container. Therefore in the past we added the Model Serving Credential Strategy to the defaultCredentials list to read this file.

Now we want to introduce a new authentication where the user's token is instead stored in a thread local variable. Agent users will initialize clients as follows:

from databricks.sdk.credentials_provider import ModelServingUserCredentials

invokers_client = WorkspaceClient(credential_strategy = ModelServingUserCredentials())
definers_client = WorkspaceClient()

Then the users can use the invoker_client to interact with resources with the invokers token or the definers_client to interact with resources using the old method of authentication.

Additionally as the users will be using these clients to test their code locally in Databricks Notebooks, if the code is not being run on model serving environments, users need to be able to authenticate using the DefaultCredential strategies.

More details: https://docs.google.com/document/d/14qLVjyxIAk581w287TWElstIeh8-DR30ab9Z6B_Vydg/edit?usp=sharing

How is this tested?

Added unit tests

@aravind-segu aravind-segu changed the title Introduce new Credential Strategies for Agents [Feature] Introduce new Credential Strategies for Agents Feb 6, 2025
@renaudhartert-db renaudhartert-db self-requested a review February 7, 2025 07:29
@renaudhartert-db
Copy link
Contributor

I've started the integration test for this change, result should appear soon.

Meanwhile, could you please run make fmt and commit the changes?

Copy link

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-py

Inputs:

  • PR number: 882
  • Commit SHA: 5fb5b1624e34bb2c1649b011bed5057136accf3b

Checks will be approved automatically on success.

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