kiln_ai.adapters.model_adapters.openai_model_adapter

  1from typing import Any, Dict
  2
  3from openai import AsyncOpenAI
  4from openai.types.chat import (
  5    ChatCompletion,
  6    ChatCompletionAssistantMessageParam,
  7    ChatCompletionSystemMessageParam,
  8    ChatCompletionUserMessageParam,
  9)
 10
 11import kiln_ai.datamodel as datamodel
 12from kiln_ai.adapters.ml_model_list import StructuredOutputMode
 13from kiln_ai.adapters.model_adapters.base_adapter import (
 14    COT_FINAL_ANSWER_PROMPT,
 15    AdapterInfo,
 16    BaseAdapter,
 17    BasePromptBuilder,
 18    RunOutput,
 19)
 20from kiln_ai.adapters.model_adapters.openai_compatible_config import (
 21    OpenAICompatibleConfig,
 22)
 23from kiln_ai.adapters.parsers.json_parser import parse_json_string
 24from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 25
 26
 27class OpenAICompatibleAdapter(BaseAdapter):
 28    def __init__(
 29        self,
 30        config: OpenAICompatibleConfig,
 31        kiln_task: datamodel.Task,
 32        prompt_builder: BasePromptBuilder | None = None,
 33        tags: list[str] | None = None,
 34    ):
 35        self.config = config
 36        self.client = AsyncOpenAI(
 37            api_key=config.api_key,
 38            base_url=config.base_url,
 39            default_headers=config.default_headers,
 40        )
 41
 42        super().__init__(
 43            kiln_task,
 44            model_name=config.model_name,
 45            model_provider_name=config.provider_name,
 46            prompt_builder=prompt_builder,
 47            tags=tags,
 48        )
 49
 50    async def _run(self, input: Dict | str) -> RunOutput:
 51        provider = self.model_provider()
 52        intermediate_outputs: dict[str, str] = {}
 53        prompt = self.build_prompt()
 54        user_msg = self.prompt_builder.build_user_message(input)
 55        messages = [
 56            ChatCompletionSystemMessageParam(role="system", content=prompt),
 57            ChatCompletionUserMessageParam(role="user", content=user_msg),
 58        ]
 59
 60        run_strategy, cot_prompt = self.run_strategy()
 61
 62        if run_strategy == "cot_as_message":
 63            if not cot_prompt:
 64                raise ValueError("cot_prompt is required for cot_as_message strategy")
 65            messages.append(
 66                ChatCompletionSystemMessageParam(role="system", content=cot_prompt)
 67            )
 68        elif run_strategy == "cot_two_call":
 69            if not cot_prompt:
 70                raise ValueError("cot_prompt is required for cot_two_call strategy")
 71            messages.append(
 72                ChatCompletionSystemMessageParam(role="system", content=cot_prompt)
 73            )
 74
 75            # First call for chain of thought
 76            cot_response = await self.client.chat.completions.create(
 77                model=provider.provider_options["model"],
 78                messages=messages,
 79            )
 80            cot_content = cot_response.choices[0].message.content
 81            if cot_content is not None:
 82                intermediate_outputs["chain_of_thought"] = cot_content
 83
 84            messages.extend(
 85                [
 86                    ChatCompletionAssistantMessageParam(
 87                        role="assistant", content=cot_content
 88                    ),
 89                    ChatCompletionUserMessageParam(
 90                        role="user",
 91                        content=COT_FINAL_ANSWER_PROMPT,
 92                    ),
 93                ]
 94            )
 95
 96        # OpenRouter specific options for reasoning models
 97        extra_body = {}
 98        require_or_reasoning = (
 99            self.config.openrouter_style_reasoning and provider.reasoning_capable
100        )
101        if require_or_reasoning:
102            extra_body["include_reasoning"] = True
103            # Filter to providers that support the reasoning parameter
104            extra_body["provider"] = {
105                "require_parameters": True,
106                # Ugly to have these here, but big range of quality of R1 providers
107                "order": ["Fireworks", "Together"],
108                # fp8 quants are awful
109                "ignore": ["DeepInfra"],
110            }
111
112        # Main completion call
113        response_format_options = await self.response_format_options()
114        response = await self.client.chat.completions.create(
115            model=provider.provider_options["model"],
116            messages=messages,
117            extra_body=extra_body,
118            **response_format_options,
119        )
120
121        if not isinstance(response, ChatCompletion):
122            raise RuntimeError(
123                f"Expected ChatCompletion response, got {type(response)}."
124            )
125
126        if hasattr(response, "error") and response.error:  # pyright: ignore
127            raise RuntimeError(
128                f"OpenAI compatible API returned status code {response.error.get('code')}: {response.error.get('message') or 'Unknown error'}.\nError: {response.error}"  # pyright: ignore
129            )
130        if not response.choices or len(response.choices) == 0:
131            raise RuntimeError(
132                "No message content returned in the response from OpenAI compatible API"
133            )
134
135        message = response.choices[0].message
136
137        # Save reasoning if it exists (OpenRouter specific format)
138        if require_or_reasoning:
139            if (
140                hasattr(message, "reasoning") and message.reasoning  # pyright: ignore
141            ):
142                intermediate_outputs["reasoning"] = message.reasoning  # pyright: ignore
143            else:
144                raise RuntimeError(
145                    "Reasoning is required for this model, but no reasoning was returned from OpenRouter."
146                )
147
148        # the string content of the response
149        response_content = message.content
150
151        # Fallback: Use args of first tool call to task_response if it exists
152        if not response_content and message.tool_calls:
153            tool_call = next(
154                (
155                    tool_call
156                    for tool_call in message.tool_calls
157                    if tool_call.function.name == "task_response"
158                ),
159                None,
160            )
161            if tool_call:
162                response_content = tool_call.function.arguments
163
164        if not isinstance(response_content, str):
165            raise RuntimeError(f"response is not a string: {response_content}")
166
167        if self.has_structured_output():
168            structured_response = parse_json_string(response_content)
169            return RunOutput(
170                output=structured_response,
171                intermediate_outputs=intermediate_outputs,
172            )
173
174        return RunOutput(
175            output=response_content,
176            intermediate_outputs=intermediate_outputs,
177        )
178
179    def adapter_info(self) -> AdapterInfo:
180        return AdapterInfo(
181            model_name=self.model_name,
182            model_provider=self.model_provider_name,
183            adapter_name="kiln_openai_compatible_adapter",
184            prompt_builder_name=self.prompt_builder.__class__.prompt_builder_name(),
185            prompt_id=self.prompt_builder.prompt_id(),
186        )
187
188    async def response_format_options(self) -> dict[str, Any]:
189        # Unstructured if task isn't structured
190        if not self.has_structured_output():
191            return {}
192
193        provider = self.model_provider()
194        match provider.structured_output_mode:
195            case StructuredOutputMode.json_mode:
196                return {"response_format": {"type": "json_object"}}
197            case StructuredOutputMode.json_schema:
198                output_schema = self.kiln_task.output_schema()
199                return {
200                    "response_format": {
201                        "type": "json_schema",
202                        "json_schema": {
203                            "name": "task_response",
204                            "schema": output_schema,
205                        },
206                    }
207                }
208            case StructuredOutputMode.function_calling:
209                return self.tool_call_params()
210            case StructuredOutputMode.json_instructions:
211                # JSON done via instructions in prompt, not the API response format. Do not ask for json_object (see option below).
212                return {}
213            case StructuredOutputMode.json_instruction_and_object:
214                # We set response_format to json_object and also set json instructions in the prompt
215                return {"response_format": {"type": "json_object"}}
216            case StructuredOutputMode.default:
217                # Default to function calling -- it's older than the other modes. Higher compatibility.
218                return self.tool_call_params()
219            case _:
220                raise_exhaustive_enum_error(provider.structured_output_mode)
221
222    def tool_call_params(self) -> dict[str, Any]:
223        # Add additional_properties: false to the schema (OpenAI requires this for some models)
224        output_schema = self.kiln_task.output_schema()
225        if not isinstance(output_schema, dict):
226            raise ValueError(
227                "Invalid output schema for this task. Can not use tool calls."
228            )
229        output_schema["additionalProperties"] = False
230
231        return {
232            "tools": [
233                {
234                    "type": "function",
235                    "function": {
236                        "name": "task_response",
237                        "parameters": output_schema,
238                        "strict": True,
239                    },
240                }
241            ],
242            "tool_choice": {
243                "type": "function",
244                "function": {"name": "task_response"},
245            },
246        }
class OpenAICompatibleAdapter(kiln_ai.adapters.model_adapters.base_adapter.BaseAdapter):
 28class OpenAICompatibleAdapter(BaseAdapter):
 29    def __init__(
 30        self,
 31        config: OpenAICompatibleConfig,
 32        kiln_task: datamodel.Task,
 33        prompt_builder: BasePromptBuilder | None = None,
 34        tags: list[str] | None = None,
 35    ):
 36        self.config = config
 37        self.client = AsyncOpenAI(
 38            api_key=config.api_key,
 39            base_url=config.base_url,
 40            default_headers=config.default_headers,
 41        )
 42
 43        super().__init__(
 44            kiln_task,
 45            model_name=config.model_name,
 46            model_provider_name=config.provider_name,
 47            prompt_builder=prompt_builder,
 48            tags=tags,
 49        )
 50
 51    async def _run(self, input: Dict | str) -> RunOutput:
 52        provider = self.model_provider()
 53        intermediate_outputs: dict[str, str] = {}
 54        prompt = self.build_prompt()
 55        user_msg = self.prompt_builder.build_user_message(input)
 56        messages = [
 57            ChatCompletionSystemMessageParam(role="system", content=prompt),
 58            ChatCompletionUserMessageParam(role="user", content=user_msg),
 59        ]
 60
 61        run_strategy, cot_prompt = self.run_strategy()
 62
 63        if run_strategy == "cot_as_message":
 64            if not cot_prompt:
 65                raise ValueError("cot_prompt is required for cot_as_message strategy")
 66            messages.append(
 67                ChatCompletionSystemMessageParam(role="system", content=cot_prompt)
 68            )
 69        elif run_strategy == "cot_two_call":
 70            if not cot_prompt:
 71                raise ValueError("cot_prompt is required for cot_two_call strategy")
 72            messages.append(
 73                ChatCompletionSystemMessageParam(role="system", content=cot_prompt)
 74            )
 75
 76            # First call for chain of thought
 77            cot_response = await self.client.chat.completions.create(
 78                model=provider.provider_options["model"],
 79                messages=messages,
 80            )
 81            cot_content = cot_response.choices[0].message.content
 82            if cot_content is not None:
 83                intermediate_outputs["chain_of_thought"] = cot_content
 84
 85            messages.extend(
 86                [
 87                    ChatCompletionAssistantMessageParam(
 88                        role="assistant", content=cot_content
 89                    ),
 90                    ChatCompletionUserMessageParam(
 91                        role="user",
 92                        content=COT_FINAL_ANSWER_PROMPT,
 93                    ),
 94                ]
 95            )
 96
 97        # OpenRouter specific options for reasoning models
 98        extra_body = {}
 99        require_or_reasoning = (
100            self.config.openrouter_style_reasoning and provider.reasoning_capable
101        )
102        if require_or_reasoning:
103            extra_body["include_reasoning"] = True
104            # Filter to providers that support the reasoning parameter
105            extra_body["provider"] = {
106                "require_parameters": True,
107                # Ugly to have these here, but big range of quality of R1 providers
108                "order": ["Fireworks", "Together"],
109                # fp8 quants are awful
110                "ignore": ["DeepInfra"],
111            }
112
113        # Main completion call
114        response_format_options = await self.response_format_options()
115        response = await self.client.chat.completions.create(
116            model=provider.provider_options["model"],
117            messages=messages,
118            extra_body=extra_body,
119            **response_format_options,
120        )
121
122        if not isinstance(response, ChatCompletion):
123            raise RuntimeError(
124                f"Expected ChatCompletion response, got {type(response)}."
125            )
126
127        if hasattr(response, "error") and response.error:  # pyright: ignore
128            raise RuntimeError(
129                f"OpenAI compatible API returned status code {response.error.get('code')}: {response.error.get('message') or 'Unknown error'}.\nError: {response.error}"  # pyright: ignore
130            )
131        if not response.choices or len(response.choices) == 0:
132            raise RuntimeError(
133                "No message content returned in the response from OpenAI compatible API"
134            )
135
136        message = response.choices[0].message
137
138        # Save reasoning if it exists (OpenRouter specific format)
139        if require_or_reasoning:
140            if (
141                hasattr(message, "reasoning") and message.reasoning  # pyright: ignore
142            ):
143                intermediate_outputs["reasoning"] = message.reasoning  # pyright: ignore
144            else:
145                raise RuntimeError(
146                    "Reasoning is required for this model, but no reasoning was returned from OpenRouter."
147                )
148
149        # the string content of the response
150        response_content = message.content
151
152        # Fallback: Use args of first tool call to task_response if it exists
153        if not response_content and message.tool_calls:
154            tool_call = next(
155                (
156                    tool_call
157                    for tool_call in message.tool_calls
158                    if tool_call.function.name == "task_response"
159                ),
160                None,
161            )
162            if tool_call:
163                response_content = tool_call.function.arguments
164
165        if not isinstance(response_content, str):
166            raise RuntimeError(f"response is not a string: {response_content}")
167
168        if self.has_structured_output():
169            structured_response = parse_json_string(response_content)
170            return RunOutput(
171                output=structured_response,
172                intermediate_outputs=intermediate_outputs,
173            )
174
175        return RunOutput(
176            output=response_content,
177            intermediate_outputs=intermediate_outputs,
178        )
179
180    def adapter_info(self) -> AdapterInfo:
181        return AdapterInfo(
182            model_name=self.model_name,
183            model_provider=self.model_provider_name,
184            adapter_name="kiln_openai_compatible_adapter",
185            prompt_builder_name=self.prompt_builder.__class__.prompt_builder_name(),
186            prompt_id=self.prompt_builder.prompt_id(),
187        )
188
189    async def response_format_options(self) -> dict[str, Any]:
190        # Unstructured if task isn't structured
191        if not self.has_structured_output():
192            return {}
193
194        provider = self.model_provider()
195        match provider.structured_output_mode:
196            case StructuredOutputMode.json_mode:
197                return {"response_format": {"type": "json_object"}}
198            case StructuredOutputMode.json_schema:
199                output_schema = self.kiln_task.output_schema()
200                return {
201                    "response_format": {
202                        "type": "json_schema",
203                        "json_schema": {
204                            "name": "task_response",
205                            "schema": output_schema,
206                        },
207                    }
208                }
209            case StructuredOutputMode.function_calling:
210                return self.tool_call_params()
211            case StructuredOutputMode.json_instructions:
212                # JSON done via instructions in prompt, not the API response format. Do not ask for json_object (see option below).
213                return {}
214            case StructuredOutputMode.json_instruction_and_object:
215                # We set response_format to json_object and also set json instructions in the prompt
216                return {"response_format": {"type": "json_object"}}
217            case StructuredOutputMode.default:
218                # Default to function calling -- it's older than the other modes. Higher compatibility.
219                return self.tool_call_params()
220            case _:
221                raise_exhaustive_enum_error(provider.structured_output_mode)
222
223    def tool_call_params(self) -> dict[str, Any]:
224        # Add additional_properties: false to the schema (OpenAI requires this for some models)
225        output_schema = self.kiln_task.output_schema()
226        if not isinstance(output_schema, dict):
227            raise ValueError(
228                "Invalid output schema for this task. Can not use tool calls."
229            )
230        output_schema["additionalProperties"] = False
231
232        return {
233            "tools": [
234                {
235                    "type": "function",
236                    "function": {
237                        "name": "task_response",
238                        "parameters": output_schema,
239                        "strict": True,
240                    },
241                }
242            ],
243            "tool_choice": {
244                "type": "function",
245                "function": {"name": "task_response"},
246            },
247        }

