Classification: db models and migration script

backend/app/alembic/versions/e3c74fab4356_add_fine_tuning_table.py

@@ -0,0 +1,68 @@
+"""add fine tuning table
+
+Revision ID: e3c74fab4356
+Revises: e9dd35eff62c
+Create Date: 2025-08-04 22:03:54.552069
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+from sqlalchemy.dialects import postgresql
+
+
+revision = "e3c74fab4356"
+down_revision = "e9dd35eff62c"
+branch_labels = None
+depends_on = None
+
+finetuning_status_enum = postgresql.ENUM(
+    "pending",
+    "running",
+    "completed",
+    "failed",
+    name="finetuningstatus",
+    create_type=False,
+)
+
+
+def upgrade():
+    finetuning_status_enum.create(op.get_bind(), checkfirst=True)
+    op.create_table(
+        "fine_tuning",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+        sa.Column("base_model", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
+        sa.Column("split_ratio", sa.Float(), nullable=False),
+        sa.Column("document_id", sa.Uuid(), nullable=False),
+        sa.ForeignKeyConstraint(["document_id"], ["document.id"]),
+        sa.Column(
+            "training_file_id", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+        ),
+        sa.Column("testing_file_id", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+        sa.Column("provider_job_id", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+        sa.Column(
+            "status",
+            finetuning_status_enum,
+            nullable=False,
+            server_default="pending",
+        ),
+        sa.Column(
+            "fine_tuned_model", sqlmodel.sql.sqltypes.AutoString(), nullable=True
+        ),
+        sa.Column("error_message", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+        sa.Column("project_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(["project_id"], ["project.id"], ondelete="CASCADE"),
+        sa.Column("organization_id", sa.Integer(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["organization_id"], ["organization.id"], ondelete="CASCADE"
+        ),
+        sa.Column("is_deleted", sa.Boolean(), nullable=False),
+        sa.Column("inserted_at", sa.DateTime(), nullable=False),
+        sa.Column("updated_at", sa.DateTime(), nullable=False),
+        sa.Column("deleted_at", sa.DateTime(), nullable=True),
+    )
+
+
+def downgrade():
+    op.drop_table("fine_tuning")

backend/app/models/__init__.py

@@ -56,6 +56,13 @@
 
 from .assistants import Assistant, AssistantBase, AssistantCreate, AssistantUpdate
 
+from .fine_tuning import (
+    FineTuningJobBase,
+    Fine_Tuning,
+    FineTuningJobCreate,
+    FineTuningJobPublic,
+)
+
 from .openai_conversation import (
     OpenAIConversationPublic,
     OpenAIConversation,

backend/app/models/fine_tuning.py

@@ -0,0 +1,88 @@
+from typing import Optional
+from uuid import UUID
+from enum import Enum
+from datetime import datetime
+from sqlmodel import SQLModel, Field, Relationship
+
+from app.core.util import now
+
+
+class FineTuningStatus(str, Enum):
+    pending = "pending"
+    running = "running"
+    completed = "completed"
+    failed = "failed"
+
+
+class FineTuningJobBase(SQLModel):
+    base_model: str = Field(nullable=False, description="Base model for fine-tuning")
+    split_ratio: float = Field(nullable=False)
+    document_id: UUID = Field(foreign_key="document.id", nullable=False)
+    training_file_id: Optional[str] = Field(default=None)
+    testing_file_id: Optional[str] = Field(default=None)
+
+
+class FineTuningJobCreate(SQLModel):
+    base_model: str
+    split_ratio: list[float]
+    document_id: UUID
+
+
+class Fine_Tuning(FineTuningJobBase, table=True):
+    """Database model for tracking fine-tuning jobs."""
+
+    id: int = Field(primary_key=True)
+    provider_job_id: str | None = Field(
+        default=None, description="Fine tuning Job ID returned by OpenAI"
+    )
+    status: FineTuningStatus = (
+        Field(default=FineTuningStatus.pending, description="Fine tuning status"),
+    )
+    fine_tuned_model: str | None = Field(
+        default=None, description="Final fine tuned model name from OpenAI"
+    )
+    error_message: str | None = Field(
+        default=None, description="error message for when something failed"
+    )
+    project_id: int = Field(
+        foreign_key="project.id", nullable=False, ondelete="CASCADE"
+    )
+    organization_id: int = Field(
+        foreign_key="organization.id", nullable=False, ondelete="CASCADE"
+    )
+    is_deleted: bool = Field(default=False, nullable=False)
+
+    inserted_at: datetime = Field(default_factory=now, nullable=False)
+    updated_at: datetime = Field(default_factory=now, nullable=False)
+    deleted_at: datetime | None = Field(default=None, nullable=True)
+
+    project: "Project" = Relationship(back_populates="fine_tuning")
+
+
+class FineTuningUpdate(SQLModel):
+    training_file_id: Optional[str] = None
+    testing_file_id: Optional[str] = None
+    split_ratio: Optional[float] = None
+    provider_job_id: Optional[str] = None
+    fine_tuned_model: Optional[str] = None
+    status: Optional[str] = None
+    error_message: Optional[str] = None
+
+
+class FineTuningJobPublic(SQLModel):
+    """Public response model with job status and metadata."""
+
+    id: int
+    split_ratio: float
+    base_model: str
+    document_id: UUID
+    provider_job_id: str | None = None
+    status: str
+    error_message: str | None = None
+    fine_tuned_model: str | None = None
+    training_file_id: str | None = None
+    testing_file_id: str | None = None
+
+    inserted_at: datetime
+    updated_at: datetime
+    deleted_at: datetime | None = None

backend/app/models/organization.py

@@ -1,7 +1,6 @@
 from datetime import datetime
 from typing import List, TYPE_CHECKING
 from sqlmodel import Field, Relationship, SQLModel
-from sqlalchemy.orm import relationship
 
 from app.core.util import now
 

backend/app/models/project.py

@@ -49,6 +49,9 @@ class Project(ProjectBase, table=True):
     collections: list["Collection"] = Relationship(
         back_populates="project", cascade_delete=True
     )
+    fine_tuning: list["Fine_Tuning"] = Relationship(
+        back_populates="project", cascade_delete=True
+    )
     openai_conversations: list["OpenAIConversation"] = Relationship(
         back_populates="project", cascade_delete=True
     )