kiln-ai 0.15.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

kiln_ai/adapters/model_adapters/base_adapter.py

@@ -1,17 +1,32 @@
 import json
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass
-from typing import Dict, Literal, Tuple
-
-import jsonschema
-
-from kiln_ai.adapters.ml_model_list import KilnModelProvider, StructuredOutputMode
+from typing import Dict, Tuple
+
+from kiln_ai.adapters.chat.chat_formatter import (
+    ChatFormatter,
+    get_chat_formatter,
+)
+from kiln_ai.adapters.ml_model_list import (
+    KilnModelProvider,
+    StructuredOutputMode,
+    default_structured_output_mode_for_model_provider,
+)
 from kiln_ai.adapters.parsers.json_parser import parse_json_string
 from kiln_ai.adapters.parsers.parser_registry import model_parser_from_id
+from kiln_ai.adapters.parsers.request_formatters import request_formatter_from_id
 from kiln_ai.adapters.prompt_builders import prompt_builder_from_id
 from kiln_ai.adapters.provider_tools import kiln_model_provider_from
 from kiln_ai.adapters.run_output import RunOutput
-from kiln_ai.datamodel import DataSource, DataSourceType, Task, TaskOutput, TaskRun
+from kiln_ai.datamodel import (
+    DataSource,
+    DataSourceType,
+    Task,
+    TaskOutput,
+    TaskRun,
+    Usage,
+)
+from kiln_ai.datamodel.datamodel_enums import ChatStrategy
 from kiln_ai.datamodel.json_schema import validate_schema_with_value_error
 from kiln_ai.datamodel.task import RunConfig
 from kiln_ai.utils.config import Config
@@ -30,9 +45,6 @@ class AdapterConfig:
     default_tags: list[str] | None = None
 
 
