From cd0ad211114478ed07690e2f39afef966de9fe49 Mon Sep 17 00:00:00 2001
From: Mustafa Elbehery <melbeher@redhat.com>
Date: Wed, 9 Jul 2025 12:55:16 +0200
Subject: [PATCH] chore(api): add `mypy` coverage to `apis` (#2648)

# What does this PR do?
<!-- Provide a short summary of what this PR does and why. Link to
relevant issues if applicable. -->
This PR adds static type coverage to `llama-stack/apis`

Part of https://github.com/meta-llama/llama-stack/issues/2647

<!-- If resolving an issue, uncomment and update the line below -->
<!-- Closes #[issue-number] -->

## Test Plan
<!-- Describe the tests you ran to verify your changes with result
summaries. *Provide clear instructions so the plan can be easily
re-executed.* -->

Signed-off-by: Mustafa Elbehery <melbeher@redhat.com>
---
 docs/_static/llama-stack-spec.html        | 59 ++++++++++++++++++++++-
 docs/_static/llama-stack-spec.yaml        | 42 +++++++++++++++-
 llama_stack/apis/common/training_types.py |  4 +-
 pyproject.toml                            |  1 -
 4 files changed, 101 insertions(+), 5 deletions(-)

diff --git a/docs/_static/llama-stack-spec.html b/docs/_static/llama-stack-spec.html
index ae9ad5d4c..8021e0e55 100644
--- a/docs/_static/llama-stack-spec.html
+++ b/docs/_static/llama-stack-spec.html
@@ -11132,8 +11132,38 @@
                 "title": "Trace"
             },
             "Checkpoint": {
-                "description": "Checkpoint created during training runs",
-                "title": "Checkpoint"
+                "type": "object",
+                "properties": {
+                    "identifier": {
+                        "type": "string"
+                    },
+                    "created_at": {
+                        "type": "string",
+                        "format": "date-time"
+                    },
+                    "epoch": {
+                        "type": "integer"
+                    },
+                    "post_training_job_id": {
+                        "type": "string"
+                    },
+                    "path": {
+                        "type": "string"
+                    },
+                    "training_metrics": {
+                        "$ref": "#/components/schemas/PostTrainingMetric"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "identifier",
+                    "created_at",
+                    "epoch",
+                    "post_training_job_id",
+                    "path"
+                ],
+                "title": "Checkpoint",
+                "description": "Checkpoint created during training runs"
             },
             "PostTrainingJobArtifactsResponse": {
                 "type": "object",
@@ -11156,6 +11186,31 @@
                 "title": "PostTrainingJobArtifactsResponse",
                 "description": "Artifacts of a finetuning job."
             },
+            "PostTrainingMetric": {
+                "type": "object",
+                "properties": {
+                    "epoch": {
+                        "type": "integer"
+                    },
+                    "train_loss": {
+                        "type": "number"
+                    },
+                    "validation_loss": {
+                        "type": "number"
+                    },
+                    "perplexity": {
+                        "type": "number"
+                    }
+                },
+                "additionalProperties": false,
+                "required": [
+                    "epoch",
+                    "train_loss",
+                    "validation_loss",
+                    "perplexity"
+                ],
+                "title": "PostTrainingMetric"
+            },
             "PostTrainingJobStatusResponse": {
                 "type": "object",
                 "properties": {
diff --git a/docs/_static/llama-stack-spec.yaml b/docs/_static/llama-stack-spec.yaml
index 48cefe12b..a18474646 100644
--- a/docs/_static/llama-stack-spec.yaml
+++ b/docs/_static/llama-stack-spec.yaml
@@ -7838,8 +7838,30 @@ components:
         - start_time
       title: Trace
     Checkpoint:
-      description: Checkpoint created during training runs
+      type: object
+      properties:
+        identifier:
+          type: string
+        created_at:
+          type: string
+          format: date-time
+        epoch:
+          type: integer
+        post_training_job_id:
+          type: string
+        path:
+          type: string
+        training_metrics:
+          $ref: '#/components/schemas/PostTrainingMetric'
+      additionalProperties: false
+      required:
+        - identifier
+        - created_at
+        - epoch
+        - post_training_job_id
+        - path
       title: Checkpoint
+      description: Checkpoint created during training runs
     PostTrainingJobArtifactsResponse:
       type: object
       properties:
@@ -7855,6 +7877,24 @@ components:
         - checkpoints
       title: PostTrainingJobArtifactsResponse
       description: Artifacts of a finetuning job.
+    PostTrainingMetric:
+      type: object
+      properties:
+        epoch:
+          type: integer
+        train_loss:
+          type: number
+        validation_loss:
+          type: number
+        perplexity:
+          type: number
+      additionalProperties: false
+      required:
+        - epoch
+        - train_loss
+        - validation_loss
+        - perplexity
+      title: PostTrainingMetric
     PostTrainingJobStatusResponse:
       type: object
       properties:
diff --git a/llama_stack/apis/common/training_types.py b/llama_stack/apis/common/training_types.py
index 46cd101af..a2c3b78f1 100644
--- a/llama_stack/apis/common/training_types.py
+++ b/llama_stack/apis/common/training_types.py
@@ -19,8 +19,10 @@ class PostTrainingMetric(BaseModel):
     perplexity: float
 
 
-@json_schema_type(schema={"description": "Checkpoint created during training runs"})
+@json_schema_type
 class Checkpoint(BaseModel):
+    """Checkpoint created during training runs"""
+
     identifier: str
     created_at: datetime
     epoch: int
diff --git a/pyproject.toml b/pyproject.toml
index e800ed689..b41e03615 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -225,7 +225,6 @@ follow_imports = "silent"
 # to exclude the entire directory.
 exclude = [
     # As we fix more and more of these, we should remove them from the list
-    "^llama_stack/apis/common/training_types\\.py$",
     "^llama_stack/cli/download\\.py$",
     "^llama_stack/cli/stack/_build\\.py$",
     "^llama_stack/distribution/build\\.py$",