rishi

rishi is a thin chat layer over Google’s on-device litert_lm engine. You give it a Gemma model id, it downloads the weights once, and then you talk to the model with a plain function call. It keeps the conversation history where you can read it, streams tokens into a notebook, shows the model’s thinking, tracks how full the context is getting, runs tools behind an approval gate, executes python from replies, and turns answers into structured objects or graded results.

Everything runs on your machine. No API keys, and no network once the model is cached.

Install

rishi is built with nbdev, so the notebooks in nbs/ are the source. Install from source:

pip install rishi

To work on it, clone the repo, install it editable, and use nbdev to export and test:

pip install -e '.[dev]'
nbdev-prepare

Quickstart

Build a Chat and call it. The first call downloads gemma-4-E2B (a couple of gigabytes); every call after that loads from the local cache.

from rishi.core import Chat, resp_text

chat = Chat()
r = chat("Give me one fact about lobsters.")
print(resp_text(r))     # in a notebook, `r` also renders as markdown on its own

Calling chat runs one turn and returns litert’s response dict wrapped in Resp. resp_text(r) pulls the text out; in a notebook r renders itself as markdown, including any thinking and tool calls. The turn is appended to chat.hist, which you can show with chat.print_hist(). Call chat again and it continues the same conversation.

Streaming

Pass stream=True and iterate. You get markdown chunks as the model decodes them. display_stream renders those chunks live in a notebook.

# print chunks as they decode:
for chunk in chat("Write a haiku about the sea.", stream=True):
    print(chunk, end='', flush=True)

# or render it live in a notebook cell:
display_stream(chat("Say hello in three languages.", stream=True))

Thinking

Set think=True to turn on the model’s thinking channel. resp_text returns just the answer, thought(r) returns the reasoning, and in a notebook r shows the thinking as a quoted block above the reply. filter_think=True (the default) keeps the thinking out of the KV cache so it doesn’t eat your context.

chat = Chat(think=True)
r = chat("A bat and ball cost $1.10, and the bat is $1 more than the ball. How much is the ball?")
print(resp_text(r))     # the answer; thought(r) has the reasoning

Tools and approval

Pass plain Python functions as tools. litert reads their signatures and docstrings to build the schema and calls them during a turn. rishi records every call in the history, and if you give it an approve function it checks before running each one.

hitl_policy builds an approve function from a per-tool rule: approved runs the tool, dont_run always blocks it, and check asks you on the console.

from rishi.core import hitl_policy

def add(a: int, b: int) -> int:
    "Add two integers."
    return a + b

def delete_files(path: str) -> str:
    "Delete everything under a path."
    return f"wiped {path}"

approve = hitl_policy({'add': 'approved', 'delete_files': 'dont_run'})
chat = Chat(tools=[add, delete_files], approve=approve)
chat("Add 2 and 3, then delete /tmp/data.")

The blocked call never runs. It is recorded as “Denied by human operator” and handed back to the model, which finishes the turn without it. For anything fancier than a fixed policy, pass your own approve(tool_call) -> bool to log the request, rate-limit it, or pop up a prompt.

Running python from replies

Add PyFenceCallback and the chat becomes a code interpreter. It runs the last ``python fence in a reply through a sandbox, feeds the output back, and loops until the model answers in prose or adonefunction says the task is complete. [output_matches](https://vedicreader.github.io/rishi/core.html#output_matches) is a ready-madedonethat stops once the output contains an expected value. Code runs through the sameapprove` gate as a tool.

from rishi.core import PyFenceCallback, output_matches

chat = Chat(cbs=[PyFenceCallback], sp="Use a ```python fence to compute the answer, then reply in prose.")
chat("What is 2**100?")

# stop as soon as the code output matches an expected value:
chat("Sort [3, 1, 2] ascending and print it.",
     cbs=[PyFenceCallback(done=output_matches('[1, 2, 3]'))])

# or run a snippet yourself in the chat's persistent sandbox:
chat.run_py("sum(range(10))")

Structured output and classification

chat.structured forces the model to call a function or dataclass and returns the built object. chat.classify picks one label from a list. Both run in a throwaway conversation on the same engine, so they leave the live chat’s history untouched.

from dataclasses import dataclass

@dataclass
class Person: name: str; age: int

chat.structured("Extract the person: John Smith is 30 years old.", Person)
# -> Person(name='John Smith', age=30)

chat.classify("I loved this film!", ['positive', 'negative'])
# -> 'positive'

Grading answers

chat.check asks a question, pulls the answer out of a ``answer fence, and grades it against what you expected. The default grade is a deterministic match. Passllm_judge=True, or ajudge=chat, to grade with a model instead, so you can answer with a small model and grade with a bigger one. Pass your owngrade_fn(answer, expected) -> bool` for custom logic.

# deterministic: the answer must contain the expected value
chat.check("What is the capital of France?", "Paris").ok      # -> True

# grade with a bigger model as the judge:
judge = Chat(model_id=gemma4_12b, multimodal=False, cache_dir='.cache/litertlm')
chat.check("Name a primary colour.", "red, blue, or yellow", judge=judge).ok
judge.close()

Knowing when to compress

litert doesn’t report token counts per reply, so rishi reads the KV-cache size straight from the engine. After each turn chat.use holds that turn’s input and output tokens, chat.token_count is the live context size, and chat.pct_full is that size over ctx_limit.

chat("...")
print(chat.use)            # total=... | in=... | out=... | turns=1
if chat.pct_full > 0.8:
    pass                   # summarise the history and start a fresh Chat

Custom callbacks

Everything above is built from callbacks. Subclass ChatCallback, hook an event (before_send, after_response, before_tool_calls, after_tool_calls), and read live turn state off the chat (self.turn_res is chat.turn_res). order sets when it runs. Register with chat.add_cb for every turn, pass cbs= to a single call to run it for that turn only, and drop one with chat.remove_cb (by instance or by class).

class Logger(ChatCallback):
    order = 40
    def after_response(self): print('reply tokens:', self.chat.use.completion_tokens)

chat.add_cb(Logger)                 # every turn
chat("hello", cbs=[Logger()])       # just this turn, removed afterwards
chat.remove_cb(Logger)              # by class or instance

Installing the skill

rishi bundles skill.md, an agent skill describing the API. A harness can install it into the standard skill directories:

from rishi.core import mv_skill_md
mv_skill_md(dry_run=False)   # writes SKILL.md under .claude/skills/rishi/ and .agents/skills/rishi/

Call it with no arguments for a dry run that just prints where it would write.

Sharing a model, backends, and benchmarks

Loading a model costs a few seconds and a couple of gigabytes of RAM. To run several conversations off one load, build the engine once and hand it to each chat:

eng = Chat.create_engine(cache_dir='.cache/litertlm')
a, b = Chat(engine=eng), Chat(engine=eng)

A Chat you build owns its engine and frees it on close(). A Chat you hand an engine to leaves it alone, so the other chats keep working. The default backend is CPU; for GPU pass backend=Backend.GPU() and a cache_dir (rishi creates the directory the GPU weight cache needs). To measure raw speed, bench() reports init time, time to first token, and prefill and decode tokens per second. Browse models at huggingface.co/litert-community.