Every team that ships an LLM feature eventually hits the same wall: the model can talk about your order database, but it cannot read it. The fix is tool calling — you describe a set of functions to the model, it emits a structured request, your code executes it, and the result goes back into the context. That loop is what people mean by "agentic AI," and in 2025 it is the single most requested addition to the generative-AI systems we are asked to review.
Most tool-calling tutorials assume a hosted frontier API that returns a guaranteed-valid tool_calls object. This one does not. We are running an open-weight Gemma model with KerasHub on the TensorFlow stack, self-hosted, which means we own the part that guarantees valid JSON. That is the interesting engineering, and it is also where self-hosted agents usually fail in production.
This tutorial builds a small, auditable tool-calling agent: a tool registry, a constrained decoder that cannot emit malformed JSON, an execution loop with hard limits, and an eval harness that tells you whether any of it works.
What you need
pip install -U keras keras-hub tensorflow
Set your backend before importing Keras, and use the Gemma preset you already have licensed (see our Gemma fine-tuning tutorial for preset access and Kaggle credentials):
import os
os.environ["KERAS_BACKEND"] = "jax" # or "tensorflow"
import json, re, keras, keras_hub
gemma = keras_hub.models.GemmaCausalLM.from_preset("gemma2_2b_en")
tokenizer = gemma.preprocessor.tokenizer
An instruction-tuned preset (*_instruct_*) follows the tool protocol far better than a base preset. If your agent is limited to five or ten tools in a narrow domain, a 2B-class instruct model is usually enough — and it is the difference between a GPU you can afford to keep warm and one you cannot.
Step 1: define tools as data, not prose
A tool is a callable plus a JSON Schema. Keep both in one place so the prompt and the executor can never drift apart:
TOOLS = {}
def tool(name, description, parameters):
def wrap(fn):
TOOLS[name] = {
"name": name,
"description": description,
"parameters": parameters, # JSON Schema
"fn": fn,
}
return fn
return wrap
@tool(
"lookup_order",
"Look up the status and line items of a customer order by its ID.",
{
"type": "object",
"properties": {"order_id": {"type": "string", "pattern": "^ORD-[0-9]{6}$"}},
"required": ["order_id"],
},
)
def lookup_order(order_id: str):
row = db.orders.get(order_id)
if row is None:
return {"error": "not_found", "order_id": order_id}
return {"status": row.status, "eta": row.eta, "items": row.items[:10]}
Three rules we apply on every engagement:
- Tools return small, bounded payloads. Truncate lists, drop blobs. A tool that can return 40 KB of JSON will eventually blow the context window at the worst possible moment.
- Tools are the authorization boundary. The model never receives a customer ID it did not already have; the tool takes the authenticated session from your application scope, not from the model's arguments. Prompt injection in retrieved text is a real attack path, and the only durable defence is that the dangerous action simply isn't reachable.
- Read tools and write tools are separate registries. Writes get a confirmation step. Always.
Step 2: prompt the model with the schemas
Render the registry into the system turn using Gemma's chat template:
def render_tools():
return json.dumps(
[{k: t[k] for k in ("name", "description", "parameters")} for t in TOOLS.values()],
indent=2,
)
SYSTEM = """You are a support agent. You may call tools.
Available tools (JSON Schema):
{tools}
Reply with exactly one JSON object and nothing else:
{{"tool": "<tool name>", "arguments": {{...}}}} to call a tool
{{"answer": "<text>"}} to answer the user
"""
def build_prompt(history):
parts = [f"<start_of_turn>user\n{SYSTEM.format(tools=render_tools())}\n"]
for turn in history:
parts.append(f"{turn['role']}: {turn['content']}\n")
parts.append("<end_of_turn>\n<start_of_turn>model\n")
return "".join(parts)
The schema block is the expensive part of the prompt and it is identical on every call, so it is worth caching — see the KV-cache note at the end.
Step 3: make invalid JSON impossible
Asking politely for JSON gets you valid JSON perhaps 90–98% of the time with a small model. In an agent loop that runs three tool calls per request, a 3% per-call failure rate is roughly a 9% request failure rate. Not acceptable.
The robust answer is constrained decoding: at each step, mask the logits so that only tokens which keep the output a valid prefix of the grammar can be sampled. KerasHub exposes this through a sampler's logit-modification hook:
class GrammarSampler(keras_hub.samplers.GreedySampler):
"""Mask logits to tokens permitted by a JSON-schema-driven state machine."""
def __init__(self, validator, tokenizer, **kwargs):
super().__init__(**kwargs)
self.validator = validator # incremental prefix checker
self.tokenizer = tokenizer
def get_next_token(self, logits, prompt, cache, index):
allowed = self.validator.allowed_token_mask(prompt, index) # bool [vocab]
logits = keras.ops.where(allowed, logits, -1e9)
return super().get_next_token(logits, prompt, cache, index)
allowed_token_mask is the real work: it tracks where you are in the object (key, value, string, number) and returns the vocabulary mask for the next token. Do not hand-roll it. Use a maintained JSON/grammar engine — outlines, llguidance, or xgrammar all compile a JSON Schema into exactly this incremental mask, and they cache the compiled automaton so the per-token cost is a lookup.
If you cannot take that dependency, the cheap 80% version is a two-stage decode:
def constrained_call(prompt, retries=2):
for attempt in range(retries + 1):
raw = gemma.generate(prompt, max_length=512)
text = raw[len(prompt):].strip()
m = re.search(r"\{.*\}", text, re.S) # first balanced-ish object
if m:
try:
obj = json.loads(m.group(0))
jsonschema.validate(obj, AGENT_OUTPUT_SCHEMA)
return obj
except Exception as e:
prompt += f"\n{text}\nInvalid: {e}. Reply with only the JSON object.\n"
raise ValueError("model failed to produce valid tool JSON")
Log every retry as a metric. A rising retry rate is the earliest signal that a prompt change or a preset upgrade has regressed your agent, long before users complain.
Step 4: the loop, with the guardrails that matter
MAX_STEPS = 5
MAX_TOOL_SECONDS = 8
def run_agent(user_message, session):
history = [{"role": "user", "content": user_message}]
trace = []
for step in range(MAX_STEPS):
obj = constrained_call(build_prompt(history))
if "answer" in obj:
return obj["answer"], trace
spec = TOOLS.get(obj.get("tool"))
if spec is None:
history.append({"role": "tool", "content": '{"error":"unknown_tool"}'})
continue
args = obj.get("arguments", {})
try:
jsonschema.validate(args, spec["parameters"]) # validate again, always
result = call_with_timeout(spec["fn"], args, session, MAX_TOOL_SECONDS)
except Exception as e:
result = {"error": type(e).__name__, "detail": str(e)[:200]}
trace.append({"step": step, "tool": spec["name"], "args": args, "result": result})
history.append({"role": "tool", "content": json.dumps(result)[:2000]})
return "I couldn't complete that request. Escalating to a human.", trace
Note what is deliberately in there: a step ceiling (an unbounded agent is an unbounded bill), a tool timeout, re-validation of arguments at the executor even though the decoder was constrained, truncated tool output, and a trace. The trace is not a debugging luxury — it is what you show an auditor, what you replay in evals, and what you attach to the support ticket when the answer was wrong.
Step 5: evaluate the agent, not just the model
Perplexity tells you nothing about whether an agent works. Build a fixed suite of task cases and score four things:
| Metric | What it catches |
|---|---|
| Tool-selection accuracy | Model picks the right tool for the request |
| Argument exact-match / schema validity | Constrained decoding and prompt clarity |
| Task success (final answer correct) | The loop end to end |
| Abstention rate on unanswerable cases | Whether it invents tool calls |
def evaluate(cases):
scores = {"tool": 0, "args": 0, "success": 0, "abstain": 0}
for c in cases:
answer, trace = run_agent(c["message"], session=c["session"])
called = [t["tool"] for t in trace]
scores["tool"] += called[:1] == c["expected_tools"][:1]
scores["args"] += any(t["args"] == c.get("expected_args") for t in trace)
scores["success"] += c["check"](answer)
if c.get("unanswerable"):
scores["abstain"] += not called
return {k: v / len(cases) for k, v in scores.items()}
Fifty hand-written cases covering your real request mix are worth more than a public agent benchmark. Freeze them, run them in CI on every prompt or preset change, and pair them with the safety checks from our generative AI evaluation and guardrailing guide.
When accuracy stalls, the fix is usually one of three things, in this order: fewer and better-described tools; a short set of worked examples in the prompt; then, only if those fail, a LoRA fine-tune on a few thousand traced tool calls harvested from the eval suite and from logged production traffic.
Serving it
The agent loop is ordinary Python around a generation endpoint, so deploy it as a stateless service in front of the model rather than inside it:
- Export the model once and serve it behind TensorFlow Serving or a dedicated generation server; the orchestration service holds the registry, the loop, and the timeouts. Our production serving tutorial covers the model-side deployment.
- Enable the KV cache in
generate()(KerasHub does this by default) and keep the system/tool block byte-identical across calls so the cached prefix stays valid. Reordering your tool schemas per request silently destroys prefix reuse. - Quantize for inference. An int8 or 4-bit Gemma typically holds tool-selection accuracy while roughly halving latency and memory — measure it on your eval suite, not on a leaderboard. See our model-optimization guide.
- Budget by step, not by request. Alert on p95 steps per request and on tool error rate; both drift upward as the tool surface grows.
The honest summary
A self-hosted tool-calling agent on Gemma and KerasHub is very achievable — for a bounded domain, a handful of well-described tools, and a team willing to own the JSON-validity problem instead of renting it. What makes it production-grade is not the model. It is the schema-constrained decoder, the step and timeout ceilings, the authorization boundary drawn at the tool, and an eval suite that runs in CI.
If you are planning an agentic feature on the TensorFlow stack and want a second opinion on tool design, decoding strategy, or whether you need a fine-tune at all, get in touch — our consultants do this work every week.