Gemini: raise a NON-retryable error on a deterministically answer-less generation.

langfun/assistant/capabilities/gui/bounding_box_parser_test.py

@@ -61,18 +61,18 @@ def test_bbox_basic_functionality(self):
 
     # Test the `to_gui_bbox` method
     gui_bbox = bbox.to_gui_bbox()
-    self.assertEqual(gui_bbox.x, 10)
-    self.assertEqual(gui_bbox.y, 20)
-    self.assertEqual(gui_bbox.right, 50)
-    self.assertEqual(gui_bbox.bottom, 80)
+    self.assertEqual(gui_bbox.x, 10)  # pyrefly: ignore[missing-attribute]
+    self.assertEqual(gui_bbox.y, 20)  # pyrefly: ignore[missing-attribute]
+    self.assertEqual(gui_bbox.right, 50)  # pyrefly: ignore[missing-attribute]
+    self.assertEqual(gui_bbox.bottom, 80)  # pyrefly: ignore[missing-attribute]
 
   def test_simple_json(self):
     json_text = '{"search button": [10, 20, 100, 200]}'
     expected = {'search button': (16, 6, 160, 60)}
     result = bounding_box_parser.parse_and_convert_json(
         json_text, screen_size=(800, 600)
     )
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
   def test_multiple_objects(self):
     json_text = (
@@ -82,7 +82,7 @@ def test_multiple_objects(self):
     result = bounding_box_parser.parse_and_convert_json(
         json_text, screen_size=(800, 600)
     )
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
   def test_nested_json(self):
     json_text = (
@@ -93,15 +93,15 @@ def test_nested_json(self):
     result = bounding_box_parser.parse_and_convert_json(
         json_text, screen_size=(800, 600)
     )
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
   def test_json_in_code_block(self):
     json_text = '```\n{"search button": [10, 20, 100, 200]}\n```'
     expected = {'search button': (16, 6, 160, 60)}
     result = bounding_box_parser.parse_and_convert_json(
         json_text, screen_size=(800, 600)
     )
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
   def test_extract_json_candidate_from_text(self):
     test_cases = [
@@ -174,7 +174,7 @@ def test_dict_in_list(self):
     json_text = '```json\n[\n  {"box_2d": [61, 22, 160, 95]}\n]\n```'
     expected = {'box_2d': (22, 61, 95, 160)}
     result = bounding_box_parser.parse_and_convert_json(json_text)
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
   def test_dict_with_label(self):
     json_text = """```json
@@ -186,7 +186,7 @@ def test_dict_with_label(self):
     expected = {'box_2d': (328, 820, 352, 872)}
     result = bounding_box_parser.parse_and_convert_json(json_text)
     print('result: ', result)
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
   def test_invalid_json(self):
     json_text = 'This is not a valid JSON'
@@ -218,21 +218,21 @@ def test_list_input(self):
     json_text = '[10, 20, 100, 200]'
     expected = {'element': (20, 10, 200, 100)}
     result = bounding_box_parser.parse_and_convert_json(json_text)
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
   def test_default_screen_size(self):
     json_text = '{"button": [10, 20, 100, 200]}'
     expected = {'button': (20, 10, 200, 100)}
     result = bounding_box_parser.parse_and_convert_json(json_text)
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
   def test_float_numbers(self):
     json_text = '{"button": [10.5, 20.2, 100.7, 200.9]}'
     expected = {'button': (20, 10, 200, 100)}  # Expected integer coordinates
     result = bounding_box_parser.parse_and_convert_json(
         json_text, screen_size=(1000, 1000)
     )
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
   def test_type_error_handling(self):
     json_text = '{"button": ["text", 20, 100, 200]}'
@@ -269,7 +269,7 @@ def test_mixed_data_types(self):
     result = bounding_box_parser.parse_and_convert_json(
         json_text, screen_size=(1000, 1000)
     )
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
   def test_none_values(self):
     json_text = '{"button": [null, 20, 100, 200]}'
@@ -290,7 +290,7 @@ def test_large_coordinates(self):
     result = bounding_box_parser.parse_and_convert_json(
         json_text, screen_size=(1000, 1000)
     )
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
   def test_different_string_formats(self):
     # Newlines
@@ -299,15 +299,15 @@ def test_different_string_formats(self):
     result = bounding_box_parser.parse_and_convert_json(
         json_text, screen_size=(800, 600)
     )
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
   def test_deeply_nested_json(self):
     json_text = (
         '{"layer1": {"layer2": {"layer3": {"button": [10, 20, 100, 200]}}}}'
     )
     expected = {'button': (20, 10, 200, 100)}
     result = bounding_box_parser.parse_and_convert_json(json_text)
-    self.assert_bbox_equal(expected, result)
+    self.assert_bbox_equal(expected, result)  # pyrefly: ignore[bad-argument-type]
 
 if __name__ == '__main__':
   unittest.main()

langfun/assistant/capabilities/gui/location.py

@@ -77,9 +77,9 @@ def random(cls,
     Returns:
       A random coordinate.
     """
-    rand = rand or random
-    x = rand.randint(bound.left, bound.right)
-    y = rand.randint(bound.top, bound.bottom)
+    rand = rand or random  # pyrefly: ignore[bad-assignment]
+    x = rand.randint(bound.left, bound.right)  # pyrefly: ignore[missing-attribute]
+    y = rand.randint(bound.top, bound.bottom)  # pyrefly: ignore[missing-attribute]
     return cls(x, y)
 
   def distance_to(self, point: 'Coordinate') -> float:
@@ -241,19 +241,19 @@ def random(
     if min_width > bound.width or min_height > bound.height:
       raise ValueError('Minimum width or height is larger than the bound.')
 
-    rand = rand or random
+    rand = rand or random  # pyrefly: ignore[bad-assignment]
 
     max_width = min(max_width, bound.width)
     max_height = min(max_height, bound.height)
 
-    width = rand.randint(min_width, max_width)
-    height = rand.randint(min_height, max_height)
+    width = rand.randint(min_width, max_width)  # pyrefly: ignore[missing-attribute]
+    height = rand.randint(min_height, max_height)  # pyrefly: ignore[missing-attribute]
 
     max_left = bound.right - width
     max_top = bound.bottom - height
 
-    left = rand.randint(bound.left, max_left)
-    top = rand.randint(bound.top, max_top)
+    left = rand.randint(bound.left, max_left)  # pyrefly: ignore[missing-attribute]
+    top = rand.randint(bound.top, max_top)  # pyrefly: ignore[missing-attribute]
 
     right = left + width
     bottom = top + height

langfun/assistant/capabilities/gui/location_test.py

@@ -111,7 +111,7 @@ def test_contains(self):
     self.assertNotIn(location.BBox(0, 0, 400, 500), bbox)
 
     with self.assertRaisesRegex(ValueError, 'Invalid tuple size'):
-      _ = (1, 2, 3) in bbox
+      _ = (1, 2, 3) in bbox  # pyrefly: ignore[unsupported-operation]
 
     with self.assertRaisesRegex(ValueError, 'Invalid type'):
       _ = 'abc' in bbox  # pytype: disable=unsupported-operands

langfun/core/__init__.py

@@ -125,6 +125,7 @@
 from langfun.core.language_model import LMInputError
 from langfun.core.language_model import ContextLimitError
 from langfun.core.language_model import ContentFilteredError
+from langfun.core.language_model import ResponseSizeLimitError
 from langfun.core.language_model import EmptyGenerationError
 from langfun.core.language_model import RetryableLMError
 from langfun.core.language_model import RateLimitError

langfun/core/agentic/action.py

@@ -291,12 +291,12 @@ def __call__(
         # Early terminate the action if the execution time is exceeded.
         session.check_execution_time()
         result = self.call(session=session, **kwargs)
-        self._invocation.end(result)
+        self._invocation.end(result)  # pyrefly: ignore[missing-attribute]
       except BaseException as e:
         error = pg.ErrorInfo.from_exception(e)
-        self._invocation.end(result=None, error=error)
+        self._invocation.end(result=None, error=error)  # pyrefly: ignore[missing-attribute]
         if self._session is not None:
-          self._session.end(result=None, error=error)
+          self._session.end(result=None, error=error)  # pyrefly: ignore[bad-argument-type]
         raise
 
     if self._session is not None:
@@ -393,7 +393,7 @@ def __repr__(self) -> str:
     def __str__(self) -> str:
       return self.to_str()
 
-    def __eq__(self, other: 'ExecutionUnit.Position') -> bool:
+    def __eq__(self, other: 'ExecutionUnit.Position') -> bool:  # pyrefly: ignore[bad-override]
       if isinstance(other, ExecutionUnit.Position):
         return self.indices() == other.indices()
       if isinstance(other, tuple):
@@ -402,7 +402,7 @@ def __eq__(self, other: 'ExecutionUnit.Position') -> bool:
         return str(self) == other
       return False
 
-    def __ne__(self, other: 'ExecutionUnit.Position') -> bool:
+    def __ne__(self, other: 'ExecutionUnit.Position') -> bool:  # pyrefly: ignore[bad-override]
       return not self == other
 
     def __hash__(self) -> int:
@@ -434,10 +434,10 @@ def position(self) -> Position:
     """Returns the execution position of the action."""
     parent_trace = self.sym_ancestor(lambda x: isinstance(x, ExecutionTrace))
     while parent_trace is not None:
-      parent_position = parent_trace.position
+      parent_position = parent_trace.position  # pyrefly: ignore[missing-attribute]
       if parent_position is not None:
         return ExecutionUnit.Position(
-            parent_position, parent_trace.indexof(self, ExecutionUnit)
+            parent_position, parent_trace.indexof(self, ExecutionUnit)  # pyrefly: ignore[missing-attribute]
         )
       parent_trace = parent_trace.sym_ancestor(
           lambda x: isinstance(x, ExecutionTrace)
@@ -666,37 +666,37 @@ def elapse(self) -> float:
   @property
   def queries(self) -> list[lf_structured.QueryInvocation]:
     """Returns queries from the sequence."""
-    return list(self._iter_children(lf_structured.QueryInvocation))
+    return list(self._iter_children(lf_structured.QueryInvocation))  # pyrefly: ignore[bad-return]
 
   @property
   def actions(self) -> list['ActionInvocation']:
     """Returns action invocations from the sequence."""
-    return list(self._iter_children(ActionInvocation))
+    return list(self._iter_children(ActionInvocation))  # pyrefly: ignore[bad-return]
 
   @property
   def execution_units(self) -> list[ExecutionUnit]:
     """Returns parallel executions from the sequence."""
-    return list(self._iter_children(ExecutionUnit))
+    return list(self._iter_children(ExecutionUnit))  # pyrefly: ignore[bad-return]
 
   @property
   def logs(self) -> list[lf.logging.LogEntry]:
     """Returns logs from the sequence."""
-    return list(self._iter_children(lf.logging.LogEntry))
+    return list(self._iter_children(lf.logging.LogEntry))  # pyrefly: ignore[bad-return]
 
   @property
   def all_queries(self) -> list[lf_structured.QueryInvocation]:
     """Returns all queries from current trace and its child execution items."""
-    return list(self._iter_subtree(lf_structured.QueryInvocation))
+    return list(self._iter_subtree(lf_structured.QueryInvocation))  # pyrefly: ignore[bad-return]
 
   @property
   def all_actions(self) -> list['ActionInvocation']:
     """Returns all actions from current trace and its child execution items."""
-    return list(self._iter_subtree(ActionInvocation))
+    return list(self._iter_subtree(ActionInvocation))  # pyrefly: ignore[bad-return]
 
   @property
   def all_logs(self) -> list[lf.logging.LogEntry]:
     """Returns all logs from current trace and its child execution items."""
-    return list(self._iter_subtree(lf.logging.LogEntry))
+    return list(self._iter_subtree(lf.logging.LogEntry))  # pyrefly: ignore[bad-return]
 
   def _iter_children(
       self, item_cls: Type[Any] | tuple[Type[Any], ...]
@@ -775,7 +775,7 @@ def append(self, item: TracedItem) -> None:
       sub_task_label = self._execution_item_label(item)
       self._time_badge.update(
           text=sub_task_label.text,
-          tooltip=sub_task_label.tooltip.content,
+          tooltip=sub_task_label.tooltip.content,  # pyrefly: ignore[missing-attribute]
           add_class=['running'],
           remove_class=['not-started'],
       )
@@ -1244,7 +1244,7 @@ def _on_parent_change(self, *args, **kwargs):
   @property
   def parent_action(self) -> Optional['ActionInvocation']:
     """Returns the parent action invocation."""
-    return self.sym_ancestor(lambda x: isinstance(x, ActionInvocation))
+    return self.sym_ancestor(lambda x: isinstance(x, ActionInvocation))  # pyrefly: ignore[bad-return]
 
   @property
   def max_remaining_execution_time(self) -> float | None:
@@ -1825,7 +1825,7 @@ def _sym_clone(self, deep: bool, memo: Any = None) -> 'Session':
     else:
       event_handler = self.event_handler
     other._event_handler = event_handler  # pylint: disable=protected-access
-    return other
+    return other  # pyrefly: ignore[bad-return]
 
   #
   # Shortcut methods for accessing the root action invocation.
@@ -1992,7 +1992,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         metadata = actions[-1].metadata
       else:
         result, error, metadata = None, None, None
-    self.end(result, error, metadata)
+    self.end(result, error, metadata)  # pyrefly: ignore[bad-argument-type]
 
   #
   # Context-manager for information tracking.
@@ -2224,7 +2224,7 @@ def _map_single(input_value):
         _map_single,
         parallel_inputs,
         max_workers=max_workers,
-        timeout=self._child_max_execution_time(timeout),
+        timeout=self._child_max_execution_time(timeout),  # pyrefly: ignore[bad-argument-type]
         max_duration=max_duration,
         silence_on_errors=silence_on_errors,
         ordered=ordered,

langfun/core/agentic/action_eval.py

@@ -80,7 +80,7 @@ class ActionEvalV1(lf_eval.Matching):
   """
   # We override the schema and prompt to dummy values since they are not used.
   schema_fn = _dummy_schema()
-  prompt = '<unused>'
+  prompt = '<unused>'  # pyrefly: ignore[bad-assignment]
 
   def process(self, example: pg.Dict, **kwargs):
     action = example.action
@@ -130,7 +130,7 @@ def _render_mismatches(self, s: io.StringIO) -> None:
         example_idx for example_idx, *_ in self.mismatches
     ])
     for example_idx in mismatched_ids:
-      url = os.path.join(self.dir, f'example_{example_idx}.html')
+      url = os.path.join(self.dir, f'example_{example_idx}.html')  # pyrefly: ignore[no-matching-overload]
       if first_url is None:
         first_url = url
       s.write(
@@ -152,7 +152,7 @@ def _render_matches(self, s: io.StringIO) -> None:
         example_idx for example_idx, *_ in self.matches
     ])
     for example_idx in matched_ids:
-      url = os.path.join(self.dir, f'example_{example_idx}.html')
+      url = os.path.join(self.dir, f'example_{example_idx}.html')  # pyrefly: ignore[no-matching-overload]
       if first_url is None:
         first_url = url
       s.write(

langfun/core/agentic/action_eval_test.py

@@ -24,6 +24,32 @@
 import pyglove as pg
 
 
+_TIMING_KEYS = ('start_time', 'end_time', 'wall_clock_s')
+
+
+def _strip_timing(result):
+  """Removes non-deterministic per-leaf timing fields for stable eq checks.
+
+  `Evaluation.finalize()` now emits `start_time`/`end_time`/`wall_clock_s` into
+  every leaf result. Those wall-clock values are non-deterministic, so tests
+  that assert the full result dict must drop them before comparing. Their
+  presence and sanity are covered separately by `test_run_timing`.
+
+  Args:
+    result: The leaf result dict (or None) to strip timing fields from.
+
+  Returns:
+    A shallow copy of `result` without the timing fields, or None if `result`
+    is None.
+  """
+  if result is None:
+    return None
+  stripped = dict(result)
+  for key in _TIMING_KEYS:
+    stripped.pop(key, None)
+  return stripped
+
+
 class Foo(action_lib.Action):
   x: int
 
@@ -69,7 +95,7 @@ class FooEval(action_eval.ActionEvalV1):
     s = FooEval()
     result = s.run(summary=False)
     self.assertEqual(
-        result,
+        _strip_timing(result),
         dict(
             experiment_setup=dict(
                 id=s.id,