kiln-ai 0.8.0__py3-none-any.whl → 0.11.1__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/{base_adapter.py → model_adapters/base_adapter.py}

@@ -1,8 +1,13 @@
 import json
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass
-from typing import Dict
+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 BasePromptBuilder, SimplePromptBuilder
+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,
@@ -13,8 +18,6 @@ from kiln_ai.datamodel import (
 from kiln_ai.datamodel.json_schema import validate_schema
 from kiln_ai.utils.config import Config
 
-from .prompt_builders import BasePromptBuilder, SimplePromptBuilder
-
 
 @dataclass
 class AdapterInfo:
@@ -22,12 +25,10 @@ class AdapterInfo:
     model_name: str
     model_provider: str
     prompt_builder_name: str
+    prompt_id: str | None = None
 
 
-@dataclass
-class RunOutput:
-    output: Dict | str
-    intermediate_outputs: Dict[str, str] | None
+COT_FINAL_ANSWER_PROMPT = "Considering the above, return a final result."
 
 
 class BaseAdapter(metaclass=ABCMeta):
@@ -47,6 +48,8 @@ class BaseAdapter(metaclass=ABCMeta):
     def __init__(
         self,
         kiln_task: Task,
+        model_name: str,
+        model_provider_name: str,
         prompt_builder: BasePromptBuilder | None = None,
         tags: list[str] | None = None,
     ):
@@ -55,6 +58,26 @@ class BaseAdapter(metaclass=ABCMeta):
         self.output_schema = self.kiln_task.output_json_schema
         self.input_schema = self.kiln_task.input_json_schema
         self.default_tags = tags
+        self.model_name = model_name
+        self.model_provider_name = model_provider_name
+        self._model_provider: KilnModelProvider | None = None
+
+    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.model_name or not self.model_provider_name:
+            raise ValueError("model_name and model_provider_name must be provided")
+        self._model_provider = kiln_model_provider_from(
+            self.model_name, self.model_provider_name
+        )
+        if not self._model_provider:
+            raise ValueError(
+                f"model_provider_name {self.model_provider_name} not found for model {self.model_name}"
+            )
+        return self._model_provider
 
     async def invoke_returning_raw(
         self,
@@ -81,21 +104,28 @@ class BaseAdapter(metaclass=ABCMeta):
         # 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(run_output.output, dict):
+            if not isinstance(parsed_output.output, dict):
                 raise RuntimeError(
-                    f"structured response is not a dict: {run_output.output}"
+                    f"structured response is not a dict: {parsed_output.output}"
                 )
-            validate_schema(run_output.output, self.output_schema)
+            validate_schema(parsed_output.output, self.output_schema)
         else:
-            if not isinstance(run_output.output, str):
+            if not isinstance(parsed_output.output, str):
                 raise RuntimeError(
-                    f"response is not a string for non-structured task: {run_output.output}"
+                    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, run_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 Config.shared().autosave_runs and self.kiln_task.path is not None:
@@ -118,16 +148,49 @@ class BaseAdapter(metaclass=ABCMeta):
         pass
 
     def build_prompt(self) -> str:
-        return self.prompt_builder.build_prompt()
+        # 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) if isinstance(input, dict) else input
+        input_str = (
+            json.dumps(input, ensure_ascii=False) if isinstance(input, dict) else input
+        )
         output_str = (
-            json.dumps(run_output.output)
+            json.dumps(run_output.output, ensure_ascii=False)
             if isinstance(run_output.output, dict)
             else run_output.output
         )
@@ -155,27 +218,6 @@ class BaseAdapter(metaclass=ABCMeta):
             tags=self.default_tags or [],
         )
 
-        exclude_fields = {
-            "id": True,
-            "created_at": True,
-            "updated_at": True,
-            "path": True,
-            "output": {"id": True, "created_at": True, "updated_at": True},
-        }
-        new_run_dump = new_task_run.model_dump(exclude=exclude_fields)
-
-        # Check if the same run already exists
-        existing_task_run = next(
-            (
-                task_run
-                for task_run in self.kiln_task.runs()
-                if task_run.model_dump(exclude=exclude_fields) == new_run_dump
-            ),
-            None,
-        )
-        if existing_task_run:
-            return existing_task_run
-
         return new_task_run
 
     def _properties_for_task_output(self) -> Dict[str, str | int | float]:
@@ -187,5 +229,7 @@ class BaseAdapter(metaclass=ABCMeta):
         props["model_name"] = adapter_info.model_name
         props["model_provider"] = adapter_info.model_provider
         props["prompt_builder_name"] = adapter_info.prompt_builder_name
+        if adapter_info.prompt_id is not None:
+            props["prompt_id"] = adapter_info.prompt_id
 
         return props

kiln_ai/adapters/{langchain_adapters.py → model_adapters/langchain_adapters.py}

@@ -1,5 +1,4 @@
 import os
-from os import getenv
 from typing import Any, Dict
 
 from langchain_aws import ChatBedrockConverse
@@ -11,20 +10,28 @@ from langchain_core.runnables import Runnable
 from langchain_fireworks import ChatFireworks
 from langchain_groq import ChatGroq
 from langchain_ollama import ChatOllama
-from langchain_openai import ChatOpenAI
 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,
+    AdapterInfo,
+    BaseAdapter,
+    BasePromptBuilder,
+    RunOutput,
+)
 from kiln_ai.adapters.ollama_tools import (
     get_ollama_connection,
     ollama_base_url,
     ollama_model_installed,
 )
 from kiln_ai.utils.config import Config
-
-from .base_adapter import AdapterInfo, BaseAdapter, BasePromptBuilder, RunOutput
-from .ml_model_list import KilnModelProvider, ModelProviderName
-from .provider_tools import kiln_model_provider_from
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 LangChainModelType = BaseChatModel | Runnable[LanguageModelInput, Dict | BaseModel]
 
@@ -41,39 +48,62 @@ class LangchainAdapter(BaseAdapter):
         prompt_builder: BasePromptBuilder | None = None,
         tags: list[str] | None = None,
     ):
