fix: update es client

opengovern · Dec 28, 2024 · 088b389 · 088b389
1 parent e3c7d5c
commit 088b389
Show file tree

Hide file tree

Showing 5 changed files with 84 additions and 156 deletions.
diff --git 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
@@ -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,96 +41,87 @@ 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)
 	if err != nil {
 		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 {
 		return &models.Resource{}, err
 	}
 	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)
 	if err != nil {
 		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
@@ -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
 }