From 088b3896a652b1ab8fd70d7bf29dd9b7188c2ce1 Mon Sep 17 00:00:00 2001 From: artaasadi Date: Sat, 28 Dec 2024 20:40:03 +0100 Subject: [PATCH] fix: update es client --- provider/describer/connectors.go | 26 +++--- provider/describer/datasets.go | 55 +++++------- provider/describer/embed-jobs.go | 33 +++---- provider/describer/fine-tuned_models.go | 88 +++++++++---------- .../describer/json_all_fields_marshaller.go | 38 -------- 5 files changed, 84 insertions(+), 156 deletions(-) delete mode 100644 provider/describer/json_all_fields_marshaller.go diff --git a/provider/describer/connectors.go b/provider/describer/connectors.go index ecca6845..05d18923 100644 --- a/provider/describer/connectors.go +++ b/provider/describer/connectors.go @@ -14,11 +14,9 @@ import ( "github.com/opengovern/og-describer-cohereai/provider/model" ) - func ListConnectors(ctx context.Context, handler *CohereAIAPIHandler, stream *models.StreamSender) ([]models.Resource, error) { var wg sync.WaitGroup cohereaiChan := make(chan models.Resource) - go func() { processConnectors(ctx, handler, cohereaiChan, &wg) @@ -45,22 +43,22 @@ func processConnectors(ctx context.Context, handler *CohereAIAPIHandler, cohereA baseURL := "https://api.cohere.com/v1/connectors" params := url.Values{} params.Set("limit", "100") - + finalURL := baseURL + "?" + params.Encode() req, err := http.NewRequest("GET", finalURL, nil) if err != nil { - return + return } requestFunc := func(req *http.Request) (*http.Response, error) { var e error - + resp, e = handler.Client.Do(req) // fmt.Printf(json.NewDecoder(resp.Body)) if e = json.NewDecoder(resp.Body).Decode(&connectorResponse); e != nil { return nil, e } connectors = append(connectors, connectorResponse.Connectors...) - + return resp, e } err = handler.DoRequest(ctx, req, requestFunc) @@ -72,11 +70,9 @@ func processConnectors(ctx context.Context, handler *CohereAIAPIHandler, cohereA go func(connector model.Connector) { defer wg.Done() value := models.Resource{ - ID: connector.ID, - Name: connector.Name, - Description: JSONAllFieldsMarshaller{ - Value: connector, - }, + ID: connector.ID, + Name: connector.Name, + Description: connector, } cohereAiChan <- value }(connector) @@ -106,11 +102,9 @@ func GetConnector(ctx context.Context, handler *CohereAIAPIHandler, connectorID return nil, err } value := &models.Resource{ - ID: connector.ID, - Name: connector.Name, - Description: JSONAllFieldsMarshaller{ - Value: connector, - }, + ID: connector.ID, + Name: connector.Name, + Description: connector, } return value, nil } diff --git a/provider/describer/datasets.go b/provider/describer/datasets.go index 467e1a0e..0ca936e6 100644 --- a/provider/describer/datasets.go +++ b/provider/describer/datasets.go @@ -14,11 +14,9 @@ import ( "github.com/opengovern/og-describer-cohereai/provider/model" ) - func ListDatasets(ctx context.Context, handler *CohereAIAPIHandler, stream *models.StreamSender) ([]models.Resource, error) { var wg sync.WaitGroup cohereaiChan := make(chan models.Resource) - go func() { processDatasets(ctx, handler, cohereaiChan, &wg) @@ -43,22 +41,21 @@ func processDatasets(ctx context.Context, handler *CohereAIAPIHandler, cohereAiC var datasets []model.DatasetDescription var resp *http.Response baseURL := "https://api.cohere.com/v1/datasets" - - + finalURL := baseURL req, err := http.NewRequest("GET", finalURL, nil) if err != nil { - return + return } requestFunc := func(req *http.Request) (*http.Response, error) { var e error - + resp, e = handler.Client.Do(req) // fmt.Printf(json.NewDecoder(resp.Body)) if e = json.NewDecoder(resp.Body).Decode(&datasetResponse); e != nil { return nil, e } - datasets =append(datasets, datasetResponse.Datasets...) + datasets = append(datasets, datasetResponse.Datasets...) return resp, nil } err = handler.DoRequest(ctx, req, requestFunc) @@ -66,49 +63,43 @@ func processDatasets(ctx context.Context, handler *CohereAIAPIHandler, cohereAiC return } // get dataset usage for each dataset - finalURL1 := baseURL + "/usage" - req1, err1:= http.NewRequest("GET", finalURL1, nil) + finalURL1 := baseURL + "/usage" + req1, err1 := http.NewRequest("GET", finalURL1, nil) var usage model.OrganizationUsage requestFunc1 := func(req *http.Request) (*http.Response, error) { - var e error - resp, e = handler.Client.Do(req) - if e = json.NewDecoder(resp.Body).Decode(usage); e != nil { - return nil, e - } - return resp, nil + var e error + resp, e = handler.Client.Do(req) + if e = json.NewDecoder(resp.Body).Decode(usage); e != nil { + return nil, e } + return resp, nil + } err = handler.DoRequest(ctx, req1, requestFunc1) if err1 != nil { return -} + } for _, dataset := range datasets { dataset.TotalUsage = float64(usage.OrganizationUsage) } - - for _, dataset := range datasets { wg.Add(1) go func(dataset model.DatasetDescription) { defer wg.Done() value := models.Resource{ - ID: dataset.ID, - Name: dataset.Name, - Description: JSONAllFieldsMarshaller{ - Value: dataset, - }, + ID: dataset.ID, + Name: dataset.Name, + Description: dataset, } cohereAiChan <- value }(dataset) } } - func GetDataset(ctx context.Context, handler *CohereAIAPIHandler, datasetID string) (*models.Resource, error) { var datasetResponse model.DatasetDescription baseURL := "https://api.cohere.com/v1/datasets" - - + finalURL := baseURL + "/" + datasetID req, err := http.NewRequest("GET", finalURL, nil) if err != nil { @@ -116,12 +107,12 @@ func GetDataset(ctx context.Context, handler *CohereAIAPIHandler, datasetID stri } requestFunc := func(req *http.Request) (*http.Response, error) { var e error - + resp, e := handler.Client.Do(req) if e = json.NewDecoder(resp.Body).Decode(&datasetResponse); e != nil { return nil, e } - + return resp, e } err = handler.DoRequest(ctx, req, requestFunc) @@ -129,10 +120,8 @@ func GetDataset(ctx context.Context, handler *CohereAIAPIHandler, datasetID stri return &models.Resource{}, err } return &models.Resource{ - ID: datasetResponse.ID, - Name: datasetResponse.Name, - Description: JSONAllFieldsMarshaller{ - Value: datasetResponse, - }, + ID: datasetResponse.ID, + Name: datasetResponse.Name, + Description: datasetResponse, }, nil } diff --git a/provider/describer/embed-jobs.go b/provider/describer/embed-jobs.go index 3d327be2..71cd7dc2 100644 --- a/provider/describer/embed-jobs.go +++ b/provider/describer/embed-jobs.go @@ -13,11 +13,9 @@ import ( "github.com/opengovern/og-describer-cohereai/provider/model" ) - func ListEmbedJobs(ctx context.Context, handler *CohereAIAPIHandler, stream *models.StreamSender) ([]models.Resource, error) { var wg sync.WaitGroup cohereaiChan := make(chan models.Resource) - go func() { processEmbedJobs(ctx, handler, cohereaiChan, &wg) @@ -43,26 +41,25 @@ func processEmbedJobs(ctx context.Context, handler *CohereAIAPIHandler, cohereAi var resp *http.Response baseURL := "https://api.cohere.com/v1/embed-jobs" - - finalURL := baseURL + finalURL := baseURL req, err := http.NewRequest("GET", finalURL, nil) if err != nil { - return + return } requestFunc := func(req *http.Request) (*http.Response, error) { var e error - + resp, e = handler.Client.Do(req) // fmt.Printf(json.NewDecoder(resp.Body)) if e = json.NewDecoder(resp.Body).Decode(&embedJobResponse); e != nil { return nil, e } embedJobs = append(embedJobs, embedJobResponse.EmbedJobs...) - + return resp, nil } - err = handler.DoRequest(ctx, req, requestFunc) + err = handler.DoRequest(ctx, req, requestFunc) if err != nil { return } @@ -71,19 +68,15 @@ func processEmbedJobs(ctx context.Context, handler *CohereAIAPIHandler, cohereAi go func(mod model.EmbedJobDescription) { defer wg.Done() value := models.Resource{ - ID: mod.JobID, - Name: mod.Name, - Description: JSONAllFieldsMarshaller{ - Value: mod, - }, + ID: mod.JobID, + Name: mod.Name, + Description: mod, } cohereAiChan <- value }(mod) } - -} - +} func GetEmbedJob(ctx context.Context, handler *CohereAIAPIHandler, embedJobID string) (*models.Resource, error) { var embedJobResponse model.EmbedJobDescription @@ -108,11 +101,9 @@ func GetEmbedJob(ctx context.Context, handler *CohereAIAPIHandler, embedJobID st return nil, err } value := &models.Resource{ - ID: embedJob.JobID, - Name: embedJob.Name, - Description: JSONAllFieldsMarshaller{ - Value: embedJob, - }, + ID: embedJob.JobID, + Name: embedJob.Name, + Description: embedJob, } return value, nil } diff --git a/provider/describer/fine-tuned_models.go b/provider/describer/fine-tuned_models.go index 6d19d996..8ff70ea5 100644 --- a/provider/describer/fine-tuned_models.go +++ b/provider/describer/fine-tuned_models.go @@ -16,11 +16,9 @@ import ( "github.com/opengovern/og-describer-cohereai/provider/model" ) - func ListFineTunedModels(ctx context.Context, handler *CohereAIAPIHandler, stream *models.StreamSender) ([]models.Resource, error) { var wg sync.WaitGroup cohereaiChan := make(chan models.Resource) - go func() { processFineTunedModels(ctx, handler, cohereaiChan, &wg) @@ -47,49 +45,48 @@ func processFineTunedModels(ctx context.Context, handler *CohereAIAPIHandler, co var fineTunedModels []model.FineTunedModelDescription var resp *http.Response baseURL := "https://api.cohere.com/v1/finetuning/finetuned-models" - req,err1 := http.NewRequest("GET", baseURL, nil) - if(err1 != nil){ - return + req, err1 := http.NewRequest("GET", baseURL, nil) + if err1 != nil { + return } requestFunc := func(req *http.Request) (*http.Response, error) { var e error pageToken := "" for { - params := url.Values{} - params.Set("page_size", "1000") - if(pageToken != ""){ - params.Set("pageToken", pageToken) - } + params := url.Values{} + params.Set("page_size", "1000") + if pageToken != "" { + params.Set("pageToken", pageToken) + } - - finalURL := baseURL + "?" + params.Encode() - req, err := http.NewRequest("GET", finalURL, nil) - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", handler.APIKey)) - if handler.ClientName != "" { - req.Header.Set("X-Client-Name", handler.ClientName) - } - if err != nil { - return nil,e - } - - resp, e = handler.Client.Do(req) - // fmt.Printf(json.NewDecoder(resp.Body)) - if e = json.NewDecoder(resp.Body).Decode(&fineTunedModelResponse); e != nil { - return nil, e - } - fineTunedModels = append(fineTunedModels, fineTunedModelResponse.FinetunedModels...) - - if(fineTunedModelResponse.NextPageToken == ""){ - break - } - pageToken = fineTunedModelResponse.NextPageToken + finalURL := baseURL + "?" + params.Encode() + req, err := http.NewRequest("GET", finalURL, nil) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", handler.APIKey)) + if handler.ClientName != "" { + req.Header.Set("X-Client-Name", handler.ClientName) + } + if err != nil { + return nil, e + } - } - return resp, e + resp, e = handler.Client.Do(req) + // fmt.Printf(json.NewDecoder(resp.Body)) + if e = json.NewDecoder(resp.Body).Decode(&fineTunedModelResponse); e != nil { + return nil, e + } + fineTunedModels = append(fineTunedModels, fineTunedModelResponse.FinetunedModels...) + + if fineTunedModelResponse.NextPageToken == "" { + break + } + pageToken = fineTunedModelResponse.NextPageToken + + } + return resp, e } - err := handler.DoRequest(ctx, req, requestFunc) + err := handler.DoRequest(ctx, req, requestFunc) if err != nil { return } @@ -98,18 +95,15 @@ func processFineTunedModels(ctx context.Context, handler *CohereAIAPIHandler, co go func(mod model.FineTunedModelDescription) { defer wg.Done() value := models.Resource{ - ID: mod.ID, - Name: mod.Name, - Description: JSONAllFieldsMarshaller{ - Value: mod, - }, + ID: mod.ID, + Name: mod.Name, + Description: mod, } cohereAiChan <- value }(mod) } } - func GetFineTunedModel(ctx context.Context, handler *CohereAIAPIHandler, fineTunedModelID string) (*models.Resource, error) { var fineTunedModelResponse model.FineTunedModelDescription var fineTunedModel model.FineTunedModelDescription @@ -133,11 +127,9 @@ func GetFineTunedModel(ctx context.Context, handler *CohereAIAPIHandler, fineTun return nil, err } value := models.Resource{ - ID: fineTunedModel.ID, - Name: fineTunedModel.Name, - Description: JSONAllFieldsMarshaller{ - Value: fineTunedModel, - }, + ID: fineTunedModel.ID, + Name: fineTunedModel.Name, + Description: fineTunedModel, } return &value, nil -} \ No newline at end of file +} diff --git a/provider/describer/json_all_fields_marshaller.go b/provider/describer/json_all_fields_marshaller.go deleted file mode 100644 index 89f10291..00000000 --- a/provider/describer/json_all_fields_marshaller.go +++ /dev/null @@ -1,38 +0,0 @@ -package describer - -import ( - "encoding/json" - "reflect" -) - -type JSONAllFieldsMarshaller struct { - Value interface{} -} - -func (x JSONAllFieldsMarshaller) MarshalJSON() (res []byte, err error) { - var val = x.Value - v := reflect.ValueOf(x.Value) - if !v.IsValid() { - return json.Marshal(val) - } - return json.Marshal(val) -} - -func (x *JSONAllFieldsMarshaller) UnmarshalJSON(data []byte) (err error) { - v := reflect.ValueOf(x.Value) - if !v.IsValid() { - return nil - } - val := reflect.New(v.Type()) - err = json.Unmarshal(data, val.Interface()) - if err != nil { - return err - } - newVal := reflect.New(v.Type()) - if !val.Elem().Type().AssignableTo(newVal.Elem().Type()) { - return nil - } - newVal.Elem().Set(val.Elem()) - x.Value = newVal.Elem().Interface() - return nil -}