react-native-gemma-agent 0.1.0

package/src/SkillSandbox.tsx

@@ -0,0 +1,155 @@
+import React, {
+  useRef,
+  useCallback,
+  useImperativeHandle,
+  forwardRef,
+  useState,
+} from 'react';
+import { View, StyleSheet } from 'react-native';
+import { WebView, type WebViewMessageEvent } from 'react-native-webview';
+import type { SkillResult } from './types';
+
+export type SkillSandboxHandle = {
+  execute: (
+    html: string,
+    params: Record<string, unknown>,
+    timeout?: number,
+  ) => Promise<SkillResult>;
+};
+
+type PendingExecution = {
+  resolve: (result: SkillResult) => void;
+  reject: (error: Error) => void;
+  timer: ReturnType<typeof setTimeout>;
+};
+
+export const SkillSandbox = forwardRef<SkillSandboxHandle>(
+  function SkillSandbox(_props, ref) {
+    const pendingRef = useRef<PendingExecution | null>(null);
+    const [source, setSource] = useState<{ html: string } | undefined>(
+      undefined,
+    );
+
+    const handleMessage = useCallback((event: WebViewMessageEvent) => {
+      const pending = pendingRef.current;
+      if (!pending) return;
+
+      pendingRef.current = null;
+      clearTimeout(pending.timer);
+
+      try {
+        const data = JSON.parse(event.nativeEvent.data);
+        if (data.type === 'skill_error') {
+          pending.resolve({ error: data.error });
+        } else if (data.type === 'skill_result') {
+          pending.resolve(data.data);
+        } else {
+          pending.resolve({ error: 'Unknown response type from skill' });
+        }
+      } catch {
+        pending.resolve({ error: 'Failed to parse skill response' });
+      }
+    }, []);
+
+    useImperativeHandle(
+      ref,
+      () => ({
+        execute: (
+          html: string,
+          params: Record<string, unknown>,
+          timeout = 30_000,
+        ): Promise<SkillResult> => {
+          // Cancel any pending execution
+          if (pendingRef.current) {
+            clearTimeout(pendingRef.current.timer);
+            pendingRef.current.reject(
+              new Error('Execution cancelled — new skill started'),
+            );
+            pendingRef.current = null;
+          }
+
+          return new Promise<SkillResult>((resolve, reject) => {
+            const timer = setTimeout(() => {
+              pendingRef.current = null;
+              resolve({ error: `Skill timed out after ${timeout}ms` });
+            }, timeout);
+
+            pendingRef.current = { resolve, reject, timer };
+
+            const injectedHtml = injectExecutionScript(html, params);
+            setSource({ html: injectedHtml });
+          });
+        },
+      }),
+      [],
+    );
+
+    return (
+      <View style={styles.container} pointerEvents="none">
+        {source && (
+          <WebView
+            source={source}
+            onMessage={handleMessage}
+            originWhitelist={['*']}
+            javaScriptEnabled
+            domStorageEnabled={false}
+            incognito
+            style={styles.webview}
+          />
+        )}
+      </View>
+    );
+  },
+);
+
+/**
+ * Inject the execution bridge script into skill HTML.
+ * Calls window['ai_edge_gallery_get_result'] with the params
+ * and posts the result back via ReactNativeWebView.postMessage.
+ */
+function injectExecutionScript(
+  html: string,
+  params: Record<string, unknown>,
+): string {
+  const escapedParams = JSON.stringify(JSON.stringify(params));
+
+  const executionScript = `
+<script>
+(async () => {
+  await new Promise(r => setTimeout(r, 100));
+  try {
+    const fn = window['ai_edge_gallery_get_result'];
+    if (!fn) throw new Error('Skill function ai_edge_gallery_get_result not found');
+    const result = await fn(${escapedParams});
+    window.ReactNativeWebView.postMessage(JSON.stringify({
+      type: 'skill_result',
+      data: JSON.parse(result)
+    }));
+  } catch (e) {
+    window.ReactNativeWebView.postMessage(JSON.stringify({
+      type: 'skill_error',
+      error: e.message || 'Unknown skill error'
+    }));
+  }
+})();
+</script>`;
+
+  if (html.includes('</body>')) {
+    return html.replace('</body>', `${executionScript}\n</body>`);
+  }
+  return html + executionScript;
+}
+
+const styles = StyleSheet.create({
+  container: {
+    position: 'absolute',
+    width: 0,
+    height: 0,
+    overflow: 'hidden',
+  },
+  webview: {
+    width: 1,
+    height: 1,
+    opacity: 0,
+  },
+});

