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 }
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 )
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 }