kiln-ai 0.7.1__tar.gz → 0.8.1__tar.gz

{kiln_ai-0.7.1 → kiln_ai-0.8.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kiln-ai
-Version: 0.7.1
+Version: 0.8.1
 Summary: Kiln AI
 Project-URL: Homepage, https://getkiln.ai
 Project-URL: Repository, https://github.com/Kiln-AI/kiln

{kiln_ai-0.7.1 → kiln_ai-0.8.1}/kiln_ai/adapters/adapter_registry.py

@@ -9,6 +9,7 @@ def adapter_for_task(
     model_name: str | None = None,
     provider: str | None = None,
     prompt_builder: BasePromptBuilder | None = None,
+    tags: list[str] | None = None,
 ) -> BaseAdapter:
     # We use langchain for everything right now, but can add any others here
     return LangchainAdapter(
@@ -16,4 +17,5 @@ def adapter_for_task(
         model_name=model_name,
         provider=provider,
         prompt_builder=prompt_builder,
+        tags=tags,
     )

{kiln_ai-0.7.1 → kiln_ai-0.8.1}/kiln_ai/adapters/base_adapter.py

@@ -45,12 +45,16 @@ class BaseAdapter(metaclass=ABCMeta):
     """
 
     def __init__(
-        self, kiln_task: Task, prompt_builder: BasePromptBuilder | None = None
+        self,
+        kiln_task: Task,
+        prompt_builder: BasePromptBuilder | None = None,
+        tags: list[str] | 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
+        self.default_tags = tags
 
     async def invoke_returning_raw(
         self,
@@ -148,6 +152,7 @@ class BaseAdapter(metaclass=ABCMeta):
                 ),
             ),
             intermediate_outputs=run_output.intermediate_outputs,
+            tags=self.default_tags or [],
         )
 
         exclude_fields = {

{kiln_ai-0.7.1 → kiln_ai-0.8.1}/kiln_ai/adapters/langchain_adapters.py

@@ -39,8 +39,9 @@ class LangchainAdapter(BaseAdapter):
         model_name: str | None = None,
         provider: str | None = None,
         prompt_builder: BasePromptBuilder | None = None,
+        tags: list[str] | None = None,
     ):
-        super().__init__(kiln_task, prompt_builder=prompt_builder)
+        super().__init__(kiln_task, prompt_builder=prompt_builder, tags=tags)
         if custom_model is not None:
             self._model = custom_model
 
@@ -198,6 +199,9 @@ async def langchain_model_from_provider(
     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]
+    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]
     elif provider.name == ModelProviderName.groq:
         api_key = Config.shared().groq_api_key
         if api_key is None:

{kiln_ai-0.7.1 → kiln_ai-0.8.1}/kiln_ai/adapters/ml_model_list.py

@@ -23,6 +23,7 @@ class ModelProviderName(str, Enum):
     fireworks_ai = "fireworks_ai"
     kiln_fine_tune = "kiln_fine_tune"
     kiln_custom_registry = "kiln_custom_registry"
+    openai_compatible = "openai_compatible"
 
 
 class ModelFamily(str, Enum):
@@ -522,6 +523,12 @@ built_in_models: List[KilnModel] = [
                     }
                 },
             ),
+            KilnModelProvider(
+                name=ModelProviderName.groq,
+                supports_structured_output=True,
+                supports_data_gen=True,
+                provider_options={"model": "llama-3.3-70b-versatile"},
+            ),
             KilnModelProvider(
                 name=ModelProviderName.ollama,
                 provider_options={"model": "llama3.3"},
@@ -530,6 +537,8 @@ built_in_models: List[KilnModel] = [
                 name=ModelProviderName.fireworks_ai,
                 # Finetuning not live yet
                 # provider_finetune_id="accounts/fireworks/models/llama-v3p3-70b-instruct",
+                supports_structured_output=True,
+                supports_data_gen=True,
                 provider_options={
                     "model": "accounts/fireworks/models/llama-v3p3-70b-instruct"
                 },

{kiln_ai-0.7.1 → kiln_ai-0.8.1}/kiln_ai/adapters/provider_tools.py

@@ -108,6 +108,9 @@ async def kiln_model_provider_from(
     if provider_name == ModelProviderName.kiln_fine_tune:
         return finetune_provider_model(name)
 
+    if provider_name == ModelProviderName.openai_compatible:
+        return openai_compatible_provider_model(name)
+
     built_in_model = await builtin_model_from(name, provider_name)
     if built_in_model:
         return built_in_model
@@ -136,6 +139,45 @@ async def kiln_model_provider_from(
 finetune_cache: dict[str, KilnModelProvider] = {}
 
 
+def openai_compatible_provider_model(
+    model_id: str,
+) -> KilnModelProvider:
+    try:
+        openai_provider_name, model_id = model_id.split("::")
+    except Exception:
+        raise ValueError(f"Invalid openai compatible model ID: {model_id}")
+
+    openai_compatible_providers = Config.shared().openai_compatible_providers or []
+    provider = next(
+        filter(
+            lambda p: p.get("name") == openai_provider_name, openai_compatible_providers
+        ),
+        None,
+    )
+    if provider is None:
+        raise ValueError(f"OpenAI compatible provider {openai_provider_name} not found")
+
+    # API key optional some providers don't use it
+    api_key = provider.get("api_key")
+    base_url = provider.get("base_url")
+    if base_url is None:
+        raise ValueError(
+            f"OpenAI compatible provider {openai_provider_name} has no base URL"
+        )
+
+    return KilnModelProvider(
+        name=ModelProviderName.openai_compatible,
+        provider_options={
+            "model": model_id,
+            "api_key": api_key,
+            "openai_api_base": base_url,
+        },
+        supports_structured_output=False,
+        supports_data_gen=False,
+        untested_model=True,
+    )
+
+
 def finetune_provider_model(
     model_id: str,
 ) -> KilnModelProvider:
@@ -228,6 +270,8 @@ def provider_name_from_id(id: str) -> str:
                 return "Fireworks AI"
             case ModelProviderName.kiln_custom_registry:
                 return "Custom Models"
+            case ModelProviderName.openai_compatible:
+                return "OpenAI Compatible"
             case _:
                 # triggers pyright warning if I miss a case
                 raise_exhaustive_error(enum_id)
@@ -266,6 +310,10 @@ def provider_options_for_custom_model(
             raise ValueError(
                 "Fine tuned models should populate provider options via another path"
             )
+        case ModelProviderName.openai_compatible:
+            raise ValueError(
+                "OpenAI compatible models should populate provider options via another path"
+            )
         case _:
             # triggers pyright warning if I miss a case
             raise_exhaustive_error(enum_id)

{kiln_ai-0.7.1 → kiln_ai-0.8.1}/kiln_ai/adapters/test_provider_tools.py

@@ -15,6 +15,7 @@ from kiln_ai.adapters.provider_tools import (
     finetune_provider_model,
     get_model_and_provider,
     kiln_model_provider_from,
+    openai_compatible_provider_model,
     provider_enabled,
     provider_name_from_id,
     provider_options_for_custom_model,
@@ -64,6 +65,25 @@ def mock_finetune():
         yield mock
 
 
+@pytest.fixture
+def mock_shared_config():
+    with patch("kiln_ai.adapters.provider_tools.Config.shared") as mock:
+        config = Mock()
+        config.openai_compatible_providers = [
+            {
+                "name": "test_provider",
+                "base_url": "https://api.test.com",
+                "api_key": "test-key",
+            },
+            {
+                "name": "no_key_provider",
+                "base_url": "https://api.nokey.com",
+            },
+        ]
+        mock.return_value = config
+        yield mock
+
+
 def test_check_provider_warnings_no_warning(mock_config):
     mock_config.return_value = "some_value"
 
@@ -529,3 +549,78 @@ def test_finetune_provider_model_fireworks_provider(
     assert provider.adapter_options == {
         "langchain": {"with_structured_output_options": {"method": "json_mode"}}
     }
+
+
+def test_openai_compatible_provider_model_success(mock_shared_config):
+    """Test successful creation of an OpenAI compatible provider"""
+    model_id = "test_provider::gpt-4"
+
+    provider = openai_compatible_provider_model(model_id)
+
+    assert provider.name == ModelProviderName.openai_compatible
+    assert provider.provider_options == {
+        "model": "gpt-4",
+        "api_key": "test-key",
+        "openai_api_base": "https://api.test.com",
+    }
+    assert provider.supports_structured_output is False
+    assert provider.supports_data_gen is False
+    assert provider.untested_model is True
+
+
+def test_openai_compatible_provider_model_no_api_key(mock_shared_config):
+    """Test provider creation without API key (should work as some providers don't require it)"""
+    model_id = "no_key_provider::gpt-4"
+
+    provider = openai_compatible_provider_model(model_id)
+
+    assert provider.name == ModelProviderName.openai_compatible
+    assert provider.provider_options == {
+        "model": "gpt-4",
+        "api_key": None,
+        "openai_api_base": "https://api.nokey.com",
+    }
+
+
+def test_openai_compatible_provider_model_invalid_id():
+    """Test handling of invalid model ID format"""
+    with pytest.raises(ValueError) as exc_info:
+        openai_compatible_provider_model("invalid-id-format")
+    assert (
+        str(exc_info.value) == "Invalid openai compatible model ID: invalid-id-format"
+    )
+
+
+def test_openai_compatible_provider_model_no_providers(mock_shared_config):
+    """Test handling when no providers are configured"""
+    mock_shared_config.return_value.openai_compatible_providers = None
+
+    with pytest.raises(ValueError) as exc_info:
+        openai_compatible_provider_model("test_provider::gpt-4")
+    assert str(exc_info.value) == "OpenAI compatible provider test_provider not found"
+
+
+def test_openai_compatible_provider_model_provider_not_found(mock_shared_config):
+    """Test handling of non-existent provider"""
+    with pytest.raises(ValueError) as exc_info:
+        openai_compatible_provider_model("unknown_provider::gpt-4")
+    assert (
+        str(exc_info.value) == "OpenAI compatible provider unknown_provider not found"
+    )
+
+
+def test_openai_compatible_provider_model_no_base_url(mock_shared_config):
+    """Test handling of provider without base URL"""
+    mock_shared_config.return_value.openai_compatible_providers = [
+        {
+            "name": "test_provider",
+            "api_key": "test-key",
+        }
+    ]
+
+    with pytest.raises(ValueError) as exc_info:
+        openai_compatible_provider_model("test_provider::gpt-4")
+    assert (
+        str(exc_info.value)
+        == "OpenAI compatible provider test_provider has no base URL"
+    )

{kiln_ai-0.7.1 → kiln_ai-0.8.1}/kiln_ai/datamodel/__init__.py

@@ -49,9 +49,15 @@ __all__ = [
     "DataSource",
     "DataSourceType",
     "DataSourceProperty",
+    "Finetune",
+    "FineTuneStatusType",
     "TaskOutputRatingType",
     "TaskRequirement",
     "TaskDeterminism",
+    "DatasetSplitDefinition",
+    "DatasetSplit",
+    "RequirementRating",
+    "TaskRequirement",
     "strict_mode",
     "set_strict_mode",
 ]
@@ -85,30 +91,71 @@ class TaskOutputRatingType(str, Enum):
     """Defines the types of rating systems available for task outputs."""
 
     five_star = "five_star"
+    pass_fail = "pass_fail"
+    pass_fail_critical = "pass_fail_critical"
     custom = "custom"
 
 
+class RequirementRating(BaseModel):
+    """Rating for a specific requirement within a task output."""
+
+    value: float = Field(
+        description="The rating value. Interpretation depends on rating type"
+    )
+    type: TaskOutputRatingType = Field(description="The type of rating")
+
+
 class TaskOutputRating(KilnBaseModel):
     """
     A rating for a task output, including an overall rating and ratings for each requirement.
 
-    Only supports five star ratings for now, but extensible for custom values.
+    Supports:
+    - five_star: 1-5 star ratings
+    - pass_fail: boolean pass/fail (1.0 = pass, 0.0 = fail)
+    - pass_fail_critical: tri-state (1.0 = pass, 0.0 = fail, -1.0 = critical fail)
     """
 
     type: TaskOutputRatingType = Field(default=TaskOutputRatingType.five_star)
     value: float | None = Field(
-        description="The overall rating value (typically 1-5 stars).",
+        description="The rating value. Interpretation depends on rating type:\n- five_star: 1-5 stars\n- pass_fail: 1.0 (pass) or 0.0 (fail)\n- pass_fail_critical: 1.0 (pass), 0.0 (fail), or -1.0 (critical fail)",
         default=None,
     )
-    requirement_ratings: Dict[ID_TYPE, float] = Field(
+    requirement_ratings: Dict[ID_TYPE, RequirementRating] = Field(
         default={},
-        description="The ratings of the requirements of the task. The keys are the ids of the requirements. The values are the ratings (typically 1-5 stars).",
+        description="The ratings of the requirements of the task.",
     )
 
+    # Previously we stored rating values as a dict of floats, but now we store them as RequirementRating objects.
+    @model_validator(mode="before")
+    def upgrade_old_format(cls, data: dict) -> dict:
+        if not isinstance(data, dict):
+            return data
+
+        # Check if we have the old format (dict of floats)
+        req_ratings = data.get("requirement_ratings", {})
+        if req_ratings and all(
+            isinstance(v, (int, float)) for v in req_ratings.values()
+        ):
+            # Convert each float to a RequirementRating object
+            # all ratings are five star at the point we used this format
+            data["requirement_ratings"] = {
+                k: {"value": v, "type": TaskOutputRatingType.five_star}
+                for k, v in req_ratings.items()
+            }
+
+        return data
+
     # Used to select high quality outputs for example selection (MultiShotPromptBuilder, etc)
     def is_high_quality(self) -> bool:
+        if self.value is None:
+            return False
+
         if self.type == TaskOutputRatingType.five_star:
-            return self.value is not None and self.value >= 4
+            return self.value >= 4
+        elif self.type == TaskOutputRatingType.pass_fail:
+            return self.value == 1.0
+        elif self.type == TaskOutputRatingType.pass_fail_critical:
+            return self.value == 1.0
         return False
 
     @model_validator(mode="after")
@@ -116,24 +163,61 @@ class TaskOutputRating(KilnBaseModel):
         if self.type not in TaskOutputRatingType:
             raise ValueError(f"Invalid rating type: {self.type}")
 
-        if self.type == TaskOutputRatingType.five_star:
-            if self.value is not None:
-                self._validate_five_star(self.value, "overall rating")
-            for req_id, req_rating in self.requirement_ratings.items():
-                self._validate_five_star(req_rating, f"requirement rating for {req_id}")
+        # Overall rating is optional
+        if self.value is not None:
+            self._validate_rating(self.type, self.value, "overall rating")
+
+        for req_id, req_rating in self.requirement_ratings.items():
+            self._validate_rating(
+                req_rating.type,
+                req_rating.value,
+                f"requirement rating for req ID: {req_id}",
+            )
 
         return self
 
-    def _validate_five_star(self, rating: float, rating_name: str) -> None:
-        if not isinstance(rating, float) or not rating.is_integer():
+    def _validate_rating(
+        self, type: TaskOutputRatingType, rating: float | None, rating_name: str
+    ) -> None:
+        if type == TaskOutputRatingType.five_star:
+            self._validate_five_star(rating, rating_name)
+        elif type == TaskOutputRatingType.pass_fail:
+            self._validate_pass_fail(rating, rating_name)
+        elif type == TaskOutputRatingType.pass_fail_critical:
+            self._validate_pass_fail_critical(rating, rating_name)
+
+    def _validate_five_star(self, rating: float | None, rating_name: str) -> None:
+        if rating is None or not isinstance(rating, float) or not rating.is_integer():
             raise ValueError(
-                f"{rating_name.capitalize()} of type five_star must be an integer value (1.0, 2.0, 3.0, 4.0, or 5.0)"
+                f"{rating_name.capitalize()} of type five_star must be an integer value (1-5)"
             )
         if rating < 1 or rating > 5:
             raise ValueError(
                 f"{rating_name.capitalize()} of type five_star must be between 1 and 5 stars"
             )
 
+    def _validate_pass_fail(self, rating: float | None, rating_name: str) -> None:
+        if rating is None or not isinstance(rating, float) or not rating.is_integer():
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail must be an integer value (0 or 1)"
+            )
+        if rating not in [0, 1]:
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail must be 0 (fail) or 1 (pass)"
+            )
+
+    def _validate_pass_fail_critical(
+        self, rating: float | None, rating_name: str
+    ) -> None:
+        if rating is None or not isinstance(rating, float) or not rating.is_integer():
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail_critical must be an integer value (-1, 0, or 1)"
+            )
+        if rating not in [-1, 0, 1]:
+            raise ValueError(
+                f"{rating_name.capitalize()} of type pass_fail_critical must be -1 (critical fail), 0 (fail), or 1 (pass)"
+            )
+
 
 class TaskOutput(KilnBaseModel):
     """
@@ -381,6 +465,10 @@ class TaskRun(KilnParentedModel):
         default=None,
         description="Intermediate outputs from the task run. Keys are the names of the intermediate output steps (cot=chain of thought, etc), values are the output data.",
     )
+    tags: List[str] = Field(
+        default=[],
+        description="Tags for the task run. Tags are used to categorize task runs for filtering and reporting.",
+    )
 
     def parent_task(self) -> Task | None:
         if not isinstance(self.parent, Task):
@@ -442,6 +530,16 @@ class TaskRun(KilnParentedModel):
             raise ValueError("input_source is required when strict mode is enabled")
         return self
 
+    @model_validator(mode="after")
+    def validate_tags(self) -> Self:
+        for tag in self.tags:
+            if not tag:
+                raise ValueError("Tags cannot be empty strings")
+            if " " in tag:
+                raise ValueError("Tags cannot contain spaces. Try underscores.")
+
+        return self
+
 
 # Define the type alias for clarity
 DatasetFilter = Callable[[TaskRun], bool]
@@ -602,7 +700,7 @@ class TaskRequirement(BaseModel):
     Defines a specific requirement that should be met by task outputs.
 
     Includes an identifier, name, description, instruction for meeting the requirement,
-    and priority level.
+    priority level, and rating type (five_star, pass_fail, pass_fail_critical, custom).
     """
 
     id: ID_TYPE = ID_FIELD
@@ -610,6 +708,7 @@ class TaskRequirement(BaseModel):
     description: str | None = Field(default=None)
     instruction: str = Field(min_length=1)
     priority: Priority = Field(default=Priority.p2)
+    type: TaskOutputRatingType = Field(default=TaskOutputRatingType.five_star)
 
 
 class TaskDeterminism(str, Enum):

{kiln_ai-0.7.1 → kiln_ai-0.8.1}/kiln_ai/datamodel/basemodel.py

@@ -142,14 +142,8 @@ class KilnBaseModel(BaseModel):
             # modified time of file for cache invalidation. From file descriptor so it's atomic w read.
             mtime_ns = os.fstat(file.fileno()).st_mtime_ns
             file_data = file.read()
-            # TODO P2 perf: parsing the JSON twice here.
-            # Once for model_type, once for model. Can't call model_validate with parsed json because enum types break; they get strings instead of enums.
             parsed_json = json.loads(file_data)
-            m = cls.model_validate_json(
-                file_data,
-                strict=True,
-                context={"loading_from_file": True},
-            )
+            m = cls.model_validate(parsed_json, context={"loading_from_file": True})
             if not isinstance(m, cls):
                 raise ValueError(f"Loaded model is not of type {cls.__name__}")
             m._loaded_from_file = True
@@ -471,7 +465,7 @@ class KilnParentModel(KilnBaseModel, metaclass=ABCMeta):
         validation_errors = []
 
         try:
-            instance = cls.model_validate(data, strict=True)
+            instance = cls.model_validate(data)
             if path is not None:
                 instance.path = path
             if parent is not None and isinstance(instance, KilnParentedModel):
@@ -499,7 +493,7 @@ class KilnParentModel(KilnBaseModel, metaclass=ABCMeta):
                             parent_type._validate_nested(**kwargs)
                         elif issubclass(parent_type, KilnParentedModel):
                             # Root node
-                            subinstance = parent_type.model_validate(value, strict=True)
+                            subinstance = parent_type.model_validate(value)
                             if instance is not None:
                                 subinstance.parent = instance
                             if save:

{kiln_ai-0.7.1 → kiln_ai-0.8.1}/kiln_ai/datamodel/test_dataset_split.py

@@ -84,7 +84,7 @@ def task_run():
                 type=DataSourceType.human,
                 properties={"created_by": "test-user"},
             ),
