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
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 )