kiln_ai.adapters.model_adapters.litellm_adapter

   1import asyncio
   2import copy
   3import json
   4import logging
   5import time
   6from dataclasses import dataclass
   7from typing import Any, Dict, List, Tuple
   8
   9import litellm
  10from litellm.types.utils import (
  11    ChatCompletionMessageToolCall,
  12    ChoiceLogprobs,
  13    Choices,
  14    ModelResponse,
  15)
  16from litellm.types.utils import Message as LiteLLMMessage
  17from litellm.types.utils import Usage as LiteLlmUsage
  18from openai.types.chat.chat_completion_message_tool_call_param import (
  19    ChatCompletionMessageToolCallParam,
  20)
  21
  22import kiln_ai.datamodel as datamodel
  23from kiln_ai.adapters.chat import ChatCompletionMessageIncludingLiteLLM
  24from kiln_ai.adapters.chat.chat_formatter import chat_message_to_dict
  25from kiln_ai.adapters.ml_model_list import (
  26    KilnModelProvider,
  27    ModelProviderName,
  28    StructuredOutputMode,
  29)
  30from kiln_ai.adapters.model_adapters.adapter_stream import (
  31    EMPTY_RESPONSE_ERROR_MESSAGE,
  32    AdapterStream,
  33    raise_for_empty_model_response,
  34)
  35from kiln_ai.adapters.model_adapters.base_adapter import (
  36    AdapterConfig,
  37    BaseAdapter,
  38    MessageUsage,
  39    RunOutput,
  40    Usage,
  41)
  42from kiln_ai.adapters.model_adapters.litellm_config import LiteLlmConfig
  43from kiln_ai.datamodel.datamodel_enums import InputType
  44from kiln_ai.datamodel.json_schema import (
  45    close_object_schemas,
  46    strip_numeric_bounds,
  47    validate_schema_with_value_error,
  48)
  49from kiln_ai.datamodel.run_config import (
  50    KilnAgentRunConfigProperties,
  51    as_kiln_agent_run_config,
  52)
  53from kiln_ai.tools.base_tool import (
  54    KilnToolInterface,
  55    ToolCallContext,
  56    ToolCallDefinition,
  57)
  58from kiln_ai.tools.kiln_task_tool import KilnTaskToolResult
  59from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
  60from kiln_ai.utils.litellm import get_litellm_provider_info
  61from kiln_ai.utils.open_ai_types import (
  62    ChatCompletionAssistantMessageParamWrapper,
  63    ChatCompletionMessageParam,
  64    ChatCompletionToolMessageParamWrapper,
  65    sanitize_messages_for_provider,
  66)
  67
  68MAX_CALLS_PER_TURN = 10
  69MAX_TOOL_CALLS_PER_TURN = 30
  70
  71logger = logging.getLogger(__name__)
  72
  73
  74def _validate_unmanaged_tools(tools: list[KilnToolInterface]) -> None:
  75    for i, tool in enumerate(tools):
  76        if not isinstance(tool, KilnToolInterface):
  77            raise TypeError(
  78                f"unmanaged_tools[{i}] must be a KilnToolInterface instance, got {type(tool).__name__}"
  79            )
  80
  81
  82@dataclass
  83class ModelTurnResult:
  84    assistant_message: str
  85    all_messages: list[ChatCompletionMessageIncludingLiteLLM]
  86    model_response: ModelResponse | None
  87    model_choice: Choices | None
  88    usage: Usage
  89    interrupted_by_tool_calls: list[ChatCompletionMessageToolCall] | None = None
  90    message_latency: dict[int, int] | None = None
  91    message_usage: dict[int, MessageUsage] | None = None
  92
  93
  94class LiteLlmAdapter(BaseAdapter):
  95    def __init__(
  96        self,
  97        config: LiteLlmConfig,
  98        kiln_task: datamodel.Task,
  99        base_adapter_config: AdapterConfig | None = None,
 100    ):
 101        if not isinstance(config.run_config_properties, KilnAgentRunConfigProperties):
 102            raise ValueError("LiteLlmAdapter requires KilnAgentRunConfigProperties")
 103        self.config = config
 104        self._additional_body_options = config.additional_body_options
 105        self._api_base = config.base_url
 106        self._headers = config.default_headers
 107        self._litellm_model_id: str | None = None
 108        self._cached_available_tools: list[KilnToolInterface] | None = None
 109
 110        super().__init__(
 111            task=kiln_task,
 112            run_config=config.run_config_properties,
 113            config=base_adapter_config,
 114        )
 115
 116        unmanaged_tools = self.base_adapter_config.unmanaged_tools
 117        if unmanaged_tools:
 118            _validate_unmanaged_tools(unmanaged_tools)
 119
 120    async def _run_model_turn(
 121        self,
 122        provider: KilnModelProvider,
 123        messages: list[ChatCompletionMessageIncludingLiteLLM],
 124        top_logprobs: int | None,
 125        skip_response_format: bool,
 126    ) -> ModelTurnResult:
 127        """
 128        Call the model for a single top level turn: from user message to agent message.
 129
 130        It may make handle iterations of tool calls between the user/agent message if needed.
 131
 132        `messages` is the caller-owned list and is mutated in place (append/extend
 133        only — never rebound) so the partial trace survives any exception
 134        escaping this method.
 135        """
 136
 137        usage = Usage()
 138        tool_calls_count = 0
 139        # Per-LLM-call latency / usage, keyed by index in the messages list.
 140        # Kept separate because we don't own the LiteLLM message objects.
 141        message_latency: dict[int, int] = {}
 142        message_usage: dict[int, MessageUsage] = {}
 143
 144        while tool_calls_count < MAX_TOOL_CALLS_PER_TURN:
 145            # Build completion kwargs for tool calls
 146            completion_kwargs = await self.build_completion_kwargs(
 147                provider,
 148                # Pass a copy, as acompletion mutates objects and breaks types.
 149                copy.deepcopy(messages),
 150                top_logprobs,
 151                skip_response_format,
 152            )
 153
 154            # Make the completion call (timed)
 155            start = time.monotonic()
 156            model_response, response_choice = await self.acompletion_checking_response(
 157                **completion_kwargs
 158            )
 159            call_latency_ms = int((time.monotonic() - start) * 1000)
 160
 161            # count the usage
 162            call_usage = self.usage_from_response(model_response)
 163            usage += call_usage
 164            usage.total_llm_latency_ms = (
 165                usage.total_llm_latency_ms or 0
 166            ) + call_latency_ms
 167
 168            # Extract content and tool calls
 169            if not hasattr(response_choice, "message"):
 170                raise ValueError("Response choice has no message")
 171            content = response_choice.message.content
 172            tool_calls = response_choice.message.tool_calls
 173            if not content and not tool_calls:
 174                raise_for_empty_model_response(response_choice)
 175
 176            # Add message to messages, so it can be used in the next turn
 177            messages.append(response_choice.message)
 178            message_latency[len(messages) - 1] = call_latency_ms
 179            message_usage[len(messages) - 1] = call_usage
 180
 181            # Process tool calls if any
 182            if tool_calls and len(tool_calls) > 0:
 183                # check if we should return control to caller
 184                if self.base_adapter_config.return_on_tool_call:
 185                    # filter out task_response tool (task_response tools are internal)
 186                    standard_tool_calls = [
 187                        tc for tc in tool_calls if tc.function.name != "task_response"
 188                    ]
 189                    has_task_response = any(
 190                        tc.function.name == "task_response" for tc in tool_calls
 191                    )
 192                    if standard_tool_calls and not has_task_response:
 193                        return ModelTurnResult(
 194                            # we don't have any content, we are waiting for toolcall output to come back from client
 195                            assistant_message="",
 196                            all_messages=messages,
 197                            model_response=model_response,
 198                            model_choice=response_choice,
 199                            usage=usage,
 200                            interrupted_by_tool_calls=standard_tool_calls,
 201                            message_latency=message_latency,
 202                            message_usage=message_usage,
 203                        )
 204
 205                # otherwise: process tool calls internally until final output
 206                (
 207                    assistant_message_from_toolcall,
 208                    tool_call_messages,
 209                ) = await self.process_tool_calls(tool_calls)
 210
 211                # Add tool call results to messages
 212                messages.extend(tool_call_messages)
 213
 214                # If task_response tool was called, we're done
 215                if assistant_message_from_toolcall is not None:
 216                    return ModelTurnResult(
 217                        assistant_message=assistant_message_from_toolcall,
 218                        all_messages=messages,
 219                        model_response=model_response,
 220                        model_choice=response_choice,
 221                        usage=usage,
 222                        message_latency=message_latency,
 223                        message_usage=message_usage,
 224                    )
 225
 226                # If there were tool calls, increment counter and continue
 227                if tool_call_messages:
 228                    tool_calls_count += 1
 229                    continue
 230
 231            # If no tool calls, return the content as final output
 232            if content:
 233                return ModelTurnResult(
 234                    assistant_message=content,
 235                    all_messages=messages,
 236                    model_response=model_response,
 237                    model_choice=response_choice,
 238                    usage=usage,
 239                    message_latency=message_latency,
 240                    message_usage=message_usage,
 241                )
 242
 243            # If we get here with no content and no tool calls, break
 244            raise RuntimeError(
 245                "Model returned neither content nor tool calls. It must return at least one of these."
 246            )
 247
 248        raise RuntimeError(
 249            f"Too many tool calls ({tool_calls_count}). Stopping iteration to avoid using too many tokens."
 250        )
 251
 252    async def _run(
 253        self,
 254        input: InputType,
 255        trace_ref: list[ChatCompletionMessageParam],
 256        prior_trace: list[ChatCompletionMessageParam] | None = None,
 257    ) -> tuple[RunOutput, Usage | None]:
 258        usage = Usage()
 259
 260        provider = self.model_provider()
 261        if not provider.model_id:
 262            raise ValueError("Model ID is required for OpenAI compatible models")
 263
 264        # build_chat_formatter returns MultiturnFormatter when prior_trace is set, else prompt-based formatter
 265        chat_formatter = self.build_chat_formatter(input, prior_trace)
 266        # `trace_ref` is typed as `ChatCompletionMessageParam` for the
 267        # caller's API surface, but internally we store LiteLLM `Message`
 268        # objects transiently (converted by `all_messages_to_trace` before
 269        # export). We widen the type alias once here so the internal
 270        # mutations don't each need their own suppression.
 271        messages_internal: list[ChatCompletionMessageIncludingLiteLLM] = trace_ref  # type: ignore[assignment]
 272        messages_internal.extend(copy.deepcopy(chat_formatter.initial_messages()))
 273
 274        prior_output: str | None = None
 275        final_choice: Choices | None = None
 276        turns = 0
 277        message_latency: dict[int, int] = {}
 278        message_usage: dict[int, MessageUsage] = {}
 279
 280        # Same loop for both fresh runs and prior_trace continuation.
 281        # _run_model_turn has its own internal loop for tool calls (model calls tool -> we run it -> model continues).
 282        while True:
 283            turns += 1
 284            if turns > MAX_CALLS_PER_TURN:
 285                raise RuntimeError(
 286                    f"Too many turns ({turns}). Stopping iteration to avoid using too many tokens."
 287                )
 288
 289            turn = chat_formatter.next_turn(prior_output)
 290            if turn is None:
 291                # No next turn, we're done
 292                break
 293
 294            # Add messages from the turn to chat history
 295            for message in turn.messages:
 296                if message.content is None:
 297                    raise ValueError("Empty message content isn't allowed")
 298                messages_internal.append(chat_message_to_dict(message))  # type: ignore
 299
 300            skip_response_format = not turn.final_call
 301            turn_result = await self._run_model_turn(
 302                provider,
 303                messages_internal,
 304                self.base_adapter_config.top_logprobs if turn.final_call else None,
 305                skip_response_format,
 306            )
 307
 308            usage += turn_result.usage
 309            if turn_result.message_latency:
 310                message_latency.update(turn_result.message_latency)
 311            if turn_result.message_usage:
 312                message_usage.update(turn_result.message_usage)
 313
 314            prior_output = turn_result.assistant_message
 315            # Update messages_internal in place from turn_result so trace_ref
 316            # stays in sync (required by the base adapter error-wrapping contract).
 317            messages_internal[:] = turn_result.all_messages
 318            final_choice = turn_result.model_choice
 319
 320            # Check if we were interrupted by tool calls
 321            if turn_result.interrupted_by_tool_calls:
 322                trace = self.all_messages_to_trace(
 323                    messages_internal, message_latency, message_usage
 324                )
 325                intermediate_outputs = chat_formatter.intermediate_outputs()
 326                output = RunOutput(
 327                    output=prior_output or "",
 328                    intermediate_outputs=intermediate_outputs,
 329                    output_logprobs=None,
 330                    trace=trace,
 331                )
 332                return output, usage
 333
 334            if not prior_output:
 335                raise RuntimeError("No assistant message/output returned from model")
 336
 337        logprobs = self._extract_and_validate_logprobs(final_choice)
 338
 339        # Save COT/reasoning if it exists. May be a message, or may be parsed by LiteLLM (or openrouter, or anyone upstream)
 340        intermediate_outputs = chat_formatter.intermediate_outputs()
 341        self._extract_reasoning_to_intermediate_outputs(
 342            final_choice, intermediate_outputs
 343        )
 344
 345        if not isinstance(prior_output, str):
 346            raise RuntimeError(f"assistant message is not a string: {prior_output}")
 347
 348        trace = self.all_messages_to_trace(
 349            messages_internal, message_latency, message_usage
 350        )
 351        output = RunOutput(
 352            output=prior_output,
 353            intermediate_outputs=intermediate_outputs,
 354            output_logprobs=logprobs,
 355            trace=trace,
 356        )
 357
 358        return output, usage
 359
 360    def _create_run_stream(
 361        self,
 362        input: InputType,
 363        prior_trace: list[ChatCompletionMessageParam] | None = None,
 364    ) -> AdapterStream:
 365        provider = self.model_provider()
 366        if not provider.model_id:
 367            raise ValueError("Model ID is required for OpenAI compatible models")
 368
 369        chat_formatter = self.build_chat_formatter(input, prior_trace)
 370        initial_messages: list[ChatCompletionMessageIncludingLiteLLM] = copy.deepcopy(
 371            chat_formatter.initial_messages()
 372        )
 373
 374        return AdapterStream(
 375            adapter=self,
 376            provider=provider,
 377            chat_formatter=chat_formatter,
 378            initial_messages=initial_messages,
 379            top_logprobs=self.base_adapter_config.top_logprobs,
 380        )
 381
 382    def _extract_and_validate_logprobs(
 383        self, final_choice: Choices | None
 384    ) -> ChoiceLogprobs | None:
 385        """
 386        Extract logprobs from the final choice and validate they exist if required.
 387        """
 388        logprobs = None
 389        if (
 390            final_choice is not None
 391            and hasattr(final_choice, "logprobs")
 392            and isinstance(final_choice.logprobs, ChoiceLogprobs)
 393        ):
 394            logprobs = final_choice.logprobs
 395
 396        # Check logprobs worked, if required
 397        if self.base_adapter_config.top_logprobs is not None and logprobs is None:
 398            raise RuntimeError("Logprobs were required, but no logprobs were returned.")
 399
 400        return logprobs
 401
 402    def _extract_reasoning_to_intermediate_outputs(
 403        self, final_choice: Choices | None, intermediate_outputs: Dict[str, Any]
 404    ) -> None:
 405        """Extract reasoning content from model choice and add to intermediate outputs if present."""
 406        if (
 407            final_choice is not None
 408            and hasattr(final_choice, "message")
 409            and hasattr(final_choice.message, "reasoning_content")
 410        ):
 411            reasoning_content = final_choice.message.reasoning_content
 412            if reasoning_content is not None:
 413                stripped_reasoning_content = reasoning_content.strip()
 414                if len(stripped_reasoning_content) > 0:
 415                    intermediate_outputs["reasoning"] = stripped_reasoning_content
 416
 417    async def acompletion_checking_response(
 418        self, **kwargs: Any
 419    ) -> Tuple[ModelResponse, Choices]:
 420        response = await litellm.acompletion(**kwargs)
 421
 422        if (
 423            not isinstance(response, ModelResponse)
 424            or not response.choices
 425            or len(response.choices) == 0
 426            or not isinstance(response.choices[0], Choices)
 427        ):
 428            raise RuntimeError(
 429                f"Expected ModelResponse with Choices, got {type(response)}."
 430            )
 431        return response, response.choices[0]
 432
 433    def adapter_name(self) -> str:
 434        return "kiln_openai_compatible_adapter"
 435
 436    async def response_format_options(self) -> dict[str, Any]:
 437        # Unstructured if task isn't structured
 438        if not self.has_structured_output():
 439            return {}
 440
 441        run_config = as_kiln_agent_run_config(self.run_config)
 442        structured_output_mode: StructuredOutputMode = run_config.structured_output_mode
 443
 444        match structured_output_mode:
 445            case StructuredOutputMode.json_mode:
 446                return {"response_format": {"type": "json_object"}}
 447            case StructuredOutputMode.json_schema:
 448                return self.json_schema_response_format()
 449            case StructuredOutputMode.function_calling_weak:
 450                return self.tool_call_params(strict=False)
 451            case StructuredOutputMode.function_calling:
 452                return self.tool_call_params(strict=True)
 453            case StructuredOutputMode.json_instructions:
 454                # JSON instructions dynamically injected in prompt, not the API response format. Do not ask for json_object (see option below).
 455                return {}
 456            case StructuredOutputMode.json_custom_instructions:
 457                # JSON instructions statically injected in system prompt, not the API response format. Do not ask for json_object (see option above).
 458                return {}
 459            case StructuredOutputMode.json_instruction_and_object:
 460                # We set response_format to json_object and also set json instructions in the prompt
 461                return {"response_format": {"type": "json_object"}}
 462            case StructuredOutputMode.default:
 463                provider_name = run_config.model_provider_name
 464                if provider_name == ModelProviderName.ollama:
 465                    # Ollama added json_schema to all models: https://ollama.com/blog/structured-outputs
 466                    return self.json_schema_response_format()
 467                elif provider_name == ModelProviderName.docker_model_runner:
 468                    # Docker Model Runner uses OpenAI-compatible API with JSON schema support
 469                    return self.json_schema_response_format()
 470                else:
 471                    # Default to function calling -- it's older than the other modes. Higher compatibility.
 472                    # Strict isn't widely supported yet, so we don't use it by default unless it's OpenAI.
 473                    strict = provider_name == ModelProviderName.openai
 474                    return self.tool_call_params(strict=strict)
 475            case StructuredOutputMode.unknown:
 476                # See above, but this case should never happen.
 477                raise ValueError("Structured output mode is unknown.")
 478            case _:
 479                raise_exhaustive_enum_error(structured_output_mode)  # type: ignore[arg-type]
 480
 481    def json_schema_response_format(self) -> dict[str, Any]:
 482        output_schema = self.task.output_schema()
 483        if output_schema is None:
 484            raise ValueError(
 485                "Invalid output schema for this task. Cannot use JSON schema response format."
 486            )
 487        output_schema = close_object_schemas(output_schema, strict=True)
 488        # Strip numeric bounds (min/max/etc.) from integer/number nodes for the
 489        # json_schema wire format. Some providers (e.g. Claude via OpenRouter,
 490        # which maps onto Anthropic's newer output_config.format.schema API)
 491        # reject numeric bounds on integer/number types and return HTTP 400.
 492        # The valid ranges are still enforced by the prompt + post-hoc
 493        # validation, so this only affects the schema sent over the wire.
 494        output_schema = strip_numeric_bounds(output_schema)
 495        return {
 496            "response_format": {
 497                "type": "json_schema",
 498                "json_schema": {
 499                    "name": "task_response",
 500                    "schema": output_schema,
 501                },
 502            }
 503        }
 504
 505    def tool_call_params(self, strict: bool) -> dict[str, Any]:
 506        # Add additional_properties: false to the schema (OpenAI requires this for some models)
 507        output_schema = self.task.output_schema()
 508        if not isinstance(output_schema, dict):
 509            raise ValueError(
 510                "Invalid output schema for this task. Can not use tool calls."
 511            )
 512        output_schema = close_object_schemas(output_schema, strict=strict)
 513
 514        function_params = {
 515            "name": "task_response",
 516            "parameters": output_schema,
 517        }
 518        # This should be on, but we allow setting function_calling_weak for APIs that don't support it.
 519        if strict:
 520            function_params["strict"] = True
 521
 522        return {
 523            "tools": [
 524                {
 525                    "type": "function",
 526                    "function": function_params,
 527                }
 528            ],
 529            "tool_choice": {
 530                "type": "function",
 531                "function": {"name": "task_response"},
 532            },
 533        }
 534
 535    def build_extra_body(self, provider: KilnModelProvider) -> dict[str, Any]:
 536        # Don't love having this logic here. But it's worth the usability improvement
 537        # so better to keep it than exclude it. Should figure out how I want to isolate
 538        # this sort of logic so it's config driven and can be overridden
 539        extra_body: dict[str, Any] = {}
 540        provider_options = {}
 541
 542        run_config = as_kiln_agent_run_config(self.run_config)
 543        # For legacy config 'thinking_level' is not set, default to provider's default
 544        if "thinking_level" in run_config.model_fields_set:
 545            thinking_level = run_config.thinking_level
 546        else:
 547            thinking_level = provider.default_thinking_level
 548
 549        # Skip if provider doesn't support thinking levels (stale configs may still have one set)
 550        if (
 551            thinking_level is not None
 552            and provider.available_thinking_levels is not None
 553        ):
 554            # Anthropic models in OpenRouter uses reasoning object. See https://openrouter.ai/docs/use-cases/reasoning-tokens
 555            if (
 556                provider.name == ModelProviderName.openrouter
 557                and provider.openrouter_reasoning_object
 558            ):
 559                extra_body["reasoning"] = {"effort": thinking_level}
 560            elif (
 561                provider.name == ModelProviderName.anthropic
 562                and thinking_level == "none"
 563            ):
 564                # Anthropic's native API has no reasoning_effort="none"; passing it makes
 565                # litellm map thinking to None and then crash. Omitting reasoning_effort
 566                # disables extended thinking, which also frees temperature from the
 567                # temperature=1 requirement that applies whenever thinking is enabled.
 568                pass
 569            else:
 570                extra_body["reasoning_effort"] = thinking_level
 571                # Opus 4.7/4.8 default thinking display to "omitted", returning empty
 572                # thinking text. Request the summary so reasoning is surfaced. litellm
 573                # still maps reasoning_effort to output_config.effort; this only adds the
 574                # display to the adaptive thinking object.
 575                if (
 576                    provider.name == ModelProviderName.anthropic
 577                    and provider.anthropic_summarized_thinking
 578                ):
 579                    extra_body["thinking"] = {
 580                        "type": "adaptive",
 581                        "display": "summarized",
 582                    }
 583
 584        if provider.require_openrouter_reasoning:
 585            # https://openrouter.ai/docs/use-cases/reasoning-tokens
 586            extra_body["reasoning"] = {
 587                "exclude": False,
 588            }
 589
 590        if provider.gemini_reasoning_enabled:
 591            extra_body["reasoning"] = {
 592                "enabled": True,
 593            }
 594
 595        if provider.name == ModelProviderName.openrouter:
 596            # Ask OpenRouter to include usage in the response (cost)
 597            extra_body["usage"] = {"include": True}
 598
 599            # Set a default provider order for more deterministic routing.
 600            # OpenRouter will ignore providers that don't support the model.
 601            # Special cases below (like R1) can override this order.
 602            # allow_fallbacks is true by default, but we can override it here.
 603            provider_options["order"] = [
 604                "fireworks",
 605                "parasail",
 606                "together",
 607                "deepinfra",
 608                "novita",
 609                "groq",
 610                "amazon-bedrock",
 611                "azure",
 612                "nebius",
 613            ]
 614
 615        if provider.anthropic_extended_thinking and "thinking" not in extra_body:
 616            extra_body["thinking"] = {"type": "enabled", "budget_tokens": 4000}
 617
 618        if provider.r1_openrouter_options:
 619            # Require providers that support the reasoning parameter
 620            provider_options["require_parameters"] = True
 621            # Prefer R1 providers with reasonable perf/quants
 622            provider_options["order"] = ["fireworks", "together"]
 623            # R1 providers with unreasonable quants
 624            provider_options["ignore"] = ["deepinfra"]
 625
 626        # Only set of this request is to get logprobs.
 627        if (
 628            provider.logprobs_openrouter_options
 629            and self.base_adapter_config.top_logprobs is not None
 630        ):
 631            # Don't let OpenRouter choose a provider that doesn't support logprobs.
 632            provider_options["require_parameters"] = True
 633            # DeepInfra silently fails to return logprobs consistently.
 634            provider_options["ignore"] = ["deepinfra"]
 635
 636        if provider.openrouter_skip_required_parameters:
 637            # Oddball case, R1 14/8/1.5B fail with this param, even though they support thinking params.
 638            provider_options["require_parameters"] = False
 639
 640        # Siliconflow uses a bool flag for thinking, for some models
 641        if provider.siliconflow_enable_thinking is not None:
 642            extra_body["enable_thinking"] = provider.siliconflow_enable_thinking
 643
 644        if len(provider_options) > 0:
 645            extra_body["provider"] = provider_options
 646
 647        return extra_body
 648
 649    def litellm_model_id(self) -> str:
 650        # The model ID is an interesting combination of format and url endpoint.
 651        # It specifics the provider URL/host, but this is overridden if you manually set an api url
 652        if self._litellm_model_id:
 653            return self._litellm_model_id
 654
 655        litellm_provider_info = get_litellm_provider_info(self.model_provider())
 656        if litellm_provider_info.is_custom and self._api_base is None:
 657            raise ValueError(
 658                "Explicit Base URL is required for OpenAI compatible APIs (custom models, ollama, fine tunes, and custom registry models)"
 659            )
 660
 661        self._litellm_model_id = litellm_provider_info.litellm_model_id
 662        return self._litellm_model_id
 663
 664    def _allowed_openai_params_for_completion_kwargs(
 665        self, completion_kwargs: dict[str, Any]
 666    ) -> list[str]:
 667        """
 668        LiteLLM drops params it thinks are not supported by the model when drop_params is True. Sometimes it is wrong
 669        and we know it is supported, so we whitelist them here and pass that as an allowed_openai_params parameter.
 670        """
 671        # callers could have set allowed_openai_params in the additional_body_options, so we need to check for that
 672        explicit_allowed_params: Any | list = completion_kwargs.get(
 673            "allowed_openai_params", []
 674        )
 675        if not isinstance(explicit_allowed_params, list):
 676            raise ValueError(
 677                f"Unexpected allowed_openai_params format: {explicit_allowed_params} - expected list, got {type(explicit_allowed_params)}"
 678            )
 679        explicit_allowed_params_validated = [
 680            param for param in explicit_allowed_params if isinstance(param, str)
 681        ]
 682        invalid_count = len(explicit_allowed_params) - len(
 683            explicit_allowed_params_validated
 684        )
 685        if invalid_count > 0:
 686            raise ValueError(
 687                f"Unexpected allowed_openai_params format: {explicit_allowed_params} - {invalid_count} items are not strings"
 688            )
 689
 690        # these are our own logic
 691        automatic_allowed_params: list[str] = []
 692        if "tools" in completion_kwargs:
 693            automatic_allowed_params.append("tools")
 694        if "tool_choice" in completion_kwargs:
 695            automatic_allowed_params.append("tool_choice")
 696
 697        return list(set(explicit_allowed_params_validated + automatic_allowed_params))
 698
 699    async def build_completion_kwargs(
 700        self,
 701        provider: KilnModelProvider,
 702        messages: list[ChatCompletionMessageIncludingLiteLLM],
 703        top_logprobs: int | None,
 704        skip_response_format: bool = False,
 705    ) -> dict[str, Any]:
 706        run_config = as_kiln_agent_run_config(self.run_config)
 707        extra_body = self.build_extra_body(provider)
 708
 709        # Merge all parameters into a single kwargs dict for litellm
 710        completion_kwargs = {
 711            "model": self.litellm_model_id(),
 712            "messages": messages,
 713            "api_base": self._api_base,
 714            "headers": self._headers,
 715            "temperature": run_config.temperature,
 716            "top_p": run_config.top_p,
 717            # This drops params that are not supported by the model. Only openai params like top_p, temperature -- not litellm params like model, etc.
 718            # Not all models and providers support all openai params (for example, o3 doesn't support top_p)
 719            # Better to ignore them than to fail the model call.
 720            # https://docs.litellm.ai/docs/completion/input
 721            "drop_params": True,
 722            **extra_body,
 723            **self._additional_body_options,
 724        }
 725
 726        if self.base_adapter_config.automatic_prompt_caching:
 727            # Mark the last message for cache control. Litellm's AnthropicCacheControlHook
 728            # handles provider-specific injection. Providers auto-cache matching prefixes,
 729            # so marking the last message is sufficient for multi-turn conversations.
 730            completion_kwargs["cache_control_injection_points"] = [
 731                {"location": "message", "index": -1}
 732            ]
 733
 734        tool_calls = await self.litellm_tools()
 735        has_tools = len(tool_calls) > 0
 736        if has_tools:
 737            completion_kwargs["tools"] = tool_calls
 738            completion_kwargs["tool_choice"] = "auto"
 739
 740        # Special condition for Claude Opus 4.1 and Sonnet 4.5, where we can only specify top_p or temp, not both.
 741        # Remove default values (1.0) prioritizing anything the user customized, then error with helpful message if they are both custom.
 742        if provider.temp_top_p_exclusive:
 743            if "top_p" in completion_kwargs and completion_kwargs["top_p"] == 1.0:
 744                del completion_kwargs["top_p"]
 745            if (
 746                "temperature" in completion_kwargs
 747                and completion_kwargs["temperature"] == 1.0
 748            ):
 749                del completion_kwargs["temperature"]
 750            if "top_p" in completion_kwargs and "temperature" in completion_kwargs:
 751                raise ValueError(
 752                    "top_p and temperature can not both have custom values for this model. This is a restriction from the model provider. Please set only one of them to a custom value (not 1.0)."
 753                )
 754
 755        if not skip_response_format:
 756            # Response format: json_schema, json_instructions, json_mode, function_calling, etc
 757            response_format_options = await self.response_format_options()
 758
 759            # Check for a conflict between tools and response format using tools
 760            # We could reconsider this. Model could be able to choose between a final answer or a tool call on any turn. However, good models for tools tend to also support json_schea, so do we need to support both? If we do, merge them, and consider auto vs forced when merging (only forced for final, auto for merged).
 761            if has_tools and "tools" in response_format_options:
 762                raise ValueError(
 763                    "Function calling/tools can't be used as the JSON response format if you're also using tools. Please select a different structured output mode."
 764                )
 765
 766            completion_kwargs.update(response_format_options)
 767
 768        if top_logprobs is not None:
 769            completion_kwargs["logprobs"] = True
 770            completion_kwargs["top_logprobs"] = top_logprobs
 771
 772        # any params listed in this list will be passed to the model regardless of LiteLLM's own validation
 773        allowed_openai_params = self._allowed_openai_params_for_completion_kwargs(
 774            completion_kwargs
 775        )
 776        if len(allowed_openai_params) > 0:
 777            completion_kwargs["allowed_openai_params"] = allowed_openai_params
 778
 779        completion_kwargs["messages"] = sanitize_messages_for_provider(messages)
 780
 781        return completion_kwargs
 782
 783    def usage_from_response(self, response: ModelResponse) -> MessageUsage:
 784        litellm_usage = response.get("usage", None)
 785
 786        # LiteLLM isn't consistent in how it returns the cost.
 787        cost = response._hidden_params.get("response_cost", None)
 788        if cost is None and litellm_usage:
 789            cost = litellm_usage.get("cost", None)
 790
 791        usage = MessageUsage()
 792
 793        if not litellm_usage and not cost:
 794            return usage
 795
 796        if litellm_usage and isinstance(litellm_usage, LiteLlmUsage):
 797            usage.input_tokens = litellm_usage.get("prompt_tokens", None)
 798            usage.output_tokens = litellm_usage.get("completion_tokens", None)
 799            usage.total_tokens = litellm_usage.get("total_tokens", None)
 800            prompt_details = litellm_usage.get("prompt_tokens_details", None)
 801            if prompt_details and hasattr(prompt_details, "cached_tokens"):
 802                usage.cached_tokens = prompt_details.cached_tokens
 803            elif prompt_details:
 804                logger.warning(
 805                    f"prompt_tokens_details has unexpected type {type(prompt_details)}, cached_tokens not extracted"
 806                )
 807        else:
 808            logger.warning(
 809                f"Unexpected usage format from litellm: {litellm_usage}. Expected Usage object, got {type(litellm_usage)}"
 810            )
 811
 812        if isinstance(cost, float):
 813            usage.cost = cost
 814        elif cost is not None:
 815            # None is allowed, but no other types are expected
 816            logger.warning(
 817                f"Unexpected cost format from litellm: {cost}. Expected float, got {type(cost)}"
 818            )
 819
 820        return usage
 821
 822    async def cached_available_tools(self) -> list[KilnToolInterface]:
 823        if self._cached_available_tools is None:
 824            self._cached_available_tools = await self.available_tools()
 825        return self._cached_available_tools
 826
 827    async def _tools_for_execution(self) -> list[KilnToolInterface]:
 828        """Registry-resolved tools plus :attr:`AdapterConfig.unmanaged_tools` (same order as ``litellm_tools``)."""
 829        registry = await self.cached_available_tools()
 830        unmanaged = self.base_adapter_config.unmanaged_tools or []
 831        return registry + unmanaged
 832
 833    async def litellm_tools(self) -> list[ToolCallDefinition]:
 834        available_tools = await self.cached_available_tools()
 835
 836        registry_defs = [await tool.toolcall_definition() for tool in available_tools]
 837        unmanaged = self.base_adapter_config.unmanaged_tools
 838        unmanaged_defs = (
 839            [await t.toolcall_definition() for t in unmanaged] if unmanaged else []
 840        )
 841
 842        merged = registry_defs + unmanaged_defs
 843        seen_names: set[str] = set()
 844        for d in merged:
 845            name = d["function"]["name"]
 846            if name in seen_names:
 847                raise ValueError(
 848                    f"Duplicate tool name {name!r}: unmanaged and registry tools must have unique names."
 849                )
 850            seen_names.add(name)
 851
 852        return merged
 853
 854    async def process_tool_calls(
 855        self, tool_calls: list[ChatCompletionMessageToolCall] | None
 856    ) -> tuple[str | None, list[ChatCompletionToolMessageParamWrapper]]:
 857        if tool_calls is None:
 858            return None, []
 859
 860        assistant_output_from_toolcall: str | None = None
 861        tool_call_response_messages: list[ChatCompletionToolMessageParamWrapper] = []
 862        tool_run_coroutines = []
 863
 864        for tool_call in tool_calls:
 865            # Kiln "task_response" tool is used for returning structured output via tool calls.
 866            # Load the output from the tool call. Also
 867            if tool_call.function.name == "task_response":
 868                assistant_output_from_toolcall = tool_call.function.arguments
 869                continue
 870
 871            # Process normal tool calls (not the "task_response" tool)
 872            tool_name = tool_call.function.name
 873            tool = None
 874            for tool_option in await self._tools_for_execution():
 875                if await tool_option.name() == tool_name:
 876                    tool = tool_option
 877                    break
 878            if not tool:
 879                raise RuntimeError(
 880                    f"A tool named '{tool_name}' was invoked by a model, but was not available."
 881                )
 882
 883            # Parse the arguments and validate them against the tool's schema
 884            try:
 885                parsed_args = json.loads(tool_call.function.arguments)
 886            except json.JSONDecodeError:
 887                raise RuntimeError(
 888                    f"Failed to parse arguments for tool '{tool_name}' (should be JSON): {tool_call.function.arguments}"
 889                )
 890            try:
 891                tool_call_definition = await tool.toolcall_definition()
 892                json_schema = json.dumps(tool_call_definition["function"]["parameters"])
 893                validate_schema_with_value_error(parsed_args, json_schema)
 894            except Exception as e:
 895                raise RuntimeError(
 896                    f"Failed to validate arguments for tool '{tool_name}'. The arguments didn't match the tool's schema. The arguments were: {parsed_args}\n The error was: {e}"
 897                ) from e
 898
 899            # Create context with the calling task's allow_saving setting
 900            context = ToolCallContext(
 901                allow_saving=self.base_adapter_config.allow_saving
 902            )
 903
 904            async def run_tool_and_format(
 905                t=tool, c=context, args=parsed_args, tc_id=tool_call.id
 906            ):
 907                result = await t.run(c, **args)
 908                return ChatCompletionToolMessageParamWrapper(
 909                    role="tool",
 910                    tool_call_id=tc_id,
 911                    content=result.output,
 912                    kiln_task_tool_data=result.kiln_task_tool_data
 913                    if isinstance(result, KilnTaskToolResult)
 914                    else None,
 915                    is_error=result.is_error if result.is_error else None,
 916                    error_message=result.error_message
 917                    if result.error_message
 918                    else None,
 919                )
 920
 921            tool_run_coroutines.append(run_tool_and_format())
 922
 923        if tool_run_coroutines:
 924            tool_call_response_messages = await asyncio.gather(*tool_run_coroutines)
 925
 926        if (
 927            assistant_output_from_toolcall is not None
 928            and len(tool_call_response_messages) > 0
 929        ):
 930            raise RuntimeError(
 931                "Model asked for impossible combination: task_response tool call and other tool calls were both provided in the same turn. This is not supported as it means the model asked us to both return task_response results (ending the turn) and run new tools calls to send back to the model. If the model makes this mistake often, try a difference structured data model like JSON schema, where this is impossible."
 932            )
 933
 934        return assistant_output_from_toolcall, tool_call_response_messages
 935
 936    def litellm_message_to_trace_message(
 937        self,
 938        raw_message: LiteLLMMessage,
 939        latency_ms: int | None = None,
 940        usage: MessageUsage | None = None,
 941    ) -> ChatCompletionAssistantMessageParamWrapper:
 942        """
 943        Convert a LiteLLM Message object to an OpenAI compatible message, our ChatCompletionAssistantMessageParamWrapper
 944        """
 945        message: ChatCompletionAssistantMessageParamWrapper = {
 946            "role": "assistant",
 947        }
 948        if raw_message.role != "assistant":
 949            raise ValueError(
 950                "Model returned a message with a role other than assistant. This is not supported."
 951            )
 952
 953        if hasattr(raw_message, "content"):
 954            message["content"] = raw_message.content
 955        if hasattr(raw_message, "reasoning_content"):
 956            message["reasoning_content"] = raw_message.reasoning_content
 957        if hasattr(raw_message, "tool_calls"):
 958            # Convert ChatCompletionMessageToolCall to ChatCompletionMessageToolCallParam
 959            open_ai_tool_calls: List[ChatCompletionMessageToolCallParam] = []
 960            for litellm_tool_call in raw_message.tool_calls or []:
 961                # Optional in the SDK for streaming responses, but should never be None at this point.
 962                if litellm_tool_call.function.name is None:
 963                    raise ValueError(
 964                        "The model requested a tool call, without providing a function name (required)."
 965                    )
 966                open_ai_tool_calls.append(
 967                    ChatCompletionMessageToolCallParam(
 968                        id=litellm_tool_call.id,
 969                        type="function",
 970                        function={
 971                            "name": litellm_tool_call.function.name,
 972                            "arguments": litellm_tool_call.function.arguments,
 973                        },
 974                    )
 975                )
 976            if len(open_ai_tool_calls) > 0:
 977                message["tool_calls"] = open_ai_tool_calls
 978
 979        if latency_ms is not None:
 980            message["latency_ms"] = latency_ms
 981
 982        if usage is not None:
 983            message["usage"] = usage
 984
 985        if not message.get("content") and not message.get("tool_calls"):
 986            raise ValueError(EMPTY_RESPONSE_ERROR_MESSAGE)
 987
 988        return message
 989
 990    def all_messages_to_trace(
 991        self,
 992        messages: list[ChatCompletionMessageIncludingLiteLLM],
 993        message_latency: dict[int, int] | None = None,
 994        message_usage: dict[int, MessageUsage] | None = None,
 995    ) -> list[ChatCompletionMessageParam]:
 996        """
 997        Internally we allow LiteLLM Message objects, but for trace we need OpenAI compatible types. Replace LiteLLM Message objects with OpenAI compatible types.
 998
 999        Non-LiteLLM dict messages pass through unchanged. Any per-message
1000        ``usage``/``latency_ms`` already attached to those dicts (e.g. from a
1001        seeded prior trace) is preserved.
1002        """
1003        trace: list[ChatCompletionMessageParam] = []
1004        for i, message in enumerate(messages):
1005            if isinstance(message, LiteLLMMessage):
1006                latency_ms = message_latency.get(i) if message_latency else None
1007                usage = message_usage.get(i) if message_usage else None
1008                trace.append(
1009                    self.litellm_message_to_trace_message(message, latency_ms, usage)
1010                )
1011            else:
1012                trace.append(message)
1013        return trace
1014
1015    def _messages_to_trace(
1016        self,
1017        messages: list[ChatCompletionMessageParam],
1018    ) -> list[ChatCompletionMessageParam]:
1019        """Override: the `messages` list may transiently contain LiteLLM
1020        Message objects. Normalize to API-safe shapes for export on error.
1021        """
1022        return self.all_messages_to_trace(messages)  # type: ignore[arg-type]