mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-25 18:03:29 +08:00
c575164695
### Summary As title --------- Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
495 lines
15 KiB
Go
495 lines
15 KiB
Go
package models
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"reflect"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
|
|
"google.golang.org/genai"
|
|
)
|
|
|
|
var googleListModelsMu sync.Mutex
|
|
|
|
func withGoogleListModelsStub(t *testing.T, fn func(context.Context, *genai.ClientConfig) ([]ListModelResponse, error)) {
|
|
t.Helper()
|
|
|
|
googleListModelsMu.Lock()
|
|
original := googleListModels
|
|
googleListModels = fn
|
|
t.Cleanup(func() {
|
|
googleListModels = original
|
|
googleListModelsMu.Unlock()
|
|
})
|
|
}
|
|
|
|
func TestGoogleModelListModelsRequiresAPIKey(t *testing.T) {
|
|
ctx := t.Context()
|
|
model := &GoogleModel{}
|
|
cases := []struct {
|
|
name string
|
|
apiConfig *APIConfig
|
|
}{
|
|
{
|
|
name: "nil config",
|
|
apiConfig: nil,
|
|
},
|
|
{
|
|
name: "nil api key",
|
|
apiConfig: &APIConfig{},
|
|
},
|
|
{
|
|
name: "empty api key",
|
|
apiConfig: &APIConfig{
|
|
ApiKey: stringPtr(""),
|
|
},
|
|
},
|
|
{
|
|
name: "blank api key",
|
|
apiConfig: &APIConfig{
|
|
ApiKey: stringPtr(" \t\n "),
|
|
},
|
|
},
|
|
}
|
|
|
|
calls := 0
|
|
withGoogleListModelsStub(t, func(context.Context, *genai.ClientConfig) ([]ListModelResponse, error) {
|
|
calls++
|
|
return nil, nil
|
|
})
|
|
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
models, err := model.ListModels(ctx, tc.apiConfig)
|
|
if err == nil {
|
|
t.Fatal("expected an API key error")
|
|
}
|
|
if !strings.Contains(err.Error(), "api key is required") {
|
|
t.Fatalf("expected API key error, got %v", err)
|
|
}
|
|
if models != nil {
|
|
t.Fatalf("expected no models, got %v", models)
|
|
}
|
|
})
|
|
}
|
|
|
|
if calls != 0 {
|
|
t.Fatalf("expected no ListModels calls without an API key, got %d", calls)
|
|
}
|
|
}
|
|
|
|
func TestGoogleModelListModelsReturnsModelNames(t *testing.T) {
|
|
ctx := t.Context()
|
|
model := &GoogleModel{}
|
|
apiKey := "test-api-key"
|
|
configuredAPIKey := " " + apiKey + " "
|
|
expected := []ListModelResponse{{Name: "models/gemini-2.5-flash"}, {Name: "models/gemini-2.5-pro"}}
|
|
|
|
withGoogleListModelsStub(t, func(_ context.Context, config *genai.ClientConfig) ([]ListModelResponse, error) {
|
|
if config.APIKey != apiKey {
|
|
t.Fatalf("expected API key %q, got %q", apiKey, config.APIKey)
|
|
}
|
|
return expected, nil
|
|
})
|
|
|
|
models, err := model.ListModels(ctx, &APIConfig{ApiKey: &configuredAPIKey})
|
|
if err != nil {
|
|
t.Fatalf("expected no error, got %v", err)
|
|
}
|
|
if !reflect.DeepEqual(models, expected) {
|
|
t.Fatalf("expected models %v, got %v", expected, models)
|
|
}
|
|
}
|
|
|
|
func TestGoogleModelCheckConnectionUsesListModels(t *testing.T) {
|
|
ctx := t.Context()
|
|
customBaseURL := "https://check-connection.example.test/google"
|
|
model := NewGoogleModel(map[string]string{"default": customBaseURL}, URLSuffix{})
|
|
apiKey := "test-api-key"
|
|
calls := 0
|
|
|
|
withGoogleListModelsStub(t, func(_ context.Context, config *genai.ClientConfig) ([]ListModelResponse, error) {
|
|
calls++
|
|
if config.APIKey != apiKey {
|
|
t.Fatalf("expected API key %q, got %q", apiKey, config.APIKey)
|
|
}
|
|
if config.HTTPOptions.BaseURL != customBaseURL {
|
|
t.Fatalf("expected base URL %q, got %q", customBaseURL, config.HTTPOptions.BaseURL)
|
|
}
|
|
return []ListModelResponse{{Name: "models/gemini-2.5-flash"}}, nil
|
|
})
|
|
|
|
if err := model.CheckConnection(ctx, &APIConfig{ApiKey: &apiKey}); err != nil {
|
|
t.Fatalf("expected no error, got %v", err)
|
|
}
|
|
if calls != 1 {
|
|
t.Fatalf("expected one ListModels call, got %d", calls)
|
|
}
|
|
}
|
|
|
|
func TestGoogleModelCheckConnectionRequiresAPIKey(t *testing.T) {
|
|
ctx := t.Context()
|
|
model := &GoogleModel{}
|
|
calls := 0
|
|
|
|
withGoogleListModelsStub(t, func(context.Context, *genai.ClientConfig) ([]ListModelResponse, error) {
|
|
calls++
|
|
return nil, nil
|
|
})
|
|
|
|
cases := []struct {
|
|
name string
|
|
apiConfig *APIConfig
|
|
}{
|
|
{
|
|
name: "nil config",
|
|
apiConfig: nil,
|
|
},
|
|
{
|
|
name: "nil api key",
|
|
apiConfig: &APIConfig{},
|
|
},
|
|
{
|
|
name: "empty api key",
|
|
apiConfig: &APIConfig{
|
|
ApiKey: stringPtr(""),
|
|
},
|
|
},
|
|
{
|
|
name: "blank api key",
|
|
apiConfig: &APIConfig{
|
|
ApiKey: stringPtr(" \t\n "),
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
err := model.CheckConnection(ctx, tc.apiConfig)
|
|
if err == nil {
|
|
t.Fatal("expected an API key error")
|
|
}
|
|
if !strings.Contains(err.Error(), "api key is required") {
|
|
t.Fatalf("expected API key error, got %v", err)
|
|
}
|
|
})
|
|
}
|
|
if calls != 0 {
|
|
t.Fatalf("expected no ListModels calls without an API key, got %d", calls)
|
|
}
|
|
}
|
|
|
|
func TestGoogleModelCheckConnectionReturnsListModelsError(t *testing.T) {
|
|
ctx := t.Context()
|
|
model := &GoogleModel{}
|
|
apiKey := "test-api-key"
|
|
listErr := errors.New("list models failed")
|
|
|
|
withGoogleListModelsStub(t, func(context.Context, *genai.ClientConfig) ([]ListModelResponse, error) {
|
|
return nil, listErr
|
|
})
|
|
|
|
err := model.CheckConnection(ctx, &APIConfig{ApiKey: &apiKey})
|
|
if !errors.Is(err, listErr) {
|
|
t.Fatalf("expected ListModels error %v, got %v", listErr, err)
|
|
}
|
|
}
|
|
|
|
func TestGoogleModelChatStreamlyRequiresAPIKey(t *testing.T) {
|
|
ctx := t.Context()
|
|
model := &GoogleModel{}
|
|
messages := []Message{{Role: "user", Content: "hello"}}
|
|
cases := []struct {
|
|
name string
|
|
apiConfig *APIConfig
|
|
}{
|
|
{name: "nil config"},
|
|
{name: "nil api key", apiConfig: &APIConfig{}},
|
|
{name: "empty api key", apiConfig: &APIConfig{ApiKey: stringPtr("")}},
|
|
{name: "blank api key", apiConfig: &APIConfig{ApiKey: stringPtr(" \t\n ")}},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
err := model.ChatStreamlyWithSender(ctx, "gemini-2.5-flash", messages, tc.apiConfig, nil, nil, func(*string, *string) error {
|
|
t.Errorf("sender should not be called without an API key")
|
|
return nil
|
|
})
|
|
if err == nil {
|
|
t.Fatal("expected an API key error")
|
|
}
|
|
if !strings.Contains(err.Error(), "api key is required") {
|
|
t.Fatalf("expected API key error, got %v", err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGoogleModelChatRequiresModelName(t *testing.T) {
|
|
ctx := t.Context()
|
|
model := &GoogleModel{}
|
|
apiKey := "test-api-key"
|
|
messages := []Message{{Role: "user", Content: "hello"}}
|
|
|
|
response, err := model.ChatWithMessages(ctx, "", messages, &APIConfig{ApiKey: &apiKey}, nil, nil)
|
|
if err == nil {
|
|
t.Fatal("expected a model name error")
|
|
}
|
|
if !strings.Contains(err.Error(), "model name is empty") {
|
|
t.Fatalf("expected model name error, got %v", err)
|
|
}
|
|
if response != nil {
|
|
t.Fatalf("expected no response, got %v", response)
|
|
}
|
|
|
|
err = model.ChatStreamlyWithSender(ctx, "", messages, &APIConfig{ApiKey: &apiKey}, nil, nil, func(*string, *string) error {
|
|
t.Errorf("sender should not be called without a model name")
|
|
return nil
|
|
})
|
|
if err == nil {
|
|
t.Fatal("expected a model name error")
|
|
}
|
|
if !strings.Contains(err.Error(), "model name is empty") {
|
|
t.Fatalf("expected model name error, got %v", err)
|
|
}
|
|
|
|
err = model.ChatStreamlyWithSender(ctx, "gemini-2.5-flash", messages, &APIConfig{ApiKey: &apiKey}, nil, nil, nil)
|
|
if err == nil {
|
|
t.Fatal("expected a sender error")
|
|
}
|
|
if !strings.Contains(err.Error(), "sender is nil") {
|
|
t.Fatalf("expected sender error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestGoogleModelNewInstancePreservesCustomBaseURL(t *testing.T) {
|
|
model := NewGoogleModel(map[string]string{"default": "https://generativelanguage.googleapis.com"}, URLSuffix{Models: "v1beta/models"})
|
|
customBaseURL := map[string]string{"default": "https://example.test/google"}
|
|
|
|
instance := model.NewInstance(customBaseURL)
|
|
google, ok := instance.(*GoogleModel)
|
|
if !ok {
|
|
t.Fatalf("expected *GoogleModel, got %T", instance)
|
|
}
|
|
if google.baseModel.BaseURL["default"] != customBaseURL["default"] {
|
|
t.Fatalf("expected base URL %q, got %q", customBaseURL["default"], google.baseModel.BaseURL["default"])
|
|
}
|
|
if google.baseModel.URLSuffix != model.baseModel.URLSuffix {
|
|
t.Fatalf("expected URL suffix %v, got %v", model.baseModel.URLSuffix, google.baseModel.URLSuffix)
|
|
}
|
|
}
|
|
|
|
func TestGoogleModelListModelsPassesBaseURL(t *testing.T) {
|
|
apiKey := "test-api-key"
|
|
cases := []struct {
|
|
name string
|
|
baseURL map[string]string
|
|
region *string
|
|
expectedBaseURL string
|
|
}{
|
|
{
|
|
name: "default custom base URL",
|
|
baseURL: map[string]string{"default": "https://example.test/google"},
|
|
expectedBaseURL: "https://example.test/google",
|
|
},
|
|
{
|
|
name: "regional custom base URL",
|
|
baseURL: map[string]string{"east": "https://east.example.test/google", "default": "https://default.example.test/google"},
|
|
region: stringPtr("east"),
|
|
expectedBaseURL: "https://east.example.test/google",
|
|
},
|
|
{
|
|
name: "empty region custom base URL",
|
|
baseURL: map[string]string{"": "https://empty-region.example.test/google"},
|
|
region: stringPtr(""),
|
|
expectedBaseURL: "https://empty-region.example.test/google",
|
|
},
|
|
{
|
|
name: "missing region falls back to default base URL",
|
|
baseURL: map[string]string{"default": "https://default.example.test/google"},
|
|
region: stringPtr("missing"),
|
|
expectedBaseURL: "https://default.example.test/google",
|
|
},
|
|
{
|
|
name: "SDK default base URL",
|
|
},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
ctx := t.Context()
|
|
model := NewGoogleModel(tc.baseURL, URLSuffix{})
|
|
withGoogleListModelsStub(t, func(_ context.Context, config *genai.ClientConfig) ([]ListModelResponse, error) {
|
|
if config.HTTPOptions.BaseURL != tc.expectedBaseURL {
|
|
t.Fatalf("expected base URL %q, got %q", tc.expectedBaseURL, config.HTTPOptions.BaseURL)
|
|
}
|
|
return []ListModelResponse{{Name: "models/gemini-2.5-flash"}}, nil
|
|
})
|
|
|
|
if _, err := model.ListModels(ctx, &APIConfig{ApiKey: &apiKey, Region: tc.region}); err != nil {
|
|
t.Fatalf("expected no error, got %v", err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCollectGoogleModelNamesPaginates(t *testing.T) {
|
|
pages := []googleModelPage{
|
|
{items: []ModelListItem{{ID: "Gemini 2.5 Flash", OwnedBy: "Google"}}, nextPageToken: "page-2"},
|
|
{items: []ModelListItem{{ID: "Gemini 2.5 Pro", OwnedBy: "Google"}}, nextPageToken: ""},
|
|
}
|
|
var pageTokens []string
|
|
|
|
models, err := collectGoogleModelNames(context.Background(), func(_ context.Context, pageToken string) (googleModelPage, error) {
|
|
pageTokens = append(pageTokens, pageToken)
|
|
if len(pageTokens) > len(pages) {
|
|
t.Fatalf("unexpected extra page request with token %q", pageToken)
|
|
}
|
|
return pages[len(pageTokens)-1], nil
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("expected no error, got %v", err)
|
|
}
|
|
|
|
expectedModels := []ListModelResponse{{Name: "Gemini 2.5 Flash@Google"}, {Name: "Gemini 2.5 Pro@Google"}}
|
|
if !reflect.DeepEqual(models, expectedModels) {
|
|
t.Fatalf("expected models %v, got %v", expectedModels, models)
|
|
}
|
|
expectedPageTokens := []string{"", "page-2"}
|
|
if !reflect.DeepEqual(pageTokens, expectedPageTokens) {
|
|
t.Fatalf("expected page tokens %v, got %v", expectedPageTokens, pageTokens)
|
|
}
|
|
}
|
|
|
|
func TestCollectGoogleModelNamesPreservesEmptyResult(t *testing.T) {
|
|
models, err := collectGoogleModelNames(context.Background(), func(context.Context, string) (googleModelPage, error) {
|
|
return googleModelPage{}, nil
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("expected no error, got %v", err)
|
|
}
|
|
if models != nil {
|
|
t.Fatalf("expected nil models, got %v", models)
|
|
}
|
|
}
|
|
|
|
func TestCollectGoogleModelNamesReturnsPageError(t *testing.T) {
|
|
pageErr := errors.New("next page failed")
|
|
calls := 0
|
|
|
|
_, err := collectGoogleModelNames(context.Background(), func(context.Context, string) (googleModelPage, error) {
|
|
calls++
|
|
if calls == 1 {
|
|
return googleModelPage{items: []ModelListItem{{ID: "Gemini 2.5 Flash", OwnedBy: "Google"}}, nextPageToken: "page-2"}, nil
|
|
}
|
|
return googleModelPage{}, pageErr
|
|
})
|
|
if !errors.Is(err, pageErr) {
|
|
t.Fatalf("expected page error %v, got %v", pageErr, err)
|
|
}
|
|
}
|
|
|
|
func TestGoogleGenerateContentConfigConvertsTools(t *testing.T) {
|
|
toolChoice := "required"
|
|
cfg := googleGenerateContentConfig(&ChatConfig{
|
|
Tools: []map[string]interface{}{{
|
|
"type": "function",
|
|
"function": map[string]interface{}{
|
|
"name": "search_my_dataset",
|
|
"description": "Search dataset.",
|
|
"parameters": map[string]interface{}{
|
|
"type": "object",
|
|
"properties": map[string]interface{}{
|
|
"query": map[string]interface{}{"type": "string"},
|
|
},
|
|
"required": []string{"query"},
|
|
},
|
|
},
|
|
}},
|
|
ToolChoice: &toolChoice,
|
|
})
|
|
if cfg == nil || len(cfg.Tools) != 1 || len(cfg.Tools[0].FunctionDeclarations) != 1 {
|
|
t.Fatalf("tools = %#v, want one function declaration", cfg)
|
|
}
|
|
declaration := cfg.Tools[0].FunctionDeclarations[0]
|
|
if declaration.Name != "search_my_dataset" || declaration.Description != "Search dataset." {
|
|
t.Fatalf("declaration = %#v", declaration)
|
|
}
|
|
if declaration.ParametersJsonSchema == nil {
|
|
t.Fatal("ParametersJsonSchema is nil")
|
|
}
|
|
if cfg.ToolConfig == nil || cfg.ToolConfig.FunctionCallingConfig == nil {
|
|
t.Fatalf("ToolConfig = %#v", cfg.ToolConfig)
|
|
}
|
|
if cfg.ToolConfig.FunctionCallingConfig.Mode != genai.FunctionCallingConfigModeAny {
|
|
t.Fatalf("mode = %s, want ANY", cfg.ToolConfig.FunctionCallingConfig.Mode)
|
|
}
|
|
}
|
|
|
|
func TestGoogleChatContentsConvertsToolHistory(t *testing.T) {
|
|
contents := googleChatContents([]Message{
|
|
{
|
|
Role: "assistant",
|
|
Content: nil,
|
|
ToolCalls: []map[string]interface{}{{
|
|
"id": "call-1",
|
|
"type": "function",
|
|
"function": map[string]interface{}{
|
|
"name": "search_my_dataset",
|
|
"arguments": `{"query":"marigold"}`,
|
|
},
|
|
}},
|
|
},
|
|
{Role: "tool", ToolCallID: "call-1", Content: "flower result"},
|
|
})
|
|
if len(contents) != 2 {
|
|
t.Fatalf("contents len = %d, want 2", len(contents))
|
|
}
|
|
functionCall := contents[0].Parts[0].FunctionCall
|
|
if functionCall == nil || functionCall.ID != "call-1" || functionCall.Name != "search_my_dataset" {
|
|
t.Fatalf("function call = %#v", functionCall)
|
|
}
|
|
if functionCall.Args["query"] != "marigold" {
|
|
t.Fatalf("args = %#v", functionCall.Args)
|
|
}
|
|
functionResponse := contents[1].Parts[0].FunctionResponse
|
|
if functionResponse == nil || functionResponse.ID != "call-1" || functionResponse.Name != "search_my_dataset" {
|
|
t.Fatalf("function response = %#v", functionResponse)
|
|
}
|
|
if functionResponse.Response["output"] != "flower result" {
|
|
t.Fatalf("response = %#v", functionResponse.Response)
|
|
}
|
|
}
|
|
|
|
func TestGoogleToolCallsConvertsFunctionCalls(t *testing.T) {
|
|
toolCalls := googleToolCalls([]*genai.FunctionCall{{
|
|
ID: "call-1",
|
|
Name: "search_my_dataset",
|
|
Args: map[string]any{"query": "marigold"},
|
|
}})
|
|
if len(toolCalls) != 1 {
|
|
t.Fatalf("tool calls len = %d, want 1", len(toolCalls))
|
|
}
|
|
if toolCalls[0]["id"] != "call-1" || toolCalls[0]["type"] != "function" {
|
|
t.Fatalf("tool call = %#v", toolCalls[0])
|
|
}
|
|
function, _ := toolCalls[0]["function"].(map[string]interface{})
|
|
if function["name"] != "search_my_dataset" {
|
|
t.Fatalf("function = %#v", function)
|
|
}
|
|
var args map[string]interface{}
|
|
if err := json.Unmarshal([]byte(function["arguments"].(string)), &args); err != nil {
|
|
t.Fatalf("arguments JSON: %v", err)
|
|
}
|
|
if args["query"] != "marigold" {
|
|
t.Fatalf("arguments = %#v", args)
|
|
}
|
|
}
|
|
|
|
func stringPtr(value string) *string {
|
|
return &value
|
|
}
|