kiln-ai 0.19.0__py3-none-any.whl → 0.21.0__py3-none-any.whl

kiln_ai/datamodel/task_output.py

@@ -1,6 +1,6 @@
 import json
 from enum import Enum
-from typing import TYPE_CHECKING, Dict, List, Type, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Type, Union
 
 from pydantic import BaseModel, Field, ValidationInfo, model_validator
 from typing_extensions import Self
@@ -8,6 +8,7 @@ from typing_extensions import Self
 from kiln_ai.datamodel.basemodel import ID_TYPE, KilnBaseModel
 from kiln_ai.datamodel.datamodel_enums import TaskOutputRatingType
 from kiln_ai.datamodel.json_schema import validate_schema_with_value_error
+from kiln_ai.datamodel.run_config import RunConfigProperties
 from kiln_ai.datamodel.strict_mode import strict_mode
 from kiln_ai.utils.exhaustive_error import raise_exhaustive_enum_error
 
@@ -199,6 +200,10 @@ class DataSource(BaseModel):
         default={},
         description="Properties describing the data source. For synthetic things like model. For human, the human's name.",
     )
+    run_config: Optional[RunConfigProperties] = Field(
+        default=None,
+        description="The run config used to generate the data, if generated by a running a model in Kiln (only true for type=synthetic).",
+    )
 
     _data_source_properties = [
         DataSourceProperty(

kiln_ai/datamodel/task_run.py

@@ -8,6 +8,7 @@ from kiln_ai.datamodel.basemodel import KilnParentedModel
 from kiln_ai.datamodel.json_schema import validate_schema_with_value_error
 from kiln_ai.datamodel.strict_mode import strict_mode
 from kiln_ai.datamodel.task_output import DataSource, TaskOutput
+from kiln_ai.utils.open_ai_types import ChatCompletionMessageParam
 
 if TYPE_CHECKING:
     from kiln_ai.datamodel.task import Task
@@ -35,6 +36,42 @@ class Usage(BaseModel):
         ge=0,
     )
 
+    def __add__(self, other: "Usage") -> "Usage":
+        """Add two Usage objects together, handling None values gracefully.
+
+        None + None = None
+        None + value = value
+        value + None = value
+        value1 + value2 = value1 + value2
+        """
+        if not isinstance(other, Usage):
+            raise TypeError(f"Cannot add Usage with {type(other).__name__}")
+
+        def _add_optional_int(a: int | None, b: int | None) -> int | None:
+            if a is None and b is None:
+                return None
+            if a is None:
+                return b
+            if b is None:
+                return a
+            return a + b
+
+        def _add_optional_float(a: float | None, b: float | None) -> float | None:
+            if a is None and b is None:
+                return None
+            if a is None:
+                return b
+            if b is None:
+                return a
+            return a + b
+
+        return Usage(
+            input_tokens=_add_optional_int(self.input_tokens, other.input_tokens),
+            output_tokens=_add_optional_int(self.output_tokens, other.output_tokens),
+            total_tokens=_add_optional_int(self.total_tokens, other.total_tokens),
+            cost=_add_optional_float(self.cost, other.cost),
+        )
+
 
 class TaskRun(KilnParentedModel):
     """
@@ -72,6 +109,10 @@ class TaskRun(KilnParentedModel):
         default=None,
         description="Usage information for the task run. This includes the number of input tokens, output tokens, and total tokens used.",
     )
+    trace: list[ChatCompletionMessageParam] | None = Field(
+        default=None,
+        description="The trace of the task run in OpenAI format. This is the list of messages that were sent to/from the model.",
+    )
 
     def thinking_training_data(self) -> str | None:
         """