Merge pull request #979 from Kiln-AI/sfierro/spec-persistence

Persistence moved to python

Author: sfierro

app/desktop/studio_server/api_client/kiln_ai_server_client/models/__init__.py

@@ -47,6 +47,7 @@
 from .specification_input_spec_fields import SpecificationInputSpecFields
 from .submit_answers_request import SubmitAnswersRequest
 from .target_task_info import TargetTaskInfo
+from .task_info import TaskInfo
 from .task_metadata import TaskMetadata
 from .validation_error import ValidationError
 
@@ -96,6 +97,7 @@
     "SpecSpecFields",
     "SubmitAnswersRequest",
     "TargetTaskInfo",
+    "TaskInfo",
     "TaskMetadata",
     "ValidationError",
 )

app/desktop/studio_server/api_client/kiln_ai_server_client/models/clarify_spec_input.py

@@ -1,49 +1,45 @@
 from __future__ import annotations
 
 from collections.abc import Mapping
-from typing import Any, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar
 
 from attrs import define as _attrs_define
 from attrs import field as _attrs_field
 
 from ..models.model_provider_name import ModelProviderName
 from ..types import UNSET, Unset
 
+if TYPE_CHECKING:
+    from ..models.task_info import TaskInfo
+
+
 T = TypeVar("T", bound="ClarifySpecInput")
 
 
 @_attrs_define
 class ClarifySpecInput:
     """
     Attributes:
-        target_task_prompt (str):
-        task_input_schema (str):
-        task_output_schema (str):
-        spec_rendered_prompt_template (str):
+        target_task_info (TaskInfo): Shared information about a task
+        target_specification (str):
         num_samples_per_topic (int):
         num_topics (int):
         providers (list[ModelProviderName]):
         num_exemplars (int | Unset):  Default: 10.
     """
 
-    target_task_prompt: str
-    task_input_schema: str
-    task_output_schema: str
-    spec_rendered_prompt_template: str
+    target_task_info: TaskInfo
+    target_specification: str
     num_samples_per_topic: int
     num_topics: int
     providers: list[ModelProviderName]
     num_exemplars: int | Unset = 10
     additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
 
     def to_dict(self) -> dict[str, Any]:
-        target_task_prompt = self.target_task_prompt
-
-        task_input_schema = self.task_input_schema
-
-        task_output_schema = self.task_output_schema
+        target_task_info = self.target_task_info.to_dict()
 
-        spec_rendered_prompt_template = self.spec_rendered_prompt_template
+        target_specification = self.target_specification
 
         num_samples_per_topic = self.num_samples_per_topic
 