-COT_FINAL_ANSWER_PROMPT = "Considering the above, return a final result."
-
-
 class BaseAdapter(metaclass=ABCMeta):
     """Base class for AI model adapters that handle task execution.
 
@@ -53,6 +65,7 @@ class BaseAdapter(metaclass=ABCMeta):
         config: AdapterConfig | None = None,
     ):
         self.run_config = run_config
+        self.update_run_config_unknown_structured_output_mode()
         self.prompt_builder = prompt_builder_from_id(
             run_config.prompt_id, run_config.task
         )
@@ -106,14 +119,19 @@ class BaseAdapter(metaclass=ABCMeta):
                 "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.",
             )
 
+        # Format model input for model call (we save the original input in the task without formatting)
+        formatted_input = input
+        formatter_id = self.model_provider().formatter
+        if formatter_id is not None:
+            formatter = request_formatter_from_id(formatter_id)
+            formatted_input = formatter.format_input(input)
+
         # Run
-        run_output = await self._run(input)
+        run_output, usage = await self._run(formatted_input)
 
         # Parse
         provider = self.model_provider()
-        parser = model_parser_from_id(provider.parser)(
-            structured_output=self.has_structured_output()
-        )
+        parser = model_parser_from_id(provider.parser)
         parsed_output = parser.parse_output(original_output=run_output)
 
         # validate output
@@ -147,7 +165,7 @@ class BaseAdapter(metaclass=ABCMeta):
             )
 
         # Generate the run and output
-        run = self.generate_run(input, input_source, parsed_output)
+        run = self.generate_run(input, input_source, parsed_output, usage)
 
         # Save the run if configured to do so, and we have a path to save to
         if (
@@ -170,15 +188,15 @@ class BaseAdapter(metaclass=ABCMeta):
         pass
 
     @abstractmethod
-    async def _run(self, input: Dict | str) -> RunOutput:
+    async def _run(self, input: Dict | str) -> Tuple[RunOutput, Usage | None]:
         pass
 
     def build_prompt(self) -> str:
         # The prompt builder needs to know if we want to inject formatting instructions
-        provider = self.model_provider()
+        structured_output_mode = self.run_config.structured_output_mode
         add_json_instructions = self.has_structured_output() and (
-            provider.structured_output_mode == StructuredOutputMode.json_instructions
-            or provider.structured_output_mode
+            structured_output_mode == StructuredOutputMode.json_instructions
+            or structured_output_mode
             == StructuredOutputMode.json_instruction_and_object
         )
 
@@ -186,30 +204,59 @@ class BaseAdapter(metaclass=ABCMeta):
             include_json_instructions=add_json_instructions
         )
 
-    def run_strategy(
-        self,
-    ) -> Tuple[Literal["cot_as_message", "cot_two_call", "basic"], str | None]:
-        # Determine the run strategy for COT prompting. 3 options:
-        # 1. "Thinking" LLM designed to output thinking in a structured format plus a COT prompt: we make 1 call to the LLM, which outputs thinking in a structured format. We include the thinking instuctions as a message.
-        # 2. Normal LLM with COT prompt: we make 2 calls to the LLM - one for thinking and one for the final response. This helps us use the LLM's structured output modes (json_schema, tools, etc), which can't be used in a single call. It also separates the thinking from the final response.
-        # 3. Non chain of thought: we make 1 call to the LLM, with no COT prompt.
+    def build_chat_formatter(self, input: Dict | str) -> ChatFormatter:
+        # 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.
+
         cot_prompt = self.prompt_builder.chain_of_thought_prompt()
-        reasoning_capable = self.model_provider().reasoning_capable
+        system_message = self.build_prompt()
+
+        # 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.
+        if not cot_prompt:
+            return get_chat_formatter(
+                strategy=ChatStrategy.single_turn,
+                system_message=system_message,
+                user_input=input,
+            )
 
-        if cot_prompt and reasoning_capable:
-            # 1: "Thinking" LLM designed to output thinking in a structured format
+        # Some models like finetunes are trained with a specific chat strategy. Use that.
+        # 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.
+        tuned_chat_strategy = self.model_provider().tuned_chat_strategy
+        if tuned_chat_strategy and tuned_chat_strategy != ChatStrategy.single_turn:
+            return get_chat_formatter(
+                strategy=tuned_chat_strategy,
+                system_message=system_message,
+                user_input=input,
+                thinking_instructions=cot_prompt,
+            )
+
+        # Pick the best chat strategy for the model given it has a cot prompt.
+        reasoning_capable = self.model_provider().reasoning_capable
+        if reasoning_capable:
+            # "Thinking" LLM designed to output thinking in a structured format. We'll use it's native format.
             # A simple message with the COT prompt appended to the message list is sufficient
-            return "cot_as_message", cot_prompt
-        elif cot_prompt:
-            # 2: Unstructured output with COT
-            # Two calls to separate the thinking from the final response
-            return "cot_two_call", cot_prompt
+            return get_chat_formatter(
+                strategy=ChatStrategy.single_turn_r1_thinking,
+                system_message=system_message,
+                user_input=input,
+                thinking_instructions=cot_prompt,
+            )
         else:
-            return "basic", None
+            # Unstructured output with COT
+            # Two calls to separate the thinking from the final response
+            return get_chat_formatter(
+                strategy=ChatStrategy.two_message_cot,
+                system_message=system_message,
+                user_input=input,
+                thinking_instructions=cot_prompt,
+            )
 
     # create a run and task output
     def generate_run(
-        self, input: Dict | str, input_source: DataSource | None, run_output: RunOutput
+        self,
+        input: Dict | str,
+        input_source: DataSource | None,
+        run_output: RunOutput,
+        usage: Usage | None = None,
     ) -> TaskRun:
         # Convert input and output to JSON strings if they are dictionaries
         input_str = (
@@ -242,6 +289,7 @@ class BaseAdapter(metaclass=ABCMeta):
             ),
             intermediate_outputs=run_output.intermediate_outputs,
             tags=self.base_adapter_config.default_tags or [],
+            usage=usage,
         )
 
         return new_task_run
@@ -254,5 +302,22 @@ class BaseAdapter(metaclass=ABCMeta):
         props["model_name"] = self.run_config.model_name
         props["model_provider"] = self.run_config.model_provider_name
         props["prompt_id"] = self.run_config.prompt_id
+        props["structured_output_mode"] = self.run_config.structured_output_mode
+        props["temperature"] = self.run_config.temperature
+        props["top_p"] = self.run_config.top_p
 
         return props
+
+    def update_run_config_unknown_structured_output_mode(self) -> None:
+        structured_output_mode = self.run_config.structured_output_mode
+
+        # Old datamodels didn't save the structured output mode. Some clients (tests, end users) might not set it.
+        # Look up our recommended mode from ml_model_list if we have one
+        if structured_output_mode == StructuredOutputMode.unknown:
+            new_run_config = self.run_config.model_copy(deep=True)
+            structured_output_mode = default_structured_output_mode_for_model_provider(
+                self.run_config.model_name,
+                self.run_config.model_provider_name,
+            )
+            new_run_config.structured_output_mode = structured_output_mode
+            self.run_config = new_run_config

kiln_ai/adapters/model_adapters/litellm_adapter.py

@@ -1,7 +1,9 @@
+import logging
 from typing import Any, Dict
 
 import litellm
 from litellm.types.utils import ChoiceLogprobs, Choices, ModelResponse
+from litellm.types.utils import Usage as LiteLlmUsage
 
 import kiln_ai.datamodel as datamodel
 from kiln_ai.adapters.ml_model_list import (
@@ -10,25 +12,23 @@ from kiln_ai.adapters.ml_model_list import (
     StructuredOutputMode,
 )
 from kiln_ai.adapters.model_adapters.base_adapter import (
-    COT_FINAL_ANSWER_PROMPT,
     AdapterConfig,
     BaseAdapter,
     RunOutput,
+    Usage,
 )
-from kiln_ai.adapters.model_adapters.litellm_config import (
-    LiteLlmConfig,
-)
-from kiln_ai.datamodel import PromptGenerators, PromptId
-from kiln_ai.datamodel.task import RunConfig
+from kiln_ai.adapters.model_adapters.litellm_config import LiteLlmConfig
+from kiln_ai.datamodel.task import run_config_from_run_config_properties
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
+logger = logging.getLogger(__name__)
+
 
 class LiteLlmAdapter(BaseAdapter):
     def __init__(
         self,
         config: LiteLlmConfig,
         kiln_task: datamodel.Task,
-        prompt_id: PromptId | None = None,
         base_adapter_config: AdapterConfig | None = None,
     ):
         self.config = config
@@ -37,11 +37,10 @@ class LiteLlmAdapter(BaseAdapter):
         self._headers = config.default_headers
         self._litellm_model_id: str | None = None
 
-        run_config = RunConfig(
+        # Create a RunConfig, adding the task to the RunConfigProperties
+        run_config = run_config_from_run_config_properties(
             task=kiln_task,
-            model_name=config.model_name,
-            model_provider_name=config.provider_name,
-            prompt_id=prompt_id or PromptGenerators.SIMPLE,
+            run_config_properties=config.run_config_properties,
         )
 
         super().__init__(
@@ -49,84 +48,74 @@ class LiteLlmAdapter(BaseAdapter):
             config=base_adapter_config,
         )
 
-    async def _run(self, input: Dict | str) -> RunOutput:
+    async def _run(self, input: Dict | str) -> tuple[RunOutput, Usage | None]:
         provider = self.model_provider()
         if not provider.model_id:
             raise ValueError("Model ID is required for OpenAI compatible models")
 
-        intermediate_outputs: dict[str, str] = {}
-        prompt = self.build_prompt()
-        user_msg = self.prompt_builder.build_user_message(input)
-        messages = [
-            {"role": "system", "content": prompt},
-            {"role": "user", "content": user_msg},
-        ]
-
-        run_strategy, cot_prompt = self.run_strategy()
-
-        if run_strategy == "cot_as_message":
-            # Used for reasoning-capable models that can output thinking and structured format
-            if not cot_prompt:
-                raise ValueError("cot_prompt is required for cot_as_message strategy")
-            messages.append({"role": "system", "content": cot_prompt})
-        elif run_strategy == "cot_two_call":
-            if not cot_prompt:
-                raise ValueError("cot_prompt is required for cot_two_call strategy")
-            messages.append({"role": "system", "content": cot_prompt})
-
-            # First call for chain of thought
-            # No response format as this request is for "thinking" in plain text
-            # No logprobs as only needed for final answer
+        chat_formatter = self.build_chat_formatter(input)
+
+        prior_output = None
+        prior_message = None
+        response = None
+        turns = 0
+        while True:
+            turns += 1
+            if turns > 10:
+                raise RuntimeError(
+                    "Too many turns. Stopping iteration to avoid using too many tokens."
+                )
+
+            turn = chat_formatter.next_turn(prior_output)
+            if turn is None:
+                break
+
+            skip_response_format = not turn.final_call
+            all_messages = chat_formatter.message_dicts()
             completion_kwargs = await self.build_completion_kwargs(
-                provider, messages, None, skip_response_format=True
+                provider,
+                all_messages,
+                self.base_adapter_config.top_logprobs if turn.final_call else None,
+                skip_response_format,
             )
-            cot_response = await litellm.acompletion(**completion_kwargs)
+            response = await litellm.acompletion(**completion_kwargs)
             if (
-                not isinstance(cot_response, ModelResponse)
-                or not cot_response.choices
-                or len(cot_response.choices) == 0
-                or not isinstance(cot_response.choices[0], Choices)
+                not isinstance(response, ModelResponse)
+                or not response.choices
+                or len(response.choices) == 0
+                or not isinstance(response.choices[0], Choices)
             ):
                 raise RuntimeError(
-                    f"Expected ModelResponse with Choices, got {type(cot_response)}."
+                    f"Expected ModelResponse with Choices, got {type(response)}."
                 )
-            cot_content = cot_response.choices[0].message.content
-            if cot_content is not None:
-                intermediate_outputs["chain_of_thought"] = cot_content
-
-            messages.extend(
-                [
-                    {"role": "assistant", "content": cot_content or ""},
-                    {"role": "user", "content": COT_FINAL_ANSWER_PROMPT},
-                ]
-            )
+            prior_message = response.choices[0].message
+            prior_output = prior_message.content
 
-        # Make the API call using litellm
-        completion_kwargs = await self.build_completion_kwargs(
-            provider, messages, self.base_adapter_config.top_logprobs
-        )
-        response = await litellm.acompletion(**completion_kwargs)
+            # Fallback: Use args of first tool call to task_response if it exists
+            if (
+                not prior_output
+                and hasattr(prior_message, "tool_calls")
+                and prior_message.tool_calls
+            ):
+                tool_call = next(
+                    (
+                        tool_call
+                        for tool_call in prior_message.tool_calls
+                        if tool_call.function.name == "task_response"
+                    ),
+                    None,
+                )
+                if tool_call:
+                    prior_output = tool_call.function.arguments
 
-        if not isinstance(response, ModelResponse):
-            raise RuntimeError(f"Expected ModelResponse, got {type(response)}.")
+            if not prior_output:
+                raise RuntimeError("No output returned from model")
 
-        # Maybe remove this? There is no error attribute on the response object.
-        # # Keeping in typesafe way as we added it for a reason, but should investigate what that was and if it still applies.
-        if hasattr(response, "error") and response.__getattribute__("error"):
-            raise RuntimeError(
-                f"LLM API returned an error: {response.__getattribute__('error')}"
-            )
+        if response is None or prior_message is None:
+            raise RuntimeError("No response returned from model")
 
-        if (
-            not response.choices
-            or len(response.choices) == 0
-            or not isinstance(response.choices[0], Choices)
-        ):
-            raise RuntimeError(
-                "No message content returned in the response from LLM API"
-            )
+        intermediate_outputs = chat_formatter.intermediate_outputs()
 
-        message = response.choices[0].message
         logprobs = (
             response.choices[0].logprobs
             if hasattr(response.choices[0], "logprobs")
@@ -139,28 +128,16 @@ class LiteLlmAdapter(BaseAdapter):
             raise RuntimeError("Logprobs were required, but no logprobs were returned.")
 
         # Save reasoning if it exists and was parsed by LiteLLM (or openrouter, or anyone upstream)
-        if hasattr(message, "reasoning_content") and message.reasoning_content:
-            intermediate_outputs["reasoning"] = message.reasoning_content
-
-        # the string content of the response
-        response_content = message.content
-
-        # Fallback: Use args of first tool call to task_response if it exists
         if (
-            not response_content
-            and hasattr(message, "tool_calls")
-            and message.tool_calls
+            prior_message is not None
+            and hasattr(prior_message, "reasoning_content")
+            and prior_message.reasoning_content
+            and len(prior_message.reasoning_content.strip()) > 0
         ):
-            tool_call = next(
-                (
-                    tool_call
-                    for tool_call in message.tool_calls
-                    if tool_call.function.name == "task_response"
-                ),
-                None,
-            )
-            if tool_call:
-                response_content = tool_call.function.arguments
+            intermediate_outputs["reasoning"] = prior_message.reasoning_content.strip()
+
+        # the string content of the response
+        response_content = prior_output
 
         if not isinstance(response_content, str):
             raise RuntimeError(f"response is not a string: {response_content}")
@@ -169,7 +146,7 @@ class LiteLlmAdapter(BaseAdapter):
             output=response_content,
             intermediate_outputs=intermediate_outputs,
             output_logprobs=logprobs,
-        )
+        ), self.usage_from_response(response)
 
     def adapter_name(self) -> str:
         return "kiln_openai_compatible_adapter"
@@ -179,8 +156,9 @@ class LiteLlmAdapter(BaseAdapter):
         if not self.has_structured_output():
             return {}
 
-        provider = self.model_provider()
-        match provider.structured_output_mode:
+        structured_output_mode = self.run_config.structured_output_mode
+
+        match structured_output_mode:
             case StructuredOutputMode.json_mode:
                 return {"response_format": {"type": "json_object"}}
             case StructuredOutputMode.json_schema:
@@ -199,16 +177,20 @@ class LiteLlmAdapter(BaseAdapter):
                 # We set response_format to json_object and also set json instructions in the prompt
                 return {"response_format": {"type": "json_object"}}
             case StructuredOutputMode.default:
-                if provider.name == ModelProviderName.ollama:
+                provider_name = self.run_config.model_provider_name
+                if provider_name == ModelProviderName.ollama:
                     # Ollama added json_schema to all models: https://ollama.com/blog/structured-outputs
                     return self.json_schema_response_format()
                 else:
                     # Default to function calling -- it's older than the other modes. Higher compatibility.
                     # Strict isn't widely supported yet, so we don't use it by default unless it's OpenAI.
-                    strict = provider.name == ModelProviderName.openai
+                    strict = provider_name == ModelProviderName.openai
                     return self.tool_call_params(strict=strict)
+            case StructuredOutputMode.unknown:
+                # See above, but this case should never happen.
+                raise ValueError("Structured output mode is unknown.")
             case _:
-                raise_exhaustive_enum_error(provider.structured_output_mode)
+                raise_exhaustive_enum_error(structured_output_mode)
 
     def json_schema_response_format(self) -> dict[str, Any]:
         output_schema = self.task().output_schema()
@@ -380,6 +362,13 @@ class LiteLlmAdapter(BaseAdapter):
             "messages": messages,
             "api_base": self._api_base,
             "headers": self._headers,
+            "temperature": self.run_config.temperature,
+            "top_p": self.run_config.top_p,
+            # This drops params that are not supported by the model. Only openai params like top_p, temperature -- not litellm params like model, etc.
+            # Not all models and providers support all openai params (for example, o3 doesn't support top_p)
+            # Better to ignore them than to fail the model call.
+            # https://docs.litellm.ai/docs/completion/input
+            "drop_params": True,
             **extra_body,
             **self._additional_body_options,
         }
@@ -394,3 +383,30 @@ class LiteLlmAdapter(BaseAdapter):
             completion_kwargs["top_logprobs"] = top_logprobs
 
         return completion_kwargs
+
+    def usage_from_response(self, response: ModelResponse) -> Usage | None:
+        litellm_usage = response.get("usage", None)
+        cost = response._hidden_params.get("response_cost", None)
+        if not litellm_usage and not cost:
+            return None
+
+        usage = Usage()
+
+        if litellm_usage and isinstance(litellm_usage, LiteLlmUsage):
+            usage.input_tokens = litellm_usage.get("prompt_tokens", None)
+            usage.output_tokens = litellm_usage.get("completion_tokens", None)
+            usage.total_tokens = litellm_usage.get("total_tokens", None)
+        else:
+            logger.warning(
+                f"Unexpected usage format from litellm: {litellm_usage}. Expected Usage object, got {type(litellm_usage)}"
+            )
+
+        if isinstance(cost, float):
+            usage.cost = cost
+        elif cost is not None:
+            # None is allowed, but no other types are expected
+            logger.warning(
+                f"Unexpected cost format from litellm: {cost}. Expected float, got {type(cost)}"
+            )
+
+        return usage

kiln_ai/adapters/model_adapters/litellm_config.py

@@ -1,10 +1,11 @@
 from dataclasses import dataclass, field
 
+from kiln_ai.datamodel.task import RunConfigProperties
+
 
 @dataclass
 class LiteLlmConfig:
-    model_name: str
-    provider_name: str
+    run_config_properties: RunConfigProperties
     # If set, over rides the provider-name based URL from litellm
     base_url: str | None = None
     # Headers to send with every request