From 0d089121d280e663c36529426e4518411b58f6c2 Mon Sep 17 00:00:00 2001
From: ChunHao <64747455+chuang8511@users.noreply.github.com>
Date: Mon, 8 Jul 2024 14:25:44 +0100
Subject: [PATCH] feat(openai): add dimensions in openai component (#200)

Because

- in many use cases, users may not need to output vectors of such high
dimensionality and might only require the first 'x' dimensions.

This commit

- add dimensions in api caller
---
 ai/openai/v0/config/tasks.json  | 13 +++++++++++++
 ai/openai/v0/main.go            | 20 ++++++++++++++++----
 ai/openai/v0/text_embeddings.go | 10 ++++++----
 3 files changed, 35 insertions(+), 8 deletions(-)

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 {