feat: Update agent simulator by improving prompts and add environment data

PiperOrigin-RevId: 868324749

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

@@ -131,7 +131,7 @@ class AgentSimulatorConfig(BaseModel):
   simulation_model_configuration: genai_types.GenerateContentConfig = Field(
       default_factory=lambda: genai_types.GenerateContentConfig(
           thinking_config=genai_types.ThinkingConfig(
-              include_thoughts=True,
+              include_thoughts=False,
               thinking_budget=10240,
           )
       ),
@@ -142,6 +142,11 @@ class AgentSimulatorConfig(BaseModel):
   """The path to the tracing file to be used for mocking. Only used if the
   mock_strategy_type is MOCK_STRATEGY_TRACING."""
 
+  environment_data: Optional[str] = None
+  """Environment-specific data (e.g., a minimal database dump in JSON string
+   format). This data is passed directly to mock strategies for contextual
+   mock generation."""
+
   @field_validator("tool_simulation_configs")
   @classmethod
   def check_tool_simulation_configs(cls, v: List[ToolSimulationConfig]):

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

@@ -65,6 +65,7 @@ class AgentSimulatorEngine:
     )
     self._state_store = {}
     self._random_generator = random.Random()
+    self._environment_data = config.environment_data
 
   async def simulate(
       self, tool: BaseTool, args: Dict[str, Any], tool_context: Any
@@ -128,5 +129,10 @@ class AgentSimulatorEngine:
         self._config.simulation_model_configuration,
     )
     return await mock_strategy.mock(
-        tool, args, tool_context, self._tool_connection_map, self._state_store
+        tool,
+        args,
+        tool_context,
+        self._tool_connection_map,
+        self._state_store,
+        self._environment_data,
     )

src/google/adk/tools/agent_simulator/strategies/base.py

@@ -31,6 +31,7 @@ class MockStrategy:
       tool_context: Any,
       tool_connection_map: Optional[ToolConnectionMap],
       state_store: Dict[str, Any],
+      environment_data: Optional[str] = None,
   ) -> Dict[str, Any]:
     """Generates a mock response for a tool call."""
     raise NotImplementedError()
@@ -52,6 +53,7 @@ class TracingMockStrategy(MockStrategy):
       tool_context: Any,
       tool_connection_map: Optional[ToolConnectionMap],
       state_store: Dict[str, Any],
+      environment_data: Optional[str] = None,
   ) -> Dict[str, Any]:
     # TODO: Implement tracing LLM-based mocking.
     return {"status": "error", "error_message": "Not implemented"}

src/google/adk/tools/agent_simulator/strategies/tool_spec_mock_strategy.py

@@ -35,6 +35,8 @@ _TOOL_SPEC_MOCK_PROMPT_TEMPLATE = """
   realistic JSON response for a tool call, maintaining consistency based
   on a shared state.
 
+  {environment_data_snippet}
+
   Here is the map of how tools connect via stateful parameters:
   {tool_connection_map_json}
 
@@ -53,17 +55,37 @@ _TOOL_SPEC_MOCK_PROMPT_TEMPLATE = """
   2.  If it's a "consuming" tool, check the provided arguments against
       the state store. If an ID is provided that does not exist in the
       state, return a realistic error (e.g., a 404 Not Found error).
-      Otherwise, use the data from the state to generate the response.
+      Otherwise, use the data from the state and the provided environment data
+      to generate the response.
   3.  If it's a "creating" tool, generate a new, unique ID for the
       stateful parameter (e.g., a random string for a ticket_id). Include
       this new ID in your response. I will then update the state with it.
-  4.  Generate a convincing, valid JSON object that mocks the tool's
+  4.  Leverage the provided environment data (if any) to make your response
+      more realistic and consistent with the simulated environment.
+  5.  Generate a convincing, valid JSON object that mocks the tool's
       response. The response must be only the JSON object, without any
       additional text or formatting.
-  5.  The response must start with '{{' and end with '}}'.
+  6.  The response must start with '{{' and end with '}}'.
   """
 
 
+def _find_value_by_key(data: Any, target_key: str) -> Optional[Any]:
+  """Recursively searches for a value by key in a nested structure."""
+  if isinstance(data, dict):
+    if target_key in data:
+      return data[target_key]
+    for key, value in data.items():
+      result = _find_value_by_key(value, target_key)
+      if result is not None:
+        return result
+  elif isinstance(data, list):
+    for item in data:
+      result = _find_value_by_key(item, target_key)
+      if result is not None:
+        return result
+  return None
+
+
 class ToolSpecMockStrategy(MockStrategy):
   """Mocks a tool response based on the tool's specification."""
 
@@ -83,6 +105,7 @@ class ToolSpecMockStrategy(MockStrategy):
       tool_context: Any,
       tool_connection_map: Optional[ToolConnectionMap],
       state_store: Dict[str, Any],
+      environment_data: Optional[str] = None,
   ) -> Dict[str, Any]:
     declaration = tool._get_declaration()
     if not declaration:
@@ -97,10 +120,23 @@ class ToolSpecMockStrategy(MockStrategy):
         else "''"
     )
     state_store_json = json.dumps(state_store, indent=2)
