package dev.langchain4j.model.vertexai;

import com.google.cloud.vertexai.api.FunctionCall;
import com.google.cloud.vertexai.api.FunctionDeclaration;
import com.google.cloud.vertexai.api.Schema;
import com.google.cloud.vertexai.api.Tool;
import com.google.cloud.vertexai.api.Type;
import com.google.protobuf.ListValue;
import com.google.protobuf.NullValue;
import com.google.protobuf.Struct;
import com.google.protobuf.Value;
import dev.langchain4j.agent.tool.JsonSchemaProperty;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import org.junit.jupiter.api.Test;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;

import static dev.langchain4j.model.vertexai.FunctionCallHelper.unwrapProtoValue;
import static org.assertj.core.api.Assertions.assertThat;

public class FunctionCallHelperTest {
    @Test
    void should_unwrap_proto_values() {
        // check basic values
        assertThat(unwrapProtoValue(Value.newBuilder().setStringValue("hello").build())).isEqualTo("hello");
        assertThat(unwrapProtoValue(Value.newBuilder().setBoolValue(false).build())).isEqualTo(false);
        assertThat(unwrapProtoValue(Value.newBuilder().setNumberValue(1.23).build())).isEqualTo(1.23);
        assertThat(unwrapProtoValue(Value.newBuilder().setNullValue(NullValue.NULL_VALUE).build())).isEqualTo(null);

        // check list unwrapping
        ListValue listValue = ListValue.newBuilder()
            .addValues(Value.newBuilder().setStringValue("hello"))
            .addValues(Value.newBuilder().setBoolValue(true))
            .addValues(Value.newBuilder().setNumberValue(3.14))
            .build();
        assertThat(unwrapProtoValue(Value.newBuilder().setListValue(listValue).build())).isEqualTo(
            Arrays.asList("hello", true, 3.14)
        );

        // check struct unwrapping
        Struct struct = Struct.newBuilder()
            .putFields("name", Value.newBuilder().setStringValue("Guillaume").build())
            .putFields("numberOfKids", Value.newBuilder().setNumberValue(2).build())
            .putFields("kids", Value.newBuilder().setListValue(ListValue.newBuilder()
                .addValues(Value.newBuilder().setStringValue("Marion").build())
                .addValues(Value.newBuilder().setStringValue("Érine").build())
                .build()).build())
            .putFields("flag", Value.newBuilder().setBoolValue(false).build())
            .build();
        HashMap<Object, Object> map = new HashMap<>();
        map.put("name", "Guillaume");
        map.put("numberOfKids", 2.0);
        map.put("kids", Arrays.asList("Marion", "Érine"));
        map.put("flag", false);
        assertThat(unwrapProtoValue(Value.newBuilder().setStructValue(struct).build())).isEqualTo(map);
    }

    @Test
    void should_convert_tool_specs() {
        // given
        ToolSpecification toolSpec = ToolSpecification.builder()
            .description("Give the weather forecast for a location")
            .name("getWeatherForecast")
            .addParameter("location", JsonSchemaProperty.STRING,
                JsonSchemaProperty.description("the location to get the weather forecast for"))
            .addOptionalParameter("days", JsonSchemaProperty.INTEGER,
                JsonSchemaProperty.description("the number of days in the forecast"))
            .build();

        // when
        Tool tool = FunctionCallHelper.convertToolSpecifications(Collections.singletonList(toolSpec));

        // then
        assertThat(tool.getFunctionDeclarationsCount()).isEqualTo(1);

        FunctionDeclaration funDecl = tool.getFunctionDeclarations(0);
        assertThat(funDecl.getDescription()).isEqualTo("Give the weather forecast for a location");
        assertThat(funDecl.getName()).isEqualTo("getWeatherForecast");

        Schema parameters = funDecl.getParameters();
        assertThat(parameters.getPropertiesCount()).isEqualTo(2);
        assertThat(parameters.getPropertiesMap().get("location").getType()).isEqualTo(Type.STRING);
        assertThat(parameters.getPropertiesMap().get("location").getDescription())
            .isEqualTo("the location to get the weather forecast for");
        assertThat(parameters.getRequiredCount()).isEqualTo(1);
        assertThat(parameters.getRequired(0)).isEqualTo("location");
        assertThat(parameters.getPropertiesMap().get("days").getType()).isEqualTo(Type.INTEGER);
    }

    @Test
    void should_convert_function_calls_to_tool_execution_requests_and_back() {
        // given
        FunctionCall functionCall = FunctionCall.newBuilder()
            .setName("getWeatherForecast")
            .setArgs(
                Struct.newBuilder()
                    .putFields("location", Value.newBuilder().setStringValue("Paris").build())
                    .build()
            )
            .build();

        // when
        List<ToolExecutionRequest> toolExecutionRequest =
            FunctionCallHelper.fromFunctionCalls(Collections.singletonList(functionCall));
        FunctionCall sameFunctionCall = FunctionCallHelper.fromToolExecutionRequest(toolExecutionRequest.get(0));

        // then
        assertThat(functionCall).isEqualTo(sameFunctionCall);

        // given
        ToolExecutionRequest newExecutionRequest = ToolExecutionRequest.builder()
            .name("getWeatherForecast")
            .arguments("{\"location\":\"Paris\"}")
            .build();

        // when
        FunctionCall newFunctionCall = FunctionCallHelper.fromToolExecutionRequest(newExecutionRequest);
        ToolExecutionRequest sameExecutionRequest =
            FunctionCallHelper.fromFunctionCalls(Collections.singletonList(newFunctionCall)).get(0);

        // then
        assertThat(newExecutionRequest).isEqualTo(sameExecutionRequest);
    }
}

