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