Começar a usar o LiteRT-LM no Android

A API Kotlin do LiteRT-LM para Android e JVM (Linux, MacOS, Windows) com recursos como aceleração de GPU e NPU, multimodalidade e uso de ferramentas.

Introdução

Confira um exemplo de app de chat de terminal criado com a API 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) }
      }
    }
  }
}

Demonstração do exemplo de código Kotlin

Para testar o exemplo acima, clone o repositório e execute com example/Main.kt:

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

Os modelos .litertlm disponíveis estão na comunidade do HuggingFace LiteRT. A animação acima usou o Gemma3-1B-IT.

Para ver um exemplo do Android, confira o app Galeria do Google AI Edge.

Como começar a usar o Gradle

Embora o LiteRT-LM seja desenvolvido com o Bazel, fornecemos os pacotes Maven para usuários do Gradle/Maven.

1. Adicionar a dependência do 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")
}

Você pode encontrar as versões disponíveis no Google Maven em litertlm-android e litertlm-jvm.

latest.release pode ser usado para acessar a versão mais recente.

2. Inicializar o mecanismo

O Engine é o ponto de entrada da API. Inicialize com o caminho do modelo e a configuração. Lembre-se de fechar o mecanismo para liberar recursos.

Observação:o método engine.initialize() pode levar um tempo considerável (por exemplo, até 10 segundos) para carregar o modelo. É altamente recomendável chamar isso em uma linha de execução ou corrotina em segundo plano para evitar o bloqueio da linha de execução da UI.

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()

No Android, para usar o back-end da GPU, o app precisa solicitar as bibliotecas nativas dependentes explicitamente adicionando o seguinte ao seu AndroidManifest.xml na tag <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 o back-end da NPU, talvez seja necessário especificar o diretório que contém as bibliotecas da NPU. No Android, se as bibliotecas forem agrupadas com seu app, defina como context.applicationInfo.nativeLibraryDir. Consulte LiteRT-LM NPU para mais detalhes sobre as bibliotecas nativas da NPU.

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

3. Criar uma conversa

Depois que o mecanismo for inicializado, crie uma instância Conversation. Você pode fornecer um ConversationConfig para personalizar o comportamento dele.

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. Assim, você pode usar o bloco use para gerenciamento automático de recursos em conversas únicas ou de curta duração:

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

4. Como enviar mensagens

Há três maneiras de enviar mensagens:

  • sendMessage(contents): Message: chamada síncrona que bloqueia até que o modelo retorne uma resposta completa. Isso é mais simples para interações básicas de solicitação/resposta.
  • sendMessageAsync(contents, callback): chamada assíncrona para respostas de streaming. Isso é melhor para solicitações de longa duração ou quando você quer mostrar a resposta à medida que ela é gerada.
  • sendMessageAsync(contents): Flow<Message>: chamada assíncrona que retorna um fluxo do Kotlin para respostas de streaming. Essa é a abordagem recomendada para usuários de corrotinas.

Exemplo 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?"))

Exemplo assíncrono com callback:

Use sendMessageAsync para enviar uma mensagem ao modelo e receber respostas por um callback.

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)

Exemplo assíncrono com fluxo:

Use sendMessageAsync (sem o argumento de callback) para enviar uma mensagem ao modelo e receber respostas por um fluxo 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. Multimodalidade

Os objetos Message podem conter diferentes tipos de Content, incluindo Text, ImageBytes, ImageFile, AudioBytes e 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. Como definir e usar ferramentas

Há duas maneiras de definir ferramentas:

  1. Com funções do Kotlin (recomendado para a maioria dos casos)
  2. Com a especificação OpenAPI (controle total da especificação e execução da ferramenta)

Definir ferramentas com funções do Kotlin

Você pode definir funções personalizadas do Kotlin como ferramentas que o modelo pode chamar para realizar ações ou buscar informações.

Crie uma classe que implemente ToolSet e adicione anotações aos métodos com @Tool e parâmetros com @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()
    }
}

Nos bastidores, a API inspeciona essas anotações e a assinatura da função para gerar um esquema no estilo OpenAPI. Esse esquema descreve a funcionalidade, os parâmetros (incluindo tipos e descrições de @ToolParam) e o tipo de retorno da ferramenta para o modelo de linguagem.

Tipos de parâmetros

Os tipos de parâmetros anotados com @ToolParam podem ser String, Int, Boolean, Float, Double ou um List desses tipos (por exemplo, List<String>). Use tipos anuláveis (por exemplo, String?) para indicar parâmetros anuláveis. Defina um valor padrão para indicar que o parâmetro é opcional e mencione o valor padrão na descrição em @ToolParam.

Tipo de retorno

O tipo de retorno da função de ferramenta pode ser qualquer tipo do Kotlin. O resultado será convertido em um elemento JSON antes de ser enviado de volta ao modelo.

  • Os tipos List são convertidos em matrizes JSON.
  • Os tipos Map são convertidos em objetos JSON.
  • Os tipos primitivos (String, Number, Boolean) são convertidos para o primitivo JSON correspondente.
  • Outros tipos são convertidos em strings com o método toString().

Para dados estruturados, é recomendável retornar Map ou uma classe de dados que será convertida em um objeto JSON.

Definir ferramentas com a especificação OpenAPI

Como alternativa, você pode definir uma ferramenta implementando a classe OpenApiTool e fornecendo a descrição da ferramenta como uma string JSON em conformidade com a especificação da API aberta. Esse método é útil se você já tem um esquema OpenAPI para sua ferramenta ou se precisa de controle refinado sobre a definição dela.

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}"""
    }
}

Ferramentas de registro

Inclua instâncias das suas ferramentas no 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)

O modelo decide quando chamar a ferramenta com base na conversa. Os resultados da execução da ferramenta são enviados automaticamente de volta ao modelo para gerar a resposta final.

Chamadas de ferramentas manuais

Por padrão, as chamadas de função geradas pelo modelo são executadas automaticamente pelo LiteRT-LM, e os resultados da execução da função são enviados automaticamente de volta ao modelo para gerar a próxima resposta.

Se você quiser executar ferramentas manualmente e enviar os resultados de volta ao modelo, defina automaticToolCalling em ConversationConfig como false.

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

Se você desativar a chamada automática de função, será necessário executar manualmente as ferramentas e enviar os resultados de volta ao modelo no código do aplicativo. O método execute de OpenApiTool não será chamado automaticamente quando automaticToolCalling for definido como 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."
}

Exemplo

Para testar o uso de ferramentas, clone o repositório e execute com example/ToolMain.kt:

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

Tratamento de erros

Os métodos de API podem gerar LiteRtLmJniException para erros da camada nativa ou exceções padrão do Kotlin, como IllegalStateException para problemas de ciclo de vida. Sempre agrupe as chamadas de API em blocos try-catch. O callback onError em MessageCallback também vai informar erros durante operações assíncronas.