diff --git a/src/backend/tests/unit/components/models_and_agents/policies/test_policies_component.py b/src/backend/tests/unit/components/models_and_agents/policies/test_policies_component.py index 9ae278a3f5..9e26d14605 100644 --- a/src/backend/tests/unit/components/models_and_agents/policies/test_policies_component.py +++ b/src/backend/tests/unit/components/models_and_agents/policies/test_policies_component.py @@ -36,40 +36,63 @@ def mock_component(mock_tool): return component +def _make_fake_tg(**overrides): + """Build a fake toolguard-imports dict for use with `_import_toolguard` patching. + + Mirrors the keys returned by `PoliciesComponent._import_toolguard()`; tests + can override individual entries to assert specific call wiring. + """ + fake = { + "PolicySpecOptions": MagicMock(), + "ToolGuardsCodeGenerationResult": MagicMock(), + "generate_guard_specs": MagicMock(), + "generate_guards_code": MagicMock(), + "langchain_tools_to_openapi": MagicMock(), + "load_toolguards": MagicMock(), + "load_toolguards_from_memory": MagicMock(), + "RESULTS_FILENAME": "results.json", + "sync_generated_guard_code_inputs": MagicMock(), + "GuardedTool": MagicMock(), + "LangchainModelWrapper": MagicMock(), + } + fake.update(overrides) + return fake + + @pytest.mark.asyncio async def test_cache_mode_success(mock_component, mock_tool): """Test PoliciesComponent in cache mode with valid cached guards.""" code_dir = mock_component.work_dir / STEP2 + fake_tg = _make_fake_tg() + mock_tg_result = MagicMock() + mock_tg_runtime = MagicMock() + fake_tg["load_toolguards_from_memory"].return_value = mock_tg_runtime + mock_guarded_instance = MagicMock() + fake_tg["GuardedTool"].return_value = mock_guarded_instance + # Mock the cache directory exists and toolguard loading with ( patch.object(Path, "exists", return_value=True), - patch("lfx.components.models_and_agents.policies_component.load_toolguards") as mock_load_guards, + patch.object(PoliciesComponent, "_import_toolguard", return_value=fake_tg), patch.object(mock_component, "make_toolguard_result") as mock_make_result, - patch("lfx.components.models_and_agents.policies_component.load_toolguards_from_memory") as mock_load_memory, - patch("lfx.components.models_and_agents.policies_component.GuardedTool") as mock_guarded_tool, ): - mock_tg_result = MagicMock() mock_make_result.return_value = mock_tg_result - mock_tg_runtime = MagicMock() - mock_load_memory.return_value = mock_tg_runtime - mock_guarded_instance = MagicMock() - mock_guarded_tool.return_value = mock_guarded_instance result = await mock_component.guard_tools() # Verify load_toolguards was called during validation - mock_load_guards.assert_called_once_with(code_dir) + fake_tg["load_toolguards"].assert_called_once_with(code_dir) # Verify make_toolguard_result was called mock_make_result.assert_called_once() # Verify load_toolguards_from_memory was called with the result - mock_load_memory.assert_called_once_with(mock_tg_result) + fake_tg["load_toolguards_from_memory"].assert_called_once_with(mock_tg_result) # Verify GuardedTool was created for each tool - assert mock_guarded_tool.call_count == len(mock_component.in_tools) - mock_guarded_tool.assert_called_with(mock_tool, mock_component.in_tools, mock_tg_runtime) + assert fake_tg["GuardedTool"].call_count == len(mock_component.in_tools) + fake_tg["GuardedTool"].assert_called_with(mock_tool, mock_component.in_tools, mock_tg_runtime) # Verify result contains guarded tools assert len(result) == 1 @@ -79,9 +102,11 @@ async def test_cache_mode_success(mock_component, mock_tool): @pytest.mark.asyncio async def test_cache_mode_directory_not_found(mock_component): """Test PoliciesComponent in cache mode when cache directory doesn't exist.""" + fake_tg = _make_fake_tg() # Mock the cache directory does not exist with ( patch.object(Path, "exists", return_value=False), + patch.object(PoliciesComponent, "_import_toolguard", return_value=fake_tg), pytest.raises(ValueError, match="Cache directory not found"), ): await mock_component.guard_tools() @@ -90,29 +115,29 @@ async def test_cache_mode_directory_not_found(mock_component): @pytest.mark.asyncio async def test_cache_mode_file_not_found(mock_component): """Test PoliciesComponent in cache mode when required files are missing.""" + fake_tg = _make_fake_tg() + fake_tg["load_toolguards"].side_effect = FileNotFoundError("Guard file not found") # Mock the cache directory exists but files are missing with ( patch.object(Path, "exists", return_value=True), - patch("lfx.components.models_and_agents.policies_component.load_toolguards") as mock_load_guards, + patch.object(PoliciesComponent, "_import_toolguard", return_value=fake_tg), + pytest.raises(ValueError, match="Required guard code files missing"), ): - mock_load_guards.side_effect = FileNotFoundError("Guard file not found") - - with pytest.raises(ValueError, match="Required guard code files missing"): - await mock_component.guard_tools() + await mock_component.guard_tools() @pytest.mark.asyncio async def test_cache_mode_corrupted_cache(mock_component): """Test PoliciesComponent in cache mode when cached code is corrupted.""" + fake_tg = _make_fake_tg() + fake_tg["load_toolguards"].side_effect = Exception("Invalid Python syntax") # Mock the cache directory exists but code is corrupted with ( patch.object(Path, "exists", return_value=True), - patch("lfx.components.models_and_agents.policies_component.load_toolguards") as mock_load_guards, + patch.object(PoliciesComponent, "_import_toolguard", return_value=fake_tg), + pytest.raises(ValueError, match="Failed to load guard code"), ): - mock_load_guards.side_effect = Exception("Invalid Python syntax") - - with pytest.raises(ValueError, match="Failed to load guard code"): - await mock_component.guard_tools() + await mock_component.guard_tools() # @pytest.mark.asyncio @@ -163,29 +188,31 @@ async def test_inenabled_returns_original_tools(mock_component, mock_tool): async def test_generate_mode_validation_errors(mock_component): """Test PoliciesComponent in generate mode with validation errors.""" mock_component.mode = MODE_GENERATE + fake_tg = _make_fake_tg() - # Test empty project - mock_component.project = "" - with pytest.raises(ValueError): # noqa: PT011 - await mock_component.guard_tools() + with patch.object(PoliciesComponent, "_import_toolguard", return_value=fake_tg): + # Test empty project + mock_component.project = "" + with pytest.raises(ValueError): # noqa: PT011 + await mock_component.guard_tools() - # Test empty policies - mock_component.project = "test_project" - mock_component.policies = [] - with pytest.raises(ValueError, match="policies cannot be empty"): - await mock_component.guard_tools() + # Test empty policies + mock_component.project = "test_project" + mock_component.policies = [] + with pytest.raises(ValueError, match="policies cannot be empty"): + await mock_component.guard_tools() - # Test empty tools - mock_component.policies = ["Policy 1"] - mock_component.in_tools = [] - with pytest.raises(ValueError, match="in_tools cannot be empty"): - await mock_component.guard_tools() + # Test empty tools + mock_component.policies = ["Policy 1"] + mock_component.in_tools = [] + with pytest.raises(ValueError, match="in_tools cannot be empty"): + await mock_component.guard_tools() - # Test missing model - mock_component.in_tools = [MagicMock()] - mock_component.model = None - with pytest.raises(ValueError, match="model or api_key cannot be empty"): - await mock_component.guard_tools() + # Test missing model + mock_component.in_tools = [MagicMock()] + mock_component.model = None + with pytest.raises(ValueError, match="model or api_key cannot be empty"): + await mock_component.guard_tools() # # Test non-recommended model # mock_component.model = [{"name": "gpt-3.5-turbo", "provider": "OpenAI"}] @@ -250,7 +277,11 @@ async def test_verify_cached_guards_error_messages(mock_component): code_dir = mock_component.work_dir / STEP2 # Test directory not found error message - with patch.object(Path, "exists", return_value=False): + fake_tg = _make_fake_tg() + with ( + patch.object(Path, "exists", return_value=False), + patch.object(PoliciesComponent, "_import_toolguard", return_value=fake_tg), + ): with pytest.raises(ValueError, match="Cache directory not found") as exc_info: mock_component._verify_cached_guards(code_dir) @@ -258,24 +289,24 @@ async def test_verify_cached_guards_error_messages(mock_component): assert str(code_dir) in str(exc_info.value) # Test file not found error message + fake_tg = _make_fake_tg() + fake_tg["load_toolguards"].side_effect = FileNotFoundError("Missing file") with ( patch.object(Path, "exists", return_value=True), - patch("lfx.components.models_and_agents.policies_component.load_toolguards") as mock_load, + patch.object(PoliciesComponent, "_import_toolguard", return_value=fake_tg), ): - mock_load.side_effect = FileNotFoundError("Missing file") - with pytest.raises(ValueError, match="Required guard code files missing") as exc_info: mock_component._verify_cached_guards(code_dir) assert "Generate" in str(exc_info.value) # Test general error message + fake_tg = _make_fake_tg() + fake_tg["load_toolguards"].side_effect = RuntimeError("Unexpected error") with ( patch.object(Path, "exists", return_value=True), - patch("lfx.components.models_and_agents.policies_component.load_toolguards") as mock_load, + patch.object(PoliciesComponent, "_import_toolguard", return_value=fake_tg), ): - mock_load.side_effect = RuntimeError("Unexpected error") - with pytest.raises(ValueError, match="Failed to load guard code") as exc_info: mock_component._verify_cached_guards(code_dir) diff --git a/src/lfx/src/lfx/components/models_and_agents/policies_component.py b/src/lfx/src/lfx/components/models_and_agents/policies_component.py index a35ef4ea91..2b9febbc7f 100644 --- a/src/lfx/src/lfx/components/models_and_agents/policies_component.py +++ b/src/lfx/src/lfx/components/models_and_agents/policies_component.py @@ -1,29 +1,17 @@ +from __future__ import annotations + import os import re import shutil from pathlib import Path from typing import TYPE_CHECKING, cast -from toolguard.buildtime import ( - PolicySpecOptions, - ToolGuardsCodeGenerationResult, - ToolGuardSpec, - generate_guard_specs, - generate_guards_code, -) -from toolguard.extra.langchain_to_oas import langchain_tools_to_openapi -from toolguard.runtime import load_toolguards, load_toolguards_from_memory -from toolguard.runtime.runtime import RESULTS_FILENAME - from lfx.base.models import LCModelComponent from lfx.base.models.unified_models import ( get_language_model_options, get_llm, update_model_options_in_build_config, ) -from lfx.components.models_and_agents.policies.guard_sync_utils import sync_generated_guard_code_inputs -from lfx.components.models_and_agents.policies.guarded_tool import GuardedTool -from lfx.components.models_and_agents.policies.llm_wrapper import LangchainModelWrapper from lfx.components.models_and_agents.policies.module_utils import unload_module from lfx.field_typing import LanguageModel, Tool from lfx.io import ( @@ -39,6 +27,8 @@ from lfx.io import ( from lfx.log.logger import logger if TYPE_CHECKING: + from toolguard.buildtime import ToolGuardsCodeGenerationResult, ToolGuardSpec + from lfx.inputs.inputs import InputTypes @@ -50,6 +40,11 @@ MODE_GENERATE = "🛠️ Generate" MODE_GUARD = "🛡️ Guard" GENERATED_GUARD_INFO_PREFIX = "Auto-generated ToolGuard code for " +_TOOLGUARD_INSTALL_HINT = ( + "The 'toolguard' package is required to use PoliciesComponent. " + "Install the optional extra: `pip install 'langflow-base[toolguard]'`." +) + class PoliciesComponent(LCModelComponent): """Component for building tool protection code from textual business policies and instructions. @@ -57,6 +52,10 @@ class PoliciesComponent(LCModelComponent): This component uses ToolGuard to generate and apply policy-based guards to tools, ensuring that tool execution complies with defined business policies. Powered by ALTK ToolGuard (https://github.com/AgentToolkit/toolguard). + + `toolguard` is an optional extra (`langflow-base[toolguard]`); imports happen + lazily inside methods so this component can be discovered and inspected even + when the extra isn't installed. """ display_name = "Policies" @@ -142,6 +141,44 @@ Powered by [ALTK ToolGuard](https://github.com/AgentToolkit/toolguard )""" ), ] + @staticmethod + def _import_toolguard(): + """Lazily import `toolguard` and the sibling helpers that depend on it. + + Defined as a static method so it survives custom-component re-execution + via `create_class`, which only re-executes the class body, not arbitrary + module-level statements such as `try/except` import guards. + """ + try: + from toolguard.buildtime import ( + PolicySpecOptions, + ToolGuardsCodeGenerationResult, + generate_guard_specs, + generate_guards_code, + ) + from toolguard.extra.langchain_to_oas import langchain_tools_to_openapi + from toolguard.runtime import load_toolguards, load_toolguards_from_memory + from toolguard.runtime.runtime import RESULTS_FILENAME + + from lfx.components.models_and_agents.policies.guard_sync_utils import sync_generated_guard_code_inputs + from lfx.components.models_and_agents.policies.guarded_tool import GuardedTool + from lfx.components.models_and_agents.policies.llm_wrapper import LangchainModelWrapper + except ModuleNotFoundError as e: + raise ImportError(_TOOLGUARD_INSTALL_HINT) from e + return { + "PolicySpecOptions": PolicySpecOptions, + "ToolGuardsCodeGenerationResult": ToolGuardsCodeGenerationResult, + "generate_guard_specs": generate_guard_specs, + "generate_guards_code": generate_guards_code, + "langchain_tools_to_openapi": langchain_tools_to_openapi, + "load_toolguards": load_toolguards, + "load_toolguards_from_memory": load_toolguards_from_memory, + "RESULTS_FILENAME": RESULTS_FILENAME, + "sync_generated_guard_code_inputs": sync_generated_guard_code_inputs, + "GuardedTool": GuardedTool, + "LangchainModelWrapper": LangchainModelWrapper, + } + @property def work_dir(self) -> Path: return TOOLGUARD_WORK_DIR / self._to_snake_case(self.project) @@ -168,8 +205,9 @@ Powered by [ALTK ToolGuard](https://github.com/AgentToolkit/toolguard )""" field_name=field_name, field_value=field_value, ) + tg = self._import_toolguard() py_module = self._to_snake_case(self.project) - return sync_generated_guard_code_inputs( + return tg["sync_generated_guard_code_inputs"]( build_config=updated_build_config, work_dir=self.work_dir, step2_subdir=STEP2, @@ -177,32 +215,34 @@ Powered by [ALTK ToolGuard](https://github.com/AgentToolkit/toolguard )""" ) async def _generate_guard_specs(self) -> list[ToolGuardSpec]: + tg = self._import_toolguard() logger.debug("Starting step 1") logger.debug(f"model = {self.model}") - llm = LangchainModelWrapper(self.build_model()) + llm = tg["LangchainModelWrapper"](self.build_model()) out_dir = self.work_dir / STEP1 if out_dir.exists(): shutil.rmtree(out_dir) policy_text = "\n * ".join(self.policies) - open_api = langchain_tools_to_openapi(self.in_tools) + open_api = tg["langchain_tools_to_openapi"](self.in_tools) - options = PolicySpecOptions(example_number=4) - specs = await generate_guard_specs( + options = tg["PolicySpecOptions"](example_number=4) + specs = await tg["generate_guard_specs"]( policy_text=policy_text, tools=open_api, llm=llm, work_dir=out_dir, options=options ) logger.debug("Step 1 Done") return specs async def _generate_guard_code(self, specs: list[ToolGuardSpec]) -> ToolGuardsCodeGenerationResult: + tg = self._import_toolguard() logger.debug("Starting step 2") out_dir = self.work_dir / STEP2 if out_dir.exists(): shutil.rmtree(out_dir) - llm = LangchainModelWrapper(self.build_model()) + llm = tg["LangchainModelWrapper"](self.build_model()) app_name = self._to_snake_case(self.project) - open_api = langchain_tools_to_openapi(self.in_tools) + open_api = tg["langchain_tools_to_openapi"](self.in_tools) - gen_result = await generate_guards_code( + gen_result = await tg["generate_guards_code"]( tools=open_api, tool_specs=specs, work_dir=out_dir, llm=llm, app_name=app_name ) logger.debug("Step 2 Done") @@ -242,6 +282,7 @@ Powered by [ALTK ToolGuard](https://github.com/AgentToolkit/toolguard )""" unload_module(res.domain.app_name) def _verify_cached_guards(self, code_dir: Path) -> None: + tg = self._import_toolguard() # Validate cache exists before attempting to load if not code_dir.exists(): msg = ( @@ -252,7 +293,7 @@ Powered by [ALTK ToolGuard](https://github.com/AgentToolkit/toolguard )""" raise ValueError(msg) try: - load_toolguards(code_dir) + tg["load_toolguards"](code_dir) except FileNotFoundError as exc: msg = ( f"Policies: Required guard code files missing in '{code_dir}'. " @@ -276,12 +317,13 @@ Powered by [ALTK ToolGuard](https://github.com/AgentToolkit/toolguard )""" self._verify_cached_guards(code_dir) def make_toolguard_result(self) -> ToolGuardsCodeGenerationResult: + tg = self._import_toolguard() attrs = self.get_vertex().data["node"]["template"] if not attrs: raise ValueError - result_str = attrs[str(RESULTS_FILENAME)]["value"] - result = ToolGuardsCodeGenerationResult.model_validate_json(result_str) + result_str = attrs[str(tg["RESULTS_FILENAME"])]["value"] + result = tg["ToolGuardsCodeGenerationResult"].model_validate_json(result_str) result.domain.app_types.content = attrs.get(str(result.domain.app_types.file_name))["value"] result.domain.app_api.content = attrs.get(str(result.domain.app_api.file_name))["value"] @@ -296,6 +338,7 @@ Powered by [ALTK ToolGuard](https://github.com/AgentToolkit/toolguard )""" async def guard_tools(self) -> list[Tool]: if self.enabled: + tg = self._import_toolguard() mode = getattr(self, "mode", MODE_GENERATE) if mode == MODE_GENERATE: self.log(f"Start generating guard code at {self.work_dir}", name="info") @@ -310,8 +353,8 @@ Powered by [ALTK ToolGuard](https://github.com/AgentToolkit/toolguard )""" self._validate_before_using_cache(code_dir) try: tg_result = self.make_toolguard_result() - tg_runtime = load_toolguards_from_memory(tg_result) - guarded_tools = [GuardedTool(tool, self.in_tools, tg_runtime) for tool in self.in_tools] + tg_runtime = tg["load_toolguards_from_memory"](tg_result) + guarded_tools = [tg["GuardedTool"](tool, self.in_tools, tg_runtime) for tool in self.in_tools] return cast("list[Tool]", guarded_tools) except Exception as e: logger.exception(e)