Schema

The schema is an OpenAPI 3.0.2 specification that describes a model's interface. It's the contract between the model and everything that interacts with it.

Why the Schema Exists

Every Cog model uses the same Prediction API envelope format, but the input and output fields are model-specific. The schema captures what each model expects and produces.

flowchart TB
    subgraph envelope ["PredictionRequest (fixed envelope)"]
        input[""input"#colon; { ... } — model-specific"]
    end
    envelope -.- note["Schema defines this part"]

Without the schema, consumers would have no way to know:

• What inputs the model accepts
• What types those inputs should be
• What constraints apply (required fields, min/max values, allowed choices)
• What the output looks like

How It's Used Today

How It's Generated

Cog generates schemas statically. The Go schema generator parses Python source code at cog build time using tree-sitter. No Python process is invoked and no container boots to discover the model interface. The schema is produced from the model's source files before Docker build begins, which keeps schema generation deterministic, fast, and independent of the model's installed dependencies.

If the static parser encounters a type it can't resolve, the build fails with a typed schema error. Hard user errors such as parse failures and unsupported features like default_factory also fail before Docker build starts.

flowchart LR
    subgraph source["Model Source"]
        predict["run.py"]
        types["output_types.py"]
    end

    subgraph parser["Go Static Parser"]
        ts["tree-sitter Python"]
        resolve["Type Resolver"]
        cross["Cross-File Resolver"]
    end

    subgraph output["Schema"]
        spec["OpenAPI 3.0.2 JSON"]
    end

    predict --> ts
    types --> cross
    ts --> resolve
    cross --> resolve
    resolve --> spec

Static Path Pipeline Steps

1. Parse the predictor file with tree-sitter (concrete syntax tree, not AST)
2. Collect imports -- track where each name came from (from cog import Path, from cog import BaseModel)
3. Collect module scope -- resolve module-level variable assignments (for default values, choices lists)
4. Collect BaseModel subclasses -- find all classes that inherit from BaseModel (cog's dataclass-based version; pydantic BaseModel also supported for compatibility)
5. Resolve cross-file models — for imported names not found locally, find the .py file on disk, parse it, and extract its BaseModel definitions
6. Extract inputs — walk the run() method parameters, resolve types, defaults, and Input() metadata. Legacy class predict() methods are still accepted with a warning.
7. Resolve output type — recursively resolve the return type annotation into a SchemaType
8. Generate OpenAPI — convert the extracted PredictorInfo into a full OpenAPI 3.0.2 JSON document

If any step fails, the build stops before Docker starts.

Cross-File Resolution

When a predictor imports types from other project files, the schema generator resolves them automatically:

# output_types.py
from cog import BaseModel

class Prediction(BaseModel):
    text: str
    score: float
    tags: list[str]

# run.py
from cog import BaseRunner
from output_types import Prediction

class Runner(BaseRunner):
    def run(self, prompt: str) -> Prediction:
        ...

The resolver handles every permutation of local imports:

How it distinguishes local from external: the resolver converts the module path to a filesystem path and checks if the file exists. If output_types.py exists in the project directory, it's local. If not (e.g., from transformers import ...), it's external. Known external packages (stdlib, torch, numpy, etc.) are skipped without a filesystem check.

Error messages: when a type can't be resolved, the error includes the import source:

cannot resolve output type 'WeirdType' (imported from 'some_package') —
external types cannot be statically analyzed. Define it as a BaseModel
subclass in your predict file, or provide a .pyi stub

For external values that are already JSON-shaped but not visible to the schema resolver, Annotated[..., cog.Opaque] is the escape hatch. It tells Cog to treat the value as an opaque JSON object while preserving container shape: Annotated[ExternalType, cog.Opaque] becomes an object, and Annotated[list[ExternalType], cog.Opaque] becomes an array of objects. The same shape is preserved for fields inside cog.BaseModel outputs and supported pydantic models.

SchemaType: The Type System

Output types are represented as a recursive algebraic data type (SchemaType) that composes arbitrarily:

flowchart TD
    root["SchemaType"] --> prim["SchemaPrimitive — str, int, float, bool, Path"]
    root --> any["SchemaAny — untyped (bare dict, Any)"]
    root --> arr["SchemaArray — list#lsqb;T#rsqb;, with Items → SchemaType"]
    root --> dict["SchemaDict — dict#lsqb;str, V#rsqb;, with ValueType → SchemaType"]
    root --> obj["SchemaObject — BaseModel subclass, with Fields → OrderedMap"]
    root --> iter["SchemaIterator — Iterator#lsqb;T#rsqb;, with Elem → SchemaType"]
    root --> concat["SchemaConcatIterator — ConcatenateIterator#lsqb;str#rsqb;"]

This recursive structure means nested types like dict[str, list[dict[str, int]]] are fully representable and produce correct JSON Schema:

{
  "type": "object",
  "additionalProperties": {
    "type": "array",
    "items": {
      "type": "object",
      "additionalProperties": {
        "type": "integer"
      }
    }
  }
}

JSON Schema Generation

Each SchemaType produces its JSON Schema fragment via JSONSchema():

Type Mappings

Input Types

Output Types

Unsupported Output Types

Cog-Specific Extensions

Where the Schema Lives

In the Image

Embedded as a Docker label during build:

docker inspect my-model | jq -r '.[0].Config.Labels["run.cog.openapi_schema"]'

Also written to .cog/openapi_schema.json inside the image for the runtime to serve.

At Runtime

Override and Configuration

# Default: static schema generation
cog build -t my-model

# Use a pre-built schema file
COG_OPENAPI_SCHEMA=my_schema.json cog build

Schema Structure

A simplified example showing a multi-file predictor with structured output:

{
  "openapi": "3.0.2",
  "info": { "title": "Cog", "version": "0.1.0" },
  "paths": {
    "/predictions": {
      "post": {
        "requestBody": {
          "content": {
            "application/json": {
              "schema": { "$ref": "#/components/schemas/PredictionRequest" }
            }
          }
        }
      }
    }
  },
  "components": {
    "schemas": {
      "Input": {
        "type": "object",
        "properties": {
          "prompt": {
            "type": "string",
            "description": "Text prompt",
            "x-order": 0
          },
          "steps": {
            "type": "integer",
            "default": 50,
            "minimum": 1,
            "maximum": 100,
            "x-order": 1
          }
        },
        "required": ["prompt"]
      },
      "Output": {
        "type": "object",
        "properties": {
          "text": { "type": "string", "title": "Text" },
          "score": { "type": "number", "title": "Score" }
        },
        "required": ["text", "score"]
      },
      "PredictionRequest": { "...": "..." },
      "PredictionResponse": { "...": "..." }
    }
  }
}

Code References