github.com/weaviate/weaviate@v1.24.6/modules/generative-anyscale/config/class_settings_test.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package config
    13  
    14  import (
    15  	"errors"
    16  	"testing"
    17  
    18  	"github.com/stretchr/testify/assert"
    19  	"github.com/weaviate/weaviate/entities/moduletools"
    20  )
    21  
    22  func Test_classSettings_Validate(t *testing.T) {
    23  	tests := []struct {
    24  		name            string
    25  		cfg             moduletools.ClassConfig
    26  		wantModel       string
    27  		wantTemperature int
    28  		wantBaseURL     string
    29  		wantErr         error
    30  	}{
    31  		{
    32  			name: "default settings",
    33  			cfg: fakeClassConfig{
    34  				classConfig: map[string]interface{}{},
    35  			},
    36  			wantModel:       "meta-llama/Llama-2-70b-chat-hf",
    37  			wantTemperature: 0,
    38  			wantBaseURL:     "https://api.endpoints.anyscale.com",
    39  			wantErr:         nil,
    40  		},
    41  		{
    42  			name: "everything non default configured",
    43  			cfg: fakeClassConfig{
    44  				classConfig: map[string]interface{}{
    45  					"model":       "meta-llama/Llama-2-70b-chat-hf",
    46  					"temperature": 1,
    47  				},
    48  			},
    49  			wantModel:       "meta-llama/Llama-2-70b-chat-hf",
    50  			wantTemperature: 1,
    51  			wantBaseURL:     "https://api.endpoints.anyscale.com",
    52  			wantErr:         nil,
    53  		},
    54  		{
    55  			name: "everything non default configured and base url",
    56  			cfg: fakeClassConfig{
    57  				classConfig: map[string]interface{}{
    58  					"model":       "meta-llama/Llama-2-70b-chat-hf",
    59  					"temperature": 1,
    60  					"baseURL":     "https://custom.endpoint.com",
    61  				},
    62  			},
    63  			wantModel:       "meta-llama/Llama-2-70b-chat-hf",
    64  			wantTemperature: 1,
    65  			wantBaseURL:     "https://custom.endpoint.com",
    66  			wantErr:         nil,
    67  		},
    68  		{
    69  			name: "unsupported model",
    70  			cfg: fakeClassConfig{
    71  				classConfig: map[string]interface{}{
    72  					"model":       "unsupported",
    73  					"temperature": 1,
    74  					"baseURL":     "https://custom.endpoint.com",
    75  				},
    76  			},
    77  			wantErr: errors.New("wrong Anyscale model name, available model names are: [meta-llama/Llama-2-70b-chat-hf meta-llama/Llama-2-13b-chat-hf meta-llama/Llama-2-7b-chat-hf codellama/CodeLlama-34b-Instruct-hf mistralai/Mistral-7B-Instruct-v0.1 mistralai/Mixtral-8x7B-Instruct-v0.1]"),
    78  		},
    79  	}
    80  	for _, tt := range tests {
    81  		t.Run(tt.name, func(t *testing.T) {
    82  			ic := NewClassSettings(tt.cfg)
    83  			if tt.wantErr != nil {
    84  				assert.Equal(t, tt.wantErr.Error(), ic.Validate(nil).Error())
    85  			} else {
    86  				assert.Equal(t, tt.wantBaseURL, ic.BaseURL())
    87  				assert.Equal(t, tt.wantModel, ic.Model())
    88  				assert.Equal(t, tt.wantTemperature, ic.Temperature())
    89  			}
    90  		})
    91  	}
    92  }
    93  
    94  type fakeClassConfig struct {
    95  	classConfig map[string]interface{}
    96  }
    97  
    98  func (f fakeClassConfig) Class() map[string]interface{} {
    99  	return f.classConfig
   100  }
   101  
   102  func (f fakeClassConfig) Tenant() string {
   103  	return ""
   104  }
   105  
   106  func (f fakeClassConfig) ClassByModuleName(moduleName string) map[string]interface{} {
   107  	return f.classConfig
   108  }
   109  
   110  func (f fakeClassConfig) Property(propName string) map[string]interface{} {
   111  	return nil
   112  }
   113  
   114  func (f fakeClassConfig) TargetVector() string {
   115  	return ""
   116  }