Function Calling and Structured Output for On-Device Agents with Gemma 4

On-device agents live or die by one thing: how reliably the model’s output can be turned into action. If a cloud API hallucinates a malformed JSON tool call, you retry and move on. If a small model on a phone does it, your agent loop turns into a parser hell, and users uninstall the app.

Gemma 4 gives on-device developers two tools that make the agent loop trustworthy: native function calling (the model natively requests tool execution) and constrained decoding (the runtime only emits tokens that keep the output valid). Used together, they make structured output reliable enough for production.

In this tutorial, you will learn how to design a robust function-calling loop for on-device agents, apply constrained decoding, and build fallbacks for when models still get it wrong.

Prerequisites

• A Gemma 4 E2B or E4B model (function-calling capable) in .litertlm format
• The flutter_gemma package, or the LiteRT-LM CLI for experiments
• A device or browser with GPU support

How Native Function Calling Works

Unlike generic JSON-in-prompt tricks, Gemma 4 has function calling built into its architecture. The runtime pauses generation, returns a structured tool-call request to your application layer, and resumes once the tool’s output is available:

User: "Create a reminder to call mom at 5pm."

Gemma 4:  {tool: "schedule_reminder", args: {title: "Call mom", at: "17:00"}}
App:      [schedules the reminder, returns "scheduled"]
Gemma 4:  "Done — I'll remind you to call mom at 5 PM."

The model never touches the OS. It emits a request; your code is the only thing with the power to act. That separation is both a reliability pattern and a security boundary.

Defining the Tool Schema

Declare tools with a schema the model can reason about. On Flutter with flutter_gemma:

const tools = [
  ToolSpec(
    name: 'schedule_reminder',
    description: 'Create a reminder for the user.',
    parameters: JsonSchema(
      type: 'object',
      properties: {
        'title': JsonSchema(type: 'string', description: 'Reminder text'),
        'at': JsonSchema(type: 'string', description: 'Time in 24h HH:mm'),
      },
      required: ['title', 'at'],
    ),
  ),
  ToolSpec(
    name: 'read_notes',
    description: 'Read the user\'s notes matching a keyword.',
    parameters: JsonSchema(
      type: 'object',
      properties: {
        'keyword': JsonSchema(type: 'string'),
      },
      required: ['keyword'],
    ),
  ),
];

Small models perform better with fewer, simpler tools. If you have dozens, use two-stage discovery (like Agent Skills): give the model names and descriptions first, and load full schemas only for the tools it picks.

The Tool-Calling Loop

Here’s the core loop, hardened for on-device reality:

class ToolLoop {
  final Map<String, ToolHandler> _tools = {};
  static const maxSteps = 6;

  Future<String> run(Conversation conversation, String task) async {
    var input = task;

    for (var step = 0; step < maxSteps; step++) {
      final response = await conversation.sendMessage(
        input,
        tools: _toolSpecs,
        forceToolCall: false,
      );

      final call = response.toolCall;
      if (call == null) return response.text; // final answer

      final handler = _tools[call.name];
      if (handler == null) {
        input = 'Unknown tool "${call.name}". Choose from: ${_tools.keys}';
        continue;
      }

      final result = await handler(call.arguments);
      input = 'Result of ${call.name}: $result\nContinue. If done, answer.';
    }

    return 'I could not finish this task.';
  }
}

Three design decisions make this robust:

1. Bounded iteration. Never let the loop spin forever.
2. Corrective feedback. On unknown tool or parse failure, tell the model what went wrong and let it retry—don’t crash.
3. Explicit termination. The loop ends only when the model stops requesting tools.

Constrained Decoding: Guaranteed-Valid JSON

Free-form generation on small models produces malformed JSON surprisingly often—missing quotes, wrong nesting, truncated objects. Constrained decoding eliminates that entire class of bugs by restricting sampling to tokens that keep the output valid against your schema.

With LiteRT-LM and flutter_gemma:

final response = await conversation.sendMessage(
  'Add a meeting at 9am tomorrow',
  tools: [addMeetingTool],
  responseFormat: JsonSchema(
    schema: {
      'type': 'object',
      'properties': {
        'tool': {'type': 'string'},
        'args': {'type': 'object'},
      },
      'required': ['tool'],
    },
  ),
);

Because every emitted token provably moves toward a valid document, your parser can trust the structure. This is the single biggest reliability win available for on-device agents.

Parse-Error Fallbacks: Assume the Worst

Even with constrained decoding, assume structured output can fail—especially if you disable constraints for speed or use an older model. The resilient pattern:

