github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-09-09 06:53:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

community[patch]: Databricks - fix scope of dangerous deserialization error in Databricks LLM connector (#20368)

Browse Source 
fix scope of dangerous deserialization error in Databricks LLM connector

---------

Signed-off-by: dbczumar <corey.zumar@databricks.com>
This commit is contained in:
Corey Zumar
2024-04-12 14:27:26 -07:00
committed by GitHub
parent f1248f8d9a
commit 3a068b26f3
2 changed files with 63 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
libs/community/langchain_community/llms/databricks.py
View File

@@ -221,8 +221,21 @@ def _is_hex_string(data: str) -> bool:
return bool(re.match(pattern, data))
def _load_pickled_fn_from_hex_string(data: str) -> Callable:
def _load_pickled_fn_from_hex_string(
data: str, allow_dangerous_deserialization: Optional[bool]
) -> Callable:
"""Loads a pickled function from a hexadecimal string."""
if not allow_dangerous_deserialization:
raise ValueError(
"This code relies on the pickle module. "
"You will need to set allow_dangerous_deserialization=True "
"if you want to opt-in to allow deserialization of data using pickle."
"Data can be compromised by a malicious actor if "
"not handled properly to include "
"a malicious payload that when deserialized with "
"pickle can execute arbitrary code on your machine."
)
try:
import cloudpickle
except Exception as e:
@@ -443,25 +456,21 @@ class Databricks(LLM):
return v
def __init__(self, **data: Any):
if not data.get("allow_dangerous_deserialization"):
raise ValueError(
"This code relies on the pickle module. "
"You will need to set allow_dangerous_deserialization=True "
"if you want to opt-in to allow deserialization of data using pickle."
"Data can be compromised by a malicious actor if "
"not handled properly to include "
"a malicious payload that when deserialized with "
"pickle can execute arbitrary code on your machine."
)
if "transform_input_fn" in data and _is_hex_string(data["transform_input_fn"]):
data["transform_input_fn"] = _load_pickled_fn_from_hex_string(
data["transform_input_fn"]
data=data["transform_input_fn"],
allow_dangerous_deserialization=data.get(
"allow_dangerous_deserialization"
),
)
if "transform_output_fn" in data and _is_hex_string(
data["transform_output_fn"]
):
data["transform_output_fn"] = _load_pickled_fn_from_hex_string(
data["transform_output_fn"]
data=data["transform_output_fn"],
allow_dangerous_deserialization=data.get(
"allow_dangerous_deserialization"
),
)
super().__init__(**data)