From 92a09a7bf7a4ef1b73bfd88ea41d5df324c5254c Mon Sep 17 00:00:00 2001 From: Mohamad Choupan Date: Mon, 2 Dec 2024 03:50:02 +0330 Subject: [PATCH] feat: remove unused tables --- steampipe-plugin-cohereai/cohereai/plugin.go | 13 +- .../cohereai/table_cohereai_classify.go | 121 --------------- .../table_cohereai_detect_language.go | 88 ----------- .../cohereai/table_cohereai_detokenize.go | 82 ---------- .../cohereai/table_cohereai_embed.go | 111 -------------- .../cohereai/table_cohereai_generation.go | 141 ------------------ .../cohereai/table_cohereai_summarize.go | 112 -------------- .../cohereai/table_cohereai_tokenize.go | 85 ----------- 8 files changed, 6 insertions(+), 747 deletions(-) delete mode 100644 steampipe-plugin-cohereai/cohereai/table_cohereai_classify.go delete mode 100644 steampipe-plugin-cohereai/cohereai/table_cohereai_detect_language.go delete mode 100644 steampipe-plugin-cohereai/cohereai/table_cohereai_detokenize.go delete mode 100644 steampipe-plugin-cohereai/cohereai/table_cohereai_embed.go delete mode 100644 steampipe-plugin-cohereai/cohereai/table_cohereai_generation.go delete mode 100644 steampipe-plugin-cohereai/cohereai/table_cohereai_summarize.go delete mode 100644 steampipe-plugin-cohereai/cohereai/table_cohereai_tokenize.go diff --git a/steampipe-plugin-cohereai/cohereai/plugin.go b/steampipe-plugin-cohereai/cohereai/plugin.go index b125b9ea..22fcc8e9 100644 --- a/steampipe-plugin-cohereai/cohereai/plugin.go +++ b/steampipe-plugin-cohereai/cohereai/plugin.go @@ -19,13 +19,12 @@ func Plugin(ctx context.Context) *plugin.Plugin { ShouldIgnoreError: isNotFoundError, }, TableMap: map[string]*plugin.Table{ - "cohereai_classify": tableCohereClassification(ctx), - "cohereai_detect_language": tableCohereDetectLanguage(ctx), - "cohereai_detokenize": tableCohereDetokenize(ctx), - "cohereai_embed": tableCohereEmbed(ctx), - "cohereai_generation": tableCohereGeneration(ctx), - "cohereai_summarize": tableCohereSummarize(ctx), - "cohereai_tokenize": tableCohereTokenize(ctx), + "cohereai_connectors": tableCohereConnectors(ctx), + "cohereai_models": tableCohereModels(ctx), + "cohereai_datasets": tableCohereDatasets(ctx), + "cohereai_fine_tuned_models": tableCohereFineTunedModels(ctx), + "cohereai_embed_jobs": tableCohereEmbedJobs(ctx), + }, } for key, table := range p.TableMap { diff --git a/steampipe-plugin-cohereai/cohereai/table_cohereai_classify.go b/steampipe-plugin-cohereai/cohereai/table_cohereai_classify.go deleted file mode 100644 index cdd77c24..00000000 --- a/steampipe-plugin-cohereai/cohereai/table_cohereai_classify.go +++ /dev/null @@ -1,121 +0,0 @@ -package cohere - -import ( - "context" - "encoding/json" - - coherego "github.com/cohere-ai/cohere-go" - "github.com/turbot/steampipe-plugin-sdk/v5/grpc/proto" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin/transform" -) - -// Defines the cohereai_classification table -func tableCohereClassification(ctx context.Context) *plugin.Table { - return &plugin.Table{ - Name: "cohereai_classify", - Description: "Classification in Cohere AI.", - List: &plugin.ListConfig{ - Hydrate: listClassification, - KeyColumns: []*plugin.KeyColumn{ - {Name: "inputs", Require: plugin.Optional}, - {Name: "examples", Require: plugin.Optional}, - {Name: "settings", Require: plugin.Optional}, - }, - }, - Columns: []*plugin.Column{ - // Columns returned from the Cohere API - {Name: "id", Type: proto.ColumnType_STRING, Transform: transform.FromField("Classification.ID"), Description: "The ID of the classification."}, - {Name: "classification", Type: proto.ColumnType_STRING, Transform: transform.FromField("Classification.Prediction"), Description: "The classification results for the given input text(s)."}, - {Name: "confidence", Type: proto.ColumnType_DOUBLE, Transform: transform.FromField("Classification.Confidence"), Description: "The confidence score of the classification."}, - {Name: "labels", Type: proto.ColumnType_JSON, Transform: transform.FromField("Classification.Labels"), Description: "The labels of the classification."}, - - // Qual columns to provide input to the API - {Name: "inputs", Type: proto.ColumnType_STRING, Transform: transform.FromQual("inputs"), Description: "The input text that was classified."}, - {Name: "examples", Type: proto.ColumnType_STRING, Transform: transform.FromQual("examples"), Description: "The example text classified."}, - {Name: "settings", Type: proto.ColumnType_JSON, Transform: transform.FromQual("settings"), Description: "Settings is a JSONB object that accepts any of the classify API request parameters."}, - }, - } -} - -// ClassificationRequestQual defines the structure of the settings qual -type ClassificationRequestQual struct { - Model *string `json:"model,omitempty"` - Inputs *[]string `json:"inputs,omitempty"` - Examples *[]coherego.Example `json:"examples,omitempty"` - Preset *string `json:"preset,omitempty"` - Truncate string `json:"truncate,omitempty"` -} - -// ClassificationRow defines the row structure returned from the API -type ClassificationRow struct { - coherego.Classification - Input []string -} - -// listClassification handles querying the Cohere AI API and returning a list of labels -func listClassification(ctx context.Context, d *plugin.QueryData, _ *plugin.HydrateData) (interface{}, error) { - - conn, err := connect(ctx, d) - if err != nil { - plugin.Logger(ctx).Error("cohereai_classification.listClassification", "connection_error", err) - return nil, err - } - - // Default settings taken from the Cohere API docs - // https://docs.cohere.ai/reference/classify - var inputs []string - var examples []coherego.Example - inputsString := d.EqualsQuals["inputs"].GetStringValue() - if inputsString != "" { - err := json.Unmarshal([]byte(inputsString), &inputs) - if err != nil { - plugin.Logger(ctx).Error("cohereai_classification.listClassification", "connection_error", err) - return nil, err - } - } - exampleString := d.EqualsQuals["examples"].GetStringValue() - if exampleString != "" { - err := json.Unmarshal([]byte(exampleString), &examples) - if err != nil { - plugin.Logger(ctx).Error("cohereai_classification.listClassification", "connection_error", err) - return nil, err - } - } - cr := coherego.ClassifyOptions{ - Model: "large", - Inputs: inputs, - Examples: examples, - } - - settingsString := d.EqualsQuals["settings"].GetJsonbValue() - if settingsString != "" { - var crQual ClassificationRequestQual - err := json.Unmarshal([]byte(settingsString), &crQual) - if err != nil { - plugin.Logger(ctx).Error("cohereai_classification.listClassification", "connection_error", err) - return nil, err - } - if crQual.Model != nil { - cr.Model = *crQual.Model - } - if crQual.Preset != nil { - cr.Preset = *crQual.Preset - } - } - - // Query the Cohere API - resp, err := conn.Classify(cr) - if err != nil { - plugin.Logger(ctx).Error("cohereai_classification.listClassification", "api_error", err) - return nil, err - } - - plugin.Logger(ctx).Trace("cohereai_classification.listClassification", "response", resp) - // Return tokenize data - for _, c := range resp.Classifications { - row := ClassificationRow{c, cr.Inputs} - d.StreamListItem(ctx, row) - } - return nil, nil -} diff --git a/steampipe-plugin-cohereai/cohereai/table_cohereai_detect_language.go b/steampipe-plugin-cohereai/cohereai/table_cohereai_detect_language.go deleted file mode 100644 index a983cca1..00000000 --- a/steampipe-plugin-cohereai/cohereai/table_cohereai_detect_language.go +++ /dev/null @@ -1,88 +0,0 @@ -package cohere - -import ( - "context" - "encoding/json" - - coherego "github.com/cohere-ai/cohere-go" - "github.com/turbot/steampipe-plugin-sdk/v5/grpc/proto" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin/transform" -) - -// Defines the cohere_detect_language table -func tableCohereDetectLanguage(ctx context.Context) *plugin.Table { - return &plugin.Table{ - Name: "cohereai_detect_language", - Description: "Detect languages of texts using Cohere AI.", - List: &plugin.ListConfig{ - Hydrate: detectLanguage, - KeyColumns: []*plugin.KeyColumn{ - {Name: "texts", Require: plugin.Optional}, - }, - }, - Columns: []*plugin.Column{ - // Columns returned from the Cohere API - {Name: "language_name", Type: proto.ColumnType_STRING, Transform: transform.FromField("LanguageDetectResult.LanguageName"), Description: "The name of the detected language."}, - {Name: "language_code", Type: proto.ColumnType_STRING, Transform: transform.FromField("LanguageDetectResult.LanguageCode"), Description: "The ISO 639-1 code for the detected language."}, - - // Qual columns to provide input to the API - {Name: "text", Type: proto.ColumnType_STRING, Transform: transform.FromField("Text"), Description: "Text from which to detect it's language."}, - {Name: "texts", Type: proto.ColumnType_STRING, Transform: transform.FromQual("texts"), Description: "The texts to detect languages for, encoded as a JSON array."}, - }, - } -} - -// DetectLanguageRequestQual defines the structure of the settings qual -type DetectLanguageRequestQual struct { - Texts []string `json:"texts"` -} - -// DetectLanguageRow defines the row structure returned from the API -type DetectLanguageRow struct { - coherego.LanguageDetectResult - Texts []string - Text string -} - -// detectLanguage handles querying the Cohere AI API and returning detected language name and code -func detectLanguage(ctx context.Context, d *plugin.QueryData, _ *plugin.HydrateData) (interface{}, error) { - // Create the API client - client, err := connect(ctx, d) - if err != nil { - return nil, err - } - - textString := d.EqualsQuals["texts"].GetStringValue() - var texts []string - if textString != "" { - err := json.Unmarshal([]byte(textString), &texts) - if err != nil { - plugin.Logger(ctx).Error("cohereai_detect_language.detectLanguage", "unmarshal_error", err) - return nil, err - } - } - // Build the DetectLanguageOptions from the request - opts := coherego.DetectLanguageOptions{ - Texts: texts, - } - - // Make the DetectLanguage API call - resp, err := client.DetectLanguage(opts) - if err != nil { - plugin.Logger(ctx).Error("cohereai_detect_language.detectLanguage", "api_error", err) - return nil, err - } - - plugin.Logger(ctx).Debug("cohereai_detect_language.detectLanguage", "response", resp) - // Return detect language data - for i, result := range resp.Results { - rows := DetectLanguageRow{ - result, - opts.Texts, - opts.Texts[i], - } - d.StreamListItem(ctx, rows) - } - return nil, nil -} diff --git a/steampipe-plugin-cohereai/cohereai/table_cohereai_detokenize.go b/steampipe-plugin-cohereai/cohereai/table_cohereai_detokenize.go deleted file mode 100644 index f9527d21..00000000 --- a/steampipe-plugin-cohereai/cohereai/table_cohereai_detokenize.go +++ /dev/null @@ -1,82 +0,0 @@ -package cohere - -import ( - "context" - "encoding/json" - - coherego "github.com/cohere-ai/cohere-go" - "github.com/turbot/steampipe-plugin-sdk/v5/grpc/proto" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin/transform" -) - -// Defines the cohere_detokenize table -func tableCohereDetokenize(ctx context.Context) *plugin.Table { - return &plugin.Table{ - Name: "cohereai_detokenize", - Description: "Detokenize tokens into text using Cohere AI.", - List: &plugin.ListConfig{ - Hydrate: detokenize, - KeyColumns: []*plugin.KeyColumn{ - {Name: "tokens", Require: plugin.Optional}, - }, - }, - Columns: []*plugin.Column{ - // Columns returned from the Cohere API - {Name: "text", Type: proto.ColumnType_STRING, Transform: transform.FromField("Text"), Description: "The detokenized text."}, - - // Qual columns to provide input to the API - {Name: "tokens", Type: proto.ColumnType(proto.ColumnType_STRING.Number()), Transform: transform.FromQual("tokens"), Description: "The tokens to detokenize, encoded as a JSON array."}, - }, - } -} - -// DetokenizeRequestQual defines the structure of the settings qual -type DetokenizeRequestQual struct { - Tokens []int64 `json:"tokens"` -} - -// DetokenizeRow defines the row structure returned from the API -type DetokenizeRow struct { - Text string - Tokens []int64 -} - -// detokenize handles querying the Cohere AI API and returning detokenized data as text -func detokenize(ctx context.Context, d *plugin.QueryData, _ *plugin.HydrateData) (interface{}, error) { - // Create the API client - client, err := connect(ctx, d) - if err != nil { - return nil, err - } - - // Build the DetokenizeOptions from the request - tokenList := d.EqualsQuals["tokens"].GetStringValue() - var tokens []int64 - if tokenList != "" { - err := json.Unmarshal([]byte(tokenList), &tokens) - if err != nil { - plugin.Logger(ctx).Error("cohereai_detokenize.detokenize", "connection_error", err) - return nil, err - } - } - opts := coherego.DetokenizeOptions{ - Tokens: tokens, - } - - // Make the Detokenize API call - resp, err := client.Detokenize(opts) - if err != nil { - plugin.Logger(ctx).Error("cohereai_detokenize.detokenize", "api_error", err) - return nil, err - } - - plugin.Logger(ctx).Debug("cohereai_detokenize.detokenize", "text", resp.Text) - // Return detokenize data - row := DetokenizeRow{ - Text: resp.Text, - Tokens: opts.Tokens, - } - d.StreamListItem(ctx, row) - return nil, nil -} diff --git a/steampipe-plugin-cohereai/cohereai/table_cohereai_embed.go b/steampipe-plugin-cohereai/cohereai/table_cohereai_embed.go deleted file mode 100644 index f306c158..00000000 --- a/steampipe-plugin-cohereai/cohereai/table_cohereai_embed.go +++ /dev/null @@ -1,111 +0,0 @@ -package cohere - -import ( - "context" - "encoding/json" - - coherego "github.com/cohere-ai/cohere-go" - "github.com/turbot/steampipe-plugin-sdk/v5/grpc/proto" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin/transform" -) - -// Defines the cohere_embed table -func tableCohereEmbed(ctx context.Context) *plugin.Table { - return &plugin.Table{ - Name: "cohereai_embed", - Description: "Get embeddings from Cohere AI.", - List: &plugin.ListConfig{ - Hydrate: embed, - KeyColumns: []*plugin.KeyColumn{ - {Name: "texts", Require: plugin.Optional}, - {Name: "settings", Require: plugin.Optional}, - }, - }, - Columns: []*plugin.Column{ - // Columns returned from the Cohere API - {Name: "embeddings", Type: proto.ColumnType_JSON, Transform: transform.FromField("Embeddings"), Description: "Embeddings for the given texts."}, - {Name: "text", Type: proto.ColumnType_STRING, Transform: transform.FromField("Text"), Description: "The texts to embed, encoded as a JSON array."}, - - // Qual columns to provide input to the API - {Name: "texts", Type: proto.ColumnType_STRING, Transform: transform.FromQual("texts"), Description: "The texts to embed, encoded as a JSON array."}, - {Name: "settings", Type: proto.ColumnType_JSON, Transform: transform.FromQual("settings"), Description: "Settings is a JSONB object that accepts any of the embed API request parameters."}, - }, - } -} - -// EmbedRequestQual defines the structure of the settings qual -type EmbedRequestQual struct { - Model *string `json:"model,omitempty"` - Texts *[]string `json:"texts,omitempty"` - Truncate *string `json:"truncate,omitempty"` -} - -// EmbedRow defines the row structure returned from the API -type EmbedRow struct { - Embeddings []float64 - Texts []string - Text string -} - -// embed handles querying the Cohere AI API and returning embedings -func embed(ctx context.Context, d *plugin.QueryData, _ *plugin.HydrateData) (interface{}, error) { - // Create the API client - client, err := connect(ctx, d) - if err != nil { - return nil, err - } - - // Build the EmbedOptions from the request - textString := d.EqualsQuals["texts"].GetStringValue() - var texts []string - if textString != "" { - err := json.Unmarshal([]byte(textString), &texts) - if err != nil { - plugin.Logger(ctx).Error("cohereai_embed.embed", "connection_error", err) - return nil, err - } - } - opts := coherego.EmbedOptions{ - Model: "small", - Texts: texts, - } - settingString := d.EqualsQuals["settings"].GetJsonbValue() - if settingString != "" { - var crQual EmbedRequestQual - err := json.Unmarshal([]byte(settingString), &crQual) - if err != nil { - plugin.Logger(ctx).Error("cohereai_embed.embed", "connection_error", err) - return nil, err - } - - if crQual.Model != nil { - opts.Model = *crQual.Model - } - if crQual.Truncate != nil { - opts.Truncate = *crQual.Truncate - } - if crQual.Texts != nil { - opts.Texts = *crQual.Texts - } - } - - // Make the Embed API call - resp, err := client.Embed(opts) - if err != nil { - plugin.Logger(ctx).Error("cohereai_embed.embed", "api_error", err) - return nil, err - } - - plugin.Logger(ctx).Debug("cohereai_embed.embed", "embeddings", resp.Embeddings) - // Return embed data - for i, result := range resp.Embeddings { - rows := EmbedRow{ - result, - texts, - texts[i], - } - d.StreamListItem(ctx, rows) - } - return nil, nil -} diff --git a/steampipe-plugin-cohereai/cohereai/table_cohereai_generation.go b/steampipe-plugin-cohereai/cohereai/table_cohereai_generation.go deleted file mode 100644 index 4b63771d..00000000 --- a/steampipe-plugin-cohereai/cohereai/table_cohereai_generation.go +++ /dev/null @@ -1,141 +0,0 @@ -package cohere - -import ( - "context" - "encoding/json" - - coherego "github.com/cohere-ai/cohere-go" - "github.com/turbot/steampipe-plugin-sdk/v5/grpc/proto" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin/transform" -) - -// Defines the cohereai_generation table -func tableCohereGeneration(ctx context.Context) *plugin.Table { - return &plugin.Table{ - Name: "cohereai_generation", - Description: "Generation in Cohere AI.", - List: &plugin.ListConfig{ - Hydrate: listGeneration, - KeyColumns: []*plugin.KeyColumn{ - {Name: "prompt", Require: plugin.Optional}, - {Name: "settings", Require: plugin.Optional}, - }, - }, - Columns: []*plugin.Column{ - // Columns returned from the Cohere API - {Name: "generation", Type: proto.ColumnType_STRING, Transform: transform.FromField("Generation.Text"), Description: "Generation for a given text prompt."}, - {Name: "likelihood", Type: proto.ColumnType_DOUBLE, Transform: transform.FromField("Generation.Likelihood"), Description: "The likelihood of the generated text."}, - {Name: "token_likelihoods", Type: proto.ColumnType_JSON, Transform: transform.FromField("Generation.TokenLikelihoods"), Description: "The likelihood of the generated tokens/prompt."}, - - // Qual columns to provide input to the API - {Name: "prompt", Type: proto.ColumnType_STRING, Transform: transform.FromQual("prompt"), Description: "The prompt to get generations for, encoded as a string."}, - {Name: "settings", Type: proto.ColumnType_JSON, Transform: transform.FromQual("settings"), Description: "Settings is a JSONB object that accepts any of the generate API request parameters."}, - }, - } -} - -// GenerationRequestQual defines the structure of the settings qual -type GenerationRequestQual struct { - Model *string `json:"model,omitempty"` - Prompt *string `json:"prompt,omitempty"` - MaxTokens *uint `json:"max_tokens,omitempty"` - NumGenerations *int `json:"num_generations,omitempty"` - Temperature *float64 `json:"temperature,omitempty"` - TopP *float64 `json:"top_p,omitempty"` - FrequencyPenalty *float64 `json:"frequency_penalty,omitempty"` - PresencePenalty *float64 `json:"presence_penalty,omitempty"` - Stop []string `json:"stop,omitempty"` - Preset string `json:"preset,omitempty"` - ReturnLikelihoods *string `json:"return_likelihoods,omitempty"` -} - -// GenerationRow defines the row structure returned from the API -type GenerationRow struct { - coherego.Generation - Prompt string -} - -// listGeneration handles querying the Cohere AI API and returning generation data -func listGeneration(ctx context.Context, d *plugin.QueryData, _ *plugin.HydrateData) (interface{}, error) { - - conn, err := connect(ctx, d) - if err != nil { - plugin.Logger(ctx).Error("cohereai_generation.listCompletion", "connection_error", err) - return nil, err - } - - // Default settings taken from the Cohere API docs - // https://docs.cohere.ai/reference/generate - var maxTokens uint = 100 - var numGenerations int = 3 - cr := coherego.GenerateOptions{ - Prompt: d.EqualsQuals["prompt"].GetStringValue(), - MaxTokens: &maxTokens, - NumGenerations: &numGenerations, - ReturnLikelihoods: "GENERATION", - } - - settingsString := d.EqualsQuals["settings"].GetJsonbValue() - if settingsString != "" { - // Overwrite any settings provided in the settings qual. If a field - // is not passed in the settings, then default to the settings above. - var crQual GenerationRequestQual - err := json.Unmarshal([]byte(settingsString), &crQual) - if err != nil { - plugin.Logger(ctx).Error("cohereai_generation.listGeneration", "connection_error", err) - return nil, err - } - if crQual.Model != nil { - cr.Model = *crQual.Model - } - if crQual.Prompt != nil { - cr.Prompt = *crQual.Prompt - } - if crQual.MaxTokens != nil { - cr.MaxTokens = crQual.MaxTokens - } - if crQual.Temperature != nil { - cr.Temperature = crQual.Temperature - } - if crQual.TopP != nil { - cr.P = crQual.TopP - } - if crQual.FrequencyPenalty != nil { - cr.FrequencyPenalty = crQual.FrequencyPenalty - } - if crQual.PresencePenalty != nil { - cr.PresencePenalty = crQual.PresencePenalty - } - if crQual.Stop != nil { - cr.StopSequences = crQual.Stop - } - if crQual.NumGenerations != nil { - cr.NumGenerations = crQual.NumGenerations - } - if crQual.Preset != "" { - cr.Preset = crQual.Preset - } - if cr.Prompt == "" { - cr.Prompt = *crQual.Prompt - } - if crQual.ReturnLikelihoods != nil { - cr.ReturnLikelihoods = *crQual.ReturnLikelihoods - } - } - - // Query the Cohere API - resp, err := conn.Generate(cr) - if err != nil { - plugin.Logger(ctx).Error("cohereai_generation.listGeneration", "api_error", err) - return nil, err - } - - plugin.Logger(ctx).Trace("cohereai_generation.listGeneration", "response", resp) - // Return generation data - for _, c := range resp.Generations { - row := GenerationRow{c, cr.Prompt} - d.StreamListItem(ctx, row) - } - return nil, nil -} diff --git a/steampipe-plugin-cohereai/cohereai/table_cohereai_summarize.go b/steampipe-plugin-cohereai/cohereai/table_cohereai_summarize.go deleted file mode 100644 index 2b2459b5..00000000 --- a/steampipe-plugin-cohereai/cohereai/table_cohereai_summarize.go +++ /dev/null @@ -1,112 +0,0 @@ -package cohere - -import ( - "context" - "encoding/json" - - coherego "github.com/cohere-ai/cohere-go" - "github.com/turbot/steampipe-plugin-sdk/v5/grpc/proto" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin/transform" -) - -// Defines the cohere_summarize table -func tableCohereSummarize(ctx context.Context) *plugin.Table { - return &plugin.Table{ - Name: "cohereai_summarize", - Description: "Summarize text using Cohere AI.", - List: &plugin.ListConfig{ - Hydrate: summarize, - KeyColumns: []*plugin.KeyColumn{ - {Name: "text", Require: plugin.Optional}, - {Name: "settings", Require: plugin.Optional}, - }, - }, - Columns: []*plugin.Column{ - // Columns returned from the Cohere API - {Name: "id", Type: proto.ColumnType_STRING, Transform: transform.FromField("SummarizeResponse.ID"), Description: "ID for the given text."}, - {Name: "summary", Type: proto.ColumnType_STRING, Transform: transform.FromField("SummarizeResponse.Summary"), Description: "Summary for the given text."}, - - // Qual columns to provide input to the API - {Name: "text", Type: proto.ColumnType_STRING, Transform: transform.FromQual("text"), Description: "The text to summarize, encoded as a string."}, - {Name: "settings", Type: proto.ColumnType_JSON, Transform: transform.FromQual("settings"), Description: "Settings is a JSONB object that accepts any of the summarize API request parameters."}, - }, - } -} - -// SummarizeRequestQual defines the structure of the settings qual -type SummarizeRequestQual struct { - Text *string `json:"text,omitempty"` - Format *string `json:"format,omitempty"` - Length *string `json:"length,omitempty"` - Extractiveness *string `json:"extractiveness,omitempty"` - Temperature *float64 `json:"temperature,omitempty"` - AdditionalCommand *string `json:"additional_command,omitempty"` - Model *string `json:"model,omitempty"` -} - -// SummarizeRow defines the row structure returned from the API -type SummarizeRow struct { - coherego.SummarizeResponse - Text string -} - -// summarize handles querying the Cohere AI API and returning summarized text as summary -func summarize(ctx context.Context, d *plugin.QueryData, _ *plugin.HydrateData) (interface{}, error) { - // Create the API client - client, err := connect(ctx, d) - if err != nil { - return nil, err - } - - // Build the SummarizeOptions from the request - opts := coherego.SummarizeOptions{ - Text: d.EqualsQuals["text"].GetStringValue(), - Model: "summarize-xlarge", - Format: "paragraph", - } - settingString := d.EqualsQuals["settings"].GetJsonbValue() - if settingString != "" { - var crQual SummarizeRequestQual - err := json.Unmarshal([]byte(settingString), &crQual) - if err != nil { - plugin.Logger(ctx).Error("cohereai_summarize.summarize", "connection_error", err) - return nil, err - } - - if crQual.Length != nil { - opts.Length = *crQual.Length - } - if crQual.Extractiveness != nil { - opts.Extractiveness = *crQual.Extractiveness - } - if crQual.Temperature != nil { - opts.Temperature = crQual.Temperature - } - if crQual.Format != nil { - opts.Format = *crQual.Format - } - if crQual.Model != nil { - opts.Model = *crQual.Model - } - if crQual.AdditionalCommand != nil { - opts.AdditionalCommand = *crQual.AdditionalCommand - } - if crQual.Text != nil { - opts.Text = *crQual.Text - } - } - - // Make the Summarize API call - resp, err := client.Summarize(opts) - if err != nil { - plugin.Logger(ctx).Error("cohereai_summarize.summarize", "api_error", err) - return nil, err - } - - plugin.Logger(ctx).Debug("cohereai_summarize.summarize", "summary", resp.Summary) - // Return summarize data - row := SummarizeRow{*resp, opts.Text} - d.StreamListItem(ctx, row) - return nil, nil -} diff --git a/steampipe-plugin-cohereai/cohereai/table_cohereai_tokenize.go b/steampipe-plugin-cohereai/cohereai/table_cohereai_tokenize.go deleted file mode 100644 index e565bb20..00000000 --- a/steampipe-plugin-cohereai/cohereai/table_cohereai_tokenize.go +++ /dev/null @@ -1,85 +0,0 @@ -package cohere - -import ( - "context" - "encoding/json" - - coherego "github.com/cohere-ai/cohere-go" - "github.com/turbot/steampipe-plugin-sdk/v5/grpc/proto" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin" - "github.com/turbot/steampipe-plugin-sdk/v5/plugin/transform" -) - -// Defines the cohere_tokenize table -func tableCohereTokenize(ctx context.Context) *plugin.Table { - return &plugin.Table{ - Name: "cohereai_tokenize", - Description: "Tokenize in Cohere AI.", - List: &plugin.ListConfig{ - Hydrate: tokenize, - KeyColumns: []*plugin.KeyColumn{ - {Name: "text", Require: plugin.Optional}, - }, - }, - Columns: []*plugin.Column{ - // Columns returned from the Cohere API - {Name: "tokens", Type: proto.ColumnType_STRING, Transform: transform.FromField("Tokens.Tokens"), Description: "Tokens for a given text prompt."}, - {Name: "token_strings", Type: proto.ColumnType_STRING, Transform: transform.FromField("Tokens.TokenStrings"), Description: "Tokens in the form of input string."}, - - // Qual columns to provide input to the API - {Name: "text", Type: proto.ColumnType_STRING, Transform: transform.FromQual("text"), Description: "The text to tokenize for, encoded as a string."}, - }, - } -} - -// CompletionRequestQual defines the structure of the settings qual -type TokenizeRequestQual struct { - Text *string `json:"text,omitempty"` -} - -// SummarizeRow defines the row structure returned from the API -type TokenizeRow struct { - Tokens coherego.TokenizeResponse - Text string -} - -// tokenize handles querying the Cohere AI API and returning tokens from provided text -func tokenize(ctx context.Context, d *plugin.QueryData, _ *plugin.HydrateData) (interface{}, error) { - - conn, err := connect(ctx, d) - if err != nil { - plugin.Logger(ctx).Error("cohereai_tokenize.tokenize", "connection_error", err) - return nil, err - } - - // Default settings taken from the Cohere API docs - // https://docs.cohere.ai/reference/generate - cr := coherego.TokenizeOptions{ - Text: d.EqualsQuals["text"].GetStringValue(), - } - - settingsString := d.EqualsQuals["settings"].GetJsonbValue() - if settingsString != "" { - // Overwrite any settings provided in the settings qual. If a field - // is not passed in the settings, then default to the settings above. - var crQual TokenizeRequestQual - err := json.Unmarshal([]byte(settingsString), &crQual) - if err != nil { - plugin.Logger(ctx).Error("cohereai_tokenize.tokenize", "connection_error", err) - return nil, err - } - } - - // Query the Cohere API - resp, err := conn.Tokenize(cr) - if err != nil { - plugin.Logger(ctx).Error("cohereai_tokenize.tokenize", "api_error", err) - return nil, err - } - - plugin.Logger(ctx).Trace("cohereai_tokenize.tokenize", "response", resp) - // Return tokenize data - row := TokenizeRow{*resp, cr.Text} - d.StreamListItem(ctx, row) - return nil, nil -}