Skip to content

Commit

Permalink
[Internal] Update Model Serving http_request mixin to correctly use…
Browse files Browse the repository at this point in the history
… the underlying API. (#876)

## What changes are proposed in this pull request?

This PR updates the Model Serving `http_request` function so that it
properly uses the underlying generated API.

This PR also updates the API of a few unrelated services. 

## How is this tested?

Added unit tests from PR #874
  • Loading branch information
renaudhartert-db authored Jan 29, 2025
1 parent 762c57b commit 5339396
Show file tree
Hide file tree
Showing 31 changed files with 1,225 additions and 225 deletions.
2 changes: 1 addition & 1 deletion .codegen/_openapi_sha
Original file line number Diff line number Diff line change
@@ -1 +1 @@
58905570a9928fc9ed31fba14a2edaf9a7c55b08
840c660106f820a1a5dff931d51fa5f65cd9fdd9
13 changes: 10 additions & 3 deletions databricks/sdk/__init__.py

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

35 changes: 25 additions & 10 deletions databricks/sdk/mixins/open_ai_client.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
import json as js
from typing import Dict, Optional

from requests import Response

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


Expand Down Expand Up @@ -63,7 +64,7 @@ def http_request(self,
*,
headers: Optional[Dict[str, str]] = None,
json: Optional[Dict[str, str]] = None,
params: Optional[Dict[str, str]] = None) -> ExternalFunctionResponse:
params: Optional[Dict[str, str]] = None) -> Response:
"""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
Expand All @@ -79,13 +80,27 @@ def http_request(self,
JSON payload for the request.
:param params: Dict[str,str] (optional)
Query parameters for the request.
:returns: :class:`ExternalFunctionResponse`
:returns: :class:`Response`
"""
response = Response()
response.status_code = 200
server_response = super().http_request(connection_name=conn,
method=method,
path=path,
headers=js.dumps(headers) if headers is not None else None,
json=js.dumps(json) if json is not None else None,
params=js.dumps(params) if params is not None else None)

# Read the content from the HttpRequestResponse object
if hasattr(server_response, "contents") and hasattr(server_response.contents, "read"):
raw_content = server_response.contents.read() # Read the bytes
else:
raise ValueError("Invalid response from the server.")

# Set the raw content
if isinstance(raw_content, bytes):
response._content = raw_content
else:
raise ValueError("Contents must be bytes.")

return super.http_request(connection_name=conn,
method=method,
path=path,
headers=js.dumps(headers),
json=js.dumps(json),
params=js.dumps(params),
)
return response
Loading

0 comments on commit 5339396

Please sign in to comment.