Skip to content

Commit 274ec52

Browse files
asdekasdek
committed
refactor: streamline test case structure and enhance message handling
- Updated test functions in bedrockllm_test.go to improve readability and maintainability by using slice initialization for message content. - Added `//nolint:funlen` comments to several test functions to suppress linting warnings related to function length. - Refactored message initialization in multiple tests to utilize `make` and `append`, enhancing clarity and performance. - Adjusted documentation in doc.go for consistent formatting of API mode descriptions and model lists. - Improved handling of reasoning model checks in reasoning.go and reasoning_test.go to include new Gemini model versions.
1 parent 8d98658 commit 274ec52

8 files changed

Lines changed: 85 additions & 82 deletions

File tree

‎llms/bedrock/bedrockllm_test.go‎

Lines changed: 45 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ func setUpTestWithTransport(rr *httprr.RecordReplay) (*bedrockruntime.Client, er
4343
return client, nil
4444
}
4545

46-
func TestAmazonOutputConverseAPI(t *testing.T) {
46+
func TestAmazonOutputConverseAPI(t *testing.T) { //nolint:funlen
4747
ctx := t.Context()
4848

4949
httprr.SkipIfNoCredentialsAndRecordingMissing(t, "AWS_ACCESS_KEY_ID")
@@ -1519,14 +1519,13 @@ func TestAmazonTextResponseWithThinkingConverseAPI(t *testing.T) {
15191519
t.Fatal(err)
15201520
}
15211521

1522-
messages := []llms.MessageContent{
1523-
{
1524-
Role: llms.ChatMessageTypeHuman,
1525-
Parts: []llms.ContentPart{
1526-
llms.TextPart("Solve: If x + 5 = 12, what is x?"),
1527-
},
1522+
messages := make([]llms.MessageContent, 0, 3)
1523+
messages = append(messages, llms.MessageContent{
1524+
Role: llms.ChatMessageTypeHuman,
1525+
Parts: []llms.ContentPart{
1526+
llms.TextPart("Solve: If x + 5 = 12, what is x?"),
15281527
},
1529-
}
1528+
})
15301529

15311530
// Request with thinking
15321531
resp, err := llm.GenerateContent(ctx, messages,
@@ -1603,14 +1602,13 @@ func TestAmazonTextResponseWithThinkingLegacyAPI(t *testing.T) {
16031602
t.Fatal(err)
16041603
}
16051604

1606-
messages := []llms.MessageContent{
1607-
{
1608-
Role: llms.ChatMessageTypeHuman,
1609-
Parts: []llms.ContentPart{
1610-
llms.TextPart("Solve: If x + 5 = 12, what is x?"),
1611-
},
1605+
messages := make([]llms.MessageContent, 0, 3)
1606+
messages = append(messages, llms.MessageContent{
1607+
Role: llms.ChatMessageTypeHuman,
1608+
Parts: []llms.ContentPart{
1609+
llms.TextPart("Solve: If x + 5 = 12, what is x?"),
16121610
},
1613-
}
1611+
})
16141612

16151613
// Request with thinking
16161614
resp, err := llm.GenerateContent(ctx, messages,
@@ -1666,7 +1664,7 @@ func TestAmazonTextResponseWithThinkingLegacyAPI(t *testing.T) {
16661664
}
16671665

16681666
// TestAmazonSingleToolCallWithThinkingConverseAPI tests single tool call with thinking
1669-
func TestAmazonSingleToolCallWithThinkingConverseAPI(t *testing.T) {
1667+
func TestAmazonSingleToolCallWithThinkingConverseAPI(t *testing.T) { //nolint:funlen
16701668
ctx := t.Context()
16711669

16721670
httprr.SkipIfNoCredentialsAndRecordingMissing(t, "AWS_ACCESS_KEY_ID")
@@ -1702,12 +1700,11 @@ func TestAmazonSingleToolCallWithThinkingConverseAPI(t *testing.T) {
17021700
},
17031701
}}
17041702

1705-
messages := []llms.MessageContent{
1706-
{
1707-
Role: llms.ChatMessageTypeHuman,
1708-
Parts: []llms.ContentPart{llms.TextPart("What's the weather in Boston?")},
1709-
},
1710-
}
1703+
messages := make([]llms.MessageContent, 0, 3)
1704+
messages = append(messages, llms.MessageContent{
1705+
Role: llms.ChatMessageTypeHuman,
1706+
Parts: []llms.ContentPart{llms.TextPart("What's the weather in Boston?")},
1707+
})
17111708

17121709
resp, err := llm.GenerateContent(ctx, messages,
17131710
llms.WithModel(bedrock.ModelAnthropicClaudeSonnet45),
@@ -1769,7 +1766,7 @@ func TestAmazonSingleToolCallWithThinkingConverseAPI(t *testing.T) {
17691766
}
17701767

17711768
// TestAmazonSequentialToolCallsWithThinkingConverseAPI tests sequential tool calls with thinking
1772-
func TestAmazonSequentialToolCallsWithThinkingConverseAPI(t *testing.T) {
1769+
func TestAmazonSequentialToolCallsWithThinkingConverseAPI(t *testing.T) { //nolint:funlen
17731770
ctx := t.Context()
17741771

17751772
httprr.SkipIfNoCredentialsAndRecordingMissing(t, "AWS_ACCESS_KEY_ID")
@@ -1807,14 +1804,13 @@ func TestAmazonSequentialToolCallsWithThinkingConverseAPI(t *testing.T) {
18071804
},
18081805
}
18091806

1810-
messages := []llms.MessageContent{
1811-
{
1812-
Role: llms.ChatMessageTypeHuman,
1813-
Parts: []llms.ContentPart{
1814-
llms.TextPart("Calculate (5 + 3) and then multiply by 2"),
1815-
},
1807+
messages := make([]llms.MessageContent, 0, 3)
1808+
messages = append(messages, llms.MessageContent{
1809+
Role: llms.ChatMessageTypeHuman,
1810+
Parts: []llms.ContentPart{
1811+
llms.TextPart("Calculate (5 + 3) and then multiply by 2"),
18161812
},
1817-
}
1813+
})
18181814

18191815
// First call
18201816
resp1, err := llm.GenerateContent(ctx, messages,
@@ -1960,7 +1956,7 @@ func TestAmazonTextResponseWithThinkingStreamingConverseAPI(t *testing.T) {
19601956
// Turn 1: CacheCreation > 0 (conversation history cached)
19611957
// Turn 2: CacheRead > 0 (previous turn read from cache), new content added to cache
19621958
// Turn 3+: Cache continues to be used and extended
1963-
func TestAmazonMultiTurnCachingWithToolsLegacyAPI(t *testing.T) {
1959+
func TestAmazonMultiTurnCachingWithToolsLegacyAPI(t *testing.T) { //nolint:funlen
19641960
ctx := t.Context()
19651961

19661962
httprr.SkipIfNoCredentialsAndRecordingMissing(t, "AWS_ACCESS_KEY_ID")
@@ -2190,7 +2186,7 @@ func TestAmazonMultiTurnCachingWithToolsLegacyAPI(t *testing.T) {
21902186
// Expected behavior:
21912187
// Turn 3: CacheCreation > 0 (conversation history with cachePoint)
21922188
// Turn 4: CacheRead > 0 (previous context read from cache)
2193-
func TestAmazonMultiTurnCachingWithToolsConverseAPI(t *testing.T) {
2189+
func TestAmazonMultiTurnCachingWithToolsConverseAPI(t *testing.T) { //nolint:funlen
21942190
ctx := t.Context()
21952191

21962192
httprr.SkipIfNoCredentialsAndRecordingMissing(t, "AWS_ACCESS_KEY_ID")
@@ -2252,16 +2248,17 @@ func TestAmazonMultiTurnCachingWithToolsConverseAPI(t *testing.T) {
22522248
systemPrompt := "BedrockConverseCache-v1: " + strings.Repeat("You are a helpful assistant with access to weather and flight booking capabilities. ", 15)
22532249

22542250
// Turn 1: Initial request with system prompt
2255-
messages := []llms.MessageContent{
2256-
{
2251+
messages := make([]llms.MessageContent, 0, 8)
2252+
messages = append(messages,
2253+
llms.MessageContent{
22572254
Role: llms.ChatMessageTypeSystem,
22582255
Parts: []llms.ContentPart{llms.TextPart(systemPrompt)},
22592256
},
2260-
{
2257+
llms.MessageContent{
22612258
Role: llms.ChatMessageTypeHuman,
22622259
Parts: []llms.ContentPart{llms.TextPart("What's the weather in Boston?")},
22632260
},
2264-
}
2261+
)
22652262

22662263
resp1, err := llm.GenerateContent(ctx, messages,
22672264
llms.WithModel(bedrock.ModelAnthropicClaudeSonnet45),
@@ -2418,7 +2415,7 @@ func TestAmazonMultiTurnCachingWithToolsConverseAPI(t *testing.T) {
24182415
// - Cache points are automatically added to conversation history
24192416
// - No manual bedrock.WithCacheControl() wrapper needed
24202417
// - Cache metrics show cache utilization
2421-
func TestAmazonAutomaticCachingLegacyAPI(t *testing.T) {
2418+
func TestAmazonAutomaticCachingLegacyAPI(t *testing.T) { //nolint:funlen
24222419
ctx := t.Context()
24232420

24242421
httprr.SkipIfNoCredentialsAndRecordingMissing(t, "AWS_ACCESS_KEY_ID")
@@ -2468,16 +2465,17 @@ func TestAmazonAutomaticCachingLegacyAPI(t *testing.T) {
24682465
systemPrompt := "BedrockAutoCacheTest-v1: " + strings.Repeat("You are a helpful assistant with access to weather capabilities. ", 15)
24692466

24702467
// Turn 1: Initial request with system prompt
2471-
messages := []llms.MessageContent{
2472-
{
2468+
messages := make([]llms.MessageContent, 0, 6)
2469+
messages = append(messages,
2470+
llms.MessageContent{
24732471
Role: llms.ChatMessageTypeSystem,
24742472
Parts: []llms.ContentPart{llms.TextPart(systemPrompt)},
24752473
},
2476-
{
2474+
llms.MessageContent{
24772475
Role: llms.ChatMessageTypeHuman,
24782476
Parts: []llms.ContentPart{llms.TextPart("What's the weather in Boston?")},
24792477
},
2480-
}
2478+
)
24812479

24822480
resp1, err := llm.GenerateContent(ctx, messages,
24832481
llms.WithModel(bedrock.ModelAnthropicClaudeSonnet45),
@@ -2589,7 +2587,7 @@ func TestAmazonAutomaticCachingLegacyAPI(t *testing.T) {
25892587

25902588
// TestAmazonAutomaticCachingConverseAPI tests automatic prompt caching with Converse API.
25912589
// This validates that WithAutomaticCaching() works correctly with Converse API.
2592-
func TestAmazonAutomaticCachingConverseAPI(t *testing.T) {
2590+
func TestAmazonAutomaticCachingConverseAPI(t *testing.T) { //nolint:funlen
25932591
ctx := t.Context()
25942592

25952593
httprr.SkipIfNoCredentialsAndRecordingMissing(t, "AWS_ACCESS_KEY_ID")
@@ -2640,16 +2638,17 @@ func TestAmazonAutomaticCachingConverseAPI(t *testing.T) {
26402638
systemPrompt := "BedrockConverseAutoCacheTest-v1: " + strings.Repeat("You are a helpful assistant with access to weather capabilities. ", 15)
26412639

26422640
// Turn 1: Initial request with system prompt
2643-
messages := []llms.MessageContent{
2644-
{
2641+
messages := make([]llms.MessageContent, 0, 6)
2642+
messages = append(messages,
2643+
llms.MessageContent{
26452644
Role: llms.ChatMessageTypeSystem,
26462645
Parts: []llms.ContentPart{llms.TextPart(systemPrompt)},
26472646
},
2648-
{
2647+
llms.MessageContent{
26492648
Role: llms.ChatMessageTypeHuman,
26502649
Parts: []llms.ContentPart{llms.TextPart("What's the weather in Boston?")},
26512650
},
2652-
}
2651+
)
26532652

26542653
resp1, err := llm.GenerateContent(ctx, messages,
26552654
llms.WithModel(bedrock.ModelAnthropicClaudeSonnet45),

‎llms/bedrock/doc.go‎

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,8 @@
1515
//
1616
// Two API modes are supported:
1717
//
18-
// - Legacy API: Model-specific implementations via InvokeModel/InvokeModelWithResponseStream
19-
// - Converse API: Unified implementation via Converse/ConverseStream (recommended)
18+
// - Legacy API: Model-specific implementations via InvokeModel/InvokeModelWithResponseStream
19+
// - Converse API: Unified implementation via Converse/ConverseStream (recommended)
2020
//
2121
// # Basic Usage
2222
//
@@ -129,17 +129,17 @@
129129
//
130130
// See models_list.go for complete list. Major providers:
131131
//
132-
// - Anthropic: Claude 4.6 (Opus, Sonnet), Claude 4.5, 4.1, 4, 3.7, 3.5
133-
// - Amazon: Nova 2 Lite, Nova Premier, Nova Pro, Nova Lite, Nova Micro
134-
// - Meta: Llama 4, Llama 3.3, 3.2, 3.1, 3
135-
// - Cohere: Command R, Command R+
136-
// - AI21: Jamba 1.5 Large, Mini
137-
// - DeepSeek: R1
138-
// - OpenAI: GPT-OSS-120B, GPT-OSS-20B
139-
// - Qwen: Qwen3 Next, Qwen3 VL, Qwen3 32B, Qwen3 Coder (30B, Next)
140-
// - Mistral: Large 3, Magistral Small
141-
// - Moonshot: Kimi K2.5, Kimi K2 Thinking
142-
// - Z.AI: GLM-4.7, GLM-4.7-Flash
132+
// - Anthropic: Claude 4.6 (Opus, Sonnet), Claude 4.5, 4.1, 4, 3.7, 3.5
133+
// - Amazon: Nova 2 Lite, Nova Premier, Nova Pro, Nova Lite, Nova Micro
134+
// - Meta: Llama 4, Llama 3.3, 3.2, 3.1, 3
135+
// - Cohere: Command R, Command R+
136+
// - AI21: Jamba 1.5 Large, Mini
137+
// - DeepSeek: R1
138+
// - OpenAI: GPT-OSS-120B, GPT-OSS-20B
139+
// - Qwen: Qwen3 Next, Qwen3 VL, Qwen3 32B, Qwen3 Coder (30B, Next)
140+
// - Mistral: Large 3, Magistral Small
141+
// - Moonshot: Kimi K2.5, Kimi K2 Thinking
142+
// - Z.AI: GLM-4.7, GLM-4.7-Flash
143143
//
144144
// # Error Handling
145145
//
@@ -163,16 +163,16 @@
163163
//
164164
// The client uses AWS SDK v2 configuration:
165165
//
166-
// - Credentials: From environment (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) or AWS config
167-
// - Region: From environment (AWS_REGION) or default config
168-
// - Custom configuration: Use bedrock.WithClient() with pre-configured bedrockruntime.Client
166+
// - Credentials: From environment (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) or AWS config
167+
// - Region: From environment (AWS_REGION) or default config
168+
// - Custom configuration: Use bedrock.WithClient() with pre-configured bedrockruntime.Client
169169
//
170170
// # Performance Considerations
171171
//
172-
// - Converse API is recommended for new applications (unified, better error handling)
173-
// - Automatic caching reduces costs by 90% for cached tokens (Claude 4.x only)
174-
// - Streaming reduces latency for interactive applications
175-
// - Minimum cache checkpoint: 1024 tokens (Sonnet 4.5), 4096 tokens (Haiku 4.5)
172+
// - Converse API is recommended for new applications (unified, better error handling)
173+
// - Automatic caching reduces costs by 90% for cached tokens (Claude 4.x only)
174+
// - Streaming reduces latency for interactive applications
175+
// - Minimum cache checkpoint: 1024 tokens (Sonnet 4.5), 4096 tokens (Haiku 4.5)
176176
//
177177
// # Maintenance
178178
//
@@ -194,9 +194,9 @@
194194
//
195195
// Tests use httprr for HTTP recording/replay:
196196
//
197-
// - Integration tests: bedrockllm_test.go (requires AWS credentials)
198-
// - Unit tests: bedrockllm_unit_test.go (no credentials needed)
199-
// - Tool calling: bedrock_tool_integration_test.go
197+
// - Integration tests: bedrockllm_test.go (requires AWS credentials)
198+
// - Unit tests: bedrockllm_unit_test.go (no credentials needed)
199+
// - Tool calling: bedrock_tool_integration_test.go
200200
//
201201
// Recording new HTTP interactions:
202202
//

‎llms/huggingface/huggingfacellm_test.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,9 +171,11 @@ func TestHuggingFaceLLMGenerateContent(t *testing.T) {
171171
if err != nil {
172172
t.Fatal(err)
173173
}
174+
//nolint:staticcheck // false positive: resp can be nil but t.Fatal exits immediately
174175
if resp == nil {
175176
t.Fatal("expected non-nil response")
176177
}
178+
//nolint:staticcheck // t.Fatal above ensures resp is not nil
177179
if len(resp.Choices) != 1 {
178180
t.Fatalf("expected 1 choice, got %d", len(resp.Choices))
179181
}

‎llms/ollama/ollama_cloud_test.go‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,17 +24,17 @@ type removeTimestampTransport struct {
2424
func (t *removeTimestampTransport) RoundTrip(req *http.Request) (*http.Response, error) {
2525
// Clone request to avoid modifying original
2626
clonedReq := req.Clone(req.Context())
27-
27+
2828
// Remove ts query parameter
2929
q := clonedReq.URL.Query()
3030
q.Del("ts")
3131
clonedReq.URL.RawQuery = q.Encode()
32-
32+
3333
return t.base.RoundTrip(clonedReq)
3434
}
3535

3636
// newCloudTestClient creates a test client configured for Ollama Cloud
37-
func newCloudTestClient(t *testing.T, opts ...Option) *LLM {
37+
func newCloudTestClient(t *testing.T) *LLM {
3838
t.Helper()
3939

4040
// Check for required credentials and skip if not available
@@ -79,13 +79,12 @@ func newCloudTestClient(t *testing.T, opts ...Option) *LLM {
7979
cloudModel = envModel
8080
}
8181

82-
// Always add server URL, API key, and HTTP client
83-
opts = append([]Option{
82+
opts := []Option{
8483
WithServerURL(serverURL),
8584
WithAPIKey(apiKey),
8685
WithHTTPClient(wrappedClient),
8786
WithModel(cloudModel),
88-
}, opts...)
87+
}
8988

9089
c, err := New(opts...)
9190
require.NoError(t, err)

‎llms/ollama/ollama_test.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -365,14 +365,14 @@ func TestWithPullTimeout(t *testing.T) {
365365
// This test only works in recording mode (timeout behavior cannot be replayed)
366366
// Skip if httprr file doesn't exist
367367
httprr.SkipIfNoCredentialsAndRecordingMissing(t)
368-
368+
369369
rr := httprr.OpenForTest(t, http.DefaultTransport)
370370
defer rr.Close()
371-
371+
372372
if !rr.Recording() {
373373
t.Skip("Skipping pull timeout test when not recording (timeout behavior cannot be replayed)")
374374
}
375-
375+
376376
// Scrub dynamic headers
377377
rr.ScrubReq(func(req *http.Request) error {
378378
req.Header.Del("Authorization")

‎llms/ollama/ollamallm.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,10 @@ type authRoundTripper struct {
4646
func (a *authRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
4747
// Clone the request to avoid modifying the original
4848
clonedReq := req.Clone(req.Context())
49-
49+
5050
// Add Authorization header with Bearer token
5151
clonedReq.Header.Set("Authorization", "Bearer "+a.apiKey)
52-
52+
5353
// Use the wrapped transport to perform the actual request
5454
return a.transport.RoundTrip(clonedReq)
5555
}
@@ -88,7 +88,7 @@ func New(opts ...Option) (*LLM, error) {
8888
if transport == nil {
8989
transport = http.DefaultTransport
9090
}
91-
91+
9292
// Create a new HTTP client with the auth transport
9393
httpClient = &http.Client{
9494
Transport: &authRoundTripper{

‎llms/reasoning/reasoning.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -215,7 +215,8 @@ func DefaultIsReasoningModel(model string) bool {
215215

216216
// Google Gemini reasoning models
217217
if strings.HasPrefix(modelLower, "gemini-2.5-") ||
218-
strings.HasPrefix(modelLower, "gemini-3-") {
218+
strings.HasPrefix(modelLower, "gemini-3-") ||
219+
strings.HasPrefix(modelLower, "gemini-3.1-") {
219220
return true
220221
}
221222

‎llms/reasoning/reasoning_test.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -691,7 +691,9 @@ func TestIsReasoningModel(t *testing.T) { //nolint:funlen
691691
{"gemini-2.5-pro", true},
692692
{"gemini-2.5-pro-preview", true},
693693
{"gemini-3-flash-preview", true},
694+
{"gemini-3.1-flash-lite-preview", true},
694695
{"gemini-3-pro-preview", true},
696+
{"gemini-3.1-pro-preview", true},
695697
{"google/gemini-2.5-flash", true},
696698

697699
// X-AI Grok reasoning models

0 commit comments

Comments
 (0)