package dev.ai4j.openai4j.chat;

import dev.ai4j.openai4j.DefaultOpenAiClient;
import dev.ai4j.openai4j.OpenAiClient;
import dev.ai4j.openai4j.RateLimitAwareTest;
import dev.ai4j.openai4j.ResponseHandle;
import dev.ai4j.openai4j.shared.StreamOptions;
import dev.ai4j.openai4j.shared.Usage;
import okhttp3.Dispatcher;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.EnumSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

import static dev.ai4j.openai4j.chat.ChatCompletionModel.GPT_4O;
import static dev.ai4j.openai4j.chat.ChatCompletionTest.*;
import static dev.ai4j.openai4j.chat.FunctionCallUtil.argument;
import static dev.ai4j.openai4j.chat.FunctionCallUtil.argumentsAsMap;
import static dev.ai4j.openai4j.chat.ResponseFormatType.JSON_OBJECT;
import static dev.ai4j.openai4j.chat.ResponseFormatType.TEXT;
import static dev.ai4j.openai4j.chat.ToolType.FUNCTION;
import static java.util.Collections.singletonMap;
import static java.util.concurrent.Executors.newSingleThreadExecutor;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.params.provider.EnumSource.Mode.EXCLUDE;

class ChatCompletionStreamingTest extends RateLimitAwareTest {

    private static final Logger log = LoggerFactory.getLogger(ChatCompletionStreamingTest.class);
    private final OpenAiClient client = OpenAiClient.builder()
            .baseUrl(System.getenv("OPENAI_BASE_URL"))
            .openAiApiKey(System.getenv("OPENAI_API_KEY"))
            .logRequests()
            .logResponses()
            .logStreamingResponses()
            .dispatcher(getCustomDispatcher())
            .build();

    private Dispatcher getCustomDispatcher() {
        return new Dispatcher(
            new ThreadPoolExecutor(0, 100, 60, TimeUnit.SECONDS,
                new SynchronousQueue<>(),
                runnable -> {
                    Thread thread = new Thread(runnable, "CustomDispatcher");
                    log.debug("CustomDispatcher thread created, id={}", thread.getId());
                    return thread;
                }
            )
        );
    }

