admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-06 17:18:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

enhance: refactor embedding credentials manager (#41442)

Browse Source 
https://github.com/milvus-io/milvus/issues/35856

Signed-off-by: junjie.jiang <junjie.jiang@zilliz.com>
This commit is contained in:
junjiejiangjjj 2025-04-24 14:34:38 +08:00 committed by GitHub
parent dbe54c2df8
commit e56adc121b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
34 changed files with 733 additions and 294 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
cmd/tools/config/generate.go
View File

@ -359,6 +359,20 @@ func WriteYaml(w io.Writer) {
name: "knowhere",
header: `
# Any configuration related to the knowhere vector search engine`,
},
{
name: "credential",
header: `
# credential configs, support apikey, AKSK, gcp credential
# examples:
# credential:
# your_apikey_crendential_name:
# apikey: # Your apikey credential
# your_aksk_crendential_name:
# access_key_id:
# secret_access_key:
# your_gcp_credential_name:
# credential_json:`,
},
{
name: "function",

44
configs/milvus.yaml
View File

@ -1219,32 +1219,50 @@ knowhere:
search:
beam_width_ratio: 4 # Ratio between the maximum number of IO requests per search iteration and CPU number
# credential configs, support apikey, AKSK, gcp credential
# examples:
# credential:
# your_apikey_crendential_name:
# apikey: # Your apikey credential
# your_aksk_crendential_name:
# access_key_id:
# secret_access_key:
# your_gcp_credential_name:
# credential_json:
credential:
aksk1:
access_key_id: # Your access_key_id
secret_access_key: # Your secret_access_key
apikey1:
apikey: # Your apikey credential
gcp1:
credential_json: # base64 based gcp credential data
# Any configuration related to functions
function:
textEmbedding:
enableVerifiInfoInParams: true # Controls whether to allow configuration of apikey and model service url on function parameters
providers:
azure_openai:
api_key: # Your azure openai embedding url, Default is the official embedding url
credential: # The name in the crendential configuration item
resource_name: # Your azure openai resource name
url: # Your azure openai api key
url: # Your azure openai embedding url, Default is the official embedding url
bedrock:
aws_access_key_id: # Your aws_access_key_id
aws_secret_access_key: # Your aws_secret_access_key
credential: # The name in the crendential configuration item
cohere:
api_key: # Your cohere embedding url, Default is the official embedding url
url: # Your cohere api key
credential: # The name in the crendential configuration item
url: # Your cohere embedding url, Default is the official embedding url
dashscope:
api_key: # Your dashscope embedding url, Default is the official embedding url
url: # Your dashscope api key
credential: # The name in the crendential configuration item
url: # Your dashscope embedding url, Default is the official embedding url
openai:
api_key: # Your openai embedding url, Default is the official embedding url
url: # Your openai api key
credential: # The name in the crendential configuration item
url: # Your openai embedding url, Default is the official embedding url
siliconflow:
api_key: # Your siliconflow api key
credential: # The name in the crendential configuration item
url: # Your siliconflow embedding url, Default is the official embedding url
tei:
credential: # The name in the crendential configuration item
enable: true # Whether to enable TEI model service
vertexai:
credentials_file_path: # Path to your google application credentials, change the file path to refresh the configuration
credential: # The name in the crendential configuration item
url: # Your VertexAI embedding url

15
internal/datanode/importv2/scheduler_test.go
View File

@ -437,6 +437,13 @@ func (s *SchedulerSuite) TestScheduler_ImportFile() {
}
func (s *SchedulerSuite) TestScheduler_ImportFileWithFunction() {
paramtable.Init()
paramtable.Get().CredentialCfg.Credential.GetFunc = func() map[string]string {
return map[string]string{
"mock.apikey": "mock",
}
}
s.syncMgr.EXPECT().SyncData(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, task syncmgr.Task, callbacks ...func(error) error) (*conc.Future[struct{}], error) {
future := conc.Go(func() (struct{}, error) {
return struct{}{}, nil
@ -445,6 +452,11 @@ func (s *SchedulerSuite) TestScheduler_ImportFileWithFunction() {
})
ts := function.CreateOpenAIEmbeddingServer()
defer ts.Close()
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
return map[string]string{
"openai.url": ts.URL,
}
}
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{
@ -484,8 +496,7 @@ func (s *SchedulerSuite) TestScheduler_ImportFileWithFunction() {
Params: []*commonpb.KeyValuePair{
{Key: "provider", Value: "openai"},
{Key: "model_name", Value: "text-embedding-ada-002"},
{Key: "api_key", Value: "mock"},
{Key: "url", Value: ts.URL},
{Key: "credential", Value: "mock"},
{Key: "dim", Value: "4"},
},
},

15
internal/proxy/task_insert_test.go
View File

@ -314,8 +314,20 @@ func TestMaxInsertSize(t *testing.T) {
}
func TestInsertTask_Function(t *testing.T) {
paramtable.Init()
paramtable.Get().CredentialCfg.Credential.GetFunc = func() map[string]string {
return map[string]string{
"mock.apikey": "mock",
}
}
ts := function.CreateOpenAIEmbeddingServer()
defer ts.Close()
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
return map[string]string{
"openai.url": ts.URL,
}
}
data := []*schemapb.FieldData{}
f := schemapb.FieldData{
Type: schemapb.DataType_VarChar,
@ -365,8 +377,7 @@ func TestInsertTask_Function(t *testing.T) {
Params: []*commonpb.KeyValuePair{
{Key: "provider", Value: "openai"},
{Key: "model_name", Value: "text-embedding-ada-002"},
{Key: "api_key", Value: "mock"},
{Key: "url", Value: ts.URL},
{Key: "credential", Value: "mock"},
{Key: "dim", Value: "4"},
},
},

17
internal/proxy/task_search_test.go
View File

@ -974,8 +974,19 @@ func TestSearchTask_PreExecute(t *testing.T) {
}
func TestSearchTask_WithFunctions(t *testing.T) {
paramtable.Init()
paramtable.Get().CredentialCfg.Credential.GetFunc = func() map[string]string {
return map[string]string{
"mock.apikey": "mock",
}
}
ts := function.CreateOpenAIEmbeddingServer()
defer ts.Close()
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
return map[string]string{
"openai.url": ts.URL,
}
}
collectionName := "TestSearchTask_function"
schema := &schemapb.CollectionSchema{
Name: collectionName,
@ -1016,8 +1027,7 @@ func TestSearchTask_WithFunctions(t *testing.T) {
Params: []*commonpb.KeyValuePair{
{Key: "provider", Value: "openai"},
{Key: "model_name", Value: "text-embedding-ada-002"},
{Key: "api_key", Value: "mock"},
{Key: "url", Value: ts.URL},
{Key: "credential", Value: "mock"},
{Key: "dim", Value: "4"},
},
},
@ -1031,8 +1041,7 @@ func TestSearchTask_WithFunctions(t *testing.T) {
Params: []*commonpb.KeyValuePair{
{Key: "provider", Value: "openai"},
{Key: "model_name", Value: "text-embedding-ada-002"},
{Key: "api_key", Value: "mock"},
{Key: "url", Value: ts.URL},
{Key: "credential", Value: "mock"},
{Key: "dim", Value: "4"},
},
},

14
internal/proxy/task_test.go
View File

@ -1299,8 +1299,19 @@ func TestCreateCollectionTask(t *testing.T) {
})
t.Run("collection with embedding function ", func(t *testing.T) {
paramtable.Init()
paramtable.Get().CredentialCfg.Credential.GetFunc = func() map[string]string {
return map[string]string{
"mock.apikey": "mock",
}
}
ts := function.CreateOpenAIEmbeddingServer()
defer ts.Close()
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
return map[string]string{
"openai.url": ts.URL,
}
}
schema.Functions = []*schemapb.FunctionSchema{
{
Name: "test",
@ -1310,8 +1321,7 @@ func TestCreateCollectionTask(t *testing.T) {
Params: []*commonpb.KeyValuePair{
{Key: "provider", Value: "openai"},
{Key: "model_name", Value: "text-embedding-ada-002"},
{Key: "api_key", Value: "mock"},
{Key: "url", Value: ts.URL},
{Key: "credential", Value: "mock"},
{Key: "dim", Value: "128"},
},
},

16
internal/proxy/task_upsert_test.go
View File

@ -34,6 +34,7 @@ import (
"github.com/milvus-io/milvus/pkg/v2/proto/rootcoordpb"
"github.com/milvus-io/milvus/pkg/v2/util/commonpbutil"
"github.com/milvus-io/milvus/pkg/v2/util/merr"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
"github.com/milvus-io/milvus/pkg/v2/util/testutils"
)
@ -367,8 +368,20 @@ func TestUpsertTaskForReplicate(t *testing.T) {
}
func TestUpsertTask_Function(t *testing.T) {
paramtable.Init()
paramtable.Get().CredentialCfg.Credential.GetFunc = func() map[string]string {
return map[string]string{
"mock.apikey": "mock",
}
}
ts := function.CreateOpenAIEmbeddingServer()
defer ts.Close()
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
return map[string]string{
"openai.url": ts.URL,
}
}
data := []*schemapb.FieldData{}
f1 := schemapb.FieldData{
Type: schemapb.DataType_Int64,
@ -434,8 +447,7 @@ func TestUpsertTask_Function(t *testing.T) {
Params: []*commonpb.KeyValuePair{
{Key: "provider", Value: "openai"},
{Key: "model_name", Value: "text-embedding-ada-002"},
{Key: "api_key", Value: "mock"},
{Key: "url", Value: ts.URL},
{Key: "credential", Value: "mock"},
{Key: "dim", Value: "4"},
},
},

14
internal/proxy/util_test.go
View File

@ -2851,8 +2851,19 @@ func TestValidateFunction(t *testing.T) {
func TestValidateModelFunction(t *testing.T) {
t.Run("Valid model function schema", func(t *testing.T) {
paramtable.Init()
paramtable.Get().CredentialCfg.Credential.GetFunc = func() map[string]string {
return map[string]string{
"mock.apikey": "mock",
}
}
ts := function.CreateOpenAIEmbeddingServer()
defer ts.Close()
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
return map[string]string{
"openai.url": ts.URL,
}
}
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{Name: "input_field", DataType: schemapb.DataType_VarChar, TypeParams: []*commonpb.KeyValuePair{{Key: "enable_analyzer", Value: "true"}}},
@ -2879,8 +2890,7 @@ func TestValidateModelFunction(t *testing.T) {
Params: []*commonpb.KeyValuePair{
{Key: "provider", Value: "openai"},
{Key: "model_name", Value: "text-embedding-ada-002"},
{Key: "api_key", Value: "mock"},
{Key: "url", Value: ts.URL},
{Key: "credential", Value: "mock"},
{Key: "dim", Value: "4"},
},
},

84
internal/util/credentials/credentials_manager.go Normal file
View File

@ -0,0 +1,84 @@
/*
* # Licensed to the LF AI & Data foundation under one
* # or more contributor license agreements. See the NOTICE file
* # distributed with this work for additional information
* # regarding copyright ownership. The ASF licenses this file
* # to you under the Apache License, Version 2.0 (the
* # "License"); you may not use this file except in compliance
* # with the License. You may obtain a copy of the License at
* #
* # http://www.apache.org/licenses/LICENSE-2.0
* #
* # Unless required by applicable law or agreed to in writing, software
* # distributed under the License is distributed on an "AS IS" BASIS,
* # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* # See the License for the specific language governing permissions and
* # limitations under the License.
*/
package credentials
import (
"encoding/base64"
"fmt"
)
const (
APIKey string = "apikey"
AccessKeyId string = "access_key_id"
SecretAccessKey string = "secret_access_key"
// #nosec G101
CredentialJSON string = "credential_json"
)
// The current version only supports plain text, and cipher text will be supported later.
type CredentialsManager struct {
// key formats:
// {credentialName}.api_key
// {credentialName}.access_key_id
// {credentialName}.secret_access_key
// {credentialName}.credential_json
confMap map[string]string
}
func NewCredentialsManager(conf map[string]string) *CredentialsManager {
return &CredentialsManager{conf}
}
func (c *CredentialsManager) GetAPIKeyCredential(name string) (string, error) {
k := name + "." + APIKey
apikey, exist := c.confMap[k]
if !exist {
return "", fmt.Errorf("%s is not a apikey crediential, can not find key: %s", name, k)
}
return apikey, nil
}
func (c *CredentialsManager) GetAKSKCredential(name string) (string, string, error) {
IdKey := name + "." + AccessKeyId
accessKeyId, exist := c.confMap[IdKey]
if !exist {
return "", "", fmt.Errorf("%s is not a aksk crediential, can not find key: %s", name, IdKey)
}
AccessKey := name + "." + SecretAccessKey
secretAccessKey, exist := c.confMap[AccessKey]
if !exist {
return "", "", fmt.Errorf("%s is not a aksk crediential, can not find key: %s", name, AccessKey)
}
return accessKeyId, secretAccessKey, nil
}
func (c *CredentialsManager) GetGcpCredential(name string) ([]byte, error) {
k := name + "." + CredentialJSON
jsonByte, exist := c.confMap[k]
if !exist {
return nil, fmt.Errorf("%s is not a gcp crediential, can not find key: %s ", name, k)
}
decode, err := base64.StdEncoding.DecodeString(jsonByte)
if err != nil {
return nil, fmt.Errorf("Parse gcp credential:%s faild, err: %s", name, err)
}
return decode, nil
}

10
internal/util/function/ali_embedding_provider.go
View File

@ -23,6 +23,7 @@ import (
"strings"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/ali"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -41,7 +42,7 @@ type AliEmbeddingProvider struct {
func createAliEmbeddingClient(apiKey string, url string) (*ali.AliDashScopeEmbedding, error) {
if apiKey == "" {
return nil, fmt.Errorf("Missing credentials. Please pass `api_key`, or configure the %s environment variable in the Milvus service.", dashscopeAKEnvStr)
return nil, fmt.Errorf("Missing credentials config or configure the %s environment variable in the Milvus service.", dashscopeAKEnvStr)
}
if url == "" {
@ -51,12 +52,15 @@ func createAliEmbeddingClient(apiKey string, url string) (*ali.AliDashScopeEmbed
return c, nil
}
func NewAliDashScopeEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string) (*AliEmbeddingProvider, error) {
func NewAliDashScopeEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string, credentials *credentials.CredentialsManager) (*AliEmbeddingProvider, error) {
fieldDim, err := typeutil.GetDim(fieldSchema)
if err != nil {
return nil, err
}
apiKey, url := parseAKAndURL(functionSchema.Params, params, dashscopeAKEnvStr)
apiKey, url, err := parseAKAndURL(credentials, functionSchema.Params, params, dashscopeAKEnvStr)
if err != nil {
return nil, err
}
var modelName string
var dim int64

11
internal/util/function/alitext_embedding_provider_test.go
View File

@ -29,6 +29,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/ali"
)
@ -69,14 +70,13 @@ func createAliProvider(url string, schema *schemapb.FieldSchema, providerName st
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: modelNameParamKey, Value: TestModel},
{Key: embeddingURLParamKey, Value: url},
{Key: apiKeyParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
{Key: dimParamKey, Value: "4"},
},
}
switch providerName {
case aliDashScopeProvider:
return NewAliDashScopeEmbeddingProvider(schema, functionSchema, map[string]string{})
return NewAliDashScopeEmbeddingProvider(schema, functionSchema, map[string]string{embeddingURLParamKey: url}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
default:
return nil, errors.New("Unknow provider")
}
@ -181,12 +181,11 @@ func (s *AliTextEmbeddingProviderSuite) TestNewAliDashScopeEmbeddingProvider() {
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: modelNameParamKey, Value: TestModel},
{Key: apiKeyParamKey, Value: "mock"},
{Key: dimParamKey, Value: "4"},
},
}
// invalid dim
functionSchema.Params[2] = &commonpb.KeyValuePair{Key: dimParamKey, Value: "Invalid"}
_, err := NewAliDashScopeEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
functionSchema.Params[1] = &commonpb.KeyValuePair{Key: dimParamKey, Value: "Invalid"}
_, err := NewAliDashScopeEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.Error(err)
}

45
internal/util/function/bedrock_embedding_provider.go
View File

@ -33,6 +33,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
milvusCredentials "github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -54,10 +55,10 @@ type BedrockEmbeddingProvider struct {
func createBedRockEmbeddingClient(awsAccessKeyId string, awsSecretAccessKey string, region string) (*bedrockruntime.Client, error) {
if awsAccessKeyId == "" {
return nil, fmt.Errorf("Missing credentials. Please pass `aws_access_key_id`, or configure the %s environment variable in the Milvus service.", bedrockAccessKeyId)
return nil, fmt.Errorf("Missing credentials config or configure the %s environment variable in the Milvus service.", bedrockAccessKeyId)
}
if awsSecretAccessKey == "" {
return nil, fmt.Errorf("Missing credentials. Please pass `aws_secret_access_key`, or configure the %s environment variable in the Milvus service.", bedrockSAKEnvStr)
return nil, fmt.Errorf("Missing credentials config or configure the %s environment variable in the Milvus service.", bedrockSAKEnvStr)
}
if region == "" {
return nil, errors.New("Missing AWS Service region. Please pass `region` param.")
@ -74,28 +75,29 @@ func createBedRockEmbeddingClient(awsAccessKeyId string, awsSecretAccessKey stri
return bedrockruntime.NewFromConfig(cfg), nil
}
func parseAccessInfo(params []*commonpb.KeyValuePair, confParams map[string]string) (string, string) {
// function param > env > yaml
func parseAKSKInfo(credentials *milvusCredentials.CredentialsManager, params []*commonpb.KeyValuePair, confParams map[string]string) (string, string, error) {
// function param > yaml > env
var awsAccessKeyId, awsSecretAccessKey string
var err error
// from function params
if isEnableVerifiInfoInParamsKey(confParams) {
for _, param := range params {
switch strings.ToLower(param.Key) {
case awsAKIdParamKey:
awsAccessKeyId = param.Value
case awsSAKParamKey:
awsSecretAccessKey = param.Value
for _, param := range params {
switch strings.ToLower(param.Key) {
case credentialParamKey:
credentialName := param.Value
if awsAccessKeyId, awsSecretAccessKey, err = credentials.GetAKSKCredential(credentialName); err != nil {
return "", "", err
}
}
}
// from milvus.yaml
if awsAccessKeyId == "" {
awsAccessKeyId = confParams[awsAKIdParamKey]
}
if awsSecretAccessKey == "" {
awsSecretAccessKey = confParams[awsSAKParamKey]
if awsAccessKeyId == "" && awsSecretAccessKey == "" {
credentialName := confParams[credentialParamKey]
if credentialName != "" {
if awsAccessKeyId, awsSecretAccessKey, err = credentials.GetAKSKCredential(credentialName); err != nil {
return "", "", err
}
}
}
// from env
@ -106,10 +108,10 @@ func parseAccessInfo(params []*commonpb.KeyValuePair, confParams map[string]stri
awsSecretAccessKey = os.Getenv(bedrockSAKEnvStr)
}
return awsAccessKeyId, awsSecretAccessKey
return awsAccessKeyId, awsSecretAccessKey, nil
}
func NewBedrockEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, c BedrockClient, params map[string]string) (*BedrockEmbeddingProvider, error) {
func NewBedrockEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, c BedrockClient, params map[string]string, credentials *milvusCredentials.CredentialsManager) (*BedrockEmbeddingProvider, error) {
fieldDim, err := typeutil.GetDim(fieldSchema)
if err != nil {
return nil, err
@ -142,7 +144,10 @@ func NewBedrockEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSche
}
}
awsAccessKeyId, awsSecretAccessKey := parseAccessInfo(functionSchema.Params, params)
awsAccessKeyId, awsSecretAccessKey, err := parseAKSKInfo(credentials, functionSchema.Params, params)
if err != nil {
return nil, err
}
var client BedrockClient
if c == nil {

54
internal/util/function/bedrock_text_embedding_provider_test.go
View File

@ -19,6 +19,7 @@
package function
import (
"os"
"testing"
"github.com/cockroachdb/errors"
@ -26,6 +27,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
)
func TestBedrockTextEmbeddingProvider(t *testing.T) {
@ -65,13 +67,12 @@ func createBedrockProvider(schema *schemapb.FieldSchema, providerName string, di
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: modelNameParamKey, Value: TestModel},
{Key: apiKeyParamKey, Value: "mock"},
{Key: dimParamKey, Value: "4"},
},
}
switch providerName {
case bedrockProvider:
return NewBedrockEmbeddingProvider(schema, functionSchema, &MockBedrockClient{dim: dim}, map[string]string{})
return NewBedrockEmbeddingProvider(schema, functionSchema, &MockBedrockClient{dim: dim}, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.access_key_id": "mock", "mock.secret_access_key": "mock"}))
default:
return nil, errors.New("Unknow provider")
}
@ -110,6 +111,38 @@ func (s *BedrockTextEmbeddingProviderSuite) TestEmbeddingDimNotMatch() {
}
}
func (s *BedrockTextEmbeddingProviderSuite) TestParseCredentail() {
{
cred := credentials.NewCredentialsManager(map[string]string{})
ak, sk, err := parseAKSKInfo(cred, []*commonpb.KeyValuePair{}, map[string]string{})
s.Equal(ak, "")
s.Equal(sk, "")
s.NoError(err)
}
{
cred := credentials.NewCredentialsManager(map[string]string{})
_, _, err := parseAKSKInfo(cred, []*commonpb.KeyValuePair{}, map[string]string{"credential": "NotExist"})
s.ErrorContains(err, "is not a aksk crediential, can not find key")
}
{
cred := credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"})
_, _, err := parseAKSKInfo(cred, []*commonpb.KeyValuePair{}, map[string]string{"credential": "mock"})
s.ErrorContains(err, "is not a aksk crediential, can not find key")
}
{
cred := credentials.NewCredentialsManager(map[string]string{"mock.access_key_id": "mock", "mock.secret_access_key": "mock"})
_, _, err := parseAKSKInfo(cred, []*commonpb.KeyValuePair{}, map[string]string{"credential": "mock"})
s.NoError(err)
}
{
os.Setenv(bedrockAccessKeyId, "mock")
os.Setenv(bedrockSAKEnvStr, "mock")
cred := credentials.NewCredentialsManager(map[string]string{})
_, _, err := parseAKSKInfo(cred, []*commonpb.KeyValuePair{}, map[string]string{})
s.NoError(err)
}
}
func (s *BedrockTextEmbeddingProviderSuite) TestCreateBedrockClient() {
_, err := createBedRockEmbeddingClient("", "", "")
s.Error(err)
@ -144,32 +177,31 @@ func (s *BedrockTextEmbeddingProviderSuite) TestNewBedrockEmbeddingProvider() {
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: modelNameParamKey, Value: TestModel},
{Key: awsAKIdParamKey, Value: "mock"},
{Key: awsSAKParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
{Key: regionParamKey, Value: "mock"},
{Key: dimParamKey, Value: "4"},
{Key: normalizeParamKey, Value: "false"},
},
}
provider, err := NewBedrockEmbeddingProvider(fieldSchema, functionSchema, nil, map[string]string{})
provider, err := NewBedrockEmbeddingProvider(fieldSchema, functionSchema, nil, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.access_key_id": "mock", "mock.secret_access_key": "mock"}))
s.NoError(err)
s.True(provider.MaxBatch() > 0)
s.Equal(provider.FieldDim(), int64(4))
_, err = NewBedrockEmbeddingProvider(fieldSchema, functionSchema, nil, map[string]string{awsAKIdParamKey: "mock", awsSAKParamKey: "mock"})
_, err = NewBedrockEmbeddingProvider(fieldSchema, functionSchema, nil, map[string]string{credentialParamKey: "mock"}, credentials.NewCredentialsManager(map[string]string{"mock.access_key_id": "mock", "mock.secret_access_key": "mock"}))
s.NoError(err)
functionSchema.Params[5] = &commonpb.KeyValuePair{Key: normalizeParamKey, Value: "true"}
_, err = NewBedrockEmbeddingProvider(fieldSchema, functionSchema, nil, map[string]string{})
functionSchema.Params[4] = &commonpb.KeyValuePair{Key: normalizeParamKey, Value: "true"}
_, err = NewBedrockEmbeddingProvider(fieldSchema, functionSchema, nil, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.access_key_id": "mock", "mock.secret_access_key": "mock"}))
s.NoError(err)
functionSchema.Params[5] = &commonpb.KeyValuePair{Key: normalizeParamKey, Value: "invalid"}
_, err = NewBedrockEmbeddingProvider(fieldSchema, functionSchema, nil, map[string]string{})
functionSchema.Params[4] = &commonpb.KeyValuePair{Key: normalizeParamKey, Value: "invalid"}
_, err = NewBedrockEmbeddingProvider(fieldSchema, functionSchema, nil, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.access_key_id": "mock", "mock.secret_access_key": "mock"}))
s.Error(err)
// invalid dim
functionSchema.Params[0] = &commonpb.KeyValuePair{Key: modelNameParamKey, Value: TestModel}
functionSchema.Params[0] = &commonpb.KeyValuePair{Key: dimParamKey, Value: "Invalid"}
_, err = NewBedrockEmbeddingProvider(fieldSchema, functionSchema, nil, map[string]string{})
_, err = NewBedrockEmbeddingProvider(fieldSchema, functionSchema, nil, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.access_key_id": "mock", "mock.secret_access_key": "mock"}))
s.Error(err)
}

10
internal/util/function/cohere_embedding_provider.go
View File

@ -23,6 +23,7 @@ import (
"strings"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/cohere"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -42,7 +43,7 @@ type CohereEmbeddingProvider struct {
func createCohereEmbeddingClient(apiKey string, url string) (*cohere.CohereEmbedding, error) {
if apiKey == "" {
return nil, fmt.Errorf("Missing credentials. Please pass `api_key`, or configure the %s environment variable in the Milvus service.", cohereAIAKEnvStr)
return nil, fmt.Errorf("Missing credentials config or configure the %s environment variable in the Milvus service.", cohereAIAKEnvStr)
}
if url == "" {
@ -53,12 +54,15 @@ func createCohereEmbeddingClient(apiKey string, url string) (*cohere.CohereEmbed
return c, nil
}
func NewCohereEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string) (*CohereEmbeddingProvider, error) {
func NewCohereEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string, credentials *credentials.CredentialsManager) (*CohereEmbeddingProvider, error) {
fieldDim, err := typeutil.GetDim(fieldSchema)
if err != nil {
return nil, err
}
apiKey, url := parseAKAndURL(functionSchema.Params, params, cohereAIAKEnvStr)
apiKey, url, err := parseAKAndURL(credentials, functionSchema.Params, params, cohereAIAKEnvStr)
if err != nil {
return nil, err
}
var modelName string
truncate := "END"
for _, param := range functionSchema.Params {

27
internal/util/function/cohere_embedding_provider_test.go
View File

@ -29,7 +29,9 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/cohere"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
)
func TestCohereTextEmbeddingProvider(t *testing.T) {
@ -43,6 +45,12 @@ type CohereTextEmbeddingProviderSuite struct {
}
func (s *CohereTextEmbeddingProviderSuite) SetupTest() {
paramtable.Init()
paramtable.Get().CredentialCfg.Credential.GetFunc = func() map[string]string {
return map[string]string{
"mock.apikey": "mock",
}
}
s.schema = &schemapb.CollectionSchema{
Name: "test",
Fields: []*schemapb.FieldSchema{
@ -69,13 +77,12 @@ func createCohereProvider(url string, schema *schemapb.FieldSchema, providerName
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: modelNameParamKey, Value: TestModel},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: url},
{Key: credentialParamKey, Value: "mock"},
},
}
switch providerName {
case cohereProvider:
return NewCohereEmbeddingProvider(schema, functionSchema, map[string]string{})
return NewCohereEmbeddingProvider(schema, functionSchema, map[string]string{embeddingURLParamKey: url}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
default:
return nil, errors.New("Unknow provider")
}
@ -259,22 +266,22 @@ func (s *CohereTextEmbeddingProviderSuite) TestNewCohereProvider() {
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: modelNameParamKey, Value: TestModel},
{Key: apiKeyParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
},
}
provider, err := NewCohereEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
provider, err := NewCohereEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.NoError(err)
s.Equal(provider.truncate, "END")
functionSchema.Params = append(functionSchema.Params, &commonpb.KeyValuePair{Key: truncateParamKey, Value: "START"})
provider, err = NewCohereEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
provider, err = NewCohereEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.NoError(err)
s.Equal(provider.truncate, "START")
// Invalid truncateParam
functionSchema.Params[2].Value = "Unknow"
_, err = NewCohereEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
_, err = NewCohereEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.Error(err)
}
@ -288,17 +295,17 @@ func (s *CohereTextEmbeddingProviderSuite) TestGetInputType() {
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: modelNameParamKey, Value: "model-v2.0"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
},
}
provider, err := NewCohereEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
provider, err := NewCohereEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.NoError(err)
s.Equal(provider.getInputType(InsertMode), "")
s.Equal(provider.getInputType(SearchMode), "")
functionSchema.Params[0].Value = "model-v3.0"
provider, err = NewCohereEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
provider, err = NewCohereEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.NoError(err)
s.Equal(provider.getInputType(InsertMode), "search_document")
s.Equal(provider.getInputType(SearchMode), "search_query")

60
internal/util/function/common.go
View File

@ -26,6 +26,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
)
type TextEmbeddingMode int
@ -45,12 +46,11 @@ const (
// common params
const (
modelNameParamKey string = "model_name"
dimParamKey string = "dim"
embeddingURLParamKey string = "url"
apiKeyParamKey string = "api_key"
truncateParamKey string = "truncate"
enableVerifiInfoInParamsKey string = "enableVerifiInfoInParams"
modelNameParamKey string = "model_name"
dimParamKey string = "dim"
embeddingURLParamKey string = "url"
credentialParamKey string = "credential"
truncateParamKey string = "truncate"
)
// ali text embedding
@ -72,8 +72,8 @@ const (
// bedrock emebdding
const (
awsAKIdParamKey string = "aws_access_key_id"
awsSAKParamKey string = "aws_secret_access_key"
// awsAKIdParamKey string = "aws_access_key_id"
// awsSAKParamKey string = "aws_secret_access_key"
regionParamKey string = "region"
normalizeParamKey string = "normalize"
@ -121,41 +121,29 @@ const (
enableTeiEnvStr string = "MILVUSAI_ENABLE_TEI"
)
const enableVerifiInfoInParams string = "ENABLE_VERIFI_INFO_IN_PARAMS"
func isEnableVerifiInfoInParamsKey(confParams map[string]string) bool {
enable := true
if strings.ToLower(confParams[enableVerifiInfoInParamsKey]) != "" {
// If enableVerifiInfoInParamsKey is configured in milvus.yaml, the configuration in milvus.yaml will be used.
enable, _ = strconv.ParseBool(confParams[enableVerifiInfoInParamsKey])
} else {
// If enableVerifiInfoInParamsKey is not configured in milvus.yaml, the configuration in env will be used.
if strings.ToLower(os.Getenv(enableVerifiInfoInParams)) != "" {
enable, _ = strconv.ParseBool(confParams[enableVerifiInfoInParamsKey])
}
}
return enable
}
func parseAKAndURL(params []*commonpb.KeyValuePair, confParams map[string]string, apiKeyEnv string) (string, string) {
// function param > env > yaml
func parseAKAndURL(credentials *credentials.CredentialsManager, params []*commonpb.KeyValuePair, confParams map[string]string, apiKeyEnv string) (string, string, error) {
// function param > yaml > env
var err error
var apiKey, url string
// from function params
if isEnableVerifiInfoInParamsKey(confParams) {
for _, param := range params {
switch strings.ToLower(param.Key) {
case apiKeyParamKey:
apiKey = param.Value
case embeddingURLParamKey:
url = param.Value
for _, param := range params {
switch strings.ToLower(param.Key) {
case credentialParamKey:
credentialName := param.Value
if apiKey, err = credentials.GetAPIKeyCredential(credentialName); err != nil {
return "", "", err
}
}
}
// from milvus.yaml
if apiKey == "" {
apiKey = confParams[apiKeyParamKey]
credentialName := confParams[credentialParamKey]
if credentialName != "" {
if apiKey, err = credentials.GetAPIKeyCredential(credentialName); err != nil {
return "", "", err
}
}
}
if url == "" {
@ -166,7 +154,7 @@ func parseAKAndURL(params []*commonpb.KeyValuePair, confParams map[string]string
if apiKey == "" {
apiKey = os.Getenv(apiKeyEnv)
}
return apiKey, url
return apiKey, url, nil
}
func parseAndCheckFieldDim(dimStr string, fieldDim int64, fieldName string) (int64, error) {

18
internal/util/function/function_executor_test.go
View File

@ -49,9 +49,21 @@ type FunctionExecutorSuite struct {
func (s *FunctionExecutorSuite) SetupTest() {
paramtable.Init()
paramtable.Get().CredentialCfg.Credential.GetFunc = func() map[string]string {
return map[string]string{
"mock.apikey": "mock",
}
}
}
func (s *FunctionExecutorSuite) creataSchema(url string) *schemapb.CollectionSchema {
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
key := openAIProvider + "." + embeddingURLParamKey
return map[string]string{
key: url,
}
}
return &schemapb.CollectionSchema{
Name: "test",
Fields: []*schemapb.FieldSchema{
@ -83,8 +95,7 @@ func (s *FunctionExecutorSuite) creataSchema(url string) *schemapb.CollectionSch
Params: []*commonpb.KeyValuePair{
{Key: Provider, Value: openAIProvider},
{Key: modelNameParamKey, Value: "text-embedding-ada-002"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: url},
{Key: credentialParamKey, Value: "mock"},
{Key: dimParamKey, Value: "4"},
},
},
@ -98,8 +109,7 @@ func (s *FunctionExecutorSuite) creataSchema(url string) *schemapb.CollectionSch
Params: []*commonpb.KeyValuePair{
{Key: Provider, Value: openAIProvider},
{Key: modelNameParamKey, Value: "text-embedding-ada-002"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: url},
{Key: credentialParamKey, Value: "mock"},
{Key: dimParamKey, Value: "8"},
},
},

25
internal/util/function/openai_embedding_provider.go
View File

@ -24,6 +24,7 @@ import (
"strings"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/openai"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -42,7 +43,7 @@ type OpenAIEmbeddingProvider struct {
func createOpenAIEmbeddingClient(apiKey string, url string) (*openai.OpenAIEmbeddingClient, error) {
if apiKey == "" {
return nil, fmt.Errorf("Missing credentials. Please pass `api_key`, or configure the %s environment variable in the Milvus service.", openaiAKEnvStr)
return nil, fmt.Errorf("Missing credentials config or configure the %s environment variable in the Milvus service.", openaiAKEnvStr)
}
if url == "" {
@ -55,7 +56,7 @@ func createOpenAIEmbeddingClient(apiKey string, url string) (*openai.OpenAIEmbed
func createAzureOpenAIEmbeddingClient(apiKey string, url string, resourceName string) (*openai.AzureOpenAIEmbeddingClient, error) {
if apiKey == "" {
return nil, fmt.Errorf("Missing credentials. Please pass `api_key`, or configure the %s environment variable in the Milvus service", azureOpenaiAKEnvStr)
return nil, fmt.Errorf("Missing credentials config or configure the %s environment variable in the Milvus service", azureOpenaiAKEnvStr)
}
if url == "" {
@ -73,7 +74,7 @@ func createAzureOpenAIEmbeddingClient(apiKey string, url string, resourceName st
return c, nil
}
func newOpenAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string, isAzure bool) (*OpenAIEmbeddingProvider, error) {
func newOpenAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string, isAzure bool, credentials *credentials.CredentialsManager) (*OpenAIEmbeddingProvider, error) {
fieldDim, err := typeutil.GetDim(fieldSchema)
if err != nil {
return nil, err
@ -98,13 +99,19 @@ func newOpenAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchem
var c openai.OpenAIEmbeddingInterface
if !isAzure {
apiKey, url := parseAKAndURL(functionSchema.Params, params, openaiAKEnvStr)
apiKey, url, err := parseAKAndURL(credentials, functionSchema.Params, params, openaiAKEnvStr)
if err != nil {
return nil, err
}
c, err = createOpenAIEmbeddingClient(apiKey, url)
if err != nil {
return nil, err
}
} else {
apiKey, url := parseAKAndURL(functionSchema.Params, params, azureOpenaiAKEnvStr)
apiKey, url, err := parseAKAndURL(credentials, functionSchema.Params, params, azureOpenaiAKEnvStr)
if err != nil {
return nil, err
}
resourceName := params["resource_name"]
c, err = createAzureOpenAIEmbeddingClient(apiKey, url, resourceName)
if err != nil {
@ -124,12 +131,12 @@ func newOpenAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchem
return &provider, nil
}
func NewOpenAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string) (*OpenAIEmbeddingProvider, error) {
return newOpenAIEmbeddingProvider(fieldSchema, functionSchema, params, false)
func NewOpenAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string, credentials *credentials.CredentialsManager) (*OpenAIEmbeddingProvider, error) {
return newOpenAIEmbeddingProvider(fieldSchema, functionSchema, params, false, credentials)
}
func NewAzureOpenAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string) (*OpenAIEmbeddingProvider, error) {
return newOpenAIEmbeddingProvider(fieldSchema, functionSchema, params, true)
func NewAzureOpenAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string, credentials *credentials.CredentialsManager) (*OpenAIEmbeddingProvider, error) {
return newOpenAIEmbeddingProvider(fieldSchema, functionSchema, params, true, credentials)
}
func (provider *OpenAIEmbeddingProvider) MaxBatch() int {

8
internal/util/function/openai_text_embedding_provider_test.go
View File

@ -30,6 +30,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/openai"
)
@ -70,16 +71,15 @@ func createOpenAIProvider(url string, schema *schemapb.FieldSchema, providerName
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: modelNameParamKey, Value: "text-embedding-ada-002"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
{Key: dimParamKey, Value: "4"},
{Key: embeddingURLParamKey, Value: url},
},
}
switch providerName {
case openAIProvider:
return NewOpenAIEmbeddingProvider(schema, functionSchema, map[string]string{})
return NewOpenAIEmbeddingProvider(schema, functionSchema, map[string]string{embeddingURLParamKey: url}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
case azureOpenAIProvider:
return NewAzureOpenAIEmbeddingProvider(schema, functionSchema, map[string]string{})
return NewAzureOpenAIEmbeddingProvider(schema, functionSchema, map[string]string{embeddingURLParamKey: url}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
default:
return nil, errors.New("Unknow provider")
}

10
internal/util/function/siliconflow_embedding_provider.go
View File

@ -23,6 +23,7 @@ import (
"strings"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/siliconflow"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -40,7 +41,7 @@ type SiliconflowEmbeddingProvider struct {
func createSiliconflowEmbeddingClient(apiKey string, url string) (*siliconflow.SiliconflowEmbedding, error) {
if apiKey == "" {
return nil, fmt.Errorf("Missing credentials. Please pass `api_key`, or configure the %s environment variable in the Milvus service.", siliconflowAKEnvStr)
return nil, fmt.Errorf("Missing credentials conifg or configure the %s environment variable in the Milvus service.", siliconflowAKEnvStr)
}
if url == "" {
@ -51,12 +52,15 @@ func createSiliconflowEmbeddingClient(apiKey string, url string) (*siliconflow.S
return c, nil
}
func NewSiliconflowEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string) (*SiliconflowEmbeddingProvider, error) {
func NewSiliconflowEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string, credentials *credentials.CredentialsManager) (*SiliconflowEmbeddingProvider, error) {
fieldDim, err := typeutil.GetDim(fieldSchema)
if err != nil {
return nil, err
}
apiKey, url := parseAKAndURL(functionSchema.Params, params, siliconflowAKEnvStr)
apiKey, url, err := parseAKAndURL(credentials, functionSchema.Params, params, siliconflowAKEnvStr)
if err != nil {
return nil, err
}
var modelName string
for _, param := range functionSchema.Params {

11
internal/util/function/siliconflow_embedding_provider_test.go
View File

@ -29,6 +29,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/siliconflow"
)
@ -69,13 +70,12 @@ func createSiliconflowProvider(url string, schema *schemapb.FieldSchema, provide
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: modelNameParamKey, Value: TestModel},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: url},
{Key: credentialParamKey, Value: "mock"},
},
}
switch providerName {
case siliconflowProvider:
return NewSiliconflowEmbeddingProvider(schema, functionSchema, map[string]string{})
return NewSiliconflowEmbeddingProvider(schema, functionSchema, map[string]string{embeddingURLParamKey: url}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
default:
return nil, errors.New("Unknow provider")
}
@ -182,11 +182,10 @@ func (s *SiliconflowTextEmbeddingProviderSuite) TestNewSiliconflowEmbeddingProvi
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: modelNameParamKey, Value: TestModel},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
},
}
provider, err := NewSiliconflowEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
provider, err := NewSiliconflowEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{embeddingURLParamKey: "mock"}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.NoError(err)
s.Equal(provider.FieldDim(), int64(4))
s.True(provider.MaxBatch() > 0)

11
internal/util/function/tei_embedding_provider.go
View File

@ -27,6 +27,7 @@ import (
"github.com/cockroachdb/errors"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/tei"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -54,12 +55,12 @@ func createTEIEmbeddingClient(apiKey string, endpoint string) (*tei.TEIEmbedding
return tei.NewTEIEmbeddingClient(apiKey, endpoint)
}
func NewTEIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string) (*TeiEmbeddingProvider, error) {
func NewTEIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string, credentials *credentials.CredentialsManager) (*TeiEmbeddingProvider, error) {
fieldDim, err := typeutil.GetDim(fieldSchema)
if err != nil {
return nil, err
}
var apiKey, endpoint, ingestionPrompt, searchPrompt string
var endpoint, ingestionPrompt, searchPrompt string
// TEI default client batch size
maxBatch := 32
truncate := false
@ -68,8 +69,6 @@ func NewTEIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *
for _, param := range functionSchema.Params {
switch strings.ToLower(param.Key) {
case apiKeyParamKey:
apiKey = param.Value
case endpointParamKey:
endpoint = param.Value
case ingestionPromptParamKey:
@ -92,6 +91,10 @@ func NewTEIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *
}
}
apiKey, _, err := parseAKAndURL(credentials, functionSchema.Params, params, "")
if err != nil {
return nil, err
}
c, err := createTEIEmbeddingClient(apiKey, endpoint)
if err != nil {
return nil, err

19
internal/util/function/tei_embedding_provider_test.go
View File

@ -30,6 +30,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
)
func TestTEITextEmbeddingProvider(t *testing.T) {
@ -68,7 +69,7 @@ func createTEIProvider(url string, schema *schemapb.FieldSchema, providerName st
InputFieldIds: []int64{101},
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: apiKeyParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
{Key: endpointParamKey, Value: url},
{Key: ingestionPromptParamKey, Value: "doc:"},
{Key: searchPromptParamKey, Value: "query:"},
@ -76,7 +77,7 @@ func createTEIProvider(url string, schema *schemapb.FieldSchema, providerName st
}
switch providerName {
case teiProvider:
return NewTEIEmbeddingProvider(schema, functionSchema, map[string]string{})
return NewTEIEmbeddingProvider(schema, functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
default:
return nil, errors.New("Unknow provider")
}
@ -168,11 +169,11 @@ func (s *TEITextEmbeddingProviderSuite) TestNewTEIEmbeddingProvider() {
InputFieldIds: []int64{101},
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: apiKeyParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
{Key: endpointParamKey, Value: "http://mymock.com"},
},
}
provider, err := NewTEIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
provider, err := NewTEIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.NoError(err)
s.Equal(provider.FieldDim(), int64(4))
s.True(provider.MaxBatch() == 32*5)
@ -180,35 +181,35 @@ func (s *TEITextEmbeddingProviderSuite) TestNewTEIEmbeddingProvider() {
// Invalid truncate
{
functionSchema.Params = append(functionSchema.Params, &commonpb.KeyValuePair{Key: truncateParamKey, Value: "Invalid"})
_, err := NewTEIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
_, err := NewTEIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.Error(err)
}
// Invalid truncationDirection
{
functionSchema.Params[2] = &commonpb.KeyValuePair{Key: truncateParamKey, Value: "true"}
functionSchema.Params = append(functionSchema.Params, &commonpb.KeyValuePair{Key: truncationDirectionParamKey, Value: "Invalid"})
_, err := NewTEIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
_, err := NewTEIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.Error(err)
}
// truncationDirection
{
functionSchema.Params[3] = &commonpb.KeyValuePair{Key: truncationDirectionParamKey, Value: "Left"}
_, err := NewTEIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
_, err := NewTEIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.NoError(err)
}
// Invalid max batch
{
functionSchema.Params = append(functionSchema.Params, &commonpb.KeyValuePair{Key: maxClientBatchSizeParamKey, Value: "Invalid"})
_, err := NewTEIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
_, err := NewTEIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.Error(err)
}
// Valid max batch
{
functionSchema.Params[4] = &commonpb.KeyValuePair{Key: maxClientBatchSizeParamKey, Value: "128"}
pv, err := NewTEIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
pv, err := NewTEIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.NoError(err)
s.True(pv.MaxBatch() == 128*5)
}

20
internal/util/function/text_embedding_function.go
View File

@ -28,6 +28,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/pkg/v2/util/funcutil"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
)
@ -98,25 +99,26 @@ func NewTextEmbeddingFunction(coll *schemapb.CollectionSchema, functionSchema *s
var embP textEmbeddingProvider
var newProviderErr error
conf := paramtable.Get().FunctionCfg.GetTextEmbeddingProviderConfig(base.provider)
credentials := credentials.NewCredentialsManager(paramtable.Get().CredentialCfg.GetCredentials())
switch base.provider {
case openAIProvider:
embP, newProviderErr = NewOpenAIEmbeddingProvider(base.outputFields[0], functionSchema, conf)
embP, newProviderErr = NewOpenAIEmbeddingProvider(base.outputFields[0], functionSchema, conf, credentials)
case azureOpenAIProvider:
embP, newProviderErr = NewAzureOpenAIEmbeddingProvider(base.outputFields[0], functionSchema, conf)
embP, newProviderErr = NewAzureOpenAIEmbeddingProvider(base.outputFields[0], functionSchema, conf, credentials)
case bedrockProvider:
embP, newProviderErr = NewBedrockEmbeddingProvider(base.outputFields[0], functionSchema, nil, conf)
embP, newProviderErr = NewBedrockEmbeddingProvider(base.outputFields[0], functionSchema, nil, conf, credentials)
case aliDashScopeProvider:
embP, newProviderErr = NewAliDashScopeEmbeddingProvider(base.outputFields[0], functionSchema, conf)
embP, newProviderErr = NewAliDashScopeEmbeddingProvider(base.outputFields[0], functionSchema, conf, credentials)
case vertexAIProvider:
embP, newProviderErr = NewVertexAIEmbeddingProvider(base.outputFields[0], functionSchema, nil, conf)
embP, newProviderErr = NewVertexAIEmbeddingProvider(base.outputFields[0], functionSchema, nil, conf, credentials)
case voyageAIProvider:
embP, newProviderErr = NewVoyageAIEmbeddingProvider(base.outputFields[0], functionSchema, conf)
embP, newProviderErr = NewVoyageAIEmbeddingProvider(base.outputFields[0], functionSchema, conf, credentials)
case cohereProvider:
embP, newProviderErr = NewCohereEmbeddingProvider(base.outputFields[0], functionSchema, conf)
embP, newProviderErr = NewCohereEmbeddingProvider(base.outputFields[0], functionSchema, conf, credentials)
case siliconflowProvider:
embP, newProviderErr = NewSiliconflowEmbeddingProvider(base.outputFields[0], functionSchema, conf)
embP, newProviderErr = NewSiliconflowEmbeddingProvider(base.outputFields[0], functionSchema, conf, credentials)
case teiProvider:
embP, newProviderErr = NewTEIEmbeddingProvider(base.outputFields[0], functionSchema, conf)
embP, newProviderErr = NewTEIEmbeddingProvider(base.outputFields[0], functionSchema, conf, credentials)
default:
return nil, fmt.Errorf("Unsupported text embedding service provider: [%s] , list of supported [%s, %s, %s, %s, %s, %s, %s, %s, %s]", base.provider, openAIProvider, azureOpenAIProvider, aliDashScopeProvider, bedrockProvider, vertexAIProvider, voyageAIProvider, cohereProvider, siliconflowProvider, teiProvider)
}

127
internal/util/function/text_embedding_function_test.go
View File

@ -29,6 +29,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/testutil"
"github.com/milvus-io/milvus/pkg/v2/util/funcutil"
"github.com/milvus-io/milvus/pkg/v2/util/paramtable"
@ -45,6 +46,13 @@ type TextEmbeddingFunctionSuite struct {
func (s *TextEmbeddingFunctionSuite) SetupTest() {
paramtable.Init()
paramtable.Get().CredentialCfg.Credential.GetFunc = func() map[string]string {
return map[string]string{
"mock.apikey": "mock",
"mock.access_key_id": "mock",
"mock.secret_access_key": "mock",
}
}
s.schema = &schemapb.CollectionSchema{
Name: "test",
Fields: []*schemapb.FieldSchema{
@ -92,8 +100,7 @@ func (s *TextEmbeddingFunctionSuite) TestInvalidProvider() {
{Key: Provider, Value: openAIProvider},
{Key: modelNameParamKey, Value: "text-embedding-ada-002"},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
},
}
providerName, err := getProvider(fSchema)
@ -110,6 +117,12 @@ func (s *TextEmbeddingFunctionSuite) TestProcessInsert() {
ts := CreateOpenAIEmbeddingServer()
defer ts.Close()
{
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
key := openAIProvider + "." + embeddingURLParamKey
return map[string]string{
key: ts.URL,
}
}
runner, err := NewTextEmbeddingFunction(s.schema, &schemapb.FunctionSchema{
Name: "test",
Type: schemapb.FunctionType_TextEmbedding,
@ -121,8 +134,7 @@ func (s *TextEmbeddingFunctionSuite) TestProcessInsert() {
{Key: Provider, Value: openAIProvider},
{Key: modelNameParamKey, Value: "text-embedding-ada-002"},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: ts.URL},
{Key: credentialParamKey, Value: "mock"},
},
})
s.NoError(err)
@ -142,6 +154,12 @@ func (s *TextEmbeddingFunctionSuite) TestProcessInsert() {
}
}
{
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
key := azureOpenAIProvider + "." + embeddingURLParamKey
return map[string]string{
key: ts.URL,
}
}
runner, err := NewTextEmbeddingFunction(s.schema, &schemapb.FunctionSchema{
Name: "test",
Type: schemapb.FunctionType_TextEmbedding,
@ -153,8 +171,7 @@ func (s *TextEmbeddingFunctionSuite) TestProcessInsert() {
{Key: Provider, Value: azureOpenAIProvider},
{Key: modelNameParamKey, Value: "text-embedding-ada-002"},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: ts.URL},
{Key: credentialParamKey, Value: "mock"},
},
})
s.NoError(err)
@ -178,6 +195,12 @@ func (s *TextEmbeddingFunctionSuite) TestProcessInsert() {
func (s *TextEmbeddingFunctionSuite) TestAliEmbedding() {
ts := CreateAliEmbeddingServer()
defer ts.Close()
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
key := aliDashScopeProvider + "." + embeddingURLParamKey
return map[string]string{
key: ts.URL,
}
}
runner, err := NewTextEmbeddingFunction(s.schema, &schemapb.FunctionSchema{
Name: "test",
@ -190,8 +213,7 @@ func (s *TextEmbeddingFunctionSuite) TestAliEmbedding() {
{Key: Provider, Value: aliDashScopeProvider},
{Key: modelNameParamKey, Value: TestModel},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: ts.URL},
{Key: credentialParamKey, Value: "mock"},
},
})
s.NoError(err)
@ -336,8 +358,7 @@ func (s *TextEmbeddingFunctionSuite) TestRunnerParamsErr() {
{Key: Provider, Value: openAIProvider},
{Key: modelNameParamKey, Value: "text-embedding-ada-002"},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
},
})
s.Error(err)
@ -375,8 +396,7 @@ func (s *TextEmbeddingFunctionSuite) TestRunnerParamsErr() {
{Key: Provider, Value: openAIProvider},
{Key: modelNameParamKey, Value: "text-embedding-ada-002"},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
},
})
s.Error(err)
@ -395,8 +415,7 @@ func (s *TextEmbeddingFunctionSuite) TestRunnerParamsErr() {
{Key: Provider, Value: openAIProvider},
{Key: modelNameParamKey, Value: "text-embedding-ada-002"},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
},
})
s.Error(err)
@ -432,8 +451,7 @@ func (s *TextEmbeddingFunctionSuite) TestNewTextEmbeddings() {
Params: []*commonpb.KeyValuePair{
{Key: Provider, Value: bedrockProvider},
{Key: modelNameParamKey, Value: TestModel},
{Key: awsAKIdParamKey, Value: "mock"},
{Key: awsSAKParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
{Key: regionParamKey, Value: "mock"},
},
}
@ -456,7 +474,7 @@ func (s *TextEmbeddingFunctionSuite) TestNewTextEmbeddings() {
Params: []*commonpb.KeyValuePair{
{Key: Provider, Value: aliDashScopeProvider},
{Key: modelNameParamKey, Value: TestModel},
{Key: apiKeyParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
},
}
@ -478,7 +496,7 @@ func (s *TextEmbeddingFunctionSuite) TestNewTextEmbeddings() {
Params: []*commonpb.KeyValuePair{
{Key: Provider, Value: voyageAIProvider},
{Key: modelNameParamKey, Value: TestModel},
{Key: apiKeyParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
},
}
@ -500,7 +518,7 @@ func (s *TextEmbeddingFunctionSuite) TestNewTextEmbeddings() {
Params: []*commonpb.KeyValuePair{
{Key: Provider, Value: siliconflowProvider},
{Key: modelNameParamKey, Value: TestModel},
{Key: apiKeyParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
},
}
@ -522,7 +540,7 @@ func (s *TextEmbeddingFunctionSuite) TestNewTextEmbeddings() {
Params: []*commonpb.KeyValuePair{
{Key: Provider, Value: cohereProvider},
{Key: modelNameParamKey, Value: TestModel},
{Key: apiKeyParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
},
}
@ -607,6 +625,12 @@ func (s *TextEmbeddingFunctionSuite) TestNewTextEmbeddings() {
func (s *TextEmbeddingFunctionSuite) TestProcessSearchFloat32() {
ts := CreateOpenAIEmbeddingServer()
defer ts.Close()
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
key := openAIProvider + "." + embeddingURLParamKey
return map[string]string{
key: ts.URL,
}
}
runner, err := NewTextEmbeddingFunction(s.schema, &schemapb.FunctionSchema{
Name: "test",
Type: schemapb.FunctionType_TextEmbedding,
@ -618,8 +642,7 @@ func (s *TextEmbeddingFunctionSuite) TestProcessSearchFloat32() {
{Key: Provider, Value: openAIProvider},
{Key: modelNameParamKey, Value: "text-embedding-ada-002"},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: ts.URL},
{Key: credentialParamKey, Value: "mock"},
},
})
s.NoError(err)
@ -692,6 +715,12 @@ func (s *TextEmbeddingFunctionSuite) TestProcessInsertInt8() {
},
},
}
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
key := cohereProvider + "." + embeddingURLParamKey
return map[string]string{
key: ts.URL,
}
}
runner, err := NewTextEmbeddingFunction(s.schema, &schemapb.FunctionSchema{
Name: "test",
Type: schemapb.FunctionType_TextEmbedding,
@ -703,8 +732,7 @@ func (s *TextEmbeddingFunctionSuite) TestProcessInsertInt8() {
{Key: Provider, Value: cohereProvider},
{Key: modelNameParamKey, Value: TestModel},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: ts.URL},
{Key: credentialParamKey, Value: "mock"},
},
})
s.NoError(err)
@ -750,8 +778,8 @@ func (s *TextEmbeddingFunctionSuite) TestUnsupportedVec() {
{Key: Provider, Value: cohereProvider},
{Key: modelNameParamKey, Value: TestModel},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
// {Key: embeddingURLParamKey, Value: "mock"},
},
})
s.Error(err)
@ -774,6 +802,12 @@ func (s *TextEmbeddingFunctionSuite) TestProcessSearchInt8() {
},
},
}
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
key := cohereProvider + "." + embeddingURLParamKey
return map[string]string{
key: ts.URL,
}
}
runner, err := NewTextEmbeddingFunction(s.schema, &schemapb.FunctionSchema{
Name: "test",
Type: schemapb.FunctionType_TextEmbedding,
@ -785,8 +819,7 @@ func (s *TextEmbeddingFunctionSuite) TestProcessSearchInt8() {
{Key: Provider, Value: cohereProvider},
{Key: modelNameParamKey, Value: TestModel},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: ts.URL},
{Key: credentialParamKey, Value: "mock"},
},
})
s.NoError(err)
@ -845,6 +878,12 @@ func (s *TextEmbeddingFunctionSuite) TestProcessSearchInt8() {
func (s *TextEmbeddingFunctionSuite) TestProcessBulkInsertFloat32() {
ts := CreateOpenAIEmbeddingServer()
defer ts.Close()
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
key := openAIProvider + "." + embeddingURLParamKey
return map[string]string{
key: ts.URL,
}
}
runner, err := NewTextEmbeddingFunction(s.schema, &schemapb.FunctionSchema{
Name: "test",
Type: schemapb.FunctionType_TextEmbedding,
@ -856,8 +895,7 @@ func (s *TextEmbeddingFunctionSuite) TestProcessBulkInsertFloat32() {
{Key: Provider, Value: openAIProvider},
{Key: modelNameParamKey, Value: "text-embedding-ada-002"},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: ts.URL},
{Key: credentialParamKey, Value: "mock"},
},
})
s.NoError(err)
@ -894,6 +932,26 @@ func (s *TextEmbeddingFunctionSuite) TestProcessBulkInsertFloat32() {
}
}
func (s *TextEmbeddingFunctionSuite) TestParseCredentail() {
{
cred := credentials.NewCredentialsManager(map[string]string{})
ak, url, err := parseAKAndURL(cred, []*commonpb.KeyValuePair{}, map[string]string{}, "")
s.Equal(ak, "")
s.Equal(url, "")
s.NoError(err)
}
{
cred := credentials.NewCredentialsManager(map[string]string{})
_, _, err := parseAKAndURL(cred, []*commonpb.KeyValuePair{}, map[string]string{"credential": "NotExist"}, "")
s.ErrorContains(err, "is not a apikey crediential, can not find key")
}
{
cred := credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"})
_, _, err := parseAKAndURL(cred, []*commonpb.KeyValuePair{}, map[string]string{"credential": "mock"}, "")
s.NoError(err)
}
}
func (s *TextEmbeddingFunctionSuite) TestProcessBulkInsertInt8() {
ts := CreateCohereEmbeddingServer[int8]()
defer ts.Close()
@ -910,6 +968,12 @@ func (s *TextEmbeddingFunctionSuite) TestProcessBulkInsertInt8() {
},
},
}
paramtable.Get().FunctionCfg.TextEmbeddingProviders.GetFunc = func() map[string]string {
key := cohereProvider + "." + embeddingURLParamKey
return map[string]string{
key: ts.URL,
}
}
runner, err := NewTextEmbeddingFunction(s.schema, &schemapb.FunctionSchema{
Name: "test",
Type: schemapb.FunctionType_TextEmbedding,
@ -921,8 +985,7 @@ func (s *TextEmbeddingFunctionSuite) TestProcessBulkInsertInt8() {
{Key: Provider, Value: cohereProvider},
{Key: modelNameParamKey, Value: TestModel},
{Key: dimParamKey, Value: "4"},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: ts.URL},
{Key: credentialParamKey, Value: "mock"},
},
})
s.NoError(err)

65
internal/util/function/vertexai_embedding_provider.go
View File

@ -26,7 +26,9 @@ import (
"github.com/cockroachdb/errors"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/vertexai"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -39,22 +41,15 @@ type vertexAIJsonKey struct {
var vtxKey vertexAIJsonKey
func getVertexAIJsonKey(credentialsFilePath string) ([]byte, error) {
func getVertexAIJsonKey() ([]byte, error) {
vtxKey.mu.Lock()
defer vtxKey.mu.Unlock()
var jsonKeyPath string
if credentialsFilePath == "" {
jsonKeyPath = os.Getenv(vertexServiceAccountJSONEnv)
} else {
jsonKeyPath = credentialsFilePath
}
jsonKeyPath := os.Getenv(vertexServiceAccountJSONEnv)
if jsonKeyPath == "" {
return nil, errors.New("VetexAI credentials file path is empty")
}
if vtxKey.filePath == jsonKeyPath {
// The file path remains unchanged, using the data in the cache
return vtxKey.jsonKey, nil
}
@ -65,6 +60,7 @@ func getVertexAIJsonKey(credentialsFilePath string) ([]byte, error) {
vtxKey.jsonKey = jsonKey
vtxKey.filePath = jsonKeyPath
return vtxKey.jsonKey, nil
}
@ -86,16 +82,47 @@ type VertexAIEmbeddingProvider struct {
timeoutSec int64
}
func createVertexAIEmbeddingClient(url string, credentialsFilePath string) (*vertexai.VertexAIEmbedding, error) {
jsonKey, err := getVertexAIJsonKey(credentialsFilePath)
if err != nil {
return nil, err
}
c := vertexai.NewVertexAIEmbedding(url, jsonKey, "https://www.googleapis.com/auth/cloud-platform", "")
func createVertexAIEmbeddingClient(url string, credentialsJSON []byte) (*vertexai.VertexAIEmbedding, error) {
c := vertexai.NewVertexAIEmbedding(url, credentialsJSON, "https://www.googleapis.com/auth/cloud-platform", "")
return c, nil
}
func NewVertexAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, c *vertexai.VertexAIEmbedding, params map[string]string) (*VertexAIEmbeddingProvider, error) {
func parseGcpCredentialInfo(credentials *credentials.CredentialsManager, params []*commonpb.KeyValuePair, confParams map[string]string) ([]byte, error) {
// function param > yaml > env
var credentialsJSON []byte
var err error
for _, param := range params {
switch strings.ToLower(param.Key) {
case credentialParamKey:
credentialName := param.Value
if credentialsJSON, err = credentials.GetGcpCredential(credentialName); err != nil {
return nil, err
}
}
}
// from milvus.yaml
if credentialsJSON == nil {
credentialName := confParams[credentialParamKey]
if credentialName != "" {
if credentialsJSON, err = credentials.GetGcpCredential(credentialName); err != nil {
return nil, err
}
}
}
// from env
if credentialsJSON == nil {
credentialsJSON, err = getVertexAIJsonKey()
if err != nil {
return nil, err
}
}
return credentialsJSON, nil
}
func NewVertexAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, c *vertexai.VertexAIEmbedding, params map[string]string, credentials *credentials.CredentialsManager) (*VertexAIEmbeddingProvider, error) {
fieldDim, err := typeutil.GetDim(fieldSchema)
if err != nil {
return nil, err
@ -136,7 +163,11 @@ func NewVertexAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSch
}
var client *vertexai.VertexAIEmbedding
if c == nil {
client, err = createVertexAIEmbeddingClient(url, params["credentials_file_path"])
jsonKey, err := parseGcpCredentialInfo(credentials, functionSchema.Params, params)
if err != nil {
return nil, err
}
client, err = createVertexAIEmbeddingClient(url, jsonKey)
if err != nil {
return nil, err
}

52
internal/util/function/vertexai_embedding_provider_test.go
View File

@ -29,6 +29,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/vertexai"
)
@ -76,7 +77,7 @@ func createVertexAIProvider(url string, schema *schemapb.FieldSchema) (textEmbed
},
}
mockClient := vertexai.NewVertexAIEmbedding(url, []byte{1, 2, 3}, "mock scope", "mock token")
return NewVertexAIEmbeddingProvider(schema, functionSchema, mockClient, map[string]string{})
return NewVertexAIEmbeddingProvider(schema, functionSchema, mockClient, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.credential_json": "mock"}))
}
func (s *VertexAITextEmbeddingProviderSuite) TestEmbedding() {
@ -177,7 +178,7 @@ func (s *VertexAITextEmbeddingProviderSuite) TestEmbeddingNubmerNotMatch() {
func (s *VertexAITextEmbeddingProviderSuite) TestGetVertexAIJsonKey() {
os.Setenv(vertexServiceAccountJSONEnv, "ErrorPath")
defer os.Unsetenv(vertexServiceAccountJSONEnv)
_, err := getVertexAIJsonKey("")
_, err := getVertexAIJsonKey()
s.Error(err)
}
@ -198,7 +199,7 @@ func (s *VertexAITextEmbeddingProviderSuite) TestGetTaskType() {
mockClient := vertexai.NewVertexAIEmbedding("mock_url", []byte{1, 2, 3}, "mock scope", "mock token")
{
provider, err := NewVertexAIEmbeddingProvider(s.schema.Fields[2], functionSchema, mockClient, map[string]string{})
provider, err := NewVertexAIEmbeddingProvider(s.schema.Fields[2], functionSchema, mockClient, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.credential_json": "mock"}))
s.NoError(err)
s.Equal(provider.getTaskType(InsertMode), "RETRIEVAL_DOCUMENT")
s.Equal(provider.getTaskType(SearchMode), "RETRIEVAL_QUERY")
@ -206,7 +207,7 @@ func (s *VertexAITextEmbeddingProviderSuite) TestGetTaskType() {
{
functionSchema.Params = append(functionSchema.Params, &commonpb.KeyValuePair{Key: taskTypeParamKey, Value: vertexAICodeRetrival})
provider, err := NewVertexAIEmbeddingProvider(s.schema.Fields[2], functionSchema, mockClient, map[string]string{})
provider, err := NewVertexAIEmbeddingProvider(s.schema.Fields[2], functionSchema, mockClient, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.credential_json": "mock"}))
s.NoError(err)
s.Equal(provider.getTaskType(InsertMode), "RETRIEVAL_DOCUMENT")
s.Equal(provider.getTaskType(SearchMode), "CODE_RETRIEVAL_QUERY")
@ -214,20 +215,13 @@ func (s *VertexAITextEmbeddingProviderSuite) TestGetTaskType() {
{
functionSchema.Params[3] = &commonpb.KeyValuePair{Key: taskTypeParamKey, Value: vertexAISTS}
provider, err := NewVertexAIEmbeddingProvider(s.schema.Fields[2], functionSchema, mockClient, map[string]string{})
provider, err := NewVertexAIEmbeddingProvider(s.schema.Fields[2], functionSchema, mockClient, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.credential_json": "mock"}))
s.NoError(err)
s.Equal(provider.getTaskType(InsertMode), "SEMANTIC_SIMILARITY")
s.Equal(provider.getTaskType(SearchMode), "SEMANTIC_SIMILARITY")
}
}
func (s *VertexAITextEmbeddingProviderSuite) TestCreateVertexAIEmbeddingClient() {
os.Setenv(vertexServiceAccountJSONEnv, "ErrorPath")
defer os.Unsetenv(vertexServiceAccountJSONEnv)
_, err := createVertexAIEmbeddingClient("https://mock_url.com", "")
s.Error(err)
}
func (s *VertexAITextEmbeddingProviderSuite) TestNewVertexAIEmbeddingProvider() {
functionSchema := &schemapb.FunctionSchema{
Name: "test",
@ -243,8 +237,40 @@ func (s *VertexAITextEmbeddingProviderSuite) TestNewVertexAIEmbeddingProvider()
},
}
mockClient := vertexai.NewVertexAIEmbedding("mock_url", []byte{1, 2, 3}, "mock scope", "mock token")
provider, err := NewVertexAIEmbeddingProvider(s.schema.Fields[2], functionSchema, mockClient, map[string]string{})
provider, err := NewVertexAIEmbeddingProvider(s.schema.Fields[2], functionSchema, mockClient, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.credential_json": "mock"}))
s.NoError(err)
s.True(provider.MaxBatch() > 0)
s.Equal(provider.FieldDim(), int64(4))
}
func (s *VertexAITextEmbeddingProviderSuite) TestParseCredentail() {
{
cred := credentials.NewCredentialsManager(map[string]string{})
data, err := parseGcpCredentialInfo(cred, []*commonpb.KeyValuePair{}, map[string]string{})
s.Nil(data)
s.ErrorContains(err, "VetexAI credentials file path is empty")
}
{
os.Setenv(vertexServiceAccountJSONEnv, "mock.json")
defer os.Unsetenv(vertexServiceAccountJSONEnv)
cred := credentials.NewCredentialsManager(map[string]string{})
data, err := parseGcpCredentialInfo(cred, []*commonpb.KeyValuePair{}, map[string]string{})
s.Nil(data)
s.ErrorContains(err, "Vertexai: read credentials file failed")
}
{
cred := credentials.NewCredentialsManager(map[string]string{})
_, err := parseGcpCredentialInfo(cred, []*commonpb.KeyValuePair{}, map[string]string{"credential": "noExist"})
s.ErrorContains(err, "is not a gcp crediential, can not find key")
}
{
cred := credentials.NewCredentialsManager(map[string]string{"mock.credential_json": "NotBase64"})
_, err := parseGcpCredentialInfo(cred, []*commonpb.KeyValuePair{}, map[string]string{"credential": "mock"})
s.ErrorContains(err, "Parse gcp credential")
}
{
cred := credentials.NewCredentialsManager(map[string]string{"mock.credential_json": "bW9jaw=="})
_, err := parseGcpCredentialInfo(cred, []*commonpb.KeyValuePair{}, map[string]string{"credential": "mock"})
s.NoError(err)
}
}

10
internal/util/function/voyageai_embedding_provider.go
View File

@ -24,6 +24,7 @@ import (
"strings"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/voyageai"
"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)
@ -44,7 +45,7 @@ type VoyageAIEmbeddingProvider struct {
func createVoyageAIEmbeddingClient(apiKey string, url string) (*voyageai.VoyageAIEmbedding, error) {
if apiKey == "" {
return nil, fmt.Errorf("Missing credentials. Please pass `api_key`, or configure the %s environment variable in the Milvus service.", voyageAIAKEnvStr)
return nil, fmt.Errorf("Missing credentials config or configure the %s environment variable in the Milvus service.", voyageAIAKEnvStr)
}
if url == "" {
@ -55,12 +56,15 @@ func createVoyageAIEmbeddingClient(apiKey string, url string) (*voyageai.VoyageA
return c, nil
}
func NewVoyageAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string) (*VoyageAIEmbeddingProvider, error) {
func NewVoyageAIEmbeddingProvider(fieldSchema *schemapb.FieldSchema, functionSchema *schemapb.FunctionSchema, params map[string]string, credentials *credentials.CredentialsManager) (*VoyageAIEmbeddingProvider, error) {
fieldDim, err := typeutil.GetDim(fieldSchema)
if err != nil {
return nil, err
}
apiKey, url := parseAKAndURL(functionSchema.Params, params, voyageAIAKEnvStr)
apiKey, url, err := parseAKAndURL(credentials, functionSchema.Params, params, voyageAIAKEnvStr)
if err != nil {
return nil, err
}
var modelName string
dim := int64(0)
truncate := false

25
internal/util/function/voyageai_embedding_provider_test.go
View File

@ -29,6 +29,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/util/credentials"
"github.com/milvus-io/milvus/internal/util/function/models/voyageai"
)
@ -69,14 +70,13 @@ func createVoyageAIProvider(url string, schema *schemapb.FieldSchema, providerNa
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: modelNameParamKey, Value: TestModel},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: url},
{Key: credentialParamKey, Value: "mock"},
{Key: dimParamKey, Value: "1024"},
},
}
switch providerName {
case voyageAIProvider:
return NewVoyageAIEmbeddingProvider(schema, functionSchema, map[string]string{})
return NewVoyageAIEmbeddingProvider(schema, functionSchema, map[string]string{embeddingURLParamKey: url}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
default:
return nil, errors.New("Unknow provider")
}
@ -293,36 +293,35 @@ func (s *VoyageAITextEmbeddingProviderSuite) TestNewVoyageAIEmbeddingProvider()
OutputFieldIds: []int64{102},
Params: []*commonpb.KeyValuePair{
{Key: modelNameParamKey, Value: TestModel},
{Key: apiKeyParamKey, Value: "mock"},
{Key: embeddingURLParamKey, Value: "mock"},
{Key: credentialParamKey, Value: "mock"},
{Key: dimParamKey, Value: "1024"},
{Key: truncationParamKey, Value: "true"},
},
}
provider, err := NewVoyageAIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
provider, err := NewVoyageAIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{embeddingURLParamKey: "mock"}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.NoError(err)
s.Equal(provider.FieldDim(), int64(1024))
s.True(provider.MaxBatch() > 0)
// Invalid truncation
{
functionSchema.Params[4] = &commonpb.KeyValuePair{Key: truncationParamKey, Value: "Invalid"}
_, err := NewVoyageAIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
functionSchema.Params[3] = &commonpb.KeyValuePair{Key: truncationParamKey, Value: "Invalid"}
_, err := NewVoyageAIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.Error(err)
functionSchema.Params[4] = &commonpb.KeyValuePair{Key: truncationParamKey, Value: "false"}
functionSchema.Params[3] = &commonpb.KeyValuePair{Key: truncationParamKey, Value: "false"}
}
// Invalid dim
{
functionSchema.Params[3] = &commonpb.KeyValuePair{Key: dimParamKey, Value: "9"}
_, err := NewVoyageAIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
functionSchema.Params[2] = &commonpb.KeyValuePair{Key: dimParamKey, Value: "9"}
_, err := NewVoyageAIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.Error(err)
}
// Invalid dim type
{
functionSchema.Params[3] = &commonpb.KeyValuePair{Key: dimParamKey, Value: "Invalied"}
_, err := NewVoyageAIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{})
functionSchema.Params[2] = &commonpb.KeyValuePair{Key: dimParamKey, Value: "Invalied"}
_, err := NewVoyageAIEmbeddingProvider(s.schema.Fields[2], functionSchema, map[string]string{}, credentials.NewCredentialsManager(map[string]string{"mock.apikey": "mock"}))
s.Error(err)
}
}

2
pkg/util/paramtable/component_param.go
View File

@ -86,6 +86,7 @@ type ComponentParam struct {
RbacConfig rbacConfig
StreamingCfg streamingConfig
FunctionCfg functionConfig
CredentialCfg credentialConfig
InternalTLSCfg InternalTLSConfig
@ -142,6 +143,7 @@ func (p *ComponentParam) init(bt *BaseTable) {
p.GpuConfig.init(bt)
p.KnowhereConfig.init(bt)
p.FunctionCfg.init(bt)
p.CredentialCfg.init(bt)
p.InternalTLSCfg.Init(bt)

48
pkg/util/paramtable/credential_param.go Normal file
View File

@ -0,0 +1,48 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package paramtable
type credentialConfig struct {
Credential ParamGroup `refreshable:"true"`
}
func (p *credentialConfig) init(base *BaseTable) {
p.Credential = ParamGroup{
KeyPrefix: "credential.",
Version: "2.6.0",
Export: true,
DocFunc: func(key string) string {
switch key {
case "apikey1.apikey":
return "Your apikey credential"
case "aksk1.access_key_id":
return "Your access_key_id"
case "aksk1.secret_access_key":
return "Your secret_access_key"
case "gcp1.credential_json":
return "base64 based gcp credential data"
default:
return ""
}
},
}
p.Credential.Init(base.mgr)
}
func (p *credentialConfig) GetCredentials() map[string]string {
return p.Credential.GetValue()
}

39
pkg/util/paramtable/credential_param_test.go Normal file
View File

@ -0,0 +1,39 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package paramtable
import (
"testing"
"github.com/stretchr/testify/assert"
)
func TestCredentialConfig(t *testing.T) {
params := ComponentParam{}
params.Init(NewBaseTable(SkipRemote(true)))
cfg := &params.CredentialCfg
keys := []string{
"aksk1.access_key_id",
"aksk1.secret_access_key",
"apikey1.apikey",
"gcp1.credential_json",
}
for _, key := range keys {
assert.True(t, cfg.Credential.GetDoc(key) != "")
}
assert.True(t, cfg.Credential.GetDoc("Unknow") == "")
}

59
pkg/util/paramtable/function_param.go
View File

@ -21,20 +21,10 @@ import (
)
type functionConfig struct {
TextEmbeddingEnableVerifiInfoInParams ParamItem `refreshable:"true"`
TextEmbeddingProviders ParamGroup `refreshable:"true"`
TextEmbeddingProviders ParamGroup `refreshable:"true"`
}
func (p *functionConfig) init(base *BaseTable) {
p.TextEmbeddingEnableVerifiInfoInParams = ParamItem{
Key: "function.textEmbedding.enableVerifiInfoInParams",
Version: "2.6.0",
DefaultValue: "true",
Export: true,
Doc: "Controls whether to allow configuration of apikey and model service url on function parameters",
}
p.TextEmbeddingEnableVerifiInfoInParams.Init(base.mgr)
p.TextEmbeddingProviders = ParamGroup{
KeyPrefix: "function.textEmbedding.providers.",
Version: "2.6.0",
@ -43,40 +33,40 @@ func (p *functionConfig) init(base *BaseTable) {
switch key {
case "tei.enable":
return "Whether to enable TEI model service"
case "azure_openai.api_key":
return "Your azure openai embedding url, Default is the official embedding url"
case "tei.credential":
return "The name in the crendential configuration item"
case "azure_openai.credential":
return "The name in the crendential configuration item"
case "azure_openai.url":
return "Your azure openai api key"
return "Your azure openai embedding url, Default is the official embedding url"
case "azure_openai.resource_name":
return "Your azure openai resource name"
case "openai.api_key":
return "Your openai embedding url, Default is the official embedding url"
case "openai.credential":
return "The name in the crendential configuration item"
case "openai.url":
return "Your openai api key"
case "dashscope.api_key":
return "Your dashscope embedding url, Default is the official embedding url"
return "Your openai embedding url, Default is the official embedding url"
case "dashscope.credential":
return "The name in the crendential configuration item"
case "dashscope.url":
return "Your dashscope api key"
case "cohere.api_key":
return "Your cohere embedding url, Default is the official embedding url"
return "Your dashscope embedding url, Default is the official embedding url"
case "cohere.credential":
return "The name in the crendential configuration item"
case "cohere.url":
return "Your cohere api key"
case "voyageai.api_key":
return "Your voyageai embedding url, Default is the official embedding url"
return "Your cohere embedding url, Default is the official embedding url"
case "voyageai.credential":
return "The name in the crendential configuration item"
case "voyageai.url":
return "Your voyageai api key"
return "Your voyageai embedding url, Default is the official embedding url"
case "siliconflow.url":
return "Your siliconflow embedding url, Default is the official embedding url"
case "siliconflow.api_key":
return "Your siliconflow api key"
case "bedrock.aws_access_key_id":
return "Your aws_access_key_id"
case "bedrock.aws_secret_access_key":
return "Your aws_secret_access_key"
case "siliconflow.credential":
return "The name in the crendential configuration item"
case "bedrock.credential":
return "The name in the crendential configuration item"
case "vertexai.url":
return "Your VertexAI embedding url"
case "vertexai.credentials_file_path":
return "Path to your google application credentials, change the file path to refresh the configuration"
case "vertexai.credential":
return "The name in the crendential configuration item"
default:
return ""
}
@ -100,6 +90,5 @@ func (p *functionConfig) GetTextEmbeddingProviderConfig(providerName string) map
matchedParam[strings.TrimPrefix(k, prefix)] = v
}
}
matchedParam["enableVerifiInfoInParams"] = p.TextEmbeddingEnableVerifiInfoInParams.GetValue()
return matchedParam
}

28
pkg/util/paramtable/function_param_test.go
View File

@ -26,40 +26,34 @@ func TestFunctionConfig(t *testing.T) {
params := ComponentParam{}
params.Init(NewBaseTable(SkipRemote(true)))
cfg := &params.FunctionCfg
notExistProvider := cfg.GetTextEmbeddingProviderConfig("notExist")
// Only has enableVerifiInfoInParams config
assert.Equal(t, len(notExistProvider), 1)
teiConf := cfg.GetTextEmbeddingProviderConfig("tei")
assert.Equal(t, teiConf["enable"], "true")
assert.Equal(t, teiConf["enableVerifiInfoInParams"], "true")
openaiConf := cfg.GetTextEmbeddingProviderConfig("openai")
assert.Equal(t, openaiConf["api_key"], "")
assert.Equal(t, openaiConf["credential"], "")
assert.Equal(t, openaiConf["url"], "")
assert.Equal(t, openaiConf["enableVerifiInfoInParams"], "true")
keys := []string{
"tei.enable",
"azure_openai.api_key",
"tei.credential",
"azure_openai.credential",
"azure_openai.url",
"azure_openai.resource_name",
"openai.api_key",
"openai.credential",
"openai.url",
"dashscope.api_key",
"dashscope.credential",
"dashscope.url",
"cohere.api_key",
"cohere.credential",
"cohere.url",
"voyageai.api_key",
"voyageai.credential",
"voyageai.url",
"siliconflow.url",
"siliconflow.api_key",
"bedrock.aws_access_key_id",
"bedrock.aws_secret_access_key",
"siliconflow.credential",
"bedrock.credential",
"vertexai.url",
"vertexai.credentials_file_path",
"vertexai.credential",
}
for _, key := range keys {
assert.True(t, cfg.TextEmbeddingProviders.GetDoc(key) != "")
}
assert.True(t, cfg.TextEmbeddingProviders.GetDoc("Unknow") == "")
}