kiln_ai.adapters.langchain_adapters

  1from typing import Dict
  2
  3from langchain_core.language_models.chat_models import BaseChatModel
  4from langchain_core.messages import HumanMessage, SystemMessage
  5from langchain_core.messages.base import BaseMessage
  6
  7import kiln_ai.datamodel as datamodel
  8
  9from .base_adapter import AdapterInfo, BaseAdapter, BasePromptBuilder
 10from .ml_model_list import langchain_model_from
 11
 12
 13class LangChainPromptAdapter(BaseAdapter):
 14    def __init__(
 15        self,
 16        kiln_task: datamodel.Task,
 17        custom_model: BaseChatModel | None = None,
 18        model_name: str | None = None,
 19        provider: str | None = None,
 20        prompt_builder: BasePromptBuilder | None = None,
 21    ):
 22        super().__init__(kiln_task, prompt_builder=prompt_builder)
 23        if custom_model is not None:
 24            self.model = custom_model
 25
 26            # Attempt to infer model provider and name from custom model
 27            self.model_provider = "custom.langchain:" + custom_model.__class__.__name__
 28            self.model_name = "custom.langchain:unknown_model"
 29            if hasattr(custom_model, "model_name") and isinstance(
 30                getattr(custom_model, "model_name"), str
 31            ):
 32                self.model_name = "custom.langchain:" + getattr(
 33                    custom_model, "model_name"
 34                )
 35            if hasattr(custom_model, "model") and isinstance(
 36                getattr(custom_model, "model"), str
 37            ):
 38                self.model_name = "custom.langchain:" + getattr(custom_model, "model")
 39        elif model_name is not None:
 40            self.model = langchain_model_from(model_name, provider)
 41            self.model_name = model_name
 42            self.model_provider = provider or "custom.langchain.default_provider"
 43        else:
 44            raise ValueError(
 45                "model_name and provider must be provided if custom_model is not provided"
 46            )
 47        if self.has_structured_output():
 48            if not hasattr(self.model, "with_structured_output") or not callable(
 49                getattr(self.model, "with_structured_output")
 50            ):
 51                raise ValueError(
 52                    f"model {self.model} does not support structured output, cannot use output_json_schema"
 53                )
 54            # Langchain expects title/description to be at top level, on top of json schema
 55            output_schema = self.kiln_task.output_schema()
 56            if output_schema is None:
 57                raise ValueError(
 58                    f"output_json_schema is not valid json: {self.kiln_task.output_json_schema}"
 59                )
 60            output_schema["title"] = "task_response"
 61            output_schema["description"] = "A response from the task"
 62            self.model = self.model.with_structured_output(
 63                output_schema, include_raw=True
 64            )
 65
 66    def adapter_specific_instructions(self) -> str | None:
 67        # TODO: would be better to explicitly use bind_tools:tool_choice="task_response" here
 68        if self.has_structured_output():
 69            return "Always respond with a tool call. Never respond with a human readable message."
 70        return None
 71
 72    async def _run(self, input: Dict | str) -> Dict | str:
 73        prompt = self.build_prompt()
 74        user_msg = self.prompt_builder.build_user_message(input)
 75        messages = [
 76            SystemMessage(content=prompt),
 77            HumanMessage(content=user_msg),
 78        ]
 79        response = self.model.invoke(messages)
 80
 81        if self.has_structured_output():
 82            if (
 83                not isinstance(response, dict)
 84                or "parsed" not in response
 85                or not isinstance(response["parsed"], dict)
 86            ):
 87                raise RuntimeError(f"structured response not returned: {response}")
 88            structured_response = response["parsed"]
 89            return self._munge_response(structured_response)
 90        else:
 91            if not isinstance(response, BaseMessage):
 92                raise RuntimeError(f"response is not a BaseMessage: {response}")
 93            text_content = response.content
 94            if not isinstance(text_content, str):
 95                raise RuntimeError(f"response is not a string: {text_content}")
 96            return text_content
 97
 98    def adapter_info(self) -> AdapterInfo:
 99        return AdapterInfo(
100            model_name=self.model_name,
101            model_provider=self.model_provider,
102            adapter_name="kiln_langchain_adapter",
103            prompt_builder_name=self.prompt_builder.__class__.prompt_builder_name(),
104        )
105
106    def _munge_response(self, response: Dict) -> Dict:
107        # Mistral Large tool calling format is a bit different. Convert to standard format.
108        if (
109            "name" in response
110            and response["name"] == "task_response"
111            and "arguments" in response
112        ):
113            return response["arguments"]
114        return response
class LangChainPromptAdapter(kiln_ai.adapters.base_adapter.BaseAdapter):
 14class LangChainPromptAdapter(BaseAdapter):
 15    def __init__(
 16        self,
 17        kiln_task: datamodel.Task,
 18        custom_model: BaseChatModel | None = None,
 19        model_name: str | None = None,
 20        provider: str | None = None,
 21        prompt_builder: BasePromptBuilder | None = None,
 22    ):
 23        super().__init__(kiln_task, prompt_builder=prompt_builder)
 24        if custom_model is not None:
 25            self.model = custom_model
 26
 27            # Attempt to infer model provider and name from custom model
 28            self.model_provider = "custom.langchain:" + custom_model.__class__.__name__
 29            self.model_name = "custom.langchain:unknown_model"
 30            if hasattr(custom_model, "model_name") and isinstance(
 31                getattr(custom_model, "model_name"), str
 32            ):
 33                self.model_name = "custom.langchain:" + getattr(
 34                    custom_model, "model_name"
 35                )
 36            if hasattr(custom_model, "model") and isinstance(
 37                getattr(custom_model, "model"), str
 38            ):
 39                self.model_name = "custom.langchain:" + getattr(custom_model, "model")
 40        elif model_name is not None:
 41            self.model = langchain_model_from(model_name, provider)
 42            self.model_name = model_name
 43            self.model_provider = provider or "custom.langchain.default_provider"
 44        else:
 45            raise ValueError(
 46                "model_name and provider must be provided if custom_model is not provided"
 47            )
 48        if self.has_structured_output():
 49            if not hasattr(self.model, "with_structured_output") or not callable(
 50                getattr(self.model, "with_structured_output")
 51            ):
 52                raise ValueError(
 53                    f"model {self.model} does not support structured output, cannot use output_json_schema"
 54                )
 55            # Langchain expects title/description to be at top level, on top of json schema
 56            output_schema = self.kiln_task.output_schema()
 57            if output_schema is None:
 58                raise ValueError(
 59                    f"output_json_schema is not valid json: {self.kiln_task.output_json_schema}"
 60                )
 61            output_schema["title"] = "task_response"
 62            output_schema["description"] = "A response from the task"
 63            self.model = self.model.with_structured_output(
 64                output_schema, include_raw=True
 65            )
 66
 67    def adapter_specific_instructions(self) -> str | None:
 68        # TODO: would be better to explicitly use bind_tools:tool_choice="task_response" here
 69        if self.has_structured_output():
 70            return "Always respond with a tool call. Never respond with a human readable message."
 71        return None
 72
 73    async def _run(self, input: Dict | str) -> Dict | str:
 74        prompt = self.build_prompt()
 75        user_msg = self.prompt_builder.build_user_message(input)
 76        messages = [
 77            SystemMessage(content=prompt),
 78            HumanMessage(content=user_msg),
 79        ]
 80        response = self.model.invoke(messages)
 81
 82        if self.has_structured_output():
 83            if (
 84                not isinstance(response, dict)
 85                or "parsed" not in response
 86                or not isinstance(response["parsed"], dict)
 87            ):
 88                raise RuntimeError(f"structured response not returned: {response}")
 89            structured_response = response["parsed"]
 90            return self._munge_response(structured_response)
 91        else:
 92            if not isinstance(response, BaseMessage):
 93                raise RuntimeError(f"response is not a BaseMessage: {response}")
 94            text_content = response.content
 95            if not isinstance(text_content, str):
 96                raise RuntimeError(f"response is not a string: {text_content}")
 97            return text_content
 98
 99    def adapter_info(self) -> AdapterInfo:
