From dc9d22861ac77703d5085264158b3e77fa570360 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Tue, 14 Jan 2025 13:20:02 -0800 Subject: [PATCH] Adding the missing super() for open_ai_client Signed-off-by: Sunish Sheth --- databricks/sdk/mixins/open_ai_client.py | 32 ++++++++++++++++++------- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/databricks/sdk/mixins/open_ai_client.py b/databricks/sdk/mixins/open_ai_client.py index 5f971311..88946a18 100644 --- a/databricks/sdk/mixins/open_ai_client.py +++ b/databricks/sdk/mixins/open_ai_client.py @@ -1,11 +1,23 @@ import json as js from typing import Dict, Optional +from dataclasses import dataclass from databricks.sdk.service.serving import (ExternalFunctionRequestHttpMethod, - ExternalFunctionResponse, + ExternalFunctionResponse as ExternalFunctionResponseAPI, ServingEndpointsAPI) +@dataclass +class ExternalFunctionResponse(ExternalFunctionResponseAPI): + text: Dict[str, any] = None + """The content of the response""" + + @classmethod + def from_dict(cls, d: Dict[str, any]) -> 'ExternalFunctionResponse': + """Deserializes the ExternalFunctionResponse from a dictionary.""" + return cls(status_code=200, text=d) + + class ServingEndpointsExt(ServingEndpointsAPI): # Using the HTTP Client to pass in the databricks authorization @@ -82,10 +94,14 @@ def http_request(self, :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), - ) + body = {} + if conn is not None: body['connection_name'] = conn + if headers is not None: body['headers'] = js.dumps(headers) + if json is not None: body['json'] = js.dumps(json) + if method is not None: body['method'] = method.value + if params is not None: body['params'] = js.dumps(params) + if path is not None: body['path'] = path + headers = {'Accept': 'application/json', 'Content-Type': 'application/json', } + + res = self._api.do('POST', '/api/2.0/external-function', body=body, headers=headers) + return ExternalFunctionResponse.from_dict(res)