Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions backend/cmd/server/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,10 @@ func main() {
panic(err)
}

if err := s.modelV1.InitModel(); err != nil {
panic(err)
}

svc := service.NewService(service.WithPprof())
svc.Add(s)
if err := svc.Run(); err != nil {
Expand Down
4 changes: 2 additions & 2 deletions backend/cmd/server/wire_gen.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 6 additions & 0 deletions backend/config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,12 @@ type Config struct {
VSCode struct {
VSIXFile string `mapstructure:"vsix_file"`
} `mapstructure:"vscode"`

InitModel struct {
ModelName string `mapstructure:"model_name"`
ModelKey string `mapstructure:"model_key"`
ModelURL string `mapstructure:"model_url"`
} `mapstructure:"init_model"`
}

func Init(dir string) (*Config, error) {
Expand Down
6 changes: 5 additions & 1 deletion backend/config/config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -27,4 +27,8 @@ llm_proxy:
client_pool_size: 10
request_log_path: /app/request/logs
vscode:
vsix_file: /app/static/monkeycode.vsix
vsix_file: /app/static/monkeycode.vsix
init_model:
model_name: ""
model_key: ""
model_url: ""
2 changes: 2 additions & 0 deletions backend/domain/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ type ModelUsecase interface {
Update(ctx context.Context, req *UpdateModelReq) (*Model, error)
Check(ctx context.Context, req *CheckModelReq) (*Model, error)
GetTokenUsage(ctx context.Context, modelType consts.ModelType) (*ModelTokenUsageResp, error)
InitModel(ctx context.Context) error
}

type ModelRepo interface {
Expand All @@ -25,6 +26,7 @@ type ModelRepo interface {
MyModelList(ctx context.Context, req *MyModelListReq) ([]*db.Model, error)
ModelUsage(ctx context.Context, ids []uuid.UUID) (map[uuid.UUID]ModelUsage, error)
GetTokenUsage(ctx context.Context, modelType consts.ModelType) (*ModelTokenUsageResp, error)
InitModel(ctx context.Context, modelName, modelKey, modelURL string) error
}

type MyModelListReq struct {
Expand Down
5 changes: 5 additions & 0 deletions backend/internal/model/handler/http/v1/model.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package v1

import (
"context"
"log/slog"

"github.com/GoYoko/web"
Expand Down Expand Up @@ -151,3 +152,7 @@ func (h *ModelHandler) GetTokenUsage(c *web.Context, req domain.GetTokenUsageReq
}
return c.Success(resp)
}

func (h *ModelHandler) InitModel() error {
return h.usecase.InitModel(context.Background())
}
29 changes: 29 additions & 0 deletions backend/internal/model/repo/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ import (

"github.com/chaitin/MonkeyCode/backend/consts"
"github.com/chaitin/MonkeyCode/backend/db"
"github.com/chaitin/MonkeyCode/backend/db/admin"
"github.com/chaitin/MonkeyCode/backend/db/model"
"github.com/chaitin/MonkeyCode/backend/db/task"
"github.com/chaitin/MonkeyCode/backend/domain"
Expand Down Expand Up @@ -177,3 +178,31 @@ func (r *ModelRepo) List(ctx context.Context) (*domain.AllModelResp, error) {
}
return resp, nil
}

func (r *ModelRepo) InitModel(ctx context.Context, modelName, modelKey, modelURL string) error {
n, err := r.db.Model.Query().
Where(model.ModelName(modelName)).
Where(model.Provider("百智云")).
Count(ctx)
if err != nil {
return err
}
if n > 0 {
return nil
}

a, err := r.db.Admin.Query().Where(admin.Username("admin")).Only(ctx)
if err != nil {
return err
}

return r.db.Model.Create().
SetAPIKey(modelKey).
SetModelName(modelName).
SetModelType(consts.ModelTypeCoder).
SetAPIBase(modelURL).
SetProvider("百智云").
SetStatus(consts.ModelStatusActive).
SetUserID(a.ID).
Exec(ctx)
}
19 changes: 16 additions & 3 deletions backend/internal/model/usecase/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,11 @@ package usecase

import (
"context"
"log/slog"

"github.com/google/uuid"

"github.com/chaitin/MonkeyCode/backend/config"
"github.com/chaitin/MonkeyCode/backend/db/model"
"github.com/chaitin/MonkeyCode/backend/pkg/cvt"

Expand All @@ -14,11 +16,17 @@ import (
)

type ModelUsecase struct {
repo domain.ModelRepo
logger *slog.Logger
repo domain.ModelRepo
cfg *config.Config
}

func NewModelUsecase(repo domain.ModelRepo) domain.ModelUsecase {
return &ModelUsecase{repo: repo}
func NewModelUsecase(
logger *slog.Logger,
repo domain.ModelRepo,
cfg *config.Config,
) domain.ModelUsecase {
return &ModelUsecase{repo: repo, cfg: cfg, logger: logger}
}

func (m *ModelUsecase) Check(ctx context.Context, req *domain.CheckModelReq) (*domain.Model, error) {
Expand Down Expand Up @@ -98,3 +106,8 @@ func (m *ModelUsecase) Update(ctx context.Context, req *domain.UpdateModelReq) (
}
return cvt.From(model, &domain.Model{}), nil
}

func (m *ModelUsecase) InitModel(ctx context.Context) error {
m.logger.With("init_model", m.cfg.InitModel).Debug("init model")
return m.repo.InitModel(ctx, m.cfg.InitModel.ModelName, m.cfg.InitModel.ModelKey, m.cfg.InitModel.ModelURL)
}
Loading