-        super().__init__(kiln_task, prompt_builder=prompt_builder, tags=tags)
         if custom_model is not None:
             self._model = custom_model
 
             # Attempt to infer model provider and name from custom model
-            self.model_provider = "custom.langchain:" + custom_model.__class__.__name__
-            self.model_name = "custom.langchain:unknown_model"
-            if hasattr(custom_model, "model_name") and isinstance(
-                getattr(custom_model, "model_name"), str
-            ):
-                self.model_name = "custom.langchain:" + getattr(
-                    custom_model, "model_name"
-                )
-            if hasattr(custom_model, "model") and isinstance(
-                getattr(custom_model, "model"), str
-            ):
-                self.model_name = "custom.langchain:" + getattr(custom_model, "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:
-            self.model_name = model_name
-            self.model_provider = provider or "custom.langchain.default_provider"
+            # 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")
+
+        super().__init__(
+            kiln_task,
+            model_name=model_name,
+            model_provider_name=provider,
+            prompt_builder=prompt_builder,
+            tags=tags,
+        )
+
     async def model(self) -> LangChainModelType:
         # cached model
         if self._model:
             return self._model
 
-        self._model = await langchain_model_from(self.model_name, self.model_provider)
+        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 self.has_structured_output():
+        if use_lc_structured_output:
             if not hasattr(self._model, "with_structured_output") or not callable(
                 getattr(self._model, "with_structured_output")
             ):
@@ -88,8 +118,8 @@ class LangchainAdapter(BaseAdapter):
                 )
             output_schema["title"] = "task_response"
             output_schema["description"] = "A response from the task"
