Add Langchain Open AI Client

Author: aravind-segu

databricks/sdk/mixins/open_ai_client.py

@@ -3,7 +3,7 @@
 
 class ServingEndpointsExt(ServingEndpointsAPI):
 
-    def get_open_api_client(self):
+    def get_open_ai_client(self):
         auth_headers = self._api._cfg.authenticate()
 
         try:
@@ -13,3 +13,16 @@ def get_open_api_client(self):
 
         from openai import OpenAI
         return OpenAI(base_url=self._api._cfg.host + "/serving-endpoints", api_key=token)
+
+    def get_langchain_chat_open_ai_client(self, model):
+        auth_headers = self._api._cfg.authenticate()
+
+        try:
+            token = auth_headers["Authorization"][len("Bearer "):]
+        except Exception:
+            raise ValueError("Unable to extract authorization token for Langchain OpenAI Client")
+
+        from langchain_openai import ChatOpenAI
+        return ChatOpenAI(model=model,
+                          openai_api_base=self._api._cfg.host + "/serving-endpoints",
+                          openai_api_key=token)

setup.py

@@ -17,9 +17,10 @@
       extras_require={"dev": ["pytest", "pytest-cov", "pytest-xdist", "pytest-mock",
                               "yapf", "pycodestyle", "autoflake", "isort", "wheel",
                               "ipython", "ipywidgets", "requests-mock", "pyfakefs",
-                              "databricks-connect", "pytest-rerunfailures", "openai"],
+                              "databricks-connect", "pytest-rerunfailures", "openai", 
+                              "langchain-openai"],
                       "notebook": ["ipython>=8,<9", "ipywidgets>=8,<9"],
-                      "openai": ["openai"]},
+                      "openai": ["openai", "langchain-openai"]},
       author="Serge Smertin",
       author_email="[email protected]",
       description="Databricks SDK for Python (Beta)",

tests/test_open_ai_mixin.py

@@ -7,7 +7,19 @@ def test_open_ai_client(monkeypatch):
     monkeypatch.setenv('DATABRICKS_HOST', 'test_host')
     monkeypatch.setenv('DATABRICKS_TOKEN', 'test_token')
     w = WorkspaceClient(config=Config())
-    client = w.serving_endpoints.get_open_api_client()
+    client = w.serving_endpoints.get_open_ai_client()
 
     assert client.base_url == "https://test_host/serving-endpoints/"
     assert client.api_key == "test_token"
+
+
+def test_langchain_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_langchain_chat_open_ai_client("databricks-meta-llama-3-1-70b-instruct")
+
+    assert client.openai_api_base == "https://test_host/serving-endpoints"
+    assert client.model_name == "databricks-meta-llama-3-1-70b-instruct"