kiln_ai.adapters.model_adapters.base_adapter

  1import json
  2from abc import ABCMeta, abstractmethod
  3from dataclasses import dataclass
  4from typing import Dict, Literal, Tuple
  5
  6import jsonschema
  7
  8from kiln_ai.adapters.ml_model_list import KilnModelProvider, StructuredOutputMode
  9from kiln_ai.adapters.parsers.json_parser import parse_json_string
 10from kiln_ai.adapters.parsers.parser_registry import model_parser_from_id
 11from kiln_ai.adapters.parsers.request_formatters import request_formatter_from_id
 12from kiln_ai.adapters.prompt_builders import prompt_builder_from_id
 13from kiln_ai.adapters.provider_tools import kiln_model_provider_from
 14from kiln_ai.adapters.run_output import RunOutput
 15from kiln_ai.datamodel import (
 16    DataSource,
 17    DataSourceType,
 18    Task,
 19    TaskOutput,
 20    TaskRun,
 21    Usage,
 22)
 23from kiln_ai.datamodel.json_schema import validate_schema_with_value_error
 24from kiln_ai.datamodel.task import RunConfig
 25from kiln_ai.utils.config import Config
 26
 27
 28@dataclass
 29class AdapterConfig:
 30    """
 31    An adapter config is config options that do NOT impact the output of the model.
 32
 33    For example: if it's saved, of if we request additional data like logprobs.
 34    """
 35
 36    allow_saving: bool = True
 37    top_logprobs: int | None = None
 38    default_tags: list[str] | None = None
 39
 40
 41COT_FINAL_ANSWER_PROMPT = "Considering the above, return a final result."
 42
 43
 44class BaseAdapter(metaclass=ABCMeta):
 45    """Base class for AI model adapters that handle task execution.
 46
 47    This abstract class provides the foundation for implementing model-specific adapters
 48    that can process tasks with structured or unstructured inputs/outputs. It handles
 49    input/output validation, prompt building, and run tracking.
 50
 51    Attributes:
 52        prompt_builder (BasePromptBuilder): Builder for constructing prompts for the model
 53        kiln_task (Task): The task configuration and metadata
 54        output_schema (dict | None): JSON schema for validating structured outputs
 55        input_schema (dict | None): JSON schema for validating structured inputs
 56    """
 57
 58    def __init__(
 59        self,
 60        run_config: RunConfig,
 61        config: AdapterConfig | None = None,
 62    ):
 63        self.run_config = run_config
 64        self.prompt_builder = prompt_builder_from_id(
 65            run_config.prompt_id, run_config.task
 66        )
 67        self._model_provider: KilnModelProvider | None = None
 68
 69        self.output_schema = self.task().output_json_schema
 70        self.input_schema = self.task().input_json_schema
 71        self.base_adapter_config = config or AdapterConfig()
 72
 73    def task(self) -> Task:
 74        return self.run_config.task
 75
 76    def model_provider(self) -> KilnModelProvider:
 77        """
 78        Lazy load the model provider for this adapter.
 79        """
 80        if self._model_provider is not None:
 81            return self._model_provider
 82        if not self.run_config.model_name or not self.run_config.model_provider_name:
 83            raise ValueError("model_name and model_provider_name must be provided")
 84        self._model_provider = kiln_model_provider_from(
 85            self.run_config.model_name, self.run_config.model_provider_name
 86        )
 87        if not self._model_provider:
 88            raise ValueError(
 89                f"model_provider_name {self.run_config.model_provider_name} not found for model {self.run_config.model_name}"
 90            )
 91        return self._model_provider
 92
 93    async def invoke(
 94        self,
 95        input: Dict | str,
 96        input_source: DataSource | None = None,
 97    ) -> TaskRun:
 98        run_output, _ = await self.invoke_returning_run_output(input, input_source)
 99        return run_output
