Added iterator streaming functions and typed prompts.

Author: apascal07

go/ai/generate.go

@@ -531,8 +531,56 @@ func GenerateText(ctx context.Context, r api.Registry, opts ...GenerateOption) (
 	return res.Text(), nil
 }
 
-// Generate run generate request for this model. Returns ModelResponse struct.
-// TODO: Stream GenerateData with partial JSON
+// StreamValue is either a streamed chunk or the final response of a generate request.
+type StreamValue[Out, Stream any] struct {
+	Done     bool
+	Chunk    Stream         // valid if Done is false
+	Output   Out            // valid if Done is true
+	Response *ModelResponse // valid if Done is true
+}
+
+// ModelStreamValue is a stream value for a model response.
+type ModelStreamValue = StreamValue[*ModelResponse, *ModelResponseChunk]
+
+// errGenerateStop is a sentinel error used to signal early termination of streaming.
+var errGenerateStop = errors.New("stop")
+
+// GenerateStream generates a model response and streams the output.
+// It returns a function whose argument function (the "yield function") will be repeatedly
+// called with the results.
+//
+// If the yield function is passed a non-nil error, generation has failed with that
+// error; the yield function will not be called again.
+//
+// If the yield function's [StreamValue] argument has Done == true, the value's
+// Response field contains the final response; the yield function will not be called
+// again.
+//
+// Otherwise the Chunk field of the passed [StreamValue] holds a streamed chunk.
+func GenerateStream(ctx context.Context, r api.Registry, opts ...GenerateOption) func(func(*ModelStreamValue, error) bool) {
+	return func(yield func(*ModelStreamValue, error) bool) {
+		cb := func(ctx context.Context, chunk *ModelResponseChunk) error {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+			if !yield(&ModelStreamValue{Chunk: chunk}, nil) {
+				return errGenerateStop
+			}
+			return nil
+		}
+
+		allOpts := append(slices.Clone(opts), WithStreaming(cb))
+
+		resp, err := Generate(ctx, r, allOpts...)
+		if err != nil {
+			yield(nil, err)
+		} else {
+			yield(&ModelStreamValue{Done: true, Response: resp}, nil)
+		}
+	}
+}
+
+// GenerateData runs a generate request and returns strongly-typed output.
 func GenerateData[Out any](ctx context.Context, r api.Registry, opts ...GenerateOption) (*Out, *ModelResponse, error) {
 	var value Out
 	opts = append(opts, WithOutputType(value))
@@ -550,6 +598,50 @@ func GenerateData[Out any](ctx context.Context, r api.Registry, opts ...Generate
 	return &value, resp, nil
 }
 
+// GenerateDataStream generates a model response with streaming and returns strongly-typed output.
+// It returns a function whose argument function (the "yield function") will be repeatedly
+// called with the results.
+//
+// If the yield function is passed a non-nil error, generation has failed with that
+// error; the yield function will not be called again.
+//
+// If the yield function's [StreamValue] argument has Done == true, the value's
+// Output and Response fields contain the final typed output and response; the yield function
+// will not be called again.
+//
+// Otherwise the Chunk field of the passed [StreamValue] holds a streamed chunk.
+func GenerateDataStream[Out, Stream any](ctx context.Context, r api.Registry, opts ...GenerateOption) func(func(*StreamValue[Out, Stream], error) bool) {
+	return func(yield func(*StreamValue[Out, Stream], error) bool) {
+		cb := func(ctx context.Context, chunk *ModelResponseChunk) error {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+			// TODO: Convert ModelResponseChunk to Stream type.
+			if !yield(&StreamValue[Out, Stream]{}, nil) {
+				return errGenerateStop
+			}
+			return nil
+		}
+
+		var value Out
+		allOpts := append(slices.Clone(opts), WithOutputType(value), WithStreaming(cb))
+
+		resp, err := Generate(ctx, r, allOpts...)
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+
+		err = resp.Output(&value)
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+
+		yield(&StreamValue[Out, Stream]{Done: true, Output: value, Response: resp}, nil)
+	}
+}
+
 // Generate applies the [Action] to provided request.
 func (m *model) Generate(ctx context.Context, req *ModelRequest, cb ModelStreamCallback) (*ModelResponse, error) {
 	if m == nil {

go/ai/prompt.go

@@ -40,6 +40,8 @@ type Prompt interface {
 	Name() string
 	// Execute executes the prompt with the given options and returns a [ModelResponse].
 	Execute(ctx context.Context, opts ...PromptExecuteOption) (*ModelResponse, error)
+	// ExecuteStream executes the prompt with streaming and returns an iterator.
+	ExecuteStream(ctx context.Context, opts ...PromptExecuteOption) func(func(*ModelStreamValue, error) bool)
 	// Render renders the prompt with the given input and returns a [GenerateActionOptions] to be used with [GenerateWithRequest].
 	Render(ctx context.Context, input any) (*GenerateActionOptions, error)
 }
@@ -51,6 +53,14 @@ type prompt struct {
 	registry api.Registry
 }
 
+// DataPrompt is a prompt with strongly-typed input and output.
+// It wraps an underlying Prompt and provides type-safe Execute and Render methods.
+// The Out type parameter can be string for text outputs or any struct type for JSON outputs.
+type DataPrompt[In, Out, Stream any] struct {
+	prompt   Prompt
+	registry api.Registry
+}
+
 // DefinePrompt creates a new [Prompt] and registers it.
 func DefinePrompt(r api.Registry, name string, opts ...PromptOption) Prompt {
 	if name == "" {
@@ -232,6 +242,51 @@ func (p *prompt) Execute(ctx context.Context, opts ...PromptExecuteOption) (*Mod
 	return GenerateWithRequest(ctx, r, actionOpts, execOpts.Middleware, execOpts.Stream)
 }
 
+// ExecuteStream executes the prompt with streaming and returns an iterator.
+// It returns a function whose argument function (the "yield function") will be repeatedly
+// called with the results.
+//
+// If the yield function is passed a non-nil error, execution has failed with that
+// error; the yield function will not be called again.
+//
+// If the yield function's [ModelStreamValue] argument has Done == true, the value's
+// Response field contains the final response; the yield function will not be called again.
+//
+// Otherwise the Chunk field of the passed [ModelStreamValue] holds a streamed chunk.
+func (p *prompt) ExecuteStream(ctx context.Context, opts ...PromptExecuteOption) func(func(*ModelStreamValue, error) bool) {
+	return func(yield func(*ModelStreamValue, error) bool) {
+		if p == nil {
+			yield(nil, errors.New("Prompt.ExecuteStream: execute called on a nil Prompt; check that all prompts are defined"))
+			return
+		}
+
+		cb := func(ctx context.Context, chunk *ModelResponseChunk) error {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+			if !yield(&ModelStreamValue{Chunk: chunk}, nil) {
+				return errPromptStop
+			}
+			return nil
+		}
+
+		allOpts := make([]PromptExecuteOption, 0, len(opts)+1)
+		allOpts = append(allOpts, opts...)
+		allOpts = append(allOpts, WithStreaming(cb))
+
+		resp, err := p.Execute(ctx, allOpts...)
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+
+		yield(&ModelStreamValue{Done: true, Response: resp}, nil)
+	}
+}
+
+// errPromptStop is a sentinel error used to signal early termination of streaming.
+var errPromptStop = errors.New("stop")
+
 // Render renders the prompt template based on user input.
 func (p *prompt) Render(ctx context.Context, input any) (*GenerateActionOptions, error) {
 	if p == nil {
@@ -759,3 +814,178 @@ func contentType(ct, uri string) (string, []byte, error) {
 
 	return "", nil, errors.New("uri content type not found")
 }
+
+// DefineDataPrompt creates a new data prompt and registers it.
+// It automatically infers input schema from the In type parameter and configures
+// output schema and JSON format from the Out type parameter (unless Out is string).
+func DefineDataPrompt[In, Out, Stream any](r api.Registry, name string, opts ...PromptOption) *DataPrompt[In, Out, Stream] {
+	if name == "" {
+		panic("ai.DefineDataPrompt: name is required")
+	}
+
+	allOpts := make([]PromptOption, 0, len(opts)+2)
+
+	var in In
+	allOpts = append(allOpts, WithInputType(in))
+
+	var out Out
+	switch any(out).(type) {
+	case string:
+		// String output - no schema needed
+	default:
+		allOpts = append(allOpts, WithOutputType(out))
+	}
+
+	allOpts = append(allOpts, opts...)
+
+	p := DefinePrompt(r, name, allOpts...)
+
+	return &DataPrompt[In, Out, Stream]{
+		prompt:   p,
+		registry: r,
+	}
+}
+
+// LookupDataPrompt looks up a prompt by name and wraps it with type information.
+// This is useful for wrapping prompts loaded from .prompt files with strong types.
+// It returns nil if the prompt was not found.
+func LookupDataPrompt[In, Out, Stream any](r api.Registry, name string) *DataPrompt[In, Out, Stream] {
+	p := LookupPrompt(r, name)
+	if p == nil {
+		return nil
+	}
+
+	return AsDataPrompt[In, Out, Stream](p)
+}
+
+// AsDataPrompt wraps an existing Prompt with type information, returning a DataPrompt.
+// This is useful for adding strong typing to a dynamically obtained prompt.
+func AsDataPrompt[In, Out, Stream any](p Prompt) *DataPrompt[In, Out, Stream] {
+	if p == nil {
+		return nil
+	}
+
+	return &DataPrompt[In, Out, Stream]{
+		prompt:   p,
+		registry: p.(*prompt).registry,
+	}
+}
+
+// Name returns the name of the prompt.
+func (tp *DataPrompt[In, Out, Stream]) Name() string {
+	if tp == nil || tp.prompt == nil {
+		return ""
+	}
+	return tp.prompt.Name()
+}
+
+// Execute executes the typed prompt and returns the strongly-typed output along with the full model response.
+// For structured output types (non-string Out), the prompt must be configured with the appropriate
+// output schema, either through [DefineDataPrompt] or by using [WithOutputType] when defining the prompt.
+func (tp *DataPrompt[In, Out, Stream]) Execute(ctx context.Context, input In, opts ...PromptExecuteOption) (*Out, *ModelResponse, error) {
+	if tp == nil || tp.prompt == nil {
+		return nil, nil, errors.New("TypedPrompt.Execute: called on a nil prompt; check that all prompts are defined")
+	}
+
+	allOpts := make([]PromptExecuteOption, 0, len(opts)+1)
+	allOpts = append(allOpts, WithInput(input))
+	allOpts = append(allOpts, opts...)
+
+	resp, err := tp.prompt.Execute(ctx, allOpts...)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	output, err := extractTypedOutput[Out](resp)
+	if err != nil {
+		return nil, resp, err
+	}
+
+	return output, resp, nil
+}
+
+// ExecuteStream executes the typed prompt with streaming and returns an iterator.
+// It returns a function whose argument function (the "yield function") will be repeatedly
+// called with the results.
+//
+// If the yield function is passed a non-nil error, execution has failed with that
+// error; the yield function will not be called again.
+//
+// If the yield function's DataPromptStreamValue argument has Done == true, the value's
+// Output and Response fields contain the final typed output and response; the yield function
+// will not be called again.
+//
+// Otherwise the Chunk field of the passed DataPromptStreamValue holds a streamed chunk.
+//
+// For structured output types (non-string Out), the prompt must be configured with the appropriate
+// output schema, either through [DefineDataPrompt] or by using [WithOutputType] when defining the prompt.
+func (tp *DataPrompt[In, Out, Stream]) ExecuteStream(ctx context.Context, input In, opts ...PromptExecuteOption) func(func(*StreamValue[Out, Stream], error) bool) {
+	return func(yield func(*StreamValue[Out, Stream], error) bool) {
+		if tp == nil || tp.prompt == nil {
+			yield(nil, errors.New("DataPrompt.ExecuteStream: called on a nil prompt; check that all prompts are defined"))
+			return
+		}
+
+		cb := func(ctx context.Context, chunk *ModelResponseChunk) error {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+			// TODO: Convert ModelResponseChunk to StreamValue[Out, Stream].
+			if !yield(&StreamValue[Out, Stream]{}, nil) {
+				return errTypedPromptStop
+			}
+			return nil
+		}
+
+		allOpts := make([]PromptExecuteOption, 0, len(opts)+2)
+		allOpts = append(allOpts, WithInput(input))
+		allOpts = append(allOpts, opts...)
+		allOpts = append(allOpts, WithStreaming(cb))
+
+		resp, err := tp.prompt.Execute(ctx, allOpts...)
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+
+		output, err := extractTypedOutput[Out](resp)
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+
+		yield(&StreamValue[Out, Stream]{Done: true, Output: *output, Response: resp}, nil)
+	}
+}
+
+// Render renders the typed prompt template with the given input.
+func (tp *DataPrompt[In, Out, Stream]) Render(ctx context.Context, input In) (*GenerateActionOptions, error) {
+	if tp == nil || tp.prompt == nil {
+		return nil, errors.New("TypedPrompt.Render: called on a nil prompt; check that all prompts are defined")
+	}
+
+	return tp.prompt.Render(ctx, input)
+}
+
+// errTypedPromptStop is a sentinel error used to signal early termination of streaming.
+var errTypedPromptStop = errors.New("stop")
+
+// extractTypedOutput extracts the typed output from a model response.
+func extractTypedOutput[Out any](resp *ModelResponse) (*Out, error) {
+	var output Out
+
+	switch any(output).(type) {
+	case string:
+		// String output - use Text()
+		text := resp.Text()
+		// Type assertion to convert string to Out (which we know is string)
+		result := any(text).(Out)
+		return &result, nil
+	default:
+		// Structured output - unmarshal from response
+		if err := resp.Output(&output); err != nil {
+			return nil, fmt.Errorf("TypedPrompt: failed to parse output: %w", err)
+		}
+		return &output, nil
+	}
+}