aiforge/modules/modelarts/resty.go

package modelarts

import (
	"code.gitea.io/gitea/modules/log"
	"fmt"
	"net/http"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/setting"
	"github.com/go-resty/resty/v2"
)

var (
	restyClient *resty.Client
	HOST        string
	TOKEN       string
)

const (
	methodPassword = "password"

	urlGetToken = "/v3/auth/tokens"
	urlNotebook = "/demanager/instances"
	urlQuaryNotebook = "/demanager/instances"
)
func getRestyClient() *resty.Client {
	if restyClient == nil {
		restyClient = resty.New()
	}
	return restyClient
}

func checkSetting() {
	if len(HOST) != 0 && len(TOKEN) != 0 && restyClient != nil {
		return
	}

	getToken()
}

func getToken() error {
	HOST = setting.ModelArtsHost

	client := getRestyClient()
	params := models.GetTokenParams{
		Auth: models.Auth{
			Identity: models.Identity{
				Methods: []string{methodPassword},
				Password: models.Password{
					User: models.NotebookUser{
						Name: setting.ModelArtsUsername,
						Password: setting.ModelArtsPassword,
						Domain: models.Domain{
							Name: setting.ModelArtsDomain,
						},
					},
				},
			},
			Scope: models.Scope{
				Project: models.Project{
					Name: setting.ProjectName,
				},
			},
		},
	}

	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetBody(params).
		Post(HOST + urlGetToken)
	if err != nil {
		return fmt.Errorf("resty getToken: %s", err)
	}

	if res.StatusCode() != http.StatusCreated {
		return fmt.Errorf("getToken failed:%s", res.String())
	}

	TOKEN = res.Header().Get("X-Subject-Token")

	return nil
}

func CreateJob(createJobParams models.CreateNotebookParams) (*models.CreateNotebookResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.CreateNotebookResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetBody(createJobParams).
		SetResult(&result).
		Post(HOST + "/v1/" + setting.ProjectID + urlNotebook)

	if err != nil {
		return nil, fmt.Errorf("resty create job: %s", err)
	}

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if len(result.ErrorCode) != 0 {
		log.Error("CreateJob failed(%s): %s", result.ErrorCode, result.ErrorMsg)
		return &result, fmt.Errorf("CreateJob failed(%s): %s", result.ErrorCode, result.ErrorMsg)
	}

	return &result, nil
}

func GetJob(jobID string) (*models.GetNotebookResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GetNotebookResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetResult(&result).
		Get(HOST + "/v1/" + setting.ProjectID + urlNotebook + jobID)

	if err != nil {
		return nil, fmt.Errorf("resty GetJob: %v", err)
	}

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if len(result.ErrorCode) != 0 {
		return &result, fmt.Errorf("CreateJob failed(%s): %s", result.ErrorCode, result.ErrorMsg)
	}

	return &result, nil
}

func StopJob(jobID string, param models.NotebookAction) (*models.NotebookActionResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.NotebookActionResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetBody(param).
		SetAuthToken(TOKEN).
		SetResult(&result).
		Post(HOST + "/v1/" + setting.ProjectID + urlNotebook + jobID + "/action")

	if err != nil {
		return &result, fmt.Errorf("resty StopJob: %v", err)
	}

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if len(result.ErrorCode) != 0 {
		return &result, fmt.Errorf("CreateJob failed(%s): %s", result.ErrorCode, result.ErrorMsg)
	}

	return &result, nil
}