github.com/weaviate/weaviate@v1.24.6/test/modules/many-modules/many_modules_openai_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 test
    13  
    14  import (
    15  	"encoding/json"
    16  	"testing"
    17  
    18  	"github.com/stretchr/testify/assert"
    19  	"github.com/stretchr/testify/require"
    20  	"github.com/weaviate/weaviate/entities/models"
    21  	"github.com/weaviate/weaviate/entities/schema"
    22  	"github.com/weaviate/weaviate/test/helper"
    23  )
    24  
    25  func createSchemaOpenAISanityChecks(endpoint string) func(t *testing.T) {
    26  	return func(t *testing.T) {
    27  		vectorizer := "text2vec-openai"
    28  		class := &models.Class{
    29  			Class: "OpenAI",
    30  			Properties: []*models.Property{
    31  				{
    32  					Name:     "text",
    33  					DataType: []string{schema.DataTypeText.String()},
    34  				},
    35  			},
    36  			Vectorizer: vectorizer,
    37  		}
    38  		tests := []struct {
    39  			name                  string
    40  			text2vecOpenAI        map[string]interface{}
    41  			expectDefaultSettings bool
    42  		}{
    43  			{
    44  				name: "model: text-embedding-3-large, dimensions: 256, vectorizeClassName: false",
    45  				text2vecOpenAI: map[string]interface{}{
    46  					"vectorizeClassName": false,
    47  					"model":              "text-embedding-3-large",
    48  					"dimensions":         256,
    49  				},
    50  			},
    51  			{
    52  				name: "model: text-embedding-3-large, dimensions: 1024, vectorizeClassName: false",
    53  				text2vecOpenAI: map[string]interface{}{
    54  					"vectorizeClassName": false,
    55  					"model":              "text-embedding-3-large",
    56  					"dimensions":         1024,
    57  				},
    58  			},
    59  			{
    60  				name: "model: text-embedding-3-large, dimensions: 3072, vectorizeClassName: false",
    61  				text2vecOpenAI: map[string]interface{}{
    62  					"vectorizeClassName": false,
    63  					"model":              "text-embedding-3-large",
    64  					"dimensions":         3072,
    65  				},
    66  			},
    67  			{
    68  				name: "model: text-embedding-3-small, dimensions: 512, vectorizeClassName: true",
    69  				text2vecOpenAI: map[string]interface{}{
    70  					"vectorizeClassName": true,
    71  					"model":              "text-embedding-3-small",
    72  					"dimensions":         512,
    73  				},
    74  			},
    75  			{
    76  				name: "model: text-embedding-3-small, dimensions: 1536, vectorizeClassName: false",
    77  				text2vecOpenAI: map[string]interface{}{
    78  					"vectorizeClassName": false,
    79  					"model":              "text-embedding-3-small",
    80  					"dimensions":         1536,
    81  				},
    82  			},
    83  			{
    84  				name: "model: text-embedding-3-small, dimensions: 1536, vectorizeClassName: true",
    85  				text2vecOpenAI: map[string]interface{}{
    86  					"vectorizeClassName": true,
    87  					"model":              "text-embedding-3-small",
    88  					"dimensions":         1536,
    89  				},
    90  			},
    91  			{
    92  				name: "model: ada, vectorizeClassName: false",
    93  				text2vecOpenAI: map[string]interface{}{
    94  					"vectorizeClassName": false,
    95  					"model":              "ada",
    96  				},
    97  			},
    98  			{
    99  				name:                  "nil settings",
   100  				text2vecOpenAI:        nil,
   101  				expectDefaultSettings: true,
   102  			},
   103  			{
   104  				name:                  "empty settings",
   105  				text2vecOpenAI:        map[string]interface{}{},
   106  				expectDefaultSettings: true,
   107  			},
   108  		}
   109  		for _, tt := range tests {
   110  			t.Run(tt.name, func(t *testing.T) {
   111  				class.ModuleConfig = map[string]interface{}{
   112  					vectorizer: tt.text2vecOpenAI,
   113  				}
   114  				helper.CreateClass(t, class)
   115  				defer helper.DeleteClass(t, class.Class)
   116  				verifyClass := helper.GetClass(t, class.Class)
   117  				moduleConfig, ok := verifyClass.ModuleConfig.(map[string]interface{})
   118  				require.True(t, ok)
   119  				require.NotEmpty(t, moduleConfig)
   120  				text2vecOpenAI, ok := moduleConfig[vectorizer].(map[string]interface{})
   121  				require.True(t, ok)
   122  				require.NotEmpty(t, text2vecOpenAI)
   123  				if tt.expectDefaultSettings {
   124  					assert.Equal(t, "ada", text2vecOpenAI["model"])
   125  					assert.Equal(t, true, text2vecOpenAI["vectorizeClassName"])
   126  				} else {
   127  					assert.Equal(t, tt.text2vecOpenAI["model"], text2vecOpenAI["model"])
   128  					assert.Equal(t, tt.text2vecOpenAI["vectorizeClassName"], text2vecOpenAI["vectorizeClassName"])
   129  					expectedDimensions, ok := tt.text2vecOpenAI["dimensions"]
   130  					if ok {
   131  						dimensions, ok := text2vecOpenAI["dimensions"].(json.Number)
   132  						require.True(t, ok)
   133  						dimensionsInt64, err := dimensions.Int64()
   134  						require.NoError(t, err)
   135  						assert.Equal(t, expectedDimensions, int(dimensionsInt64))
   136  					}
   137  				}
   138  			})
   139  		}
   140  	}
   141  }