-            rating=TaskOutputRating(rating=5, type=TaskOutputRatingType.five_star),
+            rating=TaskOutputRating(value=5, type=TaskOutputRatingType.five_star),
         ),
     )
 

{kiln_ai-0.7.1 → kiln_ai-0.8.1}/kiln_ai/datamodel/test_models.py

@@ -439,3 +439,52 @@ def test_task_output_source_validation(tmp_path):
         assert os.path.exists(task_missing_output_source)
         task_run = TaskRun.load_from_file(task_missing_output_source)
         assert task_run.output.source is None
+
+
+def test_task_run_tags_validation():
+    # Setup basic output for TaskRun creation
+    output = TaskOutput(
+        output="test output",
+        source=DataSource(
+            type=DataSourceType.synthetic,
+            properties={
+                "model_name": "test-model",
+                "model_provider": "test-provider",
+                "adapter_name": "test-adapter",
+            },
+        ),
+    )
+
+    # Test 1: Valid tags should work
+    task_run = TaskRun(
+        input="test input",
+        output=output,
+        tags=["test_tag", "another_tag", "tag123"],
+    )
+    assert task_run.tags == ["test_tag", "another_tag", "tag123"]
+
+    # Test 2: Empty list of tags should work
+    task_run = TaskRun(
+        input="test input",
+        output=output,
+        tags=[],
+    )
+    assert task_run.tags == []
+
+    # Test 3: Empty string tag should fail
+    with pytest.raises(ValueError) as exc_info:
+        TaskRun(
+            input="test input",
+            output=output,
+            tags=["valid_tag", ""],
+        )
+    assert "Tags cannot be empty strings" in str(exc_info.value)
+
+    # Test 4: Tag with spaces should fail
+    with pytest.raises(ValueError) as exc_info:
+        TaskRun(
+            input="test input",
+            output=output,
+            tags=["valid_tag", "invalid tag"],
+        )
+    assert "Tags cannot contain spaces. Try underscores." in str(exc_info.value)