package/src/index.ts

@@ -0,0 +1,52 @@
+// react-native-gemma-agent SDK
+export const SDK_VERSION = '0.1.0';
+
+// Core classes
+export { ModelManager } from './ModelManager';
+export { InferenceEngine } from './InferenceEngine';
+export { SkillRegistry } from './SkillRegistry';
+export { AgentOrchestrator } from './AgentOrchestrator';
+export { BM25Scorer } from './BM25Scorer';
+export type { SkillExecutor } from './AgentOrchestrator';
+
+// React hooks & provider
+export { GemmaAgentProvider } from './GemmaAgentProvider';
+export type { GemmaAgentProviderProps, GemmaAgentContextValue } from './GemmaAgentProvider';
+export { useGemmaAgent } from './useGemmaAgent';
+export type { UseGemmaAgentReturn } from './useGemmaAgent';
+export { useModelDownload } from './useModelDownload';
+export type { UseModelDownloadReturn } from './useModelDownload';
+export { useSkillRegistry } from './useSkillRegistry';
+export type { UseSkillRegistryReturn } from './useSkillRegistry';
+
+// Skill execution
+export { SkillSandbox } from './SkillSandbox';
+export type { SkillSandboxHandle } from './SkillSandbox';
+
+// Function call parsing
+export { validateToolCalls, extractToolCallsFromText } from './FunctionCallParser';
+export type { ParsedToolCall } from './FunctionCallParser';
+
+// Types
+export type {
+  ModelStatus,
+  ModelConfig,
+  ModelInfo,
+  DownloadProgress,
+  ToolDefinition,
+  ToolCall,
+  Message,
+  MessageRole,
+  CompletionResult,
+  CompletionTimings,
+  GenerateOptions,
+  TokenEvent,
+  InferenceEngineConfig,
+  SkillType,
+  SkillParameter,
+  SkillManifest,
+  SkillResult,
+  AgentEvent,
+  AgentConfig,
+  ContextUsage,
+} from './types';

package/src/types.ts

