kiln_ai.adapters.model_adapters.openai_model_adapter

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

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