Explorar o código

补充豆包的向量模型实现

liyanbo hai 1 mes
pai
achega
e2648ea317

+ 24 - 1
byzs-module-ai/src/main/java/cn/iocoder/byzs/module/ai/framework/ai/core/AiModelFactoryImpl.java

@@ -14,6 +14,7 @@ import cn.iocoder.byzs.module.ai.enums.model.AiPlatformEnum;
 import cn.iocoder.byzs.module.ai.framework.ai.core.model.baichuan.BaiChuanChatModel;
 import cn.iocoder.byzs.module.ai.framework.ai.core.model.deepseek.DeepSeekChatModel;
 import cn.iocoder.byzs.module.ai.framework.ai.core.model.doubao.DouBaoChatModel;
+import cn.iocoder.byzs.module.ai.framework.ai.core.model.doubao.DouBaoEmbeddingModel;
 import cn.iocoder.byzs.module.ai.framework.ai.core.model.hunyuan.HunYuanChatModel;
 import cn.iocoder.byzs.module.ai.framework.ai.core.model.midjourney.api.MidjourneyApi;
 import cn.iocoder.byzs.module.ai.framework.ai.core.model.siliconflow.SiliconFlowApiConstants;
@@ -32,7 +33,10 @@ import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingModel;
 import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingOptions;
 import com.alibaba.cloud.ai.dashscope.image.DashScopeImageModel;
 import com.azure.ai.openai.OpenAIClientBuilder;
+import com.volcengine.ark.runtime.service.ArkService;
 import io.micrometer.observation.ObservationRegistry;
+import okhttp3.ConnectionPool;
+import okhttp3.Dispatcher;
 import io.milvus.client.MilvusServiceClient;
 import io.qdrant.client.QdrantClient;
 import io.qdrant.client.QdrantGrpcClient;
@@ -115,6 +119,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.Timer;
 import java.util.TimerTask;
+import java.util.concurrent.TimeUnit;
 
 import static cn.iocoder.byzs.framework.common.util.collection.CollectionUtils.convertList;
 import static org.springframework.ai.retry.RetryUtils.DEFAULT_RETRY_TEMPLATE;
@@ -279,6 +284,8 @@ public class AiModelFactoryImpl implements AiModelFactory {
                     return buildZhiPuEmbeddingModel(apiKey, url, model);
                 case MINI_MAX:
                     return buildMiniMaxEmbeddingModel(apiKey, url, model);
+                case DOU_BAO:
+                    return buildDouBaoEmbeddingModel(apiKey, url, model);
                 case OPENAI:
                     return buildOpenAiEmbeddingModel(apiKey, url, model);
                 case AZURE_OPENAI:
@@ -586,6 +593,22 @@ public class AiModelFactoryImpl implements AiModelFactory {
         return OllamaEmbeddingModel.builder().ollamaApi(ollamaApi).defaultOptions(ollamaOptions).build();
     }
 
+    /**
+     * 创建豆包 Embedding 模型,使用火山引擎官方 SDK
+     */
+    private EmbeddingModel buildDouBaoEmbeddingModel(String apiKey, String url, String model) {
+        url = StrUtil.blankToDefault(url, "https://ark.cn-beijing.volces.com/api/v3");
+        ConnectionPool connectionPool = new ConnectionPool(5, 1, TimeUnit.SECONDS);
+        Dispatcher dispatcher = new Dispatcher();
+        ArkService arkService = ArkService.builder()
+                .dispatcher(dispatcher)
+                .connectionPool(connectionPool)
+                .baseUrl(url)
+                .apiKey(apiKey)
+                .build();
+        return new DouBaoEmbeddingModel(arkService, model);
+    }
+
     /**
      * 可参考 {@link OpenAiAutoConfiguration} 的 openAiEmbeddingModel 方法
      */
@@ -769,4 +792,4 @@ public class AiModelFactoryImpl implements AiModelFactory {
         return SpringUtil.getBean(FunctionCallbackResolver.class);
     }
 
-}
+}

+ 55 - 0
byzs-module-ai/src/main/java/cn/iocoder/byzs/module/ai/framework/ai/core/model/doubao/DouBaoEmbeddingModel.java

@@ -0,0 +1,55 @@
+package cn.iocoder.byzs.module.ai.framework.ai.core.model.doubao;
+
+import com.volcengine.ark.runtime.model.multimodalembeddings.MultimodalEmbeddingInput;
+import com.volcengine.ark.runtime.model.multimodalembeddings.MultimodalEmbeddingRequest;
+import com.volcengine.ark.runtime.model.multimodalembeddings.MultimodalEmbeddingResult;
+import com.volcengine.ark.runtime.service.ArkService;
+import lombok.RequiredArgsConstructor;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.embedding.Embedding;
+import org.springframework.ai.embedding.EmbeddingModel;
+import org.springframework.ai.embedding.EmbeddingRequest;
+import org.springframework.ai.embedding.EmbeddingResponse;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@RequiredArgsConstructor
+public class DouBaoEmbeddingModel implements EmbeddingModel {
+
+    private final ArkService arkService;
+    private final String model;
+
+    @Override
+    public EmbeddingResponse call(EmbeddingRequest request) {
+        List<String> inputs = request.getInstructions();
+        List<Embedding> embeddings = new ArrayList<>();
+        for (int idx = 0; idx < inputs.size(); idx++) {
+            String text = inputs.get(idx);
+            MultimodalEmbeddingInput embeddingInput = MultimodalEmbeddingInput.builder()
+                    .type("text")
+                    .text(text)
+                    .build();
+
+            MultimodalEmbeddingRequest multimodalRequest = MultimodalEmbeddingRequest.builder()
+                    .model(model)
+                    .input(List.of(embeddingInput))
+                    .build();
+
+            MultimodalEmbeddingResult result = arkService.createMultiModalEmbeddings(multimodalRequest);
+
+            List<Double> doubleEmbedding = result.getData().getEmbedding();
+            float[] floatEmbedding = new float[doubleEmbedding.size()];
+            for (int i = 0; i < doubleEmbedding.size(); i++) {
+                floatEmbedding[i] = doubleEmbedding.get(i).floatValue();
+            }
+            embeddings.add(new Embedding(floatEmbedding, idx, null));
+        }
+        return new EmbeddingResponse(embeddings);
+    }
+
+    @Override
+    public float[] embed(Document document) {
+        return embed(document.getText());
+    }
+}