-    tool_schema_json = json.dumps(declaration.model_dump(), indent=2)
+    tool_schema_json = json.dumps(
+        declaration.model_dump(exclude_none=True), indent=2
+    )
     tool_arguments_json = json.dumps(args, indent=2)
 
+    environment_data_snippet = ""
+    if environment_data:
+      environment_data_snippet = f"""
+        Here is relevant environment data (e.g., database snippet, context information):
+        <environment_data>
+        {environment_data}
+        </environment_data>
+        Use this information to generate more realistic responses.
+      """
+
     prompt = _TOOL_SPEC_MOCK_PROMPT_TEMPLATE.format(
+        environment_data_snippet=environment_data_snippet,
         tool_connection_map_json=tool_connection_map_json,
         state_store_json=state_store_json,
         tool_name=tool.name,
@@ -133,16 +169,33 @@ class ToolSpecMockStrategy(MockStrategy):
       clean_json_text = re.sub(r"^```[a-zA-Z]*\n", "", response_text)
       clean_json_text = re.sub(r"\n```$", "", clean_json_text)
       mock_response = json.loads(clean_json_text.strip())
-      # After getting the response, update the state if this was a creating tool.
+      # Determine if the current tool is mutative by checking the connection map.
+      is_mutative = False
       if tool_connection_map:
+        all_creating_tools = {
+            tool_name
+            for param in tool_connection_map.stateful_parameters
+            for tool_name in param.creating_tools
+        }
+        if tool.name in all_creating_tools:
+          is_mutative = True
+
+      # After getting the response, update the state if this was a mutative tool.
+      if is_mutative:
         for param_info in tool_connection_map.stateful_parameters:
           param_name = param_info.parameter_name
+          # Only update the state for the specific parameter this tool
+          # creates/modifies.
           if tool.name in param_info.creating_tools:
-            if param_name in mock_response:
-              param_value = mock_response[param_name]
+            param_value = _find_value_by_key(mock_response, param_name)
+            if param_value is not None:
               if param_name not in state_store:
                 state_store[param_name] = {}
+              # Store the entire response as the new state for this entity.
+              # This correctly captures creations and modifications (like
+              # cancellation).
               state_store[param_name][param_value] = mock_response
+
       return mock_response
     except json.JSONDecodeError:
       return {

src/google/adk/tools/agent_simulator/tool_connection_analyzer.py

@@ -31,26 +31,46 @@ from google.adk.utils.context_utils import Aclosing
 from google.genai import types as genai_types
 
 _TOOL_CONNECTION_ANALYSIS_PROMPT_TEMPLATE = """
-  You are a software architect analyzing a set of tools to understand
-  how they connect to enable stateful operations. Your task is to
-  identify parameters that are generated by one tool and consumed by
-  another.
+  You are an expert software architect analyzing a set of tools to understand
+  stateful dependencies. Your task is to identify parameters that act as
+  stateful identifiers (like IDs) and classify the tools that interact with
+  them.
 
-  For example, a "create_ticket" tool might output a "ticket_id", which is
-  then used as input for "get_ticket" or "close_ticket" tools.
+  **Definitions:**
+  - A **"creating tool"** is a tool that creates a new resource or makes a
+    significant state change to an existing one (e.g., creating, updating,
+    canceling, or deleting). Tool names like `create_account`, `cancel_order`,
+    or `update_price` are strong indicators. These tools are responsible for
+    generating or modifying the state associated with an ID.
+  - A **"consuming tool"** is a tool that uses a resource's ID to retrieve
+    information without changing its state. Tool names like `get_user`,
+    `list_events`, or `find_order` are strong indicators.
 
-  Analyze the following tool schemas:
+  **Your Goal:**
+  Analyze the following tool schemas and identify the shared, stateful
+  parameters (like `user_id`, `order_id`, etc.).
+
+  For each stateful parameter you identify, classify the tools into
+  `creating_tools` and `consuming_tools` based on the definitions above.
+
+  **Example:** A `create_ticket` tool would be a `creating_tool` for
+  `ticket_id`. A `get_ticket` tool would be a `consuming_tool` for
+  `ticket_id`. A `list_tickets` tool that takes a `user_id` as input is a
+  `consuming_tool` for `user_id`.
+
+  **Analyze the following tool schemas:**
   {tool_schemas_json}
 
-  Based on this analysis, generate a JSON object that describes these
-  stateful parameters. The JSON object should have a single key,
-  "stateful_parameters", which is a list. Each item in the list
-  should represent a stateful parameter and have the following keys:
+  **Output Format:**
+  Generate a JSON object with a single key, "stateful_parameters", which is a
+  list. Each item in the list must have these keys:
   - "parameter_name": The name of the shared parameter (e.g., "ticket_id").
-  - "creating_tools": A list of tools that generate this parameter.
-  - "consuming_tools": A list of tools that use this parameter as input.
+  - "creating_tools": A list of tools that create or modify this parameter's
+    state.
+  - "consuming_tools": A list of tools that use this parameter as input for
+    read-only operations.
 
-  Return only the raw JSON object.
+  ONLY return the raw JSON object.
   Your response must start with '{{' and end with '}}'.
   """