refactor: update Judge to use Runner protocol and RunnerResult

jsonbailey · claude · jsonbailey · commit 45845edb000d · 2026-04-30T08:58:16.000-05:00
- Judge now accepts Runner instead of ModelRunner
- evaluate() calls runner.run(output_type=...) instead of invoke_structured_model
- response.parsed replaces StructuredResponse.data; None guard added
- evaluate_messages() accepts RunnerResult instead of ModelResponse
- Tests updated to use RunnerResult and mock_runner.run

Co-Authored-By: Claude Sonnet 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py
@@ -329,7 +329,7 @@ def create_judge(
             if not provider:
                 return None
 
-            return Judge(judge_config, provider)  # type: ignore[arg-type]
+            return Judge(judge_config, provider)
         except Exception as error:
             return None
 
diff --git a/packages/sdk/server-ai/src/ldai/judge/__init__.py b/packages/sdk/server-ai/src/ldai/judge/__init__.py
@@ -8,8 +8,8 @@
 from ldai import log
 from ldai.judge.evaluation_schema_builder import EvaluationSchemaBuilder
 from ldai.models import AIJudgeConfig, LDMessage
-from ldai.providers.model_runner import ModelRunner
-from ldai.providers.types import JudgeResult, ModelResponse
+from ldai.providers.runner import Runner
+from ldai.providers.types import JudgeResult, RunnerResult
 
 
 class Judge:
@@ -23,7 +23,7 @@ class Judge:
     def __init__(
         self,
         ai_config: AIJudgeConfig,
-        model_runner: ModelRunner,
+        model_runner: Runner,
     ):
         """
         Initialize the Judge.
@@ -76,10 +76,14 @@ async def evaluate(
 
             response = await tracker.track_metrics_of_async(
                 lambda result: result.metrics,
-                lambda: self._model_runner.invoke_structured_model(messages, self._evaluation_response_structure),
+                lambda: self._model_runner.run(messages, output_type=self._evaluation_response_structure),
             )
 
-            parsed = self._parse_evaluation_response(response.data)
+            if response.parsed is None:
+                log.warning('Judge evaluation did not return structured output')
+                return judge_result
+
+            parsed = self._parse_evaluation_response(response.parsed)
 
             if parsed is None:
                 log.warning('Judge evaluation did not return the expected evaluation')
@@ -99,19 +103,19 @@ async def evaluate(
     async def evaluate_messages(
         self,
         messages: list[LDMessage],
-        response: ModelResponse,
+        response: RunnerResult,
         sampling_ratio: float = 1.0,
     ) -> JudgeResult:
         """
         Evaluates an AI response from chat messages and response.
 
         :param messages: Array of messages representing the conversation history
-        :param response: The AI response to be evaluated
+        :param response: The runner result to be evaluated
         :param sampling_ratio: Sampling ratio (0-1) to determine if evaluation should be processed (defaults to 1)
         :return: The result of the judge evaluation.
         """
         input_text = '\r\n'.join([msg.content for msg in messages]) if messages else ''
-        output_text = response.message.content
+        output_text = response.content
 
         return await self.evaluate(input_text, output_text, sampling_ratio)
 
@@ -123,7 +127,7 @@ def get_ai_config(self) -> AIJudgeConfig:
         """
         return self._ai_config
 
-    def get_model_runner(self) -> ModelRunner:
+    def get_model_runner(self) -> Runner:
         """
         Returns the model runner used by this judge.
 
diff --git a/packages/sdk/server-ai/tests/test_judge.py b/packages/sdk/server-ai/tests/test_judge.py
@@ -9,7 +9,7 @@
 from ldai.judge import Judge
 from ldai.judge.evaluation_schema_builder import EvaluationSchemaBuilder
 from ldai.models import AIJudgeConfig, AIJudgeConfigDefault, LDMessage, ModelConfig, ProviderConfig
-from ldai.providers.types import JudgeResult, LDAIMetrics, StructuredResponse
+from ldai.providers.types import JudgeResult, LDAIMetrics, RunnerResult
 from ldai.tracker import LDAIConfigTracker
 
 
@@ -40,9 +40,9 @@ def client(td: TestData) -> LDClient:
 
 @pytest.fixture
 def mock_runner():
-    """Create a mock AI provider."""
+    """Create a mock AI runner."""
     provider = MagicMock()
-    provider.invoke_structured_model = AsyncMock()
+    provider.run = AsyncMock()
     return provider
 
 
@@ -137,7 +137,7 @@ async def test_evaluate_returns_failure_when_evaluation_metric_key_missing(
         assert isinstance(result, JudgeResult)
         assert result.success is False
         assert result.sampled is False
-        mock_runner.invoke_structured_model.assert_not_called()
+        mock_runner.run.assert_not_called()
 
     @pytest.mark.asyncio
     async def test_evaluate_returns_failure_when_messages_missing(
@@ -151,23 +151,23 @@ async def test_evaluate_returns_failure_when_messages_missing(
         assert isinstance(result, JudgeResult)
         assert result.success is False
         assert result.sampled is False
-        mock_runner.invoke_structured_model.assert_not_called()
+        mock_runner.run.assert_not_called()
 
     @pytest.mark.asyncio
     async def test_evaluate_success_with_valid_response(
         self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner
     ):
         """Evaluate should return JudgeResponse with valid evaluation."""
-        mock_response = StructuredResponse(
-            data={
+        mock_response = RunnerResult(
+            content='',
+            metrics=LDAIMetrics(success=True),
+            parsed={
                 'score': 0.85,
                 'reasoning': 'The response is highly relevant to the input.'
             },
-            raw_response='{"score": 0.85, "reasoning": "..."}',
-            metrics=LDAIMetrics(success=True)
         )
 
-        mock_runner.invoke_structured_model.return_value = mock_response
+        mock_runner.run.return_value = mock_response
         tracker.track_metrics_of_async = AsyncMock(return_value=mock_response)
 
         judge = Judge(judge_config_with_key, mock_runner)
@@ -187,15 +187,15 @@ async def test_evaluate_success_with_evaluation_response_shape(
         self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner
     ):
         """Evaluate should accept shape { score, reasoning } and key by metric."""
-        mock_response = StructuredResponse(
-            data={
+        mock_response = RunnerResult(
+            content='',
+            metrics=LDAIMetrics(success=True),
+            parsed={
                 'score': 0.9,
                 'reasoning': 'The response is accurate and complete.',
             },
-            raw_response='{"score": 0.9, "reasoning": "..."}',
-            metrics=LDAIMetrics(success=True),
         )
-        mock_runner.invoke_structured_model.return_value = mock_response
+        mock_runner.run.return_value = mock_response
         tracker.track_metrics_of_async = AsyncMock(return_value=mock_response)
 
         judge = Judge(judge_config_with_key, mock_runner)
@@ -214,13 +214,13 @@ async def test_evaluate_handles_missing_evaluation_in_response(
         self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner
     ):
         """Evaluate should handle missing score/reasoning in response."""
-        mock_response = StructuredResponse(
-            data={},
-            raw_response='{}',
-            metrics=LDAIMetrics(success=True)
+        mock_response = RunnerResult(
+            content='',
+            metrics=LDAIMetrics(success=True),
+            parsed={},
         )
 
-        mock_runner.invoke_structured_model.return_value = mock_response
+        mock_runner.run.return_value = mock_response
         tracker.track_metrics_of_async = AsyncMock(return_value=mock_response)
 
         judge = Judge(judge_config_with_key, mock_runner)
@@ -236,16 +236,16 @@ async def test_evaluate_handles_invalid_score(
         self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner
     ):
         """Evaluate should handle invalid score values."""
-        mock_response = StructuredResponse(
-            data={
+        mock_response = RunnerResult(
+            content='',
+            metrics=LDAIMetrics(success=True),
+            parsed={
                 'score': 1.5,
-                'reasoning': 'Some reasoning'
+                'reasoning': 'Some reasoning',
             },
-            raw_response='{"score": 1.5, "reasoning": "..."}',
-            metrics=LDAIMetrics(success=True)
         )
 
-        mock_runner.invoke_structured_model.return_value = mock_response
+        mock_runner.run.return_value = mock_response
         tracker.track_metrics_of_async = AsyncMock(return_value=mock_response)
 
         judge = Judge(judge_config_with_key, mock_runner)
@@ -261,13 +261,13 @@ async def test_evaluate_handles_missing_reasoning(
         self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner
     ):
         """Evaluate should handle missing reasoning."""
-        mock_response = StructuredResponse(
-            data={'score': 0.8},
-            raw_response='{"score": 0.8}',
-            metrics=LDAIMetrics(success=True)
+        mock_response = RunnerResult(
+            content='',
+            metrics=LDAIMetrics(success=True),
+            parsed={'score': 0.8},
         )
 
-        mock_runner.invoke_structured_model.return_value = mock_response
+        mock_runner.run.return_value = mock_response
         tracker.track_metrics_of_async = AsyncMock(return_value=mock_response)
 
         judge = Judge(judge_config_with_key, mock_runner)
@@ -283,7 +283,7 @@ async def test_evaluate_handles_exception(
         self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner
     ):
         """Evaluate should handle exceptions gracefully."""
-        mock_runner.invoke_structured_model.side_effect = Exception("Provider error")
+        mock_runner.run.side_effect = Exception("Provider error")
         tracker.track_metrics_of_async = AsyncMock(side_effect=Exception("Provider error"))
 
         judge = Judge(judge_config_with_key, mock_runner)
@@ -306,7 +306,7 @@ async def test_evaluate_respects_sampling_rate(
         assert isinstance(result, JudgeResult)
         assert result.sampled is False
         assert result.success is False
-        mock_runner.invoke_structured_model.assert_not_called()
+        mock_runner.run.assert_not_called()
 
 
 class TestJudgeEvaluateMessages:
@@ -317,15 +317,13 @@ async def test_evaluate_messages_calls_evaluate(
         self, judge_config_with_key: AIJudgeConfig, tracker: LDAIConfigTracker, mock_runner
     ):
         """evaluate_messages should call evaluate with constructed input/output."""
-        from ldai.providers.types import ModelResponse
-
-        mock_response = StructuredResponse(
-            data={'score': 0.9, 'reasoning': 'Very relevant'},
-            raw_response='{"score": 0.9, "reasoning": "..."}',
-            metrics=LDAIMetrics(success=True)
+        mock_response = RunnerResult(
+            content='',
+            metrics=LDAIMetrics(success=True),
+            parsed={'score': 0.9, 'reasoning': 'Very relevant'},
         )
 
-        mock_runner.invoke_structured_model.return_value = mock_response
+        mock_runner.run.return_value = mock_response
         tracker.track_metrics_of_async = AsyncMock(return_value=mock_response)
 
         judge = Judge(judge_config_with_key, mock_runner)
@@ -334,9 +332,9 @@ async def test_evaluate_messages_calls_evaluate(
             LDMessage(role='user', content='Question 1'),
             LDMessage(role='assistant', content='Answer 1'),
         ]
-        chat_response = ModelResponse(
-            message=LDMessage(role='assistant', content='Answer 2'),
-            metrics=LDAIMetrics(success=True)
+        chat_response = RunnerResult(
+            content='Answer 2',
+            metrics=LDAIMetrics(success=True),
         )
 
         result = await judge.evaluate_messages(messages, chat_response)
