From 7e2f79b604ff323ec0094b32c27214b7f9b10b62 Mon Sep 17 00:00:00 2001
From: Alex Sherstinsky <alexsherstinsky@users.noreply.github.com>
Date: Wed, 24 Apr 2024 13:13:24 -0700
Subject: [PATCH] Support Predibase-LangChain integration for both Predibase
 SDK-v1 and Predibase SDK-v2 users.

---
 .../langchain_community/llms/predibase.py     | 21 ++++++++++++-------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/libs/community/langchain_community/llms/predibase.py b/libs/community/langchain_community/llms/predibase.py
index b0ee1c4270895..c359c3010430a 100644
--- a/libs/community/langchain_community/llms/predibase.py
+++ b/libs/community/langchain_community/llms/predibase.py
@@ -21,7 +21,7 @@ class Predibase(LLM):
     otherwise, an error is raised.  If the fine-tuned adapter is hosted at Predibase,
     then `adapter_version` in the adapter repository must be specified.
 
-    An optional `predibase_sdk_version` parameter defaults to the latest Predibase SDK version.
+    An optional `predibase_sdk_version` parameter defaults to latest SDK version.
     """
 
     model: str
@@ -87,8 +87,9 @@ def _call(
             result: GeneratedResponse
             if self.adapter_id:
                 """
-                Attempt to retrieve the fine-tuned adapter from a Predibase repository.
-                If absent, then load the fine-tuned adapter from a HuggingFace repository.
+                Attempt to retrieve the fine-tuned adapter from a Predibase
+                repository.  If absent, then load the fine-tuned adapter
+                from a HuggingFace repository.
                 """
                 adapter_model: Union[Model, HuggingFaceLLM]
                 try:
@@ -144,7 +145,8 @@ def _call(
                     )
                 except GenerationError as ge:
                     raise ValueError(
-                        f'An adapter with the ID "{pb_adapter_id}" cannot be found in the Predibase repository of fine-tuned adapters.'
+                        f"""An adapter with the ID "{pb_adapter_id}" cannot be \
+found in the Predibase repository of fine-tuned adapters."""
                     ) from ge
             else:
                 # The adapter version is omitted, hence look for the adapter ID in the HuggingFace repository.
@@ -157,8 +159,9 @@ def _call(
                     )
                 except GenerationError as ge:
                     raise ValueError(
-                        f"""Either an adapter with the ID "{self.adapter_id}" cannot be found in a HuggingFace repository, \
-or it is incompatible with the base model (please make sure that the adapter configuration is consistent).
+                        f"""Either an adapter with the ID "{self.adapter_id}" \
+cannot be found in a HuggingFace repository, or it is incompatible with the \
+base model (please make sure that the adapter configuration is consistent).
 """
                     ) from ge
         else:
@@ -169,8 +172,10 @@ def _call(
                 )
             except requests.JSONDecodeError as jde:
                 raise ValueError(
-                    f"""An LLM with the deployment ID "{self.model}" cannot be found at Predibase \
-(please refer to "https://docs.predibase.com/user-guide/inference/models" for the list of supported models).
+                    f"""An LLM with the deployment ID "{self.model}" cannot be found \
+at Predibase (please refer to \
+"https://docs.predibase.com/user-guide/inference/models" for the list of \
+supported models).
 """
                 ) from jde
         response_text = response.generated_text