Merge branch 'instill-ai:main' into feat/mongodb-component

instill-ai · Jul 8, 2024 · 9784665 · 9784665
2 parents 7811df1 + 31422cd
commit 9784665
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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
@@ -5,17 +5,19 @@ 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 {
 	Embedding []float64 `json:"embedding"`
 }
 
 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 {