Restructure files a bit

Author: markuswustenberg

README.md

@@ -2,7 +2,7 @@
 
 <img src="logo.jpg" alt="Logo" width="300" align="right">
 
-[![CI](https://github.com/maragudk/openai/actions/workflows/ci.yml/badge.svg)](https://github.com/maragudk/openai/actions/workflows/ci.yml)
+[![CI](https://github.com/maragudk/gai-openai/actions/workflows/ci.yml/badge.svg)](https://github.com/maragudk/gai-openai/actions/workflows/ci.yml)
 
 [GAI](https://github.com/maragudk/gai) client for [OpenAI](https://openai.com) models and compatible APIs, such as [LlamaCPP](https://github.com/ggml-org/llama.cpp).
 

chat_complete.go

@@ -0,0 +1,111 @@
+package openai
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+
+	"github.com/openai/openai-go"
+	"maragu.dev/gai"
+)
+
+type ChatCompleteModel string
+
+const (
+	ChatCompleteModelGPT4o     = ChatCompleteModel(openai.ChatModelGPT4o)
+	ChatCompleteModelGPT4oMini = ChatCompleteModel(openai.ChatModelGPT4oMini)
+)
+
+type ChatCompleter struct {
+	Client *openai.Client
+	log    *slog.Logger
+	model  ChatCompleteModel
+}
+
+type NewChatCompleterOptions struct {
+	Model ChatCompleteModel
+}
+
+func (c *Client) NewChatCompleter(opts NewChatCompleterOptions) *ChatCompleter {
+	return &ChatCompleter{
+		Client: c.Client,
+		log:    c.log,
+		model:  opts.Model,
+	}
+}
+
+// ChatComplete satisfies [gai.ChatCompleter].
+func (c *ChatCompleter) ChatComplete(ctx context.Context, req gai.ChatCompleteRequest) (gai.ChatCompleteResponse, error) {
+	var messages []openai.ChatCompletionMessageParamUnion
+
+	for _, m := range req.Messages {
+		switch m.Role {
+		case gai.MessageRoleUser:
+			var parts []openai.ChatCompletionContentPartUnionParam
+			for _, part := range m.Parts {
+				switch part.Type {
+				case gai.MessagePartTypeText:
+					parts = append(parts, openai.TextPart(part.Text()))
+				default:
+					panic("not implemented")
+				}
+			}
+			messages = append(messages, openai.UserMessageParts(parts...))
+
+		default:
+			panic("not implemented")
+		}
+	}
+
+	params := openai.ChatCompletionNewParams{
+		Messages: openai.F(messages),
+		Model:    openai.F(openai.ChatModel(c.model)),
+	}
+
+	if req.Temperature != nil {
+		params.Temperature = openai.F(req.Temperature.Float64())
+	}
+
+	stream := c.Client.Chat.Completions.NewStreaming(ctx, params)
+
+	return gai.NewChatCompleteResponse(func(yield func(gai.MessagePart, error) bool) {
+		defer func() {
+			if err := stream.Close(); err != nil {
+				c.log.Info("Error closing stream", "error", err)
+			}
+		}()
+
+		var acc openai.ChatCompletionAccumulator
+		for stream.Next() {
+			chunk := stream.Current()
+			acc.AddChunk(chunk)
+
+			if _, ok := acc.JustFinishedContent(); ok {
+				break
+			}
+
+			if _, ok := acc.JustFinishedToolCall(); ok {
+				continue
+				// TODO handle tool call
+				// println("Tool call stream finished:", tool.Index, tool.Name, tool.Arguments)
+			}
+
+			if refusal, ok := acc.JustFinishedRefusal(); ok {
+				yield(gai.MessagePart{}, fmt.Errorf("refusal: %v", refusal))
+				return
+			}
+
+			if len(chunk.Choices) > 0 {
+				if !yield(gai.TextMessagePart(chunk.Choices[0].Delta.Content), nil) {
+					return
+				}
+			}
+		}
+
+		if err := stream.Err(); err != nil {
+			yield(gai.MessagePart{}, err)
+		}
+	}), nil
+}
+
+var _ gai.ChatCompleter = (*ChatCompleter)(nil)

chat_complete_test.go

@@ -0,0 +1,35 @@
+package openai_test
+
+import (
+	"testing"
+
+	"maragu.dev/gai"
+	"maragu.dev/is"
+
+	openai "maragu.dev/gai-openai"
+)
+
+func TestChatCompleter_ChatComplete(t *testing.T) {
+	t.Run("can chat-complete", func(t *testing.T) {
+		c := newClient()
+
+		cc := c.NewChatCompleter(openai.NewChatCompleterOptions{Model: openai.ChatCompleteModelGPT4oMini})
+
+		req := gai.ChatCompleteRequest{
+			Messages: []gai.Message{
+				gai.NewUserTextMessage("Hi!"),
+			},
+			Temperature: gai.Ptr(gai.Temperature(0)),
+		}
+
+		res, err := cc.ChatComplete(t.Context(), req)
+		is.NotError(t, err)
+
+		var output string
+		for part, err := range res.Parts() {
+			is.NotError(t, err)
+			output += part.Text()
+		}
+		is.Equal(t, "Hello! How can I assist you today?", output)
+	})
+}

client.go

@@ -0,0 +1,44 @@
+package openai
+
+import (
+	"log/slog"
+	"strings"
+
+	"github.com/openai/openai-go"
+	"github.com/openai/openai-go/option"
+)
+
+type Client struct {
+	Client *openai.Client
+	log    *slog.Logger
+}
+
+type NewClientOptions struct {
+	BaseURL string
+	Key     string
+	Log     *slog.Logger
+}
+
+func NewClient(opts NewClientOptions) *Client {
+	if opts.Log == nil {
+		opts.Log = slog.New(slog.DiscardHandler)
+	}
+
+	var clientOpts []option.RequestOption
+
+	if opts.BaseURL != "" {
+		if !strings.HasSuffix(opts.BaseURL, "/") {
+			opts.BaseURL += "/"
+		}
+		clientOpts = append(clientOpts, option.WithBaseURL(opts.BaseURL))
+	}
+
+	if opts.Key != "" {
+		clientOpts = append(clientOpts, option.WithAPIKey(opts.Key))
+	}
+
+	return &Client{
+		Client: openai.NewClient(clientOpts...),
+		log:    opts.Log,
+	}
+}

client_test.go

@@ -0,0 +1,23 @@
+package openai_test
+
+import (
+	"testing"
+
+	"maragu.dev/env"
+	"maragu.dev/is"
+
+	openai "maragu.dev/gai-openai"
+)
+
+func TestNewClient(t *testing.T) {
+	t.Run("can create a new client with a key", func(t *testing.T) {
+		client := openai.NewClient(openai.NewClientOptions{Key: "123"})
+		is.NotNil(t, client)
+	})
+}
+
+func newClient() *openai.Client {
+	_ = env.Load(".env.test.local")
+
+	return openai.NewClient(openai.NewClientOptions{Key: env.GetStringOrDefault("OPENAI_KEY", "")})
+}

embed.go

@@ -0,0 +1,78 @@
+package openai
+
+import (
+	"context"
+	"log/slog"
+
+	"github.com/openai/openai-go"
+	"github.com/openai/openai-go/shared"
+	"maragu.dev/errors"
+	"maragu.dev/gai"
+)
+
+type EmbedModel string
+
+const (
+	EmbedModelTextEmbedding3Large = EmbedModel(openai.EmbeddingModelTextEmbedding3Large)
+	EmbedModelTextEmbedding3Small = EmbedModel(openai.EmbeddingModelTextEmbedding3Small)
+)
+
+type Embedder struct {
+	Client     *openai.Client
+	dimensions int
+	log        *slog.Logger
+	model      EmbedModel
+}
+
+type NewEmbedderOptions struct {
+	Dimensions int
+	Model      EmbedModel
+}
+
+func (c *Client) NewEmbedder(opts NewEmbedderOptions) *Embedder {
+	if opts.Dimensions <= 0 {
+		panic("dimensions must be greater than 0")
+	}
+
+	switch opts.Model {
+	case EmbedModelTextEmbedding3Large:
+		if opts.Dimensions > 3072 {
+			panic("dimensions must be less than or equal to 3072")
+		}
+	case EmbedModelTextEmbedding3Small:
+		if opts.Dimensions > 1536 {
+			panic("dimensions must be less than or equal to 1536")
+		}
+	}
+
+	return &Embedder{
+		Client:     c.Client,
+		dimensions: opts.Dimensions,
+		log:        c.log,
+		model:      opts.Model,
+	}
+}
+
+// Embed satisfies [gai.Embedder].
+func (e *Embedder) Embed(ctx context.Context, req gai.EmbedRequest) (gai.EmbedResponse[float64], error) {
+	v := gai.ReadAllString(req.Input)
+
+	res, err := e.Client.Embeddings.New(ctx, openai.EmbeddingNewParams{
+		Input:          openai.F[openai.EmbeddingNewParamsInputUnion](shared.UnionString(v)),
+		Model:          openai.F(openai.EmbeddingModel(e.model)),
+		EncodingFormat: openai.F(openai.EmbeddingNewParamsEncodingFormatFloat),
+		Dimensions:     openai.F(int64(e.dimensions)),
+	})
+	if err != nil {
+		return gai.EmbedResponse[float64]{}, errors.Wrap(err, "error embedding")
+	}
+	if len(res.Data) == 0 {
+		return gai.EmbedResponse[float64]{}, errors.New("no embeddings returned")
+	}
+
+	return gai.EmbedResponse[float64]{
+		Embedding: res.Data[0].Embedding,
+	}, nil
+}
+
+var _ gai.Embedder[float64] = (*Embedder)(nil)

embed_test.go

@@ -0,0 +1,31 @@
+package openai_test
+
+import (
+	"strings"
+	"testing"
+
+	"maragu.dev/gai"
+	"maragu.dev/is"
+
+	openai "maragu.dev/gai-openai"
+)
+
+func TestEmbedder_Embed(t *testing.T) {
+	t.Run("can embed a text", func(t *testing.T) {
+		c := newClient()
+
+		e := c.NewEmbedder(openai.NewEmbedderOptions{
+			Model:      openai.EmbedModelTextEmbedding3Small,
+			Dimensions: 1536,
+		})
+
+		req := gai.EmbedRequest{
+			Input: strings.NewReader("Embed this, please."),
+		}
+
+		res, err := e.Embed(t.Context(), req)
+		is.NotError(t, err)
+
+		is.Equal(t, 1536, len(res.Embedding))
+	})
+}