Kaapi v1.0: Database Comments (#476)

* added database comments

* sticking to pydantic

* updating migration

* first stab at comments

* cleanups

* updated migration

* cleanup

* coderabbit suggestions

* cleanups

* cleanups

* generated new migration with updated format

* added comment for config as well

* fixed minor comments

Author: AkhileshNegi

backend/app/alembic/versions/040_add_db_comments.py

@@ -1,21 +1,35 @@
-from uuid import UUID, uuid4
-import secrets
-import base64
 from datetime import datetime
-from typing import Optional, List
-from sqlmodel import SQLModel, Field, Relationship
+from uuid import UUID, uuid4
+
+from sqlmodel import Field, SQLModel
 
 from app.core.util import now
 
 
 class APIKeyBase(SQLModel):
+    """Base model for API keys with foreign key fields."""
+
+    # Foreign keys
     organization_id: int = Field(
-        foreign_key="organization.id", nullable=False, ondelete="CASCADE"
+        foreign_key="organization.id",
+        nullable=False,
+        ondelete="CASCADE",
+        sa_column_kwargs={"comment": "Reference to the organization"},
     )
     project_id: int = Field(
-        foreign_key="project.id", nullable=False, ondelete="CASCADE"
+        foreign_key="project.id",
+        nullable=False,
+        ondelete="CASCADE",
+        sa_column_kwargs={"comment": "Reference to the project"},
+    )
+    user_id: int = Field(
+        foreign_key="user.id",
+        nullable=False,
+        ondelete="CASCADE",
+        sa_column_kwargs={
+            "comment": "Reference to the user for whom the API key was created"
+        },
     )
-    user_id: int = Field(foreign_key="user.id", nullable=False, ondelete="CASCADE")
 
 
 class APIKeyPublic(APIKeyBase):
@@ -32,14 +46,44 @@ class APIKeyCreateResponse(APIKeyPublic):
 
 
 class APIKey(APIKeyBase, table=True):
-    id: UUID = Field(default_factory=uuid4, primary_key=True)
+    """Database model for API keys."""
 
+    id: UUID = Field(
+        default_factory=uuid4,
+        primary_key=True,
+        sa_column_kwargs={"comment": "Unique identifier for the API key"},
+    )
     key_prefix: str = Field(
-        unique=True, index=True, nullable=False
-    )  # Unique identifier from the key
-    key_hash: str = Field(nullable=False)  # bcrypt hash of the secret portion
-
-    inserted_at: datetime = Field(default_factory=now, nullable=False)
-    updated_at: datetime = Field(default_factory=now, nullable=False)
-    is_deleted: bool = Field(default=False, nullable=False)
-    deleted_at: Optional[datetime] = Field(default=None, nullable=True)
+        unique=True,
+        index=True,
+        nullable=False,
+        sa_column_kwargs={
+            "comment": "Unique prefix portion of the API key for identification"
+        },
+    )
+    key_hash: str = Field(
+        nullable=False,
+        sa_column_kwargs={"comment": "Bcrypt hash of the secret of the API key"},
+    )
+    is_deleted: bool = Field(
+        default=False,
+        nullable=False,
+        sa_column_kwargs={"comment": "Soft delete flag"},
+    )
+
+    # Timestamps
+    inserted_at: datetime = Field(
+        default_factory=now,
+        nullable=False,
+        sa_column_kwargs={"comment": "Timestamp when the API key was created"},
+    )
+    updated_at: datetime = Field(
+        default_factory=now,
+        nullable=False,
+        sa_column_kwargs={"comment": "Timestamp when the API key was last updated"},
+    )
+    deleted_at: datetime | None = Field(
+        default=None,
+        nullable=True,
+        sa_column_kwargs={"comment": "Timestamp when the API key was deleted"},
+    )

backend/app/models/api_key.py

@@ -1,11 +1,12 @@
 from datetime import datetime
-from typing import List, Optional
 
 from sqlalchemy import Column, String, Text
 from sqlalchemy.dialects.postgresql import ARRAY
 from sqlmodel import Field, Relationship, SQLModel, UniqueConstraint
 
 from app.core.util import now
