seekrai 0.5.26__tar.gz → 0.5.29__tar.gz

{seekrai-0.5.26 → seekrai-0.5.29}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: seekrai
-Version: 0.5.26
+Version: 0.5.29
 Summary: Python client for SeekrAI
 License: Apache-2.0
 License-File: LICENSE

{seekrai-0.5.26 → seekrai-0.5.29}/pyproject.toml

@@ -13,7 +13,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "seekrai"
-version = "0.5.26"
+version = "0.5.29"
 authors = [
     "SeekrFlow <support@seekr.com>"
 ]

{seekrai-0.5.26 → seekrai-0.5.29}/src/seekrai/types/__init__.py

@@ -1,6 +1,7 @@
 from seekrai.types.abstract import SeekrFlowClient
 from seekrai.types.agents import (
     Agent,
+    AgentAsToolLegacy,
     AgentDeleteResponse,
     CreateAgentRequest,
     EnvConfig,
@@ -47,6 +48,7 @@ from seekrai.types.agents import (
     WebSearchEnv,
 )
 from seekrai.types.agents.tools.schemas import (
+    AgentAsToolLegacy,
     FileSearch,
     FileSearchEnv,
     RunPython,
@@ -120,6 +122,9 @@ from seekrai.types.projects import (
     ProjectWithRuns,
 )
 from seekrai.types.tools import (
+    AgentAsTool,
+    AgentAsToolConfig,
+    CreateAgentAsTool,
     CreateFileSearch,
     CreateRunPython,
     CreateToolRequest,
@@ -132,6 +137,7 @@ from seekrai.types.tools import (
     ToolAgentSummaryResponse,
     ToolDeleteResponse,
     ToolResponse,
+    UpdateAgentAsTool,
     UpdateFileSearch,
     UpdateRunPython,
     UpdateToolRequest,
@@ -258,4 +264,9 @@ __all__ = [
     "RunPythonEnv",
     "WebSearch",
     "WebSearchEnv",
+    "AgentAsToolLegacy",
+    "AgentAsTool",
+    "CreateAgentAsTool",
+    "AgentAsToolConfig",
+    "UpdateAgentAsTool",
 ]

{seekrai-0.5.26 → seekrai-0.5.29}/src/seekrai/types/agents/__init__.py

@@ -51,6 +51,7 @@ from seekrai.types.agents.tools import (
     ToolBase,
 )
 from seekrai.types.agents.tools.schemas import (
+    AgentAsToolLegacy,
     FileSearch,
     FileSearchEnv,
     RunPython,
@@ -106,6 +107,7 @@ __all__ = [
     "RunPythonEnv",
     "WebSearch",
     "WebSearchEnv",
+    "AgentAsToolLegacy",
     "PythonFunctionBase",
     "PythonFunctionResponse",
     "DeletePythonFunctionResponse",

{seekrai-0.5.26 → seekrai-0.5.29}/src/seekrai/types/agents/tools/__init__.py

@@ -1,5 +1,6 @@
 from seekrai.types.agents.tools.env_model_config import EnvConfig
 from seekrai.types.agents.tools.schemas import (
+    AgentAsToolLegacy,
     FileSearch,
     FileSearchEnv,
     RunPython,
@@ -23,4 +24,5 @@ __all__ = [
     "RunPythonEnv",
     "WebSearch",
     "WebSearchEnv",
+    "AgentAsToolLegacy",
 ]

{seekrai-0.5.26 → seekrai-0.5.29}/src/seekrai/types/agents/tools/schemas/__init__.py

@@ -1,3 +1,4 @@
+from seekrai.types.agents.tools.schemas.agent_as_tool import AgentAsToolLegacy
 from seekrai.types.agents.tools.schemas.file_search import FileSearch
 from seekrai.types.agents.tools.schemas.file_search_env import FileSearchEnv
 from seekrai.types.agents.tools.schemas.run_python import RunPython
@@ -13,4 +14,5 @@ __all__ = [
     "RunPythonEnv",
     "WebSearch",
     "WebSearchEnv",
+    "AgentAsToolLegacy",
 ]

seekrai-0.5.29/src/seekrai/types/agents/tools/schemas/agent_as_tool.py

@@ -0,0 +1,16 @@
+from typing import Literal
+
+from seekrai.types.agents.tools import EnvConfig
+from seekrai.types.agents.tools.tool import ToolBase, ToolType
+
+
+class AgentAsToolLegacy(ToolBase[Literal["agent_as_tool"], EnvConfig]):
+    name: Literal["agent_as_tool"] = ToolType.AGENT_AS_TOOL.value
+    description: str
+    agent_id: str
+
+    model_config = {
+        "json_schema_extra": {
+            "deprecated": True,
+        }
+    }

{seekrai-0.5.26 → seekrai-0.5.29}/src/seekrai/types/agents/tools/tool_types.py

@@ -2,9 +2,13 @@ from typing import Annotated, Union
 
 from pydantic import Field
 
+from seekrai.types.agents.tools.schemas.agent_as_tool import AgentAsToolLegacy
 from seekrai.types.agents.tools.schemas.file_search import FileSearch
 from seekrai.types.agents.tools.schemas.run_python import RunPython
 from seekrai.types.agents.tools.schemas.web_search import WebSearch
 
 
-Tool = Annotated[Union[FileSearch, WebSearch, RunPython], Field(discriminator="name")]
+Tool = Annotated[
+    Union[FileSearch, WebSearch, RunPython, AgentAsToolLegacy],
+    Field(discriminator="name"),
+]

{seekrai-0.5.26 → seekrai-0.5.29}/src/seekrai/types/enums.py

@@ -28,3 +28,4 @@ class ToolType(str, Enum):
     FILE_SEARCH = "file_search"
     WEB_SEARCH = "web_search"
     RUN_PYTHON = "run_python"
+    AGENT_AS_TOOL = "agent_as_tool"

seekrai-0.5.29/src/seekrai/types/finetune.py

@@ -0,0 +1,459 @@
+from __future__ import annotations
+
+import warnings
+from datetime import datetime
+from enum import Enum
+from typing import Any, Callable, Dict, List, Literal, Optional
+
+from pydantic import Field, field_validator, model_serializer, model_validator
+
+from seekrai.types.abstract import BaseModel
+from seekrai.types.common import (
+    ObjectType,
+)
+from seekrai.utils._log import log_info
+
+
+class FinetuneJobStatus(str, Enum):
+    """
+    Possible fine-tune job status
+    """
+
+    STATUS_PENDING = "pending"
+    STATUS_QUEUED = "queued"
+    STATUS_RUNNING = "running"
+    # STATUS_COMPRESSING = "compressing"
+    # STATUS_UPLOADING = "uploading"
+    STATUS_CANCEL_REQUESTED = "cancel_requested"
+    STATUS_CANCELLED = "cancelled"
+    STATUS_FAILED = "failed"
+    STATUS_COMPLETED = "completed"
+    STATUS_DELETED = "deleted"
+
+
+class FinetuneEventLevels(str, Enum):
+    """
+    Fine-tune job event status levels
+    """
+
+    NULL = ""
+    INFO = "Info"
+    WARNING = "Warning"
+    ERROR = "Error"
+    LEGACY_INFO = "info"
+    LEGACY_IWARNING = "warning"
+    LEGACY_IERROR = "error"
+
+
+class FinetuneEventType(str, Enum):
+    """
+    Fine-tune job event types
+    """
+
+    JOB_PENDING = "JOB_PENDING"
+    JOB_START = "JOB_START"
+    JOB_STOPPED = "JOB_STOPPED"
+    MODEL_DOWNLOADING = "MODEL_DOWNLOADING"
+    MODEL_DOWNLOAD_COMPLETE = "MODEL_DOWNLOAD_COMPLETE"
+    TRAINING_DATA_DOWNLOADING = "TRAINING_DATA_DOWNLOADING"
+    TRAINING_DATA_DOWNLOAD_COMPLETE = "TRAINING_DATA_DOWNLOAD_COMPLETE"
+    VALIDATION_DATA_DOWNLOADING = "VALIDATION_DATA_DOWNLOADING"
+    VALIDATION_DATA_DOWNLOAD_COMPLETE = "VALIDATION_DATA_DOWNLOAD_COMPLETE"
+    WANDB_INIT = "WANDB_INIT"
+    TRAINING_START = "TRAINING_START"
+    CHECKPOINT_SAVE = "CHECKPOINT_SAVE"
+    BILLING_LIMIT = "BILLING_LIMIT"
+    EPOCH_COMPLETE = "EPOCH_COMPLETE"
+    TRAINING_COMPLETE = "TRAINING_COMPLETE"
+    MODEL_COMPRESSING = "COMPRESSING_MODEL"
+    MODEL_COMPRESSION_COMPLETE = "MODEL_COMPRESSION_COMPLETE"
+    MODEL_UPLOADING = "MODEL_UPLOADING"
+    MODEL_UPLOAD_COMPLETE = "MODEL_UPLOAD_COMPLETE"
+    JOB_COMPLETE = "JOB_COMPLETE"
+    JOB_ERROR = "JOB_ERROR"
+    CANCEL_REQUESTED = "CANCEL_REQUESTED"
+    JOB_RESTARTED = "JOB_RESTARTED"
+    REFUND = "REFUND"
+    WARNING = "WARNING"
+
+
+class FineTuneType(str, Enum):
+    STANDARD = "STANDARD"
+    GRPO = "GRPO"  # deprecated
+    PREFERENCE = "PREFERENCE"
+    REINFORCEMENT = "REINFORCEMENT"
+
+
+class GraderType(str, Enum):
+    FORMAT_CHECK = "format_check"
+    MATH_ACCURACY = "math_accuracy"
+    STRING_CHECK = "string_check"
+    TEXT_SIMILARITY = "text_similarity"
+
+
+class StringOperation(str, Enum):
+    EQUALS = "equals"
+    NOT_EQUALS = "not_equals"
+    CONTAINS = "contains"
+    CASE_INSENSITIVE_CONTAINS = "case_insensitive_contains"
+
+
+class TextSimilarityOperation(str, Enum):
+    BLEU = "bleu"
+    ROUGE = "rouge"
+
+
+class FinetuneEvent(BaseModel):
+    """
+    Fine-tune event type
+    """
+
+    # object type
+    object: Literal[ObjectType.FinetuneEvent]
+    # created at datetime stamp
+    created_at: datetime | None = None
+    # metrics that we expose
+    loss: float | None = None
+    epoch: float | None = None
+    reward: float | None = None
+
+    @model_serializer(mode="wrap")
+    def serialize_model(
+        self, handler: Callable[[Any], dict[str, Any]]
+    ) -> dict[str, Any]:
+        # Remove 'reward' if it's None
+        dump_dict = handler(self)
+        if dump_dict.get("reward") is None:
+            del dump_dict["reward"]
+        return dump_dict
+
+
+class LoRAConfig(BaseModel):
+    r: int = Field(8, gt=0, description="Rank of the update matrices.")
+    alpha: int = Field(32, gt=0, description="Scaling factor applied to LoRA updates.")
+    dropout: float = Field(
+        0.1,
+        ge=0.0,
+        le=1.0,
+        description="Fraction of LoRA neurons dropped during training.",
+    )
+    bias: Literal["none", "all", "lora_only"] = Field(
+        "none",
+        description="Bias terms to train; choose from 'none', 'all', or 'lora_only'.",
+    )
+    extras: Dict[str, Any] = Field(default_factory=dict)
+
+
+class Grader(BaseModel):
+    type: GraderType
+    weight: float | None = Field(default=None, gt=0.0, le=1.0)
+    operation: StringOperation | TextSimilarityOperation | None = Field(default=None)
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_operation(cls, data: Any) -> Any:
+        if not isinstance(data, dict):
+            return data
+
+        grader_type = data.get("type")
+        operation_value = data.get("operation")
+
+        if grader_type == GraderType.STRING_CHECK:
+            if not operation_value:
+                raise ValueError(
+                    "string_check grader is missing required StringOperation"
+                )
+            if isinstance(operation_value, str):
+                try:
+                    # Convert to enum to validate it's a valid value
+                    data["operation"] = StringOperation(operation_value.lower())
+                except ValueError:
+                    raise ValueError(
+                        f"Invalid operation for string_check grader: "
+                        f"expected StringOperation, but got type '{type(operation_value).__name__}' with value '{operation_value}'"
+                    )
+        elif grader_type == GraderType.TEXT_SIMILARITY:
+            if not operation_value:
+                raise ValueError(
+                    "text_similarity grader is missing required TextSimilarityOperation"
+                )
+            if isinstance(operation_value, str):
+                try:
+                    data["operation"] = TextSimilarityOperation(operation_value.lower())
+                except ValueError:
+                    raise ValueError(
+                        f"Invalid operation for text_similarity grader: "
+                        f"expected TextSimilarityOperation, got type '{type(operation_value).__name__}' with value '{operation_value}'"
+                    )
+
+        elif grader_type in (GraderType.FORMAT_CHECK, GraderType.MATH_ACCURACY):
+            if operation_value:
+                raise ValueError(f"{grader_type} grader cannot have an operation")
+            data["operation"] = None
+
+        return data
+
+
+class RewardComponents(BaseModel):
+    format_reward_weight: float = Field(default=0.1, gt=0.0, le=1.0)
+    graders: list[Grader] = Field(min_length=1)
+
+    @model_validator(mode="after")
+    def validate_weights(self) -> "RewardComponents":
+        is_format_weight_specified = "format_reward_weight" in self.model_fields_set
+
+        grader_weights_specified = [
+            grader.weight is not None for grader in self.graders
+        ]
+
+        all_graders_have_weights = all(grader_weights_specified)
+        some_graders_have_weights = any(grader_weights_specified) and not all(
+            grader_weights_specified
+        )
+        no_graders_have_weights = not any(grader_weights_specified)
+
+        if some_graders_have_weights:
+            raise ValueError(
+                "Only some graders have weights specified. Either all graders must have weights specified, or none of them."
+            )
+
+        if all_graders_have_weights and is_format_weight_specified:
+            self._validate_weights_sum_to_one()
+
+        elif all_graders_have_weights and not is_format_weight_specified:
+            self._normalize_grader_weights()
+
+        elif no_graders_have_weights:
+            self._initialize_grader_weights()
+            self._normalize_grader_weights()
+
+        return self
+
+    def _validate_weights_sum_to_one(self) -> None:
+        """Validate that format_reward_weight and grader weights sum to 1.0"""
+        total_weight = self.format_reward_weight + sum(  # type: ignore[operator]
+            grader.weight  # type: ignore[misc]
+            for grader in self.graders
+        )
+
+        if abs(total_weight - 1.0) > 1e-10:
+            raise ValueError(
+                f"When all weights are explicitly provided, they must sum to 1.0. "
+                f"Got format_reward_weight={self.format_reward_weight}, "
+                f"graders={self.graders}"
+            )
+
+    def _normalize_grader_weights(self) -> None:
+        """Normalize only grader weights to fill (1 - format_reward_weight)"""
+        total_grader_weight = sum(grader.weight for grader in self.graders)  # type: ignore[misc]
+        target_grader_total = 1.0 - self.format_reward_weight
+
+        # only normalize if weights aren't already properly normalized
+        if abs(total_grader_weight - target_grader_total) > 1e-10:
+            scale_factor = target_grader_total / total_grader_weight
+            for grader in self.graders:
+                original_weight = grader.weight
+                grader.weight *= scale_factor  # type: ignore[operator]
+                log_info(
+                    f"{grader.type}'s weight scaled from {original_weight} to {grader.weight:.2f}"
+                )
+
+    def _initialize_grader_weights(self) -> None:
+        """Initialize all grader weights when none are provided"""
+        for grader in self.graders:
+            grader.weight = 1.0
+
+
+class TrainingConfig(BaseModel):
+    # training file ID
+    training_files: List[str]
+    # base model string
+    model: str
+    # number of epochs to train for
+    n_epochs: int
+    # training learning rate
+    learning_rate: float
+    # number of checkpoints to save
+    n_checkpoints: int | None = None
+    # training batch size
+    batch_size: int = Field(..., ge=1, le=1024)
+    # up to 40 character suffix for output model name
+    experiment_name: str | None = None
+    # sequence length
+    max_length: int = 2500
+    # # weights & biases api key
+    # wandb_key: str | None = None
+    # IFT by default
+    pre_train: bool = False
+    # fine-tune type
+    fine_tune_type: FineTuneType = FineTuneType.STANDARD
+    # LoRA config
+    lora_config: Optional[LoRAConfig] = None
+    # reward_components are REINFORCEMENT-specific
+    reward_components: Optional[RewardComponents] = None
+
+    @model_validator(mode="after")
+    def validate_reward_components(self) -> "TrainingConfig":
+        # TODO: re-enable the below and make reward_components required for REINFORCEMENT. Disabled for now for backwards-compatibility
+        # if (
+        #     self.fine_tune_type in (FineTuneType.REINFORCEMENT, FineTuneType.GRPO)
+        #     and not self.reward_components
+        # ):
+        #     raise ValueError("REINFORCEMENT fine-tuning requires reward components")
+        if (
+            self.fine_tune_type in (FineTuneType.REINFORCEMENT, FineTuneType.GRPO)
+            and not self.reward_components
+        ):
+            self.reward_components = RewardComponents(
+                format_reward_weight=0.1,
+                graders=[Grader(type=GraderType.MATH_ACCURACY, weight=0.9)],
+            )
+        if self.fine_tune_type == FineTuneType.STANDARD and self.reward_components:
+            raise ValueError(
+                "Reward components are incompatible with standard fine-tuning"
+            )
+        if self.fine_tune_type == FineTuneType.PREFERENCE and self.reward_components:
+            raise ValueError(
+                "Reward components are incompatible with preference fine-tuning"
+            )
+
+        return self
+
+    @field_validator("fine_tune_type")
+    def validate_fine_tune_type(cls, v: Any) -> Any:
+        if v == FineTuneType.GRPO:
+            warnings.warn(
+                "FineTuneType.GRPO is deprecated and will be removed in a future version. Use FineTuneType.REINFORCEMENT",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        return v
+
+
+class AcceleratorType(str, Enum):
+    GAUDI2 = "GAUDI2"
+    GAUDI3 = "GAUDI3"
+    A100 = "A100"
+    A10 = "A10"
+    H100 = "H100"
+    MI300X = "MI300X"
+
+
+class InfrastructureConfig(BaseModel):
+    accel_type: AcceleratorType
+    n_accel: int
+    n_node: int = 1
+
+
+class FinetuneRequest(BaseModel):
+    """
+    Fine-tune request type
+    """
+
+    project_id: int
+    training_config: TrainingConfig
+    infrastructure_config: InfrastructureConfig
+
+
+class FinetuneResponse(BaseModel):
+    """
+    Fine-tune API response type
+    """
+
+    # job ID
+    id: str | None = None
+    # fine-tune type
+    fine_tune_type: FineTuneType = FineTuneType.STANDARD
+    reward_components: Optional[RewardComponents] = None
+    # training file id
+    training_files: List[str] | None = None
+    # validation file id
+    # validation_files: str | None = None TODO
+    # base model name
+    model: str | None = None
+    accel_type: AcceleratorType
+    n_accel: int
+    n_node: int | None = None
+    # number of epochs
+    n_epochs: int | None = None
+    # number of checkpoints to save
+    # n_checkpoints: int | None = None # TODO
+    # training batch size
+    batch_size: int | None = None
+    # training learning rate
+    learning_rate: float | None = None
+    # LoRA configuration returned when LoRA fine-tuning is enabled
+    lora_config: Optional[LoRAConfig] = None
+    # number of steps between evals
+    # eval_steps: int | None = None TODO
+    # created/updated datetime stamps
+    created_at: datetime | None = None
+    # updated_at: str | None = None
+    # up to 40 character suffix for output model name
+    experiment_name: str | None = None
+    # job status
+    status: FinetuneJobStatus | None = None
+    deleted_at: datetime | None = None
+
+    # list of fine-tune events
+    events: List[FinetuneEvent] | None = None
+    inference_available: bool = False
+    project_id: Optional[int] = None  # TODO - fix this
+    completed_at: datetime | None = None
+    description: str | None = None
+
+    # dataset token count
+    # TODO
+    # token_count: int | None = None
+    # # model parameter count
+    # param_count: int | None = None
+    # # fine-tune job price
+    # total_price: int | None = None
+    # # number of epochs completed (incrementing counter)
+    # epochs_completed: int | None = None
+    # # place in job queue (decrementing counter)
+    # queue_depth: int | None = None
+    # # weights & biases project name
+    # wandb_project_name: str | None = None
+    # # weights & biases job url
+    # wandb_url: str | None = None
+    # # training file metadata
+    # training_file_num_lines: int | None = Field(None, alias="TrainingFileNumLines")
+    # training_file_size: int | None = Field(None, alias="TrainingFileSize")
+
+    @model_serializer(mode="wrap")
+    def serialize_model(
+        self, handler: Callable[[Any], dict[str, Any]]
+    ) -> dict[str, Any]:
+        # Remove 'reward_components' if it's None
+        dump_dict = handler(self)
+        if dump_dict.get("reward_components") is None:
+            del dump_dict["reward_components"]
+        return dump_dict
+
+
+class FinetuneList(BaseModel):
+    # object type
+    object: Literal["list"] | None = None
+    # list of fine-tune job objects
+    data: List[FinetuneResponse] | None = None
+
+
+class FinetuneListEvents(BaseModel):
+    # object type
+    object: Literal["list"] | None = None
+    # list of fine-tune events
+    data: List[FinetuneEvent] | None = None
+
+
+class FinetuneDownloadResult(BaseModel):
+    # object type
+    object: Literal["local"] | None = None
+    # fine-tune job id
+    id: str | None = None
+    # checkpoint step number
+    checkpoint_step: int | None = None
+    # local path filename
+    filename: str | None = None
+    # size in bytes
+    size: int | None = None

{seekrai-0.5.26 → seekrai-0.5.29}/src/seekrai/types/tools.py

@@ -62,6 +62,15 @@ class CreateRunPython(CreateTool):
     config: RunPythonConfig
 
 
+class AgentAsToolConfig(ToolConfig):
+    agent_id: str
+
+
+class CreateAgentAsTool(CreateTool):
+    type: Literal[ToolType.AGENT_AS_TOOL] = ToolType.AGENT_AS_TOOL
+    config: AgentAsToolConfig
+
+
 class Tool(BaseModel):
     type: ToolType
     name: Annotated[str, AfterValidator(validate_length)]
@@ -88,13 +97,19 @@ class RunPythonTool(Tool):
     config: RunPythonConfig
 
 
+class AgentAsTool(Tool):
+    type: Literal[ToolType.AGENT_AS_TOOL] = ToolType.AGENT_AS_TOOL
+    config: AgentAsToolConfig
+
+
 CreateToolRequest = Annotated[
-    Union[CreateFileSearch, CreateRunPython, CreateWebSearch],
+    Union[CreateFileSearch, CreateRunPython, CreateWebSearch, CreateAgentAsTool],
     Field(discriminator="type"),
 ]
 
 ToolResponse = Annotated[
-    Union[FileSearchTool, WebSearchTool, RunPythonTool], Field(discriminator="type")
+    Union[FileSearchTool, WebSearchTool, RunPythonTool, AgentAsTool],
+    Field(discriminator="type"),
 ]
 
 
@@ -110,6 +125,10 @@ class UpdateRunPython(CreateRunPython):
     pass
 
 
+class UpdateAgentAsTool(CreateAgentAsTool):
+    pass
+
+
 UpdateToolRequest = Annotated[
     Union[UpdateFileSearch, UpdateWebSearch, UpdateRunPython],
     Field(discriminator="type"),
@@ -117,7 +136,7 @@ UpdateToolRequest = Annotated[
 
 
 class GetToolsResponse(BaseModel):
-    """Response schema for paginated tool list."""
+    """Response schema for a paginated tool list."""
 
     data: list[ToolResponse]
     total: Optional[int] = None

seekrai-0.5.26/src/seekrai/types/finetune.py

@@ -1,256 +0,0 @@
-from __future__ import annotations
-
-from datetime import datetime
-from enum import Enum
-from typing import Any, Dict, List, Literal, Optional
-
-from pydantic import Field
-
-from seekrai.types.abstract import BaseModel
-from seekrai.types.common import (
-    ObjectType,
-)
-
-
-class FinetuneJobStatus(str, Enum):
-    """
-    Possible fine-tune job status
-    """
-
-    STATUS_PENDING = "pending"
-    STATUS_QUEUED = "queued"
-    STATUS_RUNNING = "running"
-    # STATUS_COMPRESSING = "compressing"
-    # STATUS_UPLOADING = "uploading"
-    STATUS_CANCEL_REQUESTED = "cancel_requested"
-    STATUS_CANCELLED = "cancelled"
-    STATUS_FAILED = "failed"
-    STATUS_COMPLETED = "completed"
-    STATUS_DELETED = "deleted"
-
-
-class FinetuneEventLevels(str, Enum):
-    """
-    Fine-tune job event status levels
-    """
-
-    NULL = ""
-    INFO = "Info"
-    WARNING = "Warning"
-    ERROR = "Error"
-    LEGACY_INFO = "info"
-    LEGACY_IWARNING = "warning"
-    LEGACY_IERROR = "error"
-
-
-class FinetuneEventType(str, Enum):
-    """
-    Fine-tune job event types
-    """
-
-    JOB_PENDING = "JOB_PENDING"
-    JOB_START = "JOB_START"
-    JOB_STOPPED = "JOB_STOPPED"
-    MODEL_DOWNLOADING = "MODEL_DOWNLOADING"
-    MODEL_DOWNLOAD_COMPLETE = "MODEL_DOWNLOAD_COMPLETE"
-    TRAINING_DATA_DOWNLOADING = "TRAINING_DATA_DOWNLOADING"
-    TRAINING_DATA_DOWNLOAD_COMPLETE = "TRAINING_DATA_DOWNLOAD_COMPLETE"
-    VALIDATION_DATA_DOWNLOADING = "VALIDATION_DATA_DOWNLOADING"
-    VALIDATION_DATA_DOWNLOAD_COMPLETE = "VALIDATION_DATA_DOWNLOAD_COMPLETE"
-    WANDB_INIT = "WANDB_INIT"
-    TRAINING_START = "TRAINING_START"
-    CHECKPOINT_SAVE = "CHECKPOINT_SAVE"
-    BILLING_LIMIT = "BILLING_LIMIT"
-    EPOCH_COMPLETE = "EPOCH_COMPLETE"
-    TRAINING_COMPLETE = "TRAINING_COMPLETE"
-    MODEL_COMPRESSING = "COMPRESSING_MODEL"
-    MODEL_COMPRESSION_COMPLETE = "MODEL_COMPRESSION_COMPLETE"
-    MODEL_UPLOADING = "MODEL_UPLOADING"
-    MODEL_UPLOAD_COMPLETE = "MODEL_UPLOAD_COMPLETE"
-    JOB_COMPLETE = "JOB_COMPLETE"
-    JOB_ERROR = "JOB_ERROR"
-    CANCEL_REQUESTED = "CANCEL_REQUESTED"
-    JOB_RESTARTED = "JOB_RESTARTED"
-    REFUND = "REFUND"
-    WARNING = "WARNING"
-
-
-class FineTuneType(str, Enum):
-    STANDARD = "STANDARD"
-    PREFERENCE = "PREFERENCE"
-    GRPO = "GRPO"
-
-
-class FinetuneEvent(BaseModel):
-    """
-    Fine-tune event type
-    """
-
-    # object type
-    object: Literal[ObjectType.FinetuneEvent]
-    # created at datetime stamp
-    created_at: datetime | None = None
-    # metrics that we expose
-    loss: float | None = None
-    epoch: float | None = None
-
-
-class LoRAConfig(BaseModel):
-    r: int = Field(8, gt=0, description="Rank of the update matrices.")
-    alpha: int = Field(32, gt=0, description="Scaling factor applied to LoRA updates.")
-    dropout: float = Field(
-        0.1,
-        ge=0.0,
-        le=1.0,
-        description="Fraction of LoRA neurons dropped during training.",
-    )
-    bias: Literal["none", "all", "lora_only"] = Field(
-        "none",
-        description="Bias terms to train; choose from 'none', 'all', or 'lora_only'.",
-    )
-    extras: Dict[str, Any] = Field(default_factory=dict)
-
-
-class TrainingConfig(BaseModel):
-    # training file ID
-    training_files: List[str]
-    # base model string
-    model: str
-    # number of epochs to train for
-    n_epochs: int
-    # training learning rate
-    learning_rate: float
-    # number of checkpoints to save
-    n_checkpoints: int | None = None
-    # training batch size
-    batch_size: int = Field(..., ge=1, le=1024)
-    # up to 40 character suffix for output model name
-    experiment_name: str | None = None
-    # sequence length
-    max_length: int = 2500
-    # # weights & biases api key
-    # wandb_key: str | None = None
-    # IFT by default
-    pre_train: bool = False
-    # fine-tune type
-    fine_tune_type: FineTuneType = FineTuneType.STANDARD
-    # LoRA config
-    lora_config: Optional[LoRAConfig] = None
-
-
-class AcceleratorType(str, Enum):
-    GAUDI2 = "GAUDI2"
-    GAUDI3 = "GAUDI3"
-    A100 = "A100"
-    A10 = "A10"
-    H100 = "H100"
-    MI300X = "MI300X"
-
-
-class InfrastructureConfig(BaseModel):
-    accel_type: AcceleratorType
-    n_accel: int
-    n_node: int = 1
-
-
-class FinetuneRequest(BaseModel):
-    """
-    Fine-tune request type
-    """
-
-    project_id: int
-    training_config: TrainingConfig
-    infrastructure_config: InfrastructureConfig
-
-
-class FinetuneResponse(BaseModel):
-    """
-    Fine-tune API response type
-    """
-
-    # job ID
-    id: str | None = None
-    # fine-tune type
-    fine_tune_type: FineTuneType = FineTuneType.STANDARD
-    # training file id
-    training_files: List[str] | None = None
-    # validation file id
-    # validation_files: str | None = None TODO
-    # base model name
-    model: str | None = None
-    accel_type: AcceleratorType
-    n_accel: int
-    n_node: int | None = None
-    # number of epochs
-    n_epochs: int | None = None
-    # number of checkpoints to save
-    # n_checkpoints: int | None = None # TODO
-    # training batch size
-    batch_size: int | None = None
-    # training learning rate
-    learning_rate: float | None = None
-    # LoRA configuration returned when LoRA fine-tuning is enabled
-    lora_config: Optional[LoRAConfig] = None
-    # number of steps between evals
-    # eval_steps: int | None = None TODO
-    # created/updated datetime stamps
-    created_at: datetime | None = None
-    # updated_at: str | None = None
-    # up to 40 character suffix for output model name
-    experiment_name: str | None = None
-    # job status
-    status: FinetuneJobStatus | None = None
-    deleted_at: datetime | None = None
-
-    # list of fine-tune events
-    events: List[FinetuneEvent] | None = None
-    inference_available: bool = False
-    project_id: Optional[int] = None  # TODO - fix this
-    completed_at: datetime | None = None
-    description: str | None = None
-
-    # dataset token count
-    # TODO
-    # token_count: int | None = None
-    # # model parameter count
-    # param_count: int | None = None
-    # # fine-tune job price
-    # total_price: int | None = None
-    # # number of epochs completed (incrementing counter)
-    # epochs_completed: int | None = None
-    # # place in job queue (decrementing counter)
-    # queue_depth: int | None = None
-    # # weights & biases project name
-    # wandb_project_name: str | None = None
-    # # weights & biases job url
-    # wandb_url: str | None = None
-    # # training file metadata
-    # training_file_num_lines: int | None = Field(None, alias="TrainingFileNumLines")
-    # training_file_size: int | None = Field(None, alias="TrainingFileSize")
-
-
-class FinetuneList(BaseModel):
-    # object type
-    object: Literal["list"] | None = None
-    # list of fine-tune job objects
-    data: List[FinetuneResponse] | None = None
-
-
-class FinetuneListEvents(BaseModel):
-    # object type
-    object: Literal["list"] | None = None
-    # list of fine-tune events
-    data: List[FinetuneEvent] | None = None
-
-
-class FinetuneDownloadResult(BaseModel):
-    # object type
-    object: Literal["local"] | None = None
-    # fine-tune job id
-    id: str | None = None
-    # checkpoint step number
-    checkpoint_step: int | None = None
-    # local path filename
-    filename: str | None = None
-    # size in bytes
-    size: int | None = None