Merge branch 'dev' of https://github.com/continuedev/continue into dev

sestinj · sestinj · commit c3cbbd6fbff9 · 2024-10-08T21:30:00.000-07:00
diff --git a/core/index.d.ts b/core/index.d.ts
@@ -801,6 +801,7 @@ export interface ModelDescription {
 }
 
 export type EmbeddingsProviderName =
+  | "sagemaker"
   | "bedrock"
   | "huggingface-tei"
   | "transformers.js"
diff --git a/core/indexing/embeddings/SageMakerEmbeddingsProvider.ts b/core/indexing/embeddings/SageMakerEmbeddingsProvider.ts
@@ -0,0 +1,114 @@
+import {
+    InvokeEndpointCommand,
+    SageMakerRuntimeClient,
+  } from "@aws-sdk/client-sagemaker-runtime";
+  import { fromIni } from "@aws-sdk/credential-providers";
+  
+  import {
+    EmbeddingsProviderName,
+    EmbedOptions,
+    FetchFunction,
+  } from "../../index.js";
+  import BaseEmbeddingsProvider from "./BaseEmbeddingsProvider.js";
+  
+  
+  class SageMakerEmbeddingsProvider extends BaseEmbeddingsProvider {
+    static providerName: EmbeddingsProviderName = "sagemaker";
+  
+    static defaultOptions: Partial<EmbedOptions> | undefined = {
+      region: "us-west-2",
+      maxBatchSize: 1,
+    };
+    profile?: string | undefined;
+  
+    constructor(options: EmbedOptions, fetch: FetchFunction) {
+      super(options, fetch);
+      if (!options.apiBase) {
+        options.apiBase = `https://runtime.sagemaker.${options.region}.amazonaws.com`;
+      }
+  
+      if (options.profile) {
+        this.profile = options.profile;
+      } else {
+        this.profile = "sagemaker";
+      }
+    }
+  
+    async embed(chunks: string[]) {
+      const credentials = await this._getCredentials();
+      const client = new SageMakerRuntimeClient({
+        region: this.options.region,
+        credentials: {
+          accessKeyId: credentials.accessKeyId,
+          secretAccessKey: credentials.secretAccessKey,
+          sessionToken: credentials.sessionToken || "",
+        },
+      });
+  
+      const batchedChunks = this.getBatchedChunks(chunks);
+      return (
+        await Promise.all(
+          batchedChunks.map(async (batch) => {
+            const input = this._generateInvokeModelCommandInput(
+              batch,
+              this.options,
+            );
+            const command = new InvokeEndpointCommand(input);
+            const response = await client.send(command);
+  
+            if (response.Body) {
+              const responseBody = JSON.parse(
+                new TextDecoder().decode(response.Body),
+              );
+              // If the body contains a key called "embedding" or "embeddings", return the value, otherwise return the whole body
+              if (responseBody.embedding) {
+                return responseBody.embedding;
+              } else if (responseBody.embeddings) {
+                return responseBody.embeddings;
+              } else {
+                return responseBody;
+              }
+            }
+          }),
+        )
+      ).flat();
+    }
+    private _generateInvokeModelCommandInput(
+      prompts: string | string[],
+      options: EmbedOptions,
+    ): any {
+      const payload = {
+        inputs: prompts,
+        normalize: true,
+        // ...(options.requestOptions?.extraBodyProperties || {}),
+      };
+  
+      if (options.requestOptions?.extraBodyProperties) {
+        Object.assign(payload, options.requestOptions.extraBodyProperties);
+      }
+  
+      return {
+        EndpointName: this.options.model,
+        Body: JSON.stringify(payload),
+        ContentType: "application/json",
+        CustomAttributes: "accept_eula=false",
+      };
+    }
+  
+    private async _getCredentials() {
+      try {
+        return await fromIni({
+          profile: this.profile,
+          ignoreCache: true,
+        })();
+      } catch (e) {
+        console.warn(
+          `AWS profile with name ${this.profile} not found in ~/.aws/credentials, using default profile`,
+        );
+        return await fromIni()();
+      }
+    }
+  }
+  
+  export default SageMakerEmbeddingsProvider;
+  
diff --git a/core/indexing/embeddings/index.ts b/core/indexing/embeddings/index.ts
@@ -1,5 +1,6 @@
 import { EmbeddingsProviderName } from "../../index.js";
 import BaseEmbeddingsProvider from "./BaseEmbeddingsProvider.js";
+import SageMakerEmbeddingsProvider from "./SageMakerEmbeddingsProvider.js";
 import BedrockEmbeddingsProvider from "./BedrockEmbeddingsProvider.js";
 import CohereEmbeddingsProvider from "./CohereEmbeddingsProvider.js";
 import ContinueProxyEmbeddingsProvider from "./ContinueProxyEmbeddingsProvider.js";
