ai/ai.go at main · jetify-com/ai · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
package ai

import (
	"context"

	"go.jetify.com/ai/api"
)

// TODO: do we want to rename from GenerateText to Generate and from StreamText to Stream?

// GenerateText uses a language model to generate a text response from a given prompt.
//
// This function does not stream its output.
//
// It returns a [api.Response] containing the generated text, the results of
// any tool calls, and additional information.
//
// A prompt is a sequence of [api.Message]s:
//
//	GenerateText(ctx, []api.Message{
//		&api.UserMessage{
//			Content: []api.ContentBlock{
//				&api.TextBlock{Text: "Show me a picture of a cat"},
//			},
//		},
//		&api.AssistantMessage{
//			Content: []api.ContentBlock{
//				&api.TextBlock{Text: "Here is a picture of a cat"},
//				&api.ImageBlock{URL: "https://example.com/cat.png"},
//			},
//		},
//	})
//
// The last argument can optionally be a series of [GenerateOption] arguments:
//
//	GenerateText(ctx, messages, WithMaxTokens(100))
func GenerateText(ctx context.Context, prompt []api.Message, opts ...GenerateOption) (*api.Response, error) {
	config := buildGenerateConfig(opts)
	return generate(ctx, prompt, config)
}

// GenerateTextStr uses a language model to generate a text response from a given string prompt.
//
// It is a convenience wrapper around GenerateText for simple string-based prompts.
//
// Example usage:
//
//	GenerateTextStr(ctx, "Write a brief summary of the benefits of renewable energy")
//
// The function can optionally take [GenerateOption] arguments:
//
//	GenerateTextStr(ctx, "Explain the key differences between REST and GraphQL APIs", WithMaxTokens(500))
//
// The string prompt is automatically converted to a [api.UserMessage] before
// being passed to GenerateText.
func GenerateTextStr(ctx context.Context, prompt string, opts ...GenerateOption) (*api.Response, error) {
	msg := &api.UserMessage{
		Content: []api.ContentBlock{&api.TextBlock{Text: prompt}},
	}
	return GenerateText(ctx, []api.Message{msg}, opts...)
}

func generate(ctx context.Context, prompt []api.Message, opts GenerateOptions) (*api.Response, error) {
	return opts.Model.Generate(ctx, prompt, opts.CallOptions)
}

// StreamText uses a language model to generate a streaming text response from a given prompt.
//
// This function streams its output as a sequence of events.
//
// It returns a [api.StreamResponse] containing a stream of events from the model,
// including partial text, tool calls, and additional information.
//
// A prompt is a sequence of [api.Message]s:
//
//	StreamText(ctx, []api.Message{
//		&api.UserMessage{
//			Content: []api.ContentBlock{
//				&api.TextBlock{Text: "Show me a picture of a cat"},
//			},
//		},
//		&api.AssistantMessage{
//			Content: []api.ContentBlock{
//				&api.TextBlock{Text: "Here is a picture of a cat"},
//				&api.ImageBlock{URL: "https://example.com/cat.png"},
//			},
//		},
//	})
//
// The last argument can optionally be a series of [GenerateOption] arguments:
//
//	StreamText(ctx, messages, WithMaxTokens(100))
func StreamText(ctx context.Context, prompt []api.Message, opts ...GenerateOption) (*api.StreamResponse, error) {
	config := buildGenerateConfig(opts)
	return stream(ctx, prompt, config)
}

// StreamTextStr uses a language model to generate a streaming text response from a given string prompt.
//
// It is a convenience wrapper around StreamText for simple string-based prompts.
//
// Example usage:
//
//	StreamTextStr(ctx, "Write a brief summary of the benefits of renewable energy")
//
// The function can optionally take [GenerateOption] arguments:
//
//	StreamTextStr(ctx, "Explain the key differences between REST and GraphQL APIs", WithMaxTokens(500))
//
// The string prompt is automatically converted to a [api.UserMessage] before
// being passed to StreamText.
func StreamTextStr(ctx context.Context, prompt string, opts ...GenerateOption) (*api.StreamResponse, error) {
	msg := &api.UserMessage{
		Content: []api.ContentBlock{&api.TextBlock{Text: prompt}},
	}
	return StreamText(ctx, []api.Message{msg}, opts...)
}

func stream(ctx context.Context, prompt []api.Message, opts GenerateOptions) (*api.StreamResponse, error) {
	return opts.Model.Stream(ctx, prompt, opts.CallOptions)
}