From 4499baeab831a11cfc6fae18f4b927183fc9e86c Mon Sep 17 00:00:00 2001 From: Devin Sevilla Date: Sun, 21 Jul 2024 15:19:56 -0700 Subject: [PATCH] Add support for the LLM CLI utility --- plugins/llm/api_key.go | 242 ++++++++++++++++++++++++++++ plugins/llm/api_key_test.go | 72 +++++++++ plugins/llm/llm.go | 25 +++ plugins/llm/plugin.go | 22 +++ plugins/llm/test-fixtures/keys.json | 4 + 5 files changed, 365 insertions(+) create mode 100644 plugins/llm/api_key.go create mode 100644 plugins/llm/api_key_test.go create mode 100644 plugins/llm/llm.go create mode 100644 plugins/llm/plugin.go create mode 100644 plugins/llm/test-fixtures/keys.json diff --git a/plugins/llm/api_key.go b/plugins/llm/api_key.go new file mode 100644 index 00000000..27f76206 --- /dev/null +++ b/plugins/llm/api_key.go @@ -0,0 +1,242 @@ +package llm + +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" +) + +// LLM enables you to query many APIs with a single command. +// Therefore, we have an optional field for each API key in the 1Password entry. + +const ( + AnthropicFieldName = sdk.FieldName("Anthropic") + AnyscaleFieldName = sdk.FieldName("Anyscale") + CohereFieldName = sdk.FieldName("Cohere") + FireworksFieldName = sdk.FieldName("Fireworks") + GeminiFieldName = sdk.FieldName("Gemini") + GroqFieldName = sdk.FieldName("Groq") + MistralFieldName = sdk.FieldName("Mistral") + OpenAIFieldName = sdk.FieldName("OpenAI") + OpenRouterFieldName = sdk.FieldName("OpenRouter") + PALMFieldName = sdk.FieldName("PaLM") + PerplexityFieldName = sdk.FieldName("Perplexity") + RekaFieldName = sdk.FieldName("Reka") + ReplicateFieldName = sdk.FieldName("Replicate") + TogetherFieldName = sdk.FieldName("Together") +) + +var defaultValueComposition = &schema.ValueComposition{ + Charset: schema.Charset{ + Uppercase: true, + Lowercase: true, + Digits: true, + Symbols: true, + }, +} + +var availableServices = []ServiceDefinition{ + { + FieldName: AnthropicFieldName, + EnvVarName: "ANTHROPIC_API_KEY", + MarkdownDescription: "API Key for Anthropic.", + Composition: &schema.ValueComposition{ + Prefix: "sk-ant-", + Charset: schema.Charset{ + Uppercase: true, + Lowercase: true, + Digits: true, + Symbols: true, + }, + }, + ConfigFieldFunc: func(c Config) string { return c.Claude }, + }, + { + FieldName: AnyscaleFieldName, + EnvVarName: "LLM_ANYSCALE_ENDPOINTS_KEY", + MarkdownDescription: "API Key for Anyscale.", + Composition: defaultValueComposition, + ConfigFieldFunc: func(c Config) string { return c.Anyscale }, + }, + { + FieldName: CohereFieldName, + EnvVarName: "COHERE_API_KEY", + MarkdownDescription: "API Key for Cohere.", + Composition: defaultValueComposition, + ConfigFieldFunc: func(c Config) string { return c.Cohere }, + }, + { + FieldName: FireworksFieldName, + EnvVarName: "LLM_FIREWORKS_KEY", + MarkdownDescription: "API Key for Fireworks.", + Composition: defaultValueComposition, + ConfigFieldFunc: func(c Config) string { return c.Fireworks }, + }, + { + FieldName: GeminiFieldName, + EnvVarName: "LLM_GEMINI_KEY", + MarkdownDescription: "API Key for Google’s Gemini.", + Composition: defaultValueComposition, + ConfigFieldFunc: func(c Config) string { return c.Gemini }, + }, + { + FieldName: GroqFieldName, + EnvVarName: "LLM_GROQ_KEY", + MarkdownDescription: "API Key for Groq.", + Composition: defaultValueComposition, + ConfigFieldFunc: func(c Config) string { return c.Groq }, + }, + { + FieldName: MistralFieldName, + EnvVarName: "LLM_MISTRAL_KEY", + MarkdownDescription: "API Key for Mistral.", + Composition: defaultValueComposition, + ConfigFieldFunc: func(c Config) string { return c.Mistral }, + }, + { + FieldName: OpenAIFieldName, + EnvVarName: "OPENAI_API_KEY", + MarkdownDescription: "API Key for OpenAI.", + Composition: &schema.ValueComposition{ + Prefix: "sk-", + Charset: schema.Charset{ + Uppercase: true, + Lowercase: true, + Digits: true, + }, + }, + ConfigFieldFunc: func(c Config) string { return c.OpenAI }, + }, + { + FieldName: OpenRouterFieldName, + EnvVarName: "LLM_OPENROUTER_KEY", + MarkdownDescription: "API Key for OpenRouter.", + Composition: defaultValueComposition, + ConfigFieldFunc: func(c Config) string { return c.OpenRouter }, + }, + { + FieldName: PALMFieldName, + EnvVarName: "PALM_API_KEY", + MarkdownDescription: "API Key for PALM.", + Composition: defaultValueComposition, + ConfigFieldFunc: func(c Config) string { return c.PALM }, + }, + { + FieldName: PerplexityFieldName, + EnvVarName: "PERPLEXITY_API_KEY", + MarkdownDescription: "API Key for Perplexity.", + Composition: defaultValueComposition, + ConfigFieldFunc: func(c Config) string { return c.Perplexity }, + }, + { + FieldName: RekaFieldName, + EnvVarName: "LLM_REKA_KEY", + MarkdownDescription: "API Key for Reka.", + Composition: defaultValueComposition, + ConfigFieldFunc: func(c Config) string { return c.Reka }, + }, + { + FieldName: ReplicateFieldName, + EnvVarName: "REPLICATE_API_KEY", + MarkdownDescription: "API Key for Replicate.", + Composition: defaultValueComposition, + ConfigFieldFunc: func(c Config) string { return c.Replicate }, + }, + { + FieldName: TogetherFieldName, + EnvVarName: "TOGETHER_API_KEY", + MarkdownDescription: "API Key for Together.", + Composition: defaultValueComposition, + ConfigFieldFunc: func(c Config) string { return c.Together }, + }, +} + +func APIKey() schema.CredentialType { + // Create env variable mapping for each LLM service + var defaultEnvVarMapping = make(map[string]sdk.FieldName) + for _, field := range availableServices { + if field.EnvVarName != "" { + defaultEnvVarMapping[field.EnvVarName] = field.FieldName + } + } + + // Create schema fields for each LLM service + var schemaFields []schema.CredentialField + for _, field := range availableServices { + schemaFields = append(schemaFields, schema.CredentialField{ + Name: field.FieldName, + MarkdownDescription: field.MarkdownDescription, + Secret: true, + Optional: true, + Composition: field.Composition, + }) + } + + return schema.CredentialType{ + Name: credname.APIKey, + DocsURL: sdk.URL("https://llm.datasette.io/en/stable/setup.html"), + Fields: schemaFields, + DefaultProvisioner: provision.EnvVars(defaultEnvVarMapping), + Importer: importer.TryAll( + importer.TryEnvVarPair(defaultEnvVarMapping), + importer.MacOnly(TryLLMConfigFile("~/Library/Application Support/io.datasette.llm/keys.json")), + importer.LinuxOnly(TryLLMConfigFile("~/.config/io.datasette.llm/keys.json")), + )} +} + +func TryLLMConfigFile(path string) sdk.Importer { + return importer.TryFile(path, func(ctx context.Context, contents importer.FileContents, in sdk.ImportInput, out *sdk.ImportAttempt) { + var config Config + if err := contents.ToJSON(&config); err != nil { + out.AddError(err) + return + } + + // Add candidates for each service that has a value in the config file + candidateFields := make(map[sdk.FieldName]string) + for _, field := range availableServices { + var configValue string = field.ConfigFieldFunc(config) + if configValue != "" { + candidateFields[field.FieldName] = configValue + } + } + + if len(candidateFields) == 0 { + return + } + + out.AddCandidate(sdk.ImportCandidate{ + Fields: candidateFields, + }) + }) +} + +type ServiceDefinition struct { + FieldName sdk.FieldName + ConfigFileFieldName string + EnvVarName string + MarkdownDescription string + Composition *schema.ValueComposition + ConfigFieldFunc func(Config) string +} + +type Config struct { + Anyscale string `json:"anyscale-endpoints"` + Claude string `json:"claude"` + Cohere string `json:"cohere"` + Fireworks string `json:"fireworks"` + Gemini string `json:"gemini"` + Groq string `json:"groq"` + Mistral string `json:"mistral"` + OpenAI string `json:"openai"` + OpenRouter string `json:"openrouter"` + PALM string `json:"palm"` + Perplexity string `json:"perplexity"` + Reka string `json:"reka"` + Replicate string `json:"replicate"` + Together string `json:"together"` +} diff --git a/plugins/llm/api_key_test.go b/plugins/llm/api_key_test.go new file mode 100644 index 00000000..b13c4406 --- /dev/null +++ b/plugins/llm/api_key_test.go @@ -0,0 +1,72 @@ +package llm + +import ( + "testing" + + "github.com/1Password/shell-plugins/sdk" + "github.com/1Password/shell-plugins/sdk/plugintest" +) + +func TestAPIKeyProvisioner(t *testing.T) { + plugintest.TestProvisioner(t, APIKey().DefaultProvisioner, map[string]plugintest.ProvisionCase{ + "default": { + ItemFields: map[sdk.FieldName]string{ // TODO: Check if this is correct + OpenAIFieldName: "sk-proj-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + AnthropicFieldName: "sk-ant-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + }, + ExpectedOutput: sdk.ProvisionOutput{ + Environment: map[string]string{ + "OPENAI_API_KEY": "sk-proj-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + "ANTHROPIC_API_KEY": "sk-ant-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + }, + }, + }, + }) +} + +func TestAPIKeyImporter(t *testing.T) { + plugintest.TestImporter(t, APIKey().Importer, map[string]plugintest.ImportCase{ + "environment": { + Environment: map[string]string{ + "OPENAI_API_KEY": "sk-proj-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + "ANTHROPIC_API_KEY": "sk-ant-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + }, + ExpectedCandidates: []sdk.ImportCandidate{ + { + Fields: map[sdk.FieldName]string{ + OpenAIFieldName: "sk-proj-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + AnthropicFieldName: "sk-ant-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + }, + }, + }, + }, + "config file mac": { + OS: "darwin", + Files: map[string]string{ + "~/Library/Application Support/io.datasette.llm/keys.json": plugintest.LoadFixture(t, "keys.json"), + }, + ExpectedCandidates: []sdk.ImportCandidate{ + { + Fields: map[sdk.FieldName]string{ + OpenAIFieldName: "sk-proj-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + AnthropicFieldName: "sk-ant-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + }, + }, + }, + }, + "config file linux": { + OS: "linux", + Files: map[string]string{ + "~/.config/io.datasette.llm/keys.json": plugintest.LoadFixture(t, "keys.json"), + }, + ExpectedCandidates: []sdk.ImportCandidate{ + { + Fields: map[sdk.FieldName]string{ + OpenAIFieldName: "sk-proj-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + AnthropicFieldName: "sk-ant-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + }, + }, + }, + }, + }) +} diff --git a/plugins/llm/llm.go b/plugins/llm/llm.go new file mode 100644 index 00000000..1a7a9d6d --- /dev/null +++ b/plugins/llm/llm.go @@ -0,0 +1,25 @@ +package llm + +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 LLMCLI() schema.Executable { + return schema.Executable{ + Name: "LLM", + Runs: []string{"llm"}, + DocsURL: sdk.URL("https://llm.datasette.io/"), + NeedsAuth: needsauth.IfAll( + needsauth.NotForHelpOrVersion(), + needsauth.NotWithoutArgs(), + ), + Uses: []schema.CredentialUsage{ + { + Name: credname.APIKey, + }, + }, + } +} diff --git a/plugins/llm/plugin.go b/plugins/llm/plugin.go new file mode 100644 index 00000000..cf1d9a6f --- /dev/null +++ b/plugins/llm/plugin.go @@ -0,0 +1,22 @@ +package llm + +import ( + "github.com/1Password/shell-plugins/sdk" + "github.com/1Password/shell-plugins/sdk/schema" +) + +func New() schema.Plugin { + return schema.Plugin{ + Name: "llm", + Platform: schema.PlatformInfo{ + Name: "LLM", + Homepage: sdk.URL("https://llm.datasette.io/"), + }, + Credentials: []schema.CredentialType{ + APIKey(), + }, + Executables: []schema.Executable{ + LLMCLI(), + }, + } +} diff --git a/plugins/llm/test-fixtures/keys.json b/plugins/llm/test-fixtures/keys.json new file mode 100644 index 00000000..f7d09237 --- /dev/null +++ b/plugins/llm/test-fixtures/keys.json @@ -0,0 +1,4 @@ +{ + "openai": "sk-proj-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE", + "claude": "sk-ant-ysT1SpYOenNu805nCf3yUYIbNAfvHSNzR0rx2WGRHEXAMPLE" +} \ No newline at end of file