mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: add powerfx safe mode (#3028)
* add powerfx safe mode * improved docstring and aligned env_file loading * ensured test uses reset
This commit is contained in:
committed by
GitHub
Unverified
parent
5ab47596ff
commit
4b8a545589
@@ -37,6 +37,7 @@ from ._models import (
|
||||
RemoteConnection,
|
||||
Tool,
|
||||
WebSearchTool,
|
||||
_safe_mode_context,
|
||||
agent_schema_dispatch,
|
||||
)
|
||||
|
||||
@@ -118,7 +119,9 @@ class AgentFactory:
|
||||
client_kwargs: Mapping[str, Any] | None = None,
|
||||
additional_mappings: Mapping[str, ProviderTypeMapping] | None = None,
|
||||
default_provider: str = "AzureAIClient",
|
||||
env_file: str | None = None,
|
||||
safe_mode: bool = True,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
) -> None:
|
||||
"""Create the agent factory, with bindings.
|
||||
|
||||
@@ -151,7 +154,15 @@ class AgentFactory:
|
||||
that accepts the model.id value.
|
||||
default_provider: The default provider used when model.provider is not specified,
|
||||
default is "AzureAIClient".
|
||||
env_file: An optional path to a .env file to load environment variables from.
|
||||
safe_mode: Whether to run in safe mode, default is True.
|
||||
When safe_mode is True, environment variables are not accessible in the powerfx expressions.
|
||||
You can still use environment variables, but through the constructors of the classes.
|
||||
Which means you must make sure you are using the standard env variable names of the classes
|
||||
you are using and not custom ones and remove the powerfx statements that start with `=Env.`.
|
||||
Only when you trust the source of your yaml files, you can set safe_mode to False
|
||||
via the AgentFactory constructor.
|
||||
env_file_path: The path to the .env file to load environment variables from.
|
||||
env_file_encoding: The encoding of the .env file, defaults to 'utf-8'.
|
||||
"""
|
||||
self.chat_client = chat_client
|
||||
self.bindings = bindings
|
||||
@@ -159,7 +170,8 @@ class AgentFactory:
|
||||
self.client_kwargs = client_kwargs or {}
|
||||
self.additional_mappings = additional_mappings or {}
|
||||
self.default_provider: str = default_provider
|
||||
load_dotenv(dotenv_path=env_file)
|
||||
self.safe_mode = safe_mode
|
||||
load_dotenv(dotenv_path=env_file_path, encoding=env_file_encoding)
|
||||
|
||||
def create_agent_from_yaml_path(self, yaml_path: str | Path) -> ChatAgent:
|
||||
"""Create a ChatAgent from a YAML file path.
|
||||
@@ -215,6 +227,8 @@ class AgentFactory:
|
||||
ModuleNotFoundError: If the required module for the provider type cannot be imported.
|
||||
AttributeError: If the required class for the provider type cannot be found in the module.
|
||||
"""
|
||||
# Set safe_mode context before parsing YAML to control PowerFx environment variable access
|
||||
_safe_mode_context.set(self.safe_mode)
|
||||
prompt_agent = agent_schema_dispatch(yaml.safe_load(yaml_str))
|
||||
if not isinstance(prompt_agent, PromptAgent):
|
||||
raise DeclarativeLoaderError("Only yaml definitions for a PromptAgent are supported for agent creation.")
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import MutableMapping
|
||||
from contextvars import ContextVar
|
||||
from typing import Any, Literal, TypeVar, Union
|
||||
|
||||
from agent_framework import get_logger
|
||||
@@ -21,6 +22,11 @@ else:
|
||||
|
||||
logger = get_logger("agent_framework.declarative")
|
||||
|
||||
# Context variable for safe_mode setting.
|
||||
# When True (default), environment variables are NOT accessible in PowerFx expressions.
|
||||
# When False, environment variables CAN be accessed via Env symbol in PowerFx.
|
||||
_safe_mode_context: ContextVar[bool] = ContextVar("safe_mode", default=True)
|
||||
|
||||
|
||||
@overload
|
||||
def _try_powerfx_eval(value: None, log_value: bool = True) -> None: ...
|
||||
@@ -49,6 +55,9 @@ def _try_powerfx_eval(value: str | None, log_value: bool = True) -> str | None:
|
||||
)
|
||||
return value
|
||||
try:
|
||||
safe_mode = _safe_mode_context.get()
|
||||
if safe_mode:
|
||||
return engine.eval(value[1:])
|
||||
return engine.eval(value[1:], symbols={"Env": dict(os.environ)})
|
||||
except Exception as exc:
|
||||
if log_value:
|
||||
|
||||
@@ -454,3 +454,140 @@ def test_agent_schema_dispatch_agent_samples(yaml_file: Path, agent_samples_dir:
|
||||
result = agent_schema_dispatch(yaml.safe_load(content))
|
||||
# Result can be None for unknown kinds, but should not raise exceptions
|
||||
assert result is not None, f"agent_schema_dispatch returned None for {yaml_file.relative_to(agent_samples_dir)}"
|
||||
|
||||
|
||||
class TestAgentFactorySafeMode:
|
||||
"""Tests for AgentFactory safe_mode parameter."""
|
||||
|
||||
def test_agent_factory_safe_mode_default_is_true(self):
|
||||
"""Test that safe_mode is True by default."""
|
||||
from agent_framework_declarative._loader import AgentFactory
|
||||
|
||||
factory = AgentFactory()
|
||||
assert factory.safe_mode is True
|
||||
|
||||
def test_agent_factory_safe_mode_can_be_set_false(self):
|
||||
"""Test that safe_mode can be explicitly set to False."""
|
||||
from agent_framework_declarative._loader import AgentFactory
|
||||
|
||||
factory = AgentFactory(safe_mode=False)
|
||||
assert factory.safe_mode is False
|
||||
|
||||
def test_agent_factory_safe_mode_blocks_env_in_yaml(self, monkeypatch):
|
||||
"""Test that safe_mode=True blocks environment variable access in YAML parsing."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from agent_framework_declarative._loader import AgentFactory
|
||||
|
||||
monkeypatch.setenv("TEST_MODEL_ID", "gpt-4-from-env")
|
||||
|
||||
# Create a mock chat client to avoid needing real provider
|
||||
mock_client = MagicMock()
|
||||
|
||||
yaml_content = """
|
||||
kind: Prompt
|
||||
name: test-agent
|
||||
description: =Env.TEST_DESCRIPTION
|
||||
instructions: Hello world
|
||||
"""
|
||||
monkeypatch.setenv("TEST_DESCRIPTION", "Description from env")
|
||||
|
||||
# With safe_mode=True (default), Env access should fail and return original value
|
||||
factory = AgentFactory(chat_client=mock_client, safe_mode=True)
|
||||
agent = factory.create_agent_from_yaml(yaml_content)
|
||||
|
||||
# The description should NOT be resolved from env (PowerFx fails, returns original)
|
||||
assert agent.description == "=Env.TEST_DESCRIPTION"
|
||||
|
||||
def test_agent_factory_safe_mode_false_allows_env_in_yaml(self, monkeypatch):
|
||||
"""Test that safe_mode=False allows environment variable access in YAML parsing."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from agent_framework_declarative._loader import AgentFactory
|
||||
|
||||
monkeypatch.setenv("TEST_DESCRIPTION", "Description from env")
|
||||
|
||||
# Create a mock chat client to avoid needing real provider
|
||||
mock_client = MagicMock()
|
||||
|
||||
yaml_content = """
|
||||
kind: Prompt
|
||||
name: test-agent
|
||||
description: =Env.TEST_DESCRIPTION
|
||||
instructions: Hello world
|
||||
"""
|
||||
|
||||
# With safe_mode=False, Env access should work
|
||||
factory = AgentFactory(chat_client=mock_client, safe_mode=False)
|
||||
agent = factory.create_agent_from_yaml(yaml_content)
|
||||
|
||||
# The description should be resolved from env
|
||||
assert agent.description == "Description from env"
|
||||
|
||||
def test_agent_factory_safe_mode_with_api_key_connection(self, monkeypatch):
|
||||
"""Test safe_mode with API key connection containing env variable."""
|
||||
from agent_framework_declarative._models import _safe_mode_context
|
||||
|
||||
monkeypatch.setenv("MY_API_KEY", "secret-key-123")
|
||||
|
||||
yaml_content = """
|
||||
kind: Prompt
|
||||
name: test-agent
|
||||
description: Test agent
|
||||
instructions: Hello
|
||||
model:
|
||||
id: gpt-4
|
||||
provider: OpenAI
|
||||
apiType: Chat
|
||||
connection:
|
||||
kind: key
|
||||
apiKey: =Env.MY_API_KEY
|
||||
"""
|
||||
|
||||
# Manually trigger the YAML parsing to check the context is set correctly
|
||||
import yaml as yaml_module
|
||||
|
||||
from agent_framework_declarative._models import agent_schema_dispatch
|
||||
|
||||
token = _safe_mode_context.set(True) # Ensure we're in safe mode
|
||||
try:
|
||||
result = agent_schema_dispatch(yaml_module.safe_load(yaml_content))
|
||||
|
||||
# The API key should NOT be resolved (still has the PowerFx expression)
|
||||
assert result.model.connection.apiKey == "=Env.MY_API_KEY"
|
||||
finally:
|
||||
_safe_mode_context.reset(token)
|
||||
|
||||
def test_agent_factory_safe_mode_false_resolves_api_key(self, monkeypatch):
|
||||
"""Test safe_mode=False resolves API key from environment."""
|
||||
from agent_framework_declarative._models import _safe_mode_context
|
||||
|
||||
monkeypatch.setenv("MY_API_KEY", "secret-key-123")
|
||||
|
||||
yaml_content = """
|
||||
kind: Prompt
|
||||
name: test-agent
|
||||
description: Test agent
|
||||
instructions: Hello
|
||||
model:
|
||||
id: gpt-4
|
||||
provider: OpenAI
|
||||
apiType: Chat
|
||||
connection:
|
||||
kind: key
|
||||
apiKey: =Env.MY_API_KEY
|
||||
"""
|
||||
|
||||
# With safe_mode=False, the API key should be resolved
|
||||
import yaml as yaml_module
|
||||
|
||||
from agent_framework_declarative._models import agent_schema_dispatch
|
||||
|
||||
token = _safe_mode_context.set(False) # Disable safe mode
|
||||
try:
|
||||
result = agent_schema_dispatch(yaml_module.safe_load(yaml_content))
|
||||
|
||||
# The API key should be resolved from environment
|
||||
assert result.model.connection.apiKey == "secret-key-123"
|
||||
finally:
|
||||
_safe_mode_context.reset(token)
|
||||
|
||||
@@ -41,6 +41,7 @@ from agent_framework_declarative._models import (
|
||||
Template,
|
||||
ToolResource,
|
||||
WebSearchTool,
|
||||
_safe_mode_context,
|
||||
_try_powerfx_eval,
|
||||
)
|
||||
|
||||
@@ -874,35 +875,50 @@ class TestTryPowerfxEval:
|
||||
monkeypatch.setenv("API_KEY", "secret123")
|
||||
monkeypatch.setenv("PORT", "8080")
|
||||
|
||||
# Test basic env access
|
||||
assert _try_powerfx_eval("=Env.TEST_VAR") == "test_value"
|
||||
assert _try_powerfx_eval("=Env.API_KEY") == "secret123"
|
||||
assert _try_powerfx_eval("=Env.PORT") == "8080"
|
||||
# Set safe_mode=False to allow environment variable access
|
||||
token = _safe_mode_context.set(False)
|
||||
try:
|
||||
# Test basic env access
|
||||
assert _try_powerfx_eval("=Env.TEST_VAR") == "test_value"
|
||||
assert _try_powerfx_eval("=Env.API_KEY") == "secret123"
|
||||
assert _try_powerfx_eval("=Env.PORT") == "8080"
|
||||
finally:
|
||||
_safe_mode_context.reset(token)
|
||||
|
||||
def test_env_variable_with_string_concatenation(self, monkeypatch):
|
||||
"""Test env variables with string concatenation operator."""
|
||||
monkeypatch.setenv("BASE_URL", "https://api.example.com")
|
||||
monkeypatch.setenv("API_VERSION", "v1")
|
||||
|
||||
# Test concatenation with &
|
||||
result = _try_powerfx_eval('=Env.BASE_URL & "/" & Env.API_VERSION')
|
||||
assert result == "https://api.example.com/v1"
|
||||
# Set safe_mode=False to allow environment variable access
|
||||
token = _safe_mode_context.set(False)
|
||||
try:
|
||||
# Test concatenation with &
|
||||
result = _try_powerfx_eval('=Env.BASE_URL & "/" & Env.API_VERSION')
|
||||
assert result == "https://api.example.com/v1"
|
||||
|
||||
# Test concatenation with literals
|
||||
result = _try_powerfx_eval('="API Key: " & Env.API_VERSION')
|
||||
assert result == "API Key: v1"
|
||||
# Test concatenation with literals
|
||||
result = _try_powerfx_eval('="API Key: " & Env.API_VERSION')
|
||||
assert result == "API Key: v1"
|
||||
finally:
|
||||
_safe_mode_context.reset(token)
|
||||
|
||||
def test_string_comparison_operators(self, monkeypatch):
|
||||
"""Test PowerFx string comparison operators."""
|
||||
monkeypatch.setenv("ENV_MODE", "production")
|
||||
|
||||
# Equal to - returns bool
|
||||
assert _try_powerfx_eval('=Env.ENV_MODE = "production"') is True
|
||||
assert _try_powerfx_eval('=Env.ENV_MODE = "development"') is False
|
||||
# Set safe_mode=False to allow environment variable access
|
||||
token = _safe_mode_context.set(False)
|
||||
try:
|
||||
# Equal to - returns bool
|
||||
assert _try_powerfx_eval('=Env.ENV_MODE = "production"') is True
|
||||
assert _try_powerfx_eval('=Env.ENV_MODE = "development"') is False
|
||||
|
||||
# Not equal to - returns bool
|
||||
assert _try_powerfx_eval('=Env.ENV_MODE <> "development"') is True
|
||||
assert _try_powerfx_eval('=Env.ENV_MODE <> "production"') is False
|
||||
# Not equal to - returns bool
|
||||
assert _try_powerfx_eval('=Env.ENV_MODE <> "development"') is True
|
||||
assert _try_powerfx_eval('=Env.ENV_MODE <> "production"') is False
|
||||
finally:
|
||||
_safe_mode_context.reset(token)
|
||||
|
||||
def test_string_in_operator(self):
|
||||
"""Test PowerFx 'in' operator for substring testing (case-insensitive)."""
|
||||
@@ -958,11 +974,54 @@ class TestTryPowerfxEval:
|
||||
monkeypatch.setenv("URL_WITH_QUERY", "https://example.com?param=value")
|
||||
monkeypatch.setenv("PATH_WITH_SPACES", "C:\\Program Files\\App")
|
||||
|
||||
result = _try_powerfx_eval("=Env.URL_WITH_QUERY")
|
||||
assert result == "https://example.com?param=value"
|
||||
# Set safe_mode=False to allow environment variable access
|
||||
token = _safe_mode_context.set(False)
|
||||
try:
|
||||
result = _try_powerfx_eval("=Env.URL_WITH_QUERY")
|
||||
assert result == "https://example.com?param=value"
|
||||
|
||||
result = _try_powerfx_eval("=Env.PATH_WITH_SPACES")
|
||||
assert result == "C:\\Program Files\\App"
|
||||
result = _try_powerfx_eval("=Env.PATH_WITH_SPACES")
|
||||
assert result == "C:\\Program Files\\App"
|
||||
finally:
|
||||
_safe_mode_context.reset(token)
|
||||
|
||||
def test_safe_mode_blocks_env_access(self, monkeypatch):
|
||||
"""Test that safe_mode=True (default) blocks environment variable access."""
|
||||
monkeypatch.setenv("SECRET_VAR", "secret_value")
|
||||
|
||||
# Set safe_mode=True (default)
|
||||
token = _safe_mode_context.set(True)
|
||||
try:
|
||||
# When safe_mode=True, Env is not available and the expression fails,
|
||||
# returning the original value
|
||||
result = _try_powerfx_eval("=Env.SECRET_VAR")
|
||||
assert result == "=Env.SECRET_VAR"
|
||||
finally:
|
||||
_safe_mode_context.reset(token)
|
||||
|
||||
def test_safe_mode_context_isolation(self, monkeypatch):
|
||||
"""Test that safe_mode context variable properly isolates env access."""
|
||||
monkeypatch.setenv("TEST_VAR", "test_value")
|
||||
|
||||
# First, set safe_mode=True - should NOT allow env access
|
||||
token = _safe_mode_context.set(True)
|
||||
try:
|
||||
result_safe = _try_powerfx_eval("=Env.TEST_VAR")
|
||||
assert result_safe == "=Env.TEST_VAR"
|
||||
|
||||
# Then, set safe_mode=False - should allow env access
|
||||
token2 = _safe_mode_context.set(False)
|
||||
try:
|
||||
result_unsafe = _try_powerfx_eval("=Env.TEST_VAR")
|
||||
assert result_unsafe == "test_value"
|
||||
finally:
|
||||
_safe_mode_context.reset(token2)
|
||||
|
||||
# After reset, should block again
|
||||
result_safe_again = _try_powerfx_eval("=Env.TEST_VAR")
|
||||
assert result_safe_again == "=Env.TEST_VAR"
|
||||
finally:
|
||||
_safe_mode_context.reset(token)
|
||||
|
||||
|
||||
class TestAgentManifest:
|
||||
|
||||
Reference in New Issue
Block a user