kiln-ai 0.8.1__py3-none-any.whl → 0.12.0__py3-none-any.whl

kiln_ai/adapters/model_adapters/__init__.py

@@ -0,0 +1,18 @@
+"""
+# Model Adapters
+
+Model adapters are used to call AI models, like Ollama, OpenAI, etc.
+
+"""
+
+from . import (
+    base_adapter,
+    langchain_adapters,
+    openai_model_adapter,
+)
+
+__all__ = [
+    "base_adapter",
+    "langchain_adapters",
+    "openai_model_adapter",
+]

kiln_ai/adapters/model_adapters/base_adapter.py

@@ -0,0 +1,250 @@
+import json
+from abc import ABCMeta, abstractmethod
+from dataclasses import dataclass
+from typing import Dict, Literal, Tuple
+
+from kiln_ai.adapters.ml_model_list import KilnModelProvider, StructuredOutputMode
+from kiln_ai.adapters.parsers.parser_registry import model_parser_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.json_schema import validate_schema
+from kiln_ai.datamodel.task import RunConfig
+from kiln_ai.utils.config import Config
+
+
+@dataclass
+class AdapterConfig:
+    """
+    An adapter config is config options that do NOT impact the output of the model.
+
+    For example: if it's saved, of if we request additional data like logprobs.
+    """
+
+    allow_saving: bool = True
+    top_logprobs: int | None = None
+    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.
+
+    This abstract class provides the foundation for implementing model-specific adapters
+    that can process tasks with structured or unstructured inputs/outputs. It handles
+    input/output validation, prompt building, and run tracking.
+
+    Attributes:
+        prompt_builder (BasePromptBuilder): Builder for constructing prompts for the model
+        kiln_task (Task): The task configuration and metadata
+        output_schema (dict | None): JSON schema for validating structured outputs
+        input_schema (dict | None): JSON schema for validating structured inputs
+    """
+
+    def __init__(
+        self,
+        run_config: RunConfig,
+        config: AdapterConfig | None = None,
+    ):
+        self.run_config = run_config
+        self.prompt_builder = prompt_builder_from_id(
+            run_config.prompt_id, run_config.task
+        )
+        self._model_provider: KilnModelProvider | None = None
+
+        self.output_schema = self.task().output_json_schema
+        self.input_schema = self.task().input_json_schema
+        self.base_adapter_config = config or AdapterConfig()
+
+    def task(self) -> Task:
+        return self.run_config.task
+
+    def model_provider(self) -> KilnModelProvider:
+        """
+        Lazy load the model provider for this adapter.
+        """
+        if self._model_provider is not None:
+            return self._model_provider
+        if not self.run_config.model_name or not self.run_config.model_provider_name:
+            raise ValueError("model_name and model_provider_name must be provided")
+        self._model_provider = kiln_model_provider_from(
+            self.run_config.model_name, self.run_config.model_provider_name
+        )
+        if not self._model_provider:
+            raise ValueError(
+                f"model_provider_name {self.run_config.model_provider_name} not found for model {self.run_config.model_name}"
+            )
+        return self._model_provider
+
+    async def invoke_returning_raw(
+        self,
+        input: Dict | str,
+        input_source: DataSource | None = None,
+    ) -> Dict | str:
+        result = await self.invoke(input, input_source)
+        if self.task().output_json_schema is None:
+            return result.output.output
+        else:
+            return json.loads(result.output.output)
+
+    async def invoke(
+        self,
+        input: Dict | str,
+        input_source: DataSource | None = None,
+    ) -> TaskRun:
+        run_output, _ = await self.invoke_returning_run_output(input, input_source)
+        return run_output
+
+    async def invoke_returning_run_output(
+        self,
+        input: Dict | str,
+        input_source: DataSource | None = None,
+    ) -> Tuple[TaskRun, RunOutput]:
+        # validate input
+        if self.input_schema is not None:
+            if not isinstance(input, dict):
+                raise ValueError(f"structured input is not a dict: {input}")
+            validate_schema(input, self.input_schema)
+
+        # Run
+        run_output = await self._run(input)
+
+        # Parse
+        provider = self.model_provider()
+        parser = model_parser_from_id(provider.parser)(
+            structured_output=self.has_structured_output()
+        )
+        parsed_output = parser.parse_output(original_output=run_output)
+
+        # validate output
+        if self.output_schema is not None:
+            if not isinstance(parsed_output.output, dict):
+                raise RuntimeError(
+                    f"structured response is not a dict: {parsed_output.output}"
+                )
+            validate_schema(parsed_output.output, self.output_schema)
+        else:
+            if not isinstance(parsed_output.output, str):
+                raise RuntimeError(
+                    f"response is not a string for non-structured task: {parsed_output.output}"
+                )
+
+        # Generate the run and output
+        run = self.generate_run(input, input_source, parsed_output)
+
+        # Save the run if configured to do so, and we have a path to save to
+        if (
+            self.base_adapter_config.allow_saving
+            and Config.shared().autosave_runs
+            and self.task().path is not None
+        ):
+            run.save_to_file()
+        else:
+            # Clear the ID to indicate it's not persisted
+            run.id = None
+
+        return run, run_output
+
+    def has_structured_output(self) -> bool:
+        return self.output_schema is not None
+
+    @abstractmethod
+    def adapter_name(self) -> str:
+        pass
+
+    @abstractmethod
+    async def _run(self, input: Dict | str) -> RunOutput:
+        pass
+
+    def build_prompt(self) -> str:
+        # The prompt builder needs to know if we want to inject formatting instructions
+        provider = self.model_provider()
+        add_json_instructions = self.has_structured_output() and (
+            provider.structured_output_mode == StructuredOutputMode.json_instructions
+            or provider.structured_output_mode
+            == StructuredOutputMode.json_instruction_and_object
+        )
+
+        return self.prompt_builder.build_prompt(
+            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.
+        cot_prompt = self.prompt_builder.chain_of_thought_prompt()
+        reasoning_capable = self.model_provider().reasoning_capable
+
+        if cot_prompt and reasoning_capable:
+            # 1: "Thinking" LLM designed to output thinking in a structured 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
+        else:
+            return "basic", None
+
+    # create a run and task output
+    def generate_run(
+        self, input: Dict | str, input_source: DataSource | None, run_output: RunOutput
+    ) -> TaskRun:
+        # Convert input and output to JSON strings if they are dictionaries
+        input_str = (
+            json.dumps(input, ensure_ascii=False) if isinstance(input, dict) else input
+        )
+        output_str = (
+            json.dumps(run_output.output, ensure_ascii=False)
+            if isinstance(run_output.output, dict)
+            else run_output.output
+        )
+
+        # If no input source is provided, use the human data source
+        if input_source is None:
+            input_source = DataSource(
+                type=DataSourceType.human,
+                properties={"created_by": Config.shared().user_id},
+            )
+
+        new_task_run = TaskRun(
+            parent=self.task(),
+            input=input_str,
+            input_source=input_source,
+            output=TaskOutput(
+                output=output_str,
+                # Synthetic since an adapter, not a human, is creating this
+                source=DataSource(
+                    type=DataSourceType.synthetic,
+                    properties=self._properties_for_task_output(),
+                ),
+            ),
+            intermediate_outputs=run_output.intermediate_outputs,
+            tags=self.base_adapter_config.default_tags or [],
+        )
+
+        return new_task_run
+
+    def _properties_for_task_output(self) -> Dict[str, str | int | float]:
+        props = {}
+
+        # adapter info
+        props["adapter_name"] = self.adapter_name()
+        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
+
+        return props

kiln_ai/adapters/model_adapters/langchain_adapters.py

@@ -0,0 +1,309 @@
+import os
+from typing import Any, Dict
+
+from langchain_aws import ChatBedrockConverse
+from langchain_core.language_models import LanguageModelInput
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from langchain_core.messages.base import BaseMessage
+from langchain_core.runnables import Runnable
+from langchain_fireworks import ChatFireworks
+from langchain_groq import ChatGroq
+from langchain_ollama import ChatOllama
+from pydantic import BaseModel
+
+import kiln_ai.datamodel as datamodel
+from kiln_ai.adapters.ml_model_list import (
+    KilnModelProvider,
+    ModelProviderName,
+    StructuredOutputMode,
+)
+from kiln_ai.adapters.model_adapters.base_adapter import (
+    COT_FINAL_ANSWER_PROMPT,
+    AdapterConfig,
+    BaseAdapter,
+    RunOutput,
+)
+from kiln_ai.adapters.ollama_tools import (
+    get_ollama_connection,
+    ollama_base_url,
+    ollama_model_installed,
+)
+from kiln_ai.datamodel import PromptId
+from kiln_ai.datamodel.task import RunConfig
+from kiln_ai.utils.config import Config
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
+
+LangChainModelType = BaseChatModel | Runnable[LanguageModelInput, Dict | BaseModel]
+
+
+class LangchainAdapter(BaseAdapter):
+    _model: LangChainModelType | None = None
+
+    def __init__(
+        self,
+        kiln_task: datamodel.Task,
+        custom_model: BaseChatModel | None = None,
+        model_name: str | None = None,
+        provider: str | None = None,
+        prompt_id: PromptId | None = None,
+        base_adapter_config: AdapterConfig | None = None,
+    ):
+        if custom_model is not None:
+            self._model = custom_model
+
+            # Attempt to infer model provider and name from custom model
+            if provider is None:
+                provider = "custom.langchain:" + custom_model.__class__.__name__
+
+            if model_name is None:
+                model_name = "custom.langchain:unknown_model"
+                if hasattr(custom_model, "model_name") and isinstance(
+                    getattr(custom_model, "model_name"), str
+                ):
+                    model_name = "custom.langchain:" + getattr(
+                        custom_model, "model_name"
+                    )
+                if hasattr(custom_model, "model") and isinstance(
+                    getattr(custom_model, "model"), str
+                ):
+                    model_name = "custom.langchain:" + getattr(custom_model, "model")
+        elif model_name is not None:
+            # default provider name if not provided
+            provider = provider or "custom.langchain.default_provider"
+        else:
+            raise ValueError(
+                "model_name and provider must be provided if custom_model is not provided"
+            )
+
+        if model_name is None:
+            raise ValueError("model_name must be provided")
+
+        run_config = RunConfig(
+            task=kiln_task,
+            model_name=model_name,
+            model_provider_name=provider,
+            prompt_id=prompt_id or datamodel.PromptGenerators.SIMPLE,
+        )
+
+        super().__init__(
+            run_config=run_config,
+            config=base_adapter_config,
+        )
+
+    async def model(self) -> LangChainModelType:
+        # cached model
+        if self._model:
+            return self._model
+
+        self._model = await self.langchain_model_from()
+
+        # Decide if we want to use Langchain's structured output:
+        # 1. Only for structured tasks
+        # 2. Only if the provider's mode isn't json_instructions (only mode that doesn't use an API option for structured output capabilities)
+        provider = self.model_provider()
+        use_lc_structured_output = (
+            self.has_structured_output()
+            and provider.structured_output_mode
+            != StructuredOutputMode.json_instructions
+        )
+
+        if use_lc_structured_output:
+            if not hasattr(self._model, "with_structured_output") or not callable(
+                getattr(self._model, "with_structured_output")
+            ):
+                raise ValueError(
+                    f"model {self._model} does not support structured output, cannot use output_json_schema"
+                )
+            # Langchain expects title/description to be at top level, on top of json schema
+            output_schema = self.task().output_schema()
+            if output_schema is None:
+                raise ValueError(
+                    f"output_json_schema is not valid json: {self.task().output_json_schema}"
+                )
+            output_schema["title"] = "task_response"
+            output_schema["description"] = "A response from the task"
+            with_structured_output_options = self.get_structured_output_options(
+                self.run_config.model_name, self.run_config.model_provider_name
+            )
+            self._model = self._model.with_structured_output(
+                output_schema,
+                include_raw=True,
+                **with_structured_output_options,
+            )
+        return self._model
+
+    async def _run(self, input: Dict | str) -> RunOutput:
+        if self.base_adapter_config.top_logprobs is not None:
+            raise ValueError(
+                "Kiln's Langchain adapter does not support logprobs/top_logprobs. Select a model from an OpenAI compatible provider (openai, openrouter, etc) instead."
+            )
+
+        provider = self.model_provider()
+        model = await self.model()
+        chain = model
+        intermediate_outputs = {}
+
+        prompt = self.build_prompt()
+        user_msg = self.prompt_builder.build_user_message(input)
+        messages = [
+            SystemMessage(content=prompt),
+            HumanMessage(content=user_msg),
+        ]
+
+        run_strategy, cot_prompt = self.run_strategy()
+
+        if run_strategy == "cot_as_message":
+            if not cot_prompt:
+                raise ValueError("cot_prompt is required for cot_as_message strategy")
+            messages.append(SystemMessage(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(
+                SystemMessage(content=cot_prompt),
+            )
+
+            # Base model (without structured output) used for COT message
+            base_model = await self.langchain_model_from()
+
+            cot_messages = [*messages]
+            cot_response = await base_model.ainvoke(cot_messages)
+            intermediate_outputs["chain_of_thought"] = cot_response.content
+            messages.append(AIMessage(content=cot_response.content))
+            messages.append(HumanMessage(content=COT_FINAL_ANSWER_PROMPT))
+
+        response = await chain.ainvoke(messages)
+
+        # Langchain may have already parsed the response into structured output, so use that if available.
+        # However, a plain string may still be fixed at the parsing layer, so not being structured isn't a critical failure (yet)
+        if (
+            self.has_structured_output()
+            and isinstance(response, dict)
+            and "parsed" in response
+            and isinstance(response["parsed"], dict)
+        ):
+            structured_response = response["parsed"]
+            return RunOutput(
+                output=self._munge_response(structured_response),
+                intermediate_outputs=intermediate_outputs,
+            )
+
+        if not isinstance(response, BaseMessage):
+            raise RuntimeError(f"response is not a BaseMessage: {response}")
+
+        text_content = response.content
+        if not isinstance(text_content, str):
+            raise RuntimeError(f"response is not a string: {text_content}")
+
+        return RunOutput(
+            output=text_content,
+            intermediate_outputs=intermediate_outputs,
+        )
+
+    def adapter_name(self) -> str:
+        return "kiln_langchain_adapter"
+
+    def _munge_response(self, response: Dict) -> Dict:
+        # Mistral Large tool calling format is a bit different. Convert to standard format.
+        if (
+            "name" in response
+            and response["name"] == "task_response"
+            and "arguments" in response
+        ):
+            return response["arguments"]
+        return response
+
+    def get_structured_output_options(
+        self, model_name: str, model_provider_name: str
+    ) -> Dict[str, Any]:
+        provider = self.model_provider()
+        if not provider:
+            return {}
+
+        options = {}
+        # We may need to add some provider specific logic here if providers use different names for the same mode, but everyone is copying openai for now
+        match provider.structured_output_mode:
+            case StructuredOutputMode.function_calling_weak:
+                # Langchaing doesn't handle weak/strict separately
+                options["method"] = "function_calling"
+            case StructuredOutputMode.function_calling:
+                options["method"] = "function_calling"
+            case StructuredOutputMode.json_mode:
+                options["method"] = "json_mode"
+            case StructuredOutputMode.json_instruction_and_object:
+                # We also pass instructions
+                options["method"] = "json_mode"
+            case StructuredOutputMode.json_schema:
+                options["method"] = "json_schema"
+            case StructuredOutputMode.json_instructions:
+                # JSON done via instructions in prompt, not via API
+                pass
+            case StructuredOutputMode.default:
+                if provider.name == ModelProviderName.ollama:
+                    # Ollama has great json_schema support, so use that: https://ollama.com/blog/structured-outputs
+                    options["method"] = "json_schema"
+                else:
+                    # Let langchain decide the default
+                    pass
+            case _:
+                raise_exhaustive_enum_error(provider.structured_output_mode)
+
+        return options
+
+    async def langchain_model_from(self) -> BaseChatModel:
+        provider = self.model_provider()
+        return await langchain_model_from_provider(provider, self.run_config.model_name)
+
+
+async def langchain_model_from_provider(
+    provider: KilnModelProvider, model_name: str
+) -> BaseChatModel:
+    if provider.name == ModelProviderName.openai:
+        # We use the OpenAICompatibleAdapter for OpenAI
+        raise ValueError("OpenAI is not supported in Langchain adapter")
+    elif provider.name == ModelProviderName.openai_compatible:
+        # We use the OpenAICompatibleAdapter for OpenAI compatible
+        raise ValueError("OpenAI compatible is not supported in Langchain adapter")
+    elif provider.name == ModelProviderName.groq:
+        api_key = Config.shared().groq_api_key
+        if api_key is None:
+            raise ValueError(
+                "Attempted to use Groq without an API key set. "
+                "Get your API key from https://console.groq.com/keys"
+            )
+        return ChatGroq(**provider.provider_options, groq_api_key=api_key)  # type: ignore[arg-type]
+    elif provider.name == ModelProviderName.amazon_bedrock:
+        api_key = Config.shared().bedrock_access_key
+        secret_key = Config.shared().bedrock_secret_key
+        # langchain doesn't allow passing these, so ugly hack to set env vars
+        os.environ["AWS_ACCESS_KEY_ID"] = api_key
+        os.environ["AWS_SECRET_ACCESS_KEY"] = secret_key
+        return ChatBedrockConverse(
+            **provider.provider_options,
+        )
+    elif provider.name == ModelProviderName.fireworks_ai:
+        api_key = Config.shared().fireworks_api_key
+        return ChatFireworks(**provider.provider_options, api_key=api_key)
+    elif provider.name == ModelProviderName.ollama:
+        # Ollama model naming is pretty flexible. We try a few versions of the model name
+        potential_model_names = []
+        if "model" in provider.provider_options:
+            potential_model_names.append(provider.provider_options["model"])
+        if "model_aliases" in provider.provider_options:
+            potential_model_names.extend(provider.provider_options["model_aliases"])
+
+        # Get the list of models Ollama supports
+        ollama_connection = await get_ollama_connection()
+        if ollama_connection is None:
+            raise ValueError("Failed to connect to Ollama. Ensure Ollama is running.")
+
+        for model_name in potential_model_names:
+            if ollama_model_installed(ollama_connection, model_name):
+                return ChatOllama(model=model_name, base_url=ollama_base_url())
+
+        raise ValueError(f"Model {model_name} not installed on Ollama")
+    elif provider.name == ModelProviderName.openrouter:
+        raise ValueError("OpenRouter is not supported in Langchain adapter")
+    else:
+        raise ValueError(f"Invalid model or provider: {model_name} - {provider.name}")

kiln_ai/adapters/model_adapters/openai_compatible_config.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class OpenAICompatibleConfig:
+    api_key: str
+    model_name: str
+    provider_name: str
+    base_url: str | None = None  # Defaults to OpenAI
+    default_headers: dict[str, str] | None = None