refactor(openai): change to use github.com/openai/openai-go

Author: bububa

agents/agent.go

@@ -33,12 +33,24 @@ func WithModel(model string) Option {
 	}
 }
 
-func WithTemperature(temperature float32) Option {
+func WithTemperature(temperature float64) Option {
 	return func(c *Config) {
 		c.temperature = temperature
 	}
 }
 
+func WithTopP(topP float64) Option {
+	return func(c *Config) {
+		c.topP = topP
+	}
+}
+
+func WithTopK(topK int) Option {
+	return func(c *Config) {
+		c.topK = topK
+	}
+}
+
 func WithMaxTokens(maxTokens int) Option {
 	return func(c *Config) {
 		c.maxTokens = maxTokens

agents/option.go

@@ -61,11 +61,21 @@ func (t *ToolAgent[I, T, O]) SetModel(model string) {
 	t.end.model = model
 }
 
-func (t *ToolAgent[I, T, O]) SetTemperature(temperature float32) {
+func (t *ToolAgent[I, T, O]) SetTemperature(temperature float64) {
 	t.start.temperature = temperature
 	t.end.temperature = temperature
 }
 
+func (t *ToolAgent[I, T, O]) SetTopP(topP float64) {
+	t.start.topP = topP
+	t.end.topP = topP
+}
+
+func (t *ToolAgent[I, T, O]) SetTopK(topK int) {
+	t.start.topK = topK
+	t.end.topK = topK
+}
+
 func (t *ToolAgent[I, T, O]) SetMaxTokens(maxTokens int) {
 	t.start.maxTokens = maxTokens
 	t.end.maxTokens = maxTokens

agents/tool_agent.go

@@ -47,10 +47,10 @@ func (p *Embedder) Embed(ctx context.Context, text string, embedding *embedder.E
 	respV := resp.GetEmbeddingsFloats()
 	if usage != nil && respV.Meta != nil && respV.Meta.Tokens != nil {
 		if v := respV.Meta.Tokens.InputTokens; v != nil {
-			usage.InputTokens = int(*v)
+			usage.InputTokens = int64(*v)
 		}
 		if v := respV.Meta.Tokens.OutputTokens; v != nil {
-			usage.OutputTokens = int(*v)
+			usage.OutputTokens = int64(*v)
 		}
 	}
 	if len(respV.Embeddings) == 0 {
@@ -77,10 +77,10 @@ func (p *Embedder) BatchEmbed(ctx context.Context, parts []string, usage *compon
 	respV := resp.GetEmbeddingsFloats()
 	if usage != nil && respV.Meta != nil && respV.Meta.Tokens != nil {
 		if v := respV.Meta.Tokens.InputTokens; v != nil {
-			usage.InputTokens = int(*v)
+			usage.InputTokens = int64(*v)
 		}
 		if v := respV.Meta.Tokens.OutputTokens; v != nil {
-			usage.OutputTokens = int(*v)
+			usage.OutputTokens = int64(*v)
 		}
 	}
 	ret := make([]embedder.Embedding, 0, len(respV.Embeddings))

components/embedder/providers/cohere/embeder.go

@@ -2,8 +2,9 @@ package openai
 
 import (
 	"context"
+	"errors"
 
-	openai "github.com/sashabaranov/go-openai"
+	"github.com/openai/openai-go"
 
 	"github.com/bububa/atomic-agents/components"
 	"github.com/bububa/atomic-agents/components/embedder"
@@ -34,22 +35,24 @@ func New(client *openai.Client, opts ...embedder.Option) *Embedder {
 
 func (p *Embedder) Embed(ctx context.Context, text string, embedding *embedder.Embedding, usage *components.LLMUsage) error {
 	// Create an EmbeddingRequest for the user query
-	req := openai.EmbeddingRequest{
-		Input: []string{text},
+	req := openai.EmbeddingNewParams{
+		Input: openai.EmbeddingNewParamsInputUnion{
+			OfString: openai.String(text),
+		},
 		Model: openai.EmbeddingModel(p.Model()),
 	}
-	resp, err := p.CreateEmbeddings(ctx, &req)
+	resp, err := p.Embeddings.New(ctx, req)
 	if err != nil {
 		return err
 	}
 	if usage != nil {
-		usage.InputTokens = int(resp.Usage.TotalTokens)
+		usage.InputTokens = resp.Usage.TotalTokens
 	}
 	if len(resp.Data) == 0 {
 		return nil
 	}
 	ret := resp.Data[0]
-	embedding.Object = ret.Object
+	embedding.Object = text
 	embedding.Embedding = make([]float64, 0, len(ret.Embedding))
 	for _, v := range ret.Embedding {
 		embedding.Embedding = append(embedding.Embedding, float64(v))
@@ -60,16 +63,18 @@ func (p *Embedder) Embed(ctx context.Context, text string, embedding *embedder.E
 
 func (p *Embedder) BatchEmbed(ctx context.Context, parts []string, usage *components.LLMUsage) ([]embedder.Embedding, error) {
 	// Create an EmbeddingRequest for the user query
-	req := openai.EmbeddingRequest{
-		Input: parts,
+	req := openai.EmbeddingNewParams{
+		Input: openai.EmbeddingNewParamsInputUnion{
+			OfArrayOfStrings: parts,
+		},
 		Model: openai.EmbeddingModel(p.Model()),
 	}
-	resp, err := p.CreateEmbeddings(ctx, &req)
+	resp, err := p.Embeddings.New(ctx, req)
 	if err != nil {
 		return nil, err
 	}
 	if usage != nil {
-		usage.InputTokens = int(resp.Usage.TotalTokens)
+		usage.InputTokens = resp.Usage.TotalTokens
 	}
 	ret := make([]embedder.Embedding, 0, len(resp.Data))
 	for _, v := range resp.Data {
@@ -78,20 +83,17 @@ func (p *Embedder) BatchEmbed(ctx context.Context, parts []string, usage *compon
 			embeddings = append(embeddings, float64(e))
 		}
 		ret = append(ret, embedder.Embedding{
-			Object:    v.Object,
+			Object:    parts[int(v.Index)],
 			Embedding: embeddings,
-			Index:     v.Index,
+			Index:     int(v.Index),
 		})
 	}
 	return ret, nil
 }
 
 func convertToOpenAI(src *embedder.Embedding, dist *openai.Embedding) {
-	embeddings := make([]float32, 0, len(src.Embedding))
-	for _, e := range src.Embedding {
-		embeddings = append(embeddings, float32(e))
-	}
-	dist.Embedding = embeddings
+	dist.Embedding = make([]float64, len(src.Embedding))
+	copy(dist.Embedding, src.Embedding)
 }
 
 // DotProduct calculates the dot product of the embedding vector with another
@@ -103,9 +105,12 @@ func (p *Embedder) DotProduct(ctx context.Context, target, query *embedder.Embed
 	convertToOpenAI(target, t)
 	q := new(openai.Embedding)
 	convertToOpenAI(query, q)
-	ret, err := t.DotProduct(q)
-	if err != nil {
-		return 0, err
+	if len(t.Embedding) != len(q.Embedding) {
+		return 0, errors.New("vector length mismatch")
+	}
+	var dotProduct float64
+	for i := range t.Embedding {
+		dotProduct += t.Embedding[i] * q.Embedding[i]
 	}
-	return float64(ret), nil
+	return dotProduct, nil
 }

components/embedder/providers/openai/embedder.go

@@ -41,7 +41,7 @@ func (p *Embedder) Embed(ctx context.Context, text string, embedding *embedder.E
 		return err
 	}
 	if usage != nil {
-		usage.InputTokens = int(resp.Usage.TotalTokens)
+		usage.InputTokens = int64(resp.Usage.TotalTokens)
 	}
 	if len(resp.Data) == 0 {
 		return nil
@@ -64,7 +64,7 @@ func (p *Embedder) BatchEmbed(ctx context.Context, parts []string, usage *compon
 		return nil, err
 	}
 	if usage != nil {
-		usage.InputTokens = int(resp.Usage.TotalTokens)
+		usage.InputTokens = int64(resp.Usage.TotalTokens)
 	}
 	ret := make([]embedder.Embedding, 0, len(resp.Data))
 	for _, v := range resp.Data {

components/embedder/providers/voyageai/embedder.go

@@ -16,8 +16,9 @@ import (
 	cohere "github.com/cohere-ai/cohere-go/v2"
 	"github.com/gabriel-vasile/mimetype"
 	anthropic "github.com/liushuangls/go-anthropic/v2"
+	"github.com/openai/openai-go"
+	"github.com/openai/openai-go/shared/constant"
 	"github.com/rs/xid"
-	openai "github.com/sashabaranov/go-openai"
 	gemini "google.golang.org/genai"
 
 	"github.com/bububa/atomic-agents/schema"
@@ -50,7 +51,7 @@ type LLMResponse struct {
 }
 
 // FromOpenAI convnert response from openai
-func (r *LLMResponse) FromOpenAI(v *openai.ChatCompletionResponse) {
+func (r *LLMResponse) FromOpenAI(v *openai.ChatCompletion) {
 	r.ID = v.ID
 	r.Role = AssistantRole
 	r.Model = v.Model
@@ -67,8 +68,8 @@ func (r *LLMResponse) FromAnthropic(v *anthropic.MessagesResponse) {
 	r.Role = AssistantRole
 	r.Model = string(v.Model)
 	r.Usage = &LLMUsage{
-		InputTokens:  v.Usage.InputTokens,
-		OutputTokens: v.Usage.OutputTokens,
+		InputTokens:  int64(v.Usage.InputTokens),
+		OutputTokens: int64(v.Usage.OutputTokens),
 	}
 	r.Details = v.Content
 }
@@ -83,10 +84,10 @@ func (r *LLMResponse) FromCohere(v *cohere.NonStreamedChatResponse) {
 		if usage := meta.Tokens; usage != nil {
 			r.Usage = new(LLMUsage)
 			if usage.InputTokens != nil {
-				r.Usage.InputTokens = int(*usage.InputTokens)
+				r.Usage.InputTokens = int64(*usage.InputTokens)
 			}
 			if usage.OutputTokens != nil {
-				r.Usage.OutputTokens = int(*usage.OutputTokens)
+				r.Usage.OutputTokens = int64(*usage.OutputTokens)
 			}
 		}
 		if version := meta.ApiVersion; version != nil {
@@ -100,15 +101,15 @@ func (r *LLMResponse) FromGemini(v *gemini.GenerateContentResponse) {
 	r.Role = AssistantRole
 	if v.UsageMetadata != nil && (v.UsageMetadata.PromptTokenCount > 0 || v.UsageMetadata.CandidatesTokenCount > 0) {
 		r.Usage = new(LLMUsage)
-		r.Usage.InputTokens = int(v.UsageMetadata.PromptTokenCount)
-		r.Usage.OutputTokens = int(v.UsageMetadata.CachedContentTokenCount)
+		r.Usage.InputTokens = int64(v.UsageMetadata.PromptTokenCount)
+		r.Usage.OutputTokens = int64(v.UsageMetadata.CachedContentTokenCount)
 	}
 	r.Details = v.Candidates
 }
 
 type LLMUsage struct {
-	InputTokens  int `json:"input_tokens,omitempty"`
-	OutputTokens int `json:"output_tokens,omitempty"`
+	InputTokens  int64 `json:"input_tokens,omitempty"`
+	OutputTokens int64 `json:"output_tokens,omitempty"`
 }
 
 func (u *LLMUsage) Merge(v *LLMUsage) {
@@ -256,12 +257,12 @@ func (m Message) TryAttachChunkPrompt(idx int) string {
 }
 
 // ToOpenAI convert message to openai ChatCompletionMessage
-func (m Message) ToOpenAI(dist *openai.ChatCompletionMessage) []openai.ChatCompletionMessage {
+func (m Message) ToOpenAI(dist *openai.ChatCompletionMessageParamUnion) []openai.ChatCompletionMessageParamUnion {
 	m.toOpenAI(dist, 0)
 	if l := len(m.Chunks()); l > 0 {
-		list := make([]openai.ChatCompletionMessage, 0, l)
+		list := make([]openai.ChatCompletionMessageParamUnion, 0, l)
 		for idx := range l {
-			var llmMsg openai.ChatCompletionMessage
+			var llmMsg openai.ChatCompletionMessageParamUnion
 			if err := m.toOpenAI(&llmMsg, idx+1); err == nil {
 				list = append(list, llmMsg)
 			}
@@ -271,7 +272,7 @@ func (m Message) ToOpenAI(dist *openai.ChatCompletionMessage) []openai.ChatCompl
 	return nil
 }
 
-func (m Message) toOpenAI(dist *openai.ChatCompletionMessage, idx int) error {
+func (m Message) toOpenAI(dist *openai.ChatCompletionMessageParamUnion, idx int) error {
 	src := m
 	chunks := m.Chunks()
 	if idx > 0 {
@@ -281,35 +282,40 @@ func (m Message) toOpenAI(dist *openai.ChatCompletionMessage, idx int) error {
 			return errors.New("invalid chunk index")
 		}
 	}
-	dist.Role = m.role
 	txt := m.TryAttachChunkPrompt(idx)
-	if attachement := src.Attachement(); attachement != nil && (len(attachement.ImageURLs) > 0 || len(attachement.VideoURLs) > 0) {
-		dist.MultiContent = make([]openai.ChatMessagePart, 0, len(attachement.ImageURLs)+len(attachement.VideoURLs)+1)
-		dist.MultiContent = append(dist.MultiContent, openai.ChatMessagePart{
-			Type: openai.ChatMessagePartTypeText,
-			Text: txt,
-		})
+	if attachement := src.Attachement(); m.role == UserRole && attachement != nil && (len(attachement.ImageURLs) > 0 || len(attachement.VideoURLs) > 0) {
+		contents := make([]openai.ChatCompletionContentPartUnionParam, 0, len(attachement.ImageURLs)+len(attachement.VideoURLs)+1)
+		contents = append(contents, openai.TextContentPart(txt))
 		for _, imageURL := range attachement.ImageURLs {
-			dist.MultiContent = append(dist.MultiContent, openai.ChatMessagePart{
-				Type: openai.ChatMessagePartTypeImageURL,
-				ImageURL: &openai.ChatMessageImageURL{
-					URL: imageURL,
-				},
-			})
+			contents = append(contents, openai.ImageContentPart(openai.ChatCompletionContentPartImageImageURLParam{
+				URL: imageURL,
+			}))
 		}
 		for _, videoURL := range attachement.VideoURLs {
-			dist.MultiContent = append(dist.MultiContent, openai.ChatMessagePart{
-				Type: "video_url",
-				VideoURL: &openai.ChatMessageVideoURL{
-					URL: videoURL,
-				},
-				Video: &openai.ChatMessageVideo{
-					URL: videoURL,
+			videoParam := &openai.ChatCompletionContentPartImageParam{
+				Type: constant.ImageURL("video_url"),
+			}
+			videoParam.SetExtraFields(map[string]any{
+				"video_url": openai.ImageURL{
+					URL:    videoURL,
+					Detail: openai.ImageURLDetailAuto,
 				},
 			})
+			part := openai.ChatCompletionContentPartUnionParam{
+				OfImageURL: videoParam,
+			}
+			contents = append(contents, part)
 		}
-	} else {
-		dist.Content = txt
+		*dist = openai.UserMessage(contents)
+		return nil
+	}
+	switch m.role {
+	case SystemRole:
+		*dist = openai.SystemMessage(txt)
+	case AssistantRole:
+		*dist = openai.AssistantMessage(txt)
+	case UserRole:
+		*dist = openai.UserMessage(txt)
 	}
 	return nil
 }

components/message.go

@@ -8,7 +8,8 @@ import (
 	cohereClient "github.com/cohere-ai/cohere-go/v2/client"
 	cohereOption "github.com/cohere-ai/cohere-go/v2/option"
 	anthropic "github.com/liushuangls/go-anthropic/v2"
-	openai "github.com/sashabaranov/go-openai"
+	"github.com/openai/openai-go"
+	"github.com/openai/openai-go/option"
 )
 
 func NewInstructor(provider instructor.Provider, modes ...instructor.Mode) instructor.Instructor {
@@ -39,11 +40,13 @@ func NewInstructor(provider instructor.Provider, modes ...instructor.Mode) instr
 	default:
 		authToken := os.Getenv("OPENAI_API_KEY")
 		baseURL := os.Getenv("OPENAI_BASE_URL")
-		cfg := openai.DefaultConfig(authToken)
+		opts := make([]option.RequestOption, 0, 2)
+
+		opts = append(opts, option.WithAPIKey(authToken))
 		if baseURL != "" {
-			cfg.BaseURL = baseURL
+			opts = append(opts, option.WithBaseURL(baseURL))
 		}
-		clt := openai.NewClientWithConfig(cfg)
-		return instructors.FromOpenAI(clt, instructor.WithMode(mode), instructor.WithMaxRetries(1), instructor.WithValidation())
+		clt := openai.NewClient(opts...)
+		return instructors.FromOpenAI(&clt, instructor.WithMode(mode), instructor.WithMaxRetries(1), instructor.WithValidation(), instructor.WithVerbose())
 	}
 }