diff --git a/src/google/adk/tools/agent_simulator/agent_simulator_config.py b/src/google/adk/tools/agent_simulator/agent_simulator_config.py index 7f879e14..17051c68 100644 --- a/src/google/adk/tools/agent_simulator/agent_simulator_config.py +++ b/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]): diff --git a/src/google/adk/tools/agent_simulator/agent_simulator_engine.py b/src/google/adk/tools/agent_simulator/agent_simulator_engine.py index 5e611ee9..577e8639 100644 --- a/src/google/adk/tools/agent_simulator/agent_simulator_engine.py +++ b/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, ) diff --git a/src/google/adk/tools/agent_simulator/strategies/base.py b/src/google/adk/tools/agent_simulator/strategies/base.py index a41f18e8..d4ffe1bd 100644 --- a/src/google/adk/tools/agent_simulator/strategies/base.py +++ b/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"} diff --git a/src/google/adk/tools/agent_simulator/strategies/tool_spec_mock_strategy.py b/src/google/adk/tools/agent_simulator/strategies/tool_spec_mock_strategy.py index 9386b38a..2ad3a892 100644 --- a/src/google/adk/tools/agent_simulator/strategies/tool_spec_mock_strategy.py +++ b/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} + + 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 { diff --git a/src/google/adk/tools/agent_simulator/tool_connection_analyzer.py b/src/google/adk/tools/agent_simulator/tool_connection_analyzer.py index 28475175..91655e7b 100644 --- a/src/google/adk/tools/agent_simulator/tool_connection_analyzer.py +++ b/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 '}}'. """