diff --git a/go/ai/gen.go b/go/ai/gen.go
index 11de302171..531c6fcbcd 100644
--- a/go/ai/gen.go
+++ b/go/ai/gen.go
@@ -28,6 +28,20 @@ type Candidate struct {
 	Usage         *GenerationUsage `json:"usage,omitempty"`
 }
 
+type CandidateError struct {
+	Code    CandidateErrorCode `json:"code,omitempty"`
+	Index   float64            `json:"index,omitempty"`
+	Message string             `json:"message,omitempty"`
+}
+
+type CandidateErrorCode string
+
+const (
+	CandidateErrorCodeBlocked CandidateErrorCode = "blocked"
+	CandidateErrorCodeOther   CandidateErrorCode = "other"
+	CandidateErrorCodeUnknown CandidateErrorCode = "unknown"
+)
+
 // FinishReason is the reason why a model stopped generating tokens.
 type FinishReason string
 
@@ -39,10 +53,16 @@ const (
 	FinishReasonUnknown FinishReason = "unknown"
 )
 
+type DataPart struct {
+	Data     any            `json:"data,omitempty"`
+	Metadata map[string]any `json:"metadata,omitempty"`
+}
+
 // A GenerateRequest is a request to generate completions from a model.
 type GenerateRequest struct {
-	Candidates int `json:"candidates,omitempty"`
-	Config     any `json:"config,omitempty"`
+	Candidates int   `json:"candidates,omitempty"`
+	Config     any   `json:"config,omitempty"`
+	Context    []any `json:"context,omitempty"`
 	// Messages is a list of messages to pass to the model. The first n-1 Messages
 	// are treated as history. The last Message is the current request.
 	Messages []*Message             `json:"messages,omitempty"`
@@ -62,18 +82,26 @@ type GenerateRequestOutput struct {
 type OutputFormat string
 
 const (
-	OutputFormatJSON OutputFormat = "json"
-	OutputFormatText OutputFormat = "text"
+	OutputFormatJSON  OutputFormat = "json"
+	OutputFormatText  OutputFormat = "text"
+	OutputFormatMedia OutputFormat = "media"
 )
 
 // A GenerateResponse is a model's response to a [GenerateRequest].
 type GenerateResponse struct {
 	Candidates []*Candidate     `json:"candidates,omitempty"`
 	Custom     any              `json:"custom,omitempty"`
+	LatencyMs  float64          `json:"latencyMs,omitempty"`
 	Request    *GenerateRequest `json:"request,omitempty"`
 	Usage      *GenerationUsage `json:"usage,omitempty"`
 }
 
+type GenerateResponseChunk struct {
+	Content []*Part `json:"content,omitempty"`
+	Custom  any     `json:"custom,omitempty"`
+	Index   float64 `json:"index,omitempty"`
+}
+
 // GenerationCommonConfig holds configuration for generation.
 type GenerationCommonConfig struct {
 	MaxOutputTokens int      `json:"maxOutputTokens,omitempty"`
@@ -86,14 +114,20 @@ type GenerationCommonConfig struct {
 
 // GenerationUsage provides information about the generation process.
 type GenerationUsage struct {
-	Custom       map[string]float64 `json:"custom,omitempty"`
-	InputTokens  float64            `json:"inputTokens,omitempty"`
-	OutputTokens float64            `json:"outputTokens,omitempty"`
-	TotalTokens  float64            `json:"totalTokens,omitempty"`
+	Custom           map[string]float64 `json:"custom,omitempty"`
+	InputCharacters  float64            `json:"inputCharacters,omitempty"`
+	InputImages      float64            `json:"inputImages,omitempty"`
+	InputTokens      float64            `json:"inputTokens,omitempty"`
+	OutputCharacters float64            `json:"outputCharacters,omitempty"`
+	OutputImages     float64            `json:"outputImages,omitempty"`
+	OutputTokens     float64            `json:"outputTokens,omitempty"`
+	TotalTokens      float64            `json:"totalTokens,omitempty"`
 }
 
 type mediaPart struct {
-	Media *mediaPartMedia `json:"media,omitempty"`
+	Data     any             `json:"data,omitempty"`
+	Media    *mediaPartMedia `json:"media,omitempty"`
+	Metadata map[string]any  `json:"metadata,omitempty"`
 }
 
 type mediaPartMedia struct {
@@ -107,6 +141,21 @@ type Message struct {
 	Role    Role    `json:"role,omitempty"`
 }
 
+type ModelInfo struct {
+	Label    string             `json:"label,omitempty"`
+	Supports *ModelInfoSupports `json:"supports,omitempty"`
+	Versions []string           `json:"versions,omitempty"`
+}
+
+type ModelInfoSupports struct {
+	Context    bool         `json:"context,omitempty"`
+	Media      bool         `json:"media,omitempty"`
+	Multiturn  bool         `json:"multiturn,omitempty"`
+	Output     OutputFormat `json:"output,omitempty"`
+	SystemRole bool         `json:"systemRole,omitempty"`
+	Tools      bool         `json:"tools,omitempty"`
+}
+
 // Role indicates which entity is responsible for the content of a message.
 type Role string
 
@@ -123,11 +172,13 @@ const (
 )
 
 type textPart struct {
-	Text string `json:"text,omitempty"`
+	Metadata map[string]any `json:"metadata,omitempty"`
+	Text     string         `json:"text,omitempty"`
 }
 
 // A ToolDefinition describes a tool.
 type ToolDefinition struct {
+	Description string `json:"description,omitempty"`
 	// Valid JSON Schema representing the input of the tool.
 	InputSchema map[string]any `json:"inputSchema,omitempty"`
 	Name        string         `json:"name,omitempty"`
diff --git a/go/genkit/schemas.config b/go/genkit/schemas.config
index 5152ff3a47..3acb49ace2 100644
--- a/go/genkit/schemas.config
+++ b/go/genkit/schemas.config
@@ -124,9 +124,13 @@ by the model in a [ToolRequest].
 .
 
 Candidate						pkg ai
+CandidateError					pkg ai
+CandidateErrorCode				pkg ai
 CandidateFinishReason			pkg ai
+DataPart						pkg ai
 DocumentData					pkg ai
 GenerateResponse				pkg ai
+GenerateResponseChunk			pkg ai
 GenerateRequest					pkg ai
 GenerateRequestOutput			pkg ai
 GenerateRequestOutputFormat		pkg ai
@@ -144,16 +148,27 @@ ToolResponsePartToolResponse	pkg ai
 Part							pkg ai
 TextPart						pkg ai
 TextPart						name textPart
+TextPart.data					omit
+TextPart.media					omit
+TextPart.toolRequest			omit
+TextPart.toolResponse			omit
 MediaPart						pkg ai
 MediaPart						name mediaPart
+MediaPart.text					omit
+MediaPart.toolRequest			omit
+MediaPart.toolResponse			omit
 MediaPartMedia					pkg ai
 MediaPartMedia					name mediaPartMedia
+ModelInfo						pkg ai
+ModelInfoSupports				pkg ai
+ModelInfoSupports.output  		type OutputFormat
 Role							pkg ai
 RoleUser						pkg ai
 RoleModel						pkg ai
 RoleTool						pkg ai
 
 
+
 GenerateRequestOutput doc
 GenerateRequestOutput describes the structure that the model's output
 should conform to. If Format is [OutputFormatJSON], then Schema
@@ -180,3 +195,4 @@ ToolDefinition doc
 A ToolDefinition describes a tool.
 .
 
+DataPart/properties/metadata	name map[string]any
diff --git a/go/internal/cmd/jsonschemagen/jsonschemagen.go b/go/internal/cmd/jsonschemagen/jsonschemagen.go
index c6b92df673..d5dd2cb0b3 100644
--- a/go/internal/cmd/jsonschemagen/jsonschemagen.go
+++ b/go/internal/cmd/jsonschemagen/jsonschemagen.go
@@ -136,7 +136,7 @@ func run(infile, defaultPkgPath, configFile, outRoot string) error {
 		if !*noFormat {
 			src, err = format.Source(src)
 			if err != nil {
-				return err
+				return fmt.Errorf("file for package %s: %w", pkgPath, err)
 			}
 		}
 
@@ -338,6 +338,9 @@ func (g *generator) generateStruct(name string, s *Schema, tcfg *itemConfig) err
 		if fcfg == nil {
 			fcfg = &itemConfig{}
 		}
+		if fcfg.omit {
+			continue
+		}
 		fs := s.Properties[field]
 		// Ignore properties with a non-empty "not" constraint.
 		// They are probably the result of inheriting from a base zod type with a "never" constraint.
@@ -409,8 +412,13 @@ func (g *generator) typeExpr(s *Schema) (string, error) {
 		if !ok {
 			return "", fmt.Errorf("ref %q does not begin with prefix %q", s.Ref, refPrefix)
 		}
+		ic := g.cfg.configFor(name)
 		s2, ok := g.schemas[name]
 		if !ok {
+			// If there is no schema, perhaps there is a config value.
+			if ic != nil && ic.name != "" {
+				return ic.name, nil
+			}
 			return "", fmt.Errorf("unknown type in reference: %q", name)
 		}
 		// Apply a config that changes the name.