chat_model.go 2.5 KB
package domain

import (
	"context"
	"github.com/samber/lo"
	"github.com/sashabaranov/go-openai"
	"gitlab.fjmaimaimai.com/allied-creation/su-micro/pkg/transaction"
)

type ChatModel struct {
	Id        int64       `json:",omitempty"` // 唯一标识
	Name      string      `json:",omitempty"`
	Code      string      `json:",omitempty"`
	Logo      string      `json:",omitempty"`
	Config    ModelConfig `json:",omitempty"`
	CreatedAt int64       `json:",omitempty"`
	UpdatedAt int64       `json:",omitempty"`
	DeletedAt int64       `json:",omitempty"`
	Version   int         `json:",omitempty"`
}

type ChatModelRepository interface {
	Insert(ctx context.Context, conn transaction.Conn, dm *ChatModel) (*ChatModel, error)
	Update(ctx context.Context, conn transaction.Conn, dm *ChatModel) (*ChatModel, error)
	UpdateWithVersion(ctx context.Context, conn transaction.Conn, dm *ChatModel) (*ChatModel, error)
	Delete(ctx context.Context, conn transaction.Conn, dm *ChatModel) (*ChatModel, error)
	FindOne(ctx context.Context, conn transaction.Conn, id int64) (*ChatModel, error)
	FindOneUnscoped(ctx context.Context, conn transaction.Conn, id int64) (*ChatModel, error)
	Find(ctx context.Context, conn transaction.Conn, queryOptions map[string]interface{}) (int64, []*ChatModel, error)
}

type ChatModels []*ChatModel

func (list ChatModels) Match(id int64) (*ChatModel, bool) {
	return lo.Find(list, func(item *ChatModel) bool {
		if item.Id == id {
			return true
		}
		return false
	})
}

var DefaultChatModels ChatModels = []*ChatModel{
	NewChatModels(1, "ChatGPT 4o", openai.GPT4o, "", NewModelConfig("", "sk-proj-0odTHgsYqzGpjPMWYuw2T3BlbkFJpaIsQ9I6j8kYc8P7l60H", "")),
	NewChatModels(2, "ChatGPT 4-turbo", openai.GPT4Turbo, "", NewModelConfig("", "sk-proj-0odTHgsYqzGpjPMWYuw2T3BlbkFJpaIsQ9I6j8kYc8P7l60H", "")),
	NewChatModels(3, "ChatGPT 3.5", openai.GPT3Dot5Turbo, "", NewModelConfig("", "sk-proj-0odTHgsYqzGpjPMWYuw2T3BlbkFJpaIsQ9I6j8kYc8P7l60H", "")),
	NewChatModels(4, "星火大模型V3.5", "spark3.5", "", NewModelConfig("4fd8694e", "4a4081a20e9ba0fb1b9686ed93221989", "NTVkM2FjNzk2NzQ5MzBkNWMwYTUwNjAz")),
}

type ModelConfig struct {
	AppId     string
	AppKey    string
	AppSecret string
}

func NewModelConfig(appid, appKey, appSecret string) ModelConfig {
	return ModelConfig{
		AppId:     appid,
		AppKey:    appKey,
		AppSecret: appSecret,
	}
}

func NewChatModels(id int64, name, code string, logo string, config ModelConfig) *ChatModel {
	return &ChatModel{
		Id:     id,
		Name:   name,
		Code:   code,
		Logo:   logo,
		Config: config,
	}
}