Skip to content

Commit 99a31bf

Browse files
google-genai-botcopybara-github
google-genai-bot
authored and
copybara-github
committed
refactor: rename agent simulator to environment simulation. Also add tracing into environment simulation
PiperOrigin-RevId: 885759367
1 parent b6185f7 commit 99a31bf

25 files changed

Lines changed: 1125 additions & 745 deletions

‎src/google/adk/features/_feature_registry.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ class FeatureName(str, Enum):
3737
COMPUTER_USE = "COMPUTER_USE"
3838
DATA_AGENT_TOOL_CONFIG = "DATA_AGENT_TOOL_CONFIG"
3939
DATA_AGENT_TOOLSET = "DATA_AGENT_TOOLSET"
40+
ENVIRONMENT_SIMULATION = "ENVIRONMENT_SIMULATION"
4041
GOOGLE_CREDENTIALS_CONFIG = "GOOGLE_CREDENTIALS_CONFIG"
4142
GOOGLE_TOOL = "GOOGLE_TOOL"
4243
JSON_SCHEMA_FOR_FUNC_DECL = "JSON_SCHEMA_FOR_FUNC_DECL"
@@ -116,6 +117,9 @@ class FeatureConfig:
116117
FeatureName.DATA_AGENT_TOOLSET: FeatureConfig(
117118
FeatureStage.EXPERIMENTAL, default_on=True
118119
),
120+
FeatureName.ENVIRONMENT_SIMULATION: FeatureConfig(
121+
FeatureStage.EXPERIMENTAL, default_on=True
122+
),
119123
FeatureName.GOOGLE_CREDENTIALS_CONFIG: FeatureConfig(
120124
FeatureStage.EXPERIMENTAL, default_on=True
121125
),

‎src/google/adk/tools/agent_simulator/__init__.py‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,15 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15-
from google.adk.tools.agent_simulator.agent_simulator_factory import AgentSimulatorFactory
15+
import warnings
1616

17-
__all__ = ["AgentSimulator"]
17+
from google.adk.tools.environment_simulation import EnvironmentSimulationFactory as AgentSimulatorFactory
18+
19+
warnings.warn(
20+
"google.adk.tools.agent_simulator is moved to"
21+
" google.adk.tools.environment_simulation",
22+
DeprecationWarning,
23+
stacklevel=2,
24+
)
25+
26+
__all__ = ["AgentSimulatorFactory"]

‎src/google/adk/tools/agent_simulator/agent_simulator_config.py‎

Lines changed: 39 additions & 138 deletions
Original file line numberDiff line numberDiff line change
@@ -14,150 +14,51 @@
1414

1515
from __future__ import annotations
1616

17-
import enum
1817
from typing import Any
19-
from typing import Dict
20-
from typing import List
21-
from typing import Optional
18+
import warnings
2219

23-
from google.genai import types as genai_types
24-
from pydantic import BaseModel
25-
from pydantic import Field
26-
from pydantic import field_validator
20+
from google.adk.tools.environment_simulation.environment_simulation_config import EnvironmentSimulationConfig
21+
from google.adk.tools.environment_simulation.environment_simulation_config import InjectedError
22+
from google.adk.tools.environment_simulation.environment_simulation_config import InjectionConfig
23+
from google.adk.tools.environment_simulation.environment_simulation_config import MockStrategy
24+
from google.adk.tools.environment_simulation.environment_simulation_config import ToolSimulationConfig
2725
from pydantic import model_validator
28-
from pydantic_core import ValidationError
2926

27+
warnings.warn(
28+
"google.adk.tools.agent_simulator.agent_simulator_config is moved to"
29+
" google.adk.tools.environment_simulation.environment_simulation_config",
30+
DeprecationWarning,
31+
stacklevel=2,
32+
)
3033

31-
class InjectedError(BaseModel):
32-
"""An error to be injected into a tool call."""
3334

