Routing inference on-device first and escaping to the cloud only when it's worth it, in a Rork app

// ai/types.ts
export type Tier = 'cache' | 'on-device' | 'pcc' | 'remote';
 
// The nature of the task. The router reads this to pick the minimum required Tier.
export interface InferenceTask {
  // Feature id (used in telemetry and as part of the cache key)
  kind: 'summarize' | 'classify' | 'rewrite' | 'caption' | 'chat';
  input: string;
  // Does it carry an image? (if so, escalate to remote)
  image?: { uri: string; mime: string };
  // Same input, same acceptable output? (if true, cacheable)
  deterministic: boolean;
  // Max latency this feature tolerates (ms). If we'd exceed it, escalate or give up.
  latencyBudgetMs: number;
}
 
export interface InferenceResult {
  text: string;
  servedBy: Tier;     // which rung answered (for observation)
  costUsd: number;    // estimated charge (for budgeting)
  cached: boolean;
}

// ai/engine.ts
import type { InferenceTask, InferenceResult, Tier } from './types';
 
export interface Engine {
  tier: Tier;
  // Can this engine handle this task? (device capability, image support, etc.)
  canHandle(task: InferenceTask): Promise<boolean>;
  // Rough cost per call (used for budget decisions; 0 means free)
  estimateCost(task: InferenceTask): number;
  run(task: InferenceTask): Promise<InferenceResult>;
}

// ai/onDeviceEngine.ts
import { NativeModulesProxy } from 'expo-modules-core';
import type { Engine, InferenceTask, InferenceResult } from './types';
 
// Thin bridge to the native side (Foundation Models on iOS) built with expo-modules
const Native = NativeModulesProxy.OnDeviceAI as
  | { isAvailable(): Promise<boolean>; generate(prompt: string): Promise<string> }
  | undefined;
 
export const onDeviceEngine: Engine = {
  tier: 'on-device',
 
  async canHandle(task: InferenceTask) {
    // No native module, image input, or long text here
    if (!Native) return false;
    if (task.image) return false;
    if (task.input.length > 4000) return false;
    return Native.isAvailable();
  },
 
  estimateCost() {
    return 0; // on-device is free
  },
 
  async run(task: InferenceTask): Promise<InferenceResult> {
    const text = await Native!.generate(buildPrompt(task));
    return { text, servedBy: 'on-device', costUsd: 0, cached: false };
  },
};
 
function buildPrompt(task: InferenceTask): string {
  // Minimal instruction per task kind (short instructions are more stable on-device)
  const lead: Record<InferenceTask['kind'], string> = {
    summarize: 'Summarize the following in three sentences or fewer.',
    classify: 'Answer the main topic of the following in one word.',
    rewrite: 'Rewrite the following in clean, polite prose.',
    caption: 'Write a short headline that fits the content.',
    chat: '',
  };
  return `${lead[task.kind]}\n\n${task.input}`;
}

// ai/remoteEngine.ts
import type { Engine, InferenceTask, InferenceResult } from './types';
 
const PROXY_URL = 'https://your-worker.example.com/infer';
 
export const remoteEngine: Engine = {
  tier: 'remote',
 
  async canHandle() {
    // The last resort. Takes images and long text. Just needs a connection.
    return true;
  },
 
  estimateCost(task: InferenceTask) {
    // Approx input tokens x unit price (model choice happens server-side)
    const approxTokens = Math.ceil(task.input.length / 3);
    return (approxTokens / 1000) * 0.003; // rough USD
  },
 
  async run(task: InferenceTask): Promise<InferenceResult> {
    const res = await fetch(PROXY_URL, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        kind: task.kind,
        input: task.input,
        image: task.image ?? null,
      }),
    });
    if (!res.ok) throw new Error(`remote ${res.status}`);
    const json = (await res.json()) as { text: string; costUsd?: number };
    return {
      text: json.text,
      servedBy: 'remote',
      costUsd: json.costUsd ?? this.estimateCost(task),
      cached: false,
    };
  },
};

// ai/router.ts
import type { Engine, InferenceTask, InferenceResult, Tier } from './types';
 
const TIER_ORDER: Tier[] = ['cache', 'on-device', 'pcc', 'remote'];
 
// The floor: "this task needs at least this rung"
function minTierFor(task: InferenceTask): Tier {
  if (task.image) return 'remote';        // images aren't handled on-device
  if (task.kind === 'chat') return 'pcc';  // chat favors accuracy: PCC floor
  return 'on-device';                      // classify/summarize/rewrite start on-device
}
 
export class AIRouter {
  private spentTodayUsd = 0;
 
  constructor(
    private engines: Engine[],
    private dailyBudgetUsd: number,
  ) {}
 
  private orderedEngines(min: Tier): Engine[] {
    const minIdx = TIER_ORDER.indexOf(min);
    return this.engines
      .filter((e) => TIER_ORDER.indexOf(e.tier) >= minIdx)
      .sort((a, b) => TIER_ORDER.indexOf(a.tier) - TIER_ORDER.indexOf(b.tier));
  }
 
  async route(task: InferenceTask): Promise<InferenceResult> {
    const min = minTierFor(task);
    const candidates = this.orderedEngines(min);
    let lastError: unknown;
 
    for (const engine of candidates) {
      // Skip paid rungs that would blow the budget (free rungs always allowed)
      const cost = engine.estimateCost(task);
      if (cost > 0 && this.spentTodayUsd + cost > this.dailyBudgetUsd) {
        continue;
      }
      if (!(await engine.canHandle(task))) continue;
 
      try {
        const result = await withTimeout(engine.run(task), task.latencyBudgetMs);
        this.spentTodayUsd += result.costUsd;
        return result;
      } catch (err) {
        // This rung failed; drop to the next (fallback)
        lastError = err;
        continue;
      }
    }
    throw new Error(`all tiers failed: ${String(lastError)}`);
  }
}
 
// Timeout-bounded execution. On timeout, reject so the router drops a rung.
function withTimeout<T>(p: Promise<T>, ms: number): Promise<T> {
  return new Promise((resolve, reject) => {
    const id = setTimeout(() => reject(new Error('timeout')), ms);
    p.then(
      (v) => { clearTimeout(id); resolve(v); },
      (e) => { clearTimeout(id); reject(e); },
    );
  });
}

// ai/cacheEngine.ts
import * as Crypto from 'expo-crypto';
import type { Engine, InferenceTask, InferenceResult } from './types';
 
type Store = {
  get(key: string): Promise<string | null>;
  set(key: string, value: string): Promise<void>;
};
 
export function createCacheEngine(store: Store): Engine {
  return {
    tier: 'cache',
 
    async canHandle(task: InferenceTask) {
      // Non-deterministic or image tasks aren't cacheable
      return task.deterministic && !task.image;
    },
 
    estimateCost() {
      return 0;
    },
 
    async run(task: InferenceTask): Promise<InferenceResult> {
      const key = await keyFor(task);
      const hit = await store.get(key);
      if (hit === null) {
        // Treat a miss as a failure so the router drops to the next rung
        throw new Error('cache miss');
      }
      return { text: hit, servedBy: 'cache', costUsd: 0, cached: true };
    },
  };
}
 
async function keyFor(task: InferenceTask): Promise<string> {
  return Crypto.digestStringAsync(
    Crypto.CryptoDigestAlgorithm.SHA256,
    `${task.kind}:${task.input}`,
  );
}

// Inside router.route(), just before return
if (task.deterministic && result.servedBy !== 'cache') {
  const key = await keyFor(task);        // share the key function with cacheEngine
  await cacheStore.set(key, result.text);
}