    @Test
    void testSimpleApi() throws Exception {

        // when
        StringBuilder responseBuilder = new StringBuilder();
        CompletableFuture<String> future = new CompletableFuture<>();

        client.chatCompletion(USER_MESSAGE)
                .onPartialResponse(responseBuilder::append)
                .onComplete(() -> {
                    log.debug("Current thread id={}", Thread.currentThread().getId());
                    future.complete(responseBuilder.toString());
                })
                .onError(future::completeExceptionally)
                .execute();

        String response = future.get(30, SECONDS);

        // then
        assertThat(response).containsIgnoringCase("hello world");
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
    })
    void testCustomizableApi(ChatCompletionModel model) throws Exception {

        // given
        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .addSystemMessage(SYSTEM_MESSAGE)
                .addUserMessage(USER_MESSAGE)
                .temperature(1.0)
                .topP(0.1)
                .n(1)
                .stream(false) // intentionally setting to false in order to test that it is ignored
                .streamOptions(StreamOptions.builder()
                        .includeUsage(true)
                        .build())
                .stop("one", "two")
                .maxCompletionTokens(3)
                .presencePenalty(0.0)
                .frequencyPenalty(0.0)
                .logitBias(singletonMap("50256", -100))
                .user("Klaus")
                .responseFormat(TEXT)
                .seed(42)
                .build();

        // when
        StringBuilder responseBuilder = new StringBuilder();
        CompletableFuture<String> future = new CompletableFuture<>();
        AtomicReference<Usage> usageReference = new AtomicReference<>();

        client.chatCompletion(request)
                .onPartialResponse(partialResponse -> {
                    for (ChatCompletionChoice choice : partialResponse.choices()) {
                        String content = choice.delta().content();
                        if (content != null) {
                            responseBuilder.append(content);
                        }
                    }
                    Usage usage = partialResponse.usage();
                    if (usage != null) {
                        usageReference.set(usage);
                    }
                })
                .onComplete(() -> future.complete(responseBuilder.toString()))
                .onError(future::completeExceptionally)
                .execute();

        String response = future.get(30, SECONDS);

        // then
        assertThat(response).containsIgnoringCase("hello world");

        Usage usage = usageReference.get();
        assertThat(usage.promptTokens()).isGreaterThan(0);
        assertThat(usage.completionTokens()).isGreaterThan(0);
        assertThat(usage.completionTokensDetails().reasoningTokens()).isEqualTo(0);
        assertThat(usage.totalTokens()).isEqualTo(usage.promptTokens() + usage.completionTokens());
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
    })
    void testTools(ChatCompletionModel model) throws Exception {

        // given
        UserMessage userMessage = UserMessage.from("What is the weather in Boston in Celsius?");

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage)
                .tools(WEATHER_TOOL)
                .build();

        // when
        StringBuilder toolIdBuilder = new StringBuilder();
        StringBuilder functionNameBuilder = new StringBuilder();
        StringBuilder functionArgumentsBuilder = new StringBuilder();
        CompletableFuture<AssistantMessage> future = new CompletableFuture<>();

        client.chatCompletion(request)
                .onPartialResponse(partialResponse -> {
                    Delta delta = partialResponse.choices().get(0).delta();
                    assertThat(delta.content()).isNullOrEmpty();
                    assertThat(delta.functionCall()).isNull();

                    if (delta.toolCalls() != null) {
                        assertThat(delta.toolCalls()).hasSize(1);

                        ToolCall toolCall = delta.toolCalls().get(0);
                        assertThat(toolCall.type()).isIn(null, FUNCTION);
                        assertThat(toolCall.function()).isNotNull();

                        if (toolCall.id() != null) {
                            toolIdBuilder.append(toolCall.id());
                        }

                        FunctionCall functionCall = toolCall.function();
                        if (functionCall.name() != null) {
                            functionNameBuilder.append(functionCall.name());
                        }
                        if (functionCall.arguments() != null) {
                            functionArgumentsBuilder.append(functionCall.arguments());
                        }
                    }
                })
                .onComplete(() -> {
                    AssistantMessage assistantMessage = AssistantMessage.builder()
                            .toolCalls(ToolCall.builder()
                                    .id(toolIdBuilder.toString())
                                    .type(FUNCTION)
                                    .function(FunctionCall.builder()
                                            .name(functionNameBuilder.toString())
                                            .arguments(functionArgumentsBuilder.toString())
                                            .build())
                                    .build())
                            .build();
                    future.complete(assistantMessage);
                })
                .onError(future::completeExceptionally)
                .execute();

        AssistantMessage assistantMessage = future.get(30, SECONDS);

        // then
        assertThat(assistantMessage.content()).isNull();
        assertThat(assistantMessage.functionCall()).isNull();
        assertThat(assistantMessage.toolCalls()).isNotNull().hasSize(1);

        ToolCall toolCall = assistantMessage.toolCalls().get(0);
        assertThat(toolCall.id()).isNotBlank();
        assertThat(toolCall.type()).isEqualTo(FUNCTION);
        assertThat(toolCall.function()).isNotNull();

        FunctionCall functionCall = toolCall.function();
        assertThat(functionCall.name()).isEqualTo(WEATHER_TOOL_NAME);
        assertThat(functionCall.arguments()).isNotBlank();

        Map<String, Object> arguments = argumentsAsMap(functionCall.arguments());
        assertThat(arguments).hasSizeBetween(1, 2);
        assertThat(arguments.get("location").toString()).contains("Boston");

        // given
        String location = argument("location", functionCall);
        String unit = argument("unit", functionCall);
        String currentWeather = currentWeather(location, unit);
        ToolMessage toolMessage = ToolMessage.from(toolCall.id(), currentWeather);

        ChatCompletionRequest secondRequest = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage, assistantMessage, toolMessage)
                .build();

        // when
        StringBuilder responseBuilder = new StringBuilder();
        CompletableFuture<String> secondFuture = new CompletableFuture<>();

        client.chatCompletion(secondRequest)
                .onPartialResponse(partialResponse -> {
                    Delta delta = partialResponse.choices().get(0).delta();
                    if (delta.content() != null) {
                        responseBuilder.append(delta.content());
                    }
                })
                .onComplete(() -> secondFuture.complete(responseBuilder.toString()))
                .onError(secondFuture::completeExceptionally)
                .execute();

        String response = secondFuture.get(30, SECONDS);

        // then
        assertThat(response).contains("11");
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
    })
    void testFunctions(ChatCompletionModel model) throws Exception {

        // given
        UserMessage userMessage = UserMessage.from("What is the weather in Boston in Celsius?");

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage)
                .functions(WEATHER_FUNCTION)
                .build();

        // when
        StringBuilder functionNameBuilder = new StringBuilder();
        StringBuilder functionArgumentsBuilder = new StringBuilder();
        CompletableFuture<AssistantMessage> future = new CompletableFuture<>();

        client.chatCompletion(request)
                .onPartialResponse(partialResponse -> {
                    Delta delta = partialResponse.choices().get(0).delta();
                    assertThat(delta.content()).isNullOrEmpty();
                    assertThat(delta.toolCalls()).isNull();

                    if (delta.functionCall() != null) {
                        FunctionCall functionCall = delta.functionCall();
                        if (functionCall.name() != null) {
                            functionNameBuilder.append(functionCall.name());
                        }
                        if (functionCall.arguments() != null) {
                            functionArgumentsBuilder.append(functionCall.arguments());
                        }
                    }
                })
                .onComplete(() -> {
                    AssistantMessage assistantMessage = AssistantMessage.builder()
                            .functionCall(FunctionCall.builder()
                                    .name(functionNameBuilder.toString())
                                    .arguments(functionArgumentsBuilder.toString())
                                    .build())
                            .build();
                    future.complete(assistantMessage);
                })
                .onError(future::completeExceptionally)
                .execute();

        AssistantMessage assistantMessage = future.get(30, SECONDS);

        // then
        assertThat(assistantMessage.content()).isNull();
        assertThat(assistantMessage.toolCalls()).isNull();

        FunctionCall functionCall = assistantMessage.functionCall();
        assertThat(functionCall.name()).isEqualTo(WEATHER_TOOL_NAME);
        assertThat(functionCall.arguments()).isNotBlank();

        Map<String, Object> arguments = argumentsAsMap(functionCall.arguments());
        assertThat(arguments).hasSizeBetween(1, 2);
        assertThat(arguments.get("location").toString()).contains("Boston");

        // given
        String location = argument("location", functionCall);
        String unit = argument("unit", functionCall);
        String currentWeather = currentWeather(location, unit);
        FunctionMessage functionMessage = FunctionMessage.from(functionCall.name(), currentWeather);

        ChatCompletionRequest secondRequest = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage, assistantMessage, functionMessage)
                .build();

        // when
        StringBuilder responseBuilder = new StringBuilder();
        CompletableFuture<String> secondFuture = new CompletableFuture<>();

        client.chatCompletion(secondRequest)
                .onPartialResponse(partialResponse -> {
                    Delta delta = partialResponse.choices().get(0).delta();
                    if (delta.content() != null) {
                        responseBuilder.append(delta.content());
                    }
                })
                .onComplete(() -> secondFuture.complete(responseBuilder.toString()))
                .onError(secondFuture::completeExceptionally)
                .execute();

        String response = secondFuture.get(30, SECONDS);

        // then
        assertThat(response).contains("11");
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
    })
    void testToolChoice(ChatCompletionModel model) throws Exception {

        // given
        UserMessage userMessage = UserMessage.from("What is the weather in Boston in Celsius?");

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage)
                .tools(WEATHER_TOOL)
                .toolChoice(WEATHER_TOOL_NAME)
                .build();

        // when
        StringBuilder toolIdBuilder = new StringBuilder();
        StringBuilder functionNameBuilder = new StringBuilder();
        StringBuilder functionArgumentsBuilder = new StringBuilder();
        CompletableFuture<AssistantMessage> future = new CompletableFuture<>();

        client.chatCompletion(request)
                .onPartialResponse(partialResponse -> {
                    Delta delta = partialResponse.choices().get(0).delta();
                    assertThat(delta.content()).isNull();
                    assertThat(delta.functionCall()).isNull();

                    if (delta.toolCalls() != null) {
                        assertThat(delta.toolCalls()).hasSize(1);

                        ToolCall toolCall = delta.toolCalls().get(0);
                        assertThat(toolCall.type()).isIn(null, FUNCTION);
                        assertThat(toolCall.function()).isNotNull();

                        if (toolCall.id() != null) {
                            toolIdBuilder.append(toolCall.id());
                        }

                        FunctionCall functionCall = toolCall.function();
                        if (functionCall.name() != null) {
                            functionNameBuilder.append(functionCall.name());
                        }
                        if (functionCall.arguments() != null) {
                            functionArgumentsBuilder.append(functionCall.arguments());
                        }
                    }
                })
                .onComplete(() -> {
                    AssistantMessage assistantMessage = AssistantMessage.builder()
                            .toolCalls(ToolCall.builder()
                                    .id(toolIdBuilder.toString())
                                    .type(FUNCTION)
                                    .function(FunctionCall.builder()
                                            .name(functionNameBuilder.toString())
                                            .arguments(functionArgumentsBuilder.toString())
                                            .build())
                                    .build())
                            .build();
                    future.complete(assistantMessage);
                })
                .onError(future::completeExceptionally)
                .execute();

        AssistantMessage assistantMessage = future.get(30, SECONDS);

        // then
        assertThat(assistantMessage.content()).isNull();
        assertThat(assistantMessage.toolCalls()).isNotNull().hasSize(1);

        ToolCall toolCall = assistantMessage.toolCalls().get(0);
        assertThat(toolCall.id()).isNotBlank();
        assertThat(toolCall.type()).isEqualTo(FUNCTION);
        assertThat(toolCall.function()).isNotNull();

        FunctionCall functionCall = toolCall.function();
        assertThat(functionCall.name()).isEqualTo(WEATHER_TOOL_NAME);
        assertThat(functionCall.arguments()).isNotBlank();

        Map<String, Object> arguments = argumentsAsMap(functionCall.arguments());
        assertThat(arguments).hasSizeBetween(1, 2);
        assertThat(arguments.get("location").toString()).contains("Boston");

        // given
        String location = argument("location", functionCall);
        String unit = argument("unit", functionCall);
        String currentWeather = currentWeather(location, unit);
        ToolMessage toolMessage = ToolMessage.from(toolCall.id(), currentWeather);

        ChatCompletionRequest secondRequest = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage, assistantMessage, toolMessage)
                .build();

        // when
        StringBuilder responseBuilder = new StringBuilder();
        CompletableFuture<String> secondFuture = new CompletableFuture<>();

        client.chatCompletion(secondRequest)
                .onPartialResponse(partialResponse -> {
                    Delta delta = partialResponse.choices().get(0).delta();
                    if (delta.content() != null) {
                        responseBuilder.append(delta.content());
                    }
                })
                .onComplete(() -> secondFuture.complete(responseBuilder.toString()))
                .onError(secondFuture::completeExceptionally)
                .execute();

        String response = secondFuture.get(30, SECONDS);

        // then
        assertThat(response).contains("11");
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
    })
    void testFunctionChoice(ChatCompletionModel model) throws Exception {

        // given
        UserMessage userMessage = UserMessage.from("What is the weather in Boston in Celsius?");

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage)
                .functions(WEATHER_FUNCTION)
                .functionCall(WEATHER_TOOL_NAME)
                .build();

        // when
        StringBuilder functionNameBuilder = new StringBuilder();
        StringBuilder functionArgumentsBuilder = new StringBuilder();
        CompletableFuture<AssistantMessage> future = new CompletableFuture<>();

        client.chatCompletion(request)
                .onPartialResponse(partialResponse -> {
                    Delta delta = partialResponse.choices().get(0).delta();
                    assertThat(delta.content()).isNull();
                    assertThat(delta.toolCalls()).isNull();

                    if (delta.functionCall() != null) {
                        FunctionCall functionCall = delta.functionCall();
                        if (functionCall.name() != null) {
                            functionNameBuilder.append(functionCall.name());
                        }
                        if (functionCall.arguments() != null) {
                            functionArgumentsBuilder.append(functionCall.arguments());
                        }
                    }
                })
                .onComplete(() -> {
                    AssistantMessage assistantMessage = AssistantMessage.builder()
                            .functionCall(FunctionCall.builder()
                                    .name(functionNameBuilder.toString())
                                    .arguments(functionArgumentsBuilder.toString())
                                    .build())
                            .build();
                    future.complete(assistantMessage);
                })
                .onError(future::completeExceptionally)
                .execute();

        AssistantMessage assistantMessage = future.get(30, SECONDS);

        // then
        assertThat(assistantMessage.content()).isNull();
        assertThat(assistantMessage.toolCalls()).isNull();

        FunctionCall functionCall = assistantMessage.functionCall();
        assertThat(functionCall.name()).isEqualTo(WEATHER_TOOL_NAME);
        assertThat(functionCall.arguments()).isNotBlank();

        Map<String, Object> arguments = argumentsAsMap(functionCall.arguments());
        assertThat(arguments).hasSizeBetween(1, 2);
        assertThat(arguments.get("location").toString()).contains("Boston");

        // given
        String location = argument("location", functionCall);
        String unit = argument("unit", functionCall);
        String currentWeather = currentWeather(location, unit);
        FunctionMessage functionMessage = FunctionMessage.from(functionCall.name(), currentWeather);

        ChatCompletionRequest secondRequest = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage, assistantMessage, functionMessage)
                .build();

        // when
        StringBuilder responseBuilder = new StringBuilder();
        CompletableFuture<String> secondFuture = new CompletableFuture<>();

        client.chatCompletion(secondRequest)
                .onPartialResponse(partialResponse -> {
                    Delta delta = partialResponse.choices().get(0).delta();
                    if (delta.content() != null) {
                        responseBuilder.append(delta.content());
                    }
                })
                .onComplete(() -> secondFuture.complete(responseBuilder.toString()))
                .onError(secondFuture::completeExceptionally)
                .execute();

        String response = secondFuture.get(30, SECONDS);

        // then
        assertThat(response).contains("11");
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
            "GPT_4", "GPT_4_0613", // Does not support parallel tools
    })
    void testParallelTools(ChatCompletionModel model) throws Exception {

        // given
        UserMessage userMessage = UserMessage.from("What is the weather like in Boston and Madrid?");

        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage)
                .tools(WEATHER_TOOL)
                .build();

        // when
        StringBuilder firstToolIdBuilder = new StringBuilder();
        StringBuilder firstFunctionNameBuilder = new StringBuilder();
        StringBuilder firstFunctionArgumentsBuilder = new StringBuilder();

        StringBuilder secondToolIdBuilder = new StringBuilder();
        StringBuilder secondFunctionNameBuilder = new StringBuilder();
        StringBuilder secondFunctionArgumentsBuilder = new StringBuilder();

        CompletableFuture<AssistantMessage> future = new CompletableFuture<>();

        client.chatCompletion(request)
                .onPartialResponse(partialResponse -> {
                    Delta delta = partialResponse.choices().get(0).delta();
                    assertThat(delta.content()).isNull();
                    assertThat(delta.functionCall()).isNull();

                    if (delta.toolCalls() != null) {
                        assertThat(delta.toolCalls()).hasSize(1);

                        ToolCall toolCall = delta.toolCalls().get(0);
                        assertThat(toolCall.index()).isIn(0, 1);
                        assertThat(toolCall.type()).isIn(null, FUNCTION);
                        assertThat(toolCall.function()).isNotNull();

                        if (toolCall.id() != null) {
                            if (toolCall.index() == 0) {
                                firstToolIdBuilder.append(toolCall.id());
                            } else {
                                secondToolIdBuilder.append(toolCall.id());
                            }
                        }

                        FunctionCall functionCall = toolCall.function();
                        if (functionCall.name() != null) {
                            if (toolCall.index() == 0) {
                                firstFunctionNameBuilder.append(functionCall.name());
                            } else {
                                secondFunctionNameBuilder.append(functionCall.name());
                            }
                        }
                        if (functionCall.arguments() != null) {
                            if (toolCall.index() == 0) {
                                firstFunctionArgumentsBuilder.append(functionCall.arguments());
                            } else {
                                secondFunctionArgumentsBuilder.append(functionCall.arguments());
                            }
                        }
                    }
                })
                .onComplete(() -> {
                    AssistantMessage assistantMessage = AssistantMessage.builder()
                            .toolCalls(
                                    ToolCall.builder()
                                            .id(firstToolIdBuilder.toString())
                                            .type(FUNCTION)
                                            .function(FunctionCall.builder()
                                                    .name(firstFunctionNameBuilder.toString())
                                                    .arguments(firstFunctionArgumentsBuilder.toString())
                                                    .build())
                                            .build(),
                                    ToolCall.builder()
                                            .id(secondToolIdBuilder.toString())
                                            .type(FUNCTION)
                                            .function(FunctionCall.builder()
                                                    .name(secondFunctionNameBuilder.toString())
                                                    .arguments(secondFunctionArgumentsBuilder.toString())
                                                    .build())
                                            .build()
                            ).build();
                    future.complete(assistantMessage);
                })
                .onError(future::completeExceptionally)
                .execute();

        AssistantMessage assistantMessage = future.get(30, SECONDS);

        // then
        assertThat(assistantMessage.content()).isNull();
        assertThat(assistantMessage.toolCalls()).isNotNull().hasSize(2);

        ToolCall toolCall1 = assistantMessage.toolCalls().get(0);
        assertThat(toolCall1.id()).isNotBlank();
        assertThat(toolCall1.type()).isEqualTo(FUNCTION);
        assertThat(toolCall1.function()).isNotNull();

        FunctionCall functionCall1 = toolCall1.function();
        assertThat(functionCall1.name()).isEqualTo(WEATHER_TOOL_NAME);
        assertThat(functionCall1.arguments()).isNotBlank();

        Map<String, Object> arguments1 = argumentsAsMap(functionCall1.arguments());
        assertThat(arguments1).hasSizeBetween(1, 2);
        assertThat(arguments1.get("location").toString()).contains("Boston");

        ToolCall toolCall2 = assistantMessage.toolCalls().get(1);
        assertThat(toolCall2.id()).isNotBlank();
        assertThat(toolCall2.type()).isEqualTo(FUNCTION);
        assertThat(toolCall2.function()).isNotNull();

        FunctionCall functionCall2 = toolCall2.function();
        assertThat(functionCall2.name()).isEqualTo(WEATHER_TOOL_NAME);
        assertThat(functionCall2.arguments()).isNotBlank();

        Map<String, Object> arguments2 = argumentsAsMap(functionCall2.arguments());
        assertThat(arguments2).hasSizeBetween(1, 2);
        assertThat(arguments2.get("location").toString()).contains("Madrid");

        // given
        String location1 = argument("location", functionCall1);
        String unit1 = argument("unit", functionCall1);
        String currentWeather1 = currentWeather(location1, unit1);
        ToolMessage toolMessage1 = ToolMessage.from(toolCall1.id(), currentWeather1);

        String location2 = argument("location", functionCall2);
        String unit2 = argument("unit", functionCall2);
        String currentWeather2 = currentWeather(location2, unit2);
        ToolMessage toolMessage2 = ToolMessage.from(toolCall2.id(), currentWeather2);

        ChatCompletionRequest secondRequest = ChatCompletionRequest.builder()
                .model(model)
                .messages(userMessage, assistantMessage, toolMessage1, toolMessage2)
                .build();

        // when
        StringBuilder responseBuilder = new StringBuilder();
        CompletableFuture<String> secondFuture = new CompletableFuture<>();

        client.chatCompletion(secondRequest)
                .onPartialResponse(partialResponse -> {
                    Delta delta = partialResponse.choices().get(0).delta();
                    if (delta.content() != null) {
                        responseBuilder.append(delta.content());
                    }
                })
                .onComplete(() -> secondFuture.complete(responseBuilder.toString()))
                .onError(secondFuture::completeExceptionally)
                .execute();

        String response = secondFuture.get(30, SECONDS);

        // then
        assertThat(response).contains("11", "22");
    }

    @ParameterizedTest
    @EnumSource(value = ChatCompletionModel.class, mode = EXCLUDE, names = {
            "GPT_4_32K", "GPT_4_32K_0314", "GPT_4_32K_0613", // I don't have access to these models
            "GPT_4", "GPT_4_0613" // Does not support response_format
    })
    void testJsonResponseFormat(ChatCompletionModel model) throws Exception {

        // given
        ChatCompletionRequest request = ChatCompletionRequest.builder()
                .model(model)
                .addUserMessage("Extract information from the following text:\n" +
