增加多个ai

internal/bot/bot.go

@@ -2,7 +2,6 @@ package bot
 
 import (
 	"context"
-	"errors"
 	"fmt"
 
 	"github.com/openai/openai-go"
@@ -13,33 +12,66 @@ import (
 const maxHistory = 20
 
 type Bot struct {
-	client  *openai.Client
-	cfg     *config.Config
-	history []openai.ChatCompletionMessageParamUnion
+	clients      map[string]*openai.Client
+	cfg          *config.Config
+	provider     *config.Provider
+	model        string
+	history      []openai.ChatCompletionMessageParamUnion
+	systemPrompt string
 }
 
 func New(cfg *config.Config) *Bot {
-	client := openai.NewClient(
-		option.WithAPIKey(cfg.APIKey),
-		option.WithBaseURL(cfg.BaseURL),
-	)
-	return &Bot{
-		client: &client,
-		cfg:    cfg,
+	b := &Bot{
+		clients:      make(map[string]*openai.Client),
+		cfg:          cfg,
+		systemPrompt: cfg.SystemPrompt,
 	}
+	b.provider = config.FindProvider(cfg.DefaultProvider)
+	b.model = cfg.DefaultModel
+	return b
+}
+
+func (b *Bot) client() *openai.Client {
+	if c, ok := b.clients[b.provider.Name]; ok {
+		return c
+	}
+	c := openai.NewClient(
+		option.WithAPIKey(b.provider.APIKey),
+		option.WithBaseURL(b.provider.BaseURL),
+	)
+	b.clients[b.provider.Name] = &c
+	return &c
+}
+
+func (b *Bot) Models() []string {
+	return config.AllModels()
+}
+
+func (b *Bot) Current() (string, string) {
+	return b.provider.Name, b.model
+}
+
+func (b *Bot) SwitchModel(id string) error {
+	p, modelName, err := config.ResolveModel(id)
+	if err != nil {
+		return err
+	}
+	b.provider = p
+	b.model = modelName
+	return nil
 }
 
 func (b *Bot) Chat(ctx context.Context, userMsg string) (string, error) {
-	if b.cfg.APIKey == "" {
-		return "", errors.New("未配置 api_key，请编辑 data/config.yaml")
+	if b.provider.APIKey == "" {
+		return "", fmt.Errorf("供应商 %s 未配置 api_key，请编辑 data/config.yaml", b.provider.Name)
 	}
 	history := make([]openai.ChatCompletionMessageParamUnion, 0, len(b.history)+2)
-	history = append(history, openai.SystemMessage(b.cfg.SystemPrompt))
+	history = append(history, openai.SystemMessage(b.systemPrompt))
 	history = append(history, b.history...)
 	history = append(history, openai.UserMessage(userMsg))
 
-	stream := b.client.Chat.Completions.NewStreaming(ctx, openai.ChatCompletionNewParams{
-		Model:    b.cfg.Model,
+	stream := b.client().Chat.Completions.NewStreaming(ctx, openai.ChatCompletionNewParams{
+		Model:    b.model,
 		Messages: history,
 	})
 	answer := ""