Building an On-Device Q&A Agent with Chrome Built-in AI and the Prompt API

export type AIAvailabilityStatus =
  | 'readily'
  | 'available'
  | 'after-download'
  | 'downloadable'
  | 'downloading'
  | 'unavailable'
  | 'no'
  | 'unsupported';

export function getLanguageModelAPI() {
  if (typeof window === 'undefined') return null;

  const anyWin = window as unknown as {
    ai?: { languageModel?: any; assistant?: any };
    LanguageModel?: any;
  };

  return (
    anyWin.ai?.languageModel ||
    anyWin.ai?.assistant ||
    anyWin.LanguageModel ||
    null
  );
}

export async function checkAIAvailability(): Promise<AIAvailabilityStatus> {
  const api = getLanguageModelAPI();
  if (!api) return 'unsupported';

  if (typeof api.availability === 'function') {
    const status = await api.availability();
    return status || 'unsupported';
  }

  if (typeof api.capabilities === 'function') {
    const caps = await api.capabilities();
    return caps?.available || 'unsupported';
  }

  return 'unsupported';
}

const api = getLanguageModelAPI();

const session = await api.create({
  systemPrompt: 'You are a technical assistant specializing in software engineering telemetry.',
  temperature: 0.1,
  topK: 1,
  signal: abortController.signal,
  monitor: (monitorTarget: EventTarget) => {
    monitorTarget.addEventListener('downloadprogress', (event: Event) => {
      const prog = event as ProgressEvent;
      if (prog.total > 0) {
        const percentage = Math.round((prog.loaded / prog.total) * 100);
        console.log(`Downloading model weights: ${percentage}%`);
      }
    });
  },
});

export function buildGroundedPrompt(
  facts: string,
  userQuery: string,
  history: { role: string; content: string }[] = []
): string {
  let conversationBlock = '';
  if (history.length > 0) {
    const recent = history.slice(-3);
    conversationBlock = `RECENT CONVERSATION:\n${recent
      .map((m) => `${m.role === 'user' ? 'User' : 'Assistant'}: ${m.content}`)
      .join('\n')}\n\n`;
  }

  return `You are a verified technical assistant. Answer the user question accurately, concisely, and factually using ONLY the verified facts below. If the answer is not contained in the facts, explicitly state that the information is unavailable.

VERIFIED FACTS:
${facts}

${conversationBlock}USER QUESTION: ${userQuery}

CONCISE FACTUAL ANSWER:`;
}

export async function executePromptStreaming({
  session,
  prompt,
  onChunk,
  signal,
}: {
  session: any;
  prompt: string;
  onChunk: (accumulated: string, delta: string) => void;
  signal?: AbortSignal;
}): Promise<string> {
  if (typeof session.promptStreaming === 'function') {
    const streamResult = session.promptStreaming(prompt, { signal });
    let accumulated = '';

    // Case 1: AsyncIterable
    if (
      streamResult &&
      typeof streamResult[Symbol.asyncIterator] === 'function'
    ) {
      for await (const chunk of streamResult) {
        if (signal?.aborted) break;
        accumulated += chunk;
        onChunk(accumulated, chunk);
      }
      return accumulated;
    }

    // Case 2: ReadableStream
    if (
      streamResult &&
      typeof streamResult.getReader === 'function'
    ) {
      const reader = streamResult.getReader();
      try {
        while (true) {
          if (signal?.aborted) break;
          const { done, value } = await reader.read();
          if (done) break;
          if (value) {
            accumulated += value;
            onChunk(accumulated, value);
          }
        }
      } finally {
        reader.releaseLock();
      }
      return accumulated;
    }
  }

  // Fallback: Non-streaming execution
  const singleShotResult = await session.prompt(prompt, { signal });
  onChunk(singleShotResult, singleShotResult);
  return singleShotResult;
}

try {
  const result = await executePromptStreaming({ session, prompt, onChunk, signal });
  return result;
} finally {
  try {
    session.destroy();
  } catch (err) {
    // Suppress teardown errors
  }
}

import React, { useState, useRef, useEffect } from 'react';
import {
  checkAIAvailability,
  getLanguageModelAPI,
  buildGroundedPrompt,
  executePromptStreaming,
} from './aiAssistant';

export const OnDeviceChat: React.FC<{ datasetContext: string }> = ({ datasetContext }) => {
  const [messages, setMessages] = useState<{ role: string; content: string }[]>([]);
  const [input, setInput] = useState('');
  const [isGenerating, setIsGenerating] = useState(false);
  const [downloadProgress, setDownloadProgress] = useState<number | null>(null);
  const abortControllerRef = useRef<AbortController | null>(null);

  const handleSubmit = async (e: React.FormEvent) => {
    e.preventDefault();
    if (!input.trim() || isGenerating) return;

    const userText = input.trim();
    setInput('');
    setIsGenerating(true);

    const abortController = new AbortController();
    abortControllerRef.current = abortController;

    const newMessages = [
      ...messages,
      { role: 'user', content: userText },
      { role: 'assistant', content: '' },
    ];
    setMessages(newMessages);

    try {
      const api = getLanguageModelAPI();
      if (!api) throw new Error('Prompt API not supported on this browser.');

      const session = await api.create({
        systemPrompt: 'You are a factual Q&A assistant.',
        temperature: 0.1,
        topK: 1,
        signal: abortController.signal,
        monitor: (m: EventTarget) => {
          m.addEventListener('downloadprogress', (e: any) => {
            if (e.total) setDownloadProgress(Math.round((e.loaded / e.total) * 100));
          });
        },
      });

      try {
        const fullPrompt = buildGroundedPrompt(datasetContext, userText, messages);

        await executePromptStreaming({
          session,
          prompt: fullPrompt,
          signal: abortController.signal,
          onChunk: (accumulated) => {
            setMessages((prev) => {
              const updated = [...prev];
              updated[updated.length - 1] = { role: 'assistant', content: accumulated };
              return updated;
            });
          },
        });
      } finally {
        session.destroy();
      }
    } catch (err: any) {
      if (!abortController.signal.aborted) {
        console.error('Inference error:', err);
      }
    } finally {
      setIsGenerating(false);
      setDownloadProgress(null);
    }
  };

  const handleAbort = () => {
    abortControllerRef.current?.abort();
    setIsGenerating(false);
  };

  return (
    <div>
      {downloadProgress !== null && (
        <div>Downloading on-device weights: {downloadProgress}%</div>
      )}

      <div>
        {messages.map((m, idx) => (
          <div key={idx} className={m.role}>
            <strong>{m.role}:</strong> {m.content}
          </div>
        ))}
      </div>

      <form onSubmit={handleSubmit}>
        <input
          value={input}
          onChange={(e) => setInput(e.target.value)}
          placeholder="Ask a question..."
          disabled={isGenerating}
        />
        {isGenerating ? (
          <button type="button" onClick={handleAbort}>Stop</button>
        ) : (
          <button type="submit">Submit</button>
        )}
      </form>
    </div>
  );
};