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 HuggingFace LiteRT (link em inglês). A animação acima usou o Gemma3-1B-IT.

Para ver um exemplo do Android, confira o app Galeria do Google AI Edge (disponível no Google Play).

Como começar a usar o Gradle

Embora o LiteRT-LM seja desenvolvido com o Bazel, oferecemos os pacotes do Maven para usuários do Gradle ou do 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.

O 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-o com o caminho e a configuração do modelo. Lembre-se de fechar o mecanismo para liberar recursos.

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

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 AndroidManifest.xml dentro da 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 estiverem agrupadas com o app, defina-o como context.applicationInfo.nativeLibraryDir. Consulte NPU do LiteRT-LM 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 de 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, então você pode usar o bloco use para gerenciamento automático de recursos para conversas únicas ou de curta duração:

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

4. 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 e 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 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`, `audioBackend`, or both
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. Definir e usar ferramentas

Há duas maneiras de definir ferramentas:

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

Definir ferramentas com funções Kotlin

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

Crie uma classe que implemente ToolSet e anote 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 os 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 uma 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 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 no 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 OpenAPI. 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 or deserializer and
        // execute the tool.

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

Registrar ferramentas

Inclua instâncias das 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 vai decidir 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.

Chamada manual de ferramentas

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

Se você quiser executar ferramentas manualmente e enviar 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 ferramentas, será necessário executar ferramentas manualmente e enviar resultados de volta ao modelo no código do aplicativo. O método execute de OpenApiTool não será chamado automaticamente quando automaticToolCalling estiver 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 da API podem gerar LiteRtLmJniException para erros da camada nativa ou exceções Kotlin padrão, como IllegalStateException para problemas de ciclo de vida. Sempre envolva chamadas de API em blocos try-catch. O callback onError em MessageCallback também vai informar erros durante operações assíncronas.