Skip to content

Commit

Permalink
Adding the missing super() for open_ai_client
Browse files Browse the repository at this point in the history
Signed-off-by: Sunish Sheth <sunishsheth2009@gmail.com>
  • Loading branch information
sunishsheth2009 committed Jan 15, 2025
1 parent ee136e2 commit dc9d228
Showing 1 changed file with 24 additions and 8 deletions.
32 changes: 24 additions & 8 deletions databricks/sdk/mixins/open_ai_client.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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)

0 comments on commit dc9d228

Please sign in to comment.