|
1 | 1 | 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 |
3 | 5 |
|
4 | 6 | from databricks.sdk.service.serving import (ExternalFunctionRequestHttpMethod, |
5 | 7 | ExternalFunctionResponse, |
|
8 | 10 |
|
9 | 11 | class ServingEndpointsExt(ServingEndpointsAPI): |
10 | 12 |
|
| 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 | + |
11 | 31 | # Using the HTTP Client to pass in the databricks authorization |
12 | 32 | # This method will be called on every invocation, so when using with model serving will always get the refreshed token |
13 | 33 | def _get_authorized_http_client(self): |
@@ -82,10 +102,15 @@ def http_request(self, |
82 | 102 | :returns: :class:`ExternalFunctionResponse` |
83 | 103 | """ |
84 | 104 |
|
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