refactor(types): type the environment simulation tools for strict mypy

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 970127626

Author: GWeale

src/google/adk/tools/environment/_read_file_tool.py

@@ -143,7 +143,7 @@ async def run_async(
       numbered = ''.join(
           f'{start + i:6d}\t{line}' for i, line in enumerate(lines)
       )
-      result = {
+      result: dict[str, str | int] = {
           'status': 'ok',
           'content': _truncate(
               numbered,

src/google/adk/tools/environment_simulation/environment_simulation_config.py

@@ -26,6 +26,7 @@
 from pydantic import field_validator
 from pydantic import model_validator
 from pydantic import ValidationError  # noqa: F401
+from typing_extensions import Self
 
 from ...features import experimental
 from ...features import FeatureName
@@ -156,7 +157,9 @@ class EnvironmentSimulationConfig(BaseModel):
 
   @field_validator("tool_simulation_configs")
   @classmethod
-  def check_tool_simulation_configs(cls, v: List[ToolSimulationConfig]):
+  def check_tool_simulation_configs(
+      cls, v: List[ToolSimulationConfig]
+  ) -> List[ToolSimulationConfig]:
     """Checks that tool_simulation_configs is not empty."""
     if not v:
       raise ValueError("tool_simulation_configs must be provided.")

src/google/adk/tools/environment_simulation/environment_simulation_engine.py

@@ -33,6 +33,7 @@
 from google.adk.tools.environment_simulation.strategies import tool_spec_mock_strategy
 from google.adk.tools.environment_simulation.tool_connection_analyzer import ToolConnectionAnalyzer
 from google.adk.tools.environment_simulation.tool_connection_map import ToolConnectionMap
+from google.genai import types as genai_types
 
 from ...features import experimental
 from ...features import FeatureName
@@ -66,7 +67,7 @@ def __init__(self, config: EnvironmentSimulationConfig):
         llm_name=config.simulation_model,
         llm_config=config.simulation_model_configuration,
     )
-    self._state_store = {}
+    self._state_store: Dict[str, Any] = {}
     self._random_generator = random.Random()
     self._environment_data = config.environment_data
     self._tracing = config.tracing

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

@@ -20,6 +20,7 @@
 
 from google.adk.features import experimental
 from google.adk.features import FeatureName
+from google.adk.tools.base_tool import BaseTool
 from google.adk.tools.environment_simulation.tool_connection_map import ToolConnectionMap
 from google.genai import types as genai_types
 

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

@@ -17,8 +17,8 @@
 import json
 import re
 from typing import Any
+from typing import cast
 from typing import Dict
-from typing import Optional
 
 from google.adk.features import experimental
 from google.adk.features import FeatureName
@@ -73,12 +73,13 @@
   """
 
 
-def _find_value_by_key(data: Any, target_key: str) -> Optional[Any]:
+def _find_value_by_key(data: object, target_key: str) -> object | None:
   """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: object = data[target_key]
+      return result
+    for value in data.values():
       result = _find_value_by_key(value, target_key)
       if result is not None:
         return result
@@ -108,10 +109,10 @@ async def mock(
       tool: BaseTool,
       args: Dict[str, Any],
       tool_context: Any,
-      tool_connection_map: Optional[ToolConnectionMap],
+      tool_connection_map: ToolConnectionMap | None,
       state_store: Dict[str, Any],
-      environment_data: Optional[str] = None,
-      tracing: Optional[str] = None,
+      environment_data: str | None = None,
+      tracing: str | None = None,
   ) -> Dict[str, Any]:
     declaration = tool._get_declaration()
     if not declaration:
@@ -178,7 +179,9 @@ async def mock(
     response_text = ""
     async with Aclosing(self._llm.generate_content_async(request)) as agen:
       async for llm_response in agen:
-        generated_content: genai_types.Content = llm_response.content
+        generated_content = llm_response.content
+        if generated_content is None:
+          continue
         if generated_content.parts:
           for part in generated_content.parts:
             if part.text:
@@ -187,7 +190,19 @@ async def mock(
     try:
       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())
+      parsed_response: object = json.loads(clean_json_text.strip())
+      if not isinstance(parsed_response, dict) or not all(
+          isinstance(key, str) for key in parsed_response
+      ):
+        return {
+            "status": "error",
+            "error_message": "Generated mock response was not a JSON object.",
+            "llm_output": response_text,
+        }
+      # Keys checked to be str by the guard above; isinstance only narrows the
+      # value type as far as dict[Any, Any].
+      mock_response = cast(dict[str, Any], parsed_response)
+
       # Determine if the current tool is mutative by checking the connection map.
       is_mutative = False
       if tool_connection_map:
@@ -200,7 +215,7 @@ async def mock(
           is_mutative = True
 
       # After getting the response, update the state if this was a mutative tool.
-      if is_mutative:
+      if is_mutative and tool_connection_map:
         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

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

@@ -96,9 +96,9 @@ async def analyze(self, tools: List[BaseTool]) -> ToolConnectionMap:
     Analyzes a list of tools and returns a map of their connections.
     """
     tool_schemas = [
-        tool._get_declaration().model_dump(exclude_none=True)
+        declaration.model_dump(exclude_none=True)
         for tool in tools
-        if tool._get_declaration()
+        if (declaration := tool._get_declaration()) is not None
     ]
     tool_schemas_json = json.dumps(tool_schemas, indent=2)
     prompt = _TOOL_CONNECTION_ANALYSIS_PROMPT_TEMPLATE.format(
@@ -119,7 +119,9 @@ async def analyze(self, tools: List[BaseTool]) -> ToolConnectionMap:
     response_text = ""
     async with Aclosing(self._llm.generate_content_async(request)) as agen:
       async for llm_response in agen:
-        generated_content: genai_types.Content = llm_response.content
+        generated_content = llm_response.content
+        if generated_content is None:
+          continue
         if not generated_content.parts:
           continue
         for part in generated_content.parts:
@@ -129,7 +131,7 @@ async def analyze(self, tools: List[BaseTool]) -> ToolConnectionMap:
     try:
       clean_json_text = re.sub(r"^```[a-zA-Z]*\n", "", response_text)
       clean_json_text = re.sub(r"\n```$", "", clean_json_text)
-      response_json = json.loads(clean_json_text.strip())
+      response_json: object = json.loads(clean_json_text.strip())
     except json.JSONDecodeError as e:
       logging.warning(
           "Failed to parse tool connection analysis from LLM. Proceeding"