react-native-gemma-agent 0.1.0

package/src/AgentOrchestrator.ts

@@ -0,0 +1,315 @@
+import type {
+  Message,
+  AgentEvent,
+  AgentConfig,
+  SkillResult,
+} from './types';
+import type { InferenceEngine } from './InferenceEngine';
+import type { SkillRegistry } from './SkillRegistry';
+import { BM25Scorer } from './BM25Scorer';
+import {
+  validateToolCalls,
+  extractToolCallsFromText,
+  type ParsedToolCall,
+} from './FunctionCallParser';
+
+const DEFAULT_CONFIG: Required<AgentConfig> = {
+  maxChainDepth: 5,
+  skillTimeout: 30_000,
+  systemPrompt:
+    'You are a helpful AI assistant running on-device. Answer concisely and accurately.',
+  skillRouting: 'all',
+  maxToolsPerInvocation: 5,
+};
+
+export type SkillExecutor = (
+  html: string,
+  params: Record<string, unknown>,
+  timeout?: number,
+) => Promise<SkillResult>;
+
+export class AgentOrchestrator {
+  private engine: InferenceEngine;
+  private registry: SkillRegistry;
+  private executor: SkillExecutor | null = null;
+  private config: Required<AgentConfig>;
+  private history: Message[] = [];
+  private _isProcessing = false;
+  private bm25: BM25Scorer = new BM25Scorer();
+
+  constructor(
+    engine: InferenceEngine,
+    registry: SkillRegistry,
+    config?: AgentConfig,
+  ) {
+    this.engine = engine;
+    this.registry = registry;
+    this.config = { ...DEFAULT_CONFIG, ...config };
+  }
+
+  get isProcessing(): boolean {
+    return this._isProcessing;
+  }
+
+  get conversation(): ReadonlyArray<Message> {
+    return this.history;
+  }
+
+  /**
+   * Set the JS skill executor. Wired by the React layer to SkillSandbox.
+   * Not needed if you only use native skills.
+   */
+  setSkillExecutor(executor: SkillExecutor): void {
+    this.executor = executor;
+  }
+
+  /**
+   * Send a user message through the full agent loop:
+   *   inference → tool call detection → skill execution → re-invoke model
+   *
+   * Returns the final assistant response text.
+   */
+  async sendMessage(
+    text: string,
+    onEvent?: (event: AgentEvent) => void,
+  ): Promise<string> {
+    if (this._isProcessing) {
+      throw new Error('Already processing a message. Wait for completion.');
+    }
+
+    this._isProcessing = true;
+
+    try {
+      this.history = [...this.history, { role: 'user', content: text }];
+
+      const tools = this.getToolsForQuery(text);
+      let depth = 0;
+
+      while (depth < this.config.maxChainDepth) {
+        depth++;
+
+        const messages: Message[] = [
+          { role: 'system', content: this.config.systemPrompt },
+          ...this.history,
+        ];
+
+        onEvent?.({ type: 'thinking' });
+
+        const result = await this.engine.generate(
+          messages,
+          {
+            tools: tools.length > 0 ? tools : undefined,
+            toolChoice: tools.length > 0 ? 'auto' : undefined,
+          },
+          (tokenEvent) => {
+            onEvent?.({ type: 'token', token: tokenEvent.token });
+          },
+        );
+
+        // Check for tool calls — primary (llama.rn native) then fallback (text scan)
+        let parsedCalls = validateToolCalls(result.toolCalls, this.registry);
+        if (parsedCalls.length === 0 && result.text.trim()) {
+          parsedCalls = extractToolCallsFromText(result.text, this.registry);
+        }
+
+        // No tool calls → final response
+        if (parsedCalls.length === 0) {
+          const responseText = result.content || result.text;
+          this.history = [
+            ...this.history,
+            { role: 'assistant', content: responseText },
+          ];
+          onEvent?.({
+            type: 'response',
+            text: responseText,
+            reasoning: result.reasoning,
+          });
+          return responseText;
+        }
+
+        // Add assistant message with tool_calls to history.
+        // Strip thinking/reasoning from content — it leaks into chat UI otherwise.
+        // Empty string is safe: llama.rn's Jinja template handles empty content
+        // on assistant messages with tool_calls (OpenAI-compatible format).
+        this.history = [
+          ...this.history,
+          {
+            role: 'assistant',
+            content: '',
+            tool_calls: result.toolCalls,
+          },
+        ];
+
+        // Execute each tool call and append results
+        for (const call of parsedCalls) {
+          onEvent?.({
+            type: 'skill_called',
+            name: call.name,
+            parameters: call.parameters,
+          });
+
+          const skillResult = await this.executeSkill(call);
+
+          onEvent?.({
+            type: 'skill_result',
+            name: call.name,
+            result: skillResult,
+          });
+
+          const resultContent = skillResult.error
+            ? `Error: ${skillResult.error}`
+            : skillResult.result ?? 'No result';
+
+          // tool_call_id must be a string — generate one if llama.rn didn't provide it
+          const toolCallId =
+            call.id ??
+            result.toolCalls.find(tc => tc.function.name === call.name)?.id ??
+            `call_${call.name}_${depth}`;
+
+          this.history = [
+            ...this.history,
+            {
+              role: 'tool',
+              content: resultContent,
+              tool_call_id: toolCallId,
+              name: call.name,
+            },
+          ];
+        }
+
+        // Loop back — model will see tool results and generate a response
+      }
+
+      // Max chain depth reached
+      const fallback =
+        'I tried to use tools but reached the maximum chain depth. Here is what I know so far.';
+      this.history = [
+        ...this.history,
+        { role: 'assistant', content: fallback },
+      ];
+      onEvent?.({ type: 'response', text: fallback, reasoning: null });
+      return fallback;
+    } catch (err) {
+      const errorMsg =
+        err instanceof Error ? err.message : 'Unknown error';
+      onEvent?.({ type: 'error', error: errorMsg });
+      throw err;
+    } finally {
+      this._isProcessing = false;
+    }
+  }
+
+  reset(): void {
+    this.history = [];
+  }
+
+  setSystemPrompt(prompt: string): void {
+    this.config = { ...this.config, systemPrompt: prompt };
+  }
+
+  private getToolsForQuery(query: string) {
+    if (this.config.skillRouting !== 'bm25') {
+      return this.registry.toToolDefinitions();
+    }
+
+    const allSkills = this.registry.getSkills();
+    if (allSkills.length <= this.config.maxToolsPerInvocation) {
+      return this.registry.toToolDefinitions();
+    }
+
+    this.bm25.buildIndex(allSkills);
+    const ranked = this.bm25.topN(query, this.config.maxToolsPerInvocation);
+
+    return ranked.map(({ skill }) => ({
+      type: 'function' as const,
+      function: {
+        name: skill.name,
+        description:
+          skill.description +
+          (skill.instructions ? `\n${skill.instructions}` : ''),
+        parameters: {
+          type: 'object' as const,
+          properties: skill.parameters,
+          required: skill.requiredParameters,
+        },
+      },
+    }));
+  }
+
+  private async checkConnectivity(): Promise<boolean> {
+    try {
+      const controller = new AbortController();
+      const timeout = setTimeout(() => controller.abort(), 3000);
+      await fetch('https://www.google.com/generate_204', {
+        method: 'HEAD',
+        signal: controller.signal,
+      });
+      clearTimeout(timeout);
+      return true;
+    } catch {
+      return false;
+    }
+  }
+
+  private async executeSkill(call: ParsedToolCall): Promise<SkillResult> {
+    const { skill, parameters } = call;
+
+    if (skill.requiresNetwork) {
+      const online = await this.checkConnectivity();
+      if (!online) {
+        return {
+          error: 'No internet connection. This skill requires network access.',
+        };
+      }
+    }
+
+    if (skill.type === 'native' && skill.execute) {
+      try {
+        return await withTimeout(
+          skill.execute(parameters),
+          this.config.skillTimeout,
+        );
+      } catch (err) {
+        return {
+          error:
+            err instanceof Error ? err.message : 'Native skill failed',
+        };
+      }
+    }
+
+    if (skill.type === 'js' && skill.html) {
+      if (!this.executor) {
+        return {
+          error: 'No skill executor available. SkillSandbox not mounted.',
+        };
+      }
+      try {
+        return await this.executor(
+          skill.html,
+          parameters,
+          this.config.skillTimeout,
+        );
+      } catch (err) {
+        return {
+          error:
+            err instanceof Error
+              ? err.message
+              : 'JS skill execution failed',
+        };
+      }
+    }
+
+    return {
+      error: `Cannot execute skill "${call.name}" — unsupported type "${skill.type}"`,
+    };
+  }
+}
+
+function withTimeout<T>(promise: Promise<T>, ms: number): Promise<T> {
+  return Promise.race([
+    promise,
+    new Promise<never>((_, reject) =>
+      setTimeout(() => reject(new Error(`Skill timed out after ${ms}ms`)), ms),
+    ),
+  ]);
+}