ToolCall? tryParseToolCall(String text) {
  // 1. Fast path: model already returned clean JSON.
  final decoded = _tryDecodeJson(text);
  if (decoded != null && decoded is Map<String, dynamic>) {
    final tool = decoded['tool'] as String?;
    final args = decoded['args'];
    if (tool != null) return ToolCall(tool, args as Map<String, dynamic>? ?? {});
  }

  // 2. Fallback: search for a JSON object anywhere in the output.
  final match = RegExp(r'\{.*\}', dotAll: true).firstMatch(text);
  if (match != null) {
    final embedded = _tryDecodeJson(match.group(0)!);
    if (embedded != null) {
      return ToolCall(
        embedded['tool'] as String,
        (embedded['args'] as Map<String, dynamic>?) ?? {},
      );
    }
  }

  // 3. Last resort: return the raw text so the agent can retry or answer.
  return null;
}

The fallback chain is: exact JSON → embedded JSON → raw text. Every layer catches a different failure mode without ever throwing.

Never Round-Trip Ground Truth Through the Model

The most dangerous structured-output bug isn’t JSON syntax—it’s semantic hallucination. Small models routinely invent IDs, file paths, and references. Rule: validate every structured field against ground truth from app state; never round-trip them through the model.

final call = tryParseToolCall(response.text);
if (call == null) return retry();

// Validate against ground truth, not the model's claim.
if (call.name == 'open_note') {
  final requestedId = call.args['id'];
  final note = await notesStore.findById(requestedId);
  if (note == null) {
    return 'Note id "$requestedId" does not exist. Available ids: '
        '${await notesStore.allIds().take(5).join(', ')}...';
  }
  // Use note.id (ground truth), never the model's echoed id.
  return openNote(note);
}

When the model hallucinates an identifier, feed the correct list back as corrective context instead of failing silently.

Streaming Reasoning and Tool Calls

Gemma 4’s thinking mode gives the model a scratchpad for step-by-step reasoning. Stream it separately so the UI stays stable:

await for (final chunk in conversation.streamMessage(task, tools: tools)) {
  if (chunk.reasoning != null) {
    thinkingPanel.append(chunk.reasoning);
  } else if (chunk.toolCall != null) {
    toolLog.append('Calling ${chunk.toolCall.name}...');
  } else {
    answerArea.append(chunk.text);
  }
}

In long agentic sessions, consider stripping reasoning tokens entirely to save KV cache space—the tradeoff between transparency and memory is yours to tune.

Putting It All Together

A complete, hardened agent session:

Future<String> runAgentTask(String task) async {
  final model = await FlutterGemma.getActiveModel(maxTokens: 4096);
  final conversation = model.createConversation(
    systemInstruction: 'You are a helpful on-device assistant with tools.',
  );

  return ToolLoop(
    tools: {
      'schedule_reminder': scheduleReminder,
      'read_notes': readNotes,
    },
  ).run(conversation, task);
}

With constrained decoding for guaranteed-valid JSON, a bounded loop with corrective feedback, and ground-truth validation, your on-device agent becomes something users can actually rely on.

Testing Structured Output

Treat structured output as a first-class test concern:

void main() {
  test('tool call returns valid schema', () async {
    final call = tryParseToolCall('{"tool":"read_notes","args":{"keyword":"grocery"}}');
    expect(call?.name, 'read_notes');
    expect(call?.args['keyword'], 'grocery');
  });

  test('malformed JSON falls back gracefully', () {
    final call = tryParseToolCall('The result is {tool: read_notes}');
    expect(call, isNull);
  });

  test('embedded JSON is recovered', () {
    final call = tryParseToolCall('Sure! {"tool":"schedule_reminder","args":{"title":"x","at":"17:00"}}');
    expect(call?.name, 'schedule_reminder');
  });
}

Conclusion & Next Steps

You now know how to make on-device agents trustworthy: native function calling for action requests, constrained decoding for guaranteed-valid JSON, corrective fallbacks for the failures that still slip through, and ground-truth validation so the model can’t invent reality.

Next steps:

• Benchmark constrained decoding’s throughput cost vs free generation on your hardware.
• Persist tool-call history into session snapshots so interrupted workflows resume.
• Add a parse-rate metric to your analytics to track structured-output health.
• Combine with Agent Skills for a full on-device agent toolkit.

Small models can reason, plan, and act. With the right structured-output discipline, they can do it reliably—right on the user’s device.

References

• Bring state-of-the-art agentic skills to the edge with Gemma 4
• Blazing fast on-device GenAI with LiteRT-LM
• LiteRT-LM on GitHub
• Gemma 4 documentation