Format fixes

databricks · Oct 4, 2024 · 41a2ab7 · 41a2ab7
1 parent 16ee3f7
commit 41a2ab7
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/databricks/sdk/mixins/open_ai_client.py b/databricks/sdk/mixins/open_ai_client.py
@@ -1,6 +1,8 @@
 from databricks.sdk.service.serving import ServingEndpointsAPI
 
+
 class ServingEndpointsExt(ServingEndpointsAPI):
+
     def get_open_api_client(self):
         auth_headers = self._api._cfg.authenticate()
 
@@ -10,7 +12,4 @@ def get_open_api_client(self):
             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
-        )
+        return OpenAI(base_url=self._api._cfg.host + "/serving-endpoints", api_key=token)
diff --git a/tests/test_open_ai_mixin.py b/tests/test_open_ai_mixin.py
@@ -1,6 +1,6 @@
-import pytest
 from databricks.sdk.core import Config
 
+
 def test_open_ai_client(monkeypatch):
     from databricks.sdk import WorkspaceClient
 
@@ -10,4 +10,4 @@ def test_open_ai_client(monkeypatch):
     client = w.serving_endpoints.get_open_api_client()
 
     assert client.base_url == "https://test_host/serving-endpoints/"
-    assert client.api_key == "test_token"
+    assert client.api_key == "test_token"