package/src/BM25Scorer.ts

@@ -0,0 +1,118 @@
+import type { SkillManifest } from './types';
+
+/**
+ * BM25 (Best Matching 25) scorer for skill pre-filtering.
+ *
+ * Ranks skills by relevance to a user query using term frequency / inverse
+ * document frequency with length normalization. Pure math, no ML model.
+ *
+ * Each skill's "document" is: name + description + parameter descriptions.
+ */
+
+const K1 = 1.5;
+const B = 0.75;
+
+type DocEntry = {
+  skill: SkillManifest;
+  tokens: string[];
+  termFreqs: Map<string, number>;
+};
+
+export class BM25Scorer {
+  private docs: DocEntry[] = [];
+  private avgDl = 0;
+  /** Number of documents containing each term */
+  private df: Map<string, number> = new Map();
+
+  /**
+   * Build the index from a set of skill manifests.
+   * Call this once when skills are registered (or change).
+   */
+  buildIndex(skills: SkillManifest[]): void {
+    this.docs = [];
+    this.df = new Map();
+
+    for (const skill of skills) {
+      const text = this.skillToText(skill);
+      const tokens = this.tokenize(text);
+
+      const termFreqs = new Map<string, number>();
+      const seen = new Set<string>();
+
+      for (const t of tokens) {
+        termFreqs.set(t, (termFreqs.get(t) ?? 0) + 1);
+        if (!seen.has(t)) {
+          seen.add(t);
+          this.df.set(t, (this.df.get(t) ?? 0) + 1);
+        }
+      }
+
+      this.docs.push({ skill, tokens, termFreqs });
+    }
+
+    const totalTokens = this.docs.reduce((sum, d) => sum + d.tokens.length, 0);
+    this.avgDl = this.docs.length > 0 ? totalTokens / this.docs.length : 0;
+  }
+
+  /**
+   * Score all indexed skills against a query. Returns skills ranked by
+   * descending BM25 score.
+   */
+  score(query: string): Array<{ skill: SkillManifest; score: number }> {
+    const queryTokens = this.tokenize(query);
+    const n = this.docs.length;
+
+    const results = this.docs.map((doc) => {
+      let total = 0;
+      const dl = doc.tokens.length;
+
+      for (const qt of queryTokens) {
+        const tf = doc.termFreqs.get(qt) ?? 0;
+        if (tf === 0) continue;
+
+        const docFreq = this.df.get(qt) ?? 0;
+        // IDF with floor at 0 to avoid negative scores
+        const idf = Math.max(
+          0,
+          Math.log((n - docFreq + 0.5) / (docFreq + 0.5) + 1),
+        );
+        const tfNorm =
+          (tf * (K1 + 1)) / (tf + K1 * (1 - B + B * (dl / this.avgDl)));
+
+        total += idf * tfNorm;
+      }
+
+      return { skill: doc.skill, score: total };
+    });
+
+    return results.sort((a, b) => b.score - a.score);
+  }
+
+  /**
+   * Convenience: score and return only the top-N skills.
+   */
+  topN(
+    query: string,
+    n: number,
+  ): Array<{ skill: SkillManifest; score: number }> {
+    return this.score(query).slice(0, n);
+  }
+
+  private skillToText(skill: SkillManifest): string {
+    const parts = [skill.name, skill.description];
+    for (const [key, param] of Object.entries(skill.parameters)) {
+      parts.push(key);
+      if (param.description) parts.push(param.description);
+    }
+    if (skill.instructions) parts.push(skill.instructions);
+    return parts.join(' ');
+  }
+
+  private tokenize(text: string): string[] {
+    return text
+      .toLowerCase()
+      .replace(/[^a-z0-9\s]/g, ' ')
+      .split(/\s+/)
+      .filter((t) => t.length > 1);
+  }
+}

