575 lines
13 KiB
Go
575 lines
13 KiB
Go
package google
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/ajac-zero/latticelm/internal/api"
|
|
"github.com/ajac-zero/latticelm/internal/config"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"google.golang.org/genai"
|
|
)
|
|
|
|
func TestNew(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
cfg config.ProviderConfig
|
|
expectError bool
|
|
validate func(t *testing.T, p *Provider, err error)
|
|
}{
|
|
{
|
|
name: "creates provider with API key",
|
|
cfg: config.ProviderConfig{
|
|
APIKey: "test-api-key",
|
|
Model: "gemini-2.0-flash",
|
|
},
|
|
expectError: false,
|
|
validate: func(t *testing.T, p *Provider, err error) {
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, p)
|
|
assert.NotNil(t, p.client)
|
|
assert.Equal(t, "test-api-key", p.cfg.APIKey)
|
|
assert.Equal(t, "gemini-2.0-flash", p.cfg.Model)
|
|
},
|
|
},
|
|
{
|
|
name: "creates provider without API key",
|
|
cfg: config.ProviderConfig{
|
|
APIKey: "",
|
|
},
|
|
expectError: false,
|
|
validate: func(t *testing.T, p *Provider, err error) {
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, p)
|
|
assert.Nil(t, p.client)
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
p, err := New(tt.cfg)
|
|
tt.validate(t, p, err)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestNewVertexAI(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
cfg config.VertexAIConfig
|
|
expectError bool
|
|
validate func(t *testing.T, p *Provider, err error)
|
|
}{
|
|
{
|
|
name: "creates Vertex AI provider with project and location",
|
|
cfg: config.VertexAIConfig{
|
|
Project: "my-gcp-project",
|
|
Location: "us-central1",
|
|
},
|
|
expectError: false,
|
|
validate: func(t *testing.T, p *Provider, err error) {
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, p)
|
|
// Client creation may fail without proper GCP credentials in test env
|
|
// but provider should be created
|
|
},
|
|
},
|
|
{
|
|
name: "creates Vertex AI provider without project",
|
|
cfg: config.VertexAIConfig{
|
|
Project: "",
|
|
Location: "us-central1",
|
|
},
|
|
expectError: false,
|
|
validate: func(t *testing.T, p *Provider, err error) {
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, p)
|
|
assert.Nil(t, p.client)
|
|
},
|
|
},
|
|
{
|
|
name: "creates Vertex AI provider without location",
|
|
cfg: config.VertexAIConfig{
|
|
Project: "my-gcp-project",
|
|
Location: "",
|
|
},
|
|
expectError: false,
|
|
validate: func(t *testing.T, p *Provider, err error) {
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, p)
|
|
assert.Nil(t, p.client)
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
p, err := NewVertexAI(tt.cfg)
|
|
tt.validate(t, p, err)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestProvider_Name(t *testing.T) {
|
|
p := &Provider{}
|
|
assert.Equal(t, "google", p.Name())
|
|
}
|
|
|
|
func TestProvider_Generate_Validation(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
provider *Provider
|
|
messages []api.Message
|
|
req *api.ResponseRequest
|
|
expectError bool
|
|
errorMsg string
|
|
}{
|
|
{
|
|
name: "returns error when client not initialized",
|
|
provider: &Provider{
|
|
client: nil,
|
|
},
|
|
messages: []api.Message{
|
|
{Role: "user", Content: []api.ContentBlock{{Type: "input_text", Text: "Hello"}}},
|
|
},
|
|
req: &api.ResponseRequest{
|
|
Model: "gemini-2.0-flash",
|
|
},
|
|
expectError: true,
|
|
errorMsg: "client not initialized",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result, err := tt.provider.Generate(context.Background(), tt.messages, tt.req)
|
|
|
|
if tt.expectError {
|
|
assert.Error(t, err)
|
|
assert.Nil(t, result)
|
|
if tt.errorMsg != "" {
|
|
assert.Contains(t, err.Error(), tt.errorMsg)
|
|
}
|
|
} else {
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, result)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestProvider_GenerateStream_Validation(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
provider *Provider
|
|
messages []api.Message
|
|
req *api.ResponseRequest
|
|
expectError bool
|
|
errorMsg string
|
|
}{
|
|
{
|
|
name: "returns error when client not initialized",
|
|
provider: &Provider{
|
|
client: nil,
|
|
},
|
|
messages: []api.Message{
|
|
{Role: "user", Content: []api.ContentBlock{{Type: "input_text", Text: "Hello"}}},
|
|
},
|
|
req: &api.ResponseRequest{
|
|
Model: "gemini-2.0-flash",
|
|
},
|
|
expectError: true,
|
|
errorMsg: "client not initialized",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
deltaChan, errChan := tt.provider.GenerateStream(context.Background(), tt.messages, tt.req)
|
|
|
|
// Read from channels
|
|
var receivedError error
|
|
for {
|
|
select {
|
|
case _, ok := <-deltaChan:
|
|
if !ok {
|
|
deltaChan = nil
|
|
}
|
|
case err, ok := <-errChan:
|
|
if ok && err != nil {
|
|
receivedError = err
|
|
}
|
|
errChan = nil
|
|
}
|
|
|
|
if deltaChan == nil && errChan == nil {
|
|
break
|
|
}
|
|
}
|
|
|
|
if tt.expectError {
|
|
require.Error(t, receivedError)
|
|
if tt.errorMsg != "" {
|
|
assert.Contains(t, receivedError.Error(), tt.errorMsg)
|
|
}
|
|
} else {
|
|
assert.NoError(t, receivedError)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestConvertMessages(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
messages []api.Message
|
|
expectedContents int
|
|
expectedSystem string
|
|
validate func(t *testing.T, contents []*genai.Content, systemText string)
|
|
}{
|
|
{
|
|
name: "converts user message",
|
|
messages: []api.Message{
|
|
{
|
|
Role: "user",
|
|
Content: []api.ContentBlock{
|
|
{Type: "input_text", Text: "Hello"},
|
|
},
|
|
},
|
|
},
|
|
expectedContents: 1,
|
|
expectedSystem: "",
|
|
validate: func(t *testing.T, contents []*genai.Content, systemText string) {
|
|
require.Len(t, contents, 1)
|
|
assert.Equal(t, "user", contents[0].Role)
|
|
assert.Equal(t, "", systemText)
|
|
},
|
|
},
|
|
{
|
|
name: "extracts system message",
|
|
messages: []api.Message{
|
|
{
|
|
Role: "system",
|
|
Content: []api.ContentBlock{
|
|
{Type: "input_text", Text: "You are a helpful assistant"},
|
|
},
|
|
},
|
|
{
|
|
Role: "user",
|
|
Content: []api.ContentBlock{
|
|
{Type: "input_text", Text: "Hello"},
|
|
},
|
|
},
|
|
},
|
|
expectedContents: 1,
|
|
expectedSystem: "You are a helpful assistant",
|
|
validate: func(t *testing.T, contents []*genai.Content, systemText string) {
|
|
require.Len(t, contents, 1)
|
|
assert.Equal(t, "You are a helpful assistant", systemText)
|
|
assert.Equal(t, "user", contents[0].Role)
|
|
},
|
|
},
|
|
{
|
|
name: "converts assistant message with tool calls",
|
|
messages: []api.Message{
|
|
{
|
|
Role: "assistant",
|
|
Content: []api.ContentBlock{
|
|
{Type: "output_text", Text: "Let me check the weather"},
|
|
},
|
|
ToolCalls: []api.ToolCall{
|
|
{
|
|
ID: "call_123",
|
|
Name: "get_weather",
|
|
Arguments: `{"location": "SF"}`,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
expectedContents: 1,
|
|
validate: func(t *testing.T, contents []*genai.Content, systemText string) {
|
|
require.Len(t, contents, 1)
|
|
assert.Equal(t, "model", contents[0].Role)
|
|
// Should have text part and function call part
|
|
assert.GreaterOrEqual(t, len(contents[0].Parts), 1)
|
|
},
|
|
},
|
|
{
|
|
name: "converts tool result message",
|
|
messages: []api.Message{
|
|
{
|
|
Role: "assistant",
|
|
ToolCalls: []api.ToolCall{
|
|
{ID: "call_123", Name: "get_weather", Arguments: "{}"},
|
|
},
|
|
},
|
|
{
|
|
Role: "tool",
|
|
CallID: "call_123",
|
|
Name: "get_weather",
|
|
Content: []api.ContentBlock{
|
|
{Type: "output_text", Text: `{"temp": 72}`},
|
|
},
|
|
},
|
|
},
|
|
expectedContents: 2,
|
|
validate: func(t *testing.T, contents []*genai.Content, systemText string) {
|
|
require.Len(t, contents, 2)
|
|
// Tool result should be in user role
|
|
assert.Equal(t, "user", contents[1].Role)
|
|
require.Len(t, contents[1].Parts, 1)
|
|
assert.NotNil(t, contents[1].Parts[0].FunctionResponse)
|
|
},
|
|
},
|
|
{
|
|
name: "handles developer message as system",
|
|
messages: []api.Message{
|
|
{
|
|
Role: "developer",
|
|
Content: []api.ContentBlock{
|
|
{Type: "input_text", Text: "Developer instruction"},
|
|
},
|
|
},
|
|
},
|
|
expectedContents: 0,
|
|
expectedSystem: "Developer instruction",
|
|
validate: func(t *testing.T, contents []*genai.Content, systemText string) {
|
|
assert.Len(t, contents, 0)
|
|
assert.Equal(t, "Developer instruction", systemText)
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
contents, systemText := convertMessages(tt.messages)
|
|
assert.Len(t, contents, tt.expectedContents)
|
|
assert.Equal(t, tt.expectedSystem, systemText)
|
|
if tt.validate != nil {
|
|
tt.validate(t, contents, systemText)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestBuildConfig(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
systemText string
|
|
req *api.ResponseRequest
|
|
tools []*genai.Tool
|
|
toolConfig *genai.ToolConfig
|
|
expectNil bool
|
|
validate func(t *testing.T, cfg *genai.GenerateContentConfig)
|
|
}{
|
|
{
|
|
name: "returns nil when no config needed",
|
|
systemText: "",
|
|
req: &api.ResponseRequest{},
|
|
tools: nil,
|
|
toolConfig: nil,
|
|
expectNil: true,
|
|
},
|
|
{
|
|
name: "creates config with system text",
|
|
systemText: "You are helpful",
|
|
req: &api.ResponseRequest{},
|
|
expectNil: false,
|
|
validate: func(t *testing.T, cfg *genai.GenerateContentConfig) {
|
|
require.NotNil(t, cfg)
|
|
require.NotNil(t, cfg.SystemInstruction)
|
|
assert.Len(t, cfg.SystemInstruction.Parts, 1)
|
|
},
|
|
},
|
|
{
|
|
name: "creates config with max tokens",
|
|
systemText: "",
|
|
req: &api.ResponseRequest{
|
|
MaxOutputTokens: intPtr(1000),
|
|
},
|
|
expectNil: false,
|
|
validate: func(t *testing.T, cfg *genai.GenerateContentConfig) {
|
|
require.NotNil(t, cfg)
|
|
assert.Equal(t, int32(1000), cfg.MaxOutputTokens)
|
|
},
|
|
},
|
|
{
|
|
name: "creates config with temperature",
|
|
systemText: "",
|
|
req: &api.ResponseRequest{
|
|
Temperature: float64Ptr(0.7),
|
|
},
|
|
expectNil: false,
|
|
validate: func(t *testing.T, cfg *genai.GenerateContentConfig) {
|
|
require.NotNil(t, cfg)
|
|
require.NotNil(t, cfg.Temperature)
|
|
assert.Equal(t, float32(0.7), *cfg.Temperature)
|
|
},
|
|
},
|
|
{
|
|
name: "creates config with top_p",
|
|
systemText: "",
|
|
req: &api.ResponseRequest{
|
|
TopP: float64Ptr(0.9),
|
|
},
|
|
expectNil: false,
|
|
validate: func(t *testing.T, cfg *genai.GenerateContentConfig) {
|
|
require.NotNil(t, cfg)
|
|
require.NotNil(t, cfg.TopP)
|
|
assert.Equal(t, float32(0.9), *cfg.TopP)
|
|
},
|
|
},
|
|
{
|
|
name: "creates config with tools",
|
|
systemText: "",
|
|
req: &api.ResponseRequest{},
|
|
tools: []*genai.Tool{
|
|
{
|
|
FunctionDeclarations: []*genai.FunctionDeclaration{
|
|
{Name: "get_weather"},
|
|
},
|
|
},
|
|
},
|
|
expectNil: false,
|
|
validate: func(t *testing.T, cfg *genai.GenerateContentConfig) {
|
|
require.NotNil(t, cfg)
|
|
require.Len(t, cfg.Tools, 1)
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
cfg := buildConfig(tt.systemText, tt.req, tt.tools, tt.toolConfig)
|
|
if tt.expectNil {
|
|
assert.Nil(t, cfg)
|
|
} else {
|
|
require.NotNil(t, cfg)
|
|
if tt.validate != nil {
|
|
tt.validate(t, cfg)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestChooseModel(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
requested string
|
|
defaultModel string
|
|
expected string
|
|
}{
|
|
{
|
|
name: "returns requested model when provided",
|
|
requested: "gemini-1.5-pro",
|
|
defaultModel: "gemini-2.0-flash",
|
|
expected: "gemini-1.5-pro",
|
|
},
|
|
{
|
|
name: "returns default model when requested is empty",
|
|
requested: "",
|
|
defaultModel: "gemini-2.0-flash",
|
|
expected: "gemini-2.0-flash",
|
|
},
|
|
{
|
|
name: "returns fallback when both empty",
|
|
requested: "",
|
|
defaultModel: "",
|
|
expected: "gemini-2.0-flash-exp",
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := chooseModel(tt.requested, tt.defaultModel)
|
|
assert.Equal(t, tt.expected, result)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestExtractToolCallDelta(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
part *genai.Part
|
|
index int
|
|
expected *api.ToolCallDelta
|
|
}{
|
|
{
|
|
name: "extracts tool call delta",
|
|
part: &genai.Part{
|
|
FunctionCall: &genai.FunctionCall{
|
|
ID: "call_123",
|
|
Name: "get_weather",
|
|
Args: map[string]any{"location": "SF"},
|
|
},
|
|
},
|
|
index: 0,
|
|
expected: &api.ToolCallDelta{
|
|
Index: 0,
|
|
ID: "call_123",
|
|
Name: "get_weather",
|
|
Arguments: `{"location":"SF"}`,
|
|
},
|
|
},
|
|
{
|
|
name: "returns nil for nil part",
|
|
part: nil,
|
|
index: 0,
|
|
expected: nil,
|
|
},
|
|
{
|
|
name: "returns nil for part without function call",
|
|
part: &genai.Part{Text: "Hello"},
|
|
index: 0,
|
|
expected: nil,
|
|
},
|
|
{
|
|
name: "generates ID when not provided",
|
|
part: &genai.Part{
|
|
FunctionCall: &genai.FunctionCall{
|
|
ID: "",
|
|
Name: "get_time",
|
|
Args: map[string]any{},
|
|
},
|
|
},
|
|
index: 1,
|
|
expected: &api.ToolCallDelta{
|
|
Index: 1,
|
|
Name: "get_time",
|
|
Arguments: `{}`,
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := extractToolCallDelta(tt.part, tt.index)
|
|
if tt.expected == nil {
|
|
assert.Nil(t, result)
|
|
} else {
|
|
require.NotNil(t, result)
|
|
assert.Equal(t, tt.expected.Index, result.Index)
|
|
assert.Equal(t, tt.expected.Name, result.Name)
|
|
if tt.part != nil && tt.part.FunctionCall != nil && tt.part.FunctionCall.ID != "" {
|
|
assert.Equal(t, tt.expected.ID, result.ID)
|
|
} else if tt.expected.ID == "" {
|
|
// Generated ID should start with "call_"
|
|
assert.Contains(t, result.ID, "call_")
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// Helper functions
|
|
func intPtr(i int) *int {
|
|
return &i
|
|
}
|
|
|
|
func float64Ptr(f float64) *float64 {
|
|
return &f
|
|
}
|