100        return AdapterInfo(
101            model_name=self.model_name,
102            model_provider=self.model_provider,
103            adapter_name="kiln_langchain_adapter",
104            prompt_builder_name=self.prompt_builder.__class__.prompt_builder_name(),
105        )
106
107    def _munge_response(self, response: Dict) -> Dict:
108        # Mistral Large tool calling format is a bit different. Convert to standard format.
109        if (
110            "name" in response
111            and response["name"] == "task_response"
112            and "arguments" in response
113        ):
114            return response["arguments"]
115        return response

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

Example:

class CustomAdapter(BaseAdapter):
    async def _run(self, input: Dict | str) -> Dict | str:
        # Implementation for specific model
        pass

    def adapter_info(self) -> AdapterInfo:
        return AdapterInfo(
            adapter_name="custom",
            model_name="model-1",
            model_provider="provider",
            prompt_builder_name="simple"
        )
LangChainPromptAdapter( kiln_task: kiln_ai.datamodel.Task, custom_model: langchain_core.language_models.chat_models.BaseChatModel | None = None, model_name: str | None = None, provider: str | None = None, prompt_builder: kiln_ai.adapters.prompt_builders.BasePromptBuilder | None = None)
15    def __init__(
16        self,
17        kiln_task: datamodel.Task,
18        custom_model: BaseChatModel | None = None,
19        model_name: str | None = None,
20        provider: str | None = None,
21        prompt_builder: BasePromptBuilder | None = None,
22    ):
23        super().__init__(kiln_task, prompt_builder=prompt_builder)
24        if custom_model is not None:
25            self.model = custom_model
26
27            # Attempt to infer model provider and name from custom model
28            self.model_provider = "custom.langchain:" + custom_model.__class__.__name__
29            self.model_name = "custom.langchain:unknown_model"
30            if hasattr(custom_model, "model_name") and isinstance(
31                getattr(custom_model, "model_name"), str
32            ):
33                self.model_name = "custom.langchain:" + getattr(
34                    custom_model, "model_name"
35                )
36            if hasattr(custom_model, "model") and isinstance(
37                getattr(custom_model, "model"), str
38            ):
39                self.model_name = "custom.langchain:" + getattr(custom_model, "model")
40        elif model_name is not None:
41            self.model = langchain_model_from(model_name, provider)
42            self.model_name = model_name
43            self.model_provider = provider or "custom.langchain.default_provider"
44        else:
45            raise ValueError(
46                "model_name and provider must be provided if custom_model is not provided"
47            )
48        if self.has_structured_output():
49            if not hasattr(self.model, "with_structured_output") or not callable(
50                getattr(self.model, "with_structured_output")
51            ):
52                raise ValueError(
53                    f"model {self.model} does not support structured output, cannot use output_json_schema"
54                )
55            # Langchain expects title/description to be at top level, on top of json schema
56            output_schema = self.kiln_task.output_schema()
57            if output_schema is None:
58                raise ValueError(
59                    f"output_json_schema is not valid json: {self.kiln_task.output_json_schema}"
60                )
61            output_schema["title"] = "task_response"
62            output_schema["description"] = "A response from the task"
63            self.model = self.model.with_structured_output(
64                output_schema, include_raw=True
65            )
def adapter_specific_instructions(self) -> str | None:
67    def adapter_specific_instructions(self) -> str | None:
68        # TODO: would be better to explicitly use bind_tools:tool_choice="task_response" here
69        if self.has_structured_output():
70            return "Always respond with a tool call. Never respond with a human readable message."
71        return None
def adapter_info(self) -> kiln_ai.adapters.base_adapter.AdapterInfo:
 99    def adapter_info(self) -> AdapterInfo:
100        return AdapterInfo(
101            model_name=self.model_name,
102            model_provider=self.model_provider,
103            adapter_name="kiln_langchain_adapter",
104            prompt_builder_name=self.prompt_builder.__class__.prompt_builder_name(),
105        )