Constrain Ollama choices to supplied response values

concordia/contrib/language_models/ollama/ollama_model.py

@@ -16,14 +16,13 @@
 
 from collections.abc import Collection, Sequence
 import json
-from typing import override
+from typing import override, TypedDict
 
 from concordia.language_model import language_model
 from concordia.utils import measurements as measurements_lib
 from concordia.utils import sampling
 import ollama
 
-
 _MAX_MULTIPLE_CHOICE_ATTEMPTS = 20
 _DEFAULT_TEMPERATURE = 0.5
 _DEFAULT_TERMINATORS = ()
@@ -39,14 +38,24 @@
 )
 
 
+class _ThinkingOptions(TypedDict, total=False):
+  think: bool
+
+
 class OllamaLanguageModel(language_model.LanguageModel):
-  """Language Model that uses Ollama LLM models."""
+  """Language Model that uses Ollama LLM models.
+
+  Explicit sampling seeds are forwarded to the provider. They can help compare
+  repeated prompts on one host/model, but do not guarantee reproducibility
+  across provider versions, hardware, or concurrent inference.
+  """
 
   def __init__(
       self,
       model_name: str,
       *,
       system_message: str = _DEFAULT_SYSTEM_MESSAGE,
+      think: bool | None = None,
       measurements: measurements_lib.Measurements | None = None,
       channel: str = language_model.DEFAULT_STATS_CHANNEL,
   ) -> None:
@@ -57,9 +66,18 @@ def __init__(
           https://github.com/ollama/ollama.
         system_message: System message to prefix to requests when prompting the
           model.
+        think: Optional provider thinking control for supported Ollama models.
+          None preserves the provider default and omits the request field.
+          False can reduce response latency; it can also change answer quality.
+          Reasoning traces are never returned as the sampled answer.
         measurements: The measurements object to log usage statistics to.
         channel: The channel to write the statistics to.
     """
+    if think is not None and type(think) is not bool:
+      raise ValueError('think must be a boolean or None')
+    self._thinking_options: _ThinkingOptions = {}
+    if think is not None:
+      self._thinking_options['think'] = think
     self._model_name = model_name
     self._client = ollama.Client()
     self._system_message = system_message
@@ -81,7 +99,7 @@ def sample_text(
       timeout: float = -1,
       seed: int | None = None,
   ) -> str:
-    del max_tokens, timeout, seed  # Unused.
+    del max_tokens, timeout  # Unused.
 
     prompt_with_system_message = f'{self._system_message}\n\n{prompt}'
 
@@ -95,8 +113,10 @@ def sample_text(
             'temperature': temperature,
             'top_p': top_p,
             'top_k': top_k,
+            **({'seed': seed} if seed is not None else {}),
         },
         keep_alive='10m',
+        **self._thinking_options,
     )
     result = response['response']
 
@@ -115,9 +135,17 @@ def sample_choice(
       *,
       seed: int | None = None,
   ) -> tuple[int, str, dict[str, float]]:
-    del seed  # Unused.
     prompt_with_system_message = f'{self._system_message}\n\n{prompt}'
-    template = {'choice': '', 'single sentence explanation': ''}
+    if not responses:
+      raise ValueError('responses must not be empty')
+    choice_schema = {
+        'type': 'object',
+        'properties': {
+            'choice': {'type': 'string', 'enum': list(responses)},
+        },
+        'required': ['choice'],
+        'additionalProperties': False,
+    }
     sample = ''
     answer = ''
     for attempts in range(_MAX_MULTIPLE_CHOICE_ATTEMPTS):
@@ -130,31 +158,35 @@ def sample_choice(
           model=self._model_name,
           prompt=(
               f'{prompt_with_system_message}.\n'
-              f'Use the following json template: {json.dumps(template)}.'
+              'Return a JSON object with a choice field containing exactly one '
+              f'of these response values: {json.dumps(list(responses))}.'
           ),
-          options={'stop': (), 'temperature': temperature},
-          format='json',
+          options={
+              'stop': (),
+              'temperature': temperature,
+              **({'seed': seed} if seed is not None else {}),
+          },
+          format=choice_schema,
           keep_alive='10m',
+          **self._thinking_options,
       )
       try:
         json_data_response = json.loads(response['response'])
       except json.JSONDecodeError:
         continue
-      sample_or_none = json_data_response.get('choice', None)
-      if sample_or_none is None:
-        if isinstance(json_data_response, dict) and json_data_response:
-          sample = next(iter(json_data_response.values()))
-        elif isinstance(json_data_response, str) and json_data_response:
-          # pyrefly: ignore [missing-attribute]
-          sample = sample_or_none.strip()
-        else:
-          continue
-      else:
-        sample = sample_or_none
-        if isinstance(sample, str) and sample:
-          sample = sample.strip()
-
-      answer = sampling.extract_choice_response(sample)
+      if not isinstance(json_data_response, dict):
+        continue
+      sample = json_data_response.get('choice')
+      if not isinstance(sample, str):
+        continue
+      sample = sample.strip()
+      # Honour the LanguageModel contract for arbitrary response strings; retain
+      # the shared letter extractor for older providers' decorated responses.
+      answer = (
+          sample
+          if sample in responses
+          else sampling.extract_choice_response(sample)
+      )
       try:
         idx = responses.index(answer)
       except ValueError:

concordia/contrib/language_models/ollama/ollama_model_test.py

@@ -0,0 +1,146 @@
+# 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.
+
+"""Choice-contract regression tests; no live Ollama calls."""
+
+import json
+from unittest import mock
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from concordia.contrib.language_models.ollama import ollama_model
+from concordia.language_model import language_model
+
+
+class OllamaChoiceTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    client_patch = mock.patch.object(ollama_model.ollama, 'Client')
+    self.client = client_patch.start().return_value
+    self.addCleanup(client_patch.stop)
+    self.model = ollama_model.OllamaLanguageModel('test-model')
+
+  def test_schema_uses_response_keys_not_labels_in_prompt(self):
+    self.client.generate.return_value = {'response': '{"choice":"b"}'}
+    result = self.model.sample_choice('(a) ACCEPT (b) DECLINE', ['a', 'b'])
+    self.assertEqual(result[:2], (1, 'b'))
+    args = self.client.generate.call_args.kwargs
+    self.assertEqual(args['format']['properties']['choice']['enum'], ['a', 'b'])
+    self.assertEqual(args['format']['required'], ['choice'])
+    self.assertIn('["a", "b"]', args['prompt'])
+    self.client.generate.assert_called_once()
+
+  @parameterized.parameters('ACCEPT', 'one quiet song', 'қабылдау')
+  def test_arbitrary_exact_response_strings(self, value):
+    self.client.generate.return_value = {
+        'response': json.dumps({'choice': value})
+    }
+    self.assertEqual(
+        self.model.sample_choice('choose', ['other', value])[:2], (1, value)
+    )
+
+  @parameterized.parameters('a)', '(a)', 'Answer: (a)')
+  def test_legacy_decorated_letters_use_existing_extractor(self, value):
+    self.client.generate.return_value = {
+        'response': json.dumps({'choice': value})
+    }
+    self.assertEqual(
+        self.model.sample_choice('choose', ['a', 'b'])[:2], (0, 'a')
+    )
+
+  @parameterized.parameters(
+      'null',
+      '[]',
+      '"a"',
+      '{}',
+      '{"choice":1}',
+      '{"choice":[]}',
+      '{"choice":"ACCEPT"}',
+      'not json',
+  )
+  def test_invalid_response_retries_without_guessing(self, invalid):
+    self.client.generate.side_effect = [
+        {'response': invalid},
+        {'response': '{"choice":"b"}'},
+    ]
+    self.assertEqual(
+        self.model.sample_choice('(a) ACCEPT (b) DECLINE', ['a', 'b'])[:2],
+        (1, 'b'),
+    )
+    self.assertEqual(self.client.generate.call_count, 2)
+
+  def test_exhaustion_keeps_failure_instead_of_defaulting(self):
+    self.client.generate.return_value = {'response': '{"choice":"ACCEPT"}'}
+    with self.assertRaises(language_model.InvalidResponseError):
+      self.model.sample_choice('(a) ACCEPT (b) DECLINE', ['a', 'b'])
+    self.assertEqual(self.client.generate.call_count, 20)
+
+  @parameterized.parameters(None, 0, 20260925)
+  def test_sampling_seed_is_forwarded_without_changing_default(self, seed):
+    self.client.generate.return_value = {'response': 'Hello'}
+    self.assertEqual(self.model.sample_text('speak', seed=seed), 'Hello')
+    text_options = self.client.generate.call_args.kwargs['options']
+    self.client.generate.side_effect = [
+        {'response': 'invalid'},
+        {'response': '{"choice":"b"}'},
+    ]
+    self.assertEqual(
+        self.model.sample_choice('choose', ['a', 'b'], seed=seed)[:2],
+        (1, 'b'),
+    )
+    options = [text_options] + [
+        call.kwargs['options']
+        for call in self.client.generate.call_args_list[-2:]
+    ]
+    for option in options:
+      if seed is None:
+        self.assertNotIn('seed', option)
+      else:
+        self.assertEqual(option['seed'], seed)
+
+  @parameterized.parameters(None, False, True)
+  def test_optional_thinking_control_applies_to_text_and_choice(self, think):
+    model = ollama_model.OllamaLanguageModel('test-model', think=think)
+    self.client.generate.return_value = {
+        'response': 'Hello',
+        'thinking': 'private reasoning',
+    }
+    self.assertEqual(model.sample_text('speak'), 'Hello')
+    text_args = self.client.generate.call_args.kwargs
+    self.client.generate.return_value = {
+        'response': '{"choice":"a"}',
+        'thinking': 'private reasoning',
+    }
+    self.assertEqual(model.sample_choice('choose', ['a', 'b'])[:2], (0, 'a'))
+    choice_args = self.client.generate.call_args.kwargs
+    for args in (text_args, choice_args):
+      if think is None:
+        self.assertNotIn('think', args)
+      else:
+        self.assertIs(args['think'], think)
+
+  @parameterized.parameters({'value': 1}, {'value': 'false'}, {'value': []})
+  def test_invalid_thinking_control_is_rejected(self, value):
+    with self.assertRaisesRegex(ValueError, 'think must'):
+      ollama_model.OllamaLanguageModel('test-model', think=value)
+
+  def test_empty_choices_fail_before_model_request(self):
+    with self.assertRaises(ValueError):
+      self.model.sample_choice('choose', [])
+    self.client.generate.assert_not_called()
+
+
+if __name__ == '__main__':
+  absltest.main()