litert

Google’s litert_lm backend for .litertlm Gemma builds — litert message helpers, GPU/NPU backends, KV-cache usage, and bench().

Messages

litert’s Message / Content types — rishi wraps them with _mk_content, _mk_msg, _mk_msgs for the shared history format.

from fastcore.test import test_eq, test_fail
test_eq(_mk_msg("hello").to_json(), {"role": "user", "content": [{"type": "text", "text": "hello"}]})
test_eq(_mk_msg("hi", role="model").to_json()["role"], "model")
test_eq([x["role"] for x in _mk_msgs(["a", _mk_msg("b", role="model")])], ["user", "model"])
assert isinstance(_mk_content("x"), Text)
# bytes are sniffed: image vs audio
assert isinstance(_mk_content(b'\x89PNG\r\n\x1a\n' + b'\x00'*16), ImageBytes)
assert isinstance(_mk_content(b'RIFF\x00\x00\x00\x00WAVE' + b'\x00'*16), AudioBytes)

Portable history

fmt2hist / hist2fmt convert litert messages to the canonical rishi history dicts (and back), so a conversation can hop to llama.cpp or a hosted model via messages=.

Litert callbacks

Three litert-specific callbacks plus ChatToolHandler, which bridges litert’s native tool API to rishi’s approval gate. Token counts come from conv.token_count after each turn.

Tool calling

ChatToolHandler records tool rounds in litert’s wire format. Approval runs through the shared approve hook — see core.

source

ChatToolHandler

def ChatToolHandler(
    chat
):

Bridge litert’s in-engine tool loop to Chat callbacks, HITL approval, the tool-call budget, and history.

Loading models & LitertChat

_get_model downloads or cache-hits a .litertlm build. Pass backend=Backend.GPU() (or set RISHI_LITERT_GPU=1) for GPU/NPU; CPU is the default. LitertChat.create_engine owns the litert Conversation.

source

LitertChat

def LitertChat(
    model:NoneType=None, runtime:NoneType=None, model_path:NoneType=None, engine:NoneType=None,
    backend:NoneType=None, multimodal:bool=True, cache_dir:NoneType=None, enable_speculative_decoding:NoneType=None,
    eng_kw:NoneType=None, sp:str='', messages:NoneType=None, tools:NoneType=None, ctx_limit:NoneType=None,
    approve:NoneType=None, tool_max_len:NoneType=None, max_steps:int=10,
    final_prompt:str="You've reached the tool-call budget for this turn. Stop calling tools and answer with what you already have.",
    parallel_tools:bool=False, max_parallel_tools:NoneType=None, think:bool=False, filter_think:bool=True,
    temp:NoneType=None, top_k:NoneType=None, top_p:NoneType=None, seed:NoneType=None, sampler_config:NoneType=None,
    max_output_tokens:NoneType=None, constrain:NoneType=None, conv_kw:NoneType=None, cbs:NoneType=None,
    default_cbs:bool=True
):

Sync chat over a local litert_lm engine.

# fmt2hist / hist2fmt round-trip (model-free)
_nat = [
    'hi',
    Message.model(Contents.of('hello')).to_json(),
    {'role': 'model', 'tool_calls': [{'type': 'function', 'function': {'name': 'add', 'arguments': {'a': 2, 'b': 3}}}]},
    {'role': 'tool', 'content': [{'type': 'tool_response', 'name': 'add', 'response': 5}]},
]
_canon = LitertChat.fmt2hist(_nat)
test_eq(_canon[0], {'role': 'user', 'content': 'hi'})
test_eq(_canon[1], {'role': 'assistant', 'content': 'hello'})
test_eq(_canon[2]['role'], 'assistant')
test_eq(_canon[2]['tool_calls'][0]['function'], {'name': 'add', 'arguments': {'a': 2, 'b': 3}})
assert _canon[2]['tool_calls'][0]['id']                       # a synthetic id was assigned
test_eq(_canon[3], {'role': 'tool', 'name': 'add', 'content': '5'})

# canonical -> litert Messages ready for the engine
_ms = LitertChat.hist2fmt(_canon)
test_eq([m.role.value for m in _ms], ['user', 'model', 'model', 'tool'])
test_eq(str(_ms[1]), 'hello')
test_eq(_ms[2].tool_calls[0].name, 'add')
test_eq(_ms[3].contents.to_json(), [{'type': 'tool_response', 'name': 'add', 'response': '5'}])

# fmt2hist is idempotent on an already-canonical (llama-origin) tool round
_llama = [{'role': 'assistant', 'content': '', 'tool_calls': [{'id': 'call_x', 'type': 'function',
                                                               'function': {'name': 'add', 'arguments': {'a': 1}}}]},
          {'role': 'tool', 'tool_call_id': 'call_x', 'name': 'add', 'content': '1'}]
test_eq(LitertChat.fmt2hist(_llama), _llama)
set_min_log_severity(5)
_get_model(gemma4_12b)
'/Users/71293/.cache/huggingface/hub/models--litert-community--gemma-4-12B-it-litert-lm/snapshots/44cf85a326f79b814fa86a60af414c042755b43a/gemma-4-12B-it.litertlm'

Using Chat (manual)

GPU-backed demos below. First run downloads weights; set cache_dir to reuse them.

