From b9751314dcfcdf8e53320d66bc2573686228ae52 Mon Sep 17 00:00:00 2001 From: Yifeng Lu Date: Thu, 20 Aug 2026 09:55:17 -0700 Subject: [PATCH] Register the bare model id 'claude-opus-5' to the Vertex AI route. WHY `lf.LanguageModel.get('claude-opus-5')` returns the direct Anthropic API client instead of the Vertex AI one. `vertexai.py` binds Anthropic models by iterating `anthropic.SUPPORTED_MODELS` and registering only those entries whose `provider` is 'VertexAI'. The Claude Opus 5 entry declares `provider='Anthropic'`, so that loop skips it and the bare id stays bound to the direct-API class registered by `anthropic.py`. Callers that select a model purely by bare id therefore get the wrong transport. The neighbouring claude-opus-4-6 / 4-7 / 4-8 ids avoid this because each is listed in an explicit override block; Claude Opus 5 was absent from that block. WHAT One line added to `_register_vertexai_models()`: lf.LanguageModel.register('claude-opus-5', VertexAIClaude5Opus) HOW The line extends the existing override block that encodes the convention "bare model ids resolve as VertexAI (primary use case)". `VertexAIClaude5Opus` is already defined in this module with `model = 'claude-opus-5'`; it was declared but never registered, so no new class, import, dependency or BUILD change is required. Ordering is load-bearing: `_register_vertexai_models()` applies these explicit overrides after the `SUPPORTED_MODELS` loops, so the override replaces the earlier direct-API binding rather than being shadowed by it. Only the bare id is registered. The sibling opus-4-x entries additionally pin an '@latest' alias to the direct API; no such alias is defined for Opus 5, so no corresponding line is added here. PiperOrigin-RevId: 967886387 --- .../gui/bounding_box_parser_test.py | 34 +- .../assistant/capabilities/gui/location.py | 16 +- .../capabilities/gui/location_test.py | 2 +- langfun/core/__init__.py | 1 + langfun/core/agentic/action.py | 38 +- langfun/core/agentic/action_eval.py | 6 +- langfun/core/agentic/action_eval_test.py | 28 +- langfun/core/agentic/action_test.py | 14 +- langfun/core/coding/python/generation.py | 2 +- langfun/core/coding/python/sandboxing.py | 2 +- langfun/core/component.py | 2 +- langfun/core/concurrent.py | 16 +- langfun/core/embedding_model.py | 2 +- langfun/core/ems/openai.py | 6 +- langfun/core/ems/openai_test.py | 6 +- langfun/core/ems/rest.py | 2 +- langfun/core/ems/rest_test.py | 4 +- langfun/core/ems/vertexai.py | 16 +- langfun/core/ems/vertexai_test.py | 8 +- langfun/core/eval/base.py | 168 ++++-- langfun/core/eval/base_test.py | 80 ++- langfun/core/eval/matching.py | 8 +- langfun/core/eval/matching_test.py | 28 +- langfun/core/eval/scoring.py | 6 +- langfun/core/eval/scoring_test.py | 28 +- langfun/core/eval/v2/checkpointing.py | 8 +- langfun/core/eval/v2/eval_test_helper.py | 8 +- langfun/core/eval/v2/evaluation.py | 13 +- langfun/core/eval/v2/example.py | 4 +- langfun/core/eval/v2/experiment.py | 20 +- langfun/core/eval/v2/metrics.py | 14 +- langfun/core/eval/v2/progress.py | 4 +- langfun/core/eval/v2/progress_tracking.py | 20 +- langfun/core/eval/v2/reporting.py | 6 +- langfun/core/eval/v2/runners/base.py | 20 +- langfun/core/eval/v2/runners/beam.py | 4 +- langfun/core/eval/v2/runners/ckpt_monitor.py | 18 +- langfun/core/eval/v2/runners/parallel.py | 4 +- langfun/core/langfunc.py | 6 +- langfun/core/language_model.py | 29 +- langfun/core/llms/__init__.py | 9 + langfun/core/llms/anthropic.py | 382 ++++++++++++- langfun/core/llms/anthropic_test.py | 516 ++++++++++++++++++ langfun/core/llms/azure_openai.py | 2 +- langfun/core/llms/cache/base.py | 2 +- langfun/core/llms/cache/in_memory.py | 2 +- langfun/core/llms/compositional.py | 6 +- langfun/core/llms/deepseek.py | 4 +- langfun/core/llms/fake.py | 2 +- langfun/core/llms/gemini.py | 100 +++- langfun/core/llms/google_genai.py | 18 +- langfun/core/llms/google_genai_test.py | 14 + langfun/core/llms/groq.py | 2 +- langfun/core/llms/llama_cpp.py | 4 +- langfun/core/llms/openai.py | 4 +- langfun/core/llms/openai_compatible.py | 8 +- langfun/core/llms/rest.py | 172 +++++- langfun/core/llms/rest_test.py | 228 ++++++++ langfun/core/llms/veo.py | 8 +- langfun/core/llms/vertexai.py | 47 +- langfun/core/llms/vertexai_test.py | 24 +- langfun/core/mcp/tool.py | 6 +- langfun/core/mcp/tool_test.py | 4 +- langfun/core/message.py | 20 +- langfun/core/modalities/__init__.py | 6 +- langfun/core/modalities/image.py | 4 +- langfun/core/modalities/mime.py | 12 +- langfun/core/modality.py | 4 +- langfun/core/natural_language.py | 2 +- langfun/core/sampling.py | 2 +- langfun/core/structured/completion.py | 2 +- langfun/core/structured/description.py | 1 + .../core/structured/function_generation.py | 11 +- langfun/core/structured/mapping.py | 6 +- langfun/core/structured/parsing.py | 12 +- langfun/core/structured/querying.py | 20 +- langfun/core/structured/schema/base.py | 8 +- langfun/core/structured/schema/json.py | 2 +- langfun/core/structured/schema_generation.py | 2 +- langfun/core/subscription.py | 8 +- langfun/core/template.py | 20 +- langfun/core/templates/completion.py | 2 +- langfun/core/templates/conversation.py | 1 + langfun/core/templates/selfplay.py | 2 +- langfun/env/base_feature.py | 10 +- langfun/env/base_sandbox.py | 4 +- langfun/env/base_sandbox_service.py | 20 +- langfun/env/environment.py | 12 +- langfun/env/event_handlers/event_logger.py | 4 +- langfun/env/event_handlers/metric_writer.py | 108 ++-- langfun/env/interface.py | 26 +- langfun/env/test_utils.py | 18 +- requirements.txt | 2 +- 93 files changed, 2102 insertions(+), 514 deletions(-) diff --git a/langfun/assistant/capabilities/gui/bounding_box_parser_test.py b/langfun/assistant/capabilities/gui/bounding_box_parser_test.py index 93ff50a1..244e9ba6 100644 --- a/langfun/assistant/capabilities/gui/bounding_box_parser_test.py +++ b/langfun/assistant/capabilities/gui/bounding_box_parser_test.py @@ -61,10 +61,10 @@ 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]}' @@ -72,7 +72,7 @@ def test_simple_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_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,7 +93,7 @@ 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```' @@ -101,7 +101,7 @@ def test_json_in_code_block(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_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,13 +218,13 @@ 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]}' @@ -232,7 +232,7 @@ def test_float_numbers(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_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,7 +299,7 @@ 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 = ( @@ -307,7 +307,7 @@ def test_deeply_nested_json(self): ) 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() diff --git a/langfun/assistant/capabilities/gui/location.py b/langfun/assistant/capabilities/gui/location.py index 72a0dc10..ebd58ac6 100644 --- a/langfun/assistant/capabilities/gui/location.py +++ b/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 diff --git a/langfun/assistant/capabilities/gui/location_test.py b/langfun/assistant/capabilities/gui/location_test.py index f6368641..6f4a599e 100644 --- a/langfun/assistant/capabilities/gui/location_test.py +++ b/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 diff --git a/langfun/core/__init__.py b/langfun/core/__init__.py index dc0e0076..020e6260 100644 --- a/langfun/core/__init__.py +++ b/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 diff --git a/langfun/core/agentic/action.py b/langfun/core/agentic/action.py index 6b9f6cd0..04493c38 100644 --- a/langfun/core/agentic/action.py +++ b/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, diff --git a/langfun/core/agentic/action_eval.py b/langfun/core/agentic/action_eval.py index 905a043e..623b01ee 100644 --- a/langfun/core/agentic/action_eval.py +++ b/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 = '' + prompt = '' # 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( diff --git a/langfun/core/agentic/action_eval_test.py b/langfun/core/agentic/action_eval_test.py index 75ffc129..3352672e 100644 --- a/langfun/core/agentic/action_eval_test.py +++ b/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, diff --git a/langfun/core/agentic/action_test.py b/langfun/core/agentic/action_test.py index 64268320..76f3c936 100644 --- a/langfun/core/agentic/action_test.py +++ b/langfun/core/agentic/action_test.py @@ -29,7 +29,7 @@ class Bar(action_lib.Action): simulate_action_error: bool = False simulate_execution_time: float = 0 - def call(self, session, *, lm, **kwargs): + def call(self, session, *, lm, **kwargs): # pyrefly: ignore[bad-override] assert session.current_action.action is self session.info('Begin Bar') time.sleep(self.simulate_execution_time) @@ -48,7 +48,7 @@ class Foo(action_lib.Action): simulate_execution_time: list[float] = [0, 0, 0, 0] max_bar_execution_time: float | None = None - def call(self, session, *, lm, **kwargs): + def call(self, session, *, lm, **kwargs): # pyrefly: ignore[bad-override] assert session.current_action.action is self with session.track_phase('prepare'): session.info('Begin Foo', x=1) @@ -190,7 +190,7 @@ def test_succeeded_trajectory(self): self.assertEqual(result, 3) self.assertIsNone(foo.session) self.assertEqual(foo.state, [0, 1, 2]) - self.assertIs(foo.invocation.state, foo.state) + self.assertIs(foo.invocation.state, foo.state) # pyrefly: ignore[missing-attribute] self.assertEqual(foo.result, 3) self.assertEqual( foo.metadata, dict(note='foo', subtask_0=0, subtask_1=1, subtask_2=2) @@ -615,11 +615,11 @@ def test_clone(self): session = action_lib.Session(event_handler=event_handler) other = session.clone() self.assertIsNot(session, other) - self.assertIs(other.event_handler, event_handler) + self.assertIs(other.event_handler, event_handler) # pyrefly: ignore[missing-attribute] other = session.clone(deep=True) self.assertIsNot(session, other) - self.assertIsNot(other.event_handler, session.event_handler) + self.assertIsNot(other.event_handler, session.event_handler) # pyrefly: ignore[missing-attribute] def test_log(self): session = action_lib.Session() @@ -631,7 +631,7 @@ def test_log(self): def test_as_message(self): session = action_lib.Session() - self.assertIn('agent@', session.id) + self.assertIn('agent@', session.id) # pyrefly: ignore[bad-argument-type] self.assertIsInstance(session.as_message(), lf.AIMessage) def test_query_with_track_if(self): @@ -663,7 +663,7 @@ def test_tls_preserved_on_rebinding_in_new_thread(self): """Regression test: PyGlove re-binding must not wipe thread-local state.""" class Dummy(pg.Object): - session: pg.typing.Any() + session: pg.typing.Any() # pyrefly: ignore[invalid-annotation] class DummyAction(action_lib.Action): def call(self, session, **kwargs): diff --git a/langfun/core/coding/python/generation.py b/langfun/core/coding/python/generation.py index 63695043..81b22721 100644 --- a/langfun/core/coding/python/generation.py +++ b/langfun/core/coding/python/generation.py @@ -205,7 +205,7 @@ def _on_bound(self): @functools.cached_property def implementation(self) -> Callable[..., Any]: """Returns the function implementation based on source code.""" - return execution.run(self.source) + return execution.run(self.source) # pyrefly: ignore[bad-return] def __call__( self, diff --git a/langfun/core/coding/python/sandboxing.py b/langfun/core/coding/python/sandboxing.py index 92675809..3d344138 100644 --- a/langfun/core/coding/python/sandboxing.py +++ b/langfun/core/coding/python/sandboxing.py @@ -140,7 +140,7 @@ def _on_bound(self): @property def working_dir(self) -> str | None: """Returns the directory of the sandbox.""" - return self._working_dir + return self._working_dir # pyrefly: ignore[bad-return] def _setup(self) -> None: """Sets up the sandbox.""" diff --git a/langfun/core/component.py b/langfun/core/component.py index cb101b3a..1e421ebd 100644 --- a/langfun/core/component.py +++ b/langfun/core/component.py @@ -77,7 +77,7 @@ def __init_subclass__(cls): continue attr_value = getattr(cls, attr_name) if isinstance(attr_value, pg.Inferentiable): - value_spec = pg.typing.Any() + value_spec = pg.typing.Any() # pyrefly: ignore[bad-instantiation] elif isinstance(attr_value, Component): value_spec = pg.typing.Object(Component) else: diff --git a/langfun/core/concurrent.py b/langfun/core/concurrent.py index 0ff8484a..b9fab4c7 100644 --- a/langfun/core/concurrent.py +++ b/langfun/core/concurrent.py @@ -81,14 +81,14 @@ def __str__(self) -> str: f'Last error: {self.errors[-1]}' ) - def __eq__(self, other: 'RetryError') -> bool: + def __eq__(self, other: 'RetryError') -> bool: # pyrefly: ignore[bad-override] if not isinstance(other, RetryError): return False return (self.func is other.func and self.errors == other.errors and self.wait_intervals == other.wait_intervals) - def __ne__(self, other: 'RetryError') -> bool: + def __ne__(self, other: 'RetryError') -> bool: # pyrefly: ignore[bad-override] return not self.__eq__(other) def __hash__(self) -> int: @@ -151,7 +151,7 @@ def _func(*args, **kwargs): func, args, kwargs, - retry_on_errors=retry_on_errors, + retry_on_errors=retry_on_errors, # pyrefly: ignore[bad-argument-type] max_attempts=max_attempts, retry_interval=retry_interval, exponential_backoff=exponential_backoff, @@ -249,7 +249,7 @@ def square(x): Job( func, (inputs,), - retry_on_errors=retry_on_errors, + retry_on_errors=retry_on_errors, # pyrefly: ignore[bad-argument-type] max_attempts=max_attempts, retry_interval=retry_interval, exponential_backoff=exponential_backoff, @@ -380,7 +380,7 @@ def next_wait_interval(attempt: int) -> float: retry_entries = [] wait_interval = 0 while True: - with pg.catch_errors(self.retry_on_errors) as error_context: + with pg.catch_errors(self.retry_on_errors) as error_context: # pyrefly: ignore[bad-argument-type] begin_time = time.time() self.result = self.func(*self.args, **self.kwargs) @@ -799,7 +799,7 @@ def flaky_square(x): job = Job( func, (inputs,), - retry_on_errors=retry_on_errors, + retry_on_errors=retry_on_errors, # pyrefly: ignore[bad-argument-type] max_attempts=max_attempts, retry_interval=retry_interval, exponential_backoff=exponential_backoff, @@ -830,7 +830,7 @@ def update_progress_bar(progress: Progress) -> None: if progress.timeit_summary: status['TimeIt'] = progress.timeit_summary_str() - ProgressBar.update(bar_id, delta=1, status=status) + ProgressBar.update(bar_id, delta=1, status=status) # pyrefly: ignore[bad-argument-type] deadline = time.time() + max_duration if max_duration else None @@ -943,7 +943,7 @@ def update_progress_bar(progress: Progress) -> None: ProgressBar.refresh() finally: if show_progress and not external_bar: - ProgressBar.uninstall(bar_id) + ProgressBar.uninstall(bar_id) # pyrefly: ignore[bad-argument-type] if shutdown_after_finish: executor.shutdown( diff --git a/langfun/core/embedding_model.py b/langfun/core/embedding_model.py index 639a0d2e..8733016f 100644 --- a/langfun/core/embedding_model.py +++ b/langfun/core/embedding_model.py @@ -162,7 +162,7 @@ class EmbeddingModel(component.Component): _MODEL_FACTORY: ClassVar[dict[str, Callable[..., 'EmbeddingModel']]] = {} @classmethod - def register( + def register( # pyrefly: ignore[bad-override] cls, model_id_or_prefix: str, factory: Callable[..., 'EmbeddingModel'] ) -> None: diff --git a/langfun/core/ems/openai.py b/langfun/core/ems/openai.py index 3dba56b8..46ef8438 100644 --- a/langfun/core/ems/openai.py +++ b/langfun/core/ems/openai.py @@ -89,11 +89,11 @@ def _initialize(self) -> None: ) @property - def api_endpoint(self) -> str: + def api_endpoint(self) -> str: # pyrefly: ignore[bad-override] return 'https://api.openai.com/v1/embeddings' @property - def headers(self) -> dict[str, Any]: + def headers(self) -> dict[str, Any]: # pyrefly: ignore[bad-override] assert self._api_initialized headers = { 'Content-Type': 'application/json', @@ -115,7 +115,7 @@ def request(self, message: lf.Message) -> dict[str, Any]: } options = self.embedding_options if options.output_dimensionality is not None: - request_body['dimensions'] = options.output_dimensionality + request_body['dimensions'] = options.output_dimensionality # pyrefly: ignore[bad-assignment] return request_body def result(self, json_response: dict[str, Any]) -> lf.EmbeddingResult: diff --git a/langfun/core/ems/openai_test.py b/langfun/core/ems/openai_test.py index 0ccd027e..d6a68bbe 100644 --- a/langfun/core/ems/openai_test.py +++ b/langfun/core/ems/openai_test.py @@ -113,15 +113,15 @@ def test_result_parsing(self): def test_text_embedding_3_small_defaults(self): model = openai.TextEmbedding3Small.__schema__.get_field('model') - self.assertEqual(model.default_value, 'text-embedding-3-small') + self.assertEqual(model.default_value, 'text-embedding-3-small') # pyrefly: ignore[missing-attribute] def test_text_embedding_3_large_defaults(self): model = openai.TextEmbedding3Large.__schema__.get_field('model') - self.assertEqual(model.default_value, 'text-embedding-3-large') + self.assertEqual(model.default_value, 'text-embedding-3-large') # pyrefly: ignore[missing-attribute] def test_text_embedding_ada_002_defaults(self): model = openai.TextEmbeddingAda002.__schema__.get_field('model') - self.assertEqual(model.default_value, 'text-embedding-ada-002') + self.assertEqual(model.default_value, 'text-embedding-ada-002') # pyrefly: ignore[missing-attribute] if __name__ == '__main__': diff --git a/langfun/core/ems/rest.py b/langfun/core/ems/rest.py index b2e2eca3..a876fabd 100644 --- a/langfun/core/ems/rest.py +++ b/langfun/core/ems/rest.py @@ -142,4 +142,4 @@ def _parse_response( if response.status_code == 200: return self.result(response.json()) else: - raise self._error(response.status_code, response.content) + raise self._error(response.status_code, response.content) # pyrefly: ignore[bad-argument-type] diff --git a/langfun/core/ems/rest_test.py b/langfun/core/ems/rest_test.py index 55e70078..6a7f66a1 100644 --- a/langfun/core/ems/rest_test.py +++ b/langfun/core/ems/rest_test.py @@ -27,7 +27,7 @@ def mock_requests_post(url: str, json: dict[str, Any], **kwargs): del url, kwargs response = requests.Response() response.status_code = 200 - response._content = pg.to_json_str({ + response._content = pg.to_json_str({ # pyrefly: ignore[bad-assignment] 'embedding': [0.1, 0.2, 0.3], 'input_text': json.get('text', ''), }).encode() @@ -39,7 +39,7 @@ def _mock_requests(url: str, json: dict[str, Any], **kwargs): del url, json, kwargs response = requests.Response() response.status_code = status_code - response._content = b'error' + response._content = b'error' # pyrefly: ignore[bad-assignment] return response return _mock_requests diff --git a/langfun/core/ems/vertexai.py b/langfun/core/ems/vertexai.py index 42c51196..513c4520 100644 --- a/langfun/core/ems/vertexai.py +++ b/langfun/core/ems/vertexai.py @@ -130,9 +130,9 @@ def _initialize(self): @property def _project(self) -> str: """Returns a project ID. Randomly selects from list if multiple.""" - if len(self._projects) == 1: - return self._projects[0] - return random.choice(self._projects) + if len(self._projects) == 1: # pyrefly: ignore[bad-argument-type] + return self._projects[0] # pyrefly: ignore[unsupported-operation] + return random.choice(self._projects) # pyrefly: ignore[bad-argument-type] def session(self): assert self._api_initialized @@ -143,7 +143,7 @@ def session(self): return s @property - def api_endpoint(self) -> str: + def api_endpoint(self) -> str: # pyrefly: ignore[bad-override] assert self._api_initialized project = self._project return ( @@ -160,7 +160,7 @@ def request(self, message: lf.Message) -> dict[str, Any]: import base64 # pylint: disable=g-import-not-at-top parts.append({ - 'inline_data': { + 'inline_data': { # pyrefly: ignore[bad-assignment] 'mime_type': 'image/png', 'data': base64.b64encode(modality.to_bytes()).decode('utf-8'), } @@ -170,9 +170,9 @@ def request(self, message: lf.Message) -> dict[str, Any]: request_body = {'content': {'parts': parts}} options = self.embedding_options if options.task_type is not None: - request_body['taskType'] = options.task_type + request_body['taskType'] = options.task_type # pyrefly: ignore[bad-assignment] if options.output_dimensionality is not None: - request_body['outputDimensionality'] = options.output_dimensionality + request_body['outputDimensionality'] = options.output_dimensionality # pyrefly: ignore[bad-assignment] return request_body def result(self, json_response: dict[str, Any]) -> lf.EmbeddingResult: @@ -219,7 +219,7 @@ def request(self, message: lf.Message) -> dict[str, Any]: if options.output_dimensionality is not None: parameters['outputDimensionality'] = options.output_dimensionality if parameters: - request_body['parameters'] = parameters + request_body['parameters'] = parameters # pyrefly: ignore[bad-assignment] return request_body def result(self, json_response: dict[str, Any]) -> lf.EmbeddingResult: diff --git a/langfun/core/ems/vertexai_test.py b/langfun/core/ems/vertexai_test.py index 2ced18b7..ee3edd5c 100644 --- a/langfun/core/ems/vertexai_test.py +++ b/langfun/core/ems/vertexai_test.py @@ -148,7 +148,7 @@ def test_result_parsing_with_embedding_key(self): def test_gemini_embedding2_defaults(self): model = vertexai.VertexAIGeminiEmbedding2.__schema__.get_field('model') - self.assertEqual(model.default_value, 'gemini-embedding-2-preview') + self.assertEqual(model.default_value, 'gemini-embedding-2-preview') # pyrefly: ignore[missing-attribute] @mock.patch.object(vertexai.VertexAI, 'credentials', new=True) def test_embedding_options_passthrough(self): @@ -236,17 +236,17 @@ def test_predict_result_parsing(self): def test_gemini_embedding1_defaults(self): model = vertexai.VertexAIGeminiEmbedding1.__schema__.get_field('model') - self.assertEqual(model.default_value, 'gemini-embedding-001') + self.assertEqual(model.default_value, 'gemini-embedding-001') # pyrefly: ignore[missing-attribute] def test_text_embedding_005_defaults(self): model = vertexai.VertexAITextEmbedding005.__schema__.get_field('model') - self.assertEqual(model.default_value, 'text-embedding-005') + self.assertEqual(model.default_value, 'text-embedding-005') # pyrefly: ignore[missing-attribute] def test_text_multilingual_embedding_002_defaults(self): model = vertexai.VertexAITextMultilingualEmbedding002.__schema__.get_field( 'model' ) - self.assertEqual(model.default_value, 'text-multilingual-embedding-002') + self.assertEqual(model.default_value, 'text-multilingual-embedding-002') # pyrefly: ignore[missing-attribute] class VertexAIGoogleAuthMissingTest(unittest.TestCase): diff --git a/langfun/core/eval/base.py b/langfun/core/eval/base.py index b9dc226d..a76bc32c 100644 --- a/langfun/core/eval/base.py +++ b/langfun/core/eval/base.py @@ -220,6 +220,12 @@ def run( pivot_field: str = 'lm', from_root: bool = True, timeout: int | None = None, + # Optional cap on how many leaf evaluations run concurrently. Defaults to + # None => unbounded (one worker per leaf), preserving current behavior. + max_leaf_concurrency: int | None = None, + # Minimum seconds between throttled intermediate summary.save() calls off + # the hot per-leaf path. The final summary flush is always unconditional. + summary_save_interval_s: float = 10.0, **kwargs, ) -> Union['Summary', pg.Dict]: """Run the evaluation, which fills and returns the result.""" @@ -228,7 +234,7 @@ def run( if dryrun: self.dryrun(filter=filter, verbose=False, debug=debug) - summary = self.summary(pivot_field) if from_root and summary else None + summary = self.summary(pivot_field) if from_root and summary else None # pyrefly: ignore[bad-assignment] should_save = bool(save and self.dir) if self.is_leaf: @@ -247,7 +253,7 @@ def run( ): if show_progress: lf.concurrent.ProgressBar.update( - progress_bar, status='LOADING SAVED RESULTS...', color='yellow' + progress_bar, status='LOADING SAVED RESULTS...', color='yellow' # pyrefly: ignore[bad-argument-type] ) if self.try_load_result(): run_status = 'CACHED' @@ -259,12 +265,12 @@ def run( if self.result: if show_progress: lf.concurrent.ProgressBar.update( - progress_bar, delta=self.num_examples + progress_bar, delta=self.num_examples # pyrefly: ignore[bad-argument-type] ) else: self._run( start=start, - end=end, + end=end, # pyrefly: ignore[bad-argument-type] debug=debug, dryrun=dryrun, verbose=verbose, @@ -277,7 +283,7 @@ def run( if should_save: if show_progress: lf.concurrent.ProgressBar.update( - progress_bar, status='SAVING RESULTS...', color='yellow' + progress_bar, status='SAVING RESULTS...', color='yellow' # pyrefly: ignore[bad-argument-type] ) # Save evaluation results. @@ -285,17 +291,24 @@ def run( # Save summary if present. if summary: - summary.save(os.path.join(self.root_dir, Evaluable.SUMMARY_HTML)) + summary.save(os.path.join(self.root_dir, Evaluable.SUMMARY_HTML)) # pyrefly: ignore[missing-attribute, no-matching-overload] if show_progress: lf.concurrent.ProgressBar.update( - progress_bar, + progress_bar, # pyrefly: ignore[bad-argument-type] status=self._completion_status(run_status), color='green', ) else: assert from_root summary_lock = threading.Lock() + # Debounce: throttle the O(all-leaves) summary re-render/save off the hot + # per-leaf completion path to at most one write per + # summary_save_interval_s seconds. The unconditional final flush below + # still guarantees a complete summary at the end of the run. + # `last_summary_save` is a 1-element list so the nested closure can mutate + # it under `summary_lock`. + last_summary_save = [0.0] def _run_group(arg: tuple[int, list[_LeafNode]]) -> None: overview_bar, leaf_group = arg for leaf in leaf_group: @@ -308,33 +321,42 @@ def _run_group(arg: tuple[int, list[_LeafNode]]) -> None: debug=debug, dryrun=False, verbose=verbose, - show_progress=leaf.progress_bar, + show_progress=leaf.progress_bar, # pyrefly: ignore[bad-argument-type] summary=False, from_root=False, **kwargs, ) if should_save and summary: + now = time.time() with summary_lock: - summary.save( - os.path.join(self.root_dir, Evaluable.SUMMARY_HTML) - ) + # Only re-render/save the (expensive, whole-tree) summary if + # enough time has elapsed since the last write. This keeps the + # per-leaf hot path O(1) instead of O(all-leaves). + if now - last_summary_save[0] >= summary_save_interval_s: + summary.save( # pyrefly: ignore[missing-attribute] + os.path.join(self.root_dir, Evaluable.SUMMARY_HTML) # pyrefly: ignore[no-matching-overload] + ) + last_summary_save[0] = now # Signal sub-eval complete by setting the color green. - lf.concurrent.ProgressBar.uninstall(leaf.progress_bar) + lf.concurrent.ProgressBar.uninstall(leaf.progress_bar) # pyrefly: ignore[bad-argument-type] lf.concurrent.ProgressBar.update(overview_bar, 1, { 'LastCompleted': leaf.node.id }) - # NOTE(daiyip): Run leaf nodes grouped by model resource id. This allows - # evaluations using the same resource to run sequentially, which favors - # completing evaluations over running evaluations sparsely. + # NOTE: Fan out all independent leaf evaluations concurrently. Each leaf + # is placed in its own group (keyed by its unique leaf id) and the run is + # dispatched with max_workers == number of leaves, so every leaf can be + # in flight at once. Per-model backpressure is still enforced by the + # LM-level rate-limit semaphore (Layer C, keyed by resource_id), which + # remains the only cap on concurrent requests to a given model. filter = filter or (lambda x: True) leaf_nodes: list[_LeafNode] = [] leaf_groups: dict[str, list[_LeafNode]] = collections.defaultdict(list) for i, leaf in enumerate(self.leaf_nodes): node = _LeafNode(index=i + 1, node=leaf, enabled=filter(leaf)) - leaf_groups[leaf.lm.resource_id].append(node) + leaf_groups[leaf.id].append(node) leaf_nodes.append(node) if leaf_groups: @@ -346,7 +368,7 @@ def _run_group(arg: tuple[int, list[_LeafNode]]) -> None: f'[#{leaf.index} - {leaf.node.id}]', total=leaf.node.num_examples if leaf.enabled else 0, color='cyan' if leaf.enabled else 'yellow', - status=None if leaf.enabled else 'SKIPPED.') + status=None if leaf.enabled else 'SKIPPED.') # pyrefly: ignore[bad-argument-type] # Run leaf groups in parallel. try: @@ -354,7 +376,10 @@ def _run_group(arg: tuple[int, list[_LeafNode]]) -> None: _run_group, [(overview_bar, group) for group in leaf_groups.values()], silence_on_errors=None, - max_workers=len(leaf_groups)): + # Fan out all leaves at once by default (max_workers == number of + # leaves); `max_leaf_concurrency`, when set, caps this. + max_workers=max_leaf_concurrency or len(leaf_nodes), + ): pass # Save results for non-leaf nodes. @@ -373,7 +398,7 @@ def _run_group(arg: tuple[int, list[_LeafNode]]) -> None: overview_bar, status='FINALIZING SUMMARY...' ) - summary.save(os.path.join(self.root_dir, Evaluable.SUMMARY_HTML)) + summary.save(os.path.join(self.root_dir, Evaluable.SUMMARY_HTML)) # pyrefly: ignore[missing-attribute, no-matching-overload] lf.console.write( f'({self.summary_link})', @@ -391,7 +416,7 @@ def _run_group(arg: tuple[int, list[_LeafNode]]) -> None: # for leaf in leaf_nodes: # lf.concurrent.ProgressBar.uninstall(leaf.progress_bar) lf.concurrent.ProgressBar.uninstall(overview_bar) - return summary or self.result + return summary or self.result # pyrefly: ignore[bad-return] @abc.abstractmethod def _run( @@ -439,11 +464,11 @@ def save( ) -> None: # Save experiment definition. if definition: - pg.save(self, os.path.join(self.dir, Evaluable.EXPERIMENT_JSON)) + pg.save(self, os.path.join(self.dir, Evaluable.EXPERIMENT_JSON)) # pyrefly: ignore[no-matching-overload] # Save evaluation result. if result: - pg.save(self.result, os.path.join(self.dir, Evaluation.RESULT_JSON)) + pg.save(self.result, os.path.join(self.dir, Evaluation.RESULT_JSON)) # pyrefly: ignore[no-matching-overload] def _html( self, @@ -470,6 +495,7 @@ def _html( ) if include_cache_stats and self.is_deterministic: s.write( + # pyrefly: ignore[missing-attribute] '