34-
injected_http_error_code: int
35-
"""Inject http error code to the tool call. Will present as "error_code"
36-
in the tool response dict."""
35+
class AgentSimulatorConfig(EnvironmentSimulationConfig):
36+
"""Deprecated AgentSimulatorConfig alias.
3737
38-
error_message: str
39-
"""Inject error message to the tool call. Will present as
40-
"error_message" in the tool response dict."""
38+
Forwards tracing_path to tracing.
39+
"""
4140

42-
43-
class InjectionConfig(BaseModel):
44-
"""Injection configuration for a tool."""
45-
46-
injection_probability: float = 1.0
47-
"""Probability of injecting the injected_value."""
48-
49-
match_args: Optional[Dict[str, Any]] = None
50-
"""Only apply injection if the request matches the match_args.
51-
If match_args is not provided, the injection will be applied to all
52-
requests."""
53-
54-
injected_latency_seconds: float = Field(default=0.0, le=120.0)
55-
"""Inject latency to the tool call. Please note it may not be accurate if │
56-
the interceptor is applied as after tool callback."""
57-
58-
random_seed: Optional[int] = None
59-
"""The random seed to use for this injection."""
60-
61-
injected_error: Optional[InjectedError] = None
62-
"""The injected error."""
63-
64-
injected_response: Optional[Dict[str, Any]] = None
65-
"""The injected response."""
66-
67-
@model_validator(mode="after")
68-
def check_injected_error_or_response(self) -> Self:
69-
"""Checks that either injected_error or injected_response is set."""
70-
if bool(self.injected_error) == bool(self.injected_response):
71-
raise ValueError(
72-
"Either injected_error or injected_response must be set, but not"
73-
" both, and not neither."
74-
)
75-
return self
76-
77-
78-
class MockStrategy(enum.Enum):
79-
"""Mock strategy for a tool."""
80-
81-
MOCK_STRATEGY_UNSPECIFIED = 0
82-
83-
MOCK_STRATEGY_TOOL_SPEC = 1
84-
"""Use tool specifications to mock the tool response."""
85-
86-
MOCK_STRATEGY_TRACING = 2
87-
"""Use provided tracing and tool specifications to mock the tool
88-
response based on llm response. Need to provide tracing path in
89-
command."""
90-
91-
92-
class ToolSimulationConfig(BaseModel):
93-
"""Simulation configuration for a single tool."""
94-
95-
tool_name: str
96-
"""Name of the tool to be simulated."""
97-
98-
injection_configs: List[InjectionConfig] = Field(default_factory=list)
99-
"""Injection configuration for the tool. If provided, the tool will be
100-
injected with the injected_value with the injection_probability first,
101-
the mock_strategy will be applied if no injection config is hit."""
102-
103-
mock_strategy_type: MockStrategy = MockStrategy.MOCK_STRATEGY_UNSPECIFIED
104-
"""The mock strategy to use."""
105-
106-
@model_validator(mode="after")
107-
def check_mock_strategy_type(self) -> Self:
108-
"""Checks that mock_strategy_type is not UNSPECIFIED if no injections."""
109-
if (
110-
not self.injection_configs
111-
and self.mock_strategy_type == MockStrategy.MOCK_STRATEGY_UNSPECIFIED
112-
):
113-
raise ValueError(
114-
"If injection_configs is empty, mock_strategy_type cannot be"
115-
" MOCK_STRATEGY_UNSPECIFIED."
116-
)
117-
return self
118-
119-
120-
class AgentSimulatorConfig(BaseModel):
121-
"""Configuration for AgentSimulator."""
122-
123-
tool_simulation_configs: List[ToolSimulationConfig] = Field(
124-
default_factory=list
125-
)
126-
"""A list of tool simulation configurations."""
127-
128-
simulation_model: str = Field(default="gemini-2.5-flash")
129-
"""The model to use for internal simulator LLM calls (tool analysis, mock responses)."""
130-
131-
simulation_model_configuration: genai_types.GenerateContentConfig = Field(
132-
default_factory=lambda: genai_types.GenerateContentConfig(
133-
thinking_config=genai_types.ThinkingConfig(
134-
include_thoughts=False,
135-
thinking_budget=10240,
136-
)
137-
),
138-
)
139-
"""The configuration for the internal simulator LLM calls."""
140-
141-
tracing_path: Optional[str] = None
142-
"""The path to the tracing file to be used for mocking. Only used if the
143-
mock_strategy_type is MOCK_STRATEGY_TRACING."""
144-
145-
environment_data: Optional[str] = None
146-
"""Environment-specific data (e.g., a minimal database dump in JSON string
147-
format). This data is passed directly to mock strategies for contextual
148-
mock generation."""
149-
150-
@field_validator("tool_simulation_configs")
41+
@model_validator(mode="before")
15142
@classmethod
152-
def check_tool_simulation_configs(cls, v: List[ToolSimulationConfig]):
153-
"""Checks that tool_simulation_configs is not empty."""
154-
if not v:
155-
raise ValueError("tool_simulation_configs must be provided.")
156-
seen_tool_names = set()
157-
for tool_sim_config in v:
158-
if tool_sim_config.tool_name in seen_tool_names:
159-
raise ValueError(
160-
f"Duplicate tool_name found: {tool_sim_config.tool_name}"
161-
)
162-
seen_tool_names.add(tool_sim_config.tool_name)
163-
return v
43+
def convert_tracing_path(cls, data: Any) -> Any:
44+
"""Convert tracing_path to tracing."""
45+
if isinstance(data, dict) and "tracing_path" in data:
46+
warnings.warn(
47+
"`tracing_path` is deprecated. Use `tracing` instead.",
48+
DeprecationWarning,
49+
stacklevel=2,
50+
)
51+
if "tracing" not in data:
52+
data["tracing"] = data.pop("tracing_path")
53+
else:
54+
data.pop("tracing_path")
55+
return data
56+
57+
58+
__all__ = [
59+
"AgentSimulatorConfig",
60+
"InjectedError",
61+
"InjectionConfig",
62+
"MockStrategy",
63+
"ToolSimulationConfig",
64+
]

