Skip to content

Amir deris/configure ai agent max token and other properties #24

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
amir-deris merged 6 commits into from
Aug 29, 2025
Merged

Amir deris/configure ai agent max token and other properties #24

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
80c0cdb
Refactored types into a different file
Aug 24, 2025
6cb728a
refactored API parameters and passed extra options
Aug 25, 2025
4f4872b
Fixed gromit tests
Aug 26, 2025
44c7f08
Updated agent tests
Aug 26, 2025
f8531f8
Removed API key env variable lookup
Aug 29, 2025
ead7c01
Made AiParameters exportable
Aug 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 53 additions & 35 deletions agent.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -20,26 +23,34 @@ const (
geminiFlash = "gemini-2.5-flash"
)

const (
defaultMaxTokens = 1024

@masih masih Aug 28, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Group this with the rest of the constants above?

)

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)
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
Expand All @@ -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")

@masih masih Aug 28, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to set this? I wonder if the SDK would internally infer this if unset.
In which case, what would perhaps fit better is to only set it if the user has specified it as a CLI flag. Otherwise, do nothing. The SDK does whatever the default SDK does.

@amir-deris amir-deris Aug 29, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's a great point. The SDK by default looks at those env variables. I will fix this.

}
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)),
Expand All @@ -77,20 +96,25 @@ 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")

@masih masih Aug 28, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ditto; see https://github.com/GromitAI/gromit/pull/24/files#r2308183139

}
client, err := genai.NewClient(ctx, &genai.ClientConfig{
Backend: genai.BackendGeminiAPI,
APIKey: apiKey,
})
if err != nil {
return "", err
}
config := &genai.GenerateContentConfig{
SystemInstruction: &genai.Content{
Parts: []*genai.Part{
{Text: systemMessage},
{Text: g.aiParameters.systemPrompt},
},
},
}
Expand All @@ -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)

@masih masih Aug 28, 2025

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is generally a bad smell: having a exported interface with unexported types. I would either unexport the interface or export the aiParameters type in this case.

@amir-deris amir-deris Aug 29, 2025

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Totally missed that! Thanks for the great point.

}

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)
}
}
95 changes: 67 additions & 28 deletions agent_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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")
}
Expand All @@ -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")
}
Expand All @@ -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")
}
Expand Down Expand Up @@ -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")
})
Expand Down
Loading