Addressed PR feedback; fixed race condition

Author: knfreemLD

packages/sdk/server-ai/src/ldai/client.py

@@ -40,7 +40,7 @@ def completion_config(
         """
         self._client.track('$ld:ai:config:function:single', context, key, 1)
 
-        model, provider, messages, instructions, tracker, enabled, judge_configuration = self.__evaluate(
+        model, provider, messages, instructions, tracker, enabled, judge_configuration, _ = self.__evaluate(
             key, context, default_value.to_dict(), variables
         )
 
@@ -94,31 +94,29 @@ def judge_config(
         """
         self._client.track('$ld:ai:judge:function:single', context, key, 1)
 
-        model, provider, messages, instructions, tracker, enabled, judge_configuration = self.__evaluate(
+        model, provider, messages, instructions, tracker, enabled, judge_configuration, variation = self.__evaluate(
             key, context, default_value.to_dict(), variables
         )
 
-        variation = self._client.variation(key, context, default_value.to_dict())
-
         def _extract_evaluation_metric_key(
             variation: Dict[str, Any], default_value: AIJudgeConfigDefault
         ) -> Optional[str]:
             """
             Extract evaluation_metric_key with backward compatibility.
 
-            Priority: 1) evaluationMetricKey from variation, 2) evaluation_metric_key from default,
-                      3) first from evaluationMetricKeys in variation, 4) first from evaluation_metric_keys in default
+            Priority: 1) evaluationMetricKey from variation, 2) evaluationMetricKeys from variation,
+                      3) evaluation_metric_key from default, 4) evaluation_metric_keys from default
             """
             if evaluation_metric_key := variation.get('evaluationMetricKey'):
                 return evaluation_metric_key
 
-            if default_value.evaluation_metric_key:
-                return default_value.evaluation_metric_key
-
             variation_keys = variation.get('evaluationMetricKeys')
             if isinstance(variation_keys, list) and variation_keys:
                 return variation_keys[0]
 
+            if default_value.evaluation_metric_key:
+                return default_value.evaluation_metric_key
+
             if default_value.evaluation_metric_keys:
                 return default_value.evaluation_metric_keys[0]
 
@@ -458,7 +456,7 @@ def __evaluate(
         variables: Optional[Dict[str, Any]] = None,
     ) -> Tuple[
         Optional[ModelConfig], Optional[ProviderConfig], Optional[List[LDMessage]],
-        Optional[str], LDAIConfigTracker, bool, Optional[Any]
+        Optional[str], LDAIConfigTracker, bool, Optional[Any], Dict[str, Any]
     ]:
         """
         Internal method to evaluate a configuration and extract components.
@@ -467,7 +465,7 @@ def __evaluate(
         :param context: The evaluation context.
         :param default_dict: Default configuration as dictionary.
         :param variables: Variables for interpolation.
-        :return: Tuple of (model, provider, messages, instructions, tracker, enabled).
+        :return: Tuple of (model, provider, messages, instructions, tracker, enabled, judge_configuration, variation).
         """
         variation = self._client.variation(key, context, default_dict)
 
@@ -536,7 +534,7 @@ def __evaluate(
                 if judges:
                     judge_configuration = JudgeConfiguration(judges=judges)
 
-        return model, provider_config, messages, instructions, tracker, enabled, judge_configuration
+        return model, provider_config, messages, instructions, tracker, enabled, judge_configuration, variation
 
     def __evaluate_agent(
         self,
@@ -554,7 +552,7 @@ def __evaluate_agent(
         :param variables: Variables for interpolation.
         :return: Configured AIAgentConfig instance.
         """
-        model, provider, messages, instructions, tracker, enabled, judge_configuration = self.__evaluate(
+        model, provider, messages, instructions, tracker, enabled, judge_configuration, _ = self.__evaluate(
             key, context, default_value.to_dict(), variables
         )
 

packages/sdk/server-ai/tests/test_judge.py

@@ -595,3 +595,53 @@ def test_judge_config_prefers_evaluation_metric_key_over_keys(
 
         assert config is not None
         assert config.evaluation_metric_key == '$ld:ai:judge:preferred'
+
+    def test_judge_config_uses_same_variation_for_consistency(
+        self, context: Context
+    ):
+        """judge_config should use the same variation from __evaluate to avoid race conditions."""
+        from ldai import LDAIClient
+        from ldclient import Config, LDClient
+        from ldclient.integrations.test_data import TestData
+        from unittest.mock import patch
+
+        td = TestData.data_source()
+        td.update(
+            td.flag('judge-consistency-test')
+            .variations(
+                {
+                    'model': {'name': 'gpt-4'},
+                    'provider': {'name': 'openai'},
+                    'messages': [{'role': 'system', 'content': 'You are a judge.'}],
+                    'evaluationMetricKey': '$ld:ai:judge:from-flag',
+                    '_ldMeta': {'enabled': True, 'variationKey': 'judge-v1', 'version': 1},
+                }
+            )
+            .variation_for_all(0)
+        )
+
+        test_client = LDClient(Config('sdk-key', update_processor_class=td, send_events=False))
+        ldai_client = LDAIClient(test_client)
+
+        default_value = AIJudgeConfigDefault(
+            enabled=True,
+            evaluation_metric_key='$ld:ai:judge:from-default',
+            messages=[LDMessage(role='system', content='You are a judge.')],
+            model=ModelConfig('gpt-4'),
+            provider=ProviderConfig('openai'),
+        )
+
+        variation_calls = []
+        original_variation = test_client.variation
+
+        def tracked_variation(key, context, default):
+            result = original_variation(key, context, default)
+            variation_calls.append((key, result.get('evaluationMetricKey')))
+            return result
+
+        with patch.object(test_client, 'variation', side_effect=tracked_variation):
+            config = ldai_client.judge_config('judge-consistency-test', context, default_value)
+
+        assert len(variation_calls) == 1, f"Expected 1 variation call, got {len(variation_calls)}"
+        assert config is not None
+        assert config.evaluation_metric_key == '$ld:ai:judge:from-flag'