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