-            with_structured_output_options = await get_structured_output_options(
-                self.model_name, self.model_provider
+            with_structured_output_options = self.get_structured_output_options(
+                self.model_name, self.model_provider_name
             )
             self._model = self._model.with_structured_output(
                 output_schema,
@@ -99,6 +129,7 @@ class LangchainAdapter(BaseAdapter):
         return self._model
 
     async def _run(self, input: Dict | str) -> RunOutput:
+        provider = self.model_provider()
         model = await self.model()
         chain = model
         intermediate_outputs = {}
@@ -110,58 +141,63 @@ class LangchainAdapter(BaseAdapter):
             HumanMessage(content=user_msg),
         ]
 
-        # COT with structured output
-        cot_prompt = self.prompt_builder.chain_of_thought_prompt()
-        if cot_prompt and self.has_structured_output():
-            # Base model (without structured output) used for COT message
-            base_model = await langchain_model_from(
-                self.model_name, self.model_provider
-            )
+        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(
-                SystemMessage(content="Considering the above, return a final result.")
-            )
-        elif cot_prompt:
-            messages.append(SystemMessage(content=cot_prompt))
+            messages.append(HumanMessage(content=COT_FINAL_ANSWER_PROMPT))
 
         response = await chain.ainvoke(messages)
 
-        if self.has_structured_output():
-            if (
-                not isinstance(response, dict)
-                or "parsed" not in response
-                or not isinstance(response["parsed"], dict)
-            ):
-                raise RuntimeError(f"structured response not returned: {response}")
+        # 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,
             )
-        else:
-            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,
-            )
+
+        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_info(self) -> AdapterInfo:
         return AdapterInfo(
             model_name=self.model_name,
-            model_provider=self.model_provider,
+            model_provider=self.model_provider_name,
             adapter_name="kiln_langchain_adapter",
             prompt_builder_name=self.prompt_builder.__class__.prompt_builder_name(),
+            prompt_id=self.prompt_builder.prompt_id(),
         )
 
     def _munge_response(self, response: Dict) -> Dict:
@@ -174,34 +210,54 @@ class LangchainAdapter(BaseAdapter):
             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 {}
 
-async def get_structured_output_options(
-    model_name: str, model_provider: str
-) -> Dict[str, Any]:
-    finetune_provider = await kiln_model_provider_from(model_name, model_provider)
-    if finetune_provider and finetune_provider.adapter_options.get("langchain"):
-        return finetune_provider.adapter_options["langchain"].get(
-            "with_structured_output_options", {}
-        )
-    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:
+                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(
-    name: str, provider_name: str | None = None
-) -> BaseChatModel:
-    provider = await kiln_model_provider_from(name, provider_name)
-    return await langchain_model_from_provider(provider, name)
+    async def langchain_model_from(self) -> BaseChatModel:
+        provider = self.model_provider()
+        return await langchain_model_from_provider(provider, self.model_name)
 
 
 async def langchain_model_from_provider(
     provider: KilnModelProvider, model_name: str
 ) -> BaseChatModel:
     if provider.name == ModelProviderName.openai:
-        api_key = Config.shared().open_ai_api_key
-        return ChatOpenAI(**provider.provider_options, openai_api_key=api_key)  # type: ignore[arg-type]
+        # We use the OpenAICompatibleAdapter for OpenAI
+        raise ValueError("OpenAI is not supported in Langchain adapter")
     elif provider.name == ModelProviderName.openai_compatible:
-        # See provider_tools.py for how base_url, key and other parameters are set
-        return ChatOpenAI(**provider.provider_options)  # type: ignore[arg-type]
+        # 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:
@@ -241,16 +297,6 @@ async def langchain_model_from_provider(
 
         raise ValueError(f"Model {model_name} not installed on Ollama")
     elif provider.name == ModelProviderName.openrouter:
-        api_key = Config.shared().open_router_api_key
-        base_url = getenv("OPENROUTER_BASE_URL") or "https://openrouter.ai/api/v1"
-        return ChatOpenAI(
-            **provider.provider_options,
-            openai_api_key=api_key,  # type: ignore[arg-type]
-            openai_api_base=base_url,  # type: ignore[arg-type]
-            default_headers={
-                "HTTP-Referer": "https://getkiln.ai/openrouter",
-                "X-Title": "KilnAI",
-            },
-        )
+        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,11 @@
+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
+    openrouter_style_reasoning: bool = False