Skip to content

Commit

Permalink
[Feature] Add serving.http_request to call external functions. (#857)
Browse files Browse the repository at this point in the history
## What changes are proposed in this pull request?

This PR adds the `serving.http_request` function to call
`/external-function`.

The goal of this function is to make it easy for the AI agent authors to
create tools that can make request to external services by invoking the
/external-function API where the secrets are stored in UC Connections.

This PR is based on PR #852.

## How is this tested?

The mixin itself was not tested. Further testing will be conducted by
the Model Serving team.
  • Loading branch information
renaudhartert-db authored Jan 14, 2025
1 parent 8c4264b commit ee136e2
Show file tree
Hide file tree
Showing 22 changed files with 766 additions and 322 deletions.
2 changes: 1 addition & 1 deletion .codegen/_openapi_sha
Original file line number Diff line number Diff line change
@@ -1 +1 @@
779817ed8d63031f5ea761fbd25ee84f38feec0d
05a10af4ed43566968119b43605f0a7fecbe780f
27 changes: 17 additions & 10 deletions databricks/sdk/credentials_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,6 +167,7 @@ def oauth_service_principal(cfg: 'Config') -> Optional[CredentialsProvider]:
oidc = cfg.oidc_endpoints
if oidc is None:
return None

token_source = ClientCredentials(client_id=cfg.client_id,
client_secret=cfg.client_secret,
token_url=oidc.token_endpoint,
Expand Down Expand Up @@ -210,16 +211,22 @@ def external_browser(cfg: 'Config') -> Optional[CredentialsProvider]:
credentials = token_cache.load()
if credentials:
# Force a refresh in case the loaded credentials are expired.
credentials.token()
else:
oauth_client = OAuthClient(oidc_endpoints=oidc_endpoints,
client_id=client_id,
redirect_url=redirect_url,
client_secret=client_secret)
consent = oauth_client.initiate_consent()
if not consent:
return None
credentials = consent.launch_external_browser()
# If the refresh fails, rather than throw exception we will initiate a new OAuth login flow.
try:
credentials.token()
return credentials(cfg)
# TODO: we should ideally use more specific exceptions.
except Exception as e:
logger.warning(f'Failed to refresh cached token: {e}. Initiating new OAuth login flow')

oauth_client = OAuthClient(oidc_endpoints=oidc_endpoints,
client_id=client_id,
redirect_url=redirect_url,
client_secret=client_secret)
consent = oauth_client.initiate_consent()
if not consent:
return None
credentials = consent.launch_external_browser()
token_cache.save(credentials)
return credentials(cfg)

Expand Down
41 changes: 40 additions & 1 deletion databricks/sdk/mixins/open_ai_client.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,9 @@
from databricks.sdk.service.serving import ServingEndpointsAPI
import json as js
from typing import Dict, Optional

from databricks.sdk.service.serving import (ExternalFunctionRequestHttpMethod,
ExternalFunctionResponse,
ServingEndpointsAPI)


class ServingEndpointsExt(ServingEndpointsAPI):
Expand Down Expand Up @@ -50,3 +55,37 @@ def get_langchain_chat_open_ai_client(self, model):
openai_api_base=self._api._cfg.host + "/serving-endpoints",
api_key="no-token", # Passing in a placeholder to pass validations, this will not be used
http_client=self._get_authorized_http_client())

def http_request(self,
conn: str,
method: ExternalFunctionRequestHttpMethod,
path: str,
*,
headers: Optional[Dict[str, str]] = None,
json: Optional[Dict[str, str]] = None,
params: Optional[Dict[str, str]] = None) -> ExternalFunctionResponse:
"""Make external services call using the credentials stored in UC Connection.
**NOTE:** Experimental: This API may change or be removed in a future release without warning.
:param conn: str
The connection name to use. This is required to identify the external connection.
:param method: :class:`ExternalFunctionRequestHttpMethod`
The HTTP method to use (e.g., 'GET', 'POST'). This is required.
:param path: str
The relative path for the API endpoint. This is required.
:param headers: Dict[str,str] (optional)
Additional headers for the request. If not provided, only auth headers from connections would be
passed.
:param json: Dict[str,str] (optional)
JSON payload for the request.
:param params: Dict[str,str] (optional)
Query parameters for the request.
:returns: :class:`ExternalFunctionResponse`
"""

return super.http_request(connection_name=conn,
method=method,
path=path,
headers=js.dumps(headers),
json=js.dumps(json),
params=js.dumps(params),
)
3 changes: 2 additions & 1 deletion databricks/sdk/service/cleanrooms.py

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

9 changes: 6 additions & 3 deletions databricks/sdk/service/files.py

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

85 changes: 84 additions & 1 deletion databricks/sdk/service/jobs.py

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

12 changes: 12 additions & 0 deletions databricks/sdk/service/oauth2.py

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading

0 comments on commit ee136e2

Please sign in to comment.