@@ -22,6 +23,7 @@ export const allEmbeddingsProviders: Record<
   EmbeddingsProviderName,
   EmbeddingsProviderConstructor
 > = {
+  sagemaker: SageMakerEmbeddingsProvider,
   bedrock: BedrockEmbeddingsProvider,
   ollama: OllamaEmbeddingsProvider,
   "transformers.js": TransformersJsEmbeddingsProvider,
diff --git a/docs/docs/customize/model-providers/more/sagemaker.md b/docs/docs/customize/model-providers/more/sagemaker.md
@@ -1,8 +1,8 @@
 # AWS SageMaker
 
-SageMaker provider support SageMaker endpoint deployed with [LMI](https://docs.djl.ai/docs/serving/serving/docs/lmi/index.html)
+SageMaker can be used for both chat and embedding models. Chat models are supported for endpoints deployed with [LMI](https://docs.djl.ai/docs/serving/serving/docs/lmi/index.html), and embedding models are supported for endpoints deployed with [HuggingFace TEI](https://huggingface.co/blog/sagemaker-huggingface-embedding)
 
-To setup SageMaker, add the following to your `config.json` file:
+To setup SageMaker as a chat model provider, add the following to your `config.json` file:
 
 ```json title="config.json"
 {
@@ -13,7 +13,11 @@ To setup SageMaker, add the following to your `config.json` file:
       "model": "lmi-model-deepseek-coder-xxxxxxx",
       "region": "us-west-2"
     }
-  ]
+  ],
+  "embeddingsProvider": {
+    "provider": "sagemaker",
+    "model": "mxbai-embed-large-v1-endpoint"
+  },
 }
 ```
 
diff --git a/docs/static/schemas/config.json b/docs/static/schemas/config.json
@@ -246,7 +246,24 @@
         "region": {
           "title": "Region",
           "description": "The region where the model is hosted",
-          "type": "string"
+          "anyOf":[
+            {
+              "enum": [
+                "us-east-1",
+                "us-east-2",
+                "us-west-1",
+                "us-west-2",
+                "eu-west-1",
+                "eu-central-1",
+                "ap-southeast-1",
+                "ap-northeast-1",
+                "ap-south-1"
+              ]
+            },
+            {
+              "type": "string"
+            }
+          ]
         },
         "profile": {
           "title": "Profile",
@@ -2150,7 +2167,8 @@
                 "gemini",
                 "voyage",
                 "nvidia",
-                "bedrock"
+                "bedrock",
+                "sagemaker"
               ]
             },
             "model": {
@@ -2184,7 +2202,24 @@
             "region": {
               "title": "Region",
               "description": "The region where the model is hosted",
-              "type": "string"
+              "anyOf":[
+                {
+                  "enum": [
+                    "us-east-1",
+                    "us-east-2",
+                    "us-west-1",
+                    "us-west-2",
+                    "eu-west-1",
+                    "eu-central-1",
+                    "ap-southeast-1",
+                    "ap-northeast-1",
+                    "ap-south-1"
+                  ]
+                },
+                {
+                  "type": "string"
+                }
+              ]
             },
             "profile": {
               "title": "Profile",
@@ -2206,6 +2241,24 @@
               "then": {
                 "required": ["apiKey"]
               }
+            },
+            {
+              "if": {
+                "properties": {
+                  "provider": {
+                    "enum": ["sagemaker"]
+                  }
+                },
+                "required": ["provider"]
+              },
+              "then": {
+                "properties": {
+                  "model": {
+                    "markdownDescription": "SageMaker endpoint name"
+                  }
+                },
+                "required": ["model"]
+              }
             }
           ]
         },
diff --git a/extensions/intellij/src/main/resources/config_schema.json b/extensions/intellij/src/main/resources/config_schema.json
@@ -246,7 +246,25 @@
         "region": {
           "title": "Region",
           "description": "The region where the model is hosted",
-          "type": "string"
+          "anyOf":[
+            {
+              "enum": [
+                "us-east-1",
+                "us-east-2",
+                "us-west-1",
+                "us-west-2",
+                "eu-west-1",
+                "eu-central-1",
+                "ap-southeast-1",
+                "ap-northeast-1",
+                "ap-south-1"
+              ],
+              "type" : "string"
+            },
+            {
+              "type": "string"
+            }
+          ]
         },
         "profile": {
           "title": "Profile",
@@ -2150,7 +2168,8 @@
                 "gemini",
                 "voyage",
                 "nvidia",
-                "bedrock"
+                "bedrock",
+                "sagemaker"
               ]
             },
             "model": {
@@ -2184,7 +2203,7 @@
             "region": {
               "title": "Region",
               "description": "The region where the model is hosted",
-              "type": "string"
+              "$ref": "#/definitions/ModelDescription/properties/region"
             },
             "profile": {
               "title": "Profile",
@@ -2206,6 +2225,24 @@
               "then": {
                 "required": ["apiKey"]
               }
+            },
+            {
+              "if": {
+                "properties": {
+                  "provider": {
+                    "enum": ["sagemaker"]
+                  }
+                },
+                "required": ["provider"]
+              },
+              "then": {
+                "properties": {
+                  "model": {
+                    "markdownDescription": "SageMaker endpoint name"
+                  }
+                },
+                "required": ["model"]
+              }
             }
           ]
         },