+from app.models.organization import Organization
+from app.models.project import Project
 
 
 class AssistantBase(SQLModel):
@@ -17,7 +18,7 @@ class AssistantBase(SQLModel):
     name: str
     instructions: str = Field(sa_column=Column(Text, nullable=False))
     model: str
-    vector_store_ids: List[str] = Field(
+    vector_store_ids: list[str] = Field(
         default_factory=list, sa_column=Column(ARRAY(String))
     )
     temperature: float = 0.1
@@ -31,17 +32,88 @@ class AssistantBase(SQLModel):
 
 
 class Assistant(AssistantBase, table=True):
+    """OpenAI assistant configuration and metadata."""
+
     __tablename__ = "openai_assistant"
 
-    id: int = Field(default=None, primary_key=True)
-    inserted_at: datetime = Field(default_factory=now, nullable=False)
-    updated_at: datetime = Field(default_factory=now, nullable=False)
-    is_deleted: bool = Field(default=False, nullable=False)
-    deleted_at: Optional[datetime] = Field(default=None, nullable=True)
+    id: int = Field(
+        default=None,
+        primary_key=True,
+        sa_column_kwargs={"comment": "Unique identifier for the assistant"},
+    )
+    assistant_id: str = Field(
+        index=True,
+        sa_column_kwargs={"comment": "Unique identifier for the assistant at OpenAI"},
+    )
+    name: str = Field(
+        sa_column_kwargs={"comment": "Name of the assistant"},
+    )
+    instructions: str = Field(
+        sa_column=Column(
+            Text, nullable=False, comment="System instructions for the assistant"
+        )
+    )
+    model: str = Field(
+        sa_column_kwargs={"comment": "OpenAI model used by the assistant"},
+    )
+    vector_store_ids: list[str] = Field(
+        default_factory=list,
+        sa_column=Column(
+            ARRAY(String), comment="List of OpenAI vector store IDs attached"
+        ),
+    )
+    temperature: float = Field(
+        default=0.1,
+        sa_column_kwargs={
+            "comment": "Parameter that controls the creativity or randomness of the text generated by model"
+        },
+    )
+    max_num_results: int = Field(
+        default=20,
+        sa_column_kwargs={
+            "comment": "Parameter that controls maximum number of results to return"
+        },
+    )
+    is_deleted: bool = Field(
+        default=False,
+        nullable=False,
+        sa_column_kwargs={"comment": "Soft delete flag"},
+    )
+
+    # Foreign keys
+    project_id: int = Field(
+        foreign_key="project.id",
+        nullable=False,
+        ondelete="CASCADE",
+        sa_column_kwargs={"comment": "Reference to the project"},
+    )
+    organization_id: int = Field(
+        foreign_key="organization.id",
+        nullable=False,
+        ondelete="CASCADE",
+        sa_column_kwargs={"comment": "Reference to the organization"},
+    )
+
+    # Timestamps
+    inserted_at: datetime = Field(
+        default_factory=now,
+        nullable=False,
+        sa_column_kwargs={"comment": "Timestamp when the assistant was created"},
+    )
+    updated_at: datetime = Field(
+        default_factory=now,
+        nullable=False,
+        sa_column_kwargs={"comment": "Timestamp when the assistant was last updated"},
+    )
+    deleted_at: datetime | None = Field(
+        default=None,
+        nullable=True,
+        sa_column_kwargs={"comment": "Timestamp when the assistant was deleted"},
+    )
 
     # Relationships
-    project: "Project" = Relationship(back_populates="assistants")
-    organization: "Organization" = Relationship(back_populates="assistants")
+    project: Project = Relationship(back_populates="assistants")
+    organization: Organization = Relationship(back_populates="assistants")
 
 
 class AssistantCreate(SQLModel):

backend/app/models/assistants.py

@@ -13,49 +13,63 @@
 
 
 class BatchJob(SQLModel, table=True):
-    """Batch job table for tracking async LLM batch operations."""
+    """Database model for BatchJob operations."""
 
     __tablename__ = "batch_job"
     __table_args__ = (
         Index("idx_batch_job_status_org", "provider_status", "organization_id"),
         Index("idx_batch_job_status_project", "provider_status", "project_id"),
     )
 
-    id: int | None = Field(default=None, primary_key=True)
+    id: int | None = Field(
+        default=None,
+        primary_key=True,
+        sa_column_kwargs={"comment": "Unique identifier for the batch job"},
+    )
 
     # Provider and job type
     provider: str = Field(
         description="LLM provider name (e.g., 'openai', 'anthropic')",
+        sa_column_kwargs={"comment": "LLM provider name (e.g., openai, anthropic)"},
     )
     job_type: str = Field(
         index=True,
         description=(
             "Type of batch job (e.g., 'evaluation', 'classification', 'embedding')"
         ),
+        sa_column_kwargs={
+            "comment": "Type of batch job (e.g., evaluation, classification, embedding)"
+        },
     )
 
     # Batch configuration - stores all provider-specific config
     config: dict[str, Any] = Field(
         default_factory=dict,
-        sa_column=Column(JSONB, nullable=False),
+        sa_column=Column(
+            JSONB,
+            nullable=False,
+            comment="Complete batch configuration including model, temperature, instructions, tools, etc.",
+        ),
         description=(
-            "Complete batch configuration including model, temperature, "
-            "instructions, tools, etc."
+            "Complete batch configuration including model, temperature, instructions, tools, etc."
         ),
     )
 
     # Provider-specific batch tracking
     provider_batch_id: str | None = Field(
         default=None,
         description="Provider's batch job ID (e.g., OpenAI batch_id)",
+        sa_column_kwargs={"comment": "Provider's batch job ID (e.g., OpenAI batch_id)"},
     )
     provider_file_id: str | None = Field(
         default=None,
         description="Provider's input file ID",
+        sa_column_kwargs={"comment": "Provider's input file ID"},
     )
     provider_output_file_id: str | None = Field(
         default=None,
         description="Provider's output file ID",
+        sa_column_kwargs={"comment": "Provider's output file ID"},
     )
 
     # Provider status tracking
@@ -65,40 +79,56 @@ class BatchJob(SQLModel, table=True):
             "Provider-specific status (e.g., OpenAI: validating, in_progress, "
             "finalizing, completed, failed, expired, cancelling, cancelled)"
         ),
