Structured output

Sometimes you want the generated output to be valid JSON with predefined fields. You can use Grammar to manually specify a JSON schema for the response you want to receive.

Python
Swift
TypeScript
Rust

Create a new project

uv init demo && cd demo

Install dependencies

uv add uzu

Paste into main.py

import asyncio
import json

from pydantic import BaseModel

from uzu import (
    ChatConfig,
    ChatMessage,
    ChatReplyConfig,
    Engine,
    EngineConfig,
    Grammar,
    ReasoningEffort,
)


class Country(BaseModel):
    name: str
    capital: str


class CountryList(BaseModel):
    countries: list[Country]


def structured_response(response: str | None, model_type: type[BaseModel]) -> BaseModel | None:
    if not response:
        return None
    return model_type.model_validate_json(response)


async def main() -> None:
    engine_config = EngineConfig.create()
    engine = await Engine.create(engine_config)

    model = await engine.model("Qwen/Qwen3-0.6B")
    if model is None:
        raise RuntimeError("Model not found")
    async for update in (await engine.download(model)).iterator():
        print(f"Download progress: {update.progress}")

    schema_string = json.dumps(CountryList.model_json_schema())
    messages = [
        ChatMessage.system().with_reasoning_effort(ReasoningEffort.Disabled),
        ChatMessage.user().with_text(
            "Give me a JSON object containing a list of 3 countries, where each country has name and capital fields"
        ),
    ]

    session = await engine.chat(model, ChatConfig.create())
    replies = await session.reply(
        messages,
        ChatReplyConfig.create().with_grammar(Grammar.JsonSchema(schema_string)),
    )
    if replies:
        countries = structured_response(replies[0].message.text, CountryList)
        print(countries)


if __name__ == "__main__":
    asyncio.run(main())

Run the snippet

uv run main.py

Create a new SwiftUI project

Add the SDK

Add this package through SPM:

https://github.com/trymirai/uzu.git

Paste the snippet

import FoundationModels
import Uzu

@Generable()
struct Country: Codable {
    let name: String
    let capital: String
}

public func runChatStructuredOutput() async throws {
    let engineConfig = EngineConfig.create()
    let engine = try await Engine.create(config: engineConfig)
    
    guard let model = try await engine.model(identifier: "Qwen/Qwen3-0.6B") else {
        return
    }
    for try await update in try await engine.download(model: model).iterator() {
        print("Download progress: \(update.progress())")
    }
    
    let messages = [
        ChatMessage.system().withReasoningEffort(reasoningEffort: .disabled),
        ChatMessage.user().withText(text: "Give me a JSON object containing a list of 3 countries, where each country has name and capital fields")
    ]
    
    let session = try await engine.chat(model: model, config: .create())
    let reply = try await session.reply(input: messages, config: .create().withGrammar(grammar: .fromType([Country].self)))
    guard let message = reply.last?.message else {
        return
    }
    guard let countries: [Country] = message.textDecoded() else {
        return
    }
    print(countries)
}

Add the snippet call

var body: some View {
    VStack {
        Text("On-device AI")
    }
    .onAppear() {
        Task {
            try await runExampleName()
        }
    }
}

Make a new project folder

mkdir demo && cd demo

Initialize with pnpm

pnpm init

Install dependencies

pnpm add typescript ts-node @types/node -D
pnpm add @trymirai/uzu

Initialize a tsconfig.json

{
    "compilerOptions": {
        "target": "es2020",
        "module": "commonjs",
        "moduleResolution": "node",
        "strict": true,
        "esModuleInterop": true,
        "outDir": "dist",
        "types": [
            "node"
        ]
    },
    "include": [
        "*.ts"
    ]
}

Create main.ts

import { ChatConfig, ChatMessage, ChatReplyConfig, Engine, EngineConfig, GrammarJsonSchema, ReasoningEffort } from '@trymirai/uzu';
import * as z from "zod";

const CountryType = z.object({
    name: z.string(),
    capital: z.string(),
});
const CountryListType = z.array(CountryType);

function structuredResponse<T extends z.ZodType>(response: string | null | undefined, type: T): z.infer<T> | undefined {
    if (!response) {
        return undefined;
    }
    const data = JSON.parse(response);
    const result = type.parse(data);
    return result;
}

async function main() {
    let engineConfig = EngineConfig.create();
    let engine = await Engine.create(engineConfig);

    let model = await engine.model('Qwen/Qwen3-0.6B');
    if (!model) {
        throw new Error('Model not found');
    }
    for await (const update of await engine.download(model)) {
        console.log('Download progress:', update.progress);
    }

    let schema = z.toJSONSchema(CountryListType);
    let schemaString = JSON.stringify(schema);
    let messages = [
        ChatMessage.system().withReasoningEffort("Disabled" as ReasoningEffort),
        ChatMessage.user().withText('Give me a JSON object containing a list of 3 countries, where each country has name and capital fields')
    ];

    let session = await engine.chat(model, ChatConfig.create());
    let reply = await session.reply(messages, ChatReplyConfig.create().withGrammar(new GrammarJsonSchema(schemaString)));
    let message = reply[0]?.message;
    let countries = structuredResponse(message?.text, CountryListType);
    console.log(countries);
}

main().catch((error) => {
    console.error(error);
});

Run the snippet

pnpm ts-node main.ts

Create a new project

cargo new demo && cd demo

Install dependencies

cargo add uzu --git https://github.com/trymirai/uzu
cargo add tokio --features full

Paste into src/main.rs

use schemars::{JsonSchema, schema_for};
use serde::{Deserialize, Serialize};
use uzu::{
    engine::{Engine, EngineConfig},
    types::{
        basic::{Grammar, ReasoningEffort},
        session::chat::{ChatConfig, ChatMessage, ChatReplyConfig},
    },
};

#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct Country {
    name: String,
    capital: String,
}

#[derive(Debug, Serialize, Deserialize, JsonSchema)]
struct CountryList {
    countries: Vec<Country>,
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let engine_config = EngineConfig::default();
    let engine = Engine::new(engine_config).await?;

    let model = engine.model("Qwen/Qwen3-0.6B".to_string()).await?.ok_or("Model not found")?;
    let downloader = engine.download(&model).await?;
    while let Some(update) = downloader.next().await {
        println!("Download progress: {}", update.progress());
    }

    let schema_string = serde_json::to_string(&schema_for!(CountryList))?;
    let messages = vec![
        ChatMessage::system().with_reasoning_effort(ReasoningEffort::Disabled),
        ChatMessage::user().with_text(
            "Give me a JSON object containing a list of 3 countries, where each country has name and capital fields"
                .to_string(),
        ),
    ];

    let session = engine.chat(model, ChatConfig::default()).await?;
    let chat_reply_config = ChatReplyConfig::default().with_grammar(Some(Grammar::JsonSchema {
        schema: schema_string,
    }));
    let replies = session.reply(messages, chat_reply_config).await?;
    if let Some(reply) = replies.first()
        && let Some(text) = reply.message.text()
    {
        let parsed: CountryList = serde_json::from_str(&text)?;
        println!("{parsed:#?}");
    }

    Ok(())
}

Run the snippet

cargo run --release

Overview

Quick start