diff --git a/extensions/vscode/config_schema.json b/extensions/vscode/config_schema.json
@@ -246,7 +246,25 @@
         "region": {
           "title": "Region",
           "description": "The region where the model is hosted",
-          "type": "string"
+          "anyOf":[
+            {
+              "enum": [
+                "us-east-1",
+                "us-east-2",
+                "us-west-1",
+                "us-west-2",
+                "eu-west-1",
+                "eu-central-1",
+                "ap-southeast-1",
+                "ap-northeast-1",
+                "ap-south-1"
+              ],
+              "type" : "string"
+            },
+            {
+              "type": "string"
+            }
+          ]
         },
         "profile": {
           "title": "Profile",
@@ -2150,7 +2168,8 @@
                 "gemini",
                 "voyage",
                 "nvidia",
-                "bedrock"
+                "bedrock",
+                "sagemaker"
               ]
             },
             "model": {
@@ -2184,7 +2203,7 @@
             "region": {
               "title": "Region",
               "description": "The region where the model is hosted",
-              "type": "string"
+              "$ref": "#/definitions/ModelDescription/properties/region"
             },
             "profile": {
               "title": "Profile",
@@ -2206,6 +2225,24 @@
               "then": {
                 "required": ["apiKey"]
               }
+            },
+            {
+              "if": {
+                "properties": {
+                  "provider": {
+                    "enum": ["sagemaker"]
+                  }
+                },
+                "required": ["provider"]
+              },
+              "then": {
+                "properties": {
+                  "model": {
+                    "markdownDescription": "SageMaker endpoint name"
+                  }
+                },
+                "required": ["model"]
+              }
             }
           ]
         },
diff --git a/extensions/vscode/continue_rc_schema.json b/extensions/vscode/continue_rc_schema.json

Original file line number	Diff line number	Diff line change
`@@ -801,6 +801,7 @@ export interface ModelDescription {`
`801`	`801`	`}`
`802`	`802`
`803`	`803`	`export type EmbeddingsProviderName =`
	`804`	`+ \| "sagemaker"`
`804`	`805`	`\| "bedrock"`
`805`	`806`	`\| "huggingface-tei"`
`806`	`807`	`\| "transformers.js"`
Original file line number	Diff line number	Diff line change
`@@ -1,8 +1,8 @@`
`1`	`1`	`# AWS SageMaker`
`2`	`2`
`3`		`-SageMaker provider support SageMaker endpoint deployed with [LMI](https://docs.djl.ai/docs/serving/serving/docs/lmi/index.html)`
	`3`	`+SageMaker can be used for both chat and embedding models. Chat models are supported for endpoints deployed with [LMI](https://docs.djl.ai/docs/serving/serving/docs/lmi/index.html), and embedding models are supported for endpoints deployed with [HuggingFace TEI](https://huggingface.co/blog/sagemaker-huggingface-embedding)`
`4`	`4`
`5`		-To setup SageMaker, add the following to your `config.json` file:
	`5`	+To setup SageMaker as a chat model provider, add the following to your `config.json` file:
`6`	`6`
`7`	`7`	```json title="config.json"
`8`	`8`	`{`
@@ -13,7 +13,11 @@ To setup SageMaker, add the following to your `config.json` file:
`13`	`13`	`"model": "lmi-model-deepseek-coder-xxxxxxx",`
`14`	`14`	`"region": "us-west-2"`
`15`	`15`	`}`
`16`		`- ]`
	`16`	`+ ],`
	`17`	`+ "embeddingsProvider": {`
	`18`	`+ "provider": "sagemaker",`
	`19`	`+ "model": "mxbai-embed-large-v1-endpoint"`
	`20`	`+ },`
`17`	`21`	`}`
`18`	`22`	```
`19`	`23`