package/src/FunctionCallParser.ts

@@ -0,0 +1,113 @@
+import type { ToolCall, SkillManifest } from './types';
+import type { SkillRegistry } from './SkillRegistry';
+
+export type ParsedToolCall = {
+  name: string;
+  parameters: Record<string, unknown>;
+  skill: SkillManifest;
+  /** Original tool call ID from llama.rn (needed for tool role messages) */
+  id?: string;
+};
+
+/**
+ * Primary path: validate tool_calls from llama.rn's native parser
+ * against registered skills. Returns only calls for known skills.
+ */
+export function validateToolCalls(
+  toolCalls: ToolCall[],
+  registry: SkillRegistry,
+): ParsedToolCall[] {
+  const validated: ParsedToolCall[] = [];
+
+  for (const tc of toolCalls) {
+    const skill = registry.getSkill(tc.function.name);
+    if (!skill) continue;
+
+    let parameters: Record<string, unknown>;
+    try {
+      parameters = JSON.parse(tc.function.arguments);
+    } catch {
+      parameters = {};
+    }
+
+    validated.push({
+      name: tc.function.name,
+      parameters,
+      skill,
+      id: tc.id,
+    });
+  }
+
+  return validated;
+}
+
+/**
+ * Fallback: scan raw text for JSON tool call blocks when llama.rn's
+ * native PEG parser misses them (e.g., malformed special tokens).
+ *
+ * Looks for two patterns:
+ *   1. {"tool_call": {"name": "...", "parameters": {...}}}
+ *   2. {"name": "...", "arguments": {...}}
+ */
+export function extractToolCallsFromText(
+  text: string,
+  registry: SkillRegistry,
+): ParsedToolCall[] {
+  const results: ParsedToolCall[] = [];
+
+  // Find JSON-like blocks in the text
+  const jsonBlocks = findJsonBlocks(text);
+
+  for (const block of jsonBlocks) {
+    try {
+      const parsed = JSON.parse(block);
+
+      let name: string | undefined;
+      let parameters: Record<string, unknown> = {};
+
+      if (parsed.tool_call && typeof parsed.tool_call === 'object') {
+        name = parsed.tool_call.name;
+        parameters = parsed.tool_call.parameters ?? {};
+      } else if (parsed.name && typeof parsed.name === 'string') {
+        name = parsed.name;
+        const raw = parsed.arguments ?? parsed.parameters ?? {};
+        parameters = typeof raw === 'string' ? JSON.parse(raw) : raw;
+      }
+
+      if (!name) continue;
+
+      const skill = registry.getSkill(name);
+      if (!skill) continue;
+
+      results.push({ name, parameters, skill });
+    } catch {
+      // Skip malformed JSON
+    }
+  }
+
+  return results;
+}
+
+/**
+ * Extract balanced JSON blocks from text by tracking brace depth.
+ */
+function findJsonBlocks(text: string): string[] {
+  const blocks: string[] = [];
+  let depth = 0;
+  let start = -1;
+
+  for (let i = 0; i < text.length; i++) {
+    if (text[i] === '{') {
+      if (depth === 0) start = i;
+      depth++;
+    } else if (text[i] === '}') {
+      depth--;
+      if (depth === 0 && start >= 0) {
+        blocks.push(text.slice(start, i + 1));
+        start = -1;
+      }
+    }
+  }
+
+  return blocks;
+}

