Merge pull request #451 from chaitin/feat-model-billing

feat: 实现模型收费

Author: yokowu

backend/biz/git/usecase/identity.go

@@ -69,6 +69,8 @@ func (u *GitIdentityUsecase) gitClienter(identity *db.GitIdentity) domain.GitCli
 		return gitlab.NewGitlab(identity.BaseURL, identity.AccessToken, u.logger)
 	case consts.GitPlatformGitea:
 		return gitea.NewGitea(u.logger, identity.BaseURL)
+	case consts.GitPlatformGitee:
+		return gitee.NewGitee(identity.BaseURL, u.logger)
 	default:
 		return nil
 	}

backend/biz/task/repo/task.go

@@ -260,11 +260,29 @@ func (t *TaskRepo) Create(ctx context.Context, u *domain.User, req domain.Create
 		if err != nil {
 			return errcode.ErrModelAccessDenied.Wrap(err)
 		}
-		m, err := tx.Model.Query().WithUser().Where(model.ID(mid)).First(ctx)
+		m, err := tx.Model.Query().
+			WithPricing().
+			WithUser().
+			Where(model.ID(mid)).
+			First(ctx)
 		if err != nil {
 			return err
 		}
 
+		var mak *db.ModelApiKey
+		if p := m.Edges.Pricing; p != nil {
+			apikey := uuid.NewString()
+			mak, err = tx.ModelApiKey.Create().
+				SetAPIKey(apikey).
+				SetUserID(u.ID).
+				SetModelID(m.ID).
+				Save(ctx)
+			if err != nil {
+				return err
+			}
+			m.Edges.Apikeys = append(m.Edges.Apikeys, mak)
+		}
+
 		img, err := tx.Image.Query().Where(image.ID(req.ImageID)).First(ctx)
 		if err != nil {
 			return err
@@ -345,6 +363,14 @@ func (t *TaskRepo) Create(ctx context.Context, u *domain.User, req domain.Create
 			return err
 		}
 
+		if mak != nil {
+			if err := tx.ModelApiKey.UpdateOneID(mak.ID).
+				SetVirtualmachineID(vm.ID).
+				Exec(ctx); err != nil {
+				return err
+			}
+		}
+
 		return nil
 	})
 

backend/biz/task/usecase/task.go

@@ -317,6 +317,11 @@ func (a *TaskUsecase) Create(ctx context.Context, user *domain.User, req domain.
 			return nil, fmt.Errorf("task edge is nil")
 		}
 
+		if keys := m.Edges.Apikeys; len(keys) > 0 {
+			m.APIKey = keys[0].APIKey
+			m.BaseURL = a.cfg.LLMProxy.BaseURL + "/v1"
+		}
+
 		coding, configs, err := a.getCodingConfigs(req.CliName, m, req.Extra.SkillIDs)
 		if err != nil {
 			return nil, err

backend/config/config.go

@@ -28,6 +28,16 @@ type Config struct {
 	Session Session `mapstructure:"session"`
 	SMTP    SMTP    `mapstructure:"smtp"`
 
+	LLMProxy struct {
+		Addr                 string `mapstructure:"addr"`
+		BaseURL              string `mapstructure:"base_url"`
+		Timeout              string `mapstructure:"timeout"`
+		KeepAlive            string `mapstructure:"keep_alive"`
+		ClientPoolSize       int    `mapstructure:"client_pool_size"`
+		StreamClientPoolSize int    `mapstructure:"stream_client_pool_size"`
+		RequestLogPath       string `mapstructure:"request_log_path"`
+	} `mapstructure:"llm_proxy"`
+
 	RootPath   string        `mapstructure:"root_path"`
 	Logger     logger.Config `mapstructure:"logger"`
 	AdminToken string        `mapstructure:"admin_token"`