Llamadas de Funciones en Java y Spring AI Usando la API de Mistral AI

1. Visión general

Usando modelos de lenguaje grande, podemos recuperar mucha información útil. Podemos aprender muchos nuevos hechos sobre cualquier cosa y obtener respuestas basadas en datos existentes en Internet. Podemos pedirles que procesen datos de entrada y realicen diversas acciones. Pero, ¿qué pasa si pedimos al modelo que utilice una API para preparar la salida?

Para este propósito, podemos usar Llamadas de Función. Llamadas de función permiten a los LLM interactuar y manipular datos, realizar cálculos o recuperar información más allá de sus capacidades textuales inherentes.

En este artículo, exploraremos qué es la llamada de función y cómo podemos usarla para integrar los LLM con nuestra lógica interna. Como proveedor del modelo, usaremos la API de Mistral AI.

2. API de Mistral AI

Mistral AI se centra en proporcionar modelos generativos de IA abiertos y portátiles para desarrolladores y empresas. Podemos usarlo para prompts simples así como para integraciones de llamada a función.

2.1. Obtener la clave de API

Para comenzar a usar la API de Mistral, primero debemos recuperar la clave de API. Veamos la consola de gestión de claves API:

Para activar cualquier clave, debemos configurar la configuración de facturación o usar el período de prueba si está disponible:

Después de resolver todo, podemos pulsar el botón Crear nueva clave para obtener la clave de API de Mistral.

2.2. Ejemplo de uso

Empecemos con un prompt sencillo. Pediremos a la API de Mistral que nos devuelva una lista de estados de salud de pacientes. Implementemos una llamada así:

@Test
void givenHttpClient_whenSendTheRequestToChatAPI_thenShouldBeExpectedWordInResponse() throws IOException, InterruptedException {

    String apiKey = System.getenv("MISTRAL_API_KEY");
    String apiUrl = "https://api.mistral.ai/v1/chat/completions";
    String requestBody = "{"
      + "\"model\": \"mistral-large-latest\","
      + "\"messages\": [{\"role\": \"user\", "
      + "\"content\": \"What the patient health statuses can be?\"}]"
      + "}";

    HttpClient client = HttpClient.newHttpClient();
    HttpRequest request = HttpRequest.newBuilder()
      .uri(URI.create(apiUrl))
      .header("Content-Type", "application/json")
      .header("Accept", "application/json")
      .header("Authorization", "Bearer " + apiKey)
      .POST(HttpRequest.BodyPublishers.ofString(requestBody))
      .build();

    HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
    String responseBody = response.body();
    logger.info("Model response: " + responseBody);

    Assertions.assertThat(responseBody)
      .containsIgnoringCase("healthy");
}

Creamos una solicitud HTTP y la enviamos al endpoint /chat/completions. Luego, usamos la clave de API como valor del encabezado de autorización. Como era de esperar, en la respuesta vemos tanto los metadatos como el contenido mismo:

Model response: {"id":"585e3599275545c588cb0a502d1ab9e0","object":"chat.completion",
"created":1718308692,"model":"mistral-large-latest",
"choices":[{"index":0,"message":{"role":"assistant","content":"Patient health statuses can be
categorized in various ways, depending on the specific context or medical system being used.
However, some common health statuses include:
1.Healthy: The patient is in good health with no known medical issues.
...
10.Palliative: The patient is receiving care that is focused on relieving symptoms and improving quality of life, rather than curing the underlying disease.",
"tool_calls":null},"finish_reason":"stop","logprobs":null}],
"usage":{"prompt_tokens":12,"total_tokens":291,"completion_tokens":279}}

El ejemplo de llamada a función es más complejo y requiere mucha preparación antes de la llamada. Lo descubriremos en la siguiente sección.

3. Integración con Spring AI

Veamos algunos ejemplos de uso de la API de Mistral con llamadas a función. Con Spring AI podemos evitar mucho trabajo de preparación y dejar que el framework lo haga por nosotros.

3.1. Dependencias

La dependencia necesaria está en el repositorio Spring milestone. Añádelas a nuestro pom.xml:

<repositories>
    <repository>
        <id>spring-milestones</id>
        <name>Spring milestones</name>
        <url>https://repo.spring.io/milestone</url>
    </repository>
</repositories>

Ahora, añadamos la dependencia para la integración de la API de Mistral:

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-mistral-ai-spring-boot-starter</artifactId>
    <version>0.8.1</version>
</dependency>

3.2. Configuración

Añadamos ahora la clave de API que obtuvimos previamente al archivo de propiedades:

spring:
  ai:
    mistralai:
      api-key: ${MISTRAL_AI_API_KEY}
      chat:
        options:
          model: mistral-small-latest

Y eso es todo lo que necesitamos para comenzar a usar la API de Mistral.

3.3. Caso de uso con una función

En nuestro ejemplo de demostración, crearemos una función que devuelve el estado de salud del paciente según su ID.

Empecemos creando el registro de paciente record:

public record Patient(String patientId) {
}

Ahora creemos otro registro para el estado de salud del paciente:

public record HealthStatus(String status) {
}

En el siguiente paso, crearemos una clase de configuración:

@Configuration
public class MistralAIFunctionConfiguration {
    public static final Map<Patient, HealthStatus> HEALTH_DATA = Map.of(
      new Patient("P001"), new HealthStatus("Healthy"),
      new Patient("P002"), new HealthStatus("Has cough"),
      new Patient("P003"), new HealthStatus("Healthy"),
      new Patient("P004"), new HealthStatus("Has increased blood pressure"),
      new Patient("P005"), new HealthStatus("Healthy"));

    @Bean
    @Description("Get patient health status")
    public Function<Patient, HealthStatus> retrievePatientHealthStatus() {
        return (patient) -> new HealthStatus(HEALTH_DATA.get(patient).status());
    }
}

Aquí, hemos especificado el conjunto de datos con la información de salud de los pacientes. Además, creamos la función retrievePatientHealthStatus(), que devuelve el estado de salud para un ID de paciente dado.

Ahora probemos nuestra función llamándola dentro de una integración:

@Import(MistralAIFunctionConfiguration.class)
@ExtendWith(SpringExtension.class)
@SpringBootTest
public class MistralAIFunctionCallingManualTest {

    @Autowired
    private MistralAiChatModel chatClient;

    @Test
    void givenMistralAiChatClient_whenAskChatAPIAboutPatientHealthStatus_thenExpectedHealthStatusIsPresentInResponse() {

        var options = MistralAiChatOptions.builder()
          .withFunction("retrievePatientHealthStatus")
          .build();

        ChatResponse paymentStatusResponse = chatClient.call(
          new Prompt("What's the health status of the patient with id P004?",  options));

        String responseContent = paymentStatusResponse.getResult().getOutput().getContent();
        logger.info(responseContent);

        Assertions.assertThat(responseContent)
          .containsIgnoringCase("has increased blood pressure");
    }
}

Hemos importado nuestra clase MistralAIFunctionConfiguration para agregar nuestra función retrievePatientHealthStatus() al contexto de Spring del test. También inyectamos MistralAiChatClient, que será instanciado automáticamente por el starter de Spring AI.

En la solicitud a la API de chat, hemos especificado el texto del prompt que contiene uno de los ID de los pacientes y el nombre de la función para recuperar el estado de salud. Luego llamamos a la API y verificamos que la respuesta contenía el estado de salud esperado.

Además, registramos todo el texto de la respuesta, y esto es lo que vemos allí:

The patient with id P004 has increased blood pressure.

3.4. Caso de uso con múltiples funciones

También podemos especificar múltiples funciones y la IA decide cuál usar según el prompt que enviemos.

Para demostrarlo, extendamos nuestro registro HealthStatus:

public record HealthStatus(String status, LocalDate changeDate) {
}

Añadimos la fecha en que se cambió el estado por última vez.

Ahora modifiquemos la clase de configuración:

@Configuration
public class MistralAIFunctionConfiguration {
    public static final Map<Patient, HealthStatus> HEALTH_DATA = Map.of(
      new Patient("P001"), new HealthStatus("Healthy",
        LocalDate.of(2024,1, 20)),
      new Patient("P002"), new HealthStatus("Has cough",
        LocalDate.of(2024,3, 15)),
      new Patient("P003"), new HealthStatus("Healthy",
        LocalDate.of(2024,4, 12)),
      new Patient("P004"), new HealthStatus("Has increased blood pressure",
        LocalDate.of(2024,5, 19)),
      new Patient("P005"), new HealthStatus("Healthy",
        LocalDate.of(2024,6, 1)));

    @Bean
    @Description("Get patient health status")
    public Function<Patient, String> retrievePatientHealthStatus() {
        return (patient) -> HEALTH_DATA.get(patient).status();
    }

    @Bean
    @Description("Get when patient health status was updated")
    public Function<Patient, LocalDate> retrievePatientHealthStatusChangeDate() {
        return (patient) -> HEALTH_DATA.get(patient).changeDate();
    }
}

Hemos poblado las fechas de cambio para cada elemento de estado. También creamos la función retrievePatientHealthStatusChangeDate(), que devuelve información sobre la fecha de cambio del estado.

Veamos cómo podemos usar nuestras dos nuevas funciones con la API de Mistral:

@Test
void givenMistralAiChatClient_whenAskChatAPIAboutPatientHealthStatusAndWhenThisStatusWasChanged_thenExpectedInformationInResponse() {
    var options = MistralAiChatOptions.builder()
      .withFunctions(
        Set.of("retrievePatientHealthStatus",
          "retrievePatientHealthStatusChangeDate"))
      .build();

    ChatResponse paymentStatusResponse = chatClient.call(
      new Prompt(
        "What's the health status of the patient with id P005",
        options));

    String paymentStatusResponseContent = paymentStatusResponse.getResult()
      .getOutput().getContent();
    logger.info(paymentStatusResponseContent);

    Assertions.assertThat(paymentStatusResponseContent)
      .containsIgnoringCase("healthy");

    ChatResponse changeDateResponse = chatClient.call(
      new Prompt(
        "When health status of the patient with id P005 was changed?",
        options));

    String changeDateResponseContent = changeDateResponse.getResult().getOutput().getContent();
    logger.info(changeDateResponseContent);

    Assertions.assertThat(paymentStatusResponseContent)
      .containsIgnoringCase("June 1, 2024");
}

En este caso, especificamos dos nombres de función y enviamos dos prompts. Primero preguntamos sobre el estado de salud de un paciente. Y luego preguntamos cuándo se cambió ese estado. Verificamos que los resultados contengan la información esperada. Además, registramos todas las respuestas y esto es lo que se ve:

The patient with id P005 is currently healthy.
The health status of the patient with id P005 was changed on June 1, 2024.

4. Conclusión

La llamada de función es una excelente herramienta para ampliar la funcionalidad del LLM. También podemos usarla para integrar el LLM con nuestra lógica.

En este tutorial, exploramos cómo podemos implementar el flujo basado en LLM llamando a una o varias de nuestras funciones. Con este enfoque podemos implementar aplicaciones modernas que se integran con APIs de IA.