Add LLM interface

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

pkg/agents/agents.go

@@ -15,27 +15,17 @@
 package agents
 
 import (
-	"github.com/openai/openai-go"
-	"github.com/openai/openai-go/option"
 	"go.gearno.de/kit/log"
+	"go.probo.inc/probo/pkg/llm"
 )
 
-type (
-	Agent struct {
-		l      *log.Logger
-		cfg    Config
-		client *openai.Client
-	}
-
-	Config struct {
-		OpenAIAPIKey string
-		Temperature  float64
-		ModelName    string
-	}
-)
-
-func NewAgent(l *log.Logger, cfg Config) *Agent {
-	client := openai.NewClient(option.WithAPIKey(cfg.OpenAIAPIKey))
-
-	return &Agent{l: l, cfg: cfg, client: &client}
+type Agent struct {
+	l      *log.Logger
+	client *llm.Client
+	model  string
+	temp   float64
+}
+
+func NewAgent(l *log.Logger, client *llm.Client, model string, temp float64) *Agent {
+	return &Agent{l: l, client: client, model: model, temp: temp}
 }

pkg/agents/changelog_generator.go

@@ -18,8 +18,7 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/openai/openai-go"
-	"github.com/openai/openai-go/packages/param"
+	"go.probo.inc/probo/pkg/llm"
 )
 
 const (
@@ -49,23 +48,19 @@ const (
 )
 
 func (a *Agent) GenerateChangelog(ctx context.Context, oldContent string, newContent string) (*string, error) {
-	model := openai.ChatModel(a.cfg.ModelName)
-	chatCompletion, err := a.client.Chat.Completions.New(ctx, openai.ChatCompletionNewParams{
-		Messages: []openai.ChatCompletionMessageParamUnion{
-			openai.SystemMessage(changelogGeneratorSystemPrompt),
-			openai.UserMessage(fmt.Sprintf(`Old content: %s`, oldContent)),
-			openai.UserMessage(fmt.Sprintf(`New content: %s`, newContent)),
+	resp, err := a.client.ChatCompletion(ctx, &llm.ChatCompletionRequest{
+		Model: a.model,
+		Messages: []llm.Message{
+			{Role: llm.RoleSystem, Parts: []llm.Part{llm.TextPart{Text: changelogGeneratorSystemPrompt}}},
+			{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: fmt.Sprintf(`Old content: %s`, oldContent)}}},
+			{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: fmt.Sprintf(`New content: %s`, newContent)}}},
 		},
-		Model:       model,
-		Temperature: param.NewOpt(a.cfg.Temperature),
+		Temperature: &a.temp,
 	})
 	if err != nil {
-		return nil, fmt.Errorf("cannot parse vendor info: %w", err)
+		return nil, fmt.Errorf("cannot generate changelog: %w", err)
 	}
 
-	if len(chatCompletion.Choices) == 0 {
-		return nil, fmt.Errorf("no completion choices returned from API")
-	}
-
-	return &chatCompletion.Choices[0].Message.Content, nil
+	text := resp.Message.Text()
+	return &text, nil
 }

pkg/agents/vendor_assessment.go

@@ -19,8 +19,7 @@ import (
 	"encoding/json"
 	"fmt"
 
-	"github.com/openai/openai-go"
-	"github.com/openai/openai-go/packages/param"
+	"go.probo.inc/probo/pkg/llm"
 )
 
 type (
@@ -123,28 +122,25 @@ const (
 )
 
 func (a *Agent) AssessVendor(ctx context.Context, websiteURL string) (*vendorInfo, error) {
-	model := openai.ChatModel(a.cfg.ModelName)
-	chatCompletion, err := a.client.Chat.Completions.New(ctx, openai.ChatCompletionNewParams{
-		Messages: []openai.ChatCompletionMessageParamUnion{
-			openai.SystemMessage(assessVendorSystemPrompt),
-			openai.UserMessage(websiteURL),
+	resp, err := a.client.ChatCompletion(
+		ctx,
+		&llm.ChatCompletionRequest{
+			Model: a.model,
+			Messages: []llm.Message{
+				{Role: llm.RoleSystem, Parts: []llm.Part{llm.TextPart{Text: assessVendorSystemPrompt}}},
+				{Role: llm.RoleUser, Parts: []llm.Part{llm.TextPart{Text: websiteURL}}},
+			},
+			Temperature: &a.temp,
 		},
-		Model:       model,
-		Temperature: param.NewOpt(a.cfg.Temperature),
-	})
+	)
 	if err != nil {
+		return nil, fmt.Errorf("cannot assess vendor: %w", err)
+	}
+
+	var info vendorInfo
+	if err := json.Unmarshal([]byte(resp.Message.Text()), &info); err != nil {
 		return nil, fmt.Errorf("cannot parse vendor info: %w", err)
 	}
 
-	if len(chatCompletion.Choices) == 0 {
-		return nil, fmt.Errorf("no completion choices returned from API")
-	}
-
-	var vendorInfo vendorInfo
-	err = json.Unmarshal([]byte(chatCompletion.Choices[0].Message.Content), &vendorInfo)
-	if err != nil {
-		return nil, fmt.Errorf("cannot parse vendor info: %w", err)
-	}
-
-	return &vendorInfo, nil
+	return &info, nil
 }