@@ -0,0 +1,197 @@
+// react-native-gemma-agent types
+
+export type ModelStatus =
+  | 'not_downloaded'
+  | 'downloading'
+  | 'ready'
+  | 'loading'
+  | 'loaded'
+  | 'error';
+
+export type ModelConfig = {
+  /** HuggingFace repo ID (e.g., 'unsloth/gemma-4-E2B-it-GGUF') */
+  repoId: string;
+  /** GGUF filename within the repo */
+  filename: string;
+  /** Expected file size in bytes (for progress calculation) */
+  expectedSize?: number;
+  /** SHA256 checksum for verification */
+  checksum?: string;
+};
+
+export type DownloadProgress = {
+  bytesDownloaded: number;
+  totalBytes: number;
+  /** 0-100 */
+  percent: number;
+};
+
+export type ModelInfo = {
+  status: ModelStatus;
+  path: string | null;
+  sizeBytes: number | null;
+  description: string | null;
+  nParams: number | null;
+  nEmbd: number | null;
+};
+
+export type ToolDefinition = {
+  type: 'function';
+  function: {
+    name: string;
+    description: string;
+    parameters: {
+      type: 'object';
+      properties: Record<string, {
+        type: string;
+        description?: string;
+        enum?: string[];
+      }>;
+      required?: string[];
+    };
+  };
+};
+
+export type ToolCall = {
+  type: 'function';
+  id?: string;
+  function: {
+    name: string;
+    arguments: string;
+  };
+};
+
+export type MessageRole = 'user' | 'assistant' | 'system' | 'tool';
+
+export type Message = {
+  role: MessageRole;
+  content: string;
+  tool_calls?: ToolCall[];
+  tool_call_id?: string;
+  name?: string;
+};
+
+export type CompletionTimings = {
+  promptTokens: number;
+  promptMs: number;
+  promptPerSecond: number;
+  predictedTokens: number;
+  predictedMs: number;
+  predictedPerSecond: number;
+};
+
+export type CompletionResult = {
+  /** Raw text output (includes thinking + tool call tokens) */
+  text: string;
+  /** Filtered content (thinking and tool call tokens removed) */
+  content: string;
+  /** Model's chain-of-thought reasoning, if any */
+  reasoning: string | null;
+  toolCalls: ToolCall[];
+  timings: CompletionTimings;
+  stoppedEos: boolean;
+  stoppedLimit: boolean;
+  contextFull: boolean;
+};
+
+export type GenerateOptions = {
+  /** Max tokens to generate */
+  maxTokens?: number;
+  /** Sampling temperature (0-2) */
+  temperature?: number;
+  /** Top-p nucleus sampling */
+  topP?: number;
+  /** Top-k sampling */
+  topK?: number;
+  /** Stop sequences */
+  stop?: string[];
+  /** Tool definitions for function calling */
+  tools?: ToolDefinition[];
+  /** Tool choice strategy */
+  toolChoice?: 'auto' | 'none' | string;
+};
+
+export type TokenEvent = {
+  token: string;
+  toolCalls?: ToolCall[];
+};
+
+export type InferenceEngineConfig = {
+  /** Context window size in tokens */
+  contextSize?: number;
+  /** Batch size for prompt processing */
+  batchSize?: number;
+  /** Number of threads for inference */
+  threads?: number;
+  /** Flash attention mode */
+  flashAttn?: 'auto' | 'on' | 'off';
+  /** Lock model in memory */
+  useMlock?: boolean;
+  /** Number of GPU layers to offload (-1 = all) */
+  gpuLayers?: number;
+};
+
+// --- Skill System Types ---
+
+export type SkillType = 'js' | 'native';
+
+export type SkillParameter = {
+  type: string;
+  description?: string;
+  enum?: string[];
+};
+
+export type SkillManifest = {
+  name: string;
+  description: string;
+  version: string;
+  type: SkillType;
+  parameters: Record<string, SkillParameter>;
+  requiredParameters?: string[];
+  /** HTML content for JS skills — loaded into hidden WebView */
+  html?: string;
+  /** Native execution function — for skills that run in RN context */
+  execute?: (params: Record<string, unknown>) => Promise<SkillResult>;
+  /** Instructions for the LLM on when/how to use this skill */
+  instructions?: string;
+  /** Whether this skill requires network access. SDK checks connectivity before execution. */
+  requiresNetwork?: boolean;
+};
+
+export type SkillResult = {
+  result?: string;
+  error?: string;
+  image?: { base64: string };
+};
+
+// --- Agent Types ---
+
+export type AgentEvent =
+  | { type: 'thinking' }
+  | { type: 'token'; token: string }
+  | { type: 'skill_called'; name: string; parameters: Record<string, unknown> }
+  | { type: 'skill_result'; name: string; result: SkillResult }
+  | { type: 'response'; text: string; reasoning: string | null }
+  | { type: 'error'; error: string };
+
+export type AgentConfig = {
+  /** Max chained skill calls before stopping (prevents infinite loops). Default: 5 */
+  maxChainDepth?: number;
+  /** Timeout for each skill execution in ms. Default: 30000 */
+  skillTimeout?: number;
+  /** Base system prompt prepended to all conversations */
+  systemPrompt?: string;
+  /** Skill routing strategy. 'all' sends every skill; 'bm25' pre-filters by query relevance. Default: 'all' */
+  skillRouting?: 'all' | 'bm25';
+  /** Max skills sent to the model per invocation (only used with 'bm25' routing). Default: 5 */
+  maxToolsPerInvocation?: number;
+};
+
+export type ContextUsage = {
+  /** Tokens used so far */
+  used: number;
+  /** Total context window size in tokens */
+  total: number;
+  /** Usage percentage (0-100) */
+  percent: number;
+};

package/src/useGemmaAgent.ts