Base class for AI model adapters that handle task execution.

This abstract class provides the foundation for implementing model-specific adapters that can process tasks with structured or unstructured inputs/outputs. It handles input/output validation, prompt building, and run tracking.

Attributes: prompt_builder (BasePromptBuilder): Builder for constructing prompts for the model kiln_task (Task): The task configuration and metadata output_schema (dict | None): JSON schema for validating structured outputs input_schema (dict | None): JSON schema for validating structured inputs

OpenAICompatibleAdapter( config: kiln_ai.adapters.model_adapters.openai_compatible_config.OpenAICompatibleConfig, kiln_task: kiln_ai.datamodel.Task, prompt_builder: kiln_ai.adapters.prompt_builders.BasePromptBuilder | None = None, tags: list[str] | None = None)
29    def __init__(
30        self,
31        config: OpenAICompatibleConfig,
32        kiln_task: datamodel.Task,
33        prompt_builder: BasePromptBuilder | None = None,
34        tags: list[str] | None = None,
35    ):
36        self.config = config
37        self.client = AsyncOpenAI(
38            api_key=config.api_key,
39            base_url=config.base_url,
40            default_headers=config.default_headers,
41        )
42
43        super().__init__(
44            kiln_task,
45            model_name=config.model_name,
46            model_provider_name=config.provider_name,
47            prompt_builder=prompt_builder,
48            tags=tags,
49        )
config
client
def adapter_info(self) -> kiln_ai.adapters.model_adapters.base_adapter.AdapterInfo:
180    def adapter_info(self) -> AdapterInfo:
181        return AdapterInfo(
182            model_name=self.model_name,
183            model_provider=self.model_provider_name,
184            adapter_name="kiln_openai_compatible_adapter",
185            prompt_builder_name=self.prompt_builder.__class__.prompt_builder_name(),
186            prompt_id=self.prompt_builder.prompt_id(),
187        )
async def response_format_options(self) -> dict[str, typing.Any]:
189    async def response_format_options(self) -> dict[str, Any]:
190        # Unstructured if task isn't structured
191        if not self.has_structured_output():
192            return {}
193
194        provider = self.model_provider()
195        match provider.structured_output_mode:
196            case StructuredOutputMode.json_mode:
197                return {"response_format": {"type": "json_object"}}
198            case StructuredOutputMode.json_schema:
199                output_schema = self.kiln_task.output_schema()
200                return {
201                    "response_format": {
202                        "type": "json_schema",
203                        "json_schema": {
204                            "name": "task_response",
205                            "schema": output_schema,
206                        },
207                    }
208                }
209            case StructuredOutputMode.function_calling:
210                return self.tool_call_params()
211            case StructuredOutputMode.json_instructions:
212                # JSON done via instructions in prompt, not the API response format. Do not ask for json_object (see option below).
213                return {}
214            case StructuredOutputMode.json_instruction_and_object:
215                # We set response_format to json_object and also set json instructions in the prompt
216                return {"response_format": {"type": "json_object"}}
217            case StructuredOutputMode.default:
218                # Default to function calling -- it's older than the other modes. Higher compatibility.
219                return self.tool_call_params()
220            case _:
221                raise_exhaustive_enum_error(provider.structured_output_mode)
def tool_call_params(self) -> dict[str, typing.Any]:
223    def tool_call_params(self) -> dict[str, Any]:
224        # Add additional_properties: false to the schema (OpenAI requires this for some models)
225        output_schema = self.kiln_task.output_schema()
226        if not isinstance(output_schema, dict):
227            raise ValueError(
228                "Invalid output schema for this task. Can not use tool calls."
229            )
230        output_schema["additionalProperties"] = False
231
232        return {
233            "tools": [
234                {
235                    "type": "function",
236                    "function": {
237                        "name": "task_response",
238                        "parameters": output_schema,
239                        "strict": True,
240                    },
241                }
242            ],
243            "tool_choice": {
244                "type": "function",
245                "function": {"name": "task_response"},
246            },
247        }