@@ -60,10 +56,8 @@ def to_dict(self) -> dict[str, Any]:
         field_dict.update(self.additional_properties)
         field_dict.update(
             {
-                "target_task_prompt": target_task_prompt,
-                "task_input_schema": task_input_schema,
-                "task_output_schema": task_output_schema,
-                "spec_rendered_prompt_template": spec_rendered_prompt_template,
+                "target_task_info": target_task_info,
+                "target_specification": target_specification,
                 "num_samples_per_topic": num_samples_per_topic,
                 "num_topics": num_topics,
                 "providers": providers,
@@ -76,14 +70,12 @@ def to_dict(self) -> dict[str, Any]:
 
     @classmethod
     def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
-        d = dict(src_dict)
-        target_task_prompt = d.pop("target_task_prompt")
-
-        task_input_schema = d.pop("task_input_schema")
+        from ..models.task_info import TaskInfo
 
-        task_output_schema = d.pop("task_output_schema")
+        d = dict(src_dict)
+        target_task_info = TaskInfo.from_dict(d.pop("target_task_info"))
 
-        spec_rendered_prompt_template = d.pop("spec_rendered_prompt_template")
+        target_specification = d.pop("target_specification")
 
         num_samples_per_topic = d.pop("num_samples_per_topic")
 
@@ -99,10 +91,8 @@ def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
         num_exemplars = d.pop("num_exemplars", UNSET)
 
         clarify_spec_input = cls(
-            target_task_prompt=target_task_prompt,
-            task_input_schema=task_input_schema,
-            task_output_schema=task_output_schema,
-            spec_rendered_prompt_template=spec_rendered_prompt_template,
+            target_task_info=target_task_info,
+            target_specification=target_specification,
             num_samples_per_topic=num_samples_per_topic,
             num_topics=num_topics,
             providers=providers,

app/desktop/studio_server/api_client/kiln_ai_server_client/models/generate_batch_input.py

@@ -1,11 +1,15 @@
 from __future__ import annotations
 
 from collections.abc import Mapping
-from typing import Any, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar
 
 from attrs import define as _attrs_define
 from attrs import field as _attrs_field
 
+if TYPE_CHECKING:
+    from ..models.task_info import TaskInfo
+
+
 T = TypeVar("T", bound="GenerateBatchInput")
 
 
@@ -14,30 +18,30 @@ class GenerateBatchInput:
     """Input for batch generation (topics, inputs, outputs, optionally with scoring).
 
     Attributes:
-        target_task_prompt (str):
-        task_input_schema (str):
-        task_output_schema (str):
-        spec_rendered_prompt_template (str):
+        target_task_info (TaskInfo): Shared information about a task
+        topic_generation_task_info (TaskInfo): Shared information about a task
+        input_generation_task_info (TaskInfo): Shared information about a task
+        target_specification (str):
         num_samples_per_topic (int):
         num_topics (int):
     """
 
-    target_task_prompt: str
-    task_input_schema: str
-    task_output_schema: str
-    spec_rendered_prompt_template: str
+    target_task_info: TaskInfo
+    topic_generation_task_info: TaskInfo
+    input_generation_task_info: TaskInfo
+    target_specification: str
     num_samples_per_topic: int
     num_topics: int
     additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
 
     def to_dict(self) -> dict[str, Any]:
-        target_task_prompt = self.target_task_prompt
+        target_task_info = self.target_task_info.to_dict()
 
-        task_input_schema = self.task_input_schema
+        topic_generation_task_info = self.topic_generation_task_info.to_dict()
 
-        task_output_schema = self.task_output_schema
+        input_generation_task_info = self.input_generation_task_info.to_dict()
 
-        spec_rendered_prompt_template = self.spec_rendered_prompt_template
+        target_specification = self.target_specification
 
         num_samples_per_topic = self.num_samples_per_topic
 
@@ -47,10 +51,10 @@ def to_dict(self) -> dict[str, Any]:
         field_dict.update(self.additional_properties)
         field_dict.update(
             {
-                "target_task_prompt": target_task_prompt,
-                "task_input_schema": task_input_schema,
-                "task_output_schema": task_output_schema,
-                "spec_rendered_prompt_template": spec_rendered_prompt_template,
+                "target_task_info": target_task_info,
+                "topic_generation_task_info": topic_generation_task_info,
+                "input_generation_task_info": input_generation_task_info,
+                "target_specification": target_specification,
                 "num_samples_per_topic": num_samples_per_topic,
                 "num_topics": num_topics,
             }
@@ -60,24 +64,26 @@ def to_dict(self) -> dict[str, Any]:
 
     @classmethod
     def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        from ..models.task_info import TaskInfo
+
         d = dict(src_dict)
-        target_task_prompt = d.pop("target_task_prompt")
+        target_task_info = TaskInfo.from_dict(d.pop("target_task_info"))
 
-        task_input_schema = d.pop("task_input_schema")
+        topic_generation_task_info = TaskInfo.from_dict(d.pop("topic_generation_task_info"))
 
-        task_output_schema = d.pop("task_output_schema")
+        input_generation_task_info = TaskInfo.from_dict(d.pop("input_generation_task_info"))
 
-        spec_rendered_prompt_template = d.pop("spec_rendered_prompt_template")
+        target_specification = d.pop("target_specification")
 
         num_samples_per_topic = d.pop("num_samples_per_topic")
 
         num_topics = d.pop("num_topics")
 
         generate_batch_input = cls(
-            target_task_prompt=target_task_prompt,
-            task_input_schema=task_input_schema,
-            task_output_schema=task_output_schema,
-            spec_rendered_prompt_template=spec_rendered_prompt_template,
+            target_task_info=target_task_info,
+            topic_generation_task_info=topic_generation_task_info,
+            input_generation_task_info=input_generation_task_info,
+            target_specification=target_specification,
             num_samples_per_topic=num_samples_per_topic,
             num_topics=num_topics,
         )

app/desktop/studio_server/api_client/kiln_ai_server_client/models/task_info.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, TypeVar
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+T = TypeVar("T", bound="TaskInfo")
+
+
+@_attrs_define
+class TaskInfo:
+    """Shared information about a task
+
+    Attributes:
+        task_prompt (str):
+        task_input_schema (str):
+        task_output_schema (str):
+    """
+
+    task_prompt: str
+    task_input_schema: str
+    task_output_schema: str
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        task_prompt = self.task_prompt
+
+        task_input_schema = self.task_input_schema
+
+        task_output_schema = self.task_output_schema
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update(
+            {
+                "task_prompt": task_prompt,
+                "task_input_schema": task_input_schema,
+                "task_output_schema": task_output_schema,
+            }
+        )
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
+        d = dict(src_dict)
+        task_prompt = d.pop("task_prompt")
+
+        task_input_schema = d.pop("task_input_schema")
+
+        task_output_schema = d.pop("task_output_schema")
+
+        task_info = cls(
+            task_prompt=task_prompt,
+            task_input_schema=task_input_schema,
+            task_output_schema=task_output_schema,
+        )
+
+        task_info.additional_properties = d
+        return task_info
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties