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

kiln_ai/adapters/__init__.py

@@ -3,31 +3,31 @@
 
 Adapters are used to connect Kiln to external systems, or to add new functionality to Kiln.
 
-BaseAdapter is extensible, and used for adding adapters that provide AI functionality. There's currently a LangChain adapter which provides a bridge to LangChain.
+Model adapters are used to call AI models, like Ollama, OpenAI, etc.
 
 The ml_model_list submodule contains a list of models that can be used for machine learning tasks. More can easily be added, but we keep a list here of models that are known to work well with Kiln's structured data and tool calling systems.
 
 The prompt_builders submodule contains classes that build prompts for use with the AI agents.
 
 The repair submodule contains an adapter for the repair task.
+
+The parser submodule contains parsers for the output of the AI models.
 """
 
 from . import (
-    base_adapter,
     data_gen,
     fine_tune,
-    langchain_adapters,
     ml_model_list,
+    model_adapters,
     prompt_builders,
     repair,
 )
 
 __all__ = [
-    "base_adapter",
-    "langchain_adapters",
+    "model_adapters",
+    "data_gen",
+    "fine_tune",
     "ml_model_list",
     "prompt_builders",
     "repair",
-    "data_gen",
-    "fine_tune",
 ]

kiln_ai/adapters/adapter_registry.py

@@ -1,17 +1,89 @@
+from os import getenv
+
 from kiln_ai import datamodel
-from kiln_ai.adapters.base_adapter import BaseAdapter
-from kiln_ai.adapters.langchain_adapters import LangchainAdapter
+from kiln_ai.adapters.ml_model_list import ModelProviderName
+from kiln_ai.adapters.model_adapters.base_adapter import BaseAdapter
+from kiln_ai.adapters.model_adapters.langchain_adapters import LangchainAdapter
+from kiln_ai.adapters.model_adapters.openai_model_adapter import (
+    OpenAICompatibleAdapter,
+    OpenAICompatibleConfig,
+)
 from kiln_ai.adapters.prompt_builders import BasePromptBuilder
+from kiln_ai.adapters.provider_tools import core_provider, openai_compatible_config
+from kiln_ai.utils.config import Config
+from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
 
 def adapter_for_task(
     kiln_task: datamodel.Task,
-    model_name: str | None = None,
-    provider: str | None = None,
+    model_name: str,
+    provider: ModelProviderName,
     prompt_builder: BasePromptBuilder | None = None,
     tags: list[str] | None = None,
 ) -> BaseAdapter:
-    # We use langchain for everything right now, but can add any others here
+    # Get the provider to run. For things like the fine-tune provider, we want to run the underlying provider
+    core_provider_name = core_provider(model_name, provider)
+
+    match core_provider_name:
+        case ModelProviderName.openrouter:
+            return OpenAICompatibleAdapter(
+                kiln_task=kiln_task,
+                config=OpenAICompatibleConfig(
+                    base_url=getenv("OPENROUTER_BASE_URL")
+                    or "https://openrouter.ai/api/v1",
+                    api_key=Config.shared().open_router_api_key,
+                    model_name=model_name,
+                    provider_name=provider,
+                    openrouter_style_reasoning=True,
+                    default_headers={
+                        "HTTP-Referer": "https://getkiln.ai/openrouter",
+                        "X-Title": "KilnAI",
+                    },
+                ),
+                prompt_builder=prompt_builder,
+                tags=tags,
+            )
+        case ModelProviderName.openai:
+            return OpenAICompatibleAdapter(
+                kiln_task=kiln_task,
+                config=OpenAICompatibleConfig(
+                    api_key=Config.shared().open_ai_api_key,
+                    model_name=model_name,
+                    provider_name=provider,
+                ),
+                prompt_builder=prompt_builder,
+                tags=tags,
+            )
+        case ModelProviderName.openai_compatible:
+            config = openai_compatible_config(model_name)
+            return OpenAICompatibleAdapter(
+                kiln_task=kiln_task,
+                config=config,
+                prompt_builder=prompt_builder,
+                tags=tags,
+            )
+        # Use LangchainAdapter for the rest
+        case ModelProviderName.groq:
+            pass
+        case ModelProviderName.amazon_bedrock:
+            pass
+        case ModelProviderName.ollama:
+            pass
+        case ModelProviderName.fireworks_ai:
+            pass
+        # These are virtual providers that should have mapped to an actual provider in core_provider
+        case ModelProviderName.kiln_fine_tune:
+            raise ValueError(
+                "Fine tune is not a supported core provider. It should map to an actual provider."
+            )
+        case ModelProviderName.kiln_custom_registry:
+            raise ValueError(
+                "Custom openai compatible provider is not a supported core provider. It should map to an actual provider."
+            )
+        case _:
+            raise_exhaustive_enum_error(core_provider_name)
+
+    # We use langchain for all others right now, but moving off it as we touch anything.
     return LangchainAdapter(
         kiln_task,
         model_name=model_name,

kiln_ai/adapters/data_gen/data_gen_task.py

@@ -55,7 +55,7 @@ class DataGenCategoriesTaskInput(BaseModel):
             num_subtopics=num_subtopics,
             human_guidance=human_guidance,
             existing_topics=existing_topics,
-            system_prompt=prompt_builder.build_prompt(),
+            system_prompt=prompt_builder.build_prompt(include_json_instructions=False),
         )
 
 
@@ -132,7 +132,7 @@ class DataGenSampleTaskInput(BaseModel):
             topic=topic,
             num_samples=num_samples,
             human_guidance=human_guidance,
-            system_prompt=prompt_builder.build_prompt(),
+            system_prompt=prompt_builder.build_prompt(include_json_instructions=False),
         )
 
 
@@ -163,7 +163,7 @@ def list_json_schema_for_task(task: Task) -> str:
         "required": ["generated_samples"],
     }
 
-    return json.dumps(top_level_schema)
+    return json.dumps(top_level_schema, ensure_ascii=False)
 
 
 class DataGenSampleTask(Task, parent_of={}):

kiln_ai/adapters/data_gen/test_data_gen_task.py

@@ -180,7 +180,7 @@ def test_data_gen_sample_task_initialization(base_task):
     }
 
 
-def test_list_json_schema_for_task_with_output_schema(base_task):
+def test_list_json_schema_for_task_with_input_schema(base_task):
     # Arrange
     base_task.input_json_schema = json.dumps(
         {
@@ -202,9 +202,29 @@ def test_list_json_schema_for_task_with_output_schema(base_task):
     assert generated_samples_schema["items"]["properties"]["age"]["type"] == "integer"
 
 
-def test_list_json_schema_for_task_without_output_schema(base_task):
+def test_list_json_schema_for_task_with_input_schema_non_ascii(base_task):
     # Arrange
-    base_task.output_json_schema = None
+    base_task.input_json_schema = json.dumps(
+        {
+            "type": "object",
+            "properties": {
+                "名字": {"type": "string"},
+                "年齢": {"type": "integer"},
+            },
+        }
+    )
+
+    # Act
+    schema = list_json_schema_for_task(base_task)
+
+    # Assert
+    assert "名字" in schema
+    assert "年齢" in schema
+
+
+def test_list_json_schema_for_task_without_input_schema(base_task):
+    # Arrange
+    base_task.input_json_schema = None
 
     # Act
     schema = list_json_schema_for_task(base_task)

kiln_ai/adapters/fine_tune/base_finetune.py

@@ -4,7 +4,7 @@ from typing import Literal
 from pydantic import BaseModel
 
 from kiln_ai.adapters.ml_model_list import built_in_models
-from kiln_ai.datamodel import DatasetSplit, FineTuneStatusType
+from kiln_ai.datamodel import DatasetSplit, FinetuneDataStrategy, FineTuneStatusType
 from kiln_ai.datamodel import Finetune as FinetuneModel
 from kiln_ai.utils.name_generator import generate_memorable_name
 
@@ -56,6 +56,8 @@ class BaseFinetuneAdapter(ABC):
         provider_base_model_id: str,
         train_split_name: str,
         system_message: str,
+        thinking_instructions: str | None,
+        data_strategy: FinetuneDataStrategy,
         parameters: dict[str, str | int | float | bool] = {},
         name: str | None = None,
         description: str | None = None,
@@ -100,7 +102,9 @@ class BaseFinetuneAdapter(ABC):
             validation_split_name=validation_split_name,
             parameters=parameters,
             system_message=system_message,
+            thinking_instructions=thinking_instructions,
             parent=parent_task,
+            data_strategy=data_strategy,
         )
 
         adapter = cls(datamodel)