feat: add support for custom judges via evaluation metric key (#86)

**Requirements**

- [X] I have added test coverage for new or changed functionality
- [x] I have followed the repository's [pull request submission
guidelines](../blob/main/CONTRIBUTING.md#submitting-pull-requests)
- [X] I have validated my changes against all supported platform
versions

**Related issues**

https://launchdarkly.atlassian.net/browse/REL-11511
See tech spec at
https://docs.google.com/document/d/1lzYwQqCcTzN_2zkxJZDfJtgUcEJ4jbpx0KSsJ2bRENw/edit?tab=t.0#heading=h.69bdm7karsxh

**Describe the solution you've provided**

Updating the SDK to check the AI Config's evaluationMetricKey property
which now exists. Also added missing tests from previous implementation,
and fallback to the original evaluationMetricKeys list.

**Describe alternatives you've considered**

Provide a clear and concise description of any alternative solutions or
features you've considered.

**Additional context**

Add any other context about the pull request here.

<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> Implements single-key judge evaluation with backward compatibility and
comprehensive tests.
> 
> - Switches judge configs to use `evaluationMetricKey` (deprecated
`evaluationMetricKeys`), updating `AIJudgeConfig(Default)` serialization
> - `LDAIClient.__evaluate` now returns the raw `variation`;
`judge_config` extracts `evaluationMetricKey` with fallback to first in
`evaluationMetricKeys`
> - `Judge` updated to validate and parse a single metric;
`EvaluationSchemaBuilder` builds a single-key structured schema; minor
cleanup of unused imports/comments
> - Adds extensive unit tests for judge behavior, schema building, and
client extraction (including consistency of single variation, sampling,
error paths)
> 
> <sup>Written by [Cursor
Bugbot](https://cursor.com/dashboard?tab=bugbot) for commit
c6d086a. This will update automatically
on new commits. Configure
[here](https://cursor.com/dashboard?tab=bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

Author: knfreemLD

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

@@ -42,7 +42,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
         )
 
@@ -96,18 +96,31 @@ 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
         )
 
-        # Extract evaluation_metric_keys from the variation
-        variation = self._client.variation(key, context, default_value.to_dict())
-        evaluation_metric_keys = variation.get('evaluationMetricKeys', default_value.evaluation_metric_keys or [])
+        def _extract_evaluation_metric_key(variation: Dict[str, Any]) -> Optional[str]:
+            """
+            Extract evaluation_metric_key with backward compatibility.
+
+            Priority: 1) evaluationMetricKey from variation, 2) first from evaluationMetricKeys in variation
+            """
+            if evaluation_metric_key := variation.get('evaluationMetricKey'):
+                return evaluation_metric_key
+
+            variation_keys = variation.get('evaluationMetricKeys')
+            if isinstance(variation_keys, list) and variation_keys:
+                return variation_keys[0]
+
+            return None
+
+        evaluation_metric_key = _extract_evaluation_metric_key(variation)
 
         config = AIJudgeConfig(
             key=key,
             enabled=bool(enabled),
-            evaluation_metric_keys=evaluation_metric_keys,
+            evaluation_metric_key=evaluation_metric_key,
             model=model,
             messages=messages,
             provider=provider,
@@ -144,7 +157,7 @@ async def create_judge(
                     enabled=True,
                     model=ModelConfig("gpt-4"),
                     provider=ProviderConfig("openai"),
-                    evaluation_metric_keys=['$ld:ai:judge:relevance'],
+                    evaluation_metric_key='$ld:ai:judge:relevance',
                     messages=[LDMessage(role='system', content='You are a relevance judge.')]
                 ),
                 variables={'metric': "relevance"}
@@ -160,33 +173,27 @@ async def create_judge(
         self._client.track('$ld:ai:judge:function:createJudge', context, key, 1)
 
         try:
-            # Warn if reserved variables are provided
             if variables:
                 if 'message_history' in variables:
-                    # Note: Python doesn't have a logger on the client, but we could add one
-                    pass  # Would log warning if logger available
+                    pass
                 if 'response_to_evaluate' in variables:
-                    pass  # Would log warning if logger available
+                    pass
 
-            # Overwrite reserved variables to ensure they remain as placeholders for judge evaluation
             extended_variables = dict(variables) if variables else {}
             extended_variables['message_history'] = '{{message_history}}'
             extended_variables['response_to_evaluate'] = '{{response_to_evaluate}}'
 
             judge_config = self.judge_config(key, context, default_value, extended_variables)
 
             if not judge_config.enabled or not judge_config.tracker:
-                # Would log info if logger available
                 return None
 
-            # Create AI provider for the judge
             provider = await AIProviderFactory.create(judge_config, default_ai_provider)
             if not provider:
                 return None
 
             return Judge(judge_config, judge_config.tracker, provider)
         except Exception as error:
-            # Would log error if logger available
             return None
 
     async def _initialize_judges(
@@ -279,7 +286,6 @@ async def create_chat(
         config = self.completion_config(key, context, default_value, variables)
 
         if not config.enabled or not config.tracker:
-            # Would log info if logger available
             return None
 
         provider = await AIProviderFactory.create(config, default_ai_provider)
@@ -333,7 +339,6 @@ def agent_config(
         :param variables: Variables for interpolation.
         :return: Configured AIAgentConfig instance.
         """
-        # Track single agent usage
         self._client.track(
             "$ld:ai:agent:function:single",
             context,
@@ -399,7 +404,6 @@ def agent_configs(
         :param context: The context to evaluate the agent configurations in.
         :return: Dictionary mapping agent keys to their AIAgentConfig configurations.
         """
-        # Track multiple agents usage
         agent_count = len(agent_configs)
         self._client.track(
             "$ld:ai:agent:function:multiple",
@@ -538,7 +542,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.
@@ -547,7 +551,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)
 
@@ -556,7 +560,6 @@ def __evaluate(
             all_variables.update(variables)
         all_variables['ldctx'] = context.to_dict()
 
-        # Extract messages
         messages = None
         if 'messages' in variation and isinstance(variation['messages'], list) and all(
             isinstance(entry, dict) for entry in variation['messages']
@@ -571,18 +574,15 @@ def __evaluate(
                 for entry in variation['messages']
             ]
 
-        # Extract instructions
         instructions = None
         if 'instructions' in variation and isinstance(variation['instructions'], str):
             instructions = self.__interpolate_template(variation['instructions'], all_variables)
 
-        # Extract provider config
         provider_config = None
         if 'provider' in variation and isinstance(variation['provider'], dict):
             provider = variation['provider']
             provider_config = ProviderConfig(provider.get('name', ''))
 
-        # Extract model config
         model = None
         if 'model' in variation and isinstance(variation['model'], dict):
             parameters = variation['model'].get('parameters', None)
@@ -593,7 +593,6 @@ def __evaluate(
                 custom=custom
             )
 
-        # Create tracker
         tracker = LDAIConfigTracker(
             self._client,
             variation.get('_ldMeta', {}).get('variationKey', ''),
@@ -606,7 +605,6 @@ def __evaluate(
 
         enabled = variation.get('_ldMeta', {}).get('enabled', False)
 
-        # Extract judge configuration
         judge_configuration = None
         if 'judgeConfiguration' in variation and isinstance(variation['judgeConfiguration'], dict):
             judge_config = variation['judgeConfiguration']
@@ -622,7 +620,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,
@@ -640,7 +638,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/src/ldai/judge/__init__.py

@@ -9,8 +9,7 @@
 from ldai.judge.evaluation_schema_builder import EvaluationSchemaBuilder
 from ldai.models import AIJudgeConfig, LDMessage
 from ldai.providers.ai_provider import AIProvider
-from ldai.providers.types import (ChatResponse, EvalScore, JudgeResponse,
-                                  StructuredResponse)
+from ldai.providers.types import ChatResponse, EvalScore, JudgeResponse
 from ldai.tracker import LDAIConfigTracker
 
 
@@ -38,9 +37,7 @@ def __init__(
         self._ai_config = ai_config
         self._ai_config_tracker = ai_config_tracker
         self._ai_provider = ai_provider
-        self._evaluation_response_structure = EvaluationSchemaBuilder.build(
-            ai_config.evaluation_metric_keys
-        )
+        self._evaluation_response_structure = EvaluationSchemaBuilder.build(ai_config.evaluation_metric_key)
 
     async def evaluate(
         self,
@@ -57,9 +54,9 @@ async def evaluate(
         :return: Evaluation results or None if not sampled
         """
         try:
-            if not self._ai_config.evaluation_metric_keys or len(self._ai_config.evaluation_metric_keys) == 0:
+            if not self._ai_config.evaluation_metric_key:
                 log.warn(
-                    'Judge configuration is missing required evaluationMetricKeys'
+                    'Judge configuration is missing required evaluationMetricKey'
                 )
                 return None
 
@@ -72,8 +69,8 @@ async def evaluate(
                 return None
 
             messages = self._construct_evaluation_messages(input_text, output_text)
+            assert self._evaluation_response_structure is not None
 
-            # Track metrics of the structured model invocation
             response = await self._ai_config_tracker.track_metrics_of(
                 lambda: self._ai_provider.invoke_structured_model(messages, self._evaluation_response_structure),
                 lambda result: result.metrics,
@@ -83,8 +80,8 @@ async def evaluate(
 
             evals = self._parse_evaluation_response(response.data)
 
-            if len(evals) != len(self._ai_config.evaluation_metric_keys):
-                log.warn('Judge evaluation did not return all evaluations')
+            if self._ai_config.evaluation_metric_key not in evals:
+                log.warn('Judge evaluation did not return the expected evaluation')
                 success = False
 
             return JudgeResponse(
@@ -191,30 +188,34 @@ def _parse_evaluation_response(self, data: Dict[str, Any]) -> Dict[str, EvalScor
 
         evaluations = data['evaluations']
 
-        for metric_key in self._ai_config.evaluation_metric_keys:
-            evaluation = evaluations.get(metric_key)
+        metric_key = self._ai_config.evaluation_metric_key
+        if not metric_key:
+            log.warn('Evaluation metric key is missing')
+            return results
 
-            if not evaluation or not isinstance(evaluation, dict):
-                log.warn(f'Missing evaluation for metric key: {metric_key}')
-                continue
+        evaluation = evaluations.get(metric_key)
 
-            score = evaluation.get('score')
-            reasoning = evaluation.get('reasoning')
+        if not evaluation or not isinstance(evaluation, dict):
+            log.warn(f'Missing evaluation for metric key: {metric_key}')
+            return results
 
-            if not isinstance(score, (int, float)) or score < 0 or score > 1:
-                log.warn(
-                    f'Invalid score evaluated for {metric_key}: {score}. '
-                    'Score must be a number between 0 and 1 inclusive'
-                )
-                continue
+        score = evaluation.get('score')
+        reasoning = evaluation.get('reasoning')
 
-            if not isinstance(reasoning, str):
-                log.warn(
-                    f'Invalid reasoning evaluated for {metric_key}: {reasoning}. '
-                    'Reasoning must be a string'
-                )
-                continue
+        if not isinstance(score, (int, float)) or score < 0 or score > 1:
+            log.warn(
+                f'Invalid score evaluated for {metric_key}: {score}. '
+                'Score must be a number between 0 and 1 inclusive'
+            )
+            return results
+
+        if not isinstance(reasoning, str):
+            log.warn(
+                f'Invalid reasoning evaluated for {metric_key}: {reasoning}. '
+                'Reasoning must be a string'
+            )
+            return results
 
-            results[metric_key] = EvalScore(score=float(score), reasoning=reasoning)
+        results[metric_key] = EvalScore(score=float(score), reasoning=reasoning)
 
         return results

packages/sdk/server-ai/src/ldai/judge/evaluation_schema_builder.py

@@ -1,6 +1,6 @@
 """Internal class for building dynamic evaluation response schemas."""
 
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 
 class EvaluationSchemaBuilder:
@@ -10,26 +10,29 @@ class EvaluationSchemaBuilder:
     """
 
     @staticmethod
-    def build(evaluation_metric_keys: list[str]) -> Dict[str, Any]:
+    def build(evaluation_metric_key: Optional[str]) -> Optional[Dict[str, Any]]:
         """
-        Build an evaluation response schema from evaluation metric keys.
+        Build an evaluation response schema from evaluation metric key.
 
-        :param evaluation_metric_keys: List of evaluation metric keys
-        :return: Schema dictionary for structured output
+        :param evaluation_metric_key: Evaluation metric key, or None if not available
+        :return: Schema dictionary for structured output, or None if evaluation_metric_key is None
         """
+        if not evaluation_metric_key:
+            return None
+
         return {
             'title': 'EvaluationResponse',
-            'description': f"Response containing evaluation results for {', '.join(evaluation_metric_keys)} metrics",
+            'description': f"Response containing evaluation results for {evaluation_metric_key} metric",
             'type': 'object',
             'properties': {
                 'evaluations': {
                     'type': 'object',
                     'description': (
                         f"Object containing evaluation results for "
-                        f"{', '.join(evaluation_metric_keys)} metrics"
+                        f"{evaluation_metric_key} metric"
                     ),
-                    'properties': EvaluationSchemaBuilder._build_key_properties(evaluation_metric_keys),
-                    'required': evaluation_metric_keys,
+                    'properties': EvaluationSchemaBuilder._build_key_properties(evaluation_metric_key),
+                    'required': [evaluation_metric_key],
                     'additionalProperties': False,
                 },
             },
@@ -38,17 +41,16 @@ def build(evaluation_metric_keys: list[str]) -> Dict[str, Any]:
         }
 
     @staticmethod
-    def _build_key_properties(evaluation_metric_keys: list[str]) -> Dict[str, Any]:
+    def _build_key_properties(evaluation_metric_key: str) -> Dict[str, Any]:
         """
-        Build properties for each evaluation metric key.
+        Build properties for a single evaluation metric key.
 
-        :param evaluation_metric_keys: List of evaluation metric keys
-        :return: Dictionary of properties for each key
+        :param evaluation_metric_key: Evaluation metric key
+        :return: Dictionary of properties for the key
         """
-        result: Dict[str, Any] = {}
-        for key in evaluation_metric_keys:
-            result[key] = EvaluationSchemaBuilder._build_key_schema(key)
-        return result
+        return {
+            evaluation_metric_key: EvaluationSchemaBuilder._build_key_schema(evaluation_metric_key)
+        }
 
     @staticmethod
     def _build_key_schema(key: str) -> Dict[str, Any]: