Skip to content

Commit db98e99

Browse files
[Fix] Fixing http_request open_ai_client mixin
1 parent 4bcfb0a commit db98e99

File tree

2 files changed

+59
-50
lines changed

2 files changed

+59
-50
lines changed

databricks/sdk/mixins/open_ai_client.py

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
import json as js
22
from typing import Dict, Optional
33

4+
from requests import Response
5+
46
from databricks.sdk.service.serving import (ExternalFunctionRequestHttpMethod,
5-
ExternalFunctionResponse,
67
ServingEndpointsAPI)
78

89

@@ -63,7 +64,7 @@ def http_request(self,
6364
*,
6465
headers: Optional[Dict[str, str]] = None,
6566
json: Optional[Dict[str, str]] = None,
66-
params: Optional[Dict[str, str]] = None) -> ExternalFunctionResponse:
67+
params: Optional[Dict[str, str]] = None) -> Response:
6768
"""Make external services call using the credentials stored in UC Connection.
6869
**NOTE:** Experimental: This API may change or be removed in a future release without warning.
6970
:param conn: str
@@ -79,13 +80,27 @@ def http_request(self,
7980
JSON payload for the request.
8081
:param params: Dict[str,str] (optional)
8182
Query parameters for the request.
82-
:returns: :class:`ExternalFunctionResponse`
83+
:returns: :class:`Response`
8384
"""
85+
response = Response()
86+
response.status_code = 200
87+
server_response = super().http_request(connection_name=conn,
88+
method=method,
89+
path=path,
90+
headers=js.dumps(headers) if headers is not None else None,
91+
json=js.dumps(json) if json is not None else None,
92+
params=js.dumps(params) if params is not None else None)
93+
94+
# Read the content from the HttpRequestResponse object
95+
if hasattr(server_response, "contents") and hasattr(server_response.contents, "read"):
96+
raw_content = server_response.contents.read() # Read the bytes
97+
else:
98+
raise ValueError("Invalid response from the server.")
99+
100+
# Set the raw content
101+
if isinstance(raw_content, bytes):
102+
response._content = raw_content
103+
else:
104+
raise ValueError("Contents must be bytes.")
84105

85-
return super.http_request(connection_name=conn,
86-
method=method,
87-
path=path,
88-
headers=js.dumps(headers),
89-
json=js.dumps(json),
90-
params=js.dumps(params),
91-
)
106+
return response

databricks/sdk/service/serving.py

Lines changed: 34 additions & 40 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)