Introduce the concept of "context" to PredictionRequest

[aron]

This allows a prediction request to include additional per-prediction context alongside the input. This will then be made available to the predictor via current_scope().context. For now this is limited to dict[str, str] but we could expand it later.

The primary benefit is to allow prediction scoped configuration to be passed alongside user input without changing the input signature.

Usage:

POST /predictions
{
  "input": { ... },
  "context": {
    "procedure_source_url": "file:///foo/bar/baz",
    "replicate_api_token": "r8_foop",
  },
}

Then in the predictor:

def predict():
    context: dict[str, str] = cog.current_scope().context
    assert context["foo"] == "bar"

Important

The drawback of the current implementation is that this additional configuration is not currently reflected in the generated interfaces to the model, this would require additional annotations to the predictor.

[aron]

For context we've (@aron & @philandstuff) discussed out of band what it would look like to allow a cog model to provide the context interface so that it could be included in the OpenAPI schema of the PredictionRequest and validated by the request handler.

We imagine this will be an interface on the predictor that defines the types as a TypedDict that is then extracted for use on the PredictionRequest. This could either be a parameter to the predict function like inputs or a class/module level property.

from cog import BaseContext, Input

class Context(BaseContext):
    name: str
    runtime: str

def predict(context: Context, prompt: str = Input()) -> str:
    ...

We feel that it would still be appropriate to provide access to the context via the current_scope().context interface, so that it can be accessed outside of the predictor e.g. in the Replicate SDK. As such, this could be considered an enhancement that we can add at a future time.