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 module -- parse Python source with tree-sitter and store the concrete syntax tree.
2. Collect imports -- track local names, aliases, Cog types, typing helpers, and builtins.
3. Collect module scope -- resolve module-level constants that can be used by defaults, descriptions, and choices.
4. Collect local schema models -- find local BaseModel and TypedDict classes.
5. Resolve imported models -- parse each local imported module at most once and merge schema model definitions and aliases.
6. Collect input registry -- record reusable class-level Input() attributes and helper methods.
7. Find target callable -- resolve the configured runner target. Predict-mode class targets prefer run() and fall back to legacy predict() for backward compatibility; train-mode class targets resolve train(). Standalone targets use the configured function name first, then fall back to the mode default if absent.
8. Extract inputs -- walk the resolved callable parameters and resolve types, defaults, and Input() metadata.
9. Resolve output type -- recursively resolve the return annotation into a SchemaType.
10. Generate OpenAPI -- convert the extracted schema information 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 local imports relative to the predictor file and project root:

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

Input unions are intentionally narrower than output types. Cog supports JSON-native input unions (str, int, float, bool, dict/Any, list[T], and None) so request validation can happen at the HTTP boundary and Python normalisation can choose a deterministic value type. Cog rejects unions involving Path, File, Secret, custom coders, and BaseModel because those cases are ambiguous for clients or runtime coercion. Output unions remain unsupported (see below).

A plain single-type optional (Optional[T] or T | None) is never placed in required, regardless of whether a default is supplied. A multi-variant nullable union (A | B | None) is different: because the field carries a concrete anyOf value type, it stays in required unless a default makes it omittable. This is why the two rows above differ in their required behaviour.

Nullable behaviour matches every other optional field: nullable: true (plus omission from required) means an omitted value falls back to the default. An explicit JSON null is still validated against the field type and is rejected at the HTTP edge, because the runtime validator does not treat OpenAPI's nullable keyword as an additional accepted value. "May be null" therefore means "may be omitted", not "accepts an explicit null payload".

Runtime caveat: Cog marks optionals as not-required in the schema, but the predictor still needs a Python-level default so the omitted value resolves to None. Use value: Optional[T] = Input(...) (the Input(...) supplies an implicit None) or Input(default=None). A bare value: Optional[T] annotation with no = Input(...) generates a correct "optional" schema but raises TypeError: missing 1 required positional argument when the field is omitted at runtime.

Output Types

Unsupported Output Types

Cog-Specific Extensions

Iterator output types describe the shape of accumulated JSON output. SSE response support is a separate prediction operation capability and is only advertised when the prediction handler opts in with @cog.streaming.

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