Primeros pasos con LiteRT-LM en Android

La API de Kotlin de LiteRT-LM para Android y JVM (Linux, macOS, Windows) con funciones como aceleración de GPU y NPU, multimodalidad y uso de herramientas.

Introducción

Aquí tienes una app de chat de muestra para la terminal compilada con la API de Kotlin:

import com.google.ai.edge.litertlm.*

suspend fun main() {
  Engine.setNativeMinLogSeverity(LogSeverity.ERROR) // Hide log for TUI app

  val engineConfig = EngineConfig(modelPath = "/path/to/model.litertlm")
  Engine(engineConfig).use { engine ->
    engine.initialize()

    engine.createConversation().use { conversation ->
      while (true) {
        print("\n>>> ")
        conversation.sendMessageAsync(readln()).collect { print(it) }
      }
    }
  }
}

Demostración del código de muestra de Kotlin

Para probar el ejemplo anterior, clona el repo y ejecútalo con example/Main.kt:

bazel run -c opt //kotlin/java/com/google/ai/edge/litertlm/example:main -- <abs_model_path>

Los modelos .litertlm disponibles se encuentran en la comunidad de HuggingFace LiteRT. La animación anterior usaba Gemma3-1B-IT.

Para ver un ejemplo de Android, consulta la app de la Galería de Google AI Edge.

Comienza a usar Gradle

Si bien LiteRT-LM se desarrolla con Bazel, proporcionamos los paquetes de Maven para los usuarios de Gradle/Maven.

1. Agrega la dependencia de Gradle

dependencies {
    // For Android
    implementation("com.google.ai.edge.litertlm:litertlm-android:latest.release")

    // For JVM (Linux, MacOS, Windows)
    implementation("com.google.ai.edge.litertlm:litertlm-jvm:latest.release")
}

Puedes encontrar las versiones disponibles en Google Maven en litertlm-android y litertlm-jvm.

Se puede usar latest.release para obtener la versión más reciente.

2. Inicializa el motor

Engine es el punto de entrada a la API. Inicialízalo con la ruta de acceso y la configuración del modelo. Recuerda cerrar el motor para liberar recursos.

Nota: El método engine.initialize() puede tardar una cantidad significativa de tiempo (p.ej., hasta 10 segundos) en cargar el modelo. Se recomienda llamar a este método en un subproceso o una corrutina en segundo plano para evitar bloquear el subproceso de IU.

import com.google.ai.edge.litertlm.Backend
import com.google.ai.edge.litertlm.Engine
import com.google.ai.edge.litertlm.EngineConfig

val engineConfig = EngineConfig(
    modelPath = "/path/to/your/model.litertlm", // Replace with your model path
    backend = Backend.GPU(), // Or Backend.NPU(nativeLibraryDir = "...")
    // Optional: Pick a writable dir. This can improve 2nd load time.
    // cacheDir = "/tmp/" or context.cacheDir.path (for Android)
)

val engine = Engine(engineConfig)
engine.initialize()
// ... Use the engine to create a conversation ...

// Close the engine when done
engine.close()

En Android, para usar el backend de la GPU, la app debe solicitar las bibliotecas nativas dependientes de forma explícita agregando lo siguiente a tu AndroidManifest.xml dentro de la etiqueta <application>:

  <application>
    <uses-native-library android:name="libvndksupport.so" android:required="false"/>
    <uses-native-library android:name="libOpenCL.so" android:required="false"/>
  </application>

Para usar el backend de la NPU, es posible que debas especificar el directorio que contiene las bibliotecas de la NPU. En Android, si las bibliotecas se incluyen en tu app, configúralo en context.applicationInfo.nativeLibraryDir. Consulta NPU de LiteRT-LM para obtener más detalles sobre las bibliotecas nativas de la NPU.

val engineConfig = EngineConfig(
    modelPath = modelPath,
    backend = Backend.NPU(nativeLibraryDir = context.applicationInfo.nativeLibraryDir)
)

3. Cómo crear una conversación

Una vez que se inicializa el motor, crea una instancia de Conversation. Puedes proporcionar un ConversationConfig para personalizar su comportamiento.

import com.google.ai.edge.litertlm.ConversationConfig
import com.google.ai.edge.litertlm.Message
import com.google.ai.edge.litertlm.SamplerConfig

// Optional: Configure the system instruction, initial messages, sampling
// parameters, etc.
val conversationConfig = ConversationConfig(
    systemInstruction = Contents.of("You are a helpful assistant."),
    initialMessages = listOf(
        Message.user("What is the capital city of the United States?"),
        Message.model("Washington, D.C."),
    ),
    samplerConfig = SamplerConfig(topK = 10, topP = 0.95, temperature = 0.8),
)

val conversation = engine.createConversation(conversationConfig)
// Or with default config:
// val conversation = engine.createConversation()

// ... Use the conversation ...

// Close the conversation when done
conversation.close()

Conversation implementa AutoCloseable, por lo que puedes usar el bloque use para la administración automática de recursos en conversaciones únicas o de corta duración:

engine.createConversation(conversationConfig).use { conversation ->
    // Interact with the conversation
}

4. Envío de mensajes

Existen tres maneras de enviar mensajes:

  • sendMessage(contents): Message: Es una llamada síncrona que se bloquea hasta que el modelo devuelve una respuesta completa. Esto es más simple para las interacciones básicas de solicitud y respuesta.
  • sendMessageAsync(contents, callback): Es una llamada asíncrona para transmitir respuestas. Esto es mejor para las solicitudes de larga duración o cuando deseas mostrar la respuesta a medida que se genera.
  • sendMessageAsync(contents): Flow<Message>: Es una llamada asíncrona que devuelve un flujo de Kotlin para transmitir respuestas. Este es el enfoque recomendado para los usuarios de corrutinas.

Ejemplo síncrono:

import com.google.ai.edge.litertlm.Content
import com.google.ai.edge.litertlm.Message

print(conversation.sendMessage("What is the capital of France?"))

Ejemplo asíncrono con devolución de llamada:

Usa sendMessageAsync para enviar un mensaje al modelo y recibir respuestas a través de una devolución de llamada.

import com.google.ai.edge.litertlm.Content
import com.google.ai.edge.litertlm.Message
import com.google.ai.edge.litertlm.MessageCallback
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit

val callback = object : MessageCallback {
    override fun onMessage(message: Message) {
        print(message)
    }

    override fun onDone() {
        // Streaming completed
    }

    override fun onError(throwable: Throwable) {
        // Error during streaming
    }
}

conversation.sendMessageAsync("What is the capital of France?", callback)

Ejemplo asíncrono con flujo:

Usa sendMessageAsync (sin el argumento de devolución de llamada) para enviar un mensaje al modelo y recibir respuestas a través de un flujo de Kotlin.

import com.google.ai.edge.litertlm.Content
import com.google.ai.edge.litertlm.Message
import kotlinx.coroutines.flow.catch
import kotlinx.coroutines.launch

// Within a coroutine scope
conversation.sendMessageAsync("What is the capital of France?")
    .catch { ... } // Error during streaming
    .collect { print(it.toString()) }

5. Multimodalidad

Los objetos Message pueden contener diferentes tipos de Content, incluidos Text, ImageBytes, ImageFile, AudioBytes y AudioFile.

// Initialize the `visionBackend` and/or the `audioBackend`
val engineConfig = EngineConfig(
    modelPath = "/path/to/your/model.litertlm", // Replace with your model path
    backend = Backend.CPU(), // Or Backend.GPU() or Backend.NPU(...)
    visionBackend = Backend.GPU(), // Or Backend.NPU(...)
    audioBackend = Backend.CPU(), // Or Backend.NPU(...)
)

// Sends a message with multi-modality.
// See the Content class for other variants.
conversation.sendMessage(Contents.of(
    Content.ImageFile("/path/to/image"),
    Content.AudioBytes(audioBytes), // ByteArray of the audio
    Content.Text("Describe this image and audio."),
))

6. Cómo definir y usar herramientas

Existen dos formas de definir herramientas:

  1. Con funciones de Kotlin (recomendado en la mayoría de los casos)
  2. Con la especificación de OpenAPI (control total de la especificación y la ejecución de la herramienta)

Cómo definir herramientas con funciones de Kotlin

Puedes definir funciones personalizadas de Kotlin como herramientas que el modelo puede llamar para realizar acciones o recuperar información.

Crea una clase que implemente ToolSet y anota los métodos con @Tool y los parámetros con @ToolParam.

import com.google.ai.edge.litertlm.Tool
import com.google.ai.edge.litertlm.ToolParam

class SampleToolSet: ToolSet {
    @Tool(description = "Get the current weather for a city")
    fun getCurrentWeather(
        @ToolParam(description = "The city name, e.g., San Francisco") city: String,
        @ToolParam(description = "Optional country code, e.g., US") country: String? = null,
        @ToolParam(description = "Temperature unit (celsius or fahrenheit). Default: celsius") unit: String = "celsius"
    ): Map<String, Any> {
        // In a real application, you would call a weather API here
        return mapOf("temperature" to 25, "unit" to  unit, "condition" to "Sunny")
    }

    @Tool(description = "Get the sum of a list of numbers.")
    fun sum(
        @ToolParam(description = "The numbers, could be floating point.") numbers: List<Double>,
    ): Double {
        return numbers.sum()
    }
}

En segundo plano, la API inspecciona estas anotaciones y la firma de la función para generar un esquema de estilo OpenAPI. Este esquema describe la funcionalidad, los parámetros (incluidos sus tipos y descripciones de @ToolParam) y el tipo de datos que devuelve la herramienta al modelo de lenguaje.

Tipos de parámetros

Los tipos de parámetros anotados con @ToolParam pueden ser String, Int, Boolean, Float, Double o un List de estos tipos (p.ej., List<String>). Usa tipos anulables (p.ej., String?) para indicar parámetros que admiten valores nulos. Establece un valor predeterminado para indicar que el parámetro es opcional y menciona el valor predeterminado en la descripción en @ToolParam.