+        sa_column_kwargs={
+            "comment": "Provider-specific status (e.g., validating, in_progress, completed, failed)"
+        },
     )
 
     # Raw results (before parent-specific processing)
     raw_output_url: str | None = Field(
         default=None,
         description="S3 URL of raw batch output file",
+        sa_column_kwargs={"comment": "S3 URL of raw batch output file"},
     )
     total_items: int = Field(
         default=0,
         description="Total number of items in the batch",
+        sa_column_kwargs={"comment": "Total number of items in the batch"},
     )
 
     # Error handling
     error_message: str | None = Field(
         default=None,
-        sa_column=Column(Text, nullable=True),
+        sa_column=Column(Text, nullable=True, comment="Error message if batch failed"),
         description="Error message if batch failed",
     )
 
     # Foreign keys
     organization_id: int = Field(
-        foreign_key="organization.id", nullable=False, ondelete="CASCADE", index=True
+        foreign_key="organization.id",
+        nullable=False,
+        ondelete="CASCADE",
+        index=True,
+        sa_column_kwargs={"comment": "Reference to the organization"},
     )
     project_id: int = Field(
-        foreign_key="project.id", nullable=False, ondelete="CASCADE", index=True
+        foreign_key="project.id",
+        nullable=False,
+        ondelete="CASCADE",
+        index=True,
+        sa_column_kwargs={"comment": "Reference to the project"},
     )
 
     # Timestamps
     inserted_at: datetime = Field(
-        default_factory=now, description="The timestamp when the batch job was started"
+        default_factory=now,
+        description="The timestamp when the batch job was started",
+        sa_column_kwargs={"comment": "Timestamp when the batch job was started"},
     )
     updated_at: datetime = Field(
         default_factory=now,
         description="The timestamp when the batch job was last updated",
+        sa_column_kwargs={"comment": "Timestamp when the batch job was last updated"},
     )
 
     # Relationships