dev.langchain4j.agent.tool.ToolSpecification.builder dev.langchain4j.agent.tool.ToolExecutionRequest.builder
[((1025, 1075), 'com.google.protobuf.Value.newBuilder'), ((1025, 1067), 'com.google.protobuf.Value.newBuilder'), ((1134, 1180), 'com.google.protobuf.Value.newBuilder'), ((1134, 1172), 'com.google.protobuf.Value.newBuilder'), ((1237, 1284), 'com.google.protobuf.Value.newBuilder'), ((1237, 1276), 'com.google.protobuf.Value.newBuilder'), ((1340, 1401), 'com.google.protobuf.Value.newBuilder'), ((1340, 1393), 'com.google.protobuf.Value.newBuilder'), ((1485, 1721), 'com.google.protobuf.ListValue.newBuilder'), ((1485, 1700), 'com.google.protobuf.ListValue.newBuilder'), ((1485, 1636), 'com.google.protobuf.ListValue.newBuilder'), ((1485, 1574), 'com.google.protobuf.ListValue.newBuilder'), ((1531, 1573), 'com.google.protobuf.Value.newBuilder'), ((1598, 1635), 'com.google.protobuf.Value.newBuilder'), ((1660, 1699), 'com.google.protobuf.Value.newBuilder'), ((1759, 1809), 'com.google.protobuf.Value.newBuilder'), ((1759, 1801), 'com.google.protobuf.Value.newBuilder'), ((1941, 2513), 'com.google.protobuf.Struct.newBuilder'), ((1941, 2492), 'com.google.protobuf.Struct.newBuilder'), ((1941, 2413), 'com.google.protobuf.Struct.newBuilder'), ((1941, 2132), 'com.google.protobuf.Struct.newBuilder'), ((1941, 2047), 'com.google.protobuf.Struct.newBuilder'), ((1992, 2046), 'com.google.protobuf.Value.newBuilder'), ((1992, 2038), 'com.google.protobuf.Value.newBuilder'), ((2087, 2131), 'com.google.protobuf.Value.newBuilder'), ((2087, 2123), 'com.google.protobuf.Value.newBuilder'), ((2164, 2412), 'com.google.protobuf.Value.newBuilder'), ((2164, 2404), 'com.google.protobuf.Value.newBuilder'), ((2196, 2403), 'com.google.protobuf.ListValue.newBuilder'), ((2196, 2378), 'com.google.protobuf.ListValue.newBuilder'), ((2196, 2298), 'com.google.protobuf.ListValue.newBuilder'), ((2246, 2297), 'com.google.protobuf.Value.newBuilder'), ((2246, 2289), 'com.google.protobuf.Value.newBuilder'), ((2326, 2377), 'com.google.protobuf.Value.newBuilder'), ((2326, 2369), 'com.google.protobuf.Value.newBuilder'), ((2445, 2491), 'com.google.protobuf.Value.newBuilder'), ((2445, 2483), 'com.google.protobuf.Value.newBuilder'), ((2774, 2823), 'com.google.protobuf.Value.newBuilder'), ((2774, 2815), 'com.google.protobuf.Value.newBuilder'), ((2952, 3426), 'dev.langchain4j.agent.tool.ToolSpecification.builder'), ((2952, 3405), 'dev.langchain4j.agent.tool.ToolSpecification.builder'), ((2952, 3249), 'dev.langchain4j.agent.tool.ToolSpecification.builder'), ((2952, 3088), 'dev.langchain4j.agent.tool.ToolSpecification.builder'), ((2952, 3048), 'dev.langchain4j.agent.tool.ToolSpecification.builder'), ((4632, 4917), 'com.google.cloud.vertexai.api.FunctionCall.newBuilder'), ((4632, 4896), 'com.google.cloud.vertexai.api.FunctionCall.newBuilder'), ((4632, 4700), 'com.google.cloud.vertexai.api.FunctionCall.newBuilder'), ((4739, 4882), 'com.google.protobuf.Struct.newBuilder'), ((4739, 4853), 'com.google.protobuf.Struct.newBuilder'), ((4802, 4852), 'com.google.protobuf.Value.newBuilder'), ((4802, 4844), 'com.google.protobuf.Value.newBuilder'), ((5347, 5489), 'dev.langchain4j.agent.tool.ToolExecutionRequest.builder'), ((5347, 5468), 'dev.langchain4j.agent.tool.ToolExecutionRequest.builder'), ((5347, 5417), 'dev.langchain4j.agent.tool.ToolExecutionRequest.builder')]