From 60e7e17c8677feac8c5bbea9227a2f904732343f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B4=AA=E6=9C=94?= <i@tuccuay.com>
Date: Fri, 8 Mar 2024 19:04:20 +0800
Subject: [PATCH] feat: Add new Azure OpenAI Embedding models (#2758)

---
 api/core/hosting_configuration.py             |  2 +
 .../model_providers/azure_openai/_constant.py | 40 +++++++++++++++++++
 .../azure_openai/azure_openai.yaml            | 12 ++++++
 3 files changed, 54 insertions(+)

diff --git a/api/core/hosting_configuration.py b/api/core/hosting_configuration.py
index 880a30cdf4e73f..45ad1b51bf997b 100644
--- a/api/core/hosting_configuration.py
+++ b/api/core/hosting_configuration.py
@@ -82,6 +82,8 @@ def init_azure_openai(self, app_config: Config) -> HostingProvider:
                     RestrictModel(model="gpt-35-turbo-16k", base_model_name="gpt-35-turbo-16k", model_type=ModelType.LLM),
                     RestrictModel(model="text-davinci-003", base_model_name="text-davinci-003", model_type=ModelType.LLM),
                     RestrictModel(model="text-embedding-ada-002", base_model_name="text-embedding-ada-002", model_type=ModelType.TEXT_EMBEDDING),
+                    RestrictModel(model="text-embedding-3-small", base_model_name="text-embedding-3-small", model_type=ModelType.TEXT_EMBEDDING),
+                    RestrictModel(model="text-embedding-3-large", base_model_name="text-embedding-3-large", model_type=ModelType.TEXT_EMBEDDING),
                 ]
             )
             quotas.append(trial_quota)
diff --git a/api/core/model_runtime/model_providers/azure_openai/_constant.py b/api/core/model_runtime/model_providers/azure_openai/_constant.py
index 8d6cac3ec3ba1d..7fc0da73fb52d4 100644
--- a/api/core/model_runtime/model_providers/azure_openai/_constant.py
+++ b/api/core/model_runtime/model_providers/azure_openai/_constant.py
@@ -524,6 +524,46 @@ class AzureBaseModel(BaseModel):
                 currency='USD',
             )
         )
+    ),
+    AzureBaseModel(
+        base_model_name='text-embedding-3-small',
+        entity=AIModelEntity(
+            model='fake-deployment-name',
+            label=I18nObject(
+                en_US='fake-deployment-name-label'
+            ),
+            fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,
+            model_type=ModelType.TEXT_EMBEDDING,
+            model_properties={
+                ModelPropertyKey.CONTEXT_SIZE: 8191,
+                ModelPropertyKey.MAX_CHUNKS: 32,
+            },
+            pricing=PriceConfig(
+                input=0.00002,
+                unit=0.001,
+                currency='USD',
+            )
+        )
+    ),
+    AzureBaseModel(
+        base_model_name='text-embedding-3-large',
+        entity=AIModelEntity(
+            model='fake-deployment-name',
+            label=I18nObject(
+                en_US='fake-deployment-name-label'
+            ),
+            fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,
+            model_type=ModelType.TEXT_EMBEDDING,
+            model_properties={
+                ModelPropertyKey.CONTEXT_SIZE: 8191,
+                ModelPropertyKey.MAX_CHUNKS: 32,
+            },
+            pricing=PriceConfig(
+                input=0.00013,
+                unit=0.001,
+                currency='USD',
+            )
+        )
     )
 ]
 SPEECH2TEXT_BASE_MODELS = [
diff --git a/api/core/model_runtime/model_providers/azure_openai/azure_openai.yaml b/api/core/model_runtime/model_providers/azure_openai/azure_openai.yaml
index fe4f3538ed3a4b..6c56ccc920c77d 100644
--- a/api/core/model_runtime/model_providers/azure_openai/azure_openai.yaml
+++ b/api/core/model_runtime/model_providers/azure_openai/azure_openai.yaml
@@ -100,6 +100,18 @@ model_credential_schema:
           show_on:
             - variable: __model_type
               value: text-embedding
+        - label:
+            en_US: text-embedding-3-small
+          value: text-embedding-3-small
+          show_on:
+            - variable: __model_type
+              value: text-embedding
+        - label:
+            en_US: text-embedding-3-large
+          value: text-embedding-3-large
+          show_on:
+            - variable: __model_type
+              value: text-embedding
         - label:
             en_US: whisper-1
           value: whisper-1