‎src/google/adk/tools/agent_simulator/agent_simulator_engine.py‎

Lines changed: 9 additions & 119 deletions
Original file line numberDiff line numberDiff line change
@@ -14,125 +14,15 @@
1414

1515
from __future__ import annotations
1616

17-
import asyncio
18-
import concurrent.futures
19-
import logging
20-
import random
21-
import time
22-
from typing import Any
23-
from typing import Dict
24-
from typing import Optional
17+
import warnings
2518

26-
agent_simulator_logger = logging.getLogger("agent_simulator_logger")
19+
from google.adk.tools.environment_simulation.environment_simulation_engine import EnvironmentSimulationEngine as AgentSimulatorEngine
2720

28-
from google.adk.agents.llm_agent import LlmAgent
29-
from google.adk.tools.agent_simulator.agent_simulator_config import AgentSimulatorConfig
30-
from google.adk.tools.agent_simulator.agent_simulator_config import MockStrategy as MockStrategyEnum
31-
from google.adk.tools.agent_simulator.agent_simulator_config import ToolSimulationConfig
32-
from google.adk.tools.agent_simulator.strategies import base as base_mock_strategies
33-
from google.adk.tools.agent_simulator.strategies import tool_spec_mock_strategy
34-
from google.adk.tools.agent_simulator.tool_connection_analyzer import ToolConnectionAnalyzer
35-
from google.adk.tools.agent_simulator.tool_connection_map import ToolConnectionMap
36-
from google.adk.tools.base_tool import BaseTool
21+
warnings.warn(
22+
"google.adk.tools.agent_simulator.agent_simulator_engine is moved to"
23+
" google.adk.tools.environment_simulation.environment_simulation_engine",
24+
DeprecationWarning,
25+
stacklevel=2,
26+
)
3727