Cache Stats

' '
{self.result.cache_stats}
' @@ -493,7 +519,7 @@ def _render_navbar(self, s: io.StringIO) -> None: if i != len(links) - 1: # Add a right triangle symbol. s.write(' ▸ ') - s.write(f' [Directory]') + s.write(f' [Directory]') # pyrefly: ignore[bad-argument-type] def _render_index_page(self, s: io.StringIO) -> None: self._render_result(s) @@ -532,7 +558,7 @@ def _render_result_row(self, s: io.StringIO) -> None: def _render_dryrun_output(self, s: io.StringIO) -> None: s.write('

Dry Run

') - self._render_message(self.dryrun_output, s) + self._render_message(self.dryrun_output, s) # pyrefly: ignore[bad-argument-type] def _render_message(self, message: lf.Message, s: io.StringIO) -> None: s.write( @@ -564,7 +590,7 @@ def from_dir( def try_load_result(self) -> bool: """Try loads result from file if it's not loaded.""" if self.result is None: - result_json = os.path.join(self.dir, Evaluable.RESULT_JSON) + result_json = os.path.join(self.dir, Evaluable.RESULT_JSON) # pyrefly: ignore[no-matching-overload] if pg.io.path_exists(result_json): self._result = pg.load(result_json) return True @@ -603,7 +629,7 @@ def _on_bound(self): if k not in ('id', 'children') } for child in self.children: - child.rebind(overrides, notify_parents=False) + child.rebind(overrides, notify_parents=False) # pyrefly: ignore[bad-argument-type] self.__dict__.pop('hash', None) @functools.cached_property @@ -825,7 +851,7 @@ def failure_rate(self) -> float: @functools.cached_property def oop_failures(self) -> list[tuple[Any, lf_structured.MappingError]]: """Returns the OOP failures.""" - return [item for item in self.failures + return [item for item in self.failures # pyrefly: ignore[bad-return] if isinstance(item[1], lf_structured.MappingError)] @property @@ -1042,7 +1068,7 @@ def _dryrun( **kwargs, ) -> None: # We make a copy to avoid pollute the state of current object. - copy: Evaluation = self.clone() + copy: Evaluation = self.clone() # pyrefly: ignore[bad-assignment] # Set the example for dryrun. example = example or copy.examples[0] @@ -1113,6 +1139,11 @@ def _run( timeout: int | None = None, **kwargs, ) -> None: + # Capture per-leaf wall-clock start (epoch seconds). Persisted via + # finalize() so downstream makespan / longest-trajectory analysis has real + # measured per-leaf timing (previously the harness persisted none). + self._start_time = time.time() + # Setup examples. # Reset examples so it could be read from the input functor. self.__dict__.pop('examples', None) @@ -1156,7 +1187,8 @@ def _process(idx_and_example: Any): if self.dir and self.cache: self.cache.save() - # Summarize result. + # Capture per-leaf wall-clock end, then summarize result. + self._end_time = time.time() self._result = self.finalize() if verbose: lf.console.write( @@ -1198,7 +1230,7 @@ def process(self, example: Any, **kwargs) -> lf.Message: ) else: assert self.method == 'complete', self.method - assert isinstance(self.schema.spec, pg.typing.Object), self.schema + assert isinstance(self.schema.spec, pg.typing.Object), self.schema # pyrefly: ignore[missing-attribute] # TODO(daiyip): Currently multi-modal inputs within the prompt for # completion is not supported. input_value = self.schema.spec.cls.partial(prompt.render().text) @@ -1240,9 +1272,9 @@ def _status(self, progress: lf.concurrent.Progress) -> dict[str, Any]: status.update(self._eval_status(progress)) if progress.last_error is not None: - status['LastError'] = progress.last_error_str() + status['LastError'] = progress.last_error_str() # pyrefly: ignore[bad-assignment] if progress.timeit_summary: - status['TimeIt'] = progress.timeit_summary_str() + status['TimeIt'] = progress.timeit_summary_str() # pyrefly: ignore[bad-assignment] return status def _eval_status(self, progress: lf.concurrent.Progress) -> dict[str, Any]: @@ -1299,6 +1331,17 @@ def finalize(self) -> pg.Dict: else: usage = None + # Per-leaf wall-clock timing (epoch seconds), populated by Evaluation._run. + # Values are None when finalize() runs outside a live run (e.g. a legacy + # result.json predating timing instrumentation), so loading old results is + # safe. + start_time = getattr(self, '_start_time', None) + end_time = getattr(self, '_end_time', None) + if start_time is not None and end_time is not None: + wall_clock_s = end_time - start_time + else: + wall_clock_s = None + result = pg.Dict( experiment_setup=pg.Dict( id=self.id, @@ -1320,6 +1363,9 @@ def finalize(self) -> pg.Dict: failure_breakdown=self.failure_breakdown, ), usage=usage, + start_time=start_time, + end_time=end_time, + wall_clock_s=wall_clock_s, ) return result @@ -1333,7 +1379,7 @@ def summary_card(self) -> str: definition, self.hash, '', - lambda: self.link(self.dir), + lambda: self.link(self.dir), # pyrefly: ignore[bad-argument-type] ) if self.result is None: s.write( @@ -1355,7 +1401,7 @@ def summary_card(self) -> str: def _render_summary_usage(self, s: io.StringIO) -> None: """Renders usage in HTML.""" - usage = self.result.usage + usage = self.result.usage # pyrefly: ignore[missing-attribute] total = usage.total_prompt_tokens + usage.total_completion_tokens s.write( ' ' 'Error typeStats' ) - error_regex = re.compile(error_regex) - if self.result.metrics.failure_breakdown: - for name, count in self.result.metrics.failure_breakdown.items(): - if not error_regex.match(name): + error_regex = re.compile(error_regex) # pyrefly: ignore[bad-assignment] + if self.result.metrics.failure_breakdown: # pyrefly: ignore[missing-attribute] + for name, count in self.result.metrics.failure_breakdown.items(): # pyrefly: ignore[missing-attribute] + if not error_regex.match(name): # pyrefly: ignore[missing-attribute] continue link = f'{name}' - error_rate = self._format_rate(count / self.result.metrics.total) + error_rate = self._format_rate(count / self.result.metrics.total) # pyrefly: ignore[missing-attribute] + # pyrefly: ignore[missing-attribute] stats = (f'{error_rate} ' f'({count}/{self.result.metrics.total})') s.write(f'{link}{stats})') @@ -1647,7 +1694,7 @@ def _render_failures( failures_by_error = collections.defaultdict(list) for example, error in self.failures: error_name = _error_key(error) - if error_regex.match(error_name): + if error_regex.match(error_name): # pyrefly: ignore[missing-attribute] failures_by_error[error_name].append((example, error)) for error_key, failures in failures_by_error.items(): @@ -1696,10 +1743,10 @@ def inputs_from(path: str | list[str], **kwargs) -> list[Any]: import pandas as pd # pylint: disable=g-import-not-at-top dataset_df = pd.read_csv(path, **kwargs) dataset = [] - for i in range(dataset_df.shape[0]): + for i in range(dataset_df.shape[0]): # pyrefly: ignore[missing-attribute] row = {} - for col in dataset_df.columns: - row[col] = dataset_df.iloc[i][col] + for col in dataset_df.columns: # pyrefly: ignore[missing-attribute] + row[col] = dataset_df.iloc[i][col] # pyrefly: ignore[missing-attribute] dataset.append(row) return dataset else: @@ -1860,7 +1907,7 @@ def _repr_html_(self) -> str: @classmethod def from_evaluations( - cls, evaluations: list['Summary.Entry'], pivot_field: str = 'lm' + cls, evaluations: list['Summary.Entry'], pivot_field: str = 'lm' # pyrefly: ignore[missing-attribute] ) -> 'Summary.Table': """Creates a table from a list of evaluations.""" @@ -1966,6 +2013,15 @@ def json( dir=entry.dir, metrics=entry.result.metrics if entry.result else None, usage=entry.result.usage if entry.result else None, + start_time=( + entry.result.get('start_time') if entry.result else None + ), + end_time=( + entry.result.get('end_time') if entry.result else None + ), + wall_clock_s=( + entry.result.get('wall_clock_s') if entry.result else None + ), ) ) task_results[task.__name__] = results @@ -1992,7 +2048,7 @@ def from_dirs( for x in lf.concurrent_execute( Evaluable.from_dir, [ - os.path.join(root_dir, i) + os.path.join(root_dir, i) # pyrefly: ignore[no-matching-overload] for i in _iter_dirs(root_dir, filter) ], ) @@ -2124,7 +2180,7 @@ def _error_key(error: Exception) -> str: error_names = [] while error is not None: error_names.append(error.__class__.__name__) - error = getattr(error, 'cause', None) + error = getattr(error, 'cause', None) # pyrefly: ignore[bad-assignment] return '.'.join(error_names) @@ -2329,7 +2385,7 @@ def get( suite = Suite(matches, root_dir=root_dir) if patches: - suite = pg.patch(suite, patches) + suite = pg.patch(suite, patches) # pyrefly: ignore[bad-argument-type] if isinstance(filter, str): regex = re.compile(filter) diff --git a/langfun/core/eval/base_test.py b/langfun/core/eval/base_test.py index f7c7fb10..5feec5e1 100644 --- a/langfun/core/eval/base_test.py +++ b/langfun/core/eval/base_test.py @@ -25,6 +25,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 + + # We put class definitions outside the functors just to make it easier # to refer to them in test. class Solution(pg.Object): @@ -208,7 +234,7 @@ def test_run(self): s = eval_set('run_test', 'query', schema_fn=answer_schema(), lm=lm) s.run() self.assertEqual( - s.result, + _strip_timing(s.result), dict( experiment_setup=dict( id='Evaluation@e028b6e6', @@ -229,9 +255,7 @@ def test_run(self): oop_failure_rate=0.5, non_oop_failures=0, non_oop_failure_rate=0.0, - failure_breakdown={ - 'MappingError.SchemaError.TypeError': 1 - } + failure_breakdown={'MappingError.SchemaError.TypeError': 1}, ), usage=dict( total_prompt_tokens=856, @@ -274,6 +298,38 @@ def test_run(self): self.assertEqual(len(summary['Evaluation']), 1) self.assertIsNotNone(summary['Evaluation'][0].experiment) self.assertIsNotNone(summary['Evaluation'][0].metrics) + # Per-leaf timing is propagated into the summary.json leaf entry too. + self.assertIsNotNone(summary['Evaluation'][0].start_time) + self.assertIsNotNone(summary['Evaluation'][0].end_time) + self.assertIsNotNone(summary['Evaluation'][0].wall_clock_s) + + def test_run_timing(self): + lm = fake.StaticSequence([ + 'Solution(final_answer=2)', + '3', + ]) + s = eval_set('run_timing_test', 'query', schema_fn=answer_schema(), lm=lm) + s.run() + + # The in-memory result carries the per-leaf timing fields. + self.assertIsInstance(s.result.start_time, float) + self.assertIsInstance(s.result.end_time, float) + self.assertIsInstance(s.result.wall_clock_s, float) + self.assertGreaterEqual(s.result.end_time, s.result.start_time) + self.assertGreaterEqual(s.result.wall_clock_s, 0.0) + self.assertAlmostEqual( + s.result.wall_clock_s, + s.result.end_time - s.result.start_time, + places=3, + ) + + # And they are persisted into result.json on disk. + result_json = os.path.join(s.dir, base.Evaluation.RESULT_JSON) + self.assertTrue(os.path.exists(result_json)) + loaded = pg.load(result_json) + self.assertEqual(loaded.start_time, s.result.start_time) + self.assertEqual(loaded.end_time, s.result.end_time) + self.assertEqual(loaded.wall_clock_s, s.result.wall_clock_s) def test_run_wihtout_save(self): lm = fake.StaticSequence([ @@ -320,7 +376,7 @@ def test_run_with_filter(self): filter=lambda x: x.method == 'query', dryrun=True, summary=False ) self.assertEqual( - result, + {k: _strip_timing(v) for k, v in result.items()}, { s.children[0].id: None, s.children[1].id: dict( @@ -381,7 +437,7 @@ def test_search_space(self): summary = s.run(verbose=True) self.assertEqual(len(summary.evaluations), 2) self.assertEqual( - s.result, + {k: _strip_timing(v) for k, v in s.result.items()}, { s.children[0].id: dict( experiment_setup=dict( @@ -403,9 +459,7 @@ def test_search_space(self): oop_failure_rate=0.5, non_oop_failures=0, non_oop_failure_rate=0.0, - failure_breakdown={ - 'MappingError.SchemaError.TypeError': 1 - } + failure_breakdown={'MappingError.SchemaError.TypeError': 1}, ), usage=s.children[0].result.usage, ), @@ -429,9 +483,7 @@ def test_search_space(self): oop_failure_rate=0.5, non_oop_failures=0, non_oop_failure_rate=0.0, - failure_breakdown={ - 'MappingError.SchemaError.TypeError': 1 - } + failure_breakdown={'MappingError.SchemaError.TypeError': 1}, ), usage=s.children[1].result.usage, ), @@ -582,7 +634,9 @@ def test_run(self): usage=s.children[1].children[0].result.usage, ), } - self.assertEqual(s.result, expected) + self.assertEqual( + {k: _strip_timing(v) for k, v in s.result.items()}, expected + ) class InputsFrom(unittest.TestCase): diff --git a/langfun/core/eval/matching.py b/langfun/core/eval/matching.py index 22b0c707..27252eb8 100644 --- a/langfun/core/eval/matching.py +++ b/langfun/core/eval/matching.py @@ -77,12 +77,12 @@ def mismatch_rate(self) -> float: @property def matches_link(self) -> str: """Returns the link to the matches page.""" - return self.link(os.path.join(self.dir, Matching.MATCHES_HTML)) + return self.link(os.path.join(self.dir, Matching.MATCHES_HTML)) # pyrefly: ignore[no-matching-overload] @property def mismatches_link(self) -> str: """Returns the link to the mismatches page.""" - return self.link(os.path.join(self.dir, Matching.MISMATCHES_HTML)) + return self.link(os.path.join(self.dir, Matching.MISMATCHES_HTML)) # pyrefly: ignore[no-matching-overload] def _reset(self) -> None: super()._reset() @@ -178,12 +178,12 @@ def save( if report: pg.save( self._html([self._render_result, self._render_matches]), - os.path.join(self.dir, Matching.MATCHES_HTML), + os.path.join(self.dir, Matching.MATCHES_HTML), # pyrefly: ignore[no-matching-overload] file_format='txt', ) pg.save( self._html([self._render_result, self._render_mismatches]), - os.path.join(self.dir, Matching.MISMATCHES_HTML), + os.path.join(self.dir, Matching.MISMATCHES_HTML), # pyrefly: ignore[no-matching-overload] file_format='txt', ) diff --git a/langfun/core/eval/matching_test.py b/langfun/core/eval/matching_test.py index cfc7d562..ecdb0733 100644 --- a/langfun/core/eval/matching_test.py +++ b/langfun/core/eval/matching_test.py @@ -25,6 +25,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 + + # We put class definitions outside the functors just to make it easier # to refer to them in test. @@ -99,7 +125,7 @@ def test_run(self): s = eval_set('match_run_test', 'query', schema_fn=answer_schema(), lm=lm) s.run() - result_without_id = s.result.copy() + result_without_id = _strip_timing(s.result) result_without_id['experiment_setup'].pop('id') self.assertEqual( result_without_id, diff --git a/langfun/core/eval/scoring.py b/langfun/core/eval/scoring.py index 1952f9ce..a903e371 100644 --- a/langfun/core/eval/scoring.py +++ b/langfun/core/eval/scoring.py @@ -49,7 +49,7 @@ def score_rate(self) -> float: @property def scored_link(self) -> str: """Returns the scored examples page.""" - return self.link(os.path.join(self.dir, Scoring.SCORED_HTML)) + return self.link(os.path.join(self.dir, Scoring.SCORED_HTML)) # pyrefly: ignore[no-matching-overload] @property def avg_score(self) -> float: @@ -138,13 +138,13 @@ def save( pg.Dict(input=input, output=output, score=score) for input, output, score, _ in self.scored ], - os.path.join(self.dir, Scoring.SCORED_JSON), + os.path.join(self.dir, Scoring.SCORED_JSON), # pyrefly: ignore[no-matching-overload] ) if report: pg.save( self._html([self._render_result, self._render_scored]), - os.path.join(self.dir, Scoring.SCORED_HTML), + os.path.join(self.dir, Scoring.SCORED_HTML), # pyrefly: ignore[no-matching-overload] file_format='txt', ) diff --git a/langfun/core/eval/scoring_test.py b/langfun/core/eval/scoring_test.py index 4a9cd9ab..b21b5290 100644 --- a/langfun/core/eval/scoring_test.py +++ b/langfun/core/eval/scoring_test.py @@ -23,6 +23,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 + + @pg.functor() def float_list(): return list[float] @@ -77,7 +103,7 @@ def test_run(self): s = eval_set(lm=lm) self.assertEqual(s.avg_score, 0.0) s.run() - result_copy = s.result.copy() + result_copy = _strip_timing(s.result) del result_copy['experiment_setup']['id'] self.assertEqual( result_copy, diff --git a/langfun/core/eval/v2/checkpointing.py b/langfun/core/eval/v2/checkpointing.py index 0506055e..cca83dea 100644 --- a/langfun/core/eval/v2/checkpointing.py +++ b/langfun/core/eval/v2/checkpointing.py @@ -383,7 +383,7 @@ def on_experiment_complete( """Closes the checkpoint file.""" if not experiment.is_leaf: return - assert experiment.id in self._sequence_writer + assert experiment.id in self._sequence_writer # pyrefly: ignore[not-iterable] with self._lock: if self._sequence_writer is not None: # Make sure the writer is closed without delay so the file will be @@ -402,9 +402,9 @@ def _save_example( example: Example, ) -> None: """Saves the example to the checkpoint file.""" - assert experiment.id in self._sequence_writer + assert experiment.id in self._sequence_writer # pyrefly: ignore[not-iterable] def _save_example(example: Example): - writer = self._sequence_writer[experiment.id] + writer = self._sequence_writer[experiment.id] # pyrefly: ignore[unsupported-operation] try: writer.add(example) experiment.info( @@ -463,7 +463,7 @@ def close(self): if self._sequence_writer is None: return self._sequence_writer.close() - self._sequence_writer = None + self._sequence_writer = None # pyrefly: ignore[bad-assignment] pg.io.rename(self._tmp_path, self._path) def __enter__(self): diff --git a/langfun/core/eval/v2/eval_test_helper.py b/langfun/core/eval/v2/eval_test_helper.py index a10f7380..b1709c35 100644 --- a/langfun/core/eval/v2/eval_test_helper.py +++ b/langfun/core/eval/v2/eval_test_helper.py @@ -58,7 +58,7 @@ def _response_from(self, prompt: message_lib.Message) -> message_lib.Message: ) @property - def resource_id(self) -> str: + def resource_id(self) -> str: # pyrefly: ignore[bad-override] return f'test_llm:{self.offset}' @@ -197,7 +197,7 @@ def on_experiment_start( ) -> None: del runner with pg.notify_on_change(False), self._lock: - self.started_experiments.append(pg.Ref(experiment)) + self.started_experiments.append(pg.Ref(experiment)) # pyrefly: ignore[bad-argument-type] def on_experiment_skipped( self, @@ -206,7 +206,7 @@ def on_experiment_skipped( ) -> None: del runner with pg.notify_on_change(False), self._lock: - self.skipped_experiments.append(pg.Ref(experiment)) + self.skipped_experiments.append(pg.Ref(experiment)) # pyrefly: ignore[bad-argument-type] def on_experiment_complete( self, @@ -215,7 +215,7 @@ def on_experiment_complete( ) -> None: del runner with pg.notify_on_change(False), self._lock: - self.completed_experiments.append(pg.Ref(experiment)) + self.completed_experiments.append(pg.Ref(experiment)) # pyrefly: ignore[bad-argument-type] def on_example_start( self, diff --git a/langfun/core/eval/v2/evaluation.py b/langfun/core/eval/v2/evaluation.py index 47ea7d26..b5bc0600 100644 --- a/langfun/core/eval/v2/evaluation.py +++ b/langfun/core/eval/v2/evaluation.py @@ -131,7 +131,7 @@ def is_leaf(self) -> bool: return self.is_deterministic @functools.cached_property - def children(self) -> list['Evaluation']: + def children(self) -> list['Evaluation']: # pyrefly: ignore[bad-override] """Returns the children tasks.""" if self.is_leaf: return [] @@ -153,7 +153,7 @@ def example_inputs(self) -> Iterable[Any]: def example_input_by_id(self, example_id: int) -> Any: """Returns the example from the inputs by ID.""" - assert example_id <= len(self.example_inputs), example_id + assert example_id <= len(self.example_inputs), example_id # pyrefly: ignore[bad-argument-type] return self._example_input_by_id[example_id] @functools.cached_property @@ -170,7 +170,7 @@ def num_examples(self) -> int: if not isinstance(num_examples, int): it = self.example_inputs if hasattr(it, '__len__'): - num_examples = len(it) + num_examples = len(it) # pyrefly: ignore[bad-argument-type] else: num_examples = len(list(it)) return num_examples @@ -349,7 +349,7 @@ def _process( except BaseException as e: # pylint: disable=broad-except if raise_if_has_error: raise - example.error = pg.ErrorInfo.from_exception(e) + example.error = pg.ErrorInfo.from_exception(e) # pyrefly: ignore[bad-assignment] # # Handling evaluation scheduling. @@ -461,7 +461,7 @@ def _html_tree_view_content( ): if not self.is_leaf: return super()._html_tree_view_content( - view=view, extra_flags=extra_flags, **kwargs + view=view, extra_flags=extra_flags, **kwargs # pyrefly: ignore[bad-argument-type] ) extra_flags = extra_flags or {} @@ -719,6 +719,7 @@ def _in_progress_view( for example in in_progress_examples: if example.newly_processed: logs.append( + # pyrefly: ignore[unsupported-operation] f'Example {example.id}: In progress for ' f'{current_time - example.start_time:.2f} seconds.' ) @@ -894,7 +895,7 @@ def load( for example in example_lib.Example.iter_ckpts( state_file, example_input_by_id=example_input_by_id, - load_example_metadata=load_example_metadata, + load_example_metadata=load_example_metadata, # pyrefly: ignore[bad-argument-type] ): if filter is not None and not filter(example): continue diff --git a/langfun/core/eval/v2/example.py b/langfun/core/eval/v2/example.py index 006b3f05..1cceb353 100644 --- a/langfun/core/eval/v2/example.py +++ b/langfun/core/eval/v2/example.py @@ -100,7 +100,7 @@ def to_json(self, *, exclude_input: bool = False, **kwargs): ) @classmethod - def from_json( + def from_json( # pyrefly: ignore[bad-override] cls, json_value: dict[str, Any], *, @@ -111,7 +111,7 @@ def from_json( """Creates an example from the JSON representation.""" example_id = json_value.get('id') if example_input_by_id: - example_input = example_input_by_id(example_id) + example_input = example_input_by_id(example_id) # pyrefly: ignore[bad-argument-type] else: example_input = json_value.pop('input', pg.MISSING_VALUE) if example_input is not pg.MISSING_VALUE: diff --git a/langfun/core/eval/v2/experiment.py b/langfun/core/eval/v2/experiment.py index afd4db83..af0b03d6 100644 --- a/langfun/core/eval/v2/experiment.py +++ b/langfun/core/eval/v2/experiment.py @@ -304,8 +304,8 @@ def nodes(self) -> list['Experiment']: """Returns all the experiment nodes in the subtree (including self).""" nodes = [self] for child in self.children: - nodes.extend(child.nodes) - return nodes + nodes.extend(child.nodes) # pyrefly: ignore[bad-argument-type] + return nodes # pyrefly: ignore[bad-return] @functools.cached_property def leaf_nodes(self) -> list['Experiment']: @@ -329,8 +329,8 @@ def nonleaf_nodes(self) -> list['Experiment']: return [] nodes = [self] for child in self.children: - nodes.extend(child.nonleaf_nodes) - return nodes + nodes.extend(child.nonleaf_nodes) # pyrefly: ignore[bad-argument-type] + return nodes # pyrefly: ignore[bad-return] @functools.cached_property def parent(self) -> Optional['Experiment']: @@ -471,7 +471,7 @@ def run( """ if plugins is not None: kwargs['plugins'] = plugins - runner = Runner.create( + runner = Runner.create( # pyrefly: ignore[bad-assignment] runner, current_run=Run( root_dir=root_dir, @@ -493,8 +493,8 @@ def run( ), **kwargs ) - runner.run() - return runner.current_run + runner.run() # pyrefly: ignore[missing-attribute] + return runner.current_run # pyrefly: ignore[missing-attribute] def run_preconfigured( self, @@ -612,7 +612,7 @@ def _html_tree_view_summary( **kwargs ) - def _html_tree_view_content( + def _html_tree_view_content( # pyrefly: ignore[bad-override] self, *, view, @@ -779,13 +779,13 @@ def from_id( '`root_dir` must be provided for `latest` or `new` run ID.' ) if run_id == 'latest': - run_id = cls.get_latest(root_dir) + run_id = cls.get_latest(root_dir) # pyrefly: ignore[bad-assignment] if run_id is None: raise ValueError( f'There are no previous runs under the root directory: ' f'{root_dir}. Consider running the experiment using `new` as id.' ) - return run_id + return run_id # pyrefly: ignore[bad-return] if run_id == 'new': return cls.new(root_dir) return cls.get_latest(root_dir) or cls.new() diff --git a/langfun/core/eval/v2/metrics.py b/langfun/core/eval/v2/metrics.py index 5c6c4f20..0020fc6b 100644 --- a/langfun/core/eval/v2/metrics.py +++ b/langfun/core/eval/v2/metrics.py @@ -250,7 +250,7 @@ def update_metric_values( example_id, metric_metadata ) else: - self._update_metric_values(example_id, metric_metadata) + self._update_metric_values(example_id, metric_metadata) # pyrefly: ignore[bad-argument-count, bad-argument-type] @abc.abstractmethod def _compute_metric_metadata( @@ -281,9 +281,9 @@ def _update_metric_values_with_processing_error( assert error_tag is not None, (example_id, metric_metadata) self._error_breakdown[error_tag].append(example_id) if error_tag.startswith('MappingError'): - self.oop_errors.add(example_id, 1) + self.oop_errors.add(example_id, 1) # pyrefly: ignore[missing-attribute] else: - self.non_oop_errors.add(example_id, 1) + self.non_oop_errors.add(example_id, 1) # pyrefly: ignore[missing-attribute] self._error_breakdown[error_tag].append(example_id) def _oop_errors_breakdown(self) -> str | None: @@ -356,7 +356,7 @@ def _compute_metric_metadata( metadata['is_correct'] = is_correct return metadata - def _update_metric_values( + def _update_metric_values( # pyrefly: ignore[bad-override] self, example_id: int, metadata: dict[str, Any] ) -> None: """Update metric values based metric metadata.""" @@ -370,7 +370,7 @@ def _update_metric_values( def values(self) -> list[metric_values.MetricValue]: """Returns all the values computed by this metric.""" - return [ + return [ # pyrefly: ignore[bad-return] self.matches, self.mismatches, self.oop_errors, @@ -437,7 +437,7 @@ def _compute_metric_metadata( metadata['score'] = score return metadata - def _update_metric_values( + def _update_metric_values( # pyrefly: ignore[bad-override] self, example_id: int, metadata: dict[str, Any] ) -> None: """Update metric values based metric metadata.""" @@ -447,7 +447,7 @@ def _update_metric_values( def values(self) -> list[metric_values.MetricValue]: """Returns all the values computed by this metric.""" - return [ + return [ # pyrefly: ignore[bad-return] self.average_score, self.oop_errors, self.non_oop_errors diff --git a/langfun/core/eval/v2/progress.py b/langfun/core/eval/v2/progress.py index c9063232..99eec532 100644 --- a/langfun/core/eval/v2/progress.py +++ b/langfun/core/eval/v2/progress.py @@ -258,7 +258,7 @@ def merge_from(self, other: 'Progress') -> None: self.num_processed += other.num_processed self.num_failed += other.num_failed self.num_skipped += other.num_skipped - self.execution_summary.aggregate(other.execution_summary.breakdown) + self.execution_summary.aggregate(other.execution_summary.breakdown) # pyrefly: ignore[bad-argument-type] self._prior_elapse += other.prior_elapse # @@ -268,7 +268,7 @@ def merge_from(self, other: 'Progress') -> None: def _duration_text(self) -> str: if self.start_time is None: return '00:00:00' - return str(datetime.timedelta(seconds=self.elapse)).split('.')[0] + return str(datetime.timedelta(seconds=self.elapse)).split('.')[0] # pyrefly: ignore[bad-argument-type] def _time_tooltip(self) -> pg.Html.WritableTypes: time_info = pg.Dict( diff --git a/langfun/core/eval/v2/progress_tracking.py b/langfun/core/eval/v2/progress_tracking.py index f5182431..13509fe1 100644 --- a/langfun/core/eval/v2/progress_tracking.py +++ b/langfun/core/eval/v2/progress_tracking.py @@ -83,7 +83,7 @@ def experiment_progress( self, experiment: Experiment) -> lf.concurrent.ProgressBar: """Returns the progress of the experiment.""" assert experiment.is_leaf - return self._leaf_progresses[experiment.id] + return self._leaf_progresses[experiment.id] # pyrefly: ignore[bad-return] def on_run_start( self, @@ -115,11 +115,11 @@ def on_run_complete( ) -> None: """Called when a runner is complete.""" lf.concurrent.ProgressBar.update( - self._overall_progress, + self._overall_progress, # pyrefly: ignore[bad-argument-type] color='green', status='ALL COMPLETED.', ) - lf.concurrent.ProgressBar.uninstall(self._overall_progress) + lf.concurrent.ProgressBar.uninstall(self._overall_progress) # pyrefly: ignore[bad-argument-type] self._overall_progress = None for progress in self._leaf_progresses.values(): lf.concurrent.ProgressBar.uninstall(progress) @@ -140,12 +140,12 @@ def on_experiment_skipped( """Called when an evaluation is skipped.""" if experiment.is_leaf: lf.concurrent.ProgressBar.update( - self.experiment_progress(experiment), - delta=experiment.progress.num_total, + self.experiment_progress(experiment), # pyrefly: ignore[bad-argument-type] + delta=experiment.progress.num_total, # pyrefly: ignore[bad-argument-type] status='Skipped.', ) lf.concurrent.ProgressBar.update( - self._overall_progress, + self._overall_progress, # pyrefly: ignore[bad-argument-type] status=f'Skipped {experiment.id}.', ) @@ -157,11 +157,11 @@ def on_experiment_complete( """Called when an evaluation is complete.""" if experiment.is_leaf: lf.concurrent.ProgressBar.update( - self.experiment_progress(experiment), + self.experiment_progress(experiment), # pyrefly: ignore[bad-argument-type] color='green', ) lf.concurrent.ProgressBar.update( - self._overall_progress, + self._overall_progress, # pyrefly: ignore[bad-argument-type] delta=1, status=f'{experiment.id} COMPLETED.', ) @@ -183,7 +183,7 @@ def on_example_skipped( """Called when an evaluation example is skipped.""" del runner, example lf.concurrent.ProgressBar.update( - self.experiment_progress(experiment), + self.experiment_progress(experiment), # pyrefly: ignore[bad-argument-type] delta=1, ) @@ -195,7 +195,7 @@ def on_example_complete( ) -> None: """Called when an evaluation example is complete.""" lf.concurrent.ProgressBar.update( - self.experiment_progress(experiment), + self.experiment_progress(experiment), # pyrefly: ignore[bad-argument-type] delta=1, status=self.status(experiment), ) diff --git a/langfun/core/eval/v2/reporting.py b/langfun/core/eval/v2/reporting.py index 0679aee6..f55a3ffb 100644 --- a/langfun/core/eval/v2/reporting.py +++ b/langfun/core/eval/v2/reporting.py @@ -232,7 +232,7 @@ def _summary(): current_run=run, interactive=False, card_view=True, ) ) - with self._summary_lock: + with self._summary_lock: # pyrefly: ignore[bad-context-manager] html.save(os.path.join(run.output_root, _SUMMARY_FILE)) if force or (time.time() - self._last_summary_time > self.summary_interval): @@ -263,7 +263,7 @@ def _save(): card_view=False, ), ) - with self._experiment_index_lock[experiment.id]: + with self._experiment_index_lock[experiment.id]: # pyrefly: ignore[unsupported-operation] html.save(index_html_path) experiment.info( f'Updated {index_html_path!r} in {t.elapse:.2f} seconds.', @@ -279,7 +279,7 @@ def _save(): time.time() - self._last_experiment_report_time[experiment.id] > self.experiment_report_interval ): - self._last_experiment_report_time[experiment.id] = time.time() + self._last_experiment_report_time[experiment.id] = time.time() # pyrefly: ignore[unsupported-operation] if background: runner.background_run(_save) else: diff --git a/langfun/core/eval/v2/runners/base.py b/langfun/core/eval/v2/runners/base.py index a57d52db..166d5ddb 100644 --- a/langfun/core/eval/v2/runners/base.py +++ b/langfun/core/eval/v2/runners/base.py @@ -101,7 +101,7 @@ def _background_run(*args, **kwargs): self._background_last_error = e if self.max_background_threads > 0: - with self._io_pool_lock: + with self._io_pool_lock: # pyrefly: ignore[bad-context-manager] if self._io_pool is not None: self._io_pool.submit(_background_run, *args, **kwargs) else: @@ -147,7 +147,7 @@ def on_experiment_start(self, experiment: Experiment) -> None: plugin.on_experiment_start(self, experiment) if experiment.is_leaf: - self._set_prior_elapse(experiment) + self._set_prior_elapse(experiment) # pyrefly: ignore[bad-argument-type] if experiment.is_leaf: pg.io.mkdirs(self.current_run.output_dir(experiment)) @@ -293,12 +293,12 @@ def on_example_complete( f'in {example.elapse:.2f} seconds.' ) - experiment.usage_summary.merge(example.usage_summary) - experiment.progress.update_execution_summary(example.execution_status) + experiment.usage_summary.merge(example.usage_summary) # pyrefly: ignore[bad-argument-type] + experiment.progress.update_execution_summary(example.execution_status) # pyrefly: ignore[bad-argument-type] parent = experiment.parent while parent is not None: - parent.usage_summary.merge(example.usage_summary) + parent.usage_summary.merge(example.usage_summary) # pyrefly: ignore[bad-argument-type] parent = parent.parent for plugin in self._all_plugins(experiment): @@ -338,11 +338,11 @@ def run(self) -> None: # Evaluate the leaf evaluations if not skipped. with lf.use_settings(**global_settings): - self._run(targets) + self._run(targets) # pyrefly: ignore[bad-argument-type] self.on_run_complete() except BaseException as e: # pylint: disable=broad-except - self.on_run_abort(e) + self.on_run_abort(e) # pyrefly: ignore[bad-argument-type] raise e finally: if cache is not None: @@ -350,9 +350,9 @@ def run(self) -> None: # Wait for the background tasks to finish. if self.max_background_threads > 0: - with self._io_pool_lock: + with self._io_pool_lock: # pyrefly: ignore[bad-context-manager] self._io_pool, io_pool = None, self._io_pool - io_pool.shutdown(wait=True) + io_pool.shutdown(wait=True) # pyrefly: ignore[missing-attribute] @abc.abstractmethod def _run(self, evaluations: list[Evaluation]) -> None: @@ -382,7 +382,7 @@ def run_evaluation(self, evaluation: Evaluation) -> None: if self.current_run.shuffle_inputs: items = list(items) random.shuffle(items) - self._evaluate_items(evaluation, items) + self._evaluate_items(evaluation, items) # pyrefly: ignore[bad-argument-type] if cache: self.background_run(cache.save) diff --git a/langfun/core/eval/v2/runners/beam.py b/langfun/core/eval/v2/runners/beam.py index 915d879c..0a788280 100644 --- a/langfun/core/eval/v2/runners/beam.py +++ b/langfun/core/eval/v2/runners/beam.py @@ -300,7 +300,7 @@ def run(self) -> None: example_ids = range(1, evaluation.num_examples + 1) inputs = [ example_lib.Example(id=i, input=evaluation.example_input_by_id(i)) - for i in example_ids + for i in example_ids # pyrefly: ignore[not-iterable] ] if self.current_run.shuffle_inputs: random.shuffle(inputs) @@ -321,7 +321,7 @@ def run(self) -> None: ) | f'Evaluate-{evaluation.id}' >> beam.ParDo( - _EvaluateFn( + _EvaluateFn( # pyrefly: ignore[not-callable] pg.to_json_str(leaf_node_runner), ckpt_format=self.ckpt_format, concurrent_startup_delay=self.concurrent_startup_delay, diff --git a/langfun/core/eval/v2/runners/ckpt_monitor.py b/langfun/core/eval/v2/runners/ckpt_monitor.py index fee1275d..dbf90fc5 100644 --- a/langfun/core/eval/v2/runners/ckpt_monitor.py +++ b/langfun/core/eval/v2/runners/ckpt_monitor.py @@ -110,19 +110,19 @@ def start(self): # This is not precise, but we at least notify example start. if not self.current_run.filter or self.current_run.filter(evaluation): self.on_experiment_start(evaluation) - self._set_prior_elapse_from_checkpoints(evaluation) + self._set_prior_elapse_from_checkpoints(evaluation) # pyrefly: ignore[bad-argument-type] # Signal the start of the examples if we are not monitoring in-progress # files. if not self.monitor_inprogress_files: for example_id in self.current_run.examples_to_evaluate(evaluation): - self._mark_example_started(evaluation, example_id) + self._mark_example_started(evaluation, example_id) # pyrefly: ignore[bad-argument-type] # Create the aggregation entries for polling. output_dir = self.current_run.output_dir(evaluation) self._aggregation_entries.append( self._AggregationEntry( - evaluation=evaluation, + evaluation=evaluation, # pyrefly: ignore[bad-argument-type] output_dir=output_dir, ckpt_file_pattern=os.path.join( output_dir, self.checkpoint_pattern @@ -215,7 +215,7 @@ def _monitor_loop(self): self._mark_example_started(entry.evaluation, example_id) entry.example_ids_inprogress.add(example_id) - self._aggregator_pool.submit( + self._aggregator_pool.submit( # pyrefly: ignore[missing-attribute] self._aggregate, entry, filepath, example_id, last_modified_time ) pg.logging.info( @@ -229,7 +229,7 @@ def _monitor_loop(self): if self._error is None: self.on_run_complete() else: - self.on_run_abort(self._error) + self.on_run_abort(self._error) # pyrefly: ignore[bad-argument-type] def _aggregate( self, @@ -254,13 +254,13 @@ def _aggregate( example = loaded_examples[-1] if ( self.bypass_old_ckpt_files_with_non_oop_errors - and last_modified_time < self.ckpt_start_time + and last_modified_time < self.ckpt_start_time # pyrefly: ignore[unsupported-operation] and example.error is not None and not example.error.tag.startswith('MappingError') ): entry.example_ids_being_aggregated.remove(example_id) entry.example_ids_to_be_aggregated.add(example_id) - self._ckpt_bypass_timestamp[ckpt_filepath] = last_modified_time + self._ckpt_bypass_timestamp[ckpt_filepath] = last_modified_time # pyrefly: ignore[unsupported-operation] pg.logging.info( '[%s] Bypassing old checkpoint file with non-oop errors (%s) ' 'for example %d, last_modified_time: %s, ckpt_start_time: %s', @@ -282,7 +282,7 @@ def _aggregate( example = example_lib.Example( id=example_id, input=entry.evaluation.example_input_by_id(example_id), - error=error_info, + error=error_info, # pyrefly: ignore[bad-argument-type] ) # This will skip processing but still allow metrics to be collected. @@ -363,7 +363,7 @@ def _set_prior_elapse_from_checkpoints( total_elapse = 0.0 for filepath in pg.io.glob(ckpt_file_pattern): last_modified_time = pg.io.getmtime(filepath) - if last_modified_time >= self.ckpt_start_time: + if last_modified_time >= self.ckpt_start_time: # pyrefly: ignore[unsupported-operation] continue try: loaded_examples = evaluation.state.load( diff --git a/langfun/core/eval/v2/runners/parallel.py b/langfun/core/eval/v2/runners/parallel.py index 268288be..23925b40 100644 --- a/langfun/core/eval/v2/runners/parallel.py +++ b/langfun/core/eval/v2/runners/parallel.py @@ -98,7 +98,7 @@ def _evaluate_item(item: base.Example): for _, _, _ in lf.concurrent_map( _evaluate_item, items, - max_workers=self._max_workers(evaluation), + max_workers=self._max_workers(evaluation), # pyrefly: ignore[bad-argument-type] timeout=self.timeout, silence_on_errors=None, ): @@ -200,7 +200,7 @@ def _on_bound(self): f'Ignoring checkpointer: {plugin!r}.' ) elif plugin.is_per_example(): - worker_plugins.append(pg.Ref(plugin)) + worker_plugins.append(pg.Ref(plugin)) # pyrefly: ignore[bad-argument-type] else: monitor_plugins.append(pg.Ref(plugin)) diff --git a/langfun/core/langfunc.py b/langfun/core/langfunc.py index c878bf4c..1af94005 100644 --- a/langfun/core/langfunc.py +++ b/langfun/core/langfunc.py @@ -164,10 +164,10 @@ def _call_once( self._cached_lm_input = lm_input # Send rendered text to LM. - lm_output = self.lm(lm_input, cache_seed=cache_seed) + lm_output = self.lm(lm_input, cache_seed=cache_seed) # pyrefly: ignore[bad-argument-type] # Attach cache seed. - lm_input.metadata.cache_seed = cache_seed + lm_input.metadata.cache_seed = cache_seed # pyrefly: ignore[missing-attribute] # Transform the output message. lm_output = self.transform_output(lm_output) @@ -249,7 +249,7 @@ def transform_output( # Register converter from str to LangFunc, therefore we can always # pass strs to attributes that accept LangFunc. -pg.typing.register_converter(str, LangFunc, LangFunc) +pg.typing.register_converter(str, LangFunc, LangFunc) # pyrefly: ignore[bad-argument-type] # diff --git a/langfun/core/language_model.py b/langfun/core/language_model.py index b77f1523..189a8e5a 100644 --- a/langfun/core/language_model.py +++ b/langfun/core/language_model.py @@ -53,6 +53,17 @@ class ContentFilteredError(LMError): """Error raised when LLM output is blocked by content filtering policy.""" +class ResponseSizeLimitError(LMError): + """Error raised when a response body exceeds its configured size limit. + + This is intentionally a NON-retryable `LMError` (not a `RetryableLMError`): + an oversized response is typically deterministic (e.g. a runaway generation), + so retrying would merely re-stream the same oversized body and waste + resources. Enforcing a size bound while reading a streamed response prevents + unbounded in-memory buffering (which can OOM-kill the process). + """ + + class RetryableLMError(LMError): """Base class for LLM errors that can be solved by retrying.""" @@ -893,7 +904,7 @@ class LanguageModel(component.Component): _MODEL_FACTORY: ClassVar[dict[str, Callable[..., 'LanguageModel']]] = {} @classmethod - def register( + def register( # pyrefly: ignore[bad-override] cls, model_id_or_prefix: str, factory: Callable[..., 'LanguageModel'] ) -> None: @@ -1126,7 +1137,7 @@ def sample( def _sample_with_retry(): if self.cache is None: - results = self._sample(prompts) + results = self._sample(prompts) # pyrefly: ignore[bad-argument-type] else: results = self._sample_with_cache_lookup(prompts, cache_seed) @@ -1134,7 +1145,7 @@ def _sample_with_retry(): for sample in result.samples: if not sample.response.text: if self.cache is not None: - self.cache.delete(self, prompts[i], seed=cache_seed) + self.cache.delete(self, prompts[i], seed=cache_seed) # pyrefly: ignore[bad-argument-type] raise EmptyGenerationError( f'Empty generation encountered from model {self.model_id}.' ) @@ -1153,7 +1164,7 @@ def _sample_with_retry(): for prompt, result in zip(prompts, results): # Tag LM input. - prompt.tag(message_lib.Message.TAG_LM_INPUT) + prompt.tag(message_lib.Message.TAG_LM_INPUT) # pyrefly: ignore[missing-attribute] for sample in result.samples: # Update metadata for response message. @@ -1226,15 +1237,15 @@ def _sample_with_cache_lookup( r = None # Query cache if cache_seed is not None. if cache_seed is not None: - r = self.cache.get(self, prompt, seed=cache_seed) + r = self.cache.get(self, prompt, seed=cache_seed) # pyrefly: ignore[bad-argument-type] if r is None: request_to_result_index[len(requests)] = i requests.append(prompt) else: result = r.clone() - assert result.is_cached, result - results[i] = result + assert result.is_cached, result # pyrefly: ignore[missing-attribute] + results[i] = result # pyrefly: ignore[unsupported-operation] # Sample non-cache-hit prompts. if requests: @@ -1452,10 +1463,10 @@ def score( request_start = time.time() with component.context(override_attrs=True, **kwargs): - scoring_results = self._score(prompt, completions) + scoring_results = self._score(prompt, completions) # pyrefly: ignore[bad-argument-type] elapse = time.time() - request_start self._debug_score( - prompt, completions, scoring_results, call_counter, elapse + prompt, completions, scoring_results, call_counter, elapse # pyrefly: ignore[bad-argument-type] ) return scoring_results diff --git a/langfun/core/llms/__init__.py b/langfun/core/llms/__init__.py index c22bc027..60af1d53 100644 --- a/langfun/core/llms/__init__.py +++ b/langfun/core/llms/__init__.py @@ -42,7 +42,9 @@ # Gemini models. from langfun.core.llms.google_genai import GenAI +from langfun.core.llms.google_genai import Gemini37Flash from langfun.core.llms.google_genai import Gemini31ProPreview +from langfun.core.llms.google_genai import Gemini31FlashLite from langfun.core.llms.google_genai import Gemini3ProPreview from langfun.core.llms.google_genai import Gemini3FlashPreview from langfun.core.llms.google_genai import Gemini25Pro @@ -95,9 +97,12 @@ from langfun.core.llms.vertexai import VertexAIGemini25Flash from langfun.core.llms.vertexai import VertexAIGemini25FlashImagePreview from langfun.core.llms.vertexai import VertexAIGemini31ProPreview +from langfun.core.llms.vertexai import VertexAIGemini31FlashLite from langfun.core.llms.vertexai import VertexAIGemini3ProPreview from langfun.core.llms.vertexai import VertexAIGemini3ProImagePreview from langfun.core.llms.vertexai import VertexAIGemini3FlashPreview +from langfun.core.llms.vertexai import VertexAIGemini35Flash +from langfun.core.llms.vertexai import VertexAIGemini37Flash # Veo video generation models. from langfun.core.llms.veo import Veo @@ -168,6 +173,8 @@ # Anthropic models. +from langfun.core.llms.anthropic import Claude5Opus +from langfun.core.llms.anthropic import Claude48Opus from langfun.core.llms.anthropic import Claude47Opus from langfun.core.llms.anthropic import Claude46 from langfun.core.llms.anthropic import Claude46Opus @@ -192,6 +199,8 @@ from langfun.core.llms.anthropic import Claude3Haiku_20240307 from langfun.core.llms.vertexai import VertexAIAnthropic +from langfun.core.llms.vertexai import VertexAIClaude5Opus +from langfun.core.llms.vertexai import VertexAIClaude48Opus from langfun.core.llms.vertexai import VertexAIClaude47Opus from langfun.core.llms.vertexai import VertexAIClaude46Opus from langfun.core.llms.vertexai import VertexAIClaude45Haiku_20251001 diff --git a/langfun/core/llms/anthropic.py b/langfun/core/llms/anthropic.py index a9d28859..81ccd9e5 100644 --- a/langfun/core/llms/anthropic.py +++ b/langfun/core/llms/anthropic.py @@ -15,6 +15,7 @@ import datetime import functools +import json import os from typing import Annotated, Any, Literal @@ -53,7 +54,7 @@ class RateLimits(lf.ModelInfo.RateLimits): max_output_tokens_per_minute: int @property - def max_tokens_per_minute(self) -> int: + def max_tokens_per_minute(self) -> int: # pyrefly: ignore[bad-override] return (self.max_input_tokens_per_minute + self.max_output_tokens_per_minute) @@ -137,6 +138,37 @@ def max_tokens_per_minute(self) -> int: max_output_tokens_per_minute=400_000, ), ), + AnthropicModelInfo( + model_id='claude-opus-5', + provider='Anthropic', + in_service=True, + description='Claude Opus 5 model.', + # release_date and knowledge_cutoff intentionally omitted: Opus 5 is a + # dateless/pinned snapshot and neither date is doc-grounded. Both fields + # default to None (unknown), matching the convention used by most other + # entries in this list rather than shipping fabricated dates. + input_modalities=( + AnthropicModelInfo.INPUT_IMAGE_TYPES + + AnthropicModelInfo.INPUT_DOC_TYPES + ), + context_length=lf.ModelInfo.ContextLength( + max_input_tokens=1_000_000, + max_output_tokens=128_000, + ), + pricing=lf.ModelInfo.Pricing( + cost_per_1m_cached_input_tokens=0.5, + cost_per_1m_input_tokens=5.0, + cost_per_1m_output_tokens=25.0, + ), + # UNVERIFIED: no public/internal doc grounds Opus 5 quota; these + # rate_limits are copied from the Opus 4.8 entry as a best-effort + # placeholder. Update once official Opus 5 limits are published. + rate_limits=AnthropicModelInfo.RateLimits( + max_requests_per_minute=2000, + max_input_tokens_per_minute=1_000_000, + max_output_tokens_per_minute=400_000, + ), + ), AnthropicModelInfo( model_id='claude-haiku-4-5-20251001', provider='Anthropic', @@ -859,6 +891,20 @@ def max_tokens_per_minute(self) -> int: _SUPPORTED_MODELS_BY_MODEL_ID = {m.model_id: m for m in SUPPORTED_MODELS} +# Thinking-effort tiers ordered from cheapest to deepest. Anthropic's own +# vocabulary (`Anthropic.effort`) is a superset of the cross-provider +# `LMSamplingOptions.reasoning_effort` vocabulary (low/medium/high), so the two +# config surfaces can disagree; this ordering is what lets us tell an +# (acceptable) upgrade apart from a (reportable) downgrade. +_EFFORT_TIERS: dict[str, int] = { + 'low': 0, + 'medium': 1, + 'high': 2, + 'xhigh': 3, + 'max': 4, +} + + def _apply_cache_breakpoints( request: dict[str, Any], *, @@ -979,7 +1025,9 @@ class Anthropic(rest.REST): effort: Annotated[ Literal['low', 'medium', 'high', 'xhigh', 'max'] | None, 'Thinking depth for models supporting extended thinking (low, medium,' - + ' high, xhigh, max).', + + ' high, xhigh, max). It reaches the API only when thinking is enabled' + + ' on a model with adaptive thinking; a configured value that cannot be' + + ' honored is reported instead of being dropped.', ] = 'high' def _on_bound(self): @@ -996,7 +1044,7 @@ def _initialize(self): self._api_key = api_key @property - def headers(self) -> dict[str, Any]: + def headers(self) -> dict[str, Any]: # pyrefly: ignore[bad-override] return { 'x-api-key': self._api_key, 'anthropic-version': self.api_version, @@ -1016,7 +1064,87 @@ def model_info(self) -> lf.ModelInfo: @property def _use_adaptive_thinking(self) -> bool: return self.model is not None and ( - 'claude-opus-4-7' in self.model_id or 'claude-opus-4-8' in self.model_id + 'claude-opus-4-7' in self.model_id + or 'claude-opus-4-8' in self.model_id + or 'claude-opus-5' in self.model_id + ) + + @property + def _effort_is_user_configured(self) -> bool: + """Returns True if `effort` was set to something other than its default. + + Detection is value-based (`pg.Object.sym_nondefault`) rather than + init-based, so both `Claude5Opus(effort='max')` and a later + `rebind(effort='max')` count as user intent. An `effort` that merely + happens to equal the class default is treated as "not configured": it + carries no user decision, so reporting on it would be noise. + """ + return 'effort' in self.sym_nondefault() + + def _resolve_effort(self, options: lf.LMSamplingOptions) -> str | None: + """Resolves the effective thinking effort from the two config surfaces. + + Per-call `reasoning_effort` keeps precedence over the model-level `effort` + (existing, documented behavior). The addition here is honesty: when that + precedence silently *lowers* a user-configured effort -- the cross-provider + `reasoning_effort` vocabulary tops out at 'high', so it can never express + 'xhigh'/'max' -- we say so instead of quietly shipping the weaker setting. + + Args: + options: The sampling options of the current call. + + Returns: + The effort to send to the API, or None if no effort should be sent. + """ + per_call_effort = options.reasoning_effort + if per_call_effort is None: + return self.effort + + if ( + self.effort is not None + and self._effort_is_user_configured + and _EFFORT_TIERS[self.effort] > _EFFORT_TIERS[per_call_effort] + ): + pg.logging.warning( + '[%s] Per-call `reasoning_effort=%r` takes precedence over the ' + 'configured `effort=%r`, downgrading the thinking effort of this ' + 'request; the effective effort is %r. `reasoning_effort` is a ' + 'cross-provider setting limited to low/medium/high and cannot ' + "express %r's %r. Drop `reasoning_effort` (or raise it) to keep the " + 'configured effort.', + self.model_id, + per_call_effort, + self.effort, + per_call_effort, + self.__class__.__name__, + self.effort, + ) + return per_call_effort + + def _warn_effort_ignored(self, reason: str) -> None: + """Reports a user-configured `effort` that this request will not send. + + `effort` only reaches the API through the adaptive-thinking path, which is + gated on both `thinking` being enabled and the model supporting adaptive + thinking. Whenever a configured effort falls outside that path it used to + vanish without a trace; now it is reported. This warns rather than raises + so that existing, blessed configurations (e.g. a manual-budget model + carrying a leftover `effort`) keep working. + + Args: + reason: Why the effort cannot be honored, phrased to complete the + sentence 'is ignored because ...'. + """ + if self.effort is None or not self._effort_is_user_configured: + return + pg.logging.warning( + '[%s] Configured `effort=%r` is ignored because %s, so this request ' + 'is sent without an effort setting. Remove `effort`, or use a model ' + 'with adaptive thinking enabled, to avoid a silent mismatch between ' + 'your configuration and the request.', + self.model_id, + self.effort, + reason, ) def request( @@ -1058,12 +1186,20 @@ def _request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]: """Returns a dict as request arguments.""" # Authropic requires `max_tokens` to be specified. max_tokens = ( - options.max_tokens or self.model_info.context_length.max_output_tokens + options.max_tokens or self.model_info.context_length.max_output_tokens # pyrefly: ignore[missing-attribute] ) args = dict( model=self.model, max_tokens=max_tokens, - stream=False, + # Stream the response. Without this, Vertex/Anthropic buffers the + # ENTIRE response and emits a single JSON blob only at completion, so + # no bytes flow until the generation finishes -- any generation that + # needs longer than the read timeout to produce its first (and only) + # byte dies with a read timeout. With stream=True the server emits + # Server-Sent Events incrementally; _reassemble_sse() below rebuilds + # the full message + usage, and rest.py's inactivity/total deadline + # split lets a slow-but-live generation run for hours. + stream=True, ) if options.stop: args['stop_sequences'] = options.stop @@ -1088,14 +1224,20 @@ def _request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]: 'type': 'adaptive', } if self.model is not None and ( - 'claude-opus-4-7' in self.model or 'claude-opus-4-8' in self.model + 'claude-opus-4-7' in self.model + or 'claude-opus-4-8' in self.model + or 'claude-opus-5' in self.model ): args['thinking']['display'] = 'summarized' - effort = options.reasoning_effort or self.effort + effort = self._resolve_effort(options) if effort: args['output_config'] = {'effort': effort} else: + self._warn_effort_ignored( + f'{self.model_id} does not support adaptive thinking and uses a ' + 'manual thinking budget (`max_thinking_tokens`) instead' + ) budget = options.max_thinking_tokens if budget is None: # Default to 50% of the total capacity, ensuring at least 1024. @@ -1111,7 +1253,7 @@ def _request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]: args['max_tokens'] += budget # Ensure max_tokens does not exceed model's absolute hard capacity. - model_cap = self.model_info.context_length.max_output_tokens + model_cap = self.model_info.context_length.max_output_tokens # pyrefly: ignore[missing-attribute] if args['max_tokens'] > model_cap: args['max_tokens'] = model_cap @@ -1127,10 +1269,17 @@ def _request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]: args.pop('temperature', None) args.pop('top_k', None) args.pop('top_p', None) + else: + self._warn_effort_ignored( + 'thinking is off for this request (`thinking` is False, or unset ' + 'with no `max_thinking_tokens`), and effort only applies to thinking' + ) - # Claude Opus 4.7 and 4.8 do not support temperature, top_p, or top_k. + # Claude Opus 4.7, 4.8 and 5 do not support temperature, top_p, or top_k. if self.model is not None and ( - 'claude-opus-4-7' in self.model or 'claude-opus-4-8' in self.model + 'claude-opus-4-7' in self.model + or 'claude-opus-4-8' in self.model + or 'claude-opus-5' in self.model ): args.pop('temperature', None) args.pop('top_k', None) @@ -1140,10 +1289,10 @@ def _request_args(self, options: lf.LMSamplingOptions) -> dict[str, Any]: args.update(options.extras) return args - def result(self, json: dict[str, Any]) -> lf.LMSamplingResult: - message = lf.Message.from_value(json, format='anthropic') - input_tokens = json['usage']['input_tokens'] - output_tokens = json['usage']['output_tokens'] + def result(self, response_json: dict[str, Any]) -> lf.LMSamplingResult: + message = lf.Message.from_value(response_json, format='anthropic') + input_tokens = response_json['usage']['input_tokens'] + output_tokens = response_json['usage']['output_tokens'] return lf.LMSamplingResult( [lf.LMSample(message)], usage=lf.LMSamplingUsage( @@ -1153,10 +1302,203 @@ def result(self, json: dict[str, Any]) -> lf.LMSamplingResult: ), ) + def _response_to_message_dict(self, response: Any) -> dict[str, Any]: + """Builds the buffered Anthropic message dict consumed by `result`. + + With body `stream=True`, Vertex/Anthropic returns a Server-Sent Events + body instead of a single JSON object. This reassembles that event stream + back into the exact same dict shape the non-streaming Messages API would + have returned (`role`, `content` blocks, `stop_reason`, `usage`), so that + `result()` (and `lf.Message.from_value(..., format='anthropic')`) stay + correct and unchanged. + + For robustness (and to keep buffered-JSON unit tests working), a body that + is already a single JSON object is parsed directly. + + Args: + response: The streaming (or buffered) HTTP response from the API. + + Returns: + The reassembled Anthropic message dict (role/content/stop_reason/usage). + """ + raw = response.content + text = raw.decode('utf-8') if isinstance(raw, (bytes, bytearray)) else raw + stripped = text.lstrip() + # A buffered (non-streamed) JSON body starts with '{'. An SSE body starts + # with an 'event:' / 'data:' line. + if stripped.startswith('{'): + return json.loads(text) + return self._reassemble_sse(text) + + def _reassemble_sse(self, text: str) -> dict[str, Any]: + """Reassembles an Anthropic Messages SSE stream into a message dict. + + Handles the Anthropic streaming event sequence: + message_start -> (content_block_start, + content_block_delta*, content_block_stop)* -> + message_delta -> message_stop + + Rebuilds full text/thinking content AND token usage: + - message_start: provides the message skeleton incl. input_tokens. + - content_block_start: initializes a content block at its index. + - content_block_delta: appends text_delta / thinking_delta / + signature_delta / input_json_delta fragments to that block. + - content_block_stop: finalizes any accumulated tool-use JSON. + - message_delta: carries the final stop_reason and usage.output_tokens. + - message_stop: terminator. + + Args: + text: The full Server-Sent Events response body as text. + + Returns: + The reassembled Anthropic message dict, equivalent to the non-streaming + Messages API response. + """ + message: dict[str, Any] | None = None + blocks: dict[int, dict[str, Any]] = {} + json_buffers: dict[int, str] = {} + final_usage: dict[str, Any] = {} + saw_message_stop = False + + for raw_line in text.splitlines(): + line = raw_line.strip() + if not line.startswith('data:'): + continue + data_str = line[len('data:') :].strip() + if not data_str or data_str == '[DONE]': + continue + try: + event = json.loads(data_str) + except json.JSONDecodeError as e: + # A well-formed Anthropic SSE stream emits exactly one valid-JSON + # payload per `data:` line; keep-alives are comment (`:`) or + # `event: ping` lines (skipped by the `data:` check above) and + # `[DONE]`/empty payloads are handled above. A `data:` line that + # fails to parse here is therefore a CORRUPT/TRUNCATED event, not a + # benign keep-alive. Silently `continue`-ing would drop that + # fragment -- and if it was a content_block_delta, the reconstructed + # text would be silently truncated while message_stop/stop_reason + # still arrive, defeating the terminal sentinel below. The buffered + # path fails loudly on a malformed body (json.loads -> ValueError -> + # LMError) and never returns partial content; preserve that + # no-silent-corruption guarantee here by raising a RETRYABLE error + # (mid-stream corruption is transient, matching the empty-stream and + # truncated-stream guards). + raise lf.TemporaryLMError( + 'Anthropic SSE stream contained a malformed data event ' + f'({data_str[:120]!r}); cannot guarantee complete content, ' + 'retrying.' + ) from e + etype = event.get('type') + + if etype == 'message_start': + message = dict(event['message']) + message['content'] = [] + elif etype == 'content_block_start': + idx = event['index'] + block = dict(event.get('content_block', {})) + blocks[idx] = block + if block.get('type') == 'tool_use': + json_buffers[idx] = '' + elif etype == 'content_block_delta': + idx = event['index'] + delta = event.get('delta', {}) + dtype = delta.get('type') + block = blocks.setdefault(idx, {}) + if dtype == 'text_delta': + block.setdefault('type', 'text') + block['text'] = block.get('text', '') + delta.get('text', '') + elif dtype == 'thinking_delta': + block.setdefault('type', 'thinking') + block['thinking'] = block.get('thinking', '') + delta.get( + 'thinking', '' + ) + elif dtype == 'signature_delta': + block['signature'] = block.get('signature', '') + delta.get( + 'signature', '' + ) + elif dtype == 'input_json_delta': + json_buffers[idx] = json_buffers.get(idx, '') + delta.get( + 'partial_json', '' + ) + elif etype == 'content_block_stop': + idx = event['index'] + buf = json_buffers.get(idx) + if buf: + try: + blocks[idx]['input'] = json.loads(buf) + except json.JSONDecodeError: + blocks[idx]['input'] = {} + elif etype == 'message_delta': + delta = event.get('delta', {}) + if message is not None: + for key in ('stop_reason', 'stop_sequence'): + if key in delta: + message[key] = delta[key] + usage = event.get('usage') + if usage: + final_usage.update(usage) + elif etype == 'error': + # An in-stream `error` event arrives on an HTTP 200 body. Map the + # Anthropic error type to the equivalent HTTP status and reuse the + # buffered-path classification (self._error) so transient failures + # (overloaded_error -> 529 -> TemporaryLMError; rate_limit_error -> + # 429 -> RateLimitError; api_error -> 500 -> TemporaryLMError) stay + # RETRYABLE, while genuinely permanent ones (invalid_request_error -> + # 400) remain permanent. A bare ValueError/lf.LMError here would be + # downgraded to a PERMANENT error by _parse_response and silently lose + # the retry the buffered 529/429 path already gets. + error_type_to_status = { + 'invalid_request_error': 400, + 'authentication_error': 401, + 'permission_error': 403, + 'not_found_error': 404, + 'request_too_large': 413, + 'rate_limit_error': 429, + 'api_error': 500, + 'overloaded_error': 529, + } + err = event.get('error') or {} + status = error_type_to_status.get(err.get('type'), 500) # pyrefly: ignore[no-matching-overload] + # Anthropic._error inspects `content` as bytes, so encode it. + raise self._error(status, json.dumps(err or event).encode('utf-8')) # pyrefly: ignore[bad-argument-type] + elif etype == 'message_stop': + saw_message_stop = True + # 'ping' and other unrecognized events need no handling. + + if message is None: + # A 200 whose SSE body never produced a message_start (empty body, only + # keep-alives, or a dropped/garbled stream) is a transient anomaly, not a + # permanent client error. Raise a RETRYABLE error -- a bare ValueError + # here would be downgraded to a permanent lf.LMError by _parse_response. + raise lf.TemporaryLMError( + 'Anthropic SSE stream produced no message (empty body or no ' + 'message_start event); retrying.' + ) + # Assemble content blocks in index order. + message['content'] = [blocks[i] for i in sorted(blocks)] + # Merge usage: message_start carries input_tokens (and an initial + # output_tokens); message_delta carries the FINAL output_tokens. + usage = dict(message.get('usage') or {}) + usage.update(final_usage) + message['usage'] = usage + # TERMINAL SENTINEL: only accept a stream that actually completed. A + # cleanly-closed-but-incomplete 200 (no message_stop, or stop_reason still + # null because message_delta never arrived) would otherwise be returned as + # silently truncated text with an under-counted output_tokens. Require BOTH + # the message_stop terminator AND a non-null stop_reason; otherwise raise a + # RETRYABLE error so the request is retried rather than silently accepted. + if not saw_message_stop or message.get('stop_reason') is None: + raise lf.TemporaryLMError( + 'Anthropic SSE stream ended without a terminal message_stop and a ' + 'non-null stop_reason (truncated/incomplete 200 response); retrying.' + ) + return message + def _error(self, status_code: int, content: str) -> lf.LMError: - if status_code == 413 and b'Prompt is too long' in content: + if status_code == 413 and b'Prompt is too long' in content: # pyrefly: ignore[unsupported-operation] return lf.ContextLimitError(f'{status_code}: {content}') - if status_code == 400 and b'prompt is too long' in content: + if status_code == 400 and b'prompt is too long' in content: # pyrefly: ignore[unsupported-operation] return lf.ContextLimitError(f'{status_code}: {content}') return super()._error(status_code, content) @@ -1166,6 +1508,12 @@ class Claude46(Anthropic): # pylint: disable=invalid-name +class Claude5Opus(Anthropic): + """Claude Opus 5 model.""" + + model = 'claude-opus-5' + + class Claude48Opus(Anthropic): """Claude Opus 4.8 model.""" diff --git a/langfun/core/llms/anthropic_test.py b/langfun/core/llms/anthropic_test.py index 5f4bed68..c6a41ea4 100644 --- a/langfun/core/llms/anthropic_test.py +++ b/langfun/core/llms/anthropic_test.py @@ -17,6 +17,7 @@ import copy import datetime import os +import time from typing import Any import unittest from unittest import mock @@ -788,6 +789,176 @@ def test_request_caching_works_on_opus_subclass(self): ) +class EffortConfigHonestyTest(unittest.TestCase): + """Tests that thinking-effort config never diverges from the wire silently. + + Two config surfaces feed the Anthropic effort knob: + * `Anthropic.effort` (model-level, vocabulary low..max), and + * `LMSamplingOptions.reasoning_effort` (per-call, cross-provider, + vocabulary low..high). + Whenever the configured effort cannot reach the API unchanged -- because the + per-call knob downgrades it, or because the request takes a code path that + does not carry effort at all -- the model must say so out loud instead of + quietly sending something else. + """ + + def _warning_logs(self): + """Captures warnings emitted through langfun's logging primitive. + + Returns: + An `assertLogs` context manager bound to the exact logger object that + `pg.logging.warning` writes to. Binding to the logger object (rather + than the root logger) keeps the assertion hermetic regardless of which + logger PyGlove is configured with in a given environment. + """ + return self.assertLogs(pg.logging.get_logger(), level='WARNING') + + def _no_warning_logs(self): + return self.assertNoLogs(pg.logging.get_logger(), level='WARNING') + + # --- B1: per-call `reasoning_effort` must not silently downgrade `effort`. + + def test_reasoning_effort_downgrade_warns_and_names_both_values(self): + lm = anthropic.Claude5Opus(api_key='fake', thinking=True, effort='max') + with self._warning_logs() as logs: + args = lm._request_args( + lf.LMSamplingOptions(max_tokens=1024, reasoning_effort='high') + ) + # Per-call precedence is preserved (blessed behavior), but it is loud. + self.assertEqual(args['output_config'], {'effort': 'high'}) + message = '\n'.join(logs.output) + self.assertIn("reasoning_effort='high'", message) + self.assertIn("effort='max'", message) + + def test_reasoning_effort_upgrade_is_silent(self): + lm = anthropic.Claude5Opus(api_key='fake', thinking=True, effort='low') + with self._no_warning_logs(): + args = lm._request_args( + lf.LMSamplingOptions(max_tokens=1024, reasoning_effort='high') + ) + self.assertEqual(args['output_config'], {'effort': 'high'}) + + def test_reasoning_effort_equal_tier_is_silent(self): + lm = anthropic.Claude5Opus(api_key='fake', thinking=True, effort='low') + with self._no_warning_logs(): + args = lm._request_args( + lf.LMSamplingOptions(max_tokens=1024, reasoning_effort='low') + ) + self.assertEqual(args['output_config'], {'effort': 'low'}) + + def test_reasoning_effort_over_default_effort_is_silent(self): + """A class default is not a user decision, so overriding it is not news.""" + lm = anthropic.Claude5Opus(api_key='fake', thinking=True) + self.assertEqual(lm.effort, 'high') + with self._no_warning_logs(): + args = lm._request_args( + lf.LMSamplingOptions(max_tokens=1024, reasoning_effort='low') + ) + self.assertEqual(args['output_config'], {'effort': 'low'}) + + # --- B1: the wide vocabulary stays reachable through `effort`. + + def test_xhigh_and_max_reachable_via_model_effort(self): + for effort in ('xhigh', 'max'): + with self.subTest(effort=effort): + lm = anthropic.Claude5Opus(api_key='fake', thinking=True, effort=effort) + args = lm._request_args(lf.LMSamplingOptions(max_tokens=1024)) + self.assertEqual(args['output_config'], {'effort': effort}) + + # --- B2: an explicit `effort` must never be dropped silently. + + def test_explicit_effort_dropped_on_non_adaptive_model_warns(self): + lm = anthropic.Claude46Opus(api_key='fake', thinking=True, effort='max') + with self._warning_logs() as logs: + args = lm._request_args( + lf.LMSamplingOptions(max_tokens=1024, max_thinking_tokens=1024) + ) + # Existing behavior is untouched: manual budget, no effort on the wire. + self.assertEqual( + args['thinking'], {'type': 'enabled', 'budget_tokens': 1024} + ) + self.assertNotIn('output_config', args) + message = '\n'.join(logs.output) + self.assertIn("effort='max'", message) + self.assertIn('ignored', message) + + def test_explicit_effort_dropped_when_thinking_off_warns(self): + lm = anthropic.Claude5Opus(api_key='fake', effort='max') + self.assertIsNone(lm.thinking) + with self._warning_logs() as logs: + args = lm._request_args(lf.LMSamplingOptions(max_tokens=1024)) + self.assertNotIn('thinking', args) + self.assertNotIn('output_config', args) + message = '\n'.join(logs.output) + self.assertIn("effort='max'", message) + self.assertIn('ignored', message) + + def test_explicit_effort_dropped_when_thinking_disabled_warns(self): + """`thinking=False` wins over `max_thinking_tokens`, so effort is dropped. + + This is the sub-case that keeps the warning text honest: the thinking gate + is off while `max_thinking_tokens` IS set, so the message must state the + real reason (the gate is off) and must not claim that + `max_thinking_tokens` is unset. + """ + lm = anthropic.Claude5Opus(api_key='fake', thinking=False, effort='max') + with self._warning_logs() as logs: + args = lm._request_args( + lf.LMSamplingOptions(max_tokens=1024, max_thinking_tokens=4096) + ) + self.assertNotIn('thinking', args) + self.assertNotIn('output_config', args) + message = '\n'.join(logs.output) + self.assertIn("effort='max'", message) + self.assertIn('thinking is off for this request', message) + self.assertNotIn('`max_thinking_tokens` is unset', message) + + def test_default_effort_with_thinking_off_is_silent(self): + """The overwhelmingly common path must stay quiet, or nobody reads logs.""" + lm = anthropic.Claude5Opus(api_key='fake') + with self._no_warning_logs(): + args = lm._request_args(lf.LMSamplingOptions(max_tokens=1024)) + self.assertNotIn('output_config', args) + + def test_effort_none_is_silent(self): + """`effort=None` asks for no effort at all; dropping it is not divergence.""" + lm = anthropic.Claude5Opus(api_key='fake', effort=None) + with self._no_warning_logs(): + args = lm._request_args(lf.LMSamplingOptions(max_tokens=1024)) + self.assertNotIn('output_config', args) + + def test_explicit_effort_consumed_by_adaptive_path_is_silent(self): + lm = anthropic.Claude5Opus(api_key='fake', thinking=True, effort='max') + with self._no_warning_logs(): + args = lm._request_args(lf.LMSamplingOptions(max_tokens=1024)) + self.assertEqual(args['output_config'], {'effort': 'max'}) + + def test_explicit_effort_consumed_via_max_thinking_tokens_backcompat(self): + """`max_thinking_tokens` implies thinking=True, so effort is consumed.""" + lm = anthropic.Claude5Opus(api_key='fake', effort='max') + with self._no_warning_logs(): + args = lm._request_args( + lf.LMSamplingOptions(max_tokens=1024, max_thinking_tokens=1024) + ) + self.assertEqual(args['output_config'], {'effort': 'max'}) + + def test_effort_rebound_after_construction_is_honored(self): + """Divergence detection is value-based, so `rebind` is covered too.""" + lm = anthropic.Claude5Opus(api_key='fake', thinking=True) + lm.rebind(effort='max', skip_notification=True, raise_on_no_change=False) + with self._warning_logs() as logs: + args = lm._request_args( + lf.LMSamplingOptions(max_tokens=1024, reasoning_effort='low') + ) + self.assertEqual(args['output_config'], {'effort': 'low'}) + self.assertIn("effort='max'", '\n'.join(logs.output)) + + def test_invalid_reasoning_effort_still_fails_loud(self): + """Guards the existing loud path: bad per-call vocabulary must raise.""" + with self.assertRaises(ValueError): + lf.LMSamplingOptions(max_tokens=1024, reasoning_effort='max') + + class Claude48OpusTest(unittest.TestCase): """Tests for Claude Opus 4.8 model support.""" @@ -1058,5 +1229,350 @@ def test_adversarial_long_and_unicode(self): self.assertIn('cache_control', request['messages'][-1]['content'][-1]) +# A realistic Anthropic Messages SSE byte stream: one thinking block followed +# by one text block, then a message_delta carrying the final stop_reason and +# usage.output_tokens, then message_stop. Mirrors what Vertex/Anthropic emits +# when the request body has stream=True. +_SSE_STREAM = ( + 'event: message_start\n' + 'data: {"type":"message_start","message":{"id":"msg_1","type":"message",' + '"role":"assistant","model":"claude-opus-4-6","content":[],' + '"stop_reason":null,"stop_sequence":null,' + '"usage":{"input_tokens":25,"output_tokens":1}}}\n' + '\n' + 'event: content_block_start\n' + 'data: {"type":"content_block_start","index":0,' + '"content_block":{"type":"thinking","thinking":""}}\n' + '\n' + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":0,' + '"delta":{"type":"thinking_delta","thinking":"Let me think. "}}\n' + '\n' + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":0,' + '"delta":{"type":"thinking_delta","thinking":"Done."}}\n' + '\n' + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":0,' + '"delta":{"type":"signature_delta","signature":"sig=="}}\n' + '\n' + 'event: content_block_stop\n' + 'data: {"type":"content_block_stop","index":0}\n' + '\n' + 'event: content_block_start\n' + 'data: {"type":"content_block_start","index":1,' + '"content_block":{"type":"text","text":""}}\n' + '\n' + 'event: ping\n' + 'data: {"type":"ping"}\n' + '\n' + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":1,' + '"delta":{"type":"text_delta","text":"Hello"}}\n' + '\n' + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":1,' + '"delta":{"type":"text_delta","text":", world!"}}\n' + '\n' + 'event: content_block_stop\n' + 'data: {"type":"content_block_stop","index":1}\n' + '\n' + 'event: message_delta\n' + 'data: {"type":"message_delta",' + '"delta":{"stop_reason":"end_turn","stop_sequence":null},' + '"usage":{"output_tokens":42}}\n' + '\n' + 'event: message_stop\n' + 'data: {"type":"message_stop"}\n' + '\n' +) + + +def mock_sse_post(sse_text, chunk_size_bytes=16, delay_per_chunk=0.0): + """Returns a mock session.post that streams `sse_text` in chunks.""" + raw = sse_text.encode('utf-8') + + def _mock_post(url, json=None, timeout=None, stream=False, **kwargs): + del url, json, timeout, stream, kwargs + response = requests.Response() + response.status_code = 200 + response.headers['Content-Type'] = 'text/event-stream' + response._content = False # pylint: disable=protected-access + + def _iter(chunk_size=1, decode_unicode=False): + del chunk_size, decode_unicode + for i in range(0, len(raw), chunk_size_bytes): + if delay_per_chunk > 0: + time.sleep(delay_per_chunk) + yield raw[i : i + chunk_size_bytes] + + response.iter_content = _iter + response.close = lambda: None + return response + + return _mock_post + + +class AnthropicStreamingSSETest(unittest.TestCase): + """TDD tests for the streaming-body SSE reassembler.""" + + def test_reassemble_sse_matches_buffered_equivalent(self): + """Reassembled SSE dict must equal the buffered (non-streaming) dict.""" + lm = anthropic.Claude46Opus(api_key='fake_key') + reassembled = lm._reassemble_sse(_SSE_STREAM) + expected_buffered = { + 'id': 'msg_1', + 'type': 'message', + 'role': 'assistant', + 'model': 'claude-opus-4-6', + 'content': [ + { + 'type': 'thinking', + 'thinking': 'Let me think. Done.', + 'signature': 'sig==', + }, + {'type': 'text', 'text': 'Hello, world!'}, + ], + 'stop_reason': 'end_turn', + 'stop_sequence': None, + 'usage': {'input_tokens': 25, 'output_tokens': 42}, + } + self.assertEqual(reassembled, expected_buffered) + + def test_streaming_call_end_to_end(self): + """A full streamed call rebuilds text + usage through the client path.""" + with mock.patch('requests.Session.post') as mock_request: + mock_request.side_effect = mock_sse_post(_SSE_STREAM, chunk_size_bytes=8) + lm = anthropic.Claude46Opus(api_key='fake_key') + response = lm('hello') + # stream=True must be sent in the request body. + _, kwargs = mock_request.call_args + self.assertTrue(kwargs.get('stream')) + self.assertEqual(kwargs['json']['stream'], True) + # Text content correctly reassembled. + self.assertEqual(response.text, 'Hello, world!') + # Usage rebuilt: input_tokens from message_start, output_tokens from + # the final message_delta (NOT the initial output_tokens=1). + self.assertEqual(response.usage.prompt_tokens, 25) + self.assertEqual(response.usage.completion_tokens, 42) + self.assertEqual(response.usage.total_tokens, 67) + + def test_buffered_json_body_still_parses(self): + """A single buffered JSON body (non-SSE) must still parse (auto-detect).""" + with mock.patch('requests.Session.post') as mock_request: + mock_request.side_effect = mock_requests_post + lm = anthropic.Claude46Opus(api_key='fake_key') + response = lm('hello') + self.assertRegex(response.text, 'hello.*') + self.assertEqual(response.usage.completion_tokens, 1) + + def test_tool_use_input_json_reassembled(self): + """input_json_delta fragments reassemble into a parsed tool input dict.""" + sse = ( + 'event: message_start\n' + 'data: {"type":"message_start","message":{"id":"m","type":"message",' + '"role":"assistant","model":"claude-opus-4-6","content":[],' + '"usage":{"input_tokens":5,"output_tokens":1}}}\n\n' + 'event: content_block_start\n' + 'data: {"type":"content_block_start","index":0,' + '"content_block":{"type":"tool_use","id":"t1","name":"calc",' + '"input":{}}}\n\n' + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":0,' + '"delta":{"type":"input_json_delta","partial_json":"{\\"x\\":"}}\n\n' + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":0,' + '"delta":{"type":"input_json_delta","partial_json":"7}"}}\n\n' + 'event: content_block_stop\n' + 'data: {"type":"content_block_stop","index":0}\n\n' + 'event: message_delta\n' + 'data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},' + '"usage":{"output_tokens":9}}\n\n' + 'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + lm = anthropic.Claude46Opus(api_key='fake_key') + reassembled = lm._reassemble_sse(sse) + self.assertEqual(reassembled['content'][0]['input'], {'x': 7}) + self.assertEqual(reassembled['stop_reason'], 'tool_use') + self.assertEqual(reassembled['usage']['output_tokens'], 9) + + def test_malformed_stream_without_message_start_raises(self): + """An SSE stream missing message_start is an unusable 200 => retryable.""" + sse = ( + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":0,' + '"delta":{"type":"text_delta","text":"orphan"}}\n\n' + ) + lm = anthropic.Claude46Opus(api_key='fake_key') + # A 200 whose body never produced a message is a transient streaming + # anomaly, not a permanent client error; it must be retryable. + with self.assertRaises(lf.TemporaryLMError): + lm._reassemble_sse(sse) + # And via the full parse path it surfaces (the retryable error propagates + # with max_attempts=1; type precision is asserted directly above). + with mock.patch('requests.Session.post') as mock_request: + mock_request.side_effect = mock_sse_post(sse) + with self.assertRaises(Exception): + lm('hello', max_attempts=1) + + def test_truncated_stream_without_message_stop_raises_retryable(self): + """Cleanly-closed-but-incomplete 200 (no message_stop) => retryable. + + Defect (terminal sentinel): a stream carrying message_start + + content_block_delta but NO message_delta/message_stop leaves stop_reason + null. Such a truncated-but-200 response must be RETRYABLE, never silently + returned as truncated text with a low output_tokens count. + """ + truncated = ( + 'event: message_start\n' + 'data: {"type":"message_start","message":{"id":"m","type":"message",' + '"role":"assistant","model":"claude-opus-4-6","content":[],' + '"stop_reason":null,"stop_sequence":null,' + '"usage":{"input_tokens":10,"output_tokens":1}}}\n\n' + 'event: content_block_start\n' + 'data: {"type":"content_block_start","index":0,' + '"content_block":{"type":"text","text":""}}\n\n' + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":0,' + '"delta":{"type":"text_delta","text":"partial answer"}}\n\n' + ) + lm = anthropic.Claude46Opus(api_key='fake_key') + with self.assertRaises(lf.TemporaryLMError): + lm._reassemble_sse(truncated) + + def test_in_stream_overloaded_error_is_retryable(self): + """In-stream `event: error` overloaded_error (HTTP 200) => retryable. + + Defect (retryability): mirrors the buffered 529 -> TemporaryLMError + behavior so a transient Anthropic overload mid-stream stays retryable. + """ + sse = ( + 'event: message_start\n' + 'data: {"type":"message_start","message":{"id":"m","type":"message",' + '"role":"assistant","model":"claude-opus-4-6","content":[],' + '"usage":{"input_tokens":3,"output_tokens":1}}}\n\n' + 'event: error\n' + 'data: {"type":"error","error":{"type":"overloaded_error",' + '"message":"Overloaded"}}\n\n' + ) + lm = anthropic.Claude46Opus(api_key='fake_key') + with self.assertRaises(lf.TemporaryLMError): + lm._reassemble_sse(sse) + + def test_in_stream_rate_limit_error_is_retryable(self): + """In-stream `event: error` rate_limit_error (HTTP 200) => RateLimitError. + + Defect (retryability): mirrors the buffered 429 -> RateLimitError behavior. + """ + sse = ( + 'event: message_start\n' + 'data: {"type":"message_start","message":{"id":"m","type":"message",' + '"role":"assistant","model":"claude-opus-4-6","content":[],' + '"usage":{"input_tokens":3,"output_tokens":1}}}\n\n' + 'event: error\n' + 'data: {"type":"error","error":{"type":"rate_limit_error",' + '"message":"Rate limited"}}\n\n' + ) + lm = anthropic.Claude46Opus(api_key='fake_key') + with self.assertRaises(lf.RateLimitError): + lm._reassemble_sse(sse) + + def test_empty_sse_200_is_retryable(self): + """A 200 SSE body with no events (keep-alives only) => retryable. + + Defect (retryability): an empty/no-events SSE 200 must NOT be downgraded + to a permanent lf.LMError (which is what a bare ValueError becomes in + _parse_response); it is a transient anomaly and must be retried. + """ + sse = ': keep-alive\n\nevent: ping\ndata: {"type":"ping"}\n\n' + lm = anthropic.Claude46Opus(api_key='fake_key') + with self.assertRaises(lf.TemporaryLMError): + lm._reassemble_sse(sse) + + def test_tool_use_input_json_split_across_three_frames_reassembles(self): + """Adversarial re-verification that split tool-input JSON reassembles. + + Splits the tool input JSON across THREE input_json_delta frames at awkward + boundaries (mid-key, mid-value). This probes whether the partial-JSON + accumulation defect is live; the accumulator must reconstruct the full + input dict. + """ + sse = ( + 'event: message_start\n' + 'data: {"type":"message_start","message":{"id":"m","type":"message",' + '"role":"assistant","model":"claude-opus-4-6","content":[],' + '"usage":{"input_tokens":5,"output_tokens":1}}}\n\n' + 'event: content_block_start\n' + 'data: {"type":"content_block_start","index":0,' + '"content_block":{"type":"tool_use","id":"t1","name":"weather",' + '"input":{}}}\n\n' + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":0,' + '"delta":{"type":"input_json_delta","partial_json":"{\\"unit\\""}}\n\n' + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":0,' + '"delta":{"type":"input_json_delta","partial_json":":\\"C\\",\\"temp"}}' + '\n\n' + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":0,' + '"delta":{"type":"input_json_delta","partial_json":"erature\\":21}"}}' + '\n\n' + 'event: content_block_stop\n' + 'data: {"type":"content_block_stop","index":0}\n\n' + 'event: message_delta\n' + 'data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},' + '"usage":{"output_tokens":9}}\n\n' + 'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + lm = anthropic.Claude46Opus(api_key='fake_key') + reassembled = lm._reassemble_sse(sse) + self.assertEqual( + reassembled['content'][0]['input'], {'unit': 'C', 'temperature': 21} + ) + self.assertEqual(reassembled['stop_reason'], 'tool_use') + + def test_malformed_content_delta_does_not_silently_truncate(self): + """A malformed `data:` content delta must NOT be silently dropped. + + Defect (terminal integrity / silent corruption): a well-formed Anthropic + SSE stream emits exactly one valid-JSON payload per `data:` line. If a + `content_block_delta` line is corrupt/truncated JSON, silently skipping it + drops that text fragment while `message_stop` + a non-null `stop_reason` + still arrive -- so the terminal sentinel passes and the reassembler would + return SILENTLY TRUNCATED text (here "Hello " instead of "Hello world"). + The buffered path fails loudly on a malformed body (json.loads -> + ValueError -> LMError) and never returns partial content; the streaming + path must preserve that no-silent-corruption guarantee by raising a + RETRYABLE error instead of returning the truncated text. + """ + sse = ( + 'event: message_start\n' + 'data: {"type":"message_start","message":{"id":"m","type":"message",' + '"role":"assistant","model":"claude-opus-4-6","content":[],' + '"stop_reason":null,"stop_sequence":null,' + '"usage":{"input_tokens":10,"output_tokens":1}}}\n\n' + 'event: content_block_start\n' + 'data: {"type":"content_block_start","index":0,' + '"content_block":{"type":"text","text":""}}\n\n' + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":0,' + '"delta":{"type":"text_delta","text":"Hello "}}\n\n' + # Corrupt/truncated JSON payload (missing closing braces) that would + # have carried "world"; must not be silently discarded. + 'event: content_block_delta\n' + 'data: {"type":"content_block_delta","index":0,' + '"delta":{"type":"text_delta","text":"world"\n\n' + 'event: content_block_stop\n' + 'data: {"type":"content_block_stop","index":0}\n\n' + 'event: message_delta\n' + 'data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},' + '"usage":{"output_tokens":5}}\n\n' + 'event: message_stop\ndata: {"type":"message_stop"}\n\n' + ) + lm = anthropic.Claude46Opus(api_key='fake_key') + with self.assertRaises(lf.TemporaryLMError): + lm._reassemble_sse(sse) + + if __name__ == '__main__': unittest.main() diff --git a/langfun/core/llms/azure_openai.py b/langfun/core/llms/azure_openai.py index 56f04125..b3f84113 100644 --- a/langfun/core/llms/azure_openai.py +++ b/langfun/core/llms/azure_openai.py @@ -90,7 +90,7 @@ def _initialize(self): ) @property - def api_endpoint(self) -> str: + def api_endpoint(self) -> str: # pyrefly: ignore[bad-override] return self._api_endpoint @property diff --git a/langfun/core/llms/cache/base.py b/langfun/core/llms/cache/base.py index 8bfa6321..4c4669b5 100644 --- a/langfun/core/llms/cache/base.py +++ b/langfun/core/llms/cache/base.py @@ -136,7 +136,7 @@ def _delete(self, model_id: str, key: str) -> bool: def _sym_clone(self, deep: bool, memo: Any = None) -> 'LMCacheBase': v = super()._sym_clone(deep, memo) v._stats = self._stats # pylint: disable=protected-access - return v + return v # pyrefly: ignore[bad-return] def default_key(lm: lf.LanguageModel, prompt: lf.Message, seed: int) -> Any: diff --git a/langfun/core/llms/cache/in_memory.py b/langfun/core/llms/cache/in_memory.py index 5bdb25e2..402d9c34 100644 --- a/langfun/core/llms/cache/in_memory.py +++ b/langfun/core/llms/cache/in_memory.py @@ -147,7 +147,7 @@ def reset(self, model_id: str | None = None) -> None: def _sym_clone(self, deep: bool, memo: Any = None) -> 'InMemory': v = super()._sym_clone(deep, memo) v._cache = self._cache # pylint: disable=protected-access - return v + return v # pyrefly: ignore[bad-return] def save(self, path: str | None = None) -> None: """Saves the in-memory cache.""" diff --git a/langfun/core/llms/compositional.py b/langfun/core/llms/compositional.py index 603ff040..55b2c9e0 100644 --- a/langfun/core/llms/compositional.py +++ b/langfun/core/llms/compositional.py @@ -69,18 +69,18 @@ def _on_bound(self): if parent_non_default: for c in self.candidates: c.sampling_options.rebind( - parent_non_default, notify_parents=False, raise_on_no_change=False + parent_non_default, notify_parents=False, raise_on_no_change=False # pyrefly: ignore[bad-argument-type] ) @property - def model_id(self) -> str: + def model_id(self) -> str: # pyrefly: ignore[bad-override] model_ids = ', '.join( sorted(c.model_id for c in self.candidates) ) return f'RandomChoice({model_ids})' @property - def resource_id(self) -> str: + def resource_id(self) -> str: # pyrefly: ignore[bad-override] resource_ids = ', '.join( sorted(c.resource_id for c in self.candidates) ) diff --git a/langfun/core/llms/deepseek.py b/langfun/core/llms/deepseek.py index 0441fd49..f33b519f 100644 --- a/langfun/core/llms/deepseek.py +++ b/langfun/core/llms/deepseek.py @@ -33,7 +33,7 @@ class DeepSeekModelInfo(lf.ModelInfo): error_codes='https://api-docs.deepseek.com/quick_start/error_codes', ) - provider: Final[str] = 'DeepSeek' # pylint: disable=invalid-name + provider: Final[str] = 'DeepSeek' # pylint: disable=invalid-name # pyrefly: ignore[bad-override] api_model_name: Annotated[ str, @@ -169,7 +169,7 @@ def _request_args( return args @classmethod - def dir(cls): + def dir(cls): # pyrefly: ignore[bad-override] return [m.model_id for m in SUPPORTED_MODELS if m.in_service] diff --git a/langfun/core/llms/fake.py b/langfun/core/llms/fake.py index a8d1f101..29b55b53 100644 --- a/langfun/core/llms/fake.py +++ b/langfun/core/llms/fake.py @@ -123,7 +123,7 @@ class StaticMapping(Fake): ] def _response_from(self, prompt: lf.Message) -> lf.Message: - return lf.AIMessage.from_value(self.mapping[prompt]) + return lf.AIMessage.from_value(self.mapping[prompt]) # pyrefly: ignore[bad-index] @lf.use_init_args(['sequence']) diff --git a/langfun/core/llms/gemini.py b/langfun/core/llms/gemini.py index b2aacdcb..d5ae729c 100644 --- a/langfun/core/llms/gemini.py +++ b/langfun/core/llms/gemini.py @@ -157,7 +157,7 @@ def estimate_cost(self, usage: lf.LMSamplingUsage) -> float | None: # Add cost for output tokens cost += ( - self.cost_per_1m_output_tokens_with_prompt_longer_than_128k + self.cost_per_1m_output_tokens_with_prompt_longer_than_128k # pyrefly: ignore[unsupported-operation] * usage.completion_tokens ) @@ -234,6 +234,28 @@ def estimate_cost(self, usage: lf.LMSamplingUsage) -> float | None: max_tokens_per_minute=4_000_000, ), ), + GeminiModelInfo( + model_id='gemini-3.1-flash-lite', + in_service=True, + provider=pg.oneof(['Google GenAI', 'VertexAI']), + model_type='instruction-tuned', + description='Gemini 3.1 Flash Lite.', + release_date=datetime.datetime(2026, 5, 7), + input_modalities=GeminiModelInfo.ALL_SUPPORTED_INPUT_TYPES, + context_length=lf.ModelInfo.ContextLength( + max_input_tokens=1_048_576, + max_output_tokens=65_536, + ), + pricing=GeminiModelInfo.Pricing( + cost_per_1m_cached_input_tokens=0.025, + cost_per_1m_input_tokens=0.25, + cost_per_1m_output_tokens=1.5, + ), + rate_limits=lf.ModelInfo.RateLimits( + max_requests_per_minute=2000, + max_tokens_per_minute=4_000_000, + ), + ), # Gemini 3 Pro Preview GeminiModelInfo( model_id='gemini-3-pro-preview', @@ -309,6 +331,58 @@ def estimate_cost(self, usage: lf.LMSamplingUsage) -> float | None: max_tokens_per_minute=4_000_000, ), ), + # Gemini 3.5 Flash + GeminiModelInfo( + model_id='gemini-3.5-flash', + in_service=True, + provider=pg.oneof(['Google GenAI', 'VertexAI']), + model_type='instruction-tuned', + description=( + 'Gemini 3.5 Flash: High-efficiency, low-latency multimodal' + ' model optimized for agentic workflows.' + ), + release_date=datetime.datetime(2026, 5, 19), + input_modalities=GeminiModelInfo.ALL_SUPPORTED_INPUT_TYPES, + context_length=lf.ModelInfo.ContextLength( + max_input_tokens=1_048_576, + max_output_tokens=65_536, + ), + pricing=GeminiModelInfo.Pricing( + cost_per_1m_cached_input_tokens=0.15, + cost_per_1m_input_tokens=1.50, + cost_per_1m_output_tokens=9.00, + ), + rate_limits=lf.ModelInfo.RateLimits( + max_requests_per_minute=2_000, + max_tokens_per_minute=4_000_000, + ), + ), + # Gemini 3.7 Flash + GeminiModelInfo( + model_id='gemini-3.7-flash', + in_service=True, + provider=pg.oneof(['Google GenAI', 'VertexAI']), + model_type='instruction-tuned', + description=( + 'Gemini 3.7 Flash: High-efficiency, low-latency multimodal' + ' model optimized for agentic workflows.' + ), + release_date=datetime.datetime(2026, 8, 13), + input_modalities=GeminiModelInfo.ALL_SUPPORTED_INPUT_TYPES, + context_length=lf.ModelInfo.ContextLength( + max_input_tokens=1_048_576, + max_output_tokens=65_536, + ), + pricing=GeminiModelInfo.Pricing( + cost_per_1m_cached_input_tokens=0.075, + cost_per_1m_input_tokens=0.75, + cost_per_1m_output_tokens=3.75, + ), + rate_limits=lf.ModelInfo.RateLimits( + max_requests_per_minute=2_000, + max_tokens_per_minute=4_000_000, + ), + ), # Gemini 2.5 Flash GeminiModelInfo( model_id='gemini-2.5-flash', @@ -896,11 +970,11 @@ def model_info(self) -> GeminiModelInfo: return _SUPPORTED_MODELS_BY_ID[self.model] @classmethod - def dir(cls): + def dir(cls): # pyrefly: ignore[bad-override] return [m.model_id for m in SUPPORTED_MODELS if m.in_service] @property - def headers(self): + def headers(self): # pyrefly: ignore[bad-override] return { 'Content-Type': 'application/json; charset=utf-8', } @@ -915,7 +989,7 @@ def modality_conversion(chunk: str | lf.Modality) -> Any: if isinstance(chunk, lf_modalities.Mime): try: return chunk.make_compatible( - self.model_info.input_modalities + ['text/plain'] + self.model_info.input_modalities + ['text/plain'] # pyrefly: ignore[unsupported-operation] ) except lf.ModalityError as e: raise lf.ModalityError(f'Unsupported modality: {chunk!r}') from e @@ -931,7 +1005,7 @@ def modality_conversion(chunk: str | lf.Modality) -> Any: contents.append( prompt.as_format('gemini', chunk_preprocessor=modality_conversion) ) - request['contents'] = contents + request['contents'] = contents # pyrefly: ignore[bad-assignment] request['toolConfig'] = { 'functionCallingConfig': { 'mode': 'NONE', @@ -964,8 +1038,8 @@ def _generation_config( ) json_schema = pg.to_json(json_schema) config['responseSchema'] = json_schema - config['responseMimeType'] = 'application/json' - prompt.metadata.formatted_text = ( + config['responseMimeType'] = 'application/json' # pyrefly: ignore[bad-assignment] + prompt.metadata.formatted_text = ( # pyrefly: ignore[missing-attribute] prompt.text + '\n\n [RESPONSE FORMAT (not part of prompt)]\n' + pg.to_json_str(json_schema, json_indent=2) @@ -977,7 +1051,7 @@ def _generation_config( if options.thinking_level is not None: thinking_config_data['thinkingLevel'] = options.thinking_level if thinking_config_data: - config['thinkingConfig'] = thinking_config_data + config['thinkingConfig'] = thinking_config_data # pyrefly: ignore[bad-assignment] # This is the new feature since Gemini 3. # Skip for image generation models as they don't support mediaResolution. @@ -988,7 +1062,7 @@ def _generation_config( self.response_modalities and 'IMAGE' in self.response_modalities ) ): - config['mediaResolution'] = 'MEDIA_RESOLUTION_HIGH' + config['mediaResolution'] = 'MEDIA_RESOLUTION_HIGH' # pyrefly: ignore[bad-assignment] if self.response_modalities: config['responseModalities'] = self.response_modalities @@ -1036,10 +1110,10 @@ def result(self, json: dict[str, Any]) -> lf.LMSamplingResult: def _error(self, status_code: int, content: str) -> lf.LMError: if status_code == 400 and ( - b'exceeds the maximum number of tokens' in content - or b'Reduce the input token count and try again.' in content - or b'Request payload size exceeds the limit' in content - or b'Request contains text fields that are too large' in content + b'exceeds the maximum number of tokens' in content # pyrefly: ignore[unsupported-operation] + or b'Reduce the input token count and try again.' in content # pyrefly: ignore[unsupported-operation] + or b'Request payload size exceeds the limit' in content # pyrefly: ignore[unsupported-operation] + or b'Request contains text fields that are too large' in content # pyrefly: ignore[unsupported-operation] ): return lf.ContextLimitError(f'{status_code}: {content}') return super()._error(status_code, content) diff --git a/langfun/core/llms/google_genai.py b/langfun/core/llms/google_genai.py index 3c911084..e9ef5b7f 100644 --- a/langfun/core/llms/google_genai.py +++ b/langfun/core/llms/google_genai.py @@ -84,8 +84,8 @@ class GenAI(gemini.Gemini): ] = 'v1beta' @functools.cached_property - def model_info(self) -> lf.ModelInfo: - return super().model_info.clone( + def model_info(self) -> lf.ModelInfo: # pyrefly: ignore[bad-override] + return super().model_info.clone( # pyrefly: ignore[bad-return] override=dict(provider='Google GenAI') ) @@ -96,7 +96,7 @@ def session(self): return s @property - def api_endpoint(self) -> str: + def api_endpoint(self) -> str: # pyrefly: ignore[bad-override] api_key = self.api_key or os.environ.get('GOOGLE_API_KEY', None) if not api_key: raise ValueError( @@ -155,12 +155,24 @@ class Gemini3FlashPreview(GenAI): model = 'gemini-3-flash-preview' +class Gemini37Flash(GenAI): + """Gemini 3.7 Flash GA model.""" + + model = 'gemini-3.7-flash' + + class Gemini31FlashLitePreview(GenAI): """Gemini 3.1 Flash Lite Preview model.""" model = 'gemini-3.1-flash-lite-preview' +class Gemini31FlashLite(GenAI): + """Gemini 3.1 Flash Lite model.""" + + model = 'gemini-3.1-flash-lite' + + class Gemini25FlashImagePreview(GenAI): """Gemini 2.5 Flash Image Preview model.""" model = 'gemini-2.5-flash-image-preview' diff --git a/langfun/core/llms/google_genai_test.py b/langfun/core/llms/google_genai_test.py index 5c994855..bfa179be 100644 --- a/langfun/core/llms/google_genai_test.py +++ b/langfun/core/llms/google_genai_test.py @@ -35,6 +35,13 @@ def test_basics(self): self.assertEqual(lm.resource_id, 'google_genai://gemini-1.5-pro-001') del os.environ['GOOGLE_API_KEY'] + def test_gemini_37_flash(self): + os.environ['GOOGLE_API_KEY'] = 'abc' + lm = google_genai.Gemini37Flash() + self.assertEqual(lm.model_id, 'gemini-3.7-flash') + self.assertEqual(lm.resource_id, 'google_genai://gemini-3.7-flash') + del os.environ['GOOGLE_API_KEY'] + def test_gemini_31_flash_lite_preview(self): os.environ['GOOGLE_API_KEY'] = 'abc' lm = google_genai.Gemini31FlashLitePreview() @@ -44,6 +51,13 @@ def test_gemini_31_flash_lite_preview(self): ) del os.environ['GOOGLE_API_KEY'] + def test_gemini_31_flash_lite(self): + os.environ['GOOGLE_API_KEY'] = 'abc' + lm = google_genai.Gemini31FlashLite() + self.assertEqual(lm.model_id, 'gemini-3.1-flash-lite') + self.assertEqual(lm.resource_id, 'google_genai://gemini-3.1-flash-lite') + del os.environ['GOOGLE_API_KEY'] + def test_lm_get(self): self.assertIsInstance( lf.LanguageModel.get('google_genai://gemini-1.5-pro'), diff --git a/langfun/core/llms/groq.py b/langfun/core/llms/groq.py index 29b5c6f3..f21a0c88 100644 --- a/langfun/core/llms/groq.py +++ b/langfun/core/llms/groq.py @@ -33,7 +33,7 @@ class GroqModelInfo(lf.ModelInfo): error_codes='https://console.groq.com/docs/errors', ) - provider: Final[str] = 'Groq' # pylint: disable=invalid-name + provider: Final[str] = 'Groq' # pylint: disable=invalid-name # pyrefly: ignore[bad-override] SUPPORTED_MODELS = [ diff --git a/langfun/core/llms/llama_cpp.py b/langfun/core/llms/llama_cpp.py index 1dd83ecb..03d04df6 100644 --- a/langfun/core/llms/llama_cpp.py +++ b/langfun/core/llms/llama_cpp.py @@ -56,11 +56,11 @@ class LlamaCppRemote(openai_compatible.OpenAIChatCompletionAPI): ] = '' @property - def api_endpoint(self) -> str: + def api_endpoint(self) -> str: # pyrefly: ignore[bad-override] return self.url + '/completion' @property - def model_id(self) -> str: + def model_id(self) -> str: # pyrefly: ignore[bad-override] """Returns a string to identify the model.""" return f'LLaMAC++({self.model or ""})' diff --git a/langfun/core/llms/openai.py b/langfun/core/llms/openai.py index 49cdbc8a..2162ac2f 100644 --- a/langfun/core/llms/openai.py +++ b/langfun/core/llms/openai.py @@ -41,7 +41,7 @@ class OpenAIModelInfo(lf.ModelInfo): error_codes='https://platform.openai.com/docs/guides/error-codes', ) - provider: Final[str] = 'OpenAI' # pylint: disable=invalid-name + provider: Final[str] = 'OpenAI' # pylint: disable=invalid-name # pyrefly: ignore[bad-override] # @@ -1160,7 +1160,7 @@ def model_info(self) -> OpenAIModelInfo: return _SUPPORTED_MODELS_BY_MODEL_ID[self.model] @classmethod - def dir(cls): + def dir(cls): # pyrefly: ignore[bad-override] return [s.model_id for s in SUPPORTED_MODELS if s.in_service] def _request_args( diff --git a/langfun/core/llms/openai_compatible.py b/langfun/core/llms/openai_compatible.py index c91ea9f7..97107902 100644 --- a/langfun/core/llms/openai_compatible.py +++ b/langfun/core/llms/openai_compatible.py @@ -41,7 +41,7 @@ class OpenAIChatCompletionAPI(rest.REST): ] = '' @property - def headers(self) -> dict[str, Any]: + def headers(self) -> dict[str, Any]: # pyrefly: ignore[bad-override] return { 'Content-Type': 'application/json' } @@ -95,7 +95,7 @@ def request( json_schema=output_schema, ) ) - prompt.metadata.formatted_text = ( + prompt.metadata.formatted_text = ( # pyrefly: ignore[missing-attribute] prompt.text + '\n\n [RESPONSE FORMAT (not part of prompt)]\n' + pg.to_json_str(request_args['response_format'], json_indent=2) @@ -195,7 +195,7 @@ def result(self, json: dict[str, Any]) -> lf.LMSamplingResult: def _error(self, status_code: int, content: str) -> lf.LMError: if (status_code == 413 - or (status_code == 400 and b'string_above_max_length' in content)): + or (status_code == 400 and b'string_above_max_length' in content)): # pyrefly: ignore[unsupported-operation] return lf.ContextLimitError(f'{status_code}: {content}') return super()._error(status_code, content) @@ -235,7 +235,7 @@ def request( if output_schema is not None: output_schema['type'] = 'json_schema' request_args.update(text=dict(format=output_schema)) - prompt.metadata.formatted_text = ( + prompt.metadata.formatted_text = ( # pyrefly: ignore[missing-attribute] prompt.text + '\n\n [RESPONSE FORMAT (not part of prompt)]\n' + pg.to_json_str(request_args['text'], json_indent=2) diff --git a/langfun/core/llms/rest.py b/langfun/core/llms/rest.py index f7168bba..255762c3 100644 --- a/langfun/core/llms/rest.py +++ b/langfun/core/llms/rest.py @@ -61,6 +61,47 @@ class REST(lf.LanguageModel): 'The headers for the REST API.' ] = None + inactivity_timeout: Annotated[ + float | None, + ( + 'Short per-chunk INACTIVITY bound in seconds. This is the maximum ' + 'time allowed to elapse between two successive received chunks (it ' + 'resets every time a chunk arrives). A genuinely dead connection ' + 'fast-fails after this window, while a live-but-slow generation ' + 'that keeps emitting bytes is NOT killed. When None, falls back to ' + '`timeout`, preserving the historical single-timeout behavior. Used ' + 'both as the requests read timeout (socket idle) and as an in-loop ' + 'check in _read_response_with_deadline.' + ), + ] = None + + max_total_timeout: Annotated[ + float | None, + ( + 'Long TOTAL wall-clock budget in seconds for a single response. ' + 'This bounds the entire response duration regardless of how ' + 'steadily bytes arrive, letting a healthy multi-hour generation ' + 'complete. When None, falls back to `timeout`. Set this large ' + '(e.g. 14400 for 4h) together with a small `inactivity_timeout` ' + '(e.g. 120) to support slow-but-live long generations.' + ), + ] = None + + max_response_size: Annotated[ + int | None, + ( + 'Maximum total size, in BYTES, of a single streamed response body. ' + 'Streaming responses are buffered fully in memory before parsing; ' + 'without a cap a runaway or pathologically large generation can grow ' + 'the buffer to tens of GB and OOM-kill the process (the time-based ' + 'bounds do not help when a server keeps streaming bytes steadily). ' + 'When the cumulative body would exceed this many bytes the response ' + 'is closed and a (non-retryable) `ResponseSizeLimitError` is raised. ' + 'Must be a positive integer when set. When None, no size cap is ' + 'enforced (historical behavior).' + ), + ] = None + @functools.cached_property def _api_initialized(self) -> bool: """Returns whether the API is initialized.""" @@ -84,6 +125,11 @@ def _session(self) -> requests.Session: def _on_bound(self): super()._on_bound() self.__dict__.pop('_api_initialized', None) + if self.max_response_size is not None and self.max_response_size <= 0: + raise ValueError( + 'max_response_size must be a positive integer or None; got ' + f'{self.max_response_size!r}.' + ) def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]: assert self._api_initialized @@ -94,9 +140,10 @@ def _sample(self, prompts: list[lf.Message]) -> list[lf.LMSamplingResult]: def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult: try: with self.session() as session: + total_timeout = self._effective_total_timeout deadline = ( - (time.monotonic() + self.timeout) - if self.timeout is not None + (time.monotonic() + total_timeout) + if total_timeout is not None else None ) response = session.post( @@ -143,51 +190,115 @@ def _sample_single(self, prompt: lf.Message) -> lf.LMSamplingResult: raise lf.TemporaryLMError(error_message) from e raise lf.LMError(error_message) from e + @property + def _effective_inactivity_timeout(self) -> float | None: + """Short per-chunk inactivity bound (resets on each received chunk). + + Falls back to self.timeout when not explicitly configured, preserving the + historical single-timeout behavior for callers that only set `timeout`. + """ + if self.inactivity_timeout is not None: + return self.inactivity_timeout + return self.timeout + + @property + def _effective_total_timeout(self) -> float | None: + """Long total wall-clock budget for a single response. + + Falls back to self.timeout when not explicitly configured. + """ + if self.max_total_timeout is not None: + return self.max_total_timeout + return self.timeout + @property def _per_operation_timeout(self): """Per-operation (connect, read) timeout tuple for the requests library. - Splits self.timeout into separate connect and read timeouts: - - connect: bounded to 60s (no server needs more to accept a TCP connection) - - read: uses self.timeout (max wait for each chunk of response data) - - This is the per-socket-read timeout. The total-request deadline is enforced - separately by _read_response_with_deadline. + - connect: bounded to 60s (no server needs more to accept a TCP + connection). + - read: the per-socket idle timeout == the inactivity bound. requests + raises ReadTimeout if no bytes arrive within this window, which is + exactly the dead-connection fast-fail we want. It does NOT bound the + total response time (that is enforced by _read_response_with_deadline). """ - if self.timeout is None: + inactivity = self._effective_inactivity_timeout + if inactivity is None: return None - timeout = max(0.0, self.timeout) - return (min(60.0, timeout), timeout) + inactivity = max(0.0, inactivity) + return (min(60.0, inactivity), inactivity) def _read_response_with_deadline( self, response: requests.Response, deadline: float | None ) -> None: - """Reads response body, enforcing a total-request deadline. + """Reads response body, enforcing inactivity + total-request deadlines. When stream=True, session.post() returns after HTTP headers are received. - This method reads the body in chunks, checking the wall-clock deadline - between each chunk. If the deadline is exceeded, the response is closed - (which immediately closes the underlying socket) and TimeoutError is raised. + This method reads the body in chunks and enforces TWO independent bounds: + + - Inactivity bound (short, `_effective_inactivity_timeout`): the maximum + time allowed BETWEEN two successive chunks. It resets every time a chunk + is received, so a live-but-slow generation that keeps emitting bytes is + never killed, while a genuinely dead connection fast-fails. This is also + enforced at the socket layer via the read timeout in + `_per_operation_timeout` (which covers the case where iter_content blocks + because the server sends nothing at all); the in-loop check below + additionally covers slow trickles observed between yielded chunks. + - Total bound (long, via `deadline`): the absolute wall-clock budget for + the whole response. Lets a healthy multi-hour generation complete. + + If either bound is exceeded, the response is closed (which immediately + closes the underlying socket) and TimeoutError is raised. After successful read, sets response._content so that response.json() and response.content work normally for _parse_response(). Args: response: A streaming requests.Response (from stream=True). - deadline: Monotonic clock deadline (from time.monotonic()), or None - to disable deadline enforcement. + deadline: Monotonic clock TOTAL deadline (from time.monotonic()), or + None to disable total-deadline enforcement. """ # If body was already buffered (non-streaming or content pre-loaded), # there is nothing to read. if response._content is not False: # pylint: disable=protected-access,g-bool-id-comparison return + inactivity = self._effective_inactivity_timeout + max_size = self.max_response_size chunks = [] + total_bytes = 0 + last_chunk_time = time.monotonic() try: for chunk in response.iter_content(chunk_size=65536): + now = time.monotonic() + # Inactivity bound: time since the previous chunk (or since the start + # of the read for the first chunk). Resets on every received chunk. + if inactivity is not None and now - last_chunk_time > inactivity: + raise TimeoutError( + f'No response data received for {inactivity}s ' + '(inactivity timeout).' + ) + # Total-SIZE bound: the time-based bounds above do not fire while a + # server keeps streaming bytes steadily, so an unbounded (e.g. runaway) + # generation would otherwise buffer its entire body here and OOM-kill + # the process. Check BEFORE buffering the chunk so the buffered body + # never exceeds `max_response_size`, and fail fast with a NON-retryable + # error (`ResponseSizeLimitError` is an `LMError`, not a + # `RetryableLMError`) so retries do not re-stream the same oversized + # response. The `except BaseException` below closes the socket. + if max_size is not None and total_bytes + len(chunk) > max_size: + raise lf.ResponseSizeLimitError( + f'Response body exceeded max_response_size of {max_size} bytes ' + f'(received {total_bytes} bytes before the chunk that would ' + 'exceed the limit).' + ) chunks.append(chunk) - if deadline is not None and time.monotonic() > deadline: + total_bytes += len(chunk) + last_chunk_time = now + # Total bound: absolute wall-clock budget for the whole response. + if deadline is not None and now > deadline: raise TimeoutError( - f'Response exceeded total deadline of {self.timeout}s' + 'Response exceeded total deadline of ' + f'{self._effective_total_timeout}s.' ) except BaseException: # Close on any error to prevent TCP connection leaks. Protect close() @@ -198,7 +309,7 @@ def _read_response_with_deadline( pass raise # Set internal cache so response.json() / response.content work normally. - response._content = b''.join(chunks) # pylint: disable=protected-access + response._content = b''.join(chunks) # pylint: disable=protected-access # pyrefly: ignore[bad-assignment] # Content filtering patterns observed from various LLM providers. # These are best-effort substring heuristics derived from real API error @@ -240,18 +351,33 @@ def _error(self, status_code: int, content: str) -> lf.LMError: error_cls = lf.LMError return error_cls(f'{status_code}: {content}') + def _response_to_message_dict(self, response: requests.Response) -> Any: + """Converts an HTTP response into the message dict expected by `result`. + + Default implementation assumes a single buffered JSON body. Subclasses + whose API returns a streamed (e.g. Server-Sent Events) body override this + to reassemble the full message dict before it reaches `result()`. + + Args: + response: The HTTP response returned by the API. + + Returns: + The parsed message dict to be passed to `result`. + """ + return response.json() + def _parse_response(self, response: requests.Response) -> lf.LMSamplingResult: """Parses the LLM response.""" if response.status_code == 200: try: - return self.result(response.json()) + return self.result(self._response_to_message_dict(response)) except (ValueError, KeyError) as e: raise lf.LMError(str(e)) from e else: - raise self._error(response.status_code, response.content) + raise self._error(response.status_code, response.content) # pyrefly: ignore[bad-argument-type] @property - def max_concurrency(self) -> int | None: + def max_concurrency(self) -> int | None: # pyrefly: ignore[bad-override] """Returns the max concurrency for this model.""" rate_limits = self.model_info.rate_limits if rate_limits is not None: diff --git a/langfun/core/llms/rest_test.py b/langfun/core/llms/rest_test.py index 493504d8..4d4a2bcb 100644 --- a/langfun/core/llms/rest_test.py +++ b/langfun/core/llms/rest_test.py @@ -394,6 +394,129 @@ def test_error_response_works_with_streaming(self): lm._sample_single(lf.UserMessage('hello')) +class MaxResponseSizeTest(unittest.TestCase): + """Tests for the total-response-size cap (OOM guard) via stream=True. + + Regression coverage for a worker OOM where an unbounded streamed response + buffered its entire body into memory (tens of GB) because the time-based + bounds never fire while the server keeps streaming bytes steadily. + """ + + def _make_lm(self, max_response_size=None, timeout=120.0): + return rest.REST( + api_endpoint='https://fake-api.com', + request=lambda x, o: dict(prompt=x.text), + result=lambda x: lf.LMSamplingResult( + [lf.LMSample(c) for c in x['content']] + ), + timeout=timeout, + max_response_size=max_response_size, + ) + + def test_response_exceeding_max_size_raises_response_size_limit_error(self): + """A body larger than max_response_size fails with ResponseSizeLimitError.""" + lm = self._make_lm(max_response_size=250) + chunks = [b'x' * 100] * 5 # 500 bytes total > 250 cap. + with mock.patch('requests.Session.post') as mock_post: + mock_post.side_effect = mock_streaming_post(chunks) + with self.assertRaises(lf.ResponseSizeLimitError) as ctx: + lm._sample_single(lf.UserMessage('hello')) + self.assertIn('max_response_size', str(ctx.exception)) + + def test_size_cap_error_is_catchable_as_lm_error(self): + """Backward-compat: existing `except lf.LMError` handlers still catch it.""" + lm = self._make_lm(max_response_size=10) + with mock.patch('requests.Session.post') as mock_post: + mock_post.side_effect = mock_streaming_post([b'x' * 100]) + with self.assertRaises(lf.LMError): + lm._sample_single(lf.UserMessage('hello')) + + def test_size_cap_error_is_non_retryable(self): + """The size-cap error must NOT be a RetryableLMError. + + Retrying would re-stream the same oversized response (and with e.g. + max_attempts=80 that is catastrophic), so the cap has to fail fast. + """ + lm = self._make_lm(max_response_size=10) + with mock.patch('requests.Session.post') as mock_post: + mock_post.side_effect = mock_streaming_post([b'x' * 100]) + with self.assertRaises(lf.LMError) as ctx: + lm._sample_single(lf.UserMessage('hello')) + self.assertNotIsInstance(ctx.exception, lf.RetryableLMError) + + def test_response_within_max_size_ok(self): + """A body under the cap completes and parses normally.""" + lm = self._make_lm(max_response_size=10_000_000) + valid_json = pg.to_json_str({'content': ['hello']}).encode() + with mock.patch('requests.Session.post') as mock_post: + mock_post.side_effect = mock_streaming_post([valid_json]) + result = lm._sample_single(lf.UserMessage('test')) + self.assertEqual(result.samples[0].response.text, 'hello') + + def test_body_exactly_at_cap_is_allowed(self): + """A body whose total size equals the cap exactly must NOT be rejected.""" + valid_json = pg.to_json_str({'content': ['hello']}).encode() + lm = self._make_lm(max_response_size=len(valid_json)) + with mock.patch('requests.Session.post') as mock_post: + mock_post.side_effect = mock_streaming_post([valid_json]) + result = lm._sample_single(lf.UserMessage('test')) + self.assertEqual(result.samples[0].response.text, 'hello') + + def test_buffer_is_strictly_bounded_and_socket_closed(self): + """The tripping chunk is not buffered and the socket is closed promptly. + + Verifies the check-before-append behavior: iteration stops as soon as the + cap would be exceeded (so the buffer never grows past the cap), and the + response is closed to release the underlying socket. + """ + consumed = [] + closed = {'count': 0} + + def _mock_post(url, json=None, timeout=None, stream=False, **kwargs): + del url, json, timeout, stream, kwargs + response = requests.Response() + response.status_code = 200 + response.headers['Content-Type'] = 'application/json' + + def _iter(chunk_size=1, decode_unicode=False): + del chunk_size, decode_unicode + # 10 chunks of 100 bytes; cap=250 should stop after 2 are buffered. + for i in range(10): + consumed.append(i) + yield b'x' * 100 + + response.iter_content = _iter + response._content = False + response.close = lambda: closed.__setitem__('count', closed['count'] + 1) + return response + + lm = self._make_lm(max_response_size=250) + with mock.patch('requests.Session.post') as mock_post: + mock_post.side_effect = _mock_post + with self.assertRaises(lf.ResponseSizeLimitError): + lm._sample_single(lf.UserMessage('hello')) + # Only the chunks that fit under the cap (2 x 100 = 200 <= 250) plus the + # one that trips it (3rd) are pulled from the generator; the rest are not. + self.assertLessEqual(len(consumed), 3) + self.assertGreaterEqual(closed['count'], 1) + + def test_none_disables_size_cap(self): + """max_response_size=None preserves historical (uncapped) behavior.""" + lm = self._make_lm(max_response_size=None) + valid_json = pg.to_json_str({'content': ['hello' * 1000]}).encode() + with mock.patch('requests.Session.post') as mock_post: + mock_post.side_effect = mock_streaming_post([valid_json]) + result = lm._sample_single(lf.UserMessage('test')) + self.assertEqual(result.samples[0].response.text, 'hello' * 1000) + + def test_non_positive_max_response_size_is_rejected(self): + """A non-positive cap is a misconfiguration and must be rejected at bind.""" + for bad in (0, -1, -1024): + with self.subTest(max_response_size=bad): + with self.assertRaises(ValueError): + self._make_lm(max_response_size=bad) + + class AdversarialStreamingTest(unittest.TestCase): """Red-team tests: adversarial scenarios targeting stream+deadline logic. @@ -1326,5 +1449,110 @@ def test_timeout_zero_trawler_consistent(self): self.assertEqual(kwargs['request_deadline_ms'], 0) +class InactivityVsTotalTimeoutTest(unittest.TestCase): + """Tests the inactivity vs total deadline split. + + This split (a SHORT per-chunk inactivity bound vs a LONG total wall-clock + budget) is the core slow-but-live-generation fix. Real timescales + (120s inactivity / 4h total) are scaled down here so the tests run quickly; + the logic under test is identical. + """ + + def _make_lm(self, timeout=None, inactivity_timeout=None, + max_total_timeout=None): + return rest.REST( + api_endpoint='https://fake-api.com', + request=lambda x, o: dict(prompt=x.text), + result=lambda x: lf.LMSamplingResult( + [lf.LMSample(c) for c in x['content']] + ), + timeout=timeout, + inactivity_timeout=inactivity_timeout, + max_total_timeout=max_total_timeout, + ) + + def test_effective_timeouts_fall_back_to_timeout(self): + """When the new knobs are unset, both collapse to self.timeout (legacy).""" + lm = self._make_lm(timeout=900.0) + self.assertEqual(lm._effective_inactivity_timeout, 900.0) + self.assertEqual(lm._effective_total_timeout, 900.0) + + def test_effective_timeouts_use_explicit_knobs(self): + """Explicit knobs override the timeout fallback independently.""" + lm = self._make_lm( + timeout=900.0, inactivity_timeout=120.0, max_total_timeout=14400.0 + ) + self.assertEqual(lm._effective_inactivity_timeout, 120.0) + self.assertEqual(lm._effective_total_timeout, 14400.0) + + def test_per_operation_read_timeout_uses_inactivity_not_total(self): + """The requests read timeout is the inactivity bound, not the total.""" + lm = self._make_lm( + timeout=900.0, inactivity_timeout=120.0, max_total_timeout=14400.0 + ) + self.assertEqual(lm._per_operation_timeout, (60.0, 120.0)) + + def test_silent_stream_fails_fast_via_inactivity(self): + """A silent stream fails fast via the inactivity bound. + + It must fail well before the (large) total budget would elapse. + """ + lm = self._make_lm(inactivity_timeout=0.1, max_total_timeout=10.0) + valid_json = pg.to_json_str({'content': ['never finishes']}).encode() + half = len(valid_json) // 2 + start = time.monotonic() + with mock.patch('requests.Session.post') as mock_post: + # Gap between chunks (0.3s) exceeds the 0.1s inactivity bound. + mock_post.side_effect = mock_streaming_post( + [valid_json[:half], valid_json[half:]], delay_per_chunk=0.3 + ) + with self.assertRaises(lf.TemporaryLMError) as ctx: + lm._sample_single(lf.UserMessage('hello')) + elapsed = time.monotonic() - start + self.assertIn('inactivity', str(ctx.exception).lower()) + # Fast-fail: nowhere near the 10s total budget. + self.assertLess(elapsed, 5.0) + + def test_steady_stream_past_old_timeout_still_succeeds(self): + """A steady stream past the OLD single-timeout wall still succeeds. + + The total budget is large and each inter-chunk gap stays under the + inactivity bound. The legacy single timeout (0.3s) would have killed this + response, but with inactivity=0.2s and total=5s it completes. + """ + old_single_timeout = 0.3 + lm = self._make_lm(inactivity_timeout=0.2, max_total_timeout=5.0) + valid_json = pg.to_json_str({'content': ['streamed-done']}).encode() + # Split into 10 chunks emitted every 0.05s => ~0.5s total > 0.3s old wall. + n = 10 + size = max(1, len(valid_json) // n) + chunks = [valid_json[i:i + size] for i in range(0, len(valid_json), size)] + start = time.monotonic() + with mock.patch('requests.Session.post') as mock_post: + mock_post.side_effect = mock_streaming_post(chunks, delay_per_chunk=0.05) + result = lm._sample_single(lf.UserMessage('hello')) + elapsed = time.monotonic() - start + self.assertEqual(result.samples[0].response.text, 'streamed-done') + # Proof it genuinely crossed the old single-timeout wall. + self.assertGreater(elapsed, old_single_timeout) + + def test_total_deadline_uses_max_total_timeout(self): + """The total wall budget is governed by max_total_timeout. + + A large inactivity bound does not prevent the total deadline from firing. + """ + lm = self._make_lm(inactivity_timeout=10.0, max_total_timeout=0.2) + valid_json = pg.to_json_str({'content': ['too slow overall']}).encode() + half = len(valid_json) // 2 + with mock.patch('requests.Session.post') as mock_post: + # Each gap (0.15s) < 10s inactivity, but cumulative 0.3s > 0.2s total. + mock_post.side_effect = mock_streaming_post( + [valid_json[:half], valid_json[half:]], delay_per_chunk=0.15 + ) + with self.assertRaises(lf.TemporaryLMError) as ctx: + lm._sample_single(lf.UserMessage('hello')) + self.assertIn('total deadline', str(ctx.exception).lower()) + + if __name__ == '__main__': unittest.main() diff --git a/langfun/core/llms/veo.py b/langfun/core/llms/veo.py index 5a66d88a..84933800 100644 --- a/langfun/core/llms/veo.py +++ b/langfun/core/llms/veo.py @@ -171,13 +171,13 @@ def model_info(self) -> VeoModelInfo: return _SUPPORTED_MODELS_BY_ID[self.model] @property - def headers(self): + def headers(self): # pyrefly: ignore[bad-override] return { 'Content-Type': 'application/json; charset=utf-8', } @property - def api_endpoint(self) -> str: + def api_endpoint(self) -> str: # pyrefly: ignore[bad-override] assert self._api_initialized return ( f'https://{self._location}-aiplatform.googleapis.com/v1/projects/' @@ -200,11 +200,11 @@ def request( f'got {len(prompt.images)}.' ) first_frame = prompt.images[0] - instance['image'] = self._encode_image(first_frame) + instance['image'] = self._encode_image(first_frame) # pyrefly: ignore[bad-argument-type] if len(prompt.images) > 1: last_frame = prompt.images[1] - instance['lastFrame'] = self._encode_image(last_frame) + instance['lastFrame'] = self._encode_image(last_frame) # pyrefly: ignore[bad-argument-type] parameters: dict[str, Any] = { 'durationSeconds': self.duration_seconds, diff --git a/langfun/core/llms/vertexai.py b/langfun/core/llms/vertexai.py index 03673323..14d66d26 100644 --- a/langfun/core/llms/vertexai.py +++ b/langfun/core/llms/vertexai.py @@ -170,9 +170,9 @@ def _initialize(self): @property def _project(self) -> str: """Returns a project ID. Randomly selects from list if multiple provided.""" - if len(self._projects) == 1: - return self._projects[0] - return random.choice(self._projects) + if len(self._projects) == 1: # pyrefly: ignore[bad-argument-type] + return self._projects[0] # pyrefly: ignore[unsupported-operation] + return random.choice(self._projects) # pyrefly: ignore[bad-argument-type] def session(self): assert self._api_initialized @@ -236,7 +236,7 @@ class VertexAIGemini(VertexAI, gemini.Gemini): location = 'us-central1' @property - def api_endpoint(self) -> str: + def api_endpoint(self) -> str: # pyrefly: ignore[bad-override] assert self._api_initialized project = self._project return ( @@ -247,7 +247,7 @@ def api_endpoint(self) -> str: @functools.cached_property def model_info(self) -> gemini.GeminiModelInfo: - return super().model_info.clone(override=dict(provider='VertexAI')) + return super().model_info.clone(override=dict(provider='VertexAI')) # pyrefly: ignore[bad-return] # @@ -294,6 +294,20 @@ class VertexAIGemini3FlashPreview(VertexAIGemini): # pylint: disable=invalid-na location = 'global' +class VertexAIGemini35Flash(VertexAIGemini): # pylint: disable=invalid-name + """Gemini 3.5 Flash GA model launched on 05/19/2026.""" + + model = 'gemini-3.5-flash' + location = 'global' + + +class VertexAIGemini37Flash(VertexAIGemini): # pylint: disable=invalid-name + """Gemini 3.7 Flash GA model launched on 08/13/2026.""" + + model = 'gemini-3.7-flash' + location = 'global' + + class VertexAIGemini31FlashLitePreview(VertexAIGemini): # pylint: disable=invalid-name """Gemini 3.1 Flash Lite Preview model.""" @@ -301,6 +315,13 @@ class VertexAIGemini31FlashLitePreview(VertexAIGemini): # pylint: disable=inval location = 'global' +class VertexAIGemini31FlashLite(VertexAIGemini): # pylint: disable=invalid-name + """Gemini 3.1 Flash Lite model.""" + + model = 'gemini-3.1-flash-lite' + location = 'global' + + class VertexAIGemini25Pro(VertexAIGemini): # pylint: disable=invalid-name """Gemini 2.5 Pro GA model launched on 06/17/2025.""" @@ -463,13 +484,13 @@ def model_info(self) -> lf.ModelInfo: return mi @property - def headers(self): + def headers(self): # pyrefly: ignore[bad-override] return { 'Content-Type': 'application/json; charset=utf-8', } @property - def api_endpoint(self) -> str: + def api_endpoint(self) -> str: # pyrefly: ignore[bad-override] project = self._project model_id = str(self.model).removesuffix('@latest') host = ( @@ -497,6 +518,13 @@ def request( # pylint: disable=invalid-name +class VertexAIClaude5Opus(VertexAIAnthropic): + """Anthropic's Claude Opus 5 model on VertexAI.""" + + model = 'claude-opus-5' + location = 'global' + + class VertexAIClaude48Opus(VertexAIAnthropic): """Anthropic's Claude 4.8 Opus model on VertexAI.""" @@ -672,7 +700,7 @@ def model_info(self) -> lf.ModelInfo: return _LLAMA_MODELS_BY_MODEL_ID[self.model] @property - def api_endpoint(self) -> str: + def api_endpoint(self) -> str: # pyrefly: ignore[bad-override] assert self._api_initialized project = self._project return ( @@ -786,7 +814,7 @@ def model_info(self) -> lf.ModelInfo: return _MISTRAL_MODELS_BY_MODEL_ID[self.model] @property - def api_endpoint(self) -> str: + def api_endpoint(self) -> str: # pyrefly: ignore[bad-override] assert self._api_initialized project = self._project return ( @@ -835,6 +863,7 @@ def _register_vertexai_models(): lf.LanguageModel.register('claude-opus-4-7@latest', anthropic.Anthropic) lf.LanguageModel.register('claude-opus-4-8', VertexAIClaude48Opus) lf.LanguageModel.register('claude-opus-4-8@latest', anthropic.Anthropic) + lf.LanguageModel.register('claude-opus-5', VertexAIClaude5Opus) for m in LLAMA_MODELS: lf.LanguageModel.register(m.model_id, VertexAILlama) diff --git a/langfun/core/llms/vertexai_test.py b/langfun/core/llms/vertexai_test.py index 7de45307..f8b3cd15 100644 --- a/langfun/core/llms/vertexai_test.py +++ b/langfun/core/llms/vertexai_test.py @@ -54,6 +54,17 @@ def test_project_and_location_check(self): del os.environ['VERTEXAI_PROJECT'] del os.environ['VERTEXAI_LOCATION'] + @mock.patch.object(vertexai.VertexAI, 'credentials', new=True) + def test_gemini_37_flash(self): + os.environ['VERTEXAI_PROJECT'] = 'abc' + os.environ['VERTEXAI_LOCATION'] = 'us-central1' + model = vertexai.VertexAIGemini37Flash(location=pg.MISSING_VALUE) + self.assertEqual(model.resource_id, 'vertexai://gemini-3.7-flash') + # 3.x models default to 'global' location. + self.assertIn('global', model.api_endpoint) + del os.environ['VERTEXAI_PROJECT'] + del os.environ['VERTEXAI_LOCATION'] + @mock.patch.object(vertexai.VertexAI, 'credentials', new=True) def test_gemini_31_flash_lite_preview(self): os.environ['VERTEXAI_PROJECT'] = 'abc' @@ -67,6 +78,17 @@ def test_gemini_31_flash_lite_preview(self): del os.environ['VERTEXAI_PROJECT'] del os.environ['VERTEXAI_LOCATION'] + @mock.patch.object(vertexai.VertexAI, 'credentials', new=True) + def test_gemini_31_flash_lite(self): + os.environ['VERTEXAI_PROJECT'] = 'abc' + os.environ['VERTEXAI_LOCATION'] = 'us-central1' + model = vertexai.VertexAIGemini31FlashLite(location=pg.MISSING_VALUE) + self.assertEqual(model.resource_id, 'vertexai://gemini-3.1-flash-lite') + # 3.x models default to 'global' location. + self.assertIn('global', model.api_endpoint) + del os.environ['VERTEXAI_PROJECT'] + del os.environ['VERTEXAI_LOCATION'] + @mock.patch.object(vertexai.VertexAI, 'credentials', new=True) def test_multi_project_support(self): # Test single project (backward compatibility) @@ -202,7 +224,7 @@ def test_basics(self): }], 'role': 'user', }], - 'stream': False, + 'stream': True, 'temperature': 0.0, 'top_k': 40, }, diff --git a/langfun/core/mcp/tool.py b/langfun/core/mcp/tool.py index 0d631ba8..ba7c5636 100644 --- a/langfun/core/mcp/tool.py +++ b/langfun/core/mcp/tool.py @@ -107,7 +107,7 @@ def result_to_message( message = lf_message.ToolMessage.from_chunks(chunks) if result.structuredContent: message.metadata.update(result.structuredContent) - return message + return message # pyrefly: ignore[bad-return] def __call__( self, @@ -199,7 +199,7 @@ def make_class(cls, tool_definition: mcp.Tool) -> type['McpTool']: represents the defined tool. """ - class _McpTool(cls): + class _McpTool(cls): # pyrefly: ignore[invalid-inheritance] auto_schema = False tool_cls = _McpTool @@ -234,7 +234,7 @@ def make_class(cls, name: str, schema: pg.Schema): A dynamically generated class that inherits from `McpToolInput`. """ - class _McpToolInput(cls): + class _McpToolInput(cls): # pyrefly: ignore[invalid-inheritance] pass input_cls = _McpToolInput diff --git a/langfun/core/mcp/tool_test.py b/langfun/core/mcp/tool_test.py index 5a410224..052fe90b 100644 --- a/langfun/core/mcp/tool_test.py +++ b/langfun/core/mcp/tool_test.py @@ -95,8 +95,8 @@ def test_make_tool_class(self): s = tool_cls.__schema__ self.assertEqual(list(s.fields.keys()), ['a', 'b']) self.assertEqual(repr(tool_cls), "") - self.assertEqual(s.fields['a'].description, 'Integer a.') - self.assertEqual(s.fields['b'].description, 'String b.') + self.assertEqual(s.fields['a'].description, 'Integer a.') # pyrefly: ignore[bad-index] + self.assertEqual(s.fields['b'].description, 'String b.') # pyrefly: ignore[bad-index] self.assertEqual( tool_cls.python_definition(markdown=True), diff --git a/langfun/core/message.py b/langfun/core/message.py index 5748df1f..9a119917 100644 --- a/langfun/core/message.py +++ b/langfun/core/message.py @@ -202,13 +202,13 @@ def __init__( metadata = metadata or {} metadata.update(kwargs) if isinstance(referred_modalities, list): - referred_modalities = {m.id: pg.Ref(m) for m in referred_modalities} + referred_modalities = {m.id: pg.Ref(m) for m in referred_modalities} # pyrefly: ignore[bad-assignment] # Auto-append text markers for modalities not yet referenced in the # text, so they are included in the LM prompt. This only applies to the # list input path; dict users explicitly manage the text themselves. - for modality_id in referred_modalities: - marker = modality.Modality.text_marker(modality_id) + for modality_id in referred_modalities: # pyrefly: ignore[not-iterable] + marker = modality.Modality.text_marker(modality_id) # pyrefly: ignore[bad-argument-type] if marker not in text: text += f'\n{marker}' @@ -312,7 +312,7 @@ def convertible_formats(cls) -> list[str]: @classmethod def convertible_types(cls) -> list[str]: """Returns supported types for message conversion.""" - return MessageConverter.convertible_types() + return MessageConverter.convertible_types() # pyrefly: ignore[bad-return] # # Unified interface for accessing text, result and metadata. @@ -340,7 +340,7 @@ def set(self, key_path: str | pg.KeyPath, value: Any) -> None: if key_path == Message.PATH_TEXT: self.rebind({key_path: value}, raise_on_no_change=False) else: - self.metadata.rebind({key_path: value}, raise_on_no_change=False) + self.metadata.rebind({key_path: value}, raise_on_no_change=False) # pyrefly: ignore[missing-attribute] def get(self, key_path: str | pg.KeyPath, default: Any = None) -> Any: """Gets text or metadata by key path. @@ -360,7 +360,7 @@ def get(self, key_path: str | pg.KeyPath, default: Any = None) -> Any: if key_path == Message.PATH_TEXT: return self.text else: - return self.metadata.sym_get(key_path, default, use_inferred=True) + return self.metadata.sym_get(key_path, default, use_inferred=True) # pyrefly: ignore[missing-attribute] # # API for accessing the structured result and error. @@ -434,7 +434,7 @@ def apply_updates(self, updates: dict[pg.KeyPath, pg.FieldUpdate]) -> None: # Rebind will trigger _on_change, which inserts the updates # to current message' updates. - self.rebind(delta, raise_on_no_change=False) + self.rebind(delta, raise_on_no_change=False) # pyrefly: ignore[bad-argument-type] # # API for supporting modalities. @@ -573,7 +573,7 @@ def from_chunks( last_char = None else: assert isinstance(chunk, modality.Modality), chunk - fused_text.write(modality.Modality.text_marker(chunk.id)) + fused_text.write(modality.Modality.text_marker(chunk.id)) # pyrefly: ignore[bad-argument-type] last_char = modality.Modality.REF_END[-1] # Make a reference if the chunk is already owned by another object # to avoid copy. @@ -1009,7 +1009,7 @@ def unregister(self, converter: Type['MessageConverter']) -> None: def get_by_type(self, t: Type[Any], **kwargs) -> 'MessageConverter': """Returns a message converter for the given type.""" - t = self._type_to_converters[t] + t = self._type_to_converters[t] # pyrefly: ignore[bad-assignment] if not t: raise TypeError( f'Cannot convert Message to {t!r}.' @@ -1019,7 +1019,7 @@ def get_by_type(self, t: Type[Any], **kwargs) -> 'MessageConverter': f'More than one converters found for output type {t!r}. ' f'Please specify one for this conversion: {[x.FORMAT_ID for x in t]}.' ) - return t[0](**kwargs) + return t[0](**kwargs) # pyrefly: ignore[unsupported-operation] def get_by_format(self, format: str, **kwargs) -> 'MessageConverter': # pylint: disable=redefined-builtin """Returns a message converter for the given format.""" diff --git a/langfun/core/modalities/__init__.py b/langfun/core/modalities/__init__.py index 4b9771a9..af5fe4e4 100644 --- a/langfun/core/modalities/__init__.py +++ b/langfun/core/modalities/__init__.py @@ -28,9 +28,9 @@ # Override the `images`, `videos` and `audios` properties of `Message` to # return the modalities of the corresponding types. -_message_lib.Message.images = property(lambda self: self.modalities(Image)) -_message_lib.Message.videos = property(lambda self: self.modalities(Video)) -_message_lib.Message.audios = property(lambda self: self.modalities(Audio)) +_message_lib.Message.images = property(lambda self: self.modalities(Image)) # pyrefly: ignore[bad-assignment] +_message_lib.Message.videos = property(lambda self: self.modalities(Video)) # pyrefly: ignore[bad-assignment] +_message_lib.Message.audios = property(lambda self: self.modalities(Audio)) # pyrefly: ignore[bad-assignment] # pylint: enable=g-import-not-at-top # pylint: enable=g-bad-import-order diff --git a/langfun/core/modalities/image.py b/langfun/core/modalities/image.py index 852645ae..d4977e9a 100644 --- a/langfun/core/modalities/image.py +++ b/langfun/core/modalities/image.py @@ -113,7 +113,7 @@ def _convert_to_format(self, pil_format: str) -> 'Image': img.save(buf, format=pil_format) finally: os.chdir(cwd) - return self.from_bytes(buf.getvalue()) + return self.from_bytes(buf.getvalue()) # pyrefly: ignore[bad-return] @classmethod def from_pil_image(cls, img: PILImage) -> 'Image': # pytype: disable=invalid-annotation @@ -127,4 +127,4 @@ def from_pil_image(cls, img: PILImage) -> 'Image': # pytype: disable=invalid-an img.save(buf, format='PNG') finally: os.chdir(cwd) - return cls.from_bytes(buf.getvalue()) + return cls.from_bytes(buf.getvalue()) # pyrefly: ignore[bad-return] diff --git a/langfun/core/modalities/mime.py b/langfun/core/modalities/mime.py index dcff3ae9..52079d9a 100644 --- a/langfun/core/modalities/mime.py +++ b/langfun/core/modalities/mime.py @@ -207,10 +207,10 @@ def _on_bound(self): def to_bytes(self) -> bytes: if self.content is not None: - return self.content + return self.content # pyrefly: ignore[bad-return] - self.rebind(content=self.download(self.uri), skip_notification=True) - return self.content + self.rebind(content=self.download(self.uri), skip_notification=True) # pyrefly: ignore[bad-argument-type] + return self.content # pyrefly: ignore[bad-return] @property def content_uri(self) -> str: @@ -235,7 +235,7 @@ def from_uri(cls, uri: str, **kwargs) -> 'Mime': if 'youtube.com/watch' in uri: return Custom(mime='text/html', uri=uri, **kwargs) content = cls.download(uri) - mime = _detect_mime_type(content) + mime = _detect_mime_type(content) # pyrefly: ignore[bad-argument-type] return cls.class_from_mime_type(mime)(uri=uri, content=content, **kwargs) return cls(uri=uri, content=None, **kwargs) @@ -260,7 +260,7 @@ def _parse_data_uri(cls, uri: str) -> tuple[str, bytes]: def from_bytes(cls, content: bytes | str, **kwargs) -> 'Mime': if cls is Mime: return cls.class_from_mime_type( - _detect_mime_type(content) + _detect_mime_type(content) # pyrefly: ignore[bad-argument-type] )(content=content, **kwargs) return cls(content=content, **kwargs) @@ -283,7 +283,7 @@ def download(cls, uri: str) -> bytes | str: assert content is not None return content - def _html_tree_view_content( + def _html_tree_view_content( # pyrefly: ignore[bad-override] self, **kwargs) -> str: return self._raw_html() diff --git a/langfun/core/modality.py b/langfun/core/modality.py index fd467c09..18ce0f46 100644 --- a/langfun/core/modality.py +++ b/langfun/core/modality.py @@ -70,7 +70,7 @@ def format(self, *args, **kwargs) -> str: capture_scope = get_modality_capture_context() if capture_scope is not None: capture_scope.capture(self) - return Modality.text_marker(self.id) + return Modality.text_marker(self.id) # pyrefly: ignore[bad-argument-type] def __str_kwargs__(self) -> dict[str, Any]: # For modality objects, we don't want to use markdown format when they @@ -225,7 +225,7 @@ def __init__(self): def capture(self, modality: Modality) -> None: """Captures the modality object.""" - self._references[modality.id] = pg.Ref(modality) + self._references[modality.id] = pg.Ref(modality) # pyrefly: ignore[unsupported-operation] @property def references(self) -> dict[str, pg.Ref[Modality]]: diff --git a/langfun/core/natural_language.py b/langfun/core/natural_language.py index 86ae9bea..4b527e34 100644 --- a/langfun/core/natural_language.py +++ b/langfun/core/natural_language.py @@ -32,7 +32,7 @@ def format( ) -> str: if natural_language: return self.natural_language_format() - return super().format(*args, **kwargs) + return super().format(*args, **kwargs) # pyrefly: ignore[missing-attribute] def __str__(self): return self.natural_language_format() diff --git a/langfun/core/sampling.py b/langfun/core/sampling.py index a55ff33d..42ae2643 100644 --- a/langfun/core/sampling.py +++ b/langfun/core/sampling.py @@ -143,7 +143,7 @@ def _concurrent_sample( if pg.is_deterministic(sampling_space): num_examples = num_examples or None def repeat_example(example, num_examples=num_examples): - for _ in range(num_examples): + for _ in range(num_examples): # pyrefly: ignore[bad-argument-type] yield example.clone() pg_sample_fn = repeat_example diff --git a/langfun/core/structured/completion.py b/langfun/core/structured/completion.py index f746fd54..6ccade3a 100644 --- a/langfun/core/structured/completion.py +++ b/langfun/core/structured/completion.py @@ -106,7 +106,7 @@ def missing_type_dependencies(self, value: Any) -> list[Type[Any]]: value_specs = tuple( [v.value_spec for v in schema_lib.Missing.find_missing(value).values()] ) - return schema_lib.class_dependencies(value_specs, include_subclasses=True) + return schema_lib.class_dependencies(value_specs, include_subclasses=True) # pyrefly: ignore[bad-argument-type] def class_defs_repr(self, value: Any) -> str | None: return schema_lib.class_definitions( diff --git a/langfun/core/structured/description.py b/langfun/core/structured/description.py index bebbba35..0ce51d96 100644 --- a/langfun/core/structured/description.py +++ b/langfun/core/structured/description.py @@ -29,6 +29,7 @@ class _DescribeStructure(mapping.Mapping): context_title = 'CONTEXT_FOR_DESCRIPTION' output_title = 'NATURAL_LANGUAGE_TEXT' + # pyrefly: ignore[bad-assignment] preamble = """ Please help describe {{ input_title }} in natural language. diff --git a/langfun/core/structured/function_generation.py b/langfun/core/structured/function_generation.py index 9311ddf4..8c58e10b 100644 --- a/langfun/core/structured/function_generation.py +++ b/langfun/core/structured/function_generation.py @@ -151,17 +151,18 @@ def calculate_area_circle(radius: float) -> float: f = python.evaluate(source_code, global_vars=context) # Check whether the sigantures are the same. - if inspect.signature(f) != inspect.signature(func): + if inspect.signature(f) != inspect.signature(func): # pyrefly: ignore[bad-argument-type] raise python.CodeError( code=source_code, cause=TypeError( + # pyrefly: ignore[bad-argument-type] f"Signature mismatch: Expected: {inspect.signature(func)}, " f"Actual: {inspect.signature(f)}.", ), ) if callable(unittest): - unittest(f) + unittest(f) # pyrefly: ignore[bad-argument-type] elif unittest_examples: unittest_with_test_cases(f, unittest_examples) @@ -171,7 +172,7 @@ def calculate_area_circle(radius: float) -> float: pg.logging.warning( f"Bad code generated: {e}", ) - raise last_error + raise last_error # pyrefly: ignore[bad-raise] def _process_signature(signature): @@ -270,7 +271,7 @@ def lm_generated_func(*args, **kwargs): func.__function__ = python.evaluate( func.__source_code__, global_vars=context ) - return func.__function__(*args, **kwargs) + return func.__function__(*args, **kwargs) # pyrefly: ignore[not-callable] func.__function__, func.__source_code__ = _function_gen( func, @@ -289,7 +290,7 @@ def lm_generated_func(*args, **kwargs): lm_generated_func.__name__ = func.__name__ lm_generated_func.__qualname__ = func.__qualname__ lm_generated_func.__module__ = func.__module__ - lm_generated_func.source = lambda: func.__source_code__ + lm_generated_func.source = lambda: func.__source_code__ # pyrefly: ignore[missing-attribute] return lm_generated_func return _decorate diff --git a/langfun/core/structured/mapping.py b/langfun/core/structured/mapping.py index f7547747..b16961eb 100644 --- a/langfun/core/structured/mapping.py +++ b/langfun/core/structured/mapping.py @@ -120,7 +120,7 @@ class Flight(pg.Object): """ input: pg.typing.Annotated[ - pg.typing.Any(transform=schema_lib.mark_missing), + pg.typing.Any(transform=schema_lib.mark_missing), # pyrefly: ignore[bad-instantiation] ( 'The input object of the mapping. It could be either a natural ' 'language-based string, or a Python object.' @@ -462,9 +462,9 @@ def transform_output(self, lm_output: lf.Message) -> lf.Message: lm_output.result = self.postprocess_result(self.parse_result(lm_output)) except Exception as e: # pylint: disable=broad-exception-caught if (self.lm.cache is not None - and lm_output.lm_input.cache_seed is not None): + and lm_output.lm_input.cache_seed is not None): # pyrefly: ignore[missing-attribute] success = self.lm.cache.delete( - self.lm, lm_output.lm_input, lm_output.lm_input.cache_seed + self.lm, lm_output.lm_input, lm_output.lm_input.cache_seed # pyrefly: ignore[bad-argument-type] ) assert success if self.default == lf.RAISE_IF_HAS_ERROR: diff --git a/langfun/core/structured/parsing.py b/langfun/core/structured/parsing.py index dbf119e3..36ea2c83 100644 --- a/langfun/core/structured/parsing.py +++ b/langfun/core/structured/parsing.py @@ -41,6 +41,7 @@ class _ParseStructure(mapping.Mapping): class _ParseStructureJson(_ParseStructure): """Parses an object out from a NL text using JSON as the protocol.""" + # pyrefly: ignore[bad-assignment] preamble = """ Please help translate the last LM response into JSON based on the request and the schema: @@ -57,6 +58,7 @@ class _ParseStructureJson(_ParseStructure): class _ParseStructurePython(_ParseStructure): """Parses an object out from a NL text using Python as the protocol.""" + # pyrefly: ignore[bad-assignment] preamble = """ Please help translate the last {{ input_title }} into {{ output_title}} based on {{ schema_title }}. @@ -174,10 +176,10 @@ class Flight(pg.Object): # Setting up context. call_context = dict(cache_seed=cache_seed, autofix=autofix) if lm is not None: - call_context['lm'] = lm + call_context['lm'] = lm # pyrefly: ignore[bad-assignment] autofix_lm = autofix_lm or lm if autofix_lm is not None: - call_context['autofix_lm'] = autofix_lm + call_context['autofix_lm'] = autofix_lm # pyrefly: ignore[bad-assignment] call_context.update(kwargs) output = t(input=message, **call_context) @@ -314,7 +316,7 @@ def call( A string if `returns` is None or an instance of the return type. """ # Call `lm` for natural response. - lm_output = lf.LangFunc.from_value(prompt, **kwargs)(lm=lm) + lm_output = lf.LangFunc.from_value(prompt, **kwargs)(lm=lm) # pyrefly: ignore[not-callable] if response_postprocess is not None: postprocessed_text = response_postprocess(lm_output.text) @@ -329,7 +331,7 @@ def _chain_nl_output_message(parsing_message: lf.Message): """Chain the source of the parsed output to the LM output.""" parsing_message.root.source = lm_output parsing_message.tag('parsing-lm-output') - parsing_message.lm_input.tag('parsing-lm-input') + parsing_message.lm_input.tag('parsing-lm-input') # pyrefly: ignore[missing-attribute] # Call `parsing_lm` for structured parsing. try: @@ -337,7 +339,7 @@ def _chain_nl_output_message(parsing_message: lf.Message): lm_output.text, schema, examples=parsing_examples, - lm=parsing_lm or lm, + lm=parsing_lm or lm, # pyrefly: ignore[bad-argument-type] include_context=parsing_include_context, cache_seed=cache_seed, autofix=autofix, diff --git a/langfun/core/structured/querying.py b/langfun/core/structured/querying.py index c62e5600..e2a3d18e 100644 --- a/langfun/core/structured/querying.py +++ b/langfun/core/structured/querying.py @@ -95,13 +95,14 @@ def __init_subclass__(cls) -> Any: if version_dict is None: version_dict = {} cls._OOP_PROMPT_MAP[protocol] = version_dict - dest_cls = version_dict.get(cls.version) + dest_cls = version_dict.get(cls.version) # pyrefly: ignore[missing-attribute] if dest_cls is not None and dest_cls.__type_name__ != cls.__type_name__: raise ValueError( + # pyrefly: ignore[missing-attribute] f'Version {cls.version} is already registered for {dest_cls!r} ' f'under protocol {protocol!r}. Please use a different version.' ) - version_dict[cls.version] = cls + version_dict[cls.version] = cls # pyrefly: ignore[missing-attribute] @classmethod def from_protocol(cls, protocol: str) -> Type['LfQuery']: @@ -143,6 +144,7 @@ def from_protocol(cls, protocol: str) -> Type['LfQuery']: class _LfQueryJsonV1(LfQuery): """Query a structured value using JSON as the protocol.""" + # pyrefly: ignore[bad-assignment] preamble = """ Please respond to the last {{ input_title }} with {{ output_title}} according to {{ schema_title }}: @@ -169,6 +171,7 @@ class _LfQueryJsonV1(LfQuery): class _LfQueryPythonV1(LfQuery): """Query a structured value using Python as the protocol.""" + # pyrefly: ignore[bad-assignment] preamble = """ Please respond to the last {{ input_title }} with {{ output_title }} according to {{ schema_title }}. @@ -222,6 +225,7 @@ class Answer: class _LfQueryPythonV2(LfQuery): """Query a structured value using Python as the protocol.""" + # pyrefly: ignore[bad-assignment] preamble = """ Please respond to the last {{ input_title }} with {{ output_title }} only according to {{ schema_title }}. @@ -563,7 +567,7 @@ def _single_query(inputs): # Query with structured output. query_cls = LfQuery.from_protocol(protocol) if ':' not in protocol: - protocol = f'{protocol}:{query_cls.version}' + protocol = f'{protocol}:{query_cls.version}' # pyrefly: ignore[missing-attribute] # `skip_lm`` is True when `lf.query_prompt` is called. # and `prompt` is `pg.MISSING_VALUE` when `lf.query_output` is called. @@ -625,7 +629,7 @@ def _mark_query_completed(output_message, error, usage_summary): try: if query_cls is None: # Query with natural language output. - output_message = lf.LangFunc.from_value(query_input, **kwargs)( + output_message = lf.LangFunc.from_value(query_input, **kwargs)( # pyrefly: ignore[not-callable] lm=lm, cache_seed=cache_seed, skip_lm=skip_lm ) if response_postprocess: @@ -758,7 +762,7 @@ def query_and_reduce( Returns: The reduced output from multiple `lf.query` calls. """ - results = query(prompt, schema, lm=lm, num_samples=num_samples, **kwargs) + results = query(prompt, schema, lm=lm, num_samples=num_samples, **kwargs) # pyrefly: ignore[bad-argument-type] if isinstance(results, list): results = reduce(results) return results @@ -893,7 +897,7 @@ def query_reward( query_output(response, output_cls), mapping_example.input, mapping_example.output, - mapping_example.metadata, + mapping_example.metadata, # pyrefly: ignore[bad-argument-type] ) @@ -1077,12 +1081,12 @@ def mark_completed( if self.schema is not None: try: output = query_output( - lm_response, self.schema, + lm_response, self.schema, # pyrefly: ignore[bad-argument-type] default=self.default, protocol=self.protocol ) except mapping.MappingError as e: output = None - error = pg.ErrorInfo.from_exception(e) + error = pg.ErrorInfo.from_exception(e) # pyrefly: ignore[bad-assignment] self._output = output else: assert lm_response is not None diff --git a/langfun/core/structured/schema/base.py b/langfun/core/structured/schema/base.py index cafc3b4d..f75f1829 100644 --- a/langfun/core/structured/schema/base.py +++ b/langfun/core/structured/schema/base.py @@ -61,7 +61,7 @@ def _parse_node(v) -> pg.typing.ValueSpec: spec = pg.typing.ValueSpec.from_annotation(v, auto_typing=True) if isinstance( spec, - ( + ( # pyrefly: ignore[invalid-argument] pg.typing.Any, pg.typing.Callable, pg.typing.Tuple, @@ -397,13 +397,13 @@ def _fill_dependencies(vs: pg.typing.ValueSpec, include_subclasses: bool): _fill_dependencies(v, include_subclasses) for value_spec in value_specs: - _fill_dependencies(value_spec, include_subclasses) + _fill_dependencies(value_spec, include_subclasses) # pyrefly: ignore[bad-argument-type] return dependencies def schema_spec(noneable: bool = False) -> pg.typing.ValueSpec: # pylint: disable=unused-argument if typing.TYPE_CHECKING: - return Any + return Any # pyrefly: ignore[bad-return] return pg.typing.Object( Schema, transform=Schema.from_value, is_noneable=noneable ) # pylint: disable=unreachable-code @@ -419,7 +419,7 @@ def annotation( child_annotation_kwargs = dict( strict=strict, allowed_dependencies=allowed_dependencies ) - if isinstance(vs, pg.typing.Any): + if isinstance(vs, pg.typing.Any): # pyrefly: ignore[invalid-argument] return 'Any' elif isinstance(vs, pg.typing.Enum): candidate_str = ', '.join([repr(v) for v in vs.values]) diff --git a/langfun/core/structured/schema/json.py b/langfun/core/structured/schema/json.py index cf282a39..ef5fc33c 100644 --- a/langfun/core/structured/schema/json.py +++ b/langfun/core/structured/schema/json.py @@ -71,7 +71,7 @@ def _visit(node: Any) -> None: f'"{v}"' if isinstance(v, str) else repr(v) for v in node.values)) elif isinstance(node, pg.typing.PrimitiveType): - x = node.value_type.__name__ + x = node.value_type.__name__ # pyrefly: ignore[missing-attribute] if isinstance(node, pg.typing.Number): params = [] if node.min_value is not None: diff --git a/langfun/core/structured/schema_generation.py b/langfun/core/structured/schema_generation.py index 9a63e3f9..2e9c3e16 100644 --- a/langfun/core/structured/schema_generation.py +++ b/langfun/core/structured/schema_generation.py @@ -152,7 +152,7 @@ def generate_class( call_kwargs = dict(skip_lm=skip_lm) if lm is not None: - call_kwargs['lm'] = lm + call_kwargs['lm'] = lm # pyrefly: ignore[bad-assignment] message = GenerateClass( input=prompt, context=name, diff --git a/langfun/core/subscription.py b/langfun/core/subscription.py index 5641cf9f..80a685be 100644 --- a/langfun/core/subscription.py +++ b/langfun/core/subscription.py @@ -97,7 +97,7 @@ def _map_sender_subscriber( for subscriber in subscriber_list: for sender in senders: - func(subscriber, sender) + func(subscriber, sender) # pyrefly: ignore[bad-argument-type] def _sender_info( self, sender: Union[Any, Type[Any], None] @@ -190,7 +190,7 @@ def subscribe( ) -> None: """Subscribes one or a list subscribers to one or a list of senders.""" return self._map_sender_subscriber( - self._subscribe, + self._subscribe, # pyrefly: ignore[bad-argument-type] sender_or_senders=sender, subscriber_or_subscribers=subscriber, ) @@ -204,7 +204,7 @@ def unsubscribe( ) -> None: """Unsubscribes one or a list subscribers from one or a list of senders.""" return self._map_sender_subscriber( - self._unsubscribe, + self._unsubscribe, # pyrefly: ignore[bad-argument-type] sender_or_senders=sender, subscriber_or_subscribers=subscriber, ) @@ -230,7 +230,7 @@ def subscribers(self, sender: Any | Type[Any]) -> Iterator[EventHandler[Any]]: # Yield type_level subscribers. for registered_type, subscriber_list in self._sender_type_registry.items(): if isinstance(sender, registered_type) or ( - sender is None and issubclass(sender_type, registered_type) + sender is None and issubclass(sender_type, registered_type) # pyrefly: ignore[bad-argument-type] ): for subscriber in subscriber_list: if id(subscriber) not in visited: diff --git a/langfun/core/template.py b/langfun/core/template.py index 33348d84..3cb16059 100644 --- a/langfun/core/template.py +++ b/langfun/core/template.py @@ -524,8 +524,8 @@ def additional_metadata(self) -> dict[str, Any]: # Carry metadata from fields. for k, v in self.sym_init_args.sym_items(): - if k.startswith(_ADDITIONAL_METADATA_PREFIX): - metadata[k.removeprefix(_ADDITIONAL_METADATA_PREFIX)] = v + if k.startswith(_ADDITIONAL_METADATA_PREFIX): # pyrefly: ignore[missing-attribute] + metadata[k.removeprefix(_ADDITIONAL_METADATA_PREFIX)] = v # pyrefly: ignore[missing-attribute] return metadata # @@ -562,7 +562,7 @@ def natural_language_format(self) -> str: def _sym_clone(self, *args, **kwargs) -> 'Template': copy = super()._sym_clone(*args, **kwargs) copy._referred_modalities = self._referred_modalities # pylint: disable=protected-access - return copy + return copy # pyrefly: ignore[bad-return] def __eq__(self, other: Any) -> bool: if isinstance(other, str): @@ -669,7 +669,7 @@ def from_value( ) -> 'Template': """Creates a template object from a value.""" if isinstance(value, cls): - return value.clone(override=kwargs) if kwargs else value # pylint: disable=no-value-for-parameter + return value.clone(override=kwargs) if kwargs else value # pylint: disable=no-value-for-parameter # pyrefly: ignore[bad-return] if isinstance(value, str): return cls(template_str=value, **kwargs) if isinstance(value, Template): @@ -681,7 +681,7 @@ def from_value( if message_lib.Message.is_convertible(type(value)): value = message_lib.Message.from_value(value) if isinstance(value, message_lib.Message): - for k, v in value.metadata.sym_items(): # pylint: disable=attribute-error + for k, v in value.metadata.sym_items(): # pylint: disable=attribute-error # pyrefly: ignore[missing-attribute] kwargs[_ADDITIONAL_METADATA_PREFIX + k] = v t = cls(template_str=value.text, **kwargs) t._referred_modalities = value.referred_modalities @@ -716,10 +716,10 @@ def render_fields(): return view.complex_value( {k: v for k, v in self.sym_items()}, name='fields', - root_path=root_path, + root_path=root_path, # pyrefly: ignore[bad-argument-type] parent=self, exclude_keys=['template_str', 'clean'], - collapse_level=max( + collapse_level=max( # pyrefly: ignore[bad-specialization] collapse_template_vars_level, collapse_level ) if collapse_level is not None else None, extra_flags=extra_flags, @@ -773,7 +773,7 @@ def _html_tree_view_config(cls) -> dict[str, Any]: # Register converter from str to LangFunc, therefore we can always # pass strs to attributes that accept LangFunc. -pg.typing.register_converter(str, Template, Template) +pg.typing.register_converter(str, Template, Template) # pyrefly: ignore[bad-argument-type] @dataclasses.dataclass @@ -854,10 +854,10 @@ def __mod__(self, other: Any) -> '_UnresolvedExpression': def __rmod__(self, other: Any) -> '_UnresolvedExpression': return _UnresolvedExpression(f'{other!r} % {self.expression}') - def __eq__(self, other: Any) -> '_UnresolvedExpression': + def __eq__(self, other: Any) -> '_UnresolvedExpression': # pyrefly: ignore[bad-override] return _UnresolvedExpression(f'{self.expression} == {other!r}') - def __ne__(self, other: Any) -> '_UnresolvedExpression': + def __ne__(self, other: Any) -> '_UnresolvedExpression': # pyrefly: ignore[bad-override] return _UnresolvedExpression(f'{self.expression} != {other!r}') def __lt__(self, other: Any) -> '_UnresolvedExpression': diff --git a/langfun/core/templates/completion.py b/langfun/core/templates/completion.py index 60a0d082..b9dadd72 100644 --- a/langfun/core/templates/completion.py +++ b/langfun/core/templates/completion.py @@ -61,7 +61,7 @@ def __call__(self, **kwargs) -> lf.Message: lm_response = super().__call__(**kwargs) if self.cache_response: self.rebind(lm_response=lm_response, skip_notification=True) - return lm_response + return lm_response # pyrefly: ignore[bad-return] def clear_lm_response(self): """Clear LM response.""" diff --git a/langfun/core/templates/conversation.py b/langfun/core/templates/conversation.py index 612d7586..31b17584 100644 --- a/langfun/core/templates/conversation.py +++ b/langfun/core/templates/conversation.py @@ -23,6 +23,7 @@ class Conversation(Completion): """LM-based conversation.""" + # pyrefly: ignore[bad-assignment] prompt = """ {%- if preamble -%} {{ preamble }} diff --git a/langfun/core/templates/selfplay.py b/langfun/core/templates/selfplay.py index e991062f..3bcf1560 100644 --- a/langfun/core/templates/selfplay.py +++ b/langfun/core/templates/selfplay.py @@ -42,7 +42,7 @@ def __call__(self, **kwargs) -> lf.Message: output = self.step(**kwargs) if output is None: break - return self._last_response + return self._last_response # pyrefly: ignore[bad-return] def step(self, **kwargs) -> lf.Message | None: """Play the next step and return the response.""" diff --git a/langfun/env/base_feature.py b/langfun/env/base_feature.py index 61395a30..4bd54e52 100644 --- a/langfun/env/base_feature.py +++ b/langfun/env/base_feature.py @@ -106,7 +106,7 @@ def environment(self) -> interface.AbstractEnvironment | None: env = self.sym_ancestor( lambda v: isinstance(v, interface.AbstractEnvironment) ) - return env + return env # pyrefly: ignore[bad-return] @property def sandbox(self) -> interface.Sandbox | None: @@ -117,7 +117,7 @@ def sandbox(self) -> interface.Sandbox | None: return self._sandbox @property - def event_handler(self) -> interface.EventHandler: + def event_handler(self) -> interface.EventHandler: # pyrefly: ignore[bad-override] if hasattr(self, '_event_handler_ref'): return self._event_handler_ref return super().event_handler @@ -126,7 +126,7 @@ def event_handler(self) -> interface.EventHandler: def is_online(self) -> bool: """Returns True if the feature is online.""" if self.is_sandbox_based: - return self.sandbox.is_online + return self.sandbox.is_online # pyrefly: ignore[missing-attribute] return self._is_online @property @@ -140,7 +140,7 @@ def offline_duration(self) -> float: def working_dir(self) -> str | None: """Returns the working directory of the feature.""" if self.is_sandbox_based: - sandbox_workdir = self.sandbox.working_dir + sandbox_workdir = self.sandbox.working_dir # pyrefly: ignore[missing-attribute] if sandbox_workdir is None: return None return os.path.join(sandbox_workdir, self.name) @@ -292,7 +292,7 @@ def on_teardown_session( ) -> None: """Called when the feature is teardown for a user session.""" self.event_handler.on_feature_teardown_session( - feature=self, session_id=self.session_id, duration=duration, error=error + feature=self, session_id=self.session_id, duration=duration, error=error # pyrefly: ignore[bad-argument-type] ) def on_activity( diff --git a/langfun/env/base_sandbox.py b/langfun/env/base_sandbox.py index 0e740188..cbdae42d 100644 --- a/langfun/env/base_sandbox.py +++ b/langfun/env/base_sandbox.py @@ -259,7 +259,7 @@ def state_errors(self) -> list[interface.SandboxStateError]: @property def is_shutting_down(self) -> bool: """Returns True if the sandbox is shutting down.""" - return self._status == self.Status.SHUTTING_DOWN or ( + return self._status == self.Status.SHUTTING_DOWN or ( # pyrefly: ignore[bad-return] self._state_errors and self._status == self.Status.EXITING_SESSION ) @@ -824,6 +824,6 @@ def on_session_end( sandbox=self, session_id=session_id, duration=duration, - lifetime=time.time() - self._session_start_time, + lifetime=time.time() - self._session_start_time, # pyrefly: ignore[unsupported-operation] error=error, ) diff --git a/langfun/env/base_sandbox_service.py b/langfun/env/base_sandbox_service.py index 78d20143..2a4a2fe2 100644 --- a/langfun/env/base_sandbox_service.py +++ b/langfun/env/base_sandbox_service.py @@ -159,7 +159,7 @@ def _check_feature_requirements(self) -> None: ) @property - def event_handler(self) -> interface.EventHandler: + def event_handler(self) -> interface.EventHandler: # pyrefly: ignore[bad-override] if hasattr(self, '_event_handler_ref'): return self._event_handler_ref return super().event_handler @@ -358,7 +358,7 @@ def _bring_up_sandbox_with_retry( reusable=reusable, ) except (interface.EnvironmentError, interface.SandboxStateError) as e: - self._report_outage_or_wait(e, shutdown_env_upon_outage) + self._report_outage_or_wait(e, shutdown_env_upon_outage) # pyrefly: ignore[bad-argument-type] def _report_outage_or_wait( self, @@ -569,7 +569,7 @@ def _start(self) -> None: min_pool_size = self.min_pool_size(image_id) for i in range(min_pool_size): sandbox_startup_infos.append((image_id, i)) - self._sandbox_pool[image_id] = [None] * min_pool_size + self._sandbox_pool[image_id] = [None] * min_pool_size # pyrefly: ignore[unsupported-operation] next_sandbox_id = min_pool_size self._next_sandbox_id[image_id] = next_sandbox_id @@ -630,18 +630,18 @@ def _acquire( image_id: str | None = None ) -> interface.Sandbox: """Acquires a sandbox from the sandbox service.""" - if not self.enable_pooling(image_id): - return super()._acquire(image_id) + if not self.enable_pooling(image_id): # pyrefly: ignore[bad-argument-type] + return super()._acquire(image_id) # pyrefly: ignore[bad-argument-type] allocation_start_time = time.time() - sandbox_pool = self._sandbox_pool[image_id] + sandbox_pool = self._sandbox_pool[image_id] # pyrefly: ignore[bad-index] while True: try: # We only append or replace items in the sandbox pool, therefore # there is no need to lock the pool. return self.load_balancer.acquire(sandbox_pool) except IndexError: - if len(sandbox_pool) == self.max_pool_size(image_id): + if len(sandbox_pool) == self.max_pool_size(image_id): # pyrefly: ignore[bad-argument-type] if time.time() - allocation_start_time > self.outage_grace_period: raise interface.SandboxServiceOverloadError( # pylint: disable=raise-missing-from sandbox_service=self @@ -650,8 +650,8 @@ def _acquire( else: try: sandbox = self._bring_up_sandbox( - image_id=image_id, - sandbox_id=f'{self._increment_sandbox_id(image_id)}:0', + image_id=image_id, # pyrefly: ignore[bad-argument-type] + sandbox_id=f'{self._increment_sandbox_id(image_id)}:0', # pyrefly: ignore[bad-argument-type] set_acquired=True, reusable=True, ) @@ -661,7 +661,7 @@ def _acquire( except ( interface.EnvironmentError, interface.SandboxStateError ) as ex: - self._report_outage_or_wait(ex) + self._report_outage_or_wait(ex) # pyrefly: ignore[bad-argument-type] def _increment_sandbox_id(self, image_id: str) -> int: """Returns the next pooled sandbox ID.""" diff --git a/langfun/env/environment.py b/langfun/env/environment.py index 27c876cb..d3d2e7df 100644 --- a/langfun/env/environment.py +++ b/langfun/env/environment.py @@ -504,15 +504,15 @@ def _get_sandbox_service( """Returns the sandbox service for the given image ID.""" if sandbox_service is not None: return self.sandboxes[sandbox_service] - for sandbox_service in self.sandboxes.values(): - if image_id is None or image_id in sandbox_service.image_ids: - return sandbox_service + for sandbox_service in self.sandboxes.values(): # pyrefly: ignore[bad-assignment] + if image_id is None or image_id in sandbox_service.image_ids: # pyrefly: ignore[missing-attribute] + return sandbox_service # pyrefly: ignore[bad-return] # Returns the first sandbox service that supports dynamic image loading # if image ID is not found in pre-configured image IDs. - for sandbox_service in self.sandboxes.values(): - if sandbox_service.supports_dynamic_image_loading: - return sandbox_service + for sandbox_service in self.sandboxes.values(): # pyrefly: ignore[bad-assignment] + if sandbox_service.supports_dynamic_image_loading: # pyrefly: ignore[missing-attribute] + return sandbox_service # pyrefly: ignore[bad-return] raise ValueError( f'Environment {self.id} does not serve image ID {image_id!r}.' ) diff --git a/langfun/env/event_handlers/event_logger.py b/langfun/env/event_handlers/event_logger.py index dcb7ebe5..8f35026f 100644 --- a/langfun/env/event_handlers/event_logger.py +++ b/langfun/env/event_handlers/event_logger.py @@ -509,7 +509,7 @@ def _write_log( styles: list[str] | None = None, ): message = self._maybe_colored( - message, color if error is None else 'red', styles=styles + message, color if error is None else 'red', styles=styles # pyrefly: ignore[bad-argument-type] ) if error is not None: pg.logging.error(message) @@ -531,6 +531,6 @@ def _write_log( ): print( self._maybe_colored( - message, color if error is None else 'red', styles=styles + message, color if error is None else 'red', styles=styles # pyrefly: ignore[bad-argument-type] ) ) diff --git a/langfun/env/event_handlers/metric_writer.py b/langfun/env/event_handlers/metric_writer.py index 3832ec26..05c0f14a 100644 --- a/langfun/env/event_handlers/metric_writer.py +++ b/langfun/env/event_handlers/metric_writer.py @@ -36,7 +36,7 @@ def _get_counter( return self._metric_collection.get_counter( name=name, description=description, - parameters=parameters, + parameters=parameters, # pyrefly: ignore[bad-argument-type] ) def _get_scalar( @@ -46,7 +46,7 @@ def _get_scalar( parameters: dict[str, type[str]] | None = None, ) -> pg.monitoring.Metric: return self._metric_collection.get_scalar( - name=name, description=description, parameters=parameters + name=name, description=description, parameters=parameters # pyrefly: ignore[bad-argument-type] ) def _get_distribution( @@ -56,7 +56,7 @@ def _get_distribution( parameters: dict[str, type[str]] | None = None, ) -> pg.monitoring.Metric: return self._metric_collection.get_distribution( - name=name, description=description, parameters=parameters + name=name, description=description, parameters=parameters # pyrefly: ignore[bad-argument-type] ) def _error_tag(self, error: BaseException | None) -> str: @@ -478,7 +478,7 @@ def on_environment_housekeep( **kwargs ) -> None: """Called when the environment is housekeeping.""" - self._environment_housekeep_duration_ms.record( + self._environment_housekeep_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, environment_id=str(environment.id), @@ -502,7 +502,7 @@ def on_sandbox_service_start( sandbox_service_name=sandbox_service.id.name, error=self._error_tag(error), ) - self._sandbox_service_start_duration_ms.record( + self._sandbox_service_start_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, environment_id=env_id, @@ -510,7 +510,7 @@ def on_sandbox_service_start( error=self._error_tag(error), ) if error is None: - self._sandbox_service_count.increment( + self._sandbox_service_count.increment( # pyrefly: ignore[missing-attribute] app=self.app, environment_id=env_id, sandbox_service_name=sandbox_service.id.name, @@ -535,21 +535,21 @@ def on_sandbox_service_shutdown( sandbox_service_name=sandbox_service.id.name, error=self._error_tag(error), ) - self._sandbox_service_shutdown_duration_ms.record( + self._sandbox_service_shutdown_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, environment_id=env_id, sandbox_service_name=sandbox_service.id.name, error=self._error_tag(error), ) - self._sandbox_service_lifetime_ms.record( + self._sandbox_service_lifetime_ms.record( # pyrefly: ignore[missing-attribute] int(lifetime * 1000), app=self.app, environment_id=env_id, sandbox_service_name=sandbox_service.id.name, error=self._error_tag(error), ) - self._sandbox_service_count.increment( + self._sandbox_service_count.increment( # pyrefly: ignore[missing-attribute] delta=-1, app=self.app, environment_id=env_id, @@ -576,7 +576,7 @@ def on_sandbox_service_housekeep( sandbox_service_name=sandbox_service.id.name, error=self._error_tag(error), ) - self._sandbox_service_housekeep_duration_ms.record( + self._sandbox_service_housekeep_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, environment_id=env_id, @@ -592,14 +592,14 @@ def on_sandbox_start( ) -> None: self._sandbox_start.increment( app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, error=self._error_tag(error), ) - self._sandbox_start_duration_ms.record( + self._sandbox_start_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, error=self._error_tag(error), ) @@ -611,25 +611,25 @@ def on_sandbox_status_change( new_status: interface.Sandbox.Status, span: float, ) -> None: - self._sandbox_status_duration_ms.record( + self._sandbox_status_duration_ms.record( # pyrefly: ignore[missing-attribute] int(span * 1000), app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, status=old_status.value, ) if old_status != interface.Sandbox.Status.CREATED: - self._sandbox_count.increment( + self._sandbox_count.increment( # pyrefly: ignore[missing-attribute] delta=-1, app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, status=old_status.value, ) if new_status != interface.Sandbox.Status.OFFLINE: - self._sandbox_count.increment( + self._sandbox_count.increment( # pyrefly: ignore[missing-attribute] app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, status=new_status.value, ) @@ -643,21 +643,21 @@ def on_sandbox_shutdown( ) -> None: self._sandbox_shutdown.increment( app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, error=self._error_tag(error), ) - self._sandbox_shutdown_duration_ms.record( + self._sandbox_shutdown_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, error=self._error_tag(error), ) - self._sandbox_lifetime_ms.record( + self._sandbox_lifetime_ms.record( # pyrefly: ignore[missing-attribute] int(lifetime * 1000), app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, error=self._error_tag(error), ) @@ -670,10 +670,10 @@ def on_sandbox_session_start( error: BaseException | None, ) -> None: """Called when a sandbox session starts.""" - self._sandbox_session_start_duration_ms.record( + self._sandbox_session_start_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, error=self._error_tag(error), ) @@ -687,17 +687,17 @@ def on_sandbox_session_end( error: BaseException | None, ) -> None: """Called when a sandbox session ends.""" - self._sandbox_session_end_duration_ms.record( + self._sandbox_session_end_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, error=self._error_tag(error), ) - self._sandbox_session_lifetime_ms.record( + self._sandbox_session_lifetime_ms.record( # pyrefly: ignore[missing-attribute] int(lifetime * 1000), app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, error=self._error_tag(error), ) @@ -714,15 +714,15 @@ def on_sandbox_activity( """Called when a sandbox activity is performed.""" self._sandbox_activity.increment( app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, activity=name, error=self._error_tag(error), ) - self._sandbox_activity_duration_ms.record( + self._sandbox_activity_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, activity=name, error=self._error_tag(error), @@ -739,14 +739,14 @@ def on_sandbox_housekeep( """Called when a sandbox feature is housekeeping.""" self._sandbox_housekeep.increment( app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, error=self._error_tag(error), ) - self._sandbox_housekeep_duration_ms.record( + self._sandbox_housekeep_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, - environment_id=str(sandbox.environment.id), + environment_id=str(sandbox.environment.id), # pyrefly: ignore[missing-attribute] image_id=sandbox.image_id, error=self._error_tag(error), ) @@ -761,15 +761,15 @@ def on_feature_setup( image_id = feature.sandbox.image_id if feature.sandbox else '' self._feature_setup.increment( app=self.app, - environment_id=str(feature.environment.id), + environment_id=str(feature.environment.id), # pyrefly: ignore[missing-attribute] image_id=image_id, feature_name=feature.name, error=self._error_tag(error), ) - self._feature_setup_duration_ms.record( + self._feature_setup_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, - environment_id=str(feature.environment.id), + environment_id=str(feature.environment.id), # pyrefly: ignore[missing-attribute] image_id=image_id, feature_name=feature.name, error=self._error_tag(error), @@ -785,15 +785,15 @@ def on_feature_teardown( image_id = feature.sandbox.image_id if feature.sandbox else '' self._feature_teardown.increment( app=self.app, - environment_id=str(feature.environment.id), + environment_id=str(feature.environment.id), # pyrefly: ignore[missing-attribute] image_id=image_id, feature_name=feature.name, error=self._error_tag(error), ) - self._feature_teardown_duration_ms.record( + self._feature_teardown_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, - environment_id=str(feature.environment.id), + environment_id=str(feature.environment.id), # pyrefly: ignore[missing-attribute] image_id=image_id, feature_name=feature.name, error=self._error_tag(error), @@ -810,15 +810,15 @@ def on_feature_setup_session( image_id = feature.sandbox.image_id if feature.sandbox else '' self._feature_setup_session.increment( app=self.app, - environment_id=str(feature.environment.id), + environment_id=str(feature.environment.id), # pyrefly: ignore[missing-attribute] image_id=image_id, feature_name=feature.name, error=self._error_tag(error), ) - self._feature_setup_session_duration_ms.record( + self._feature_setup_session_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, - environment_id=str(feature.environment.id), + environment_id=str(feature.environment.id), # pyrefly: ignore[missing-attribute] image_id=image_id, feature_name=feature.name, error=self._error_tag(error), @@ -835,15 +835,15 @@ def on_feature_teardown_session( image_id = feature.sandbox.image_id if feature.sandbox else '' self._feature_teardown_session.increment( app=self.app, - environment_id=str(feature.environment.id), + environment_id=str(feature.environment.id), # pyrefly: ignore[missing-attribute] image_id=image_id, feature_name=feature.name, error=self._error_tag(error), ) - self._feature_teardown_session_duration_ms.record( + self._feature_teardown_session_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, - environment_id=str(feature.environment.id), + environment_id=str(feature.environment.id), # pyrefly: ignore[missing-attribute] image_id=image_id, feature_name=feature.name, error=self._error_tag(error), @@ -862,15 +862,15 @@ def on_feature_activity( image_id = feature.sandbox.image_id if feature.sandbox else '' self._feature_activity.increment( app=self.app, - environment_id=str(feature.environment.id), + environment_id=str(feature.environment.id), # pyrefly: ignore[missing-attribute] image_id=image_id, activity=name, error=self._error_tag(error), ) - self._feature_activity_duration_ms.record( + self._feature_activity_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, - environment_id=str(feature.environment.id), + environment_id=str(feature.environment.id), # pyrefly: ignore[missing-attribute] image_id=image_id, activity=name, error=self._error_tag(error), @@ -888,15 +888,15 @@ def on_feature_housekeep( image_id = feature.sandbox.image_id if feature.sandbox else '' self._feature_housekeep.increment( app=self.app, - environment_id=str(feature.environment.id), + environment_id=str(feature.environment.id), # pyrefly: ignore[missing-attribute] image_id=image_id, feature_name=feature.name, error=self._error_tag(error), ) - self._feature_housekeep_duration_ms.record( + self._feature_housekeep_duration_ms.record( # pyrefly: ignore[missing-attribute] int(duration * 1000), app=self.app, - environment_id=str(feature.environment.id), + environment_id=str(feature.environment.id), # pyrefly: ignore[missing-attribute] image_id=image_id, feature_name=feature.name, error=self._error_tag(error), diff --git a/langfun/env/interface.py b/langfun/env/interface.py index 5a422b95..beb6263f 100644 --- a/langfun/env/interface.py +++ b/langfun/env/interface.py @@ -708,7 +708,7 @@ def working_dir(self, root_dir: str | None) -> str | None: return None if self.container_id is None: return os.path.join(root_dir, _make_path_compatible(self.feature_name)) - return os.path.join( + return os.path.join( # pyrefly: ignore[no-matching-overload] self.container_id.working_dir(root_dir), _make_path_compatible(self.feature_name) ) @@ -778,7 +778,7 @@ def name(self) -> str: def id(self) -> Id: """Returns the identifier of the feature.""" if self.is_sandbox_based: - return Feature.Id(self.sandbox.id, self.name) + return Feature.Id(self.sandbox.id, self.name) # pyrefly: ignore[missing-attribute] if self.environment is not None: return Feature.Id(self.environment.id, self.name) return Feature.Id(None, pg.utils.camel_to_snake(self.__class__.__name__)) @@ -913,7 +913,7 @@ def track_activity( def session_id(self) -> str | None: """Returns the current user session identifier.""" if self.is_sandbox_based: - return self.sandbox.session_id + return self.sandbox.session_id # pyrefly: ignore[missing-attribute] return self._non_sandbox_based_session_id @contextlib.contextmanager @@ -986,7 +986,7 @@ def working_dir(self, root_dir: str | None) -> str | None: """Returns the download directory for the sandbox.""" if root_dir is None: return None - return os.path.join( + return os.path.join( # pyrefly: ignore[no-matching-overload] self.service_id.working_dir(root_dir), _make_path_compatible(self.image_id), _make_path_compatible(self.sandbox_id) @@ -1392,7 +1392,7 @@ def working_dir(self, root_dir: str | None) -> str | None: return None if self.environment_id is None: return os.path.join(root_dir, self.name) - return os.path.join(self.environment_id.working_dir(root_dir), self.name) + return os.path.join(self.environment_id.working_dir(root_dir), self.name) # pyrefly: ignore[no-matching-overload] image_ids: Annotated[ list[str], @@ -1482,7 +1482,7 @@ def _on_path_change(self, old_path: pg.KeyPath, new_path: pg.KeyPath) -> None: @functools.cached_property def environment(self) -> Optional['AbstractEnvironment']: """Returns the containing environment.""" - return self.sym_ancestor(lambda v: isinstance(v, AbstractEnvironment)) + return self.sym_ancestor(lambda v: isinstance(v, AbstractEnvironment)) # pyrefly: ignore[bad-return] @functools.cached_property def name(self) -> str: @@ -2076,15 +2076,15 @@ def _get_sandbox_service( """Returns the sandbox service for the given image ID.""" if sandbox_service is not None: return self.sandboxes[sandbox_service] - for sandbox_service in self.sandboxes.values(): - if image_id is None or image_id in sandbox_service.image_ids: - return sandbox_service + for sandbox_service in self.sandboxes.values(): # pyrefly: ignore[bad-assignment] + if image_id is None or image_id in sandbox_service.image_ids: # pyrefly: ignore[missing-attribute] + return sandbox_service # pyrefly: ignore[bad-return] # Returns the first sandbox service that supports dynamic image loading # if image ID is not found in pre-configured image IDs. - for sandbox_service in self.sandboxes.values(): - if sandbox_service.supports_dynamic_image_loading: - return sandbox_service + for sandbox_service in self.sandboxes.values(): # pyrefly: ignore[bad-assignment] + if sandbox_service.supports_dynamic_image_loading: # pyrefly: ignore[missing-attribute] + return sandbox_service # pyrefly: ignore[bad-return] raise ValueError( f'Environment {self.id} does not serve image ID {image_id!r}.' ) @@ -2229,7 +2229,7 @@ def method_wrapper(self, *args, **kwargs) -> Any: # Execute the service function. return func(self, *args, **kwargs) except BaseException as e: - if pg.match_error(e, errors): + if pg.match_error(e, errors): # pyrefly: ignore[bad-argument-type] state_error = SandboxStateError( 'Sandbox encountered an unexpected error executing ' f'`{func.__name__}` (args={args!r}, kwargs={kwargs!r}): {e}', diff --git a/langfun/env/test_utils.py b/langfun/env/test_utils.py index cfc15f94..4eb0d819 100644 --- a/langfun/env/test_utils.py +++ b/langfun/env/test_utils.py @@ -167,19 +167,19 @@ def do(self, code: str, raise_error: Type[BaseException] | None = None): self._sandbox.shell(code, raise_error=raise_error) def _raise_error(self, message, error_type: Type[BaseException], **kwargs): - self._sandbox._raise_error(message, error_type, **kwargs) # pylint: disable=protected-access + self._sandbox._raise_error(message, error_type, **kwargs) # pylint: disable=protected-access # pyrefly: ignore[missing-attribute] def _setup(self) -> None: if self.simulate_setup_error: self._raise_error(f'{self.name} setup error', self.simulate_setup_error) - self.sandbox.shell(f'"{self.name}" setup') + self.sandbox.shell(f'"{self.name}" setup') # pyrefly: ignore[missing-attribute] def _teardown(self) -> None: if self.simulate_teardown_error: self._raise_error( f'{self.name} teardown error', self.simulate_teardown_error ) - self.sandbox.shell(f'"{self.name}" teardown') + self.sandbox.shell(f'"{self.name}" teardown') # pyrefly: ignore[missing-attribute] def _setup_session(self) -> None: if self.setup_session_delay > 0: @@ -189,24 +189,24 @@ def _setup_session(self) -> None: self._raise_error( 'Feature session setup error', self.simulate_setup_session_error ) - self.sandbox.shell(f'"{self.name}" setup session') + self.sandbox.shell(f'"{self.name}" setup session') # pyrefly: ignore[missing-attribute] def _teardown_session(self) -> None: if self.simulate_teardown_session_error: self._raise_error( 'Feature session teardown error', self.simulate_teardown_session_error ) - self.sandbox.shell(f'"{self.name}" teardown session') + self.sandbox.shell(f'"{self.name}" teardown session') # pyrefly: ignore[missing-attribute] if self.call_end_session_on_teardown_session: - self.sandbox.end_session() + self.sandbox.end_session() # pyrefly: ignore[missing-attribute] @interface.log_activity() def num_shell_calls(self) -> int: - return len(self.sandbox._shell_history) # pylint: disable=protected-access + return len(self.sandbox._shell_history) # pylint: disable=protected-access # pyrefly: ignore[missing-attribute] @interface.log_activity() def bad_shell_call(self) -> None: - self.sandbox.shell('bad command', raise_error=RuntimeError) + self.sandbox.shell('bad command', raise_error=RuntimeError) # pyrefly: ignore[missing-attribute] @interface.log_activity() def show_session_id(self): @@ -224,7 +224,7 @@ def _on_bound(self) -> None: @contextlib.contextmanager def my_service(self) -> Iterator[Service]: try: - self._service = TestingFeature.Service(sandbox=self.sandbox) + self._service = TestingFeature.Service(sandbox=self.sandbox) # pyrefly: ignore[bad-argument-type] yield self._service finally: self._service = None diff --git a/requirements.txt b/requirements.txt index f73a512e..f3189020 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ anyio>=4.7.0 jinja2>=3.1.2 -mcp>=1.17.0 +mcp>=1.17.0,<2.0.0 puremagic>=1.20 pyglove>=0.5.0.dev202510170226 requests>=2.31.0