11import asyncio
2- import warnings
3- from typing import List , Union
2+ from typing import List
43
54from ldai import log
65from ldai .models import AICompletionConfig , LDMessage
7- from ldai .providers .model_runner import ModelRunner
86from ldai .providers .runner import Runner
9- from ldai .providers .types import JudgeResult , ManagedResult , ModelResponse , RunnerResult
7+ from ldai .providers .types import JudgeResult , ManagedResult , RunnerResult
108from ldai .tracker import LDAIConfigTracker
119
1210
1311class ManagedModel :
1412 """
1513 LaunchDarkly managed wrapper for AI model invocations.
1614
17- Holds a Runner (or legacy ModelRunner) . Handles conversation management,
18- judge evaluation dispatch, and tracking automatically via ``create_tracker()``.
15+ Holds a Runner. Handles conversation management, judge evaluation
16+ dispatch, and tracking automatically via ``create_tracker()``.
1917 Obtain an instance via ``LDAIClient.create_model()``.
2018 """
2119
2220 def __init__ (
2321 self ,
2422 ai_config : AICompletionConfig ,
25- model_runner : Union [ Runner , ModelRunner ] ,
23+ model_runner : Runner ,
2624 ):
2725 self ._ai_config = ai_config
2826 self ._model_runner = model_runner
@@ -48,86 +46,27 @@ async def run(self, prompt: str) -> ManagedResult:
4846 config_messages = self ._ai_config .messages or []
4947 all_messages = config_messages + self ._messages
5048
51- result : Union [ RunnerResult , ModelResponse ] = await tracker .track_metrics_of_async (
49+ result : RunnerResult = await tracker .track_metrics_of_async (
5250 lambda r : r .metrics ,
53- lambda : self ._invoke_runner (all_messages ),
51+ lambda : self ._model_runner . run (all_messages ),
5452 )
5553
56- # Support both new RunnerResult and legacy ModelResponse
57- if isinstance (result , RunnerResult ):
58- content = result .content
59- raw = result .raw
60- parsed = result .parsed
61- assistant_message = LDMessage (role = 'assistant' , content = content )
62- else :
63- content = result .message .content
64- raw = getattr (result , 'raw' , None )
65- parsed = getattr (result , 'parsed' , None )
66- assistant_message = result .message
54+ assistant_message = LDMessage (role = 'assistant' , content = result .content )
6755
6856 input_text = '\r \n ' .join (m .content for m in self ._messages ) if self ._messages else ''
6957
70- evaluations_task = self ._track_judge_results (tracker , input_text , content )
58+ evaluations_task = self ._track_judge_results (tracker , input_text , result . content )
7159
7260 self ._messages .append (assistant_message )
7361
7462 return ManagedResult (
75- content = content ,
63+ content = result . content ,
7664 metrics = tracker .get_summary (),
77- raw = raw ,
78- parsed = parsed ,
65+ raw = result . raw ,
66+ parsed = result . parsed ,
7967 evaluations = evaluations_task ,
8068 )
8169
82- async def _invoke_runner (
83- self , all_messages : List [LDMessage ]
84- ) -> Union [RunnerResult , ModelResponse ]:
85- """
86- Delegate to the runner. Supports both the new ``Runner`` protocol
87- (``run(messages) → RunnerResult``) and the legacy ``ModelRunner``
88- (``invoke_model(messages) → ModelResponse``).
89- """
90- if isinstance (self ._model_runner , Runner ):
91- return await self ._model_runner .run (all_messages )
92- # Legacy ModelRunner path
93- return await self ._model_runner .invoke_model (all_messages ) # type: ignore[union-attr]
94-
95- async def invoke (self , prompt : str ) -> ModelResponse :
96- """
97- Invoke the model with a prompt string.
98-
99- .. deprecated::
100- Use :meth:`run` instead. This method will be removed in a future
101- release once the migration to :class:`ManagedResult` is complete.
102-
103- :param prompt: The user prompt to send to the model
104- :return: ModelResponse containing the model's response and metrics
105- """
106- warnings .warn (
107- "ManagedModel.invoke() is deprecated. Use run() instead." ,
108- DeprecationWarning ,
109- stacklevel = 2 ,
110- )
111- tracker = self ._ai_config .create_tracker ()
112-
113- user_message = LDMessage (role = 'user' , content = prompt )
114- self ._messages .append (user_message )
115-
116- config_messages = self ._ai_config .messages or []
117- all_messages = config_messages + self ._messages
118-
119- response : ModelResponse = await tracker .track_metrics_of_async (
120- lambda result : result .metrics ,
121- lambda : self ._model_runner .invoke_model (all_messages ), # type: ignore[union-attr]
122- )
123-
124- input_text = '\r \n ' .join (m .content for m in self ._messages ) if self ._messages else ''
125- output_text = response .message .content
126- response .evaluations = self ._track_judge_results (tracker , input_text , output_text )
127-
128- self ._messages .append (response .message )
129- return response
130-
13170 def _track_judge_results (
13271 self ,
13372 tracker : LDAIConfigTracker ,
@@ -169,11 +108,11 @@ def append_messages(self, messages: List[LDMessage]) -> None:
169108 """
170109 self ._messages .extend (messages )
171110
172- def get_model_runner (self ) -> Union [ Runner , ModelRunner ] :
111+ def get_model_runner (self ) -> Runner :
173112 """
174113 Return the underlying runner for advanced use.
175114
176- :return: The Runner or legacy ModelRunner instance.
115+ :return: The Runner instance.
177116 """
178117 return self ._model_runner
179118
0 commit comments