From 80c0cdba8f24d4529ab09a82326737e61d649c43 Mon Sep 17 00:00:00 2001 From: Amir Deris Date: Sun, 24 Aug 2025 19:35:20 +0200 Subject: [PATCH 1/6] Refactored types into a different file --- gromit.go | 26 -------------------------- types.go | 29 +++++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 26 deletions(-) create mode 100644 types.go diff --git a/gromit.go b/gromit.go index 49f93be..2d9e094 100644 --- a/gromit.go +++ b/gromit.go @@ -27,13 +27,6 @@ type Gromit struct { *configuration } -type systemInfo struct { - operatingSystem string - currentShell string - delimiter string - kernelInfo string -} - func getSystemInfo() systemInfo { o := runtime.GOOS var eol, shell, kernelInfo string @@ -57,25 +50,10 @@ func getSystemInfo() systemInfo { } } -type messagePrinter struct { - w io.Writer - promptPrefix string - delimiter string -} - -type configuration struct { - promptPrefix string - w io.Writer - askForConfirmation bool - systemInfo -} - func (m *messagePrinter) print(s string) { fmt.Fprintf(m.w, "%s %s %s", m.promptPrefix, s, m.delimiter) } -type ConfigurationModifier func(*configuration) error - func WithPromptPrefix(prefix string) ConfigurationModifier { return func(c *configuration) error { c.promptPrefix = prefix @@ -175,10 +153,6 @@ func addEnvironmentInfo(systemInfo systemInfo, systemPrompt string) string { return result } -type userConfirmation struct { - confirmed bool -} - func (g *Gromit) askConfirmation(message string) (userConfirmation, error) { if !g.configuration.askForConfirmation { return userConfirmation{ diff --git a/types.go b/types.go new file mode 100644 index 0000000..e35ea46 --- /dev/null +++ b/types.go @@ -0,0 +1,29 @@ +package main + +import "io" + +type systemInfo struct { + operatingSystem string + currentShell string + delimiter string + kernelInfo string +} + +type messagePrinter struct { + w io.Writer + promptPrefix string + delimiter string +} + +type configuration struct { + promptPrefix string + w io.Writer + askForConfirmation bool + systemInfo +} + +type userConfirmation struct { + confirmed bool +} + +type ConfigurationModifier func(*configuration) error From 6cb728ae49128109c22ac255835c6c2da57da416 Mon Sep 17 00:00:00 2001 From: Amir Deris Date: Mon, 25 Aug 2025 13:45:57 +0200 Subject: [PATCH 2/6] refactored API parameters and passed extra options --- agent.go | 88 +++++++++++++++++++++++++++++++++---------------------- gromit.go | 30 ++++++++++++++----- types.go | 9 ++++++ 3 files changed, 85 insertions(+), 42 deletions(-) diff --git a/agent.go b/agent.go index 5a8c801..230ad56 100644 --- a/agent.go +++ b/agent.go @@ -4,8 +4,11 @@ import ( "context" "fmt" "github.com/anthropics/anthropic-sdk-go" + anthropicOption "github.com/anthropics/anthropic-sdk-go/option" "github.com/openai/openai-go" + openaiOption "github.com/openai/openai-go/option" "google.golang.org/genai" + "os" ) const ( @@ -20,8 +23,12 @@ const ( geminiFlash = "gemini-2.5-flash" ) +const ( + defaultMaxTokens = 1024 +) + type Assister interface { - GetTerminalCommand(ctx context.Context, userMessage string, systemMessage string) (string, error) + GetTerminalCommand(ctx context.Context, userMessage string) (string, error) } var _ Assister = (*OpenAIAssister)(nil) @@ -29,17 +36,21 @@ var _ Assister = (*AnthropicAIAssister)(nil) var _ Assister = (*GeminiAIAssister)(nil) type OpenAIAssister struct { - model string + aiParameters } -func (o *OpenAIAssister) GetTerminalCommand(ctx context.Context, userMessage string, systemMessage string) (string, error) { - client := openai.NewClient() +func (o *OpenAIAssister) GetTerminalCommand(ctx context.Context, userMessage string) (string, error) { + apiKey := o.aiParameters.apiKey + if apiKey == "" { + apiKey, _ = os.LookupEnv("OPENAI_API_KEY") + } + client := openai.NewClient(openaiOption.WithAPIKey(apiKey)) chatCompletion, err := client.Chat.Completions.New(ctx, openai.ChatCompletionNewParams{ Messages: []openai.ChatCompletionMessageParamUnion{ openai.UserMessage(userMessage), - openai.SystemMessage(systemMessage), + openai.SystemMessage(o.aiParameters.systemPrompt), }, - Model: o.model, + Model: o.aiParameters.model, }) if err != nil { return "", err @@ -48,15 +59,23 @@ func (o *OpenAIAssister) GetTerminalCommand(ctx context.Context, userMessage str } type AnthropicAIAssister struct { - model string + aiParameters } -func (c *AnthropicAIAssister) GetTerminalCommand(ctx context.Context, userMessage string, systemMessage string) (string, error) { - client := anthropic.NewClient() //defaults to os.LookupEnv("ANTHROPIC_API_KEY") +func (c *AnthropicAIAssister) GetTerminalCommand(ctx context.Context, userMessage string) (string, error) { + apiKey := c.aiParameters.apiKey + if apiKey == "" { + apiKey, _ = os.LookupEnv("ANTHROPIC_API_KEY") + } + maxTokens := c.aiParameters.maxTokens + if maxTokens == 0 { + maxTokens = defaultMaxTokens + } + client := anthropic.NewClient(anthropicOption.WithAPIKey(apiKey)) message, err := client.Messages.New(ctx, anthropic.MessageNewParams{ - MaxTokens: 1024, + MaxTokens: maxTokens, System: []anthropic.TextBlockParam{ - {Text: systemMessage}, + {Text: c.aiParameters.systemPrompt}, }, Messages: []anthropic.MessageParam{ anthropic.NewUserMessage(anthropic.NewTextBlock(userMessage)), @@ -77,12 +96,17 @@ func (c *AnthropicAIAssister) GetTerminalCommand(ctx context.Context, userMessag } type GeminiAIAssister struct { - model string + aiParameters } -func (g *GeminiAIAssister) GetTerminalCommand(ctx context.Context, userMessage string, systemMessage string) (string, error) { +func (g *GeminiAIAssister) GetTerminalCommand(ctx context.Context, userMessage string) (string, error) { + apiKey := g.aiParameters.apiKey + if apiKey == "" { + apiKey, _ = os.LookupEnv("GEMINI_API_KEY") + } client, err := genai.NewClient(ctx, &genai.ClientConfig{ Backend: genai.BackendGeminiAPI, + APIKey: apiKey, }) if err != nil { return "", err @@ -90,7 +114,7 @@ func (g *GeminiAIAssister) GetTerminalCommand(ctx context.Context, userMessage s config := &genai.GenerateContentConfig{ SystemInstruction: &genai.Content{ Parts: []*genai.Part{ - {Text: systemMessage}, + {Text: g.aiParameters.systemPrompt}, }, }, } @@ -106,38 +130,32 @@ func (g *GeminiAIAssister) GetTerminalCommand(ctx context.Context, userMessage s } type AssisterCreator interface { - GetAssister(agent string, model string) (Assister, error) + GetAssister(parameters aiParameters) (Assister, error) } var _ AssisterCreator = (*defaultAIAssisterCreator)(nil) type defaultAIAssisterCreator struct{} -func (d *defaultAIAssisterCreator) GetAssister(agent, model string) (Assister, error) { +func (d *defaultAIAssisterCreator) GetAssister(p aiParameters) (Assister, error) { switch { - case agent == "" || agent == openAIAgent: - if model == "" { - model = openai.ChatModelGPT4o + case p.agent == "" || p.agent == openAIAgent: + if p.model == "" { + p.model = openai.ChatModelGPT4o } - return &OpenAIAssister{ - model: model, - }, nil + return &OpenAIAssister{p}, nil - case agent == anthropicAIAgent: - if model == "" { - model = string(anthropic.ModelClaude3_5HaikuLatest) + case p.agent == anthropicAIAgent: + if p.model == "" { + p.model = string(anthropic.ModelClaude3_5HaikuLatest) } - return &AnthropicAIAssister{ - model: model, - }, nil - case agent == geminiAIAgent: - if model == "" { - model = geminiFlashLite + return &AnthropicAIAssister{p}, nil + case p.agent == geminiAIAgent: + if p.model == "" { + p.model = geminiFlashLite } - return &GeminiAIAssister{ - model, - }, nil + return &GeminiAIAssister{p}, nil default: - return nil, fmt.Errorf("cannot create AI agent for %s and model %s", agent, model) + return nil, fmt.Errorf("cannot create AI agent for %s and model %s", p.agent, p.model) } } diff --git a/gromit.go b/gromit.go index 2d9e094..c08c4da 100644 --- a/gromit.go +++ b/gromit.go @@ -82,6 +82,18 @@ func (g *Gromit) actionGromit(ctx context.Context, command *cli.Command) error { g.print("Please run ./gromit --help to see usage") return nil } + prompt := g.String("systemPrompt") + if prompt == "" { + prompt = systemPrompt + } + prompt = addEnvironmentInfo(g.configuration.systemInfo, prompt) + g.configuration.aiParameters = aiParameters{ + maxTokens: g.Int64("maxToken"), + apiKey: g.String("apiKey"), + agent: g.String("agent"), + model: g.String("model"), + systemPrompt: prompt, + } err := g.handleUserQuery(ctx, query) if err != nil { return err @@ -109,16 +121,11 @@ func (g *Gromit) actionGromit(ctx context.Context, command *cli.Command) error { } func (g *Gromit) handleUserQuery(ctx context.Context, query string) error { - assister, err := g.AssisterCreator.GetAssister(g.String("agent"), g.String("model")) + assister, err := g.AssisterCreator.GetAssister(g.configuration.aiParameters) if err != nil { return err } - prompt := g.String("systemPrompt") - if prompt == "" { - prompt = systemPrompt - } - prompt = addEnvironmentInfo(g.configuration.systemInfo, prompt) - exeCommand, err := assister.GetTerminalCommand(ctx, query, prompt) + exeCommand, err := assister.GetTerminalCommand(ctx, query) if err != nil { return err } @@ -231,12 +238,21 @@ func NewGromit(a AssisterCreator, mods ...ConfigurationModifier) (*Gromit, error Name: "systemPrompt", Usage: "The system prompt for the AI agent. Defaults to command line helper in a linux environment.", }, + &cli.StringFlag{ + Name: "apiKey", + Usage: "The API key to use for given AI agent. By default it is read from environment variables.", + }, + &cli.Int32Flag{ + Name: "maxTokens", + Usage: "Maximum number of tokens for AI agents to generate", + }, } config := configuration{ promptPrefix: "⚡️🐶", w: os.Stdout, askForConfirmation: true, systemInfo: getSystemInfo(), + aiParameters: aiParameters{}, } gromit := Gromit{ AssisterCreator: a, diff --git a/types.go b/types.go index e35ea46..5449bbb 100644 --- a/types.go +++ b/types.go @@ -16,6 +16,7 @@ type messagePrinter struct { } type configuration struct { + aiParameters promptPrefix string w io.Writer askForConfirmation bool @@ -27,3 +28,11 @@ type userConfirmation struct { } type ConfigurationModifier func(*configuration) error + +type aiParameters struct { + systemPrompt string + agent string + model string + apiKey string + maxTokens int64 +} From 4f4872bdd5733ecf97afd5979fb00a043ca82cec Mon Sep 17 00:00:00 2001 From: Amir Deris Date: Tue, 26 Aug 2025 11:33:08 +0200 Subject: [PATCH 3/6] Fixed gromit tests --- gromit_test.go | 31 ++++++++++++++++--------------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/gromit_test.go b/gromit_test.go index 2470d52..d278ce0 100644 --- a/gromit_test.go +++ b/gromit_test.go @@ -80,18 +80,22 @@ func TestAIAssisterFindingCorrectCommand(t *testing.T) { g, err := NewGromit(m, WithWriter(&buff), WithPromptPrefix("🐶"), WithAskForConfirmation(false)) require.NoError(t, err) - g.Run(t.Context(), []string{"gromit", "--model", "myModel", "--agent", "myAgent", "--systemPrompt", "myPrompt", "I", "want", "to", "list", "all", "files", "in", "current", "directory"}) + g.Run(t.Context(), []string{"gromit", "--model", "myModel", "--agent", "myAgent", + "--apiKey=key1234", "--maxTokens=2000", + "--systemPrompt", "myPrompt", "I", "want", "to", "list", "all", "files", "in", "current", "directory"}) result := buff.String() require.Contains(t, result, "🐶 In order to do that, you need to run") require.Contains(t, result, "🐶 ls") require.Contains(t, result, "README.md") require.Contains(t, result, "🐶 How can I help?") - require.Equal(t, "myAgent", m.actualAgent) - require.Equal(t, "myModel", m.actualModel) - require.Contains(t, m.actualSystemMessage, "myPrompt") - require.Contains(t, m.actualSystemMessage, "User's operating system is") - require.Contains(t, m.actualSystemMessage, "User's current shell is") + require.Equal(t, "myAgent", m.actualAiParameters.agent) + require.Equal(t, "myModel", m.actualAiParameters.model) + require.Equal(t, "key1234", m.actualAiParameters.apiKey) + require.Equal(t, 2000, m.actualAiParameters.maxTokens) + require.Contains(t, m.actualAiParameters.systemPrompt, "myPrompt") + require.Contains(t, m.actualAiParameters.systemPrompt, "User's operating system is") + require.Contains(t, m.actualAiParameters.systemPrompt, "User's current shell is") require.Equal(t, "I want to list all files in current directory", m.actualUserMessage) } @@ -100,23 +104,20 @@ type mockAIProvider struct { commandError error commandResult string - actualAgent string - actualModel string - actualSystemMessage string - actualUserMessage string + actualUserMessage string + + actualAiParameters aiParameters } -func (m *mockAIProvider) GetAssister(agent string, model string) (Assister, error) { - m.actualAgent = agent - m.actualModel = model +func (m *mockAIProvider) GetAssister(p aiParameters) (Assister, error) { + m.actualAiParameters = p if m.assisterError != nil { return nil, m.assisterError } return m, nil } -func (m *mockAIProvider) GetTerminalCommand(ctx context.Context, userMessage string, systemMessage string) (string, error) { - m.actualSystemMessage = systemMessage +func (m *mockAIProvider) GetTerminalCommand(ctx context.Context, userMessage string) (string, error) { m.actualUserMessage = userMessage if m.commandError != nil { return "", m.commandError From 44c7f08631940951b30f9de62a981743f5a70054 Mon Sep 17 00:00:00 2001 From: Amir Deris Date: Tue, 26 Aug 2025 11:55:51 +0200 Subject: [PATCH 4/6] Updated agent tests --- agent_test.go | 95 +++++++++++++++++++++++++++++++++++--------------- gromit.go | 4 +-- gromit_test.go | 2 +- 3 files changed, 70 insertions(+), 31 deletions(-) diff --git a/agent_test.go b/agent_test.go index 60ce144..c44143d 100644 --- a/agent_test.go +++ b/agent_test.go @@ -10,10 +10,10 @@ import ( ) type assisterTestCase struct { - name string - inputAgent, inputModel string - expectedModel string - err string + name string + inputAiParameters aiParameters + expectedModel string + err string } func TestGetOpenAIAssister(t *testing.T) { @@ -23,28 +23,38 @@ func TestGetOpenAIAssister(t *testing.T) { expectedModel: openai.ChatModelGPT4o, }, { - name: "OpenAI agent with given model should create correct assister", - inputAgent: openAIAgent, - inputModel: "gpt-5o-mini", + name: "OpenAI agent with given model should create correct assister", + inputAiParameters: aiParameters{ + agent: openAIAgent, + model: "gpt-5o-mini", + apiKey: "key1234", + maxTokens: defaultMaxTokens, + }, expectedModel: "gpt-5o-mini", }, { - name: "Unknown agent should result in error", - inputAgent: "Unknown agent", - inputModel: "unknown model", - err: "cannot create AI agent for Unknown agent and model unknown model", + name: "Unknown agent should result in error", + inputAiParameters: aiParameters{ + agent: "Unknown agent", + model: "unknown model", + apiKey: "key1234", + maxTokens: defaultMaxTokens, + }, + err: "cannot create AI agent for Unknown agent and model unknown model", }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { var creator defaultAIAssisterCreator - assister, err := creator.GetAssister(test.inputAgent, test.inputModel) + assister, err := creator.GetAssister(test.inputAiParameters) if test.err != "" { require.EqualError(t, err, test.err) } else { require.NoError(t, err) if openAIAssister, ok := assister.(*OpenAIAssister); ok { require.Equal(t, test.expectedModel, openAIAssister.model) + require.Equal(t, test.inputAiParameters.apiKey, openAIAssister.apiKey) + require.Equal(t, test.inputAiParameters.maxTokens, openAIAssister.maxTokens) } else { require.Fail(t, "Expected OpenAIAssister") } @@ -56,27 +66,37 @@ func TestGetOpenAIAssister(t *testing.T) { func TestGetAnthropicAssister(t *testing.T) { tests := []assisterTestCase{ { - name: "Anthropic agent with no model should default to Claude 3.5 Haiku latest", - inputAgent: anthropicAIAgent, + name: "Anthropic agent with no model should default to Claude 3.5 Haiku latest", + inputAiParameters: aiParameters{ + agent: anthropicAIAgent, + apiKey: "key1234", + maxTokens: defaultMaxTokens, + }, expectedModel: string(anthropic.ModelClaude3_5HaikuLatest), }, { - name: "Anthropic agent with given model should create correct assister", - inputAgent: anthropicAIAgent, - inputModel: string(anthropic.ModelClaudeOpus4_0), + name: "Anthropic agent with given model should create correct assister", + inputAiParameters: aiParameters{ + agent: anthropicAIAgent, + model: string(anthropic.ModelClaudeOpus4_0), + apiKey: "key1234", + maxTokens: defaultMaxTokens, + }, expectedModel: string(anthropic.ModelClaudeOpus4_0), }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { var creator defaultAIAssisterCreator - assister, err := creator.GetAssister(test.inputAgent, test.inputModel) + assister, err := creator.GetAssister(test.inputAiParameters) if test.err != "" { require.EqualError(t, err, test.err) } else { require.NoError(t, err) if anthropicAssister, ok := assister.(*AnthropicAIAssister); ok { require.Equal(t, test.expectedModel, anthropicAssister.model) + require.Equal(t, test.inputAiParameters.apiKey, anthropicAssister.apiKey) + require.Equal(t, test.inputAiParameters.maxTokens, anthropicAssister.maxTokens) } else { require.Fail(t, "Expected Anthropic AI assister") } @@ -88,27 +108,37 @@ func TestGetAnthropicAssister(t *testing.T) { func TestGetGeminiAssister(t *testing.T) { tests := []assisterTestCase{ { - name: "Gemini agent with no model should default to gemini-2.5-flash-lite", - inputAgent: geminiAIAgent, + name: "Gemini agent with no model should default to gemini-2.5-flash-lite", + inputAiParameters: aiParameters{ + agent: geminiAIAgent, + apiKey: "key1234", + maxTokens: defaultMaxTokens, + }, expectedModel: geminiFlashLite, }, { - name: "Gemini agent with given model should create correct assister", - inputAgent: geminiAIAgent, - inputModel: "gemini-2.5-flash-preview-tts", + name: "Gemini agent with given model should create correct assister", + inputAiParameters: aiParameters{ + agent: geminiAIAgent, + model: "gemini-2.5-flash-preview-tts", + apiKey: "key1234", + maxTokens: defaultMaxTokens, + }, expectedModel: "gemini-2.5-flash-preview-tts", }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { var creator defaultAIAssisterCreator - assister, err := creator.GetAssister(test.inputAgent, test.inputModel) + assister, err := creator.GetAssister(test.inputAiParameters) if test.err != "" { require.EqualError(t, err, test.err) } else { require.NoError(t, err) if geminiAssister, ok := assister.(*GeminiAIAssister); ok { require.Equal(t, test.expectedModel, geminiAssister.model) + require.Equal(t, test.inputAiParameters.apiKey, geminiAssister.apiKey) + require.Equal(t, test.inputAiParameters.maxTokens, geminiAssister.maxTokens) } else { require.Fail(t, "Expected Gemini AI assister") } @@ -145,18 +175,27 @@ func TestGetTerminalCommand(t *testing.T) { switch test.agent { case openAIAgent: assister = &OpenAIAssister{ - model: openai.ChatModelGPT4o, + aiParameters{ + model: openai.ChatModelGPT4o, + systemPrompt: systemPrompt, + }, } case anthropicAIAgent: assister = &AnthropicAIAssister{ - model: string(anthropic.ModelClaude3_5HaikuLatest), + aiParameters{ + model: string(anthropic.ModelClaude3_5HaikuLatest), + systemPrompt: systemPrompt, + }, } case geminiAIAgent: assister = &GeminiAIAssister{ - model: geminiFlash, + aiParameters{ + model: geminiFlash, + systemPrompt: systemPrompt, + }, } } - command, err := assister.GetTerminalCommand(t.Context(), "I want to list all files in current directory", systemPrompt) + command, err := assister.GetTerminalCommand(t.Context(), "I want to list all files in current directory") require.NoError(t, err) require.Contains(t, command, "ls") }) diff --git a/gromit.go b/gromit.go index c08c4da..edc8782 100644 --- a/gromit.go +++ b/gromit.go @@ -88,7 +88,7 @@ func (g *Gromit) actionGromit(ctx context.Context, command *cli.Command) error { } prompt = addEnvironmentInfo(g.configuration.systemInfo, prompt) g.configuration.aiParameters = aiParameters{ - maxTokens: g.Int64("maxToken"), + maxTokens: g.Int64("maxTokens"), apiKey: g.String("apiKey"), agent: g.String("agent"), model: g.String("model"), @@ -242,7 +242,7 @@ func NewGromit(a AssisterCreator, mods ...ConfigurationModifier) (*Gromit, error Name: "apiKey", Usage: "The API key to use for given AI agent. By default it is read from environment variables.", }, - &cli.Int32Flag{ + &cli.Int64Flag{ Name: "maxTokens", Usage: "Maximum number of tokens for AI agents to generate", }, diff --git a/gromit_test.go b/gromit_test.go index d278ce0..28288ce 100644 --- a/gromit_test.go +++ b/gromit_test.go @@ -92,7 +92,7 @@ func TestAIAssisterFindingCorrectCommand(t *testing.T) { require.Equal(t, "myAgent", m.actualAiParameters.agent) require.Equal(t, "myModel", m.actualAiParameters.model) require.Equal(t, "key1234", m.actualAiParameters.apiKey) - require.Equal(t, 2000, m.actualAiParameters.maxTokens) + require.Equal(t, int64(2000), m.actualAiParameters.maxTokens) require.Contains(t, m.actualAiParameters.systemPrompt, "myPrompt") require.Contains(t, m.actualAiParameters.systemPrompt, "User's operating system is") require.Contains(t, m.actualAiParameters.systemPrompt, "User's current shell is") From f8531f8d3611e59547563419dc993c05a1cabd12 Mon Sep 17 00:00:00 2001 From: Amir Deris Date: Fri, 29 Aug 2025 13:42:17 +0200 Subject: [PATCH 5/6] Removed API key env variable lookup --- agent.go | 27 ++++++++++++--------------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/agent.go b/agent.go index 230ad56..8eaad8e 100644 --- a/agent.go +++ b/agent.go @@ -8,13 +8,13 @@ import ( "github.com/openai/openai-go" openaiOption "github.com/openai/openai-go/option" "google.golang.org/genai" - "os" ) const ( openAIAgent = "openai" anthropicAIAgent = "anthropic" geminiAIAgent = "gemini" + defaultMaxTokens = 1024 ) // Gemini models @@ -23,10 +23,6 @@ const ( geminiFlash = "gemini-2.5-flash" ) -const ( - defaultMaxTokens = 1024 -) - type Assister interface { GetTerminalCommand(ctx context.Context, userMessage string) (string, error) } @@ -41,10 +37,12 @@ type OpenAIAssister struct { func (o *OpenAIAssister) GetTerminalCommand(ctx context.Context, userMessage string) (string, error) { apiKey := o.aiParameters.apiKey - if apiKey == "" { - apiKey, _ = os.LookupEnv("OPENAI_API_KEY") + var client openai.Client + if apiKey != "" { + client = openai.NewClient(openaiOption.WithAPIKey(apiKey)) + } else { + client = openai.NewClient() } - client := openai.NewClient(openaiOption.WithAPIKey(apiKey)) chatCompletion, err := client.Chat.Completions.New(ctx, openai.ChatCompletionNewParams{ Messages: []openai.ChatCompletionMessageParamUnion{ openai.UserMessage(userMessage), @@ -64,14 +62,16 @@ type AnthropicAIAssister struct { func (c *AnthropicAIAssister) GetTerminalCommand(ctx context.Context, userMessage string) (string, error) { apiKey := c.aiParameters.apiKey - if apiKey == "" { - apiKey, _ = os.LookupEnv("ANTHROPIC_API_KEY") - } maxTokens := c.aiParameters.maxTokens if maxTokens == 0 { maxTokens = defaultMaxTokens } - client := anthropic.NewClient(anthropicOption.WithAPIKey(apiKey)) + var client anthropic.Client + if apiKey != "" { + client = anthropic.NewClient(anthropicOption.WithAPIKey(apiKey)) + } else { + client = anthropic.NewClient() + } message, err := client.Messages.New(ctx, anthropic.MessageNewParams{ MaxTokens: maxTokens, System: []anthropic.TextBlockParam{ @@ -101,9 +101,6 @@ type GeminiAIAssister struct { func (g *GeminiAIAssister) GetTerminalCommand(ctx context.Context, userMessage string) (string, error) { apiKey := g.aiParameters.apiKey - if apiKey == "" { - apiKey, _ = os.LookupEnv("GEMINI_API_KEY") - } client, err := genai.NewClient(ctx, &genai.ClientConfig{ Backend: genai.BackendGeminiAPI, APIKey: apiKey, From ead7c012eb0b11a1f0956769b71fcf7bf4e4c932 Mon Sep 17 00:00:00 2001 From: Amir Deris Date: Fri, 29 Aug 2025 13:45:48 +0200 Subject: [PATCH 6/6] Made AiParameters exportable --- agent.go | 26 +++++++++++++------------- agent_test.go | 20 ++++++++++---------- gromit.go | 6 +++--- gromit_test.go | 4 ++-- types.go | 4 ++-- 5 files changed, 30 insertions(+), 30 deletions(-) diff --git a/agent.go b/agent.go index 8eaad8e..04cf8e3 100644 --- a/agent.go +++ b/agent.go @@ -32,11 +32,11 @@ var _ Assister = (*AnthropicAIAssister)(nil) var _ Assister = (*GeminiAIAssister)(nil) type OpenAIAssister struct { - aiParameters + AiParameters } func (o *OpenAIAssister) GetTerminalCommand(ctx context.Context, userMessage string) (string, error) { - apiKey := o.aiParameters.apiKey + apiKey := o.AiParameters.apiKey var client openai.Client if apiKey != "" { client = openai.NewClient(openaiOption.WithAPIKey(apiKey)) @@ -46,9 +46,9 @@ func (o *OpenAIAssister) GetTerminalCommand(ctx context.Context, userMessage str chatCompletion, err := client.Chat.Completions.New(ctx, openai.ChatCompletionNewParams{ Messages: []openai.ChatCompletionMessageParamUnion{ openai.UserMessage(userMessage), - openai.SystemMessage(o.aiParameters.systemPrompt), + openai.SystemMessage(o.AiParameters.systemPrompt), }, - Model: o.aiParameters.model, + Model: o.AiParameters.model, }) if err != nil { return "", err @@ -57,12 +57,12 @@ func (o *OpenAIAssister) GetTerminalCommand(ctx context.Context, userMessage str } type AnthropicAIAssister struct { - aiParameters + AiParameters } func (c *AnthropicAIAssister) GetTerminalCommand(ctx context.Context, userMessage string) (string, error) { - apiKey := c.aiParameters.apiKey - maxTokens := c.aiParameters.maxTokens + apiKey := c.AiParameters.apiKey + maxTokens := c.AiParameters.maxTokens if maxTokens == 0 { maxTokens = defaultMaxTokens } @@ -75,7 +75,7 @@ func (c *AnthropicAIAssister) GetTerminalCommand(ctx context.Context, userMessag message, err := client.Messages.New(ctx, anthropic.MessageNewParams{ MaxTokens: maxTokens, System: []anthropic.TextBlockParam{ - {Text: c.aiParameters.systemPrompt}, + {Text: c.AiParameters.systemPrompt}, }, Messages: []anthropic.MessageParam{ anthropic.NewUserMessage(anthropic.NewTextBlock(userMessage)), @@ -96,11 +96,11 @@ func (c *AnthropicAIAssister) GetTerminalCommand(ctx context.Context, userMessag } type GeminiAIAssister struct { - aiParameters + AiParameters } func (g *GeminiAIAssister) GetTerminalCommand(ctx context.Context, userMessage string) (string, error) { - apiKey := g.aiParameters.apiKey + apiKey := g.AiParameters.apiKey client, err := genai.NewClient(ctx, &genai.ClientConfig{ Backend: genai.BackendGeminiAPI, APIKey: apiKey, @@ -111,7 +111,7 @@ func (g *GeminiAIAssister) GetTerminalCommand(ctx context.Context, userMessage s config := &genai.GenerateContentConfig{ SystemInstruction: &genai.Content{ Parts: []*genai.Part{ - {Text: g.aiParameters.systemPrompt}, + {Text: g.AiParameters.systemPrompt}, }, }, } @@ -127,14 +127,14 @@ func (g *GeminiAIAssister) GetTerminalCommand(ctx context.Context, userMessage s } type AssisterCreator interface { - GetAssister(parameters aiParameters) (Assister, error) + GetAssister(parameters AiParameters) (Assister, error) } var _ AssisterCreator = (*defaultAIAssisterCreator)(nil) type defaultAIAssisterCreator struct{} -func (d *defaultAIAssisterCreator) GetAssister(p aiParameters) (Assister, error) { +func (d *defaultAIAssisterCreator) GetAssister(p AiParameters) (Assister, error) { switch { case p.agent == "" || p.agent == openAIAgent: if p.model == "" { diff --git a/agent_test.go b/agent_test.go index c44143d..813d2c8 100644 --- a/agent_test.go +++ b/agent_test.go @@ -11,7 +11,7 @@ import ( type assisterTestCase struct { name string - inputAiParameters aiParameters + inputAiParameters AiParameters expectedModel string err string } @@ -24,7 +24,7 @@ func TestGetOpenAIAssister(t *testing.T) { }, { name: "OpenAI agent with given model should create correct assister", - inputAiParameters: aiParameters{ + inputAiParameters: AiParameters{ agent: openAIAgent, model: "gpt-5o-mini", apiKey: "key1234", @@ -34,7 +34,7 @@ func TestGetOpenAIAssister(t *testing.T) { }, { name: "Unknown agent should result in error", - inputAiParameters: aiParameters{ + inputAiParameters: AiParameters{ agent: "Unknown agent", model: "unknown model", apiKey: "key1234", @@ -67,7 +67,7 @@ func TestGetAnthropicAssister(t *testing.T) { tests := []assisterTestCase{ { name: "Anthropic agent with no model should default to Claude 3.5 Haiku latest", - inputAiParameters: aiParameters{ + inputAiParameters: AiParameters{ agent: anthropicAIAgent, apiKey: "key1234", maxTokens: defaultMaxTokens, @@ -76,7 +76,7 @@ func TestGetAnthropicAssister(t *testing.T) { }, { name: "Anthropic agent with given model should create correct assister", - inputAiParameters: aiParameters{ + inputAiParameters: AiParameters{ agent: anthropicAIAgent, model: string(anthropic.ModelClaudeOpus4_0), apiKey: "key1234", @@ -109,7 +109,7 @@ func TestGetGeminiAssister(t *testing.T) { tests := []assisterTestCase{ { name: "Gemini agent with no model should default to gemini-2.5-flash-lite", - inputAiParameters: aiParameters{ + inputAiParameters: AiParameters{ agent: geminiAIAgent, apiKey: "key1234", maxTokens: defaultMaxTokens, @@ -118,7 +118,7 @@ func TestGetGeminiAssister(t *testing.T) { }, { name: "Gemini agent with given model should create correct assister", - inputAiParameters: aiParameters{ + inputAiParameters: AiParameters{ agent: geminiAIAgent, model: "gemini-2.5-flash-preview-tts", apiKey: "key1234", @@ -175,21 +175,21 @@ func TestGetTerminalCommand(t *testing.T) { switch test.agent { case openAIAgent: assister = &OpenAIAssister{ - aiParameters{ + AiParameters{ model: openai.ChatModelGPT4o, systemPrompt: systemPrompt, }, } case anthropicAIAgent: assister = &AnthropicAIAssister{ - aiParameters{ + AiParameters{ model: string(anthropic.ModelClaude3_5HaikuLatest), systemPrompt: systemPrompt, }, } case geminiAIAgent: assister = &GeminiAIAssister{ - aiParameters{ + AiParameters{ model: geminiFlash, systemPrompt: systemPrompt, }, diff --git a/gromit.go b/gromit.go index edc8782..d492017 100644 --- a/gromit.go +++ b/gromit.go @@ -87,7 +87,7 @@ func (g *Gromit) actionGromit(ctx context.Context, command *cli.Command) error { prompt = systemPrompt } prompt = addEnvironmentInfo(g.configuration.systemInfo, prompt) - g.configuration.aiParameters = aiParameters{ + g.configuration.AiParameters = AiParameters{ maxTokens: g.Int64("maxTokens"), apiKey: g.String("apiKey"), agent: g.String("agent"), @@ -121,7 +121,7 @@ func (g *Gromit) actionGromit(ctx context.Context, command *cli.Command) error { } func (g *Gromit) handleUserQuery(ctx context.Context, query string) error { - assister, err := g.AssisterCreator.GetAssister(g.configuration.aiParameters) + assister, err := g.AssisterCreator.GetAssister(g.configuration.AiParameters) if err != nil { return err } @@ -252,7 +252,7 @@ func NewGromit(a AssisterCreator, mods ...ConfigurationModifier) (*Gromit, error w: os.Stdout, askForConfirmation: true, systemInfo: getSystemInfo(), - aiParameters: aiParameters{}, + AiParameters: AiParameters{}, } gromit := Gromit{ AssisterCreator: a, diff --git a/gromit_test.go b/gromit_test.go index 28288ce..84c6371 100644 --- a/gromit_test.go +++ b/gromit_test.go @@ -106,10 +106,10 @@ type mockAIProvider struct { actualUserMessage string - actualAiParameters aiParameters + actualAiParameters AiParameters } -func (m *mockAIProvider) GetAssister(p aiParameters) (Assister, error) { +func (m *mockAIProvider) GetAssister(p AiParameters) (Assister, error) { m.actualAiParameters = p if m.assisterError != nil { return nil, m.assisterError diff --git a/types.go b/types.go index 5449bbb..5ffa30c 100644 --- a/types.go +++ b/types.go @@ -16,7 +16,7 @@ type messagePrinter struct { } type configuration struct { - aiParameters + AiParameters promptPrefix string w io.Writer askForConfirmation bool @@ -29,7 +29,7 @@ type userConfirmation struct { type ConfigurationModifier func(*configuration) error -type aiParameters struct { +type AiParameters struct { systemPrompt string agent string model string