kiln-ai 0.0.4__py3-none-any.whl → 0.5.1__py3-none-any.whl

kiln_ai/adapters/base_adapter.py

@@ -0,0 +1,168 @@
+import json
+from abc import ABCMeta, abstractmethod
+from dataclasses import dataclass
+from typing import Dict
+
+from kiln_ai.datamodel import (
+    DataSource,
+    DataSourceType,
+    Task,
+    TaskOutput,
+    TaskRun,
+)
+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:
+    adapter_name: str
+    model_name: str
+    model_provider: str
+    prompt_builder_name: str
+
+
+class BaseAdapter(metaclass=ABCMeta):
+    def __init__(
+        self, kiln_task: Task, prompt_builder: BasePromptBuilder | None = None
+    ):
+        self.prompt_builder = prompt_builder or SimplePromptBuilder(kiln_task)
+        self.kiln_task = kiln_task
+        self.output_schema = self.kiln_task.output_json_schema
+        self.input_schema = self.kiln_task.input_json_schema
+
+    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.kiln_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:
+        # 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
+        result = await self._run(input)
+
+        # validate output
+        if self.output_schema is not None:
+            if not isinstance(result, dict):
+                raise RuntimeError(f"structured response is not a dict: {result}")
+            validate_schema(result, self.output_schema)
+        else:
+            if not isinstance(result, str):
+                raise RuntimeError(
+                    f"response is not a string for non-structured task: {result}"
+                )
+
+        # Generate the run and output
+        run = self.generate_run(input, input_source, result)
+
+        # 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:
+            run.save_to_file()
+        else:
+            # Clear the ID to indicate it's not persisted
+            run.id = None
+
+        return run
+
+    def has_structured_output(self) -> bool:
+        return self.output_schema is not None
+
+    @abstractmethod
+    def adapter_info(self) -> AdapterInfo:
+        pass
+
+    @abstractmethod
+    async def _run(self, input: Dict | str) -> Dict | str:
+        pass
+
+    def build_prompt(self) -> str:
+        prompt = self.prompt_builder.build_prompt()
+        adapter_instructions = self.adapter_specific_instructions()
+        if adapter_instructions is not None:
+            prompt += f"# Format Instructions\n\n{adapter_instructions}\n\n"
+        return prompt
+
+    # override for adapter specific instructions (e.g. tool calling, json format, etc)
+    def adapter_specific_instructions(self) -> str | None:
+        return None
+
+    # create a run and task output
+    def generate_run(
+        self, input: Dict | str, input_source: DataSource | None, output: Dict | str
+    ) -> TaskRun:
+        # Convert input and output to JSON strings if they are dictionaries
+        input_str = json.dumps(input) if isinstance(input, dict) else input
+        output_str = json.dumps(output) if isinstance(output, dict) else 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.kiln_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(),
+                ),
+            ),
+        )
+
+        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]:
+        props = {}
+
+        # adapter info
+        adapter_info = self.adapter_info()
+        props["adapter_name"] = adapter_info.adapter_name
+        props["model_name"] = adapter_info.model_name
+        props["model_provider"] = adapter_info.model_provider
+        props["prompt_builder_name"] = adapter_info.prompt_builder_name
+
+        return props

kiln_ai/adapters/langchain_adapters.py

@@ -0,0 +1,113 @@
+from typing import Dict
+
+import kiln_ai.datamodel as datamodel
+from kiln_ai.adapters.prompt_builders import SimplePromptBuilder
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.messages.base import BaseMessage
+
+from .base_adapter import AdapterInfo, BaseAdapter, BasePromptBuilder
+from .ml_model_list import langchain_model_from
+
+
+class LangChainPromptAdapter(BaseAdapter):
+    def __init__(
+        self,
+        kiln_task: datamodel.Task,
+        custom_model: BaseChatModel | None = None,
+        model_name: str | None = None,
+        provider: str | None = None,
+        prompt_builder: BasePromptBuilder | None = None,
+    ):
+        super().__init__(kiln_task, prompt_builder=prompt_builder)
+        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")
+        elif model_name is not None:
+            self.model = langchain_model_from(model_name, provider)
+            self.model_name = model_name
+            self.model_provider = provider or "custom.langchain.default_provider"
+        else:
+            raise ValueError(
+                "model_name and provider must be provided if custom_model is not provided"
+            )
+        if self.has_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.kiln_task.output_schema()
+            if output_schema is None:
+                raise ValueError(
+                    f"output_json_schema is not valid json: {self.kiln_task.output_json_schema}"
+                )
+            output_schema["title"] = "task_response"
+            output_schema["description"] = "A response from the task"
+            self.model = self.model.with_structured_output(
+                output_schema, include_raw=True
+            )
+
+    def adapter_specific_instructions(self) -> str | None:
+        # TODO: would be better to explicitly use bind_tools:tool_choice="task_response" here
+        if self.has_structured_output():
+            return "Always respond with a tool call. Never respond with a human readable message."
+        return None
+
+    async def _run(self, input: Dict | str) -> Dict | str:
+        prompt = self.build_prompt()
+        user_msg = self.prompt_builder.build_user_message(input)
+        messages = [
+            SystemMessage(content=prompt),
+            HumanMessage(content=user_msg),
+        ]
+        response = self.model.invoke(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}")
+            structured_response = response["parsed"]
+            return self._munge_response(structured_response)
+        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 text_content
+
+    def adapter_info(self) -> AdapterInfo:
+        return AdapterInfo(
+            model_name=self.model_name,
+            model_provider=self.model_provider,
+            adapter_name="kiln_langchain_adapter",
+            prompt_builder_name=self.prompt_builder.__class__.prompt_builder_name(),
+        )
+
+    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