1package main
 2
 3import (
 4	"context"
 5	"errors"
 6	"os"
 7
 8	openai "github.com/sashabaranov/go-openai"
 9)
10
11type OpenAIClient struct {
12	client *openai.Client
13	model  string
14}
15
16var (
17	ErrNoApiKey = errors.New("no api key set")
18)
19
20func NewOpenAIClient(model string) (M8B, error) {
21	apiKey, apiKeyFound := os.LookupEnv("OPENAI_API_KEY")
22	if !apiKeyFound {
23		return nil, ErrNoApiKey
24	}
25	return &OpenAIClient{
26		client: openai.NewClient(apiKey),
27		model:  model,
28	}, nil
29}
30
31func (o *OpenAIClient) Complete(ctx context.Context, prompt string) (<-chan string, error) {
32	out := make(chan string)
33
34	req := openai.ChatCompletionRequest{
35		Model: o.model,
36		Messages: []openai.ChatCompletionMessage{
37			{Role: openai.ChatMessageRoleUser, Content: prompt},
38		},
39		Stream: true,
40	}
41
42	stream, err := o.client.CreateChatCompletionStream(ctx, req)
43	if err != nil {
44		close(out)
45		return nil, err
46	}
47
48	go func() {
49		defer stream.Close()
50		defer close(out)
51
52		for {
53			resp, err := stream.Recv()
54			if err != nil {
55				return
56			}
57			if len(resp.Choices) > 0 {
58				out <- resp.Choices[0].Delta.Content
59			}
60		}
61	}()
62
63	return out, nil
64}