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 } err := getToken() if err != nil { log.Error("getToken falied:", err.Error()) } } 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 }