chat=Chat(backend=Backend.GPU(), cache_dir='.cache/litertlm', think=True)
# chat_12 = LitertChat(engine=LitertChat.create_engine(gemma4_12b, multimodal=False, cache_dir='.cache/litertlm', be=Backend.GPU()))
# set_min_log_severity(2)
r = chat("Reply with exactly: pong")
assert 'pong' in resp_text(r).lower()
assert resp_text(chat.hist[-1]) == resp_text(r) and chat.hist[0]['role'] == 'user'
assert chat.use.total_tokens > 0 and chat.token_count > 0
print(chat.use); chat.print_hist()
total=104|in=103|out=1|turns=1

user

Reply with exactly: pong

assistant

🧠 Thinking

Thinking Process:

  1. Analyze the Request: The user has instructed me to reply with exactly the word “pong”.
  2. Determine the Constraint: The constraint is strict: “Reply with exactly: pong”.
  3. Formulate the Response: The response must be the string “pong”.
  4. Final Output Generation: pong

pong

# Real LiteRT conversation-cache test. BenchmarkInfo reports the actual last prefill performed by
# the runtime; turn two should prefill only its short new message while retaining the old KV context.
kvchat = Chat(backend=Backend.GPU(), cache_dir='.cache/litertlm', think=False, max_output_tokens=16,
              eng_kw={'enable_benchmark': True})
try:
    kvchat(('cache verification context ' * 96) + '\nReply with exactly: stored')
    first = kvchat.conv.get_benchmark_info()
    context_before_turn2 = kvchat.token_count
    r = kvchat('What exact word did I ask you to reply with? Reply with only that word.')
    second = kvchat.conv.get_benchmark_info()
    print(dict(first_prefill=first.last_prefill_token_count,
               second_prefill=second.last_prefill_token_count,
               context_before_turn2=context_before_turn2,
               context_after_turn2=kvchat.token_count,
               reported_cached=kvchat.use.cached_tokens,
               answer=resp_text(r)))
    assert context_before_turn2 > 0 and kvchat.token_count > context_before_turn2
    assert second.last_prefill_token_count < context_before_turn2, 'turn two appears to have re-prefilled the old context'
    assert kvchat.use.cached_tokens == context_before_turn2, kvchat.use
    assert 'stored' in resp_text(r).lower(), resp_text(r)
finally:
    kvchat.close()
{'first_prefill': 304, 'second_prefill': 25, 'context_before_turn2': 306, 'context_after_turn2': 333, 'reported_cached': 306, 'answer': 'stored'}

Images and audio

Gemma-4 litert builds accept images and audio beside text in one call — same message list shape as other backends.

from PIL import Image
from fastcore.all import img_bytes, Path
im = Image.open(repo_root()/'nbs'/'images.jpeg')
chat(['explain this image', img_bytes(im)])
chat(['transcribe this audio', AudioFile(str(repo_root()/'nbs'/'speech.wav'))])

🧠 Thinking

Thinking Process:

  1. Analyze the Request: The user wants me to transcribe the provided audio.

  2. Analyze the Input (Audio/Text): The input is a piece of poetry or song lyrics.

  3. Perform Transcription (Segment by Segment): I will listen/read the provided text and transcribe it accurately.

    • Original text: “Dancing in the masquerade, idle truth in plain sight jaded, pop, roll, click, dot, who will I be today or not. But such a tide as moving seems asleep, too full for sound and foam, when that drew from out the boundless deep turns again home, twilight and evening bell, and after that”

  4. Review and Finalize Transcription: Ensure the transcription matches the provided text exactly. (The provided text is already quite clean, so the task is primarily verification and output.)

  5. Output Generation.

Dancing in the masquerade, idle truth in plain sight jaded, pop, roll, click, dot, who will I be today or not. But such a tide as moving seems asleep, too full for sound and foam, when that drew from out the boundless deep turns again home, twilight and evening bell, and after that

Utilities & bench

classify, structured, and check are inherited from core. bench() is litert-only: init time, time-to-first-token, prefill/decode tok/s. LiteRT does not expose log-likelihood scoring.

source

bench

def bench(
    model_id:str='litert-community/gemma-4-E2B-it-litert-lm', model_path:NoneType=None,
    backend:CPU=CPU(thread_count=None), prefill_tokens:int=64, decode_tokens:int=64, **kw
):

Benchmark init time, TTFT, and prefill/decode tokens-per-sec via litert’s Benchmark.

# classify + structured on a real model
sm = Chat(cache_dir='.cache/litertlm')
test_eq(sm.classify("I absolutely loved this film!", ['positive', 'negative']), 'positive')

from dataclasses import dataclass
@dataclass
class Person: name:str; age:int
p = sm.structured("Extract the person: John Smith is 30 years old.", Person)
print(p); assert isinstance(p, Person) and 'John' in p.name
sm.close()
Person(name='John Smith', age=30)
from fastcore.test import test_eq
import rishi.core, rishi.litert
test_eq(rishi.core.get_runtime('litert'), rishi.litert.LitertChat)

# dispatch really happens in `Chat.__new__`, so it can be checked without loading a model
test_eq(type(rishi.core.Chat.__new__(rishi.core.Chat)), rishi.litert.LitertChat)              # the default runtime
test_eq(type(rishi.core.Chat.__new__(rishi.core.Chat, 'litert-community/gemma-4-E2B-it-litert-lm')),
        rishi.litert.LitertChat)
test_eq(rishi.litert.LitertChat._runtime, 'litert')