[Feature] Add serving.http_request to call external functions. (#857)

## 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.

Author: renaudhartert-db

.codegen/_openapi_sha

@@ -1 +1 @@
-779817ed8d63031f5ea761fbd25ee84f38feec0d
+05a10af4ed43566968119b43605f0a7fecbe780f

databricks/sdk/credentials_provider.py

@@ -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,
@@ -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)
 

databricks/sdk/mixins/open_ai_client.py

@@ -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):
@@ -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),
+                                  )