diff --git a/ai/instill/v0/main.go b/ai/instill/v0/main.go index 5aa87541..f1332e1e 100644 --- a/ai/instill/v0/main.go +++ b/ai/instill/v0/main.go @@ -68,6 +68,10 @@ func getInstillUserUID(vars map[string]any) string { return vars["__PIPELINE_USER_UID"].(string) } +func getInstillRequesterUID(vars map[string]any) string { + return vars["__PIPELINE_REQUESTER_UID"].(string) +} + func getModelServerURL(vars map[string]any) string { if v, ok := vars["__MODEL_BACKEND"]; ok { return v.(string) @@ -83,11 +87,17 @@ func getMgmtServerURL(vars map[string]any) string { } func getRequestMetadata(vars map[string]any) metadata.MD { - return metadata.Pairs( + md := metadata.Pairs( "Authorization", getHeaderAuthorization(vars), "Instill-User-Uid", getInstillUserUID(vars), "Instill-Auth-Type", "user", ) + + if requester := getInstillRequesterUID(vars); requester != "" { + md.Set("Instill-Requester-Uid", requester) + } + + return md } func (e *execution) Execute(ctx context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) { diff --git a/ai/openai/v0/config/tasks.json b/ai/openai/v0/config/tasks.json index 83c02cd7..d26a92b4 100644 --- a/ai/openai/v0/config/tasks.json +++ b/ai/openai/v0/config/tasks.json @@ -173,6 +173,19 @@ ], "title": "Text", "type": "string" + }, + "dimensions": { + "description": "The number of dimensions the resulting output embeddings should have. Only supported in text-embedding-3 and later models.", + "instillAcceptFormats": [ + "integer" + ], + "instillUIOrder": 2, + "instillUpstreamTypes": [ + "value", + "reference" + ], + "title": "Dimensions", + "type": "integer" } }, "required": [ diff --git a/ai/openai/v0/main.go b/ai/openai/v0/main.go index ea80c05b..5c7b14d9 100644 --- a/ai/openai/v0/main.go +++ b/ai/openai/v0/main.go @@ -215,10 +215,22 @@ func (e *execution) Execute(_ context.Context, inputs []*structpb.Struct) ([]*st } resp := TextEmbeddingsResp{} - req := client.R().SetBody(TextEmbeddingsReq{ - Model: inputStruct.Model, - Input: []string{inputStruct.Text}, - }).SetResult(&resp) + + var reqParams TextEmbeddingsReq + if inputStruct.Dimensions == 0 { + reqParams = TextEmbeddingsReq{ + Model: inputStruct.Model, + Input: []string{inputStruct.Text}, + } + } else { + reqParams = TextEmbeddingsReq{ + Model: inputStruct.Model, + Input: []string{inputStruct.Text}, + Dimensions: inputStruct.Dimensions, + } + } + + req := client.R().SetBody(reqParams).SetResult(&resp) if _, err := req.Post(embeddingsPath); err != nil { return inputs, err diff --git a/ai/openai/v0/text_embeddings.go b/ai/openai/v0/text_embeddings.go index fd4ef64e..4cd29d80 100644 --- a/ai/openai/v0/text_embeddings.go +++ b/ai/openai/v0/text_embeddings.go @@ -5,8 +5,9 @@ const ( ) type TextEmbeddingsInput struct { - Text string `json:"text"` - Model string `json:"model"` + Text string `json:"text"` + Model string `json:"model"` + Dimensions int `json:"dimensions"` } type TextEmbeddingsOutput struct { @@ -14,8 +15,9 @@ type TextEmbeddingsOutput struct { } type TextEmbeddingsReq struct { - Model string `json:"model"` - Input []string `json:"input"` + Model string `json:"model"` + Dimensions int `json:"dimensions,omitempty"` + Input []string `json:"input"` } type TextEmbeddingsResp struct {