kiln_ai.adapters.model_adapters.base_adapter

  1from __future__ import annotations
  2
  3import json
  4import uuid
  5from abc import ABCMeta, abstractmethod
  6from dataclasses import dataclass
  7from typing import TYPE_CHECKING, AsyncIterator, Dict, Tuple
  8
  9from litellm.types.utils import ModelResponseStream
 10
 11from kiln_ai.adapters.chat.chat_formatter import (
 12    ChatFormatter,
 13    MultiturnFormatter,
 14    chat_strategy_for_run,
 15    get_chat_formatter,
 16)
 17from kiln_ai.adapters.errors import (
 18    KilnRunError,
 19    StructuredOutputParseError,
 20    format_error_message,
 21)
 22from kiln_ai.adapters.ml_model_list import (
 23    KilnModelProvider,
 24    StructuredOutputMode,
 25    default_structured_output_mode_for_model_provider,
 26)
 27from kiln_ai.adapters.model_adapters.adapter_stream import AdapterStreamResult
 28from kiln_ai.adapters.model_adapters.stream_events import (
 29    AiSdkStreamConverter,
 30    AiSdkStreamEvent,
 31    FinishEvent,
 32    FinishMessageMetadata,
 33    FinishStepEvent,
 34    StartEvent,
 35    StartStepEvent,
 36    ToolCallEvent,
 37)
 38from kiln_ai.adapters.parsers.json_parser import parse_json_string
 39from kiln_ai.adapters.parsers.parser_registry import model_parser_from_id
 40from kiln_ai.adapters.parsers.request_formatters import request_formatter_from_id
 41from kiln_ai.adapters.prompt_builders import BasePromptBuilder, prompt_builder_from_id
 42from kiln_ai.adapters.provider_tools import kiln_model_provider_from
 43from kiln_ai.adapters.run_output import RunOutput
 44from kiln_ai.datamodel import (
 45    DataSource,
 46    DataSourceType,
 47    MessageUsage,
 48    Task,
 49    TaskOutput,
 50    TaskRun,
 51    Usage,
 52)
 53from kiln_ai.datamodel.datamodel_enums import InputType
 54from kiln_ai.datamodel.json_schema import validate_schema_with_value_error
 55from kiln_ai.datamodel.run_config import (
 56    KilnAgentRunConfigProperties,
 57    as_kiln_agent_run_config,
 58)
 59from kiln_ai.datamodel.skill import Skill
 60from kiln_ai.datamodel.task import RunConfigProperties
 61from kiln_ai.datamodel.task_output import TASK_OUTPUT_SCHEMA_ERROR_PREFIX
 62from kiln_ai.datamodel.tool_id import SKILL_TOOL_ID_PREFIX, skill_id_from_tool_id
 63from kiln_ai.tools import KilnToolInterface
 64from kiln_ai.tools.mcp_session_manager import mcp_session_scope
 65from kiln_ai.tools.skill_tool import SkillTool
 66from kiln_ai.tools.tool_registry import tool_from_id
 67from kiln_ai.utils.config import Config
 68from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 69from kiln_ai.utils.jinja_engine import render_input_transform
 70from kiln_ai.utils.open_ai_types import ChatCompletionMessageParam
 71
 72if TYPE_CHECKING:
 73    from kiln_ai.adapters.model_adapters.adapter_stream import AdapterStream
 74
 75SkillsDict = Dict[str, Skill]
 76
 77
 78@dataclass
 79class AdapterConfig:
 80    """
 81    An adapter config is config options that do NOT impact the output of the model.
 82
 83    For example: if it's saved, of if we request additional data like logprobs.
 84    """
 85
 86    allow_saving: bool = True
 87    top_logprobs: int | None = None
 88    default_tags: list[str] | None = None
 89
 90    """
 91    The ID of the TaskRunConfig that originated this run, if any. Stored on the
 92    resulting TaskRun so the run can be traced back to its originating saved config
 93    (in addition to the inline run_config snapshot). None for ad-hoc/inline runs
 94    that were not initiated from a saved TaskRunConfig.
 95    """
 96    task_run_config_id: str | None = None
 97
 98    """
 99    A custom prompt builder can be injected to override the system prompt building process.
100    If not provided, the prompt builder will be created from the run_config.prompt_id which
101    may load additional files from disk.
102    """
103    prompt_builder: BasePromptBuilder | None = None
104
105    """
106    Pre-loaded skills keyed by skill ID. When the run config references skills,
107    they are looked up from this dict instead of reading from the filesystem.
108    Use load_skills_for_task() to build this dict.
109    """
110    skills: SkillsDict | None = None
111
112    """
113    When True, the adapter will stop and return control to the caller when a tool call
114    is invoked, instead of processing tool calls internally. Default is False (process
115    tool calls internally).
116    """
117    return_on_tool_call: bool = False
118
119    """
120    Extra tools provided directly by the caller, in addition to tools resolved from the
121    task's tool registry. These are sent to the model together with registry tools, and
122    their names must not collide with registry tool names.
123
124    If ``return_on_tool_call`` is False (the default), the adapter executes these tools
125    itself just like registry tools. If True, the adapter returns as soon as the model
126    requests a tool call and the caller is responsible for running the tool and passing
127    results back via ``prior_trace``.
128    """
129    unmanaged_tools: list[KilnToolInterface] | None = None
130
131    """
132    When True, automatically inject prompt caching hints into completion
133    requests. This is a cost optimization and does not affect model output.
134    """
135    automatic_prompt_caching: bool = False
136
137
138class BaseAdapter(metaclass=ABCMeta):
139    """Base class for AI model adapters that handle task execution.
140
141    This abstract class provides the foundation for implementing model-specific adapters
142    that can process tasks with structured or unstructured inputs/outputs. It handles
143    input/output validation, prompt building, and run tracking.
144
145    Prompt building is handled internally by the adapter, which uses a prompt builder
146    based on the run config. To override the prompt building behavior, pass a custom prompt
147    builder to the adapter config.
148    """
149
150    def __init__(
151        self,
152        task: Task,
153        run_config: RunConfigProperties,
154        config: AdapterConfig | None = None,
155    ):
156        self.task = task
157        self.run_config: RunConfigProperties = run_config
158        self.base_adapter_config = config or AdapterConfig()
159
160        if isinstance(run_config, KilnAgentRunConfigProperties):
161            self.update_run_config_unknown_structured_output_mode()
162            self.prompt_builder = (
163                self.base_adapter_config.prompt_builder
164                or prompt_builder_from_id(run_config.prompt_id, task)
165            )
166        else:
167            self.prompt_builder = None
168        self._model_provider: KilnModelProvider | None = None
169        self._resolved_skills: list[Skill] | None = None
170
171        self.output_schema = task.output_json_schema
172        self.input_schema = task.input_json_schema
173
174    def model_provider(self) -> KilnModelProvider:
175        """
176        Lazy load the model provider for this adapter.
177        """
178        if self._model_provider is not None:
179            return self._model_provider
180        run_config = as_kiln_agent_run_config(self.run_config)
181        if not run_config.model_name or not run_config.model_provider_name:
182            raise ValueError("model_name and model_provider_name must be provided")
183        self._model_provider = kiln_model_provider_from(
184            run_config.model_name, run_config.model_provider_name
185        )
186        if not self._model_provider:
187            raise ValueError(
188                f"model_provider_name {run_config.model_provider_name} not found for model {run_config.model_name}"
189            )
190        return self._model_provider
191
192    @staticmethod
193    def _normalize_prior_trace(
194        prior_trace: list[ChatCompletionMessageParam] | None,
195    ) -> list[ChatCompletionMessageParam] | None:
196        if not prior_trace:
197            return None
198        return prior_trace
199
200    def _reject_multiturn_with_structured_input(
201        self,
202        prior_trace: list[ChatCompletionMessageParam] | None,
203    ) -> None:
204        if prior_trace is not None and self.input_schema is not None:
205            raise ValueError(
206                "Cannot run multiturn execution with a task that has a structured input schema. "
207                "Use an unstructured task, or call without prior_trace."
208            )
209
210    async def invoke(
211        self,
212        input: InputType,
213        input_source: DataSource | None = None,
214        prior_trace: list[ChatCompletionMessageParam] | None = None,
215        parent_task_run: TaskRun | None = None,
216    ) -> TaskRun:
217        task_run, _ = await self.invoke_returning_run_output(
218            input, input_source, prior_trace, parent_task_run
219        )
220        return task_run
221
222    async def _run_returning_run_output(
223        self,
224        input: InputType,
225        input_source: DataSource | None = None,
226        prior_trace: list[ChatCompletionMessageParam] | None = None,
227        parent_task_run: TaskRun | None = None,
228    ) -> Tuple[TaskRun, RunOutput]:
229        # Pre-run validation: these checks run before any model call, so
230        # there is no trace to preserve. They stay outside the
231        # exception-wrapping block and surface as plain exceptions.
232        prior_trace = self._normalize_prior_trace(prior_trace)
233        self._reject_multiturn_with_structured_input(prior_trace)
234
235        if self.input_schema is not None:
236            validate_schema_with_value_error(
237                input,
238                self.input_schema,
239                "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.",
240                require_object=False,
241            )
242
243        # Apply input transform if configured. The original `input` is preserved
244        # for TaskRun.input persistence; `model_input` is what the model sees.
245        model_input = self._apply_input_transform(input)
246
247        # Format model input for model call (we save the original input in the
248        # task without formatting). This runs in the adapter but before any
249        # trace is built, so it also stays outside the wrapped region.
250        formatted_input = model_input
251        formatter_id = self.model_provider().formatter
252        if formatter_id is not None:
253            formatter = request_formatter_from_id(formatter_id)
254            formatted_input = formatter.format_input(model_input)
255
256        # Allocate the trace-so-far list here so the reference survives any
257        # exception thrown from inside `_run` (or the post-processing that
258        # follows). `_run` must mutate this list in place (extend/append,
259        # or `list[:] = ...`) and never rebind the local name.
260        trace_ref: list[ChatCompletionMessageParam] = []
261        try:
262            run_output, usage = await self._run(
263                formatted_input, trace_ref, prior_trace=prior_trace
264            )
265
266            if not run_output.is_toolcall_pending:
267                # Normal completion: parse and validate output
268                provider = self.model_provider()
269                parser = model_parser_from_id(provider.parser)
270                parsed_output = parser.parse_output(original_output=run_output)
271
272                # validate output
273                if self.output_schema is not None:
274                    # Parse json to dict if we have structured output
275                    if isinstance(parsed_output.output, str):
276                        try:
277                            parsed_output.output = parse_json_string(
278                                parsed_output.output
279                            )
280                        except ValueError as e:
281                            # Re-type (message unchanged) so retries treat bad
282                            # JSON like a schema mismatch: a one-off model slip.
283                            raise StructuredOutputParseError(str(e)) from e
284
285                    if not isinstance(parsed_output.output, dict):
286                        # Valid JSON of the wrong shape (often the object wrapped
287                        # in a list) is the same one-off model slip as bad JSON,
288                        # so it carries the retryable type too.
289                        raise StructuredOutputParseError(
290                            f"structured response is not a dict: {parsed_output.output}"
291                        )
292                    validate_schema_with_value_error(
293                        parsed_output.output,
294                        self.output_schema,
295                        TASK_OUTPUT_SCHEMA_ERROR_PREFIX,
296                    )
297                else:
298                    if not isinstance(parsed_output.output, str):
299                        raise RuntimeError(
300                            f"response is not a string for non-structured task: {parsed_output.output}"
301                        )
302
303                trace_has_toolcalls = parsed_output.trace is not None and any(
304                    message.get("role", None) == "tool"
305                    for message in parsed_output.trace
306                )
307
308                # Validate reasoning content is present if required.
309                # Models often skip reasoning on the final turn when tools are involved, so we don't require it then.
310                if (
311                    provider.reasoning_capable
312                    and (
313                        not parsed_output.intermediate_outputs
314                        or "reasoning" not in parsed_output.intermediate_outputs
315                    )
316                    and not (
317                        provider.reasoning_optional_for_structured_output
318                        and self.has_structured_output()
319                    )
320                    and not trace_has_toolcalls
321                ):
322                    raise RuntimeError(
323                        "Reasoning is required for this model, but no reasoning was returned."
324                    )
325
326                run_output = parsed_output
327
328            run = self.generate_run(
329                input,
330                input_source,
331                run_output,
332                usage,
333                run_output.trace,
334                parent_task_run,
335            )
336
337            # Save the run if configured to do so, and we have a path to save to
338            if (
339                self.base_adapter_config.allow_saving
340                and Config.shared().autosave_runs
341                and self.task.path is not None
342            ):
343                run.save_to_file()
344            else:
345                # Clear the ID to indicate it's not persisted
346                run.id = None
347
348            return run, run_output
349        except KilnRunError:
350            # Already wrapped — pass through so we don't double-wrap.
351            raise
352        except Exception as e:
353            # Trace conversion can itself throw (e.g., a malformed partial
354            # assistant message was appended to `messages` just before the
355            # real failure). Never let that swallow the original exception —
356            # fall back to no trace so the user still sees the real error.
357            partial_trace: list[ChatCompletionMessageParam] | None = None
358            if trace_ref:
359                try:
360                    partial_trace = self._messages_to_trace(trace_ref)
361                except Exception:
362                    partial_trace = None
363            raise KilnRunError(
364                message=format_error_message(e),
365                partial_trace=partial_trace,
366                original=e,
367            ) from e
368
369    async def invoke_returning_run_output(
370        self,
371        input: InputType,
372        input_source: DataSource | None = None,
373        prior_trace: list[ChatCompletionMessageParam] | None = None,
374        parent_task_run: TaskRun | None = None,
375    ) -> Tuple[TaskRun, RunOutput]:
376        async with mcp_session_scope():
377            return await self._run_returning_run_output(
378                input, input_source, prior_trace, parent_task_run
379            )
380
381    def invoke_openai_stream(
382        self,
383        input: InputType,
384        input_source: DataSource | None = None,
385        prior_trace: list[ChatCompletionMessageParam] | None = None,
386        parent_task_run: TaskRun | None = None,
387    ) -> OpenAIStreamResult:
388        """Stream raw OpenAI-protocol chunks for the task execution.
389
390        Returns an async-iterable that yields ``ModelResponseStream`` chunks
391        as they arrive from the model.  After the iterator is exhausted the
392        run has been validated and saved (when configured).  The resulting
393        ``TaskRun`` is available via the ``.task_run`` property.
394
395        Tool-call rounds happen internally and are not surfaced; use
396        ``invoke_ai_sdk_stream`` if you need tool-call events.
397        """
398        return OpenAIStreamResult(
399            self, input, input_source, prior_trace, parent_task_run
400        )
401
402    def invoke_ai_sdk_stream(
403        self,
404        input: InputType,
405        input_source: DataSource | None = None,
406        prior_trace: list[ChatCompletionMessageParam] | None = None,
407        parent_task_run: TaskRun | None = None,
408    ) -> AiSdkStreamResult:
409        """Stream AI SDK protocol events for the task execution.
410
411        Returns an async-iterable that yields ``AiSdkStreamEvent`` instances
412        covering text, reasoning, tool-call lifecycle, step boundaries, and
413        control events.  After the iterator is exhausted the resulting
414        ``TaskRun`` is available via the ``.task_run`` property.
415        """
416        return AiSdkStreamResult(
417            self, input, input_source, prior_trace, parent_task_run
418        )
419
420    def _prepare_stream(
421        self,
422        input: InputType,
423        prior_trace: list[ChatCompletionMessageParam] | None,
424    ) -> AdapterStream:
425        prior_trace = self._normalize_prior_trace(prior_trace)
426        self._reject_multiturn_with_structured_input(prior_trace)
427
428        if self.input_schema is not None:
429            validate_schema_with_value_error(
430                input,
431                self.input_schema,
432                "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.",
433                require_object=False,
434            )
435
436        model_input = self._apply_input_transform(input)
437
438        formatted_input = model_input
439        formatter_id = self.model_provider().formatter
440        if formatter_id is not None:
441            formatter = request_formatter_from_id(formatter_id)
442            formatted_input = formatter.format_input(model_input)
443
444        return self._create_run_stream(formatted_input, prior_trace)
445
446    def _finalize_stream(
447        self,
448        adapter_stream: AdapterStream,
449        input: InputType,
450        input_source: DataSource | None,
451        parent_task_run: TaskRun | None = None,
452    ) -> TaskRun:
453        """Streaming invocations are only concerned with passing through events as they come in.
454        At the end of the stream, we still need to validate the output, create a run and everything
455        else that a non-streaming invocation would do.
456        """
457
458        result: AdapterStreamResult = adapter_stream.result
459        run_output = result.run_output
460        usage = result.usage
461
462        if not run_output.is_toolcall_pending:
463            # Normal completion: parse and validate output
464            provider = self.model_provider()
465            parser = model_parser_from_id(provider.parser)
466            parsed_output = parser.parse_output(original_output=run_output)
467
468            if self.output_schema is not None:
469                if isinstance(parsed_output.output, str):
470                    try:
471                        parsed_output.output = parse_json_string(parsed_output.output)
472                    except ValueError as e:
473                        # Re-type (message unchanged) so retries treat bad JSON
474                        # like a schema mismatch: a one-off model slip.
475                        raise StructuredOutputParseError(str(e)) from e
476                if not isinstance(parsed_output.output, dict):
477                    # Valid JSON of the wrong shape (often the object wrapped in
478                    # a list) is the same one-off model slip as bad JSON, so it
479                    # carries the retryable type too.
480                    raise StructuredOutputParseError(
481                        f"structured response is not a dict: {parsed_output.output}"
482                    )
483                validate_schema_with_value_error(
484                    parsed_output.output,
485                    self.output_schema,
486                    TASK_OUTPUT_SCHEMA_ERROR_PREFIX,
487                )
488            else:
489                if not isinstance(parsed_output.output, str):
490                    raise RuntimeError(
491                        f"response is not a string for non-structured task: {parsed_output.output}"
492                    )
493
494            trace_has_toolcalls = parsed_output.trace is not None and any(
495                message.get("role", None) == "tool" for message in parsed_output.trace
496            )
497            if (
498                provider.reasoning_capable
499                and (
500                    not parsed_output.intermediate_outputs
501                    or "reasoning" not in parsed_output.intermediate_outputs
502                )
503                and not (
504                    provider.reasoning_optional_for_structured_output
505                    and self.has_structured_output()
506                )
507                and not trace_has_toolcalls
508            ):
509                raise RuntimeError(
510                    "Reasoning is required for this model, but no reasoning was returned."
511                )
512
513            run_output = parsed_output
514
515        run = self.generate_run(
516            input, input_source, run_output, usage, run_output.trace, parent_task_run
517        )
518
519        if (
520            self.base_adapter_config.allow_saving
521            and Config.shared().autosave_runs
522            and self.task.path is not None
523        ):
524            run.save_to_file()
525        else:
526            run.id = None
527
528        return run
529
530    def _apply_input_transform(self, input: InputType) -> InputType:
531        """If the run config has an input_transform, render it and return the
532        resulting string. Otherwise return input unchanged.
533
534        MCP run configs (no input_transform field) are a no-op.
535        """
536        if not isinstance(self.run_config, KilnAgentRunConfigProperties):
537            return input
538        transform = self.run_config.input_transform
539        if transform is None:
540            return input
541        try:
542            return render_input_transform(transform, input)
543        except Exception as e:
544            raise ValueError(f"Input transform failed: {e}") from e
545
546    def has_structured_output(self) -> bool:
547        return self.output_schema is not None
548
549    @abstractmethod
550    def adapter_name(self) -> str:
551        pass
552
553    @abstractmethod
554    async def _run(
555        self,
556        input: InputType,
557        trace_ref: list[ChatCompletionMessageParam],
558        prior_trace: list[ChatCompletionMessageParam] | None = None,
559    ) -> Tuple[RunOutput, Usage | None]:
560        """Run the model. Implementations MUST mutate `trace_ref` in place
561        (extend/append, or `trace_ref[:] = ...`) — never rebind it — so the
562        caller keeps a live reference to the partial trace if an exception
563        escapes.
564        """
565        pass
566
567    def _messages_to_trace(
568        self,
569        messages: list[ChatCompletionMessageParam],
570    ) -> list[ChatCompletionMessageParam]:
571        """Convert the adapter's internal `messages` list to an API-safe trace.
572
573        Default implementation returns the list as-is. Adapters that store
574        internal message objects (e.g. LiteLLM's `Message`) should override
575        this to normalize to `ChatCompletionMessageParam` shapes.
576        """
577        return messages
578
579    def _create_run_stream(
580        self,
581        input: InputType,
582        prior_trace: list[ChatCompletionMessageParam] | None = None,
583    ) -> AdapterStream:
584        """Create a stream for the adapter. Implementations must override this method to support streaming."""
585        raise NotImplementedError("Streaming is not supported for this adapter type")
586
587    def build_prompt(self) -> str:
588        if self.prompt_builder is None:
589            raise ValueError("Prompt builder is not available for MCP run config")
590        # The prompt builder needs to know if we want to inject formatting instructions
591        structured_output_mode = as_kiln_agent_run_config(
592            self.run_config
593        ).structured_output_mode
594        add_json_instructions = self.has_structured_output() and (
595            structured_output_mode == StructuredOutputMode.json_instructions
596            or structured_output_mode
597            == StructuredOutputMode.json_instruction_and_object
598        )
599
600        return self.prompt_builder.build_prompt(
601            include_json_instructions=add_json_instructions,
602            skills=self._resolve_skills(),
603        )
604
605    def _resolve_skills(self) -> list[Skill]:
606        """Resolve skills from the injected skills dict.
607
608        Uses the pre-loaded skills dict from AdapterConfig. Caches the result
609        so that build_prompt and available_tools don't repeat
610        the lookup. Raises ValueError if the run config references a skill
611        that is not in the injected dict.
612        """
613        if self._resolved_skills is not None:
614            return self._resolved_skills
615
616        if self.run_config.type != "kiln_agent":
617            self._resolved_skills = []
618            return self._resolved_skills
619
620        tool_config = as_kiln_agent_run_config(self.run_config).tools_config
621        if tool_config is None or tool_config.tools is None:
622            self._resolved_skills = []
623            return self._resolved_skills
624
625        skill_tool_ids = [
626            tid for tid in tool_config.tools if tid.startswith(SKILL_TOOL_ID_PREFIX)
627        ]
628        if not skill_tool_ids:
629            self._resolved_skills = []
630            return self._resolved_skills
631
632        injected = self.base_adapter_config.skills
633        if injected is None:
634            raise ValueError(
635                "Run config references skills but no skills dict was provided via "
636                "AdapterConfig(skills=...). Use load_skills_for_task() to pre-load "
637                "skills and pass them to the adapter."
638            )
639
640        skills: list[Skill] = []
641        seen: set[str] = set()
642        for tool_id in skill_tool_ids:
643            sid = skill_id_from_tool_id(tool_id)
644            if sid not in injected:
645                raise ValueError(
646                    f"Skill {sid} referenced in run config but not found in the "
647                    "injected skills dict."
648                )
649            if sid in seen:
650                continue
651            seen.add(sid)
652            skills.append(injected[sid])
653
654        self._resolved_skills = skills
655        return self._resolved_skills
656
657    def build_chat_formatter(
658        self,
659        input: InputType,
660        prior_trace: list[ChatCompletionMessageParam] | None = None,
661    ) -> ChatFormatter:
662        prior_trace = self._normalize_prior_trace(prior_trace)
663        self._reject_multiturn_with_structured_input(prior_trace)
664        if prior_trace is not None:
665            return MultiturnFormatter(prior_trace, input)
666        if self.prompt_builder is None:
667            raise ValueError("Prompt builder is not available for MCP run config")
668        # Determine the chat strategy to use based on the prompt the user selected, the model's capabilities, and if the model was finetuned with a specific chat strategy.
669
670        cot_prompt = self.prompt_builder.chain_of_thought_prompt()
671        system_message = self.build_prompt()
672
673        # The model only enters the decision once a COT prompt is in play, so keep the
674        # provider lookup lazy: a plain prompt is single turn whatever the model is, and
675        # resolving a provider can fail.
676        provider = self.model_provider() if cot_prompt else None
677        strategy = chat_strategy_for_run(
678            cot_prompt=cot_prompt,
679            tuned_chat_strategy=provider.tuned_chat_strategy if provider else None,
680            reasoning_capable=provider.reasoning_capable if provider else False,
681        )
682        return get_chat_formatter(
683            strategy=strategy,
684            system_message=system_message,
685            user_input=input,
686            thinking_instructions=cot_prompt,
687        )
688
689    # create a run and task output
690    def generate_run(
691        self,
692        input: InputType,
693        input_source: DataSource | None,
694        run_output: RunOutput,
695        usage: Usage | None = None,
696        trace: list[ChatCompletionMessageParam] | None = None,
697        parent_task_run: TaskRun | None = None,
698    ) -> TaskRun:
699        output_str = (
700            json.dumps(run_output.output, ensure_ascii=False)
701            if isinstance(run_output.output, dict)
702            else run_output.output
703        )
704
705        output_source_type = (
706            DataSourceType.tool_call
707            if self.run_config.type == "mcp"
708            else DataSourceType.synthetic
709        )
710
711        new_output = TaskOutput(
712            output=output_str,
713            source=DataSource(
714                type=output_source_type,
715                properties=self._properties_for_task_output(),
716                run_config_id=self.base_adapter_config.task_run_config_id,
717                run_config=self.run_config,
718            ),
719        )
720
721        # Convert input and output to JSON strings if they aren't strings
722        input_str = (
723            input if isinstance(input, str) else json.dumps(input, ensure_ascii=False)
724        )
725
726        if input_source is None:
727            input_source = DataSource(
728                type=DataSourceType.human,
729                properties={"created_by": Config.shared().user_id},
730            )
731
732        parent_task_run_id: str | None = None
733        if parent_task_run is not None:
734            if parent_task_run.id is None:
735                raise ValueError(
736                    "parent_task_run must be persisted before using as parent: save the parent "
737                    "TaskRun (e.g. save_to_file()) so it has a stable id."
738                )
739            parent_task_run_id = parent_task_run.id
740
741        return TaskRun(
742            parent=self.task,
743            parent_task_run_id=parent_task_run_id,
744            input=input_str,
745            input_source=input_source,
746            output=new_output,
747            intermediate_outputs=run_output.intermediate_outputs,
748            tags=self.base_adapter_config.default_tags or [],
749            usage=usage,
750            trace=trace,
751            cumulative_usage=MessageUsage.from_trace(trace),
752        )
753
754    def _properties_for_task_output(self) -> Dict[str, str | int | float]:
755        match self.run_config.type:
756            case "mcp":
757                return {}
758            case "kiln_agent":
759                if not isinstance(self.run_config, KilnAgentRunConfigProperties):
760                    raise ValueError("Kiln agent run config is required")
761                run_config = self.run_config
762
763                props: Dict[str, str | int | float] = {}
764                props["adapter_name"] = self.adapter_name()
765                # Legacy properties where we save the run_config details into custom properties.
766                # These are now also be saved in the run_config field.
767                props["model_name"] = run_config.model_name
768                props["model_provider"] = run_config.model_provider_name
769                props["prompt_id"] = run_config.prompt_id
770                props["structured_output_mode"] = run_config.structured_output_mode
771                props["temperature"] = run_config.temperature
772                props["top_p"] = run_config.top_p
773
774                return props
775            case _:
776                raise_exhaustive_enum_error(self.run_config.type)
777
778    def update_run_config_unknown_structured_output_mode(self) -> None:
779        if self.run_config.type != "kiln_agent":
780            return
781        run_config = as_kiln_agent_run_config(self.run_config)
782        structured_output_mode = run_config.structured_output_mode
783
784        # Old datamodels didn't save the structured output mode. Some clients (tests, end users) might not set it.
785        # Look up our recommended mode from ml_model_list if we have one
786        if structured_output_mode == StructuredOutputMode.unknown:
787            new_run_config = run_config.model_copy(deep=True)
788            structured_output_mode = default_structured_output_mode_for_model_provider(
789                run_config.model_name,
790                run_config.model_provider_name,
791            )
792            new_run_config.structured_output_mode = structured_output_mode
793            self.run_config = new_run_config
794
795    async def available_tools(self) -> list[KilnToolInterface]:
796        if self.run_config.type != "kiln_agent":
797            return []
798        tool_config = as_kiln_agent_run_config(self.run_config).tools_config
799        if tool_config is None or tool_config.tools is None:
800            return []
801
802        return await assemble_unique_agent_tools(
803            self.task, tool_config.tools, self._resolve_skills()
804        )
805
806
807async def assemble_unique_agent_tools(
808    task: Task, tool_ids: list[str], skills: list[Skill]
809) -> list[KilnToolInterface]:
810    """Resolve a run config's tool IDs into tool instances, rejecting name collisions.
811
812    Names may repeat across a project, but within one run config every tool the
813    model sees must have a unique function name (and every attached skill a
814    unique skill name — skills are loaded by name through the single "skill"
815    tool). Shared by creation-time validation of run configs, which catches
816    collisions at selection time, and by the runtime adapter, which stays the
817    backstop for configs whose tools were renamed after creation (e.g. via
818    edit_kiln_task_tool) or that were created outside the API.
819    """
820    non_skill_tool_ids = [
821        tid for tid in tool_ids if not tid.startswith(SKILL_TOOL_ID_PREFIX)
822    ]
823
824    tools: list[KilnToolInterface] = [
825        tool_from_id(tool_id, task) for tool_id in non_skill_tool_ids
826    ]
827
828    if skills:
829        seen_names: set[str] = set()
830        for skill in skills:
831            if skill.name in seen_names:
832                raise ValueError(
833                    f"Duplicate skill name '{skill.name}'. Each skill must have a unique name."
834                )
835            seen_names.add(skill.name)
836        tools.append(SkillTool(f"{SKILL_TOOL_ID_PREFIX}_combined", skills))
837
838    tool_names = [await tool.name() for tool in tools]
839    if len(tool_names) != len(set(tool_names)):
840        duplicates = sorted({n for n in tool_names if tool_names.count(n) > 1})
841        message = (
842            f"Multiple selected tools share the same function name: {', '.join(duplicates)}. "
843            "Each tool must have a unique name. Either de-select the duplicate tools, or "
844            "modify their names to describe their unique purpose. The model will struggle "
845            "if tools do not have descriptive names and tool execution will be undefined."
846        )
847        if skills and "skill" in duplicates:
848            message += (
849                " The name 'skill' is reserved for the tool that loads skills, "
850                "so no other tool may use it while skills are selected."
851            )
852        raise ValueError(message)
853
854    return tools
855
856
857class OpenAIStreamResult:
858    """Async-iterable wrapper around the OpenAI streaming flow.
859
860    Yields ``ModelResponseStream`` chunks.  After iteration the resulting
861    ``TaskRun`` is available via the ``.task_run`` property.
862
863    When return_on_tool_call=True and the model requests tool calls, the stream
864    will stop and ``task_run.is_toolcall_pending`` will be True.
865    """
866
867    def __init__(
868        self,
869        adapter: BaseAdapter,
870        input: InputType,
871        input_source: DataSource | None,
872        prior_trace: list[ChatCompletionMessageParam] | None,
873        parent_task_run: TaskRun | None = None,
874    ) -> None:
875        self._adapter = adapter
876        self._input = input
877        self._input_source = input_source
878        self._prior_trace = prior_trace
879        self._parent_task_run = parent_task_run
880        self._task_run: TaskRun | None = None
881
882    @property
883    def task_run(self) -> TaskRun:
884        if self._task_run is None:
885            raise RuntimeError(
886                "Stream has not been fully consumed yet. "
887                "Iterate over the stream before accessing .task_run"
888            )
889        return self._task_run
890
891    async def __aiter__(self) -> AsyncIterator[ModelResponseStream]:
892        self._task_run = None
893        async with mcp_session_scope():
894            adapter_stream = self._adapter._prepare_stream(
895                self._input, self._prior_trace
896            )
897
898            async for event in adapter_stream:
899                if isinstance(event, ModelResponseStream):
900                    yield event
901
902            self._task_run = self._adapter._finalize_stream(
903                adapter_stream, self._input, self._input_source, self._parent_task_run
904            )
905
906
907class AiSdkStreamResult:
908    """Async-iterable wrapper around the AI SDK streaming flow.
909
910    Yields ``AiSdkStreamEvent`` instances.  After iteration the resulting
911    ``TaskRun`` is available via the ``.task_run`` property.
912
913    When return_on_tool_call=True and the model requests tool calls, the FINISH
914    event will have finishReason: "tool-calls" and ``task_run.is_toolcall_pending``
915    will be True.
916    """
917
918    def __init__(
919        self,
920        adapter: BaseAdapter,
921        input: InputType,
922        input_source: DataSource | None,
923        prior_trace: list[ChatCompletionMessageParam] | None,
924        parent_task_run: TaskRun | None = None,
925    ) -> None:
926        self._adapter = adapter
927        self._input = input
928        self._input_source = input_source
929        self._prior_trace = prior_trace
930        self._parent_task_run = parent_task_run
931        self._task_run: TaskRun | None = None
932
933    @property
934    def task_run(self) -> TaskRun:
935        if self._task_run is None:
936            raise RuntimeError(
937                "Stream has not been fully consumed yet. "
938                "Iterate over the stream before accessing .task_run"
939            )
940        return self._task_run
941
942    async def __aiter__(self) -> AsyncIterator[AiSdkStreamEvent]:
943        self._task_run = None
944        async with mcp_session_scope():
945            adapter_stream = self._adapter._prepare_stream(
946                self._input, self._prior_trace
947            )
948
949            message_id = f"msg-{uuid.uuid4().hex}"
950            converter = AiSdkStreamConverter()
951
952            yield StartEvent(messageId=message_id)
953            yield StartStepEvent()
954
955            last_event_was_tool_call = False
956            async for event in adapter_stream:
957                if isinstance(event, ModelResponseStream):
958                    if last_event_was_tool_call:
959                        converter.reset_for_next_step()
960                        last_event_was_tool_call = False
961                    for ai_event in converter.convert_chunk(event):
962                        yield ai_event
963                elif isinstance(event, ToolCallEvent):
964                    last_event_was_tool_call = True
965                    for ai_event in converter.convert_tool_event(event):
966                        yield ai_event
967
968            for ai_event in converter.close_open_blocks():
969                yield ai_event
970
971            yield FinishStepEvent()
972
973            self._task_run = self._adapter._finalize_stream(
974                adapter_stream, self._input, self._input_source, self._parent_task_run
975            )
976
977            if self._task_run.is_toolcall_pending:
978                yield FinishEvent(
979                    messageMetadata=FinishMessageMetadata(finishReason="tool-calls"),
980                )
981            else:
982                for ai_event in converter.finalize():
983                    yield ai_event