11import asyncio
2- from typing import List , Optional
2+ import warnings
3+ from typing import List , Union
34
45from ldai .models import AICompletionConfig , LDMessage
56from ldai .providers .model_runner import ModelRunner
6- from ldai .providers .types import JudgeResult , ModelResponse
7+ from ldai .providers .runner import Runner
8+ from ldai .providers .types import JudgeResult , ManagedResult , ModelResponse , RunnerResult
79from ldai .tracker import LDAIConfigTracker
810
911
1012class ManagedModel :
1113 """
1214 LaunchDarkly managed wrapper for AI model invocations.
1315
14- Holds a ModelRunner. Handles conversation management, judge evaluation
15- dispatch, and tracking automatically via ``create_tracker()``.
16+ Holds a Runner (or legacy ModelRunner) . Handles conversation management,
17+ judge evaluation dispatch, and tracking automatically via ``create_tracker()``.
1618 Obtain an instance via ``LDAIClient.create_model()``.
1719 """
1820
1921 def __init__ (
2022 self ,
2123 ai_config : AICompletionConfig ,
22- model_runner : ModelRunner ,
24+ model_runner : Union [ Runner , ModelRunner ] ,
2325 ):
2426 self ._ai_config = ai_config
2527 self ._model_runner = model_runner
2628 self ._messages : List [LDMessage ] = []
2729
28- async def invoke (self , prompt : str ) -> ModelResponse :
30+ async def run (self , prompt : str ) -> ManagedResult :
2931 """
30- Invoke the model with a prompt string.
32+ Run the model with a prompt string.
3133
3234 Appends the prompt to the conversation history, prepends any
3335 system messages from the config, delegates to the runner, and
3436 appends the response to the history.
3537
38+ :param prompt: The user prompt to send to the model
39+ :return: ManagedResult containing the model's response, metric summary,
40+ and an optional evaluations task
41+ """
42+ tracker = self ._ai_config .create_tracker ()
43+
44+ user_message = LDMessage (role = 'user' , content = prompt )
45+ self ._messages .append (user_message )
46+
47+ config_messages = self ._ai_config .messages or []
48+ all_messages = config_messages + self ._messages
49+
50+ result : Union [RunnerResult , ModelResponse ] = await tracker .track_metrics_of_async (
51+ lambda r : r .metrics ,
52+ lambda : self ._invoke_runner (all_messages ),
53+ )
54+
55+ # Support both new RunnerResult and legacy ModelResponse
56+ if isinstance (result , RunnerResult ):
57+ content = result .content
58+ raw = result .raw
59+ parsed = result .parsed
60+ assistant_message = LDMessage (role = 'assistant' , content = content )
61+ else :
62+ content = result .message .content
63+ raw = getattr (result , 'raw' , None )
64+ parsed = getattr (result , 'parsed' , None )
65+ assistant_message = result .message
66+
67+ input_text = '\r \n ' .join (m .content for m in self ._messages ) if self ._messages else ''
68+
69+ evaluations_task = self ._track_judge_results (tracker , input_text , content )
70+
71+ self ._messages .append (assistant_message )
72+
73+ return ManagedResult (
74+ content = content ,
75+ metrics = tracker .get_summary (),
76+ raw = raw ,
77+ parsed = parsed ,
78+ evaluations = evaluations_task ,
79+ )
80+
81+ async def _invoke_runner (
82+ self , all_messages : List [LDMessage ]
83+ ) -> Union [RunnerResult , ModelResponse ]:
84+ """
85+ Delegate to the runner. Supports both the new ``Runner`` protocol
86+ (``run(messages) → RunnerResult``) and the legacy ``ModelRunner``
87+ (``invoke_model(messages) → ModelResponse``).
88+ """
89+ if isinstance (self ._model_runner , Runner ):
90+ return await self ._model_runner .run (all_messages )
91+ # Legacy ModelRunner path
92+ return await self ._model_runner .invoke_model (all_messages ) # type: ignore[union-attr]
93+
94+ async def invoke (self , prompt : str ) -> ModelResponse :
95+ """
96+ Invoke the model with a prompt string.
97+
98+ .. deprecated::
99+ Use :meth:`run` instead. This method will be removed in a future
100+ release once the migration to :class:`ManagedResult` is complete.
101+
36102 :param prompt: The user prompt to send to the model
37103 :return: ModelResponse containing the model's response and metrics
38104 """
105+ warnings .warn (
106+ "ManagedModel.invoke() is deprecated. Use run() instead." ,
107+ DeprecationWarning ,
108+ stacklevel = 2 ,
109+ )
39110 tracker = self ._ai_config .create_tracker ()
40111
41112 user_message = LDMessage (role = 'user' , content = prompt )
@@ -44,9 +115,9 @@ async def invoke(self, prompt: str) -> ModelResponse:
44115 config_messages = self ._ai_config .messages or []
45116 all_messages = config_messages + self ._messages
46117
47- response = await tracker .track_metrics_of_async (
118+ response : ModelResponse = await tracker .track_metrics_of_async (
48119 lambda result : result .metrics ,
49- lambda : self ._model_runner .invoke_model (all_messages ),
120+ lambda : self ._model_runner .invoke_model (all_messages ), # type: ignore[union-attr]
50121 )
51122
52123 input_text = '\r \n ' .join (m .content for m in self ._messages ) if self ._messages else ''
0 commit comments