Text & image generation

// Ask a vision-language model about an image and stream the answer, then
// ask a follow-up question in the same conversation.
//
// start() reads the prompt; next_token() returns one piece of text per
// call. There are no callbacks: the app pulls tokens in a loop.

import "stdlib/io" as io;
import "stdlib/env" as env;
import "stdlib/fs" as fs;
import "stdlib/result" as result;
import "stdlib/option" as option;
import "stdlib/str" as _;
import "ffi/api" as slab;

fn stream(ref session: slab::Session, ref inp: slab::Input, ref out: slab::Output) -> bool {
    if !session.start(inp, out) { return false; }
    while session.next_token(out) == 1 {
        io::print(out.text());
    }
    io::println("");
    return out.finish_reason() != slab::FINISH_NONE;
}

fn main() -> i32 {
    var cfg = slab::Config::new();
    cfg.set_models_dir("models");
    guard let option::Option::Some(path) = env::var("SLAB_RUNTIME_LICENSE_FILE") else { return 1; };
    guard let result::Result::Ok(token) = fs::read_to_string(path.view()) else { return 1; };
    if !cfg.set_runtime_license(token.view()) { return 1; }

    var session = slab::Session::new(cfg);
    session.set_context_length(2048);
    if !session.open("vlm.slab") {
        io::eprintln("Cannot open the model: ${session.error_message()}");
        return 1;
    }

    var photo = slab::Frame::new();
    if !photo.load_file("images/receipt.jpg") {
        io::eprintln("Cannot load the image: ${photo.error_message()}");
        return 1;
    }

    var inp = slab::Input::new();
    var out = slab::Output::new();
    inp.set_max_tokens(64);
    inp.set_temperature(0.0);
    // Large images are split into tiles instead of being shrunk.
    inp.set_image_tiled(true, 1536);

    // Turn 1: a question about the image.
    let _ = inp.add_image_message("user", photo, "What is this document?");
    if !stream(session, inp, out) {
        io::eprintln("Generation failed: ${session.error_message()}");
        return 1;
    }

    // Turn 2: send the conversation again with the answer and a new question.
    // The shared beginning is reused, not recomputed.
    let answer = out.full_text().trim().to_text();
    let _ = inp.add_message("assistant", answer.view());
    let _ = inp.add_message("user", "What is the total amount?");
    if !stream(session, inp, out) {
        io::eprintln("Generation failed: ${session.error_message()}");
        return 1;
    }
    io::println("prompt ${out.prompt_tokens()} tokens (${out.cached_prompt_tokens()} reused), generated ${out.generated_tokens()}");

    // Forget the conversation before an unrelated one.
    let _ = session.reset();
    return 0;
}