mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-25 10:06:43 +08:00
fix: Lazy loading of policies deps
This commit is contained in:
Eric Hare
@ -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)
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
Reference in New Issue
Block a user