package/src/GemmaAgentProvider.tsx

@@ -0,0 +1,101 @@
+import React, {
+  createContext,
+  useContext,
+  useRef,
+  useLayoutEffect,
+  useMemo,
+} from 'react';
+import { ModelManager } from './ModelManager';
+import { InferenceEngine } from './InferenceEngine';
+import { SkillRegistry } from './SkillRegistry';
+import { AgentOrchestrator } from './AgentOrchestrator';
+import { SkillSandbox, type SkillSandboxHandle } from './SkillSandbox';
+import type {
+  ModelConfig,
+  SkillManifest,
+  InferenceEngineConfig,
+  AgentConfig,
+} from './types';
+
+export type GemmaAgentContextValue = {
+  modelManager: ModelManager;
+  engine: InferenceEngine;
+  registry: SkillRegistry;
+  orchestrator: AgentOrchestrator;
+};
+
+const GemmaAgentContext = createContext<GemmaAgentContextValue | null>(null);
+
+export type GemmaAgentProviderProps = {
+  /** Model download config (repoId, filename, etc.) */
+  model: ModelConfig;
+  /** Skills to register on mount */
+  skills?: SkillManifest[];
+  /** Base system prompt for the agent */
+  systemPrompt?: string;
+  /** Inference engine configuration */
+  engineConfig?: InferenceEngineConfig;
+  /** Agent orchestrator configuration */
+  agentConfig?: AgentConfig;
+  children: React.ReactNode;
+};
+
+export function GemmaAgentProvider({
+  model,
+  skills,
+  systemPrompt,
+  engineConfig,
+  agentConfig,
+  children,
+}: GemmaAgentProviderProps) {
+  const sandboxRef = useRef<SkillSandboxHandle>(null);
+
+  // Create SDK instances once (stable across re-renders)
+  const instances = useRef<GemmaAgentContextValue | null>(null);
+  if (!instances.current) {
+    const modelManager = new ModelManager(model);
+    const engine = new InferenceEngine(engineConfig);
+    const registry = new SkillRegistry();
+
+    const orchestrator = new AgentOrchestrator(engine, registry, {
+      ...agentConfig,
+      systemPrompt: systemPrompt ?? agentConfig?.systemPrompt,
+    });
+
+    if (skills) {
+      for (const skill of skills) {
+        registry.registerSkill(skill);
+      }
+    }
+
+    instances.current = { modelManager, engine, registry, orchestrator };
+  }
+
+  // Wire SkillSandbox executor into orchestrator after mount
+  useLayoutEffect(() => {
+    if (sandboxRef.current) {
+      instances.current!.orchestrator.setSkillExecutor(
+        sandboxRef.current.execute,
+      );
+    }
+  }, []);
+
+  const value = useMemo(() => instances.current!, []);
+
+  return (
+    <GemmaAgentContext.Provider value={value}>
+      {children}
+      <SkillSandbox ref={sandboxRef} />
+    </GemmaAgentContext.Provider>
+  );
+}
+
+export function useGemmaAgentContext(): GemmaAgentContextValue {
+  const ctx = useContext(GemmaAgentContext);
+  if (!ctx) {
+    throw new Error(
+      'useGemmaAgent must be used within a <GemmaAgentProvider>',
+    );
+  }
+  return ctx;
+}