38-
39-
def _create_mock_strategy(
40-
mock_strategy_type: MockStrategyEnum,
41-
llm_name: str,
42-
llm_config: genai_types.GenerateContentConfig,
43-
) -> base_mock_strategies.MockStrategy:
44-
"""Creates a mock strategy based on the given type."""
45-
if mock_strategy_type == MockStrategyEnum.MOCK_STRATEGY_TOOL_SPEC:
46-
return tool_spec_mock_strategy.ToolSpecMockStrategy(llm_name, llm_config)
47-
if mock_strategy_type == MockStrategyEnum.MOCK_STRATEGY_TRACING:
48-
return base_mock_strategies.TracingMockStrategy()
49-
raise ValueError(f"Unknown mock strategy type: {mock_strategy_type}")
50-
51-
52-
class AgentSimulatorEngine:
53-
"""Core engine to handle the simulation logic."""
54-
55-
def __init__(self, config: AgentSimulatorConfig):
56-
self._config = config
57-
self._tool_sim_configs = {
58-
c.tool_name: c for c in config.tool_simulation_configs
59-
}
60-
self._is_analyzed = False
61-
self._tool_connection_map: Optional[ToolConnectionMap] = None
62-
self._analyzer = ToolConnectionAnalyzer(
63-
llm_name=config.simulation_model,
64-
llm_config=config.simulation_model_configuration,
65-
)
66-
self._state_store = {}
67-
self._random_generator = random.Random()
68-
self._environment_data = config.environment_data
69-
70-
async def simulate(
71-
self, tool: BaseTool, args: Dict[str, Any], tool_context: Any
72-
) -> Optional[Dict[str, Any]]:
73-
"""Simulates a tool call."""
74-
if tool.name not in self._tool_sim_configs:
75-
return None
76-
77-
tool_sim_config = self._tool_sim_configs[tool.name]
78-
79-
if not self._is_analyzed and any(
80-
c.mock_strategy_type != MockStrategyEnum.MOCK_STRATEGY_UNSPECIFIED
81-
for c in self._config.tool_simulation_configs
82-
):
83-
agent = tool_context._invocation_context.agent
84-
if isinstance(agent, LlmAgent):
85-
tools = await agent.canonical_tools(tool_context)
86-
self._tool_connection_map = await self._analyzer.analyze(tools)
87-
self._is_analyzed = True
88-
89-
for injection_config in tool_sim_config.injection_configs:
90-
if injection_config.match_args:
91-
if not all(
92-
item in args.items() for item in injection_config.match_args.items()
93-
):
94-
continue
95-
96-
if injection_config.random_seed is not None:
97-
self._random_generator.seed(injection_config.random_seed)
98-
99-
if (
100-
self._random_generator.random()
101-
< injection_config.injection_probability
102-
):
103-
time.sleep(injection_config.injected_latency_seconds)
104-
if injection_config.injected_error:
105-
return {
106-
"error_code": (
107-
injection_config.injected_error.injected_http_error_code
108-
),
109-
"error_message": injection_config.injected_error.error_message,
110-
}
111-
if injection_config.injected_response:
112-
return injection_config.injected_response
113-
114-
# If no injection was applied, fall back to the mock strategy.
115-
if (
116-
tool_sim_config.mock_strategy_type
117-
== MockStrategyEnum.MOCK_STRATEGY_UNSPECIFIED
118-
):
119-
agent_simulator_logger.warning(
120-
"Tool '%s' did not hit any injection config and has no mock strategy"
121-
" configured. Returning no-op.",
122-
tool.name,
123-
)
124-
return None
125-
126-
mock_strategy = _create_mock_strategy(
127-
tool_sim_config.mock_strategy_type,
128-
self._config.simulation_model,
129-
self._config.simulation_model_configuration,
130-
)
131-
return await mock_strategy.mock(
132-
tool,
133-
args,
134-
tool_context,
135-
self._tool_connection_map,
136-
self._state_store,
137-
self._environment_data,
138-
)
28+
__all__ = ["AgentSimulatorEngine"]

0 commit comments

Comments
 (0)