Repository: langchain4j/langchain4j (Java)

Issue:
[BUG] The method 'toEmbedTexts' in DefaultZhipuAiHelper may have a bug.
When processing textSegments in the 'toEmbedTexts' method of DefaultZhipuAiHelper, only the first item of the result after streaming conversion is taken. If the size of the passed List<TextSegment> is greater than 1, this handling will cause: 1. Data loss, 2. An inconsistency between the entries of 'embeddings' and 'segments' during the subsequent embeddingStore.addAll(embeddings, segments), resulting in detection exceptions. It is unclear whether the 'return embedText.get(0)' is simply a bug or intentionally designed this way.
![bug](https://github.com/langchain4j/langchain4j/assets/11806067/74b8c531-c874-4e9a-b48b-125a8f78cca6)




Fix (patch):
diff --git a/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/DefaultZhipuAiHelper.java b/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/DefaultZhipuAiHelper.java
index 7c2816702e8..5a0c07c829b 100644
--- a/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/DefaultZhipuAiHelper.java
+++ b/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/DefaultZhipuAiHelper.java
@@ -4,25 +4,12 @@
 import dev.langchain4j.agent.tool.ToolParameters;
 import dev.langchain4j.agent.tool.ToolSpecification;
 import dev.langchain4j.data.embedding.Embedding;
-import dev.langchain4j.data.message.AiMessage;
-import dev.langchain4j.data.message.ChatMessage;
 import dev.langchain4j.data.message.SystemMessage;
-import dev.langchain4j.data.message.ToolExecutionResultMessage;
 import dev.langchain4j.data.message.UserMessage;
-import dev.langchain4j.data.segment.TextSegment;
-import dev.langchain4j.internal.Utils;
+import dev.langchain4j.data.message.*;
 import dev.langchain4j.model.output.FinishReason;
 import dev.langchain4j.model.output.TokenUsage;
-import dev.langchain4j.model.zhipu.chat.AssistantMessage;
-import dev.langchain4j.model.zhipu.chat.ChatCompletionResponse;
-import dev.langchain4j.model.zhipu.chat.Function;
-import dev.langchain4j.model.zhipu.chat.FunctionCall;
-import dev.langchain4j.model.zhipu.chat.Message;
-import dev.langchain4j.model.zhipu.chat.Parameters;
-import dev.langchain4j.model.zhipu.chat.Tool;
-import dev.langchain4j.model.zhipu.chat.ToolCall;
-import dev.langchain4j.model.zhipu.chat.ToolMessage;
-import dev.langchain4j.model.zhipu.chat.ToolType;
+import dev.langchain4j.model.zhipu.chat.*;
 import dev.langchain4j.model.zhipu.embedding.EmbeddingResponse;
 import dev.langchain4j.model.zhipu.shared.Usage;
 
@@ -32,29 +19,16 @@
 
 import static dev.langchain4j.internal.Exceptions.illegalArgument;
 import static dev.langchain4j.internal.Utils.isNullOrEmpty;
-import static dev.langchain4j.model.output.FinishReason.LENGTH;
-import static dev.langchain4j.model.output.FinishReason.OTHER;
-import static dev.langchain4j.model.output.FinishReason.STOP;
-import static dev.langchain4j.model.output.FinishReason.TOOL_EXECUTION;
+import static dev.langchain4j.model.output.FinishReason.*;
 
 class DefaultZhipuAiHelper {
 
-    public static List<Embedding> toEmbed(EmbeddingResponse response) {
-        return response.getData().stream()
+    public static List<Embedding> toEmbed(List<EmbeddingResponse> response) {
+        return response.stream()
                 .map(zhipuAiEmbedding -> Embedding.from(zhipuAiEmbedding.getEmbedding()))
                 .collect(Collectors.toList());
     }
 
-    public static String toEmbedTexts(List<TextSegment> textSegments) {
-        List<String> embedText = textSegments.stream()
-                .map(TextSegment::text)
-                .collect(Collectors.toList());
-        if (Utils.isNullOrEmpty(embedText)) {
-            return null;
-        }
-        return embedText.get(0);
-    }
-
     public static List<Tool> toTools(List<ToolSpecification> toolSpecifications) {
         return toolSpecifications.stream()
                 .map(toolSpecification -> Tool.from(toFunction(toolSpecification)))
@@ -95,7 +69,7 @@ private static Message toZhipuAiMessage(ChatMessage message) {
         if (message instanceof UserMessage) {
             UserMessage userMessage = (UserMessage) message;
             return dev.langchain4j.model.zhipu.chat.UserMessage.builder()
-                    .content(userMessage.text())
+                    .content(userMessage.singleText())
                     .build();
         }
 
@@ -137,13 +111,12 @@ private static Message toZhipuAiMessage(ChatMessage message) {
     }
 
     public static AiMessage aiMessageFrom(ChatCompletionResponse response) {
-        Message message = response.getChoices().get(0).getMessage();
-        AssistantMessage assistantMessage = (AssistantMessage) message;
-        if (isNullOrEmpty(assistantMessage.getToolCalls())) {
-            return AiMessage.from(assistantMessage.getContent());
+        AssistantMessage message = response.getChoices().get(0).getMessage();
+        if (isNullOrEmpty(message.getToolCalls())) {
+            return AiMessage.from(message.getContent());
         }
 
-        return AiMessage.from(specificationsFrom(assistantMessage.getToolCalls()));
+        return AiMessage.from(specificationsFrom(message.getToolCalls()));
     }
 
     public static List<ToolExecutionRequest> specificationsFrom(List<ToolCall> toolCalls) {
@@ -160,6 +133,19 @@ public static List<ToolExecutionRequest> specificationsFrom(List<ToolCall> toolC
         return specifications;
     }
 
+    public static Usage getEmbeddingUsage(List<EmbeddingResponse> responses) {
+        Usage tokenUsage = Usage.builder()
+                .completionTokens(0)
+                .promptTokens(0)
+                .totalTokens(0)
+                .build();
+
+        for (EmbeddingResponse response : responses) {
+            tokenUsage.add(response.getUsage());
+        }
+        return tokenUsage;
+    }
+
 
     public static TokenUsage tokenUsageFrom(Usage zhipuUsage) {
         if (zhipuUsage == null) {
diff --git a/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/ZhipuAiEmbeddingModel.java b/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/ZhipuAiEmbeddingModel.java
index e338db68424..50128634028 100644
--- a/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/ZhipuAiEmbeddingModel.java
+++ b/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/ZhipuAiEmbeddingModel.java
@@ -6,16 +6,17 @@
 import dev.langchain4j.model.output.Response;
 import dev.langchain4j.model.zhipu.embedding.EmbeddingRequest;
 import dev.langchain4j.model.zhipu.embedding.EmbeddingResponse;
+import dev.langchain4j.model.zhipu.shared.Usage;
 import dev.langchain4j.model.zhipu.spi.ZhipuAiEmbeddingModelBuilderFactory;
 import lombok.Builder;
 
+import java.util.LinkedList;
 import java.util.List;
+import java.util.stream.Collectors;
 
 import static dev.langchain4j.internal.RetryUtils.withRetry;
 import static dev.langchain4j.internal.Utils.getOrDefault;
-import static dev.langchain4j.model.zhipu.DefaultZhipuAiHelper.toEmbed;
-import static dev.langchain4j.model.zhipu.DefaultZhipuAiHelper.toEmbedTexts;
-import static dev.langchain4j.model.zhipu.DefaultZhipuAiHelper.tokenUsageFrom;
+import static dev.langchain4j.model.zhipu.DefaultZhipuAiHelper.*;
 import static dev.langchain4j.spi.ServiceHelper.loadFactories;
 
 /**
@@ -58,16 +59,20 @@ public static ZhipuAiEmbeddingModelBuilder builder() {
     @Override
     public Response<List<Embedding>> embedAll(List<TextSegment> textSegments) {
 
-        EmbeddingRequest request = EmbeddingRequest.builder()
-                .model(this.model)
-                .input(toEmbedTexts(textSegments))
-                .build();
+        List<EmbeddingResponse> embeddingRequests = textSegments.stream()
+                .map(item -> EmbeddingRequest.builder()
+                        .input(item.text())
+                        .model(this.model)
+                        .build()
+                )
+                .map(request -> withRetry(() -> client.embedAll(request), maxRetries))
+                .collect(Collectors.toList());
 
-        EmbeddingResponse response = withRetry(() -> client.embedAll(request), maxRetries);
+        Usage usage = getEmbeddingUsage(embeddingRequests);
 
         return Response.from(
-                toEmbed(response),
-                tokenUsageFrom(response.getUsage())
+                toEmbed(embeddingRequests),
+                tokenUsageFrom(usage)
         );
     }
 
diff --git a/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/embedding/EmbeddingRequest.java b/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/embedding/EmbeddingRequest.java
index fb15274178a..a830c061d1b 100644
--- a/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/embedding/EmbeddingRequest.java
+++ b/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/embedding/EmbeddingRequest.java
@@ -1,32 +1,23 @@
 package dev.langchain4j.model.zhipu.embedding;
 
-import java.util.List;
-import java.util.Objects;
+import lombok.Getter;
 
-import static java.util.Arrays.asList;
-import static java.util.Collections.unmodifiableList;
+import java.util.Objects;
 
+@Getter
 public final class EmbeddingRequest {
     private final String input;
     private final String model;
 
     private EmbeddingRequest(Builder builder) {
         this.model = builder.model;
-        this.input = builder.input.get(0);
+        this.input = builder.input;
     }
 
     public static Builder builder() {
         return new Builder();
     }
 
-    public String getInput() {
-        return input;
-    }
-
-    public String getModel() {
-        return model;
-    }
-
     @Override
     public boolean equals(Object another) {
         if (this == another) return true;
@@ -58,7 +49,7 @@ public String toString() {
     public static final class Builder {
 
         private String model = EmbeddingModel.EMBEDDING_2.toString();
-        private List<String> input;
+        private String input;
 
         private Builder() {
         }
@@ -72,18 +63,8 @@ public Builder model(String model) {
             return this;
         }
 
-        public Builder input(String... input) {
-            return input(asList(input));
-        }
-
-        public Builder input(List<String> input) {
-            if (input == null || input.isEmpty()) {
-                throw new RuntimeException();
-            }
-            if (input.size() > 1) {
-                throw new RuntimeException();
-            }
-            this.input = unmodifiableList(input);
+        public Builder input(String input) {
+            this.input = input;
             return this;
         }
 
diff --git a/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/shared/Usage.java b/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/shared/Usage.java
index 1dfda68bd21..75b6e91e06e 100644
--- a/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/shared/Usage.java
+++ b/langchain4j-zhipu-ai/src/main/java/dev/langchain4j/model/zhipu/shared/Usage.java
@@ -26,6 +26,12 @@ public static Builder builder() {
         return new Builder();
     }
 
+    public void add(Usage usage) {
+        this.promptTokens += usage.promptTokens;
+        this.completionTokens += usage.completionTokens;
+        this.totalTokens += usage.totalTokens;
+    }
+
     public static final class Builder {
 
         private Integer promptTokens;