refactor: address review feedback and remove deprecated result types

Author: jsonbailey

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

@@ -34,7 +34,6 @@
 from ldai.providers import (
     AgentGraphResult,
     AgentGraphRunner,
-    AgentResult,
     AgentRunner,
     ManagedResult,
     Runner,
@@ -49,7 +48,6 @@
     'Evaluator',
     'AgentRunner',
     'AgentGraphRunner',
-    'AgentResult',
     'AgentGraphResult',
     'ManagedResult',
     'Runner',

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

@@ -308,7 +308,7 @@ def create_judge(
             if not provider:
                 return None
 
-            return Judge(judge_config, provider)
+            return Judge(judge_config, provider)  # type: ignore[arg-type]
         except Exception as error:
             return None
 

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,13 +23,13 @@ class Judge:
     def __init__(
         self,
         ai_config: AIJudgeConfig,
-        model_runner: ModelRunner,
+        model_runner: Runner,
     ):
         """
         Initialize the Judge.
 
         :param ai_config: The judge AI configuration
-        :param model_runner: The model runner to use for evaluation
+        :param model_runner: The runner to use for evaluation
         """
         self._ai_config = ai_config
         self._model_runner = model_runner
@@ -76,10 +76,10 @@ 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)
+            parsed = self._parse_evaluation_response(response.parsed)
 
             if parsed is None:
                 log.warning('Judge evaluation did not return the expected evaluation')
@@ -99,7 +99,7 @@ async def evaluate(
     async def evaluate_messages(
         self,
         messages: list[LDMessage],
-        response: ModelResponse,
+        response: RunnerResult,
         sampling_ratio: float = 1.0,
     ) -> JudgeResult:
         """
@@ -111,7 +111,7 @@ async def evaluate_messages(
         :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,11 +123,11 @@ 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.
+        Returns the runner used by this judge.
 
-        :return: The model runner
+        :return: The runner
         """
         return self._model_runner
 
@@ -164,7 +164,7 @@ def _interpolate_message(self, content: str, variables: Dict[str, str]) -> str:
         # Use chevron (Mustache) for templating, with no escaping
         return chevron.render(content, variables)
 
-    def _parse_evaluation_response(self, data: Dict[str, Any]) -> Optional[Tuple[float, str]]:
+    def _parse_evaluation_response(self, data: Optional[Dict[str, Any]]) -> Optional[Tuple[float, str]]:
         """
         Parses the structured evaluation response. Expects {"score": n, "reasoning": "..."}.
 

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

@@ -3,7 +3,7 @@
 from typing import Union
 
 from ldai.models import AIAgentConfig
-from ldai.providers import AgentResult, AgentRunner
+from ldai.providers import AgentRunner
 from ldai.providers.runner import Runner
 from ldai.providers.types import ManagedResult, RunnerResult
 
@@ -33,14 +33,12 @@ async def run(self, input: str) -> ManagedResult:
         :return: ManagedResult containing the agent's output and metric summary
         """
         tracker = self._ai_config.create_tracker()
-        result: Union[RunnerResult, AgentResult] = await tracker.track_metrics_of_async(
+        result: RunnerResult = await tracker.track_metrics_of_async(
             lambda r: r.metrics,
             lambda: self._agent_runner.run(input),
         )
-        # Support both RunnerResult (content) and legacy AgentResult (output)
-        content = result.content if isinstance(result, RunnerResult) else result.output  # type: ignore[union-attr]
         return ManagedResult(
-            content=content,
+            content=result.content,
             metrics=tracker.get_summary(),
             raw=result.raw,
         )

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

@@ -1,12 +1,11 @@
 import asyncio
-import warnings
 from typing import List, Union
 
 from ldai import log
 from ldai.models import AICompletionConfig, LDMessage
 from ldai.providers.model_runner import ModelRunner
 from ldai.providers.runner import Runner
-from ldai.providers.types import JudgeResult, ManagedResult, ModelResponse, RunnerResult
+from ldai.providers.types import JudgeResult, ManagedResult, RunnerResult
 from ldai.tracker import LDAIConfigTracker
 
 
@@ -48,86 +47,38 @@ async def run(self, prompt: str) -> ManagedResult:
         config_messages = self._ai_config.messages or []
         all_messages = config_messages + self._messages
 
-        result: Union[RunnerResult, ModelResponse] = await tracker.track_metrics_of_async(
+        result: RunnerResult = await tracker.track_metrics_of_async(
             lambda r: r.metrics,
             lambda: self._invoke_runner(all_messages),
         )
 
-        # Support both new RunnerResult and legacy ModelResponse
-        if isinstance(result, RunnerResult):
-            content = result.content
-            raw = result.raw
-            parsed = result.parsed
-            assistant_message = LDMessage(role='assistant', content=content)
-        else:
-            content = result.message.content
-            raw = getattr(result, 'raw', None)
-            parsed = getattr(result, 'parsed', None)
-            assistant_message = result.message
+        assistant_message = LDMessage(role='assistant', content=result.content)
 
         input_text = '\r\n'.join(m.content for m in self._messages) if self._messages else ''
 
-        evaluations_task = self._track_judge_results(tracker, input_text, content)
+        evaluations_task = self._track_judge_results(tracker, input_text, result.content)
 
         self._messages.append(assistant_message)
 
         return ManagedResult(
-            content=content,
+            content=result.content,
             metrics=tracker.get_summary(),
-            raw=raw,
-            parsed=parsed,
+            raw=result.raw,
+            parsed=result.parsed,
             evaluations=evaluations_task,
         )
 
-    async def _invoke_runner(
-        self, all_messages: List[LDMessage]
-    ) -> Union[RunnerResult, ModelResponse]:
+    async def _invoke_runner(self, all_messages: List[LDMessage]) -> RunnerResult:
         """
         Delegate to the runner.  Supports both the new ``Runner`` protocol
         (``run(messages) → RunnerResult``) and the legacy ``ModelRunner``
-        (``invoke_model(messages) → ModelResponse``).
+        (``invoke_model(messages) → RunnerResult``).
         """
         if isinstance(self._model_runner, Runner):
             return await self._model_runner.run(all_messages)
         # Legacy ModelRunner path
         return await self._model_runner.invoke_model(all_messages)  # type: ignore[union-attr]
 
-    async def invoke(self, prompt: str) -> ModelResponse:
-        """
-        Invoke the model with a prompt string.
-
-        .. deprecated::
-            Use :meth:`run` instead. This method will be removed in a future
-            release once the migration to :class:`ManagedResult` is complete.
-
-        :param prompt: The user prompt to send to the model
-        :return: ModelResponse containing the model's response and metrics
-        """
-        warnings.warn(
-            "ManagedModel.invoke() is deprecated. Use run() instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
-        tracker = self._ai_config.create_tracker()
-
-        user_message = LDMessage(role='user', content=prompt)
-        self._messages.append(user_message)
-
-        config_messages = self._ai_config.messages or []
-        all_messages = config_messages + self._messages
-
-        response: ModelResponse = await tracker.track_metrics_of_async(
-            lambda result: result.metrics,
-            lambda: self._model_runner.invoke_model(all_messages),  # type: ignore[union-attr]
-        )
-
-        input_text = '\r\n'.join(m.content for m in self._messages) if self._messages else ''
-        output_text = response.message.content
-        response.evaluations = self._track_judge_results(tracker, input_text, output_text)
-
-        self._messages.append(response.message)
-        return response
-
     def _track_judge_results(
         self,
         tracker: LDAIConfigTracker,

packages/sdk/server-ai/src/ldai/providers/__init__.py

@@ -6,30 +6,24 @@
 from ldai.providers.runner_factory import RunnerFactory
 from ldai.providers.types import (
     AgentGraphResult,
-    AgentResult,
     JudgeResult,
     LDAIMetrics,
     ManagedResult,
-    ModelResponse,
     RunnerResult,
-    StructuredResponse,
     ToolRegistry,
 )
 
 __all__ = [
     'AIProvider',
     'AgentGraphResult',
     'AgentGraphRunner',
-    'AgentResult',
     'AgentRunner',
     'JudgeResult',
     'LDAIMetrics',
     'ManagedResult',
-    'ModelResponse',
     'ModelRunner',
     'Runner',
     'RunnerFactory',
     'RunnerResult',
-    'StructuredResponse',
     'ToolRegistry',
 ]

packages/sdk/server-ai/src/ldai/providers/agent_runner.py

@@ -1,6 +1,6 @@
 from typing import Any, Protocol, runtime_checkable
 
-from ldai.providers.types import AgentResult
+from ldai.providers.types import RunnerResult
 
 
 @runtime_checkable
@@ -18,11 +18,11 @@ class AgentRunner(Protocol):
     the caller just passes input.
     """
 
-    async def run(self, input: Any) -> AgentResult:
+    async def run(self, input: Any) -> RunnerResult:
         """
         Run the agent with the given input.
 
         :param input: The input to the agent (string prompt or structured input)
-        :return: AgentResult containing the output, raw response, and metrics
+        :return: RunnerResult containing the agent's content, metrics, and optional raw/parsed fields
         """
         ...

packages/sdk/server-ai/src/ldai/providers/ai_provider.py

@@ -1,9 +1,9 @@
 from abc import ABC
-from typing import Any, Dict, List, Optional
+from typing import Any, List, Optional
 
 from ldai import log
 from ldai.models import LDMessage
-from ldai.providers.types import ModelResponse, StructuredResponse, ToolRegistry
+from ldai.providers.types import RunnerResult, ToolRegistry
 
 
 class AIProvider(ABC):
@@ -16,48 +16,22 @@ class AIProvider(ABC):
     create_model(), create_agent(), and create_agent_graph().
     """
 
-    async def invoke_model(self, messages: List[LDMessage]) -> ModelResponse:
+    async def invoke_model(self, messages: List[LDMessage]) -> RunnerResult:
         """
         Invoke the chat model with an array of messages.
 
         Default implementation takes no action and returns a placeholder response.
         Provider implementations should override this method.
 
         :param messages: Array of LDMessage objects representing the conversation
-        :return: ModelResponse containing the model's response
+        :return: RunnerResult containing the model's response
         """
         log.warning('invoke_model not implemented by this provider')
 
-        from ldai.models import LDMessage
         from ldai.providers.types import LDAIMetrics
 
-        return ModelResponse(
-            message=LDMessage(role='assistant', content=''),
-            metrics=LDAIMetrics(success=False, usage=None),
-        )
-
-    async def invoke_structured_model(
-        self,
-        messages: List[LDMessage],
-        response_structure: Dict[str, Any],
-    ) -> StructuredResponse:
-        """
-        Invoke the chat model with structured output support.
-
-        Default implementation takes no action and returns a placeholder response.
-        Provider implementations should override this method.
-
-        :param messages: Array of LDMessage objects representing the conversation
-        :param response_structure: Dictionary of output configurations keyed by output name
-        :return: StructuredResponse containing the structured data
-        """
-        log.warning('invoke_structured_model not implemented by this provider')
-
-        from ldai.providers.types import LDAIMetrics
-
-        return StructuredResponse(
-            data={},
-            raw_response='',
+        return RunnerResult(
+            content='',
             metrics=LDAIMetrics(success=False, usage=None),
         )
 

packages/sdk/server-ai/src/ldai/providers/model_runner.py

@@ -1,7 +1,7 @@
-from typing import Any, Dict, List, Protocol, runtime_checkable
+from typing import List, Protocol, runtime_checkable
 
 from ldai.models import LDMessage
-from ldai.providers.types import ModelResponse, StructuredResponse
+from ldai.providers.types import RunnerResult
 
 
 @runtime_checkable
@@ -14,25 +14,11 @@ class ModelRunner(Protocol):
     and with what parameters — the caller just passes messages.
     """
 
-    async def invoke_model(self, messages: List[LDMessage]) -> ModelResponse:
+    async def invoke_model(self, messages: List[LDMessage]) -> RunnerResult:
         """
         Invoke the model with an array of messages.
 
         :param messages: Array of LDMessage objects representing the conversation
-        :return: ModelResponse containing the model's response and metrics
-        """
-        ...
-
-    async def invoke_structured_model(
-        self,
-        messages: List[LDMessage],
-        response_structure: Dict[str, Any],
-    ) -> StructuredResponse:
-        """
-        Invoke the model with structured output support.
-
-        :param messages: Array of LDMessage objects representing the conversation
-        :param response_structure: Dictionary defining the JSON schema for output structure
-        :return: StructuredResponse containing the structured data
+        :return: RunnerResult containing the model's response and metrics
         """
         ...