Merge pull request #393 from 1Password/amanda/huggingface

Amanda/huggingface
1Password · Oct 23, 2023 · 33b9329 · 33b9329
2 parents 6f8040b + 234fb2b
commit 33b9329
Show file tree

Hide file tree

Showing 6 changed files with 187 additions and 0 deletions.
diff --git a/plugins/huggingface/huggingface-cli.go b/plugins/huggingface/huggingface-cli.go
@@ -0,0 +1,27 @@
+package huggingface
+
+import (
+	"github.com/1Password/shell-plugins/sdk"
+	"github.com/1Password/shell-plugins/sdk/needsauth"
+	"github.com/1Password/shell-plugins/sdk/schema"
+	"github.com/1Password/shell-plugins/sdk/schema/credname"
+)
+
+func HuggingFaceCLI() schema.Executable {
+	return schema.Executable{
+		Name:    "HuggingFace CLI",
+		Runs:    []string{"huggingface-cli"},
+		DocsURL: sdk.URL("https://huggingface.co/docs/huggingface_hub/quick-start"),
+		NeedsAuth: needsauth.IfAll(
+			needsauth.NotForHelpOrVersion(),
+			needsauth.NotWithoutArgs(),
+			needsauth.NotWhenContainsArgs("login"),
+			needsauth.NotWhenContainsArgs("logout"),
+		),
+		Uses: []schema.CredentialUsage{
+			{
+				Name: credname.APIToken,
+			},
+		},
+	}
+}
diff --git a/plugins/huggingface/plugin.go b/plugins/huggingface/plugin.go
@@ -0,0 +1,22 @@
+package huggingface
+
+import (
+	"github.com/1Password/shell-plugins/sdk"
+	"github.com/1Password/shell-plugins/sdk/schema"
+)
+
+func New() schema.Plugin {
+	return schema.Plugin{
+		Name: "huggingface",
+		Platform: schema.PlatformInfo{
+			Name:     "Hugging Face",
+			Homepage: sdk.URL("https://huggingface.co"),
+		},
+		Credentials: []schema.CredentialType{
+			UserAccessToken(),
+		},
+		Executables: []schema.Executable{
+			HuggingFaceCLI(),
+		},
+	}
+}
diff --git a/plugins/huggingface/test-fixtures/token b/plugins/huggingface/test-fixtures/token
@@ -0,0 +1 @@
+hf_yVvZeburdKtnwkVCWPXimmNwaFuEXAMPLE
diff --git a/plugins/huggingface/user_access_token.go b/plugins/huggingface/user_access_token.go
@@ -0,0 +1,87 @@
+package huggingface
+
+import (
+	"context"
+
+	"github.com/1Password/shell-plugins/sdk"
+	"github.com/1Password/shell-plugins/sdk/importer"
+	"github.com/1Password/shell-plugins/sdk/provision"
+	"github.com/1Password/shell-plugins/sdk/schema"
+	"github.com/1Password/shell-plugins/sdk/schema/credname"
+	"github.com/1Password/shell-plugins/sdk/schema/fieldname"
+)
+
+func UserAccessToken() schema.CredentialType {
+	return schema.CredentialType{
+		Name:          credname.APIToken,
+		DocsURL:       sdk.URL("https://huggingface.co/docs/hub/security-tokens"),
+		ManagementURL: sdk.URL("https://huggingface.co/settings/tokens"),
+		Fields: []schema.CredentialField{
+			{
+				Name:                fieldname.UserAccessToken,
+				MarkdownDescription: "Token used to authenticate to HuggingFace.",
+				Secret:              true,
+				Composition: &schema.ValueComposition{
+					Length: 37,
+					Prefix: "hf_",
+					Charset: schema.Charset{
+						Uppercase: true,
+						Lowercase: true,
+						Symbols:   true,
+					},
+				},
+			},
+			{
+				Name:                fieldname.Endpoint,
+				MarkdownDescription: "Endpoint used to connect to HuggingFace CLI",
+				Optional:            true,
+				Secret:              false,
+				Composition: &schema.ValueComposition{
+					Charset: schema.Charset{
+						Uppercase: true,
+						Lowercase: true,
+						Digits:    true,
+						Symbols:   true,
+					},
+				},
+			},
+			{
+				Name:                fieldname.APIUrl,
+				MarkdownDescription: "HF Inference Endpoint used to connect to HuggingFace CLI",
+				Optional:            true,
+				Secret:              false,
+				Composition: &schema.ValueComposition{
+					Charset: schema.Charset{
+						Uppercase: true,
+						Lowercase: true,
+						Digits:    true,
+						Symbols:   true,
+					},
+				},
+			},
+		},
+		DefaultProvisioner: provision.EnvVars(defaultEnvVarMapping),
+		Importer: importer.TryAll(
+			importer.TryEnvVarPair(defaultEnvVarMapping),
+			TryHuggingFaceTokenFile(),
+		)}
+}
+
+var defaultEnvVarMapping = map[string]sdk.FieldName{
+	"HUGGING_FACE_HUB_TOKEN": fieldname.UserAccessToken,
+	"HF_ENDPOINT":            fieldname.Endpoint,
+	"HF_INFERENCE_ENDPOINT":  fieldname.APIUrl,
+}
+
+func TryHuggingFaceTokenFile() sdk.Importer {
+	return importer.TryFile("~/.cache/huggingface/token", func(ctx context.Context, contents importer.FileContents, in sdk.ImportInput, out *sdk.ImportAttempt) {
+		fileData := string(contents)
+
+		out.AddCandidate(sdk.ImportCandidate{
+			Fields: map[sdk.FieldName]string{
+				fieldname.UserAccessToken: fileData,
+			},
+		})
+	})
+
+}
diff --git a/plugins/huggingface/user_access_token_test.go b/plugins/huggingface/user_access_token_test.go
@@ -0,0 +1,48 @@
+package huggingface
+
+import (
+	"testing"
+
+	"github.com/1Password/shell-plugins/sdk"
+	"github.com/1Password/shell-plugins/sdk/plugintest"
+	"github.com/1Password/shell-plugins/sdk/schema/fieldname"
+)
+
+func TestAPITokenProvisioner(t *testing.T) {
+	plugintest.TestProvisioner(
+		t, UserAccessToken().DefaultProvisioner, map[string]plugintest.ProvisionCase{
+			"default": {
+				ItemFields: map[sdk.FieldName]string{
+					fieldname.UserAccessToken: "hf_yVvZeburdKtnwkVCWPXimmNwaFuEXAMPLE",
+					fieldname.Endpoint:        "https://huggingface.co",
+					fieldname.APIUrl:          "https://api-inference.huggingface.com",
+				},
+				ExpectedOutput: sdk.ProvisionOutput{
+					Environment: map[string]string{
+						"HUGGING_FACE_HUB_TOKEN": "hf_yVvZeburdKtnwkVCWPXimmNwaFuEXAMPLE",
+						"HF_ENDPOINT":            "https://huggingface.co",
+						"HF_INFERENCE_ENDPOINT":  "https://api-inference.huggingface.com",
+					},
+				},
+			},
+		})
+}
+
+func TestAPITokenImporter(t *testing.T) {
+	plugintest.TestImporter(
+		t, UserAccessToken().Importer, map[string]plugintest.ImportCase{
+			"config file (macOS)": {
+				Files: map[string]string{
+					"~/.cache/huggingface/token": plugintest.LoadFixture(t, "token"),
+				},
+				ExpectedCandidates: []sdk.ImportCandidate{
+					{
+						Fields: map[sdk.FieldName]string{
+							fieldname.UserAccessToken: "hf_yVvZeburdKtnwkVCWPXimmNwaFuEXAMPLE\n",
+						},
+					},
+				},
+			},
+		},
+	)
+}
diff --git a/sdk/schema/fieldname/names.go b/sdk/schema/fieldname/names.go
@@ -5,6 +5,7 @@ import "github.com/1Password/shell-plugins/sdk"
 // Credential field names.
 const (
 	APIHost         = sdk.FieldName("API Host")
+	APIUrl          = sdk.FieldName("API URL")
 	APIKey          = sdk.FieldName("API Key")
 	APIKeyID        = sdk.FieldName("API Key ID")
 	APISecret       = sdk.FieldName("API Secret")
@@ -53,6 +54,7 @@ const (
 	Token           = sdk.FieldName("Token")
 	URL             = sdk.FieldName("URL")
 	User            = sdk.FieldName("User")
+	UserAccessToken = sdk.FieldName("User Access Token")
 	Username        = sdk.FieldName("Username")
 	Website         = sdk.FieldName("Website")
 )