100
101    async def invoke_returning_run_output(
102        self,
103        input: Dict | str,
104        input_source: DataSource | None = None,
105    ) -> Tuple[TaskRun, RunOutput]:
106        # validate input
107        if self.input_schema is not None:
108            if not isinstance(input, dict):
109                raise ValueError(f"structured input is not a dict: {input}")
110
111            validate_schema_with_value_error(
112                input,
113                self.input_schema,
114                "This task requires a specific input schema. While the model produced JSON, that JSON didn't meet the schema. Search 'Troubleshooting Structured Data Issues' in our docs for more information.",
115            )
116
117        # Format model input for model call (we save the original input in the task without formatting)
118        formatted_input = input
119        formatter_id = self.model_provider().formatter
120        if formatter_id is not None:
121            formatter = request_formatter_from_id(formatter_id)
122            formatted_input = formatter.format_input(input)
123
124        # Run
125        run_output, usage = await self._run(formatted_input)
126
127        # Parse
128        provider = self.model_provider()
129        parser = model_parser_from_id(provider.parser)
130        parsed_output = parser.parse_output(original_output=run_output)
131
132        # validate output
133        if self.output_schema is not None:
134            # Parse json to dict if we have structured output
135            if isinstance(parsed_output.output, str):
136                parsed_output.output = parse_json_string(parsed_output.output)
137
138            if not isinstance(parsed_output.output, dict):
139                raise RuntimeError(
140                    f"structured response is not a dict: {parsed_output.output}"
141                )
142            validate_schema_with_value_error(
143                parsed_output.output,
144                self.output_schema,
145                "This task requires a specific output schema. While the model produced JSON, that JSON didn't meet the schema. Search 'Troubleshooting Structured Data Issues' in our docs for more information.",
146            )
147        else:
148            if not isinstance(parsed_output.output, str):
149                raise RuntimeError(
150                    f"response is not a string for non-structured task: {parsed_output.output}"
151                )
152
153        # Validate reasoning content is present (if reasoning)
154        if provider.reasoning_capable and (
155            not parsed_output.intermediate_outputs
156            or "reasoning" not in parsed_output.intermediate_outputs
157        ):
158            raise RuntimeError(
159                "Reasoning is required for this model, but no reasoning was returned."
160            )
161
162        # Generate the run and output
163        run = self.generate_run(input, input_source, parsed_output, usage)
164
165        # Save the run if configured to do so, and we have a path to save to
166        if (
167            self.base_adapter_config.allow_saving
168            and Config.shared().autosave_runs
169            and self.task().path is not None
170        ):
171            run.save_to_file()
172        else:
173            # Clear the ID to indicate it's not persisted
174            run.id = None
175
176        return run, run_output
177
178    def has_structured_output(self) -> bool:
179        return self.output_schema is not None
180
181    @abstractmethod
182    def adapter_name(self) -> str:
183        pass
184
185    @abstractmethod
186    async def _run(self, input: Dict | str) -> Tuple[RunOutput, Usage | None]:
187        pass
188
189    def build_prompt(self) -> str:
190        # The prompt builder needs to know if we want to inject formatting instructions
191        provider = self.model_provider()
192        add_json_instructions = self.has_structured_output() and (
193            provider.structured_output_mode == StructuredOutputMode.json_instructions
194            or provider.structured_output_mode
195            == StructuredOutputMode.json_instruction_and_object
196        )
197
198        return self.prompt_builder.build_prompt(
199            include_json_instructions=add_json_instructions
200        )
201
202    def run_strategy(
203        self,
204    ) -> Tuple[Literal["cot_as_message", "cot_two_call", "basic"], str | None]:
205        # Determine the run strategy for COT prompting. 3 options:
206        # 1. "Thinking" LLM designed to output thinking in a structured format plus a COT prompt: we make 1 call to the LLM, which outputs thinking in a structured format. We include the thinking instuctions as a message.
207        # 2. Normal LLM with COT prompt: 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. It also separates the thinking from the final response.
208        # 3. Non chain of thought: we make 1 call to the LLM, with no COT prompt.
209        cot_prompt = self.prompt_builder.chain_of_thought_prompt()
210        reasoning_capable = self.model_provider().reasoning_capable
211
212        if cot_prompt and reasoning_capable:
213            # 1: "Thinking" LLM designed to output thinking in a structured format
214            # A simple message with the COT prompt appended to the message list is sufficient
215            return "cot_as_message", cot_prompt
216        elif cot_prompt:
217            # 2: Unstructured output with COT
218            # Two calls to separate the thinking from the final response
219            return "cot_two_call", cot_prompt
220        else:
221            return "basic", None
222
223    # create a run and task output
224    def generate_run(
225        self,
226        input: Dict | str,
227        input_source: DataSource | None,
228        run_output: RunOutput,
229        usage: Usage | None = None,
230    ) -> TaskRun:
231        # Convert input and output to JSON strings if they are dictionaries
232        input_str = (
233            json.dumps(input, ensure_ascii=False) if isinstance(input, dict) else input
234        )
235        output_str = (
236            json.dumps(run_output.output, ensure_ascii=False)
237            if isinstance(run_output.output, dict)
238            else run_output.output
239        )
240
241        # If no input source is provided, use the human data source
242        if input_source is None:
243            input_source = DataSource(
244                type=DataSourceType.human,
245                properties={"created_by": Config.shared().user_id},
246            )
247
248        new_task_run = TaskRun(
249            parent=self.task(),
250            input=input_str,
251            input_source=input_source,
252            output=TaskOutput(
253                output=output_str,
254                # Synthetic since an adapter, not a human, is creating this
255                source=DataSource(
256                    type=DataSourceType.synthetic,
257                    properties=self._properties_for_task_output(),
258                ),
259            ),
260            intermediate_outputs=run_output.intermediate_outputs,
261            tags=self.base_adapter_config.default_tags or [],
262            usage=usage,
263        )
264
265        return new_task_run
266
267    def _properties_for_task_output(self) -> Dict[str, str | int | float]:
268        props = {}
269
270        # adapter info
271        props["adapter_name"] = self.adapter_name()
272        props["model_name"] = self.run_config.model_name
273        props["model_provider"] = self.run_config.model_provider_name
274        props["prompt_id"] = self.run_config.prompt_id
275
276        return props