diff --git a/code_review_graph/embeddings.py b/code_review_graph/embeddings.py index be29045..485b989 100644 --- a/code_review_graph/embeddings.py +++ b/code_review_graph/embeddings.py @@ -66,10 +66,12 @@ def _get_model(self): if self._model is None: try: from sentence_transformers import SentenceTransformer + # Check environment variable, default to False to prevent RCE + allow_remote_code = os.environ.get("CRG_ALLOW_REMOTE_CODE", "0").lower() in ("1", "true", "yes") + self._model = SentenceTransformer( self._model_name, - trust_remote_code=True, - model_kwargs={"trust_remote_code": True}, + trust_remote_code=allow_remote_code, ) except ImportError: raise ImportError( diff --git a/uv.lock b/uv.lock index 62a32ad..c40535b 100644 --- a/uv.lock +++ b/uv.lock @@ -411,6 +411,7 @@ requires-dist = [ { name = "pyyaml", marker = "extra == 'eval'", specifier = ">=6.0" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.3.0,<1" }, { name = "sentence-transformers", marker = "extra == 'embeddings'", specifier = ">=3.0.0,<4" }, + { name = "tomli", marker = "python_full_version < '3.11'", specifier = ">=2.0.0,<3" }, { name = "tomli", marker = "python_full_version < '3.11' and extra == 'dev'", specifier = ">=2.0" }, { name = "tree-sitter", specifier = ">=0.23.0,<1" }, { name = "tree-sitter-language-pack", specifier = ">=0.3.0,<1" },