Tipo de datos que se muestra

El tipo de datos que se muestra de la función de la herramienta puede ser cualquier tipo de Kotlin. El resultado se convertirá en un elemento JSON antes de enviarse de vuelta al modelo.

  • Los tipos List se convierten en arrays JSON.
  • Los tipos Map se convierten en objetos JSON.
  • Los tipos primitivos (String, Number, Boolean) se convierten en el primitivo JSON correspondiente.
  • Otros tipos se convierten en cadenas con el método toString().

Para los datos estructurados, se recomienda devolver Map o una clase de datos que se convertirá en un objeto JSON.

Cómo definir herramientas con la especificación de OpenAPI

Como alternativa, puedes definir una herramienta implementando la clase OpenApiTool y proporcionando la descripción de la herramienta como una cadena JSON que cumpla con la especificación de la API abierta. Este método es útil si ya tienes un esquema de OpenAPI para tu herramienta o si necesitas un control detallado sobre la definición de la herramienta.

import com.google.ai.edge.litertlm.OpenApiTool

class SampleOpenApiTool : OpenApiTool {

    override fun getToolDescriptionJsonString(): String {
        return """
        {
          "name": "addition",
          "description": "Add all numbers.",
          "parameters": {
            "type": "object",
            "properties": {
              "numbers": {
                "type": "array",
                "items": {
                  "type": "number"
                }
              },
              "description": "The list of numbers to sum."
            },
            "required": [
              "numbers"
            ]
          }
        }
        """.trimIndent() // Tip: trim to save tokens
    }

    override fun execute(paramsJsonString: String): String {
        // Parse paramsJsonString with your choice of parser/deserializer and
        // execute the tool.

        // Return the result as a JSON string
        return """{"result": 1.4142}"""
    }
}

Registro de herramientas

Incluye instancias de tus herramientas en ConversationConfig.

val conversation = engine.createConversation(
    ConversationConfig(
        tools = listOf(
            tool(SampleToolSet()),
            tool(SampleOpenApiTool()),
        ),
        // ... other configs
    )
)

// Send messages that might trigger the tool
conversation.sendMessageAsync("What's the weather like in London?", callback)

El modelo decidirá cuándo llamar a la herramienta según la conversación. Los resultados de la ejecución de la herramienta se envían automáticamente al modelo para generar la respuesta final.

Llamadas a herramientas manuales

De forma predeterminada, LiteRT-LM ejecuta automáticamente las llamadas a herramientas que genera el modelo, y los resultados de la ejecución de la herramienta se envían automáticamente al modelo para generar la siguiente respuesta.

Si deseas ejecutar herramientas de forma manual y enviar los resultados al modelo, puedes establecer automaticToolCalling en ConversationConfig como false.

val conversation = engine.createConversation(
    ConversationConfig(
        tools = listOf(
            tool(SampleOpenApiTool()),
        ),
        automaticToolCalling = false,
    )
)

Si inhabilitas las llamadas automáticas a herramientas, deberás ejecutar las herramientas de forma manual y enviar los resultados al modelo en el código de tu aplicación. No se llamará automáticamente al método execute de OpenApiTool cuando automaticToolCalling se establezca en false.

// Send a message that triggers a tool call.
val responseMessage = conversation.sendMessage("What's the weather like in London?")

// The model returns a Message with `toolCalls` populated.
if (responseMessage.toolCalls.isNotEmpty()) {
    val toolResponses = mutableListOf<Content.ToolResponse>()
    // There can be multiple tool calls in a single response.
    for (toolCall in responseMessage.toolCalls) {
        println("Model wants to call: ${toolCall.name} with arguments: ${toolCall.arguments}")

        // Execute the tool manually with your own logic. `executeTool` is just an example here.
        val toolResponseJson = executeTool(toolCall.name, toolCall.arguments)

        // Collect tool responses.
        toolResponses.add(Content.ToolResponse(toolCall.name, toolResponseJson))
    }

    // Use Message.tool to create the tool response message.
    val toolResponseMessage = Message.tool(Contents.of(toolResponses))

    // Send the tool response message to the model.
    val finalMessage = conversation.sendMessage(toolResponseMessage)
    println("Final answer: ${finalMessage.text}") // e.g., "The weather in London is 25c."
}

Ejemplo

Para probar el uso de herramientas, clona el repo y ejecútalo con example/ToolMain.kt:

bazel run -c opt //kotlin/java/com/google/ai/edge/litertlm/example:tool -- <abs_model_path>

Manejo de errores

Los métodos de la API pueden arrojar LiteRtLmJniException para los errores de la capa nativa o excepciones estándar de Kotlin, como IllegalStateException para los problemas del ciclo de vida. Siempre une las llamadas a la API en bloques try-catch. La devolución de llamada onError en MessageCallback también informará errores durante las operaciones asíncronas.