@@ -0,0 +1,222 @@
+import { useState, useCallback, useRef } from 'react';
+import { useGemmaAgentContext } from './GemmaAgentProvider';
+import type { Message, ModelStatus, AgentEvent, ContextUsage } from './types';
+
+export type UseGemmaAgentReturn = {
+  /** Send a message to the agent. Returns the final response text. */
+  sendMessage: (
+    text: string,
+    onEvent?: (event: AgentEvent) => void,
+  ) => Promise<string>;
+  /** Conversation history */
+  messages: ReadonlyArray<Message>;
+  /** Content tokens streamed so far (thinking excluded). Empty when idle. */
+  streamingText: string;
+  /** Whether the agent is currently processing */
+  isProcessing: boolean;
+  /** Whether the model is loaded and ready for inference */
+  isModelLoaded: boolean;
+  /** Current model status */
+  modelStatus: ModelStatus;
+  /** Name of the skill currently being executed, or null */
+  activeSkill: string | null;
+  /** Last error message, or null */
+  error: string | null;
+  /** Current context window usage (updated after each generation) */
+  contextUsage: ContextUsage;
+  /** Load the model into memory. Must be downloaded first. */
+  loadModel: (onProgress?: (percent: number) => void) => Promise<number>;
+  /** Unload the model from memory */
+  unloadModel: () => Promise<void>;
+  /** Clear conversation history */
+  reset: () => void;
+};
+
+export function useGemmaAgent(): UseGemmaAgentReturn {
+  const { modelManager, engine, orchestrator } = useGemmaAgentContext();
+
+  const [messages, setMessages] = useState<ReadonlyArray<Message>>([]);
+  const [streamingText, setStreamingText] = useState('');
+  const [isProcessing, setIsProcessing] = useState(false);
+  const [modelStatus, setModelStatus] = useState<ModelStatus>(
+    modelManager.status,
+  );
+  const [activeSkill, setActiveSkill] = useState<string | null>(null);
+  const [error, setError] = useState<string | null>(null);
+  const [contextUsage, setContextUsage] = useState<ContextUsage>({
+    used: 0,
+    total: 0,
+    percent: 0,
+  });
+
+  // Track model status changes
+  const unsubRef = useRef<(() => void) | null>(null);
+  if (!unsubRef.current) {
+    unsubRef.current = modelManager.onStatusChange((status) => {
+      setModelStatus(status);
+    });
+  }
+
+  // Buffer to detect special token sequences in streaming
+  const tokenBufferRef = useRef('');
+  // Whether we've seen the model produce content tokens (after thinking)
+  const seenContentRef = useRef(false);
+
+  const loadModel = useCallback(
+    async (onProgress?: (percent: number) => void): Promise<number> => {
+      let path = modelManager.modelPath;
+      if (!path) {
+        path = await modelManager.findModel();
+      }
+      if (!path) {
+        throw new Error(
+          'Model not found on device. Download it first via useModelDownload().',
+        );
+      }
+
+      setModelStatus('loading');
+      try {
+        const loadTimeMs = await engine.loadModel(path, onProgress);
+        setModelStatus('loaded');
+        return loadTimeMs;
+      } catch (err) {
+        setModelStatus('error');
+        throw err;
+      }
+    },
+    [modelManager, engine],
+  );
+
+  const unloadModel = useCallback(async () => {
+    await engine.unload();
+    setModelStatus(modelManager.modelPath ? 'ready' : 'not_downloaded');
+  }, [engine, modelManager]);
+
+  const sendMessage = useCallback(
+    async (
+      text: string,
+      onEvent?: (event: AgentEvent) => void,
+    ): Promise<string> => {
+      if (!engine.isLoaded) {
+        throw new Error('Model not loaded. Call loadModel() first.');
+      }
+
+      setError(null);
+      setIsProcessing(true);
+      setActiveSkill(null);
+      setStreamingText('');
+      tokenBufferRef.current = '';
+      seenContentRef.current = false;
+
+      try {
+        const response = await orchestrator.sendMessage(
+          text,
+          (event: AgentEvent) => {
+            switch (event.type) {
+              case 'token': {
+                // Accumulate into buffer to detect thinking vs content
+                tokenBufferRef.current += event.token;
+                const buf = tokenBufferRef.current;
+
+                // Gemma 4 outputs thinking as: thought\n<reasoning>
+                // followed by content after tool call results or directly.
+                // The `content` field in the final result has thinking stripped.
+                // For streaming: skip tokens until we see content starting.
+                // Since llama.rn gives us `content` (filtered) at the end,
+                // streamingText is a best-effort preview. We use it only
+                // when the model is generating the final answer (not thinking
+                // before a tool call).
+
+                // Don't stream during thinking phase — the model starts with
+                // "thought\n" or similar when reasoning before tool calls.
+                // We detect this by checking if the buffer starts with "thought"
+                if (!seenContentRef.current) {
+                  if (buf.length >= 7) {
+                    if (buf.trimStart().startsWith('thought')) {
+                      // In thinking mode — don't stream these tokens
+                      break;
+                    }
+                    // Not thinking — start streaming
+                    seenContentRef.current = true;
+                    setStreamingText(buf);
+                  }
+                  // Still buffering — wait for more tokens
+                  break;
+                }
+
+                // We're in content mode — stream normally
+                setStreamingText((prev) => prev + event.token);
+                break;
+              }
+              case 'thinking':
+                // Reset streaming state at the start of each generation loop.
+                // Ensures clean slate before the final answer turn (after tool results).
+                setStreamingText('');
+                tokenBufferRef.current = '';
+                seenContentRef.current = false;
+                break;
+              case 'skill_called':
+                setActiveSkill(event.name);
+                setStreamingText('');
+                tokenBufferRef.current = '';
+                seenContentRef.current = false;
+                break;
+              case 'skill_result':
+                setActiveSkill(null);
+                break;
+              case 'error':
+                setError(event.error);
+                break;
+              case 'response':
+                setStreamingText('');
+                tokenBufferRef.current = '';
+                seenContentRef.current = false;
+                setMessages([...orchestrator.conversation]);
+                break;
+            }
+            // Always forward raw events to the developer's callback
+            onEvent?.(event);
+          },
+        );
+
+        setMessages([...orchestrator.conversation]);
+        setContextUsage(engine.getContextUsage());
+        return response;
+      } catch (err) {
+        const msg = err instanceof Error ? err.message : 'Unknown error';
+        setError(msg);
+        throw err;
+      } finally {
+        setIsProcessing(false);
+        setActiveSkill(null);
+        setStreamingText('');
+        tokenBufferRef.current = '';
+        seenContentRef.current = false;
+      }
+    },
+    [engine, orchestrator],
+  );
+
+  const reset = useCallback(() => {
+    orchestrator.reset();
+    setMessages([]);
+    setStreamingText('');
+    setError(null);
+    setActiveSkill(null);
+  }, [orchestrator]);
+
+  return {
+    sendMessage,
+    messages,
+    streamingText,
+    isProcessing,
+    isModelLoaded: engine.isLoaded,
+    modelStatus,
+    activeSkill,
+    error,
+    contextUsage,
+    loadModel,
+    unloadModel,
+    reset,
+  };
+}

package/src/useModelDownload.ts

@@ -0,0 +1,80 @@
+import { useState, useCallback, useRef } from 'react';
+import { useGemmaAgentContext } from './GemmaAgentProvider';
+import type { ModelStatus, DownloadProgress } from './types';
+
+export type UseModelDownloadReturn = {
+  /** Start downloading the model. Resolves with the local file path. */
+  download: () => Promise<string>;
+  /** Cancel an in-progress download */
+  cancelDownload: () => void;
+  /** Check if model already exists on device */
+  checkModel: () => Promise<boolean>;
+  /** Set a custom model path (for pre-downloaded models) */
+  setModelPath: (path: string) => Promise<void>;
+  /** Delete the downloaded model file */
+  deleteModel: () => Promise<void>;
+  /** Current download progress */
+  progress: DownloadProgress | null;
+  /** Current model status */
+  status: ModelStatus;
+  /** Check available storage */
+  checkStorage: () => Promise<{ available: number; required: number; sufficient: boolean }>;
+};
+
+export function useModelDownload(): UseModelDownloadReturn {
+  const { modelManager } = useGemmaAgentContext();
+
+  const [status, setStatus] = useState<ModelStatus>(modelManager.status);
+  const [progress, setProgress] = useState<DownloadProgress | null>(null);
+
+  const unsubRef = useRef<(() => void) | null>(null);
+  if (!unsubRef.current) {
+    unsubRef.current = modelManager.onStatusChange((s) => {
+      setStatus(s);
+    });
+  }
+
+  const download = useCallback(async (): Promise<string> => {
+    setProgress(null);
+    const path = await modelManager.download((p) => {
+      setProgress(p);
+    });
+    return path;
+  }, [modelManager]);
+
+  const cancelDownload = useCallback(() => {
+    modelManager.cancelDownload();
+    setProgress(null);
+  }, [modelManager]);
+
+  const checkModel = useCallback(
+    () => modelManager.checkModel(),
+    [modelManager],
+  );
+
+  const setModelPath = useCallback(
+    (path: string) => modelManager.setModelPath(path),
+    [modelManager],
+  );
+
+  const deleteModel = useCallback(
+    () => modelManager.deleteModel(),
+    [modelManager],
+  );
+
+  const checkStorage = useCallback(
+    () => modelManager.checkStorage(),
+    [modelManager],
+  );
+
+  return {
+    download,
+    cancelDownload,
+    checkModel,
+    setModelPath,
+    deleteModel,
+    progress,
+    status,
+    checkStorage,
+  };
+}