Add open ai client mixin

Author: aravind-segu

.codegen/__init__.py.tmpl

@@ -5,6 +5,7 @@ from databricks.sdk.credentials_provider import CredentialsStrategy
 from databricks.sdk.mixins.files import DbfsExt
 from databricks.sdk.mixins.compute import ClustersExt
 from databricks.sdk.mixins.workspace import WorkspaceExt
+from databricks.sdk.mixins.open_ai_client import ServingEndpointsExt
 {{- range .Services}}
 from databricks.sdk.service.{{.Package.Name}} import {{.PascalName}}API{{end}}
 from databricks.sdk.service.provisioning import Workspace
@@ -17,7 +18,7 @@ from typing import Optional
   "google_credentials" "google_service_account" }}
 
 {{- define "api" -}}
-  {{- $mixins := dict "ClustersAPI" "ClustersExt" "DbfsAPI" "DbfsExt" "WorkspaceAPI" "WorkspaceExt" -}}
+  {{- $mixins := dict "ClustersAPI" "ClustersExt" "DbfsAPI" "DbfsExt" "WorkspaceAPI" "WorkspaceExt" "ServingEndpointsExt" "ServingEndpointsApi" -}}
   {{- $genApi := concat .PascalName "API" -}}
   {{- getOrDefault $mixins $genApi $genApi -}}
 {{- end -}}

databricks/sdk/__init__.py

@@ -0,0 +1,16 @@
+from databricks.sdk.service.serving import ServingEndpointsAPI
+
+class ServingEndpointsExt(ServingEndpointsAPI):
+    def get_open_api_client(self):
+        auth_headers = self._api._cfg.authenticate()
+
+        try:
+            token = auth_headers["Authorization"][len("Bearer "):]
+        except Exception:
+            raise ValueError("Unable to extract authorization token for OpenAI Client")
+
+        from openai import OpenAI
+        return OpenAI(
+            base_url=self._api._cfg.host + "/serving-endpoints", 
+            api_key=token
+        )

databricks/sdk/mixins/open_ai_client.py

@@ -0,0 +1,13 @@
+import pytest
+from databricks.sdk.core import Config
+
+def test_open_ai_client(monkeypatch):
+    from databricks.sdk import WorkspaceClient
+
+    monkeypatch.setenv('DATABRICKS_HOST', 'test_host')
+    monkeypatch.setenv('DATABRICKS_TOKEN', 'test_token')
+    w = WorkspaceClient(config=Config())
+    client = w.serving_endpoints.get_open_api_client()
+
+    assert client.base_url == "https://test_host/serving-endpoints/"
+    assert client.api_key == "test_token"