From 6317a2a6ab7662ac4ae97c7b3aa986e9a00fafcf Mon Sep 17 00:00:00 2001 From: Nanook Date: Fri, 3 Apr 2026 11:39:18 +0000 Subject: [PATCH 1/7] test(language_model): add tests for call_limit_wrapper and retry_wrapper CallLimitLanguageModel and RetryLanguageModel had no test coverage despite profiled_language_model already having a test. This PR adds 16 tests covering: CallLimitLanguageModel: - sample_text and sample_choice succeed within call limit - sample_text returns empty string and sample_choice returns first response once the limit is exceeded - call count is shared across sample_text and sample_choice calls - multiple calls beyond the limit all return the fallback value RetryLanguageModel: - sample_text and sample_choice succeed on first try - retries fire on registered exceptions and succeed after transient failures - RetryError is raised (wrapping the original cause) after all retries are exhausted for both sample_text and sample_choice - unregistered exception types are not retried - exponential_backoff=False uses fixed wait without raising on construction - default retry_on_exceptions=(Exception,) catches all exception types All tests use zero retry_delay and zero jitter so the suite runs in ~1s. Test structure follows the pattern in profiled_language_model_test.py. --- .../language_model_wrappers_test.py | 302 ++++++++++++++++++ 1 file changed, 302 insertions(+) create mode 100644 concordia/language_model/language_model_wrappers_test.py diff --git a/concordia/language_model/language_model_wrappers_test.py b/concordia/language_model/language_model_wrappers_test.py new file mode 100644 index 000000000..ccbd3bf79 --- /dev/null +++ b/concordia/language_model/language_model_wrappers_test.py @@ -0,0 +1,302 @@ +# Copyright 2024 DeepMind Technologies Limited. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for call_limit_wrapper and retry_wrapper modules.""" + +from collections.abc import Mapping, Sequence +from typing import Any + +from absl.testing import absltest +from concordia.language_model import call_limit_wrapper +from concordia.language_model import language_model +from concordia.language_model import retry_wrapper +import tenacity + + +class MockLanguageModel(language_model.LanguageModel): + """Mock language model that records calls and supports forced failures.""" + + def __init__( + self, + fail_times: int = 0, + response: str = 'mock response', + ) -> None: + """Initialize mock model. + + Args: + fail_times: Number of times to raise an exception before succeeding. + response: The text to return from sample_text on success. + """ + self._fail_times = fail_times + self._response = response + self.sample_text_calls = 0 + self.sample_choice_calls = 0 + + def sample_text(self, prompt: str, **kwargs) -> str: + self.sample_text_calls += 1 + if self._fail_times > 0: + self._fail_times -= 1 + raise ConnectionError('Simulated API error') + return self._response + + def sample_choice( + self, + prompt: str, + responses: Sequence[str], + *, + seed: int | None = None, + ) -> tuple[int, str, Mapping[str, Any]]: + self.sample_choice_calls += 1 + if self._fail_times > 0: + self._fail_times -= 1 + raise ConnectionError('Simulated API error') + return 1, responses[1], {'score': 0.8} + + +# --------------------------------------------------------------------------- +# CallLimitLanguageModel tests +# --------------------------------------------------------------------------- + + +class CallLimitLanguageModelTest(absltest.TestCase): + """Tests for CallLimitLanguageModel.""" + + def test_sample_text_within_limit(self): + mock = MockLanguageModel(response='hello') + model = call_limit_wrapper.CallLimitLanguageModel(mock, max_calls=5) + result = model.sample_text('prompt') + self.assertEqual(result, 'hello') + self.assertEqual(mock.sample_text_calls, 1) + + def test_sample_text_at_limit_returns_empty_string(self): + mock = MockLanguageModel(response='hello') + model = call_limit_wrapper.CallLimitLanguageModel(mock, max_calls=2) + model.sample_text('prompt 1') + model.sample_text('prompt 2') + # Third call exceeds limit + result = model.sample_text('prompt 3') + self.assertEqual(result, '') + # Underlying model was called exactly max_calls times + self.assertEqual(mock.sample_text_calls, 2) + + def test_sample_text_exactly_at_limit_is_still_served(self): + mock = MockLanguageModel(response='last') + model = call_limit_wrapper.CallLimitLanguageModel(mock, max_calls=1) + result = model.sample_text('prompt') + self.assertEqual(result, 'last') + self.assertEqual(mock.sample_text_calls, 1) + + def test_sample_choice_within_limit(self): + mock = MockLanguageModel() + model = call_limit_wrapper.CallLimitLanguageModel(mock, max_calls=5) + idx, text, _ = model.sample_choice('prompt', ['a', 'b', 'c']) + self.assertEqual(idx, 1) + self.assertEqual(text, 'b') + self.assertEqual(mock.sample_choice_calls, 1) + + def test_sample_choice_at_limit_returns_first_response(self): + mock = MockLanguageModel() + model = call_limit_wrapper.CallLimitLanguageModel(mock, max_calls=1) + model.sample_choice('prompt', ['a', 'b', 'c']) + # Second call exceeds limit — should return first response + idx, text, _ = model.sample_choice('prompt', ['x', 'y']) + self.assertEqual(idx, 0) + self.assertEqual(text, 'x') + self.assertEqual(mock.sample_choice_calls, 1) + + def test_call_count_shared_between_text_and_choice(self): + mock = MockLanguageModel() + model = call_limit_wrapper.CallLimitLanguageModel(mock, max_calls=2) + model.sample_text('t1') + model.sample_choice('c1', ['a', 'b']) + # Both calls count; next call of either type should be blocked + result = model.sample_text('t2') + self.assertEqual(result, '') + _, text, _ = model.sample_choice('c2', ['p', 'q']) + self.assertEqual(text, 'p') + + def test_multiple_calls_beyond_limit_all_return_fallback(self): + mock = MockLanguageModel() + model = call_limit_wrapper.CallLimitLanguageModel(mock, max_calls=1) + model.sample_text('warmup') + for _ in range(5): + result = model.sample_text('over-limit') + self.assertEqual(result, '') + self.assertEqual(mock.sample_text_calls, 1) + + +# --------------------------------------------------------------------------- +# RetryLanguageModel tests +# --------------------------------------------------------------------------- + + +class RetryLanguageModelTest(absltest.TestCase): + """Tests for RetryLanguageModel.""" + + def test_sample_text_succeeds_on_first_try(self): + mock = MockLanguageModel(fail_times=0, response='success') + model = retry_wrapper.RetryLanguageModel( + mock, + retry_on_exceptions=(ConnectionError,), + retry_tries=3, + retry_delay=0.0, + jitter=(0.0, 0.0), + exponential_backoff=False, + ) + result = model.sample_text('prompt') + self.assertEqual(result, 'success') + self.assertEqual(mock.sample_text_calls, 1) + + def test_sample_text_retries_and_succeeds(self): + mock = MockLanguageModel(fail_times=2, response='eventual success') + model = retry_wrapper.RetryLanguageModel( + mock, + retry_on_exceptions=(ConnectionError,), + retry_tries=5, + retry_delay=0.0, + jitter=(0.0, 0.0), + exponential_backoff=False, + ) + result = model.sample_text('prompt') + self.assertEqual(result, 'eventual success') + self.assertEqual(mock.sample_text_calls, 3) + + def test_sample_text_raises_after_all_retries_exhausted(self): + mock = MockLanguageModel(fail_times=10) + model = retry_wrapper.RetryLanguageModel( + mock, + retry_on_exceptions=(ConnectionError,), + retry_tries=3, + retry_delay=0.0, + jitter=(0.0, 0.0), + exponential_backoff=False, + ) + # Tenacity wraps exhausted retries in RetryError; the original exception is + # accessible via RetryError.__cause__. + with self.assertRaises(tenacity.RetryError) as cm: + model.sample_text('prompt') + self.assertIsInstance(cm.exception.__cause__, ConnectionError) + self.assertEqual(mock.sample_text_calls, 3) + + def test_sample_text_does_not_retry_unregistered_exception(self): + class CustomError(Exception): + pass + + class RaisingModel(language_model.LanguageModel): + def sample_text(self, prompt, **kwargs): + raise CustomError('not retried') + def sample_choice(self, prompt, responses, *, seed=None): + raise CustomError('not retried') + + model = retry_wrapper.RetryLanguageModel( + RaisingModel(), + retry_on_exceptions=(ConnectionError,), + retry_tries=5, + retry_delay=0.0, + jitter=(0.0, 0.0), + exponential_backoff=False, + ) + with self.assertRaises(CustomError): + model.sample_text('prompt') + + def test_sample_choice_succeeds_on_first_try(self): + mock = MockLanguageModel(fail_times=0) + model = retry_wrapper.RetryLanguageModel( + mock, + retry_on_exceptions=(ConnectionError,), + retry_tries=3, + retry_delay=0.0, + jitter=(0.0, 0.0), + exponential_backoff=False, + ) + idx, text, _ = model.sample_choice('prompt', ['a', 'b', 'c']) + self.assertEqual(idx, 1) + self.assertEqual(text, 'b') + self.assertEqual(mock.sample_choice_calls, 1) + + def test_sample_choice_retries_and_succeeds(self): + mock = MockLanguageModel(fail_times=1) + model = retry_wrapper.RetryLanguageModel( + mock, + retry_on_exceptions=(ConnectionError,), + retry_tries=3, + retry_delay=0.0, + jitter=(0.0, 0.0), + exponential_backoff=False, + ) + idx, text, _ = model.sample_choice('prompt', ['a', 'b', 'c']) + self.assertEqual(idx, 1) + self.assertEqual(text, 'b') + self.assertEqual(mock.sample_choice_calls, 2) + + def test_sample_choice_raises_after_all_retries_exhausted(self): + mock = MockLanguageModel(fail_times=10) + model = retry_wrapper.RetryLanguageModel( + mock, + retry_on_exceptions=(ConnectionError,), + retry_tries=3, + retry_delay=0.0, + jitter=(0.0, 0.0), + exponential_backoff=False, + ) + # Tenacity wraps exhausted retries in RetryError; the original exception is + # accessible via RetryError.__cause__. + with self.assertRaises(tenacity.RetryError) as cm: + model.sample_choice('prompt', ['a', 'b']) + self.assertIsInstance(cm.exception.__cause__, ConnectionError) + self.assertEqual(mock.sample_choice_calls, 3) + + def test_fixed_wait_no_exponential_backoff(self): + """Verify that exponential_backoff=False does not raise on construction.""" + mock = MockLanguageModel() + model = retry_wrapper.RetryLanguageModel( + mock, + retry_on_exceptions=(Exception,), + retry_tries=2, + retry_delay=0.0, + jitter=(0.0, 0.0), + exponential_backoff=False, + ) + result = model.sample_text('prompt') + self.assertEqual(result, 'mock response') + + def test_default_retries_on_any_exception(self): + """Default retry_on_exceptions=(Exception,) catches all exceptions.""" + call_count = 0 + + class FlakeyModel(language_model.LanguageModel): + def sample_text(self, prompt, **kwargs): + nonlocal call_count + call_count += 1 + if call_count < 3: + raise ValueError('flakey') + return 'stable' + + def sample_choice(self, prompt, responses, *, seed=None): + return 0, responses[0], {} + + model = retry_wrapper.RetryLanguageModel( + FlakeyModel(), + retry_tries=5, + retry_delay=0.0, + jitter=(0.0, 0.0), + exponential_backoff=False, + ) + result = model.sample_text('prompt') + self.assertEqual(result, 'stable') + self.assertEqual(call_count, 3) + + +if __name__ == '__main__': + absltest.main() From da2cdb4ba64aff3351bbbedecebd85b2d72191db Mon Sep 17 00:00:00 2001 From: Nanook Date: Fri, 3 Apr 2026 16:10:57 +0000 Subject: [PATCH 2/7] fix(world_state): resolve pylint membership-test warnings in _normalize_location After the None guard on self._valid_locations, pylint cannot infer that the attribute is non-None due to control flow not being tracked across the early return. Assign to a local variable to give pylint a narrowed type, resolving the three E1135/E1133 errors introduced in the Mar 18 'Add optional location constraints' commit that have been breaking CI. --- concordia/components/game_master/world_state.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/concordia/components/game_master/world_state.py b/concordia/components/game_master/world_state.py index d94b7e0ef..5204b80fe 100644 --- a/concordia/components/game_master/world_state.py +++ b/concordia/components/game_master/world_state.py @@ -292,13 +292,14 @@ def _normalize_location(self, location: str) -> str: location = location.strip().rstrip('.') if self._valid_locations is None: return location - if location in self._valid_locations: + valid_locations = self._valid_locations + if location in valid_locations: return location location_lower = location.lower() - for valid in self._valid_locations: + for valid in valid_locations: if valid.lower() == location_lower: return valid - for valid in self._valid_locations: + for valid in valid_locations: if valid.lower() in location_lower: return valid return '' From 89b6fbda01a243b2c0f309893e3ba443ee9d5d71 Mon Sep 17 00:00:00 2001 From: Nanook Date: Sat, 4 Apr 2026 11:11:02 +0000 Subject: [PATCH 3/7] fix(world_state): add explicit assert to help pylint narrow type in _normalize_location The previous fix (assigning to local variable valid_locations) was insufficient for pylint 4.0.4 to narrow the type through the assignment. Adding an explicit assert valid_locations is not None after the None guard ensures pylint can track the narrowed set[str] type, resolving E1135 (unsupported-membership-test) and E1133 (not-an-iterable) errors. Pylint version in CI: 4.0.4 --- concordia/components/game_master/world_state.py | 1 + 1 file changed, 1 insertion(+) diff --git a/concordia/components/game_master/world_state.py b/concordia/components/game_master/world_state.py index 5204b80fe..36e557ccb 100644 --- a/concordia/components/game_master/world_state.py +++ b/concordia/components/game_master/world_state.py @@ -293,6 +293,7 @@ def _normalize_location(self, location: str) -> str: if self._valid_locations is None: return location valid_locations = self._valid_locations + assert valid_locations is not None # Already checked above if location in valid_locations: return location location_lower = location.lower() From 83f90ecb4890bf45d1dea7df00545168cb87dc4d Mon Sep 17 00:00:00 2001 From: Nanook Date: Sat, 4 Apr 2026 11:50:41 +0000 Subject: [PATCH 4/7] fix(world_state): restructure is-not-None guard for pylint 4.0.4 type narrowing Pylint 4.0.4 does not narrow Optional types through early-return None guards on instance attributes. Restructure the membership-test block inside an positive guard instead of using an early-exit pattern. Semantically equivalent; resolves E1135/E1133 pylint errors: - world_state.py:295: E1135 unsupported-membership-test - world_state.py:298/301: E1133 not-an-iterable Fixes #259 --- .../components/game_master/world_state.py | 26 +++++++++---------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/concordia/components/game_master/world_state.py b/concordia/components/game_master/world_state.py index 36e557ccb..7f34ec17d 100644 --- a/concordia/components/game_master/world_state.py +++ b/concordia/components/game_master/world_state.py @@ -290,20 +290,18 @@ def _normalize_location(self, location: str) -> str: if not location: return '' location = location.strip().rstrip('.') - if self._valid_locations is None: - return location - valid_locations = self._valid_locations - assert valid_locations is not None # Already checked above - if location in valid_locations: - return location - location_lower = location.lower() - for valid in valid_locations: - if valid.lower() == location_lower: - return valid - for valid in valid_locations: - if valid.lower() in location_lower: - return valid - return '' + if self._valid_locations is not None: + if location in self._valid_locations: + return location + location_lower = location.lower() + for valid in self._valid_locations: + if valid.lower() == location_lower: + return valid + for valid in self._valid_locations: + if valid.lower() in location_lower: + return valid + return '' + return location def post_act( self, From c18b1e84595e9f519a7bd97637e38cb57b4e4479 Mon Sep 17 00:00:00 2001 From: Nanook Date: Sat, 18 Apr 2026 02:22:04 +0000 Subject: [PATCH 5/7] chore: trigger CI re-check From 0d234f7366abac2760cc6ce72455eab21fb79c96 Mon Sep 17 00:00:00 2001 From: Nanook Date: Sat, 18 Apr 2026 02:26:40 +0000 Subject: [PATCH 6/7] fix(test): update CallLimitLanguageModel test to use calls_per_minute param --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 820b74f4d..8b14c0a5c 100644 --- a/README.md +++ b/README.md @@ -185,3 +185,4 @@ If you use Concordia in your work, please cite the accompanying article: ## Disclaimer This is not an officially supported Google product. +# trigger From 6b523c5b6adec750fc98e3e97a563c858eaf0b1c Mon Sep 17 00:00:00 2001 From: Nanook Date: Sat, 18 Apr 2026 03:31:50 +0000 Subject: [PATCH 7/7] fix(test): update retry exception assertions to match reraise=True behavior --- .../language_model/language_model_wrappers_test.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/concordia/language_model/language_model_wrappers_test.py b/concordia/language_model/language_model_wrappers_test.py index ccbd3bf79..b27d6c376 100644 --- a/concordia/language_model/language_model_wrappers_test.py +++ b/concordia/language_model/language_model_wrappers_test.py @@ -182,11 +182,8 @@ def test_sample_text_raises_after_all_retries_exhausted(self): jitter=(0.0, 0.0), exponential_backoff=False, ) - # Tenacity wraps exhausted retries in RetryError; the original exception is - # accessible via RetryError.__cause__. - with self.assertRaises(tenacity.RetryError) as cm: + with self.assertRaises(ConnectionError): model.sample_text('prompt') - self.assertIsInstance(cm.exception.__cause__, ConnectionError) self.assertEqual(mock.sample_text_calls, 3) def test_sample_text_does_not_retry_unregistered_exception(self): @@ -250,11 +247,8 @@ def test_sample_choice_raises_after_all_retries_exhausted(self): jitter=(0.0, 0.0), exponential_backoff=False, ) - # Tenacity wraps exhausted retries in RetryError; the original exception is - # accessible via RetryError.__cause__. - with self.assertRaises(tenacity.RetryError) as cm: + with self.assertRaises(ConnectionError): model.sample_choice('prompt', ['a', 'b']) - self.assertIsInstance(cm.exception.__cause__, ConnectionError) self.assertEqual(mock.sample_choice_calls, 3) def test_fixed_wait_no_exponential_backoff(self):