Skip to content

Commit 77e838e

Browse files
[Fix] Fixing http_request open_ai_client mixin
1 parent 5576d32 commit 77e838e

File tree

1 file changed

+33
-8
lines changed

1 file changed

+33
-8
lines changed

databricks/sdk/mixins/open_ai_client.py

Lines changed: 33 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import json as js
2-
from typing import Dict, Optional
2+
import databricks
3+
from typing import Dict, Optional, Any
4+
from dataclasses import dataclass, asdict
35

46
from databricks.sdk.service.serving import (ExternalFunctionRequestHttpMethod,
57
ExternalFunctionResponse,
@@ -8,6 +10,24 @@
810

911
class ServingEndpointsExt(ServingEndpointsAPI):
1012

13+
@dataclass
14+
class ExternalFunctionResponseOverride(ExternalFunctionResponse):
15+
text: Dict[str, Any] = None
16+
"""The content of the response"""
17+
18+
@classmethod
19+
def from_dict(cls, d: Dict[str, Any]) -> "ExternalFunctionResponse":
20+
"""Deserializes the ExternalFunctionResponse from a dictionary."""
21+
return cls(status_code=200, text=d)
22+
23+
def to_dict(self) -> Dict[str, Any]:
24+
"""Serializes the object back into a dictionary."""
25+
result = asdict(self) # Use dataclasses.asdict to serialize fields
26+
# Ensure the text field is serialized correctly
27+
if self.text is not None and not isinstance(self.text, str):
28+
result["text"] = js.dumps(self.text) # Serialize text as JSON if it's a dict
29+
return result
30+
1131
# Using the HTTP Client to pass in the databricks authorization
1232
# This method will be called on every invocation, so when using with model serving will always get the refreshed token
1333
def _get_authorized_http_client(self):
@@ -82,10 +102,15 @@ def http_request(self,
82102
:returns: :class:`ExternalFunctionResponse`
83103
"""
84104

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-
)
105+
databricks.sdk.service.serving.ExternalFunctionResponse = (
106+
ServingEndpointsExt.ExternalFunctionResponseOverride)
107+
108+
response = super().http_request(connection_name=conn,
109+
method=method,
110+
path=path,
111+
headers=js.dumps(headers) if headers is not None else None,
112+
json=js.dumps(json) if json is not None else None,
113+
params=js.dumps(params) if params is not None else None)
114+
115+
# Convert the overridden response back to the original response type
116+
return ExternalFunctionResponse.from_dict(response.to_dict())

0 commit comments

Comments
 (0)