Merge OpenAiStreamingChatModel and OpenAiChatModel classes

docs/docs/integrations/language-models/open-ai.md

@@ -73,14 +73,15 @@ OpenAiChatModel model = OpenAiChatModel.builder()
     .logResponses(...)
     .tokenizer(...)
     .customHeaders(...)
+    .isStreaming(...)
     .build();
 ```
 See the description of some of the parameters above [here](https://platform.openai.com/docs/api-reference/chat/create).
 
-## OpenAiStreamingChatModel
+## Streaming ChatModel
 
 ```java
-OpenAiStreamingChatModel model = OpenAiStreamingChatModel.withApiKey(System.getenv("OPENAI_API_KEY"));
+OpenAiChatModel model = OpenAiChatModel.withApiKey(System.getenv("OPENAI_API_KEY"), true);
 
 model.generate("Say 'Hello World'", new StreamingResponseHandler<AiMessage>() {
 

langchain4j-ollama/src/test/java/dev/langchain4j/model/ollama/OllamaOpenAiStreamingChatModelIT.java

@@ -4,7 +4,7 @@
 import dev.langchain4j.data.message.UserMessage;
 import dev.langchain4j.model.chat.StreamingChatLanguageModel;
 import dev.langchain4j.model.chat.TestStreamingResponseHandler;
-import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
+import dev.langchain4j.model.openai.OpenAiChatModel;
 import dev.langchain4j.model.output.Response;
 import org.junit.jupiter.api.Test;
 
@@ -18,13 +18,14 @@
  */
 class OllamaOpenAiStreamingChatModelIT extends AbstractOllamaLanguageModelInfrastructure {
 
-    StreamingChatLanguageModel model = OpenAiStreamingChatModel.builder()
+    StreamingChatLanguageModel model = OpenAiChatModel.builder()
             .apiKey("does not matter") // TODO make apiKey optional when using custom baseUrl?
             .baseUrl(ollama.getEndpoint() + "/v1") // TODO add "/v1" by default?
             .modelName(TINY_DOLPHIN_MODEL)
             .temperature(0.0)
             .logRequests(true)
             .logResponses(true)
+            .isStreaming(true)
             .build();
 
     @Test

langchain4j-open-ai/src/main/java/dev/langchain4j/model/openai/OpenAiChatModel.java

@@ -1,26 +1,30 @@
 package dev.langchain4j.model.openai;
 
 import dev.ai4j.openai4j.OpenAiClient;
+import dev.ai4j.openai4j.chat.ChatCompletionChoice;
 import dev.ai4j.openai4j.chat.ChatCompletionRequest;
 import dev.ai4j.openai4j.chat.ChatCompletionResponse;
+import dev.ai4j.openai4j.chat.Delta;
 import dev.langchain4j.agent.tool.ToolSpecification;
 import dev.langchain4j.data.message.AiMessage;
 import dev.langchain4j.data.message.ChatMessage;
+import dev.langchain4j.model.StreamingResponseHandler;
 import dev.langchain4j.model.Tokenizer;
 import dev.langchain4j.model.chat.ChatLanguageModel;
+import dev.langchain4j.model.chat.StreamingChatLanguageModel;
 import dev.langchain4j.model.chat.TokenCountEstimator;
 import dev.langchain4j.model.openai.spi.OpenAiChatModelBuilderFactory;
 import dev.langchain4j.model.output.Response;
 import lombok.Builder;
 
 import java.net.Proxy;
 import java.time.Duration;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 
 import static dev.langchain4j.internal.RetryUtils.withRetry;
 import static dev.langchain4j.internal.Utils.getOrDefault;
+import static dev.langchain4j.internal.Utils.isNullOrEmpty;
 import static dev.langchain4j.model.openai.InternalOpenAiHelper.*;
 import static dev.langchain4j.model.openai.OpenAiModelName.GPT_3_5_TURBO;
 import static dev.langchain4j.spi.ServiceHelper.loadFactories;
@@ -31,7 +35,7 @@
  * Represents an OpenAI language model with a chat completion interface, such as gpt-3.5-turbo and gpt-4.
  * You can find description of parameters <a href="https://platform.openai.com/docs/api-reference/chat/create">here</a>.
  */
-public class OpenAiChatModel implements ChatLanguageModel, TokenCountEstimator {
+public class OpenAiChatModel implements ChatLanguageModel, StreamingChatLanguageModel, TokenCountEstimator {
 
     private final OpenAiClient client;
     private final String modelName;
@@ -47,6 +51,7 @@ public class OpenAiChatModel implements ChatLanguageModel, TokenCountEstimator {
     private final String user;
     private final Integer maxRetries;
     private final Tokenizer tokenizer;
+    private final boolean isOpenAiModel;
 
     @Builder
     public OpenAiChatModel(String baseUrl,
@@ -69,7 +74,8 @@ public OpenAiChatModel(String baseUrl,
                            Boolean logRequests,
                            Boolean logResponses,
                            Tokenizer tokenizer,
-                           Map<String, String> customHeaders) {
+                           Map<String, String> customHeaders,
+                           boolean isStreaming) {
 
         baseUrl = getOrDefault(baseUrl, OPENAI_URL);
         if (OPENAI_DEMO_API_KEY.equals(apiKey)) {
@@ -78,7 +84,7 @@ public OpenAiChatModel(String baseUrl,
 
         timeout = getOrDefault(timeout, ofSeconds(60));
 
-        this.client = OpenAiClient.builder()
+        OpenAiClient.Builder openAiClientBuilder = OpenAiClient.builder()
                 .openAiApiKey(apiKey)
                 .baseUrl(baseUrl)
                 .organizationId(organizationId)
@@ -90,9 +96,18 @@ public OpenAiChatModel(String baseUrl,
                 .logRequests(logRequests)
                 .logResponses(logResponses)
                 .userAgent(DEFAULT_USER_AGENT)
-                .customHeaders(customHeaders)
-                .build();
+                .customHeaders(customHeaders);
+
+        if (isStreaming) {
+            openAiClientBuilder.logResponses(logResponses);
+        } else {
+            openAiClientBuilder.logStreamingResponses(logResponses);
+        }
+        this.client = openAiClientBuilder.build();
+
         this.modelName = getOrDefault(modelName, GPT_3_5_TURBO);
+        this.isOpenAiModel = isOpenAiModel(modelName);
+
         this.temperature = getOrDefault(temperature, 0.7);
         this.topP = topP;
         this.stop = stop;
@@ -113,22 +128,22 @@ public String modelName() {
 
     @Override
     public Response<AiMessage> generate(List<ChatMessage> messages) {
-        return generate(messages, null, null);
+        return generateMessage(messages, null, null);
     }
 
     @Override
     public Response<AiMessage> generate(List<ChatMessage> messages, List<ToolSpecification> toolSpecifications) {
-        return generate(messages, toolSpecifications, null);
+        return generateMessage(messages, toolSpecifications, null);
     }
 
     @Override
     public Response<AiMessage> generate(List<ChatMessage> messages, ToolSpecification toolSpecification) {
-        return generate(messages, singletonList(toolSpecification), toolSpecification);
+        return generateMessage(messages, singletonList(toolSpecification), toolSpecification);
     }
 
-    private Response<AiMessage> generate(List<ChatMessage> messages,
-                                         List<ToolSpecification> toolSpecifications,
-                                         ToolSpecification toolThatMustBeExecuted
+    private Response<AiMessage> generateMessage(List<ChatMessage> messages,
+                                                List<ToolSpecification> toolSpecifications,
+                                                ToolSpecification toolThatMustBeExecuted
     ) {
         ChatCompletionRequest.Builder requestBuilder = ChatCompletionRequest.builder()
                 .model(modelName)
@@ -168,7 +183,11 @@ public int estimateTokenCount(List<ChatMessage> messages) {
     }
 
     public static OpenAiChatModel withApiKey(String apiKey) {
-        return builder().apiKey(apiKey).build();
+        return withApiKey(apiKey, false);
+    }
+
+    public static OpenAiChatModel withApiKey(String apiKey, boolean isStreaming) {
+        return builder().apiKey(apiKey).isStreaming(isStreaming).build();
     }
 
     public static OpenAiChatModelBuilder builder() {
@@ -195,4 +214,92 @@ public OpenAiChatModelBuilder modelName(OpenAiChatModelName modelName) {
             return this;
         }
     }
+
+    @Override
+    public void generate(List<ChatMessage> messages, StreamingResponseHandler<AiMessage> handler) {
+        generate(messages, null, null, handler);
+    }
+
+    @Override
+    public void generate(List<ChatMessage> messages, List<ToolSpecification> toolSpecifications, StreamingResponseHandler<AiMessage> handler) {
+        generate(messages, toolSpecifications, null, handler);
+    }
+
+    @Override
+    public void generate(List<ChatMessage> messages, ToolSpecification toolSpecification, StreamingResponseHandler<AiMessage> handler) {
+        generate(messages, null, toolSpecification, handler);
+    }
+
+    private void generate(List<ChatMessage> messages,
+                          List<ToolSpecification> toolSpecifications,
+                          ToolSpecification toolThatMustBeExecuted,
+                          StreamingResponseHandler<AiMessage> handler
+    ) {
+        ChatCompletionRequest.Builder requestBuilder = ChatCompletionRequest.builder()
+                .stream(true)
+                .model(modelName)
+                .messages(toOpenAiMessages(messages))
+                .temperature(temperature)
+                .topP(topP)
+                .stop(stop)
+                .maxTokens(maxTokens)
+                .presencePenalty(presencePenalty)
+                .frequencyPenalty(frequencyPenalty)
+                .logitBias(logitBias)
+                .responseFormat(responseFormat)
+                .seed(seed)
+                .user(user);
+
+        if (toolThatMustBeExecuted != null) {
+            requestBuilder.tools(toTools(singletonList(toolThatMustBeExecuted)));
+            requestBuilder.toolChoice(toolThatMustBeExecuted.name());
+        } else if (!isNullOrEmpty(toolSpecifications)) {
+            requestBuilder.tools(toTools(toolSpecifications));
+        }
+
+        ChatCompletionRequest request = requestBuilder.build();
+
+        int inputTokenCount = countInputTokens(messages, toolSpecifications, toolThatMustBeExecuted);
+        OpenAiStreamingResponseBuilder responseBuilder = new OpenAiStreamingResponseBuilder(inputTokenCount);
+
+        client.chatCompletion(request)
+                .onPartialResponse(partialResponse -> {
+                    responseBuilder.append(partialResponse);
+                    handle(partialResponse, handler);
+                })
+                .onComplete(() -> {
+                    Response<AiMessage> response = responseBuilder.build(tokenizer, toolThatMustBeExecuted != null);
+                    if (!isOpenAiModel) {
+                        response = removeTokenUsage(response);
+                    }
+                    handler.onComplete(response);
+                })
+                .onError(handler::onError)
+                .execute();
+    }
+
+    private int countInputTokens(List<ChatMessage> messages,
+                                 List<ToolSpecification> toolSpecifications,
+                                 ToolSpecification toolThatMustBeExecuted) {
+        int inputTokenCount = tokenizer.estimateTokenCountInMessages(messages);
+        if (toolThatMustBeExecuted != null) {
+            inputTokenCount += tokenizer.estimateTokenCountInForcefulToolSpecification(toolThatMustBeExecuted);
+        } else if (!isNullOrEmpty(toolSpecifications)) {
+            inputTokenCount += tokenizer.estimateTokenCountInToolSpecifications(toolSpecifications);
+        }
+        return inputTokenCount;
+    }
+
+    private static void handle(ChatCompletionResponse partialResponse,
+                               StreamingResponseHandler<AiMessage> handler) {
+        List<ChatCompletionChoice> choices = partialResponse.choices();
+        if (choices == null || choices.isEmpty()) {
+            return;
+        }
+        Delta delta = choices.get(0).delta();
+        String content = delta.content();
+        if (content != null) {
+            handler.onNext(content);
+        }
+    }
 }