11import json
2- from typing import Any , Dict , List
2+ from typing import Any , Dict , List , Optional
33
44from ldai import LDMessage , log
5- from ldai .providers .model_runner import ModelRunner
6- from ldai .providers .types import LDAIMetrics , ModelResponse , StructuredResponse
5+ from ldai .providers .types import LDAIMetrics , RunnerResult
76from openai import AsyncOpenAI
87
98from ldai_openai .openai_helper import (
1211)
1312
1413
15- class OpenAIModelRunner ( ModelRunner ) :
14+ class OpenAIModelRunner :
1615 """
17- ModelRunner implementation for OpenAI.
16+ Runner implementation for OpenAI chat completions .
1817
1918 Holds a fully-configured AsyncOpenAI client, model name, and parameters.
20- Returned by OpenAIConnector.create_model(config).
19+ Returned by ``OpenAIRunnerFactory.create_model(config)``.
20+
21+ Implements the unified :class:`~ldai.providers.runner.Runner` protocol via
22+ :meth:`run`.
2123 """
2224
2325 def __init__ (
@@ -30,13 +32,38 @@ def __init__(
3032 self ._model_name = model_name
3133 self ._parameters = parameters
3234
33- async def invoke_model (self , messages : List [LDMessage ]) -> ModelResponse :
35+ async def run (
36+ self ,
37+ input : Any ,
38+ output_type : Optional [Dict [str , Any ]] = None ,
39+ ) -> RunnerResult :
3440 """
35- Invoke the OpenAI model with an array of messages.
36-
37- :param messages: Array of LDMessage objects representing the conversation
38- :return: ModelResponse containing the model's response and metrics
41+ Run the OpenAI model with the given input.
42+
43+ :param input: A string prompt or a list of :class:`LDMessage` objects
44+ :param output_type: Optional JSON schema dict requesting structured output.
45+ When provided, ``parsed`` on the returned :class:`RunnerResult` is
46+ populated with the parsed JSON document.
47+ :return: :class:`RunnerResult` containing ``content``, ``metrics``,
48+ ``raw`` and (when ``output_type`` is set) ``parsed``.
3949 """
50+ messages = self ._coerce_input (input )
51+
52+ if output_type is not None :
53+ return await self ._run_structured (messages , output_type )
54+ return await self ._run_completion (messages )
55+
56+ @staticmethod
57+ def _coerce_input (input : Any ) -> List [LDMessage ]:
58+ if isinstance (input , str ):
59+ return [LDMessage (role = 'user' , content = input )]
60+ if isinstance (input , list ):
61+ return input
62+ raise TypeError (
63+ f"Unsupported input type for OpenAIModelRunner.run: { type (input ).__name__ } "
64+ )
65+
66+ async def _run_completion (self , messages : List [LDMessage ]) -> RunnerResult :
4067 try :
4168 response = await self ._client .chat .completions .create (
4269 model = self ._model_name ,
@@ -45,40 +72,29 @@ async def invoke_model(self, messages: List[LDMessage]) -> ModelResponse:
4572 )
4673
4774 metrics = get_ai_metrics_from_response (response )
48-
49- content = ''
50- if response .choices and len (response .choices ) > 0 :
51- message = response .choices [0 ].message
52- if message and message .content :
53- content = message .content
75+ content = self ._extract_content (response )
5476
5577 if not content :
5678 log .warning ('OpenAI response has no content available' )
57- metrics = LDAIMetrics (success = False , usage = metrics .usage )
79+ return RunnerResult (
80+ content = '' ,
81+ metrics = LDAIMetrics (success = False , usage = metrics .usage ),
82+ raw = response ,
83+ )
5884
59- return ModelResponse (
60- message = LDMessage (role = 'assistant' , content = content ),
61- metrics = metrics ,
62- )
85+ return RunnerResult (content = content , metrics = metrics , raw = response )
6386 except Exception as error :
6487 log .warning (f'OpenAI model invocation failed: { error } ' )
65- return ModelResponse (
66- message = LDMessage ( role = 'assistant' , content = '' ) ,
88+ return RunnerResult (
89+ content = '' ,
6790 metrics = LDAIMetrics (success = False , usage = None ),
6891 )
6992
70- async def invoke_structured_model (
93+ async def _run_structured (
7194 self ,
7295 messages : List [LDMessage ],
73- response_structure : Dict [str , Any ],
74- ) -> StructuredResponse :
75- """
76- Invoke the OpenAI model with structured output support.
77-
78- :param messages: Array of LDMessage objects representing the conversation
79- :param response_structure: Dictionary defining the JSON schema for output structure
80- :return: StructuredResponse containing the structured data
81- """
96+ output_type : Dict [str , Any ],
97+ ) -> RunnerResult :
8298 try :
8399 response = await self ._client .chat .completions .create (
84100 model = self ._model_name ,
@@ -87,43 +103,50 @@ async def invoke_structured_model(
87103 'type' : 'json_schema' ,
88104 'json_schema' : {
89105 'name' : 'structured_output' ,
90- 'schema' : response_structure ,
106+ 'schema' : output_type ,
91107 'strict' : True ,
92108 },
93109 },
94110 ** self ._parameters ,
95111 )
96112
97113 metrics = get_ai_metrics_from_response (response )
98-
99- content = ''
100- if response .choices and len (response .choices ) > 0 :
101- message = response .choices [0 ].message
102- if message and message .content :
103- content = message .content
114+ content = self ._extract_content (response )
104115
105116 if not content :
106117 log .warning ('OpenAI structured response has no content available' )
107- return StructuredResponse (
108- data = {},
109- raw_response = '' ,
118+ return RunnerResult (
119+ content = '' ,
110120 metrics = LDAIMetrics (success = False , usage = metrics .usage ),
121+ raw = response ,
111122 )
112123
113124 try :
114- data = json .loads (content )
115- return StructuredResponse (data = data , raw_response = content , metrics = metrics )
125+ parsed = json .loads (content )
126+ return RunnerResult (
127+ content = content ,
128+ metrics = metrics ,
129+ raw = response ,
130+ parsed = parsed ,
131+ )
116132 except json .JSONDecodeError as parse_error :
117133 log .warning (f'OpenAI structured response contains invalid JSON: { parse_error } ' )
118- return StructuredResponse (
119- data = {},
120- raw_response = content ,
134+ return RunnerResult (
135+ content = content ,
121136 metrics = LDAIMetrics (success = False , usage = metrics .usage ),
137+ raw = response ,
122138 )
123139 except Exception as error :
124140 log .warning (f'OpenAI structured model invocation failed: { error } ' )
125- return StructuredResponse (
126- data = {},
127- raw_response = '' ,
141+ return RunnerResult (
142+ content = '' ,
128143 metrics = LDAIMetrics (success = False , usage = None ),
129144 )
145+
146+ @staticmethod
147+ def _extract_content (response : Any ) -> str :
148+ if response .choices and len (response .choices ) > 0 :
149+ message = response .choices [0 ].message
150+ if message and message .content :
151+ return message .content
152+ return ''
0 commit comments