WaveCut/Bonsai-Chat-WebGPU

1import type { ModelManifestV2, ModelTierId } from './manifest';2import type { BackendReport } from './native-log';3 4export type RuntimeBackend = 'webgpu' | 'wasm';5export type RequestedBackend = 'auto' | RuntimeBackend;6 7export interface StorageEstimate {8  usageBytes: number | null;9  quotaBytes: number | null;10  persisted: boolean;11}12 13export interface WebGpuAdapterSnapshot {14  available: boolean;15  name: string | null;16  vendor: string | null;17  architecture: string | null;18  device: string | null;19  description: string | null;20  limits: Record<string, number>;21  features: string[];22}23 24export interface EngineCapabilities {25  crossOriginIsolated: boolean;26  sharedArrayBuffer: boolean;27  hardwareConcurrency: number;28  webgpu: WebGpuAdapterSnapshot;29  storage: StorageEstimate;30  browser: {31    userAgent: string;32    platform: string;33    language: string;34  };35  runtime: {36    implementation: 'bonsai-wllama';37    wllamaRevision: string;38    llamaCppRevision: string;39    patchSetSha256: string;40    moduleSha256: string;41    wasmFlavor: 'jspi' | 'compat';42    wasmSha256: string;43    compatWorkerSha256: string | null;44    tokenEmbeddingOnWebGPU: true;45    tensorPlacementOverrides: false;46  };47}48 49export interface GateDecision {50  allowed: boolean;51  requestedBackend: RequestedBackend;52  selectedBackend: RuntimeBackend | null;53  reasons: string[];54  warnings: string[];55  requiredBytes: number;56  availableStorageBytes: number | null;57}58 59export interface FutureTensorPlacement {60  engine: 'future-custom-build';61  overrides: Array<{62    tensor: string;63    backend: RuntimeBackend;64  }>;65}66 67export type BenchmarkFlashMode = 'off' | 'auto';68export type BenchmarkKvCacheType = 'f16' | 'q8_0' | 'q4_0';69export type BenchmarkWasmFlavor = 'auto' | 'jspi' | 'compat';70 71export interface BenchmarkLoadTuning {72  nBatch?: number;73  nUbatch?: number;74  flashMode: BenchmarkFlashMode;75  kvCacheType: BenchmarkKvCacheType;76  wasmFlavor: BenchmarkWasmFlavor;77}78 79export interface EngineChatToolCall {80  id: string;81  type: 'function';82  function: {83    name: string;84    arguments: string;85  };86}87 88export type EngineChatMessage =89  | { role: 'system' | 'user'; content: string; name?: string }90  | { role: 'assistant'; content?: string | null; name?: string; tool_calls?: EngineChatToolCall[] }91  | { role: 'tool'; content: string; tool_call_id: string };92 93export interface EngineChatTool {94  type: 'function';95  function: {96    name: string;97    description?: string;98    parameters: {99      type: 'object';100      properties: Record<string, {101        type: string;102        description?: string;103        enum?: string[];104        [key: string]: unknown;105      }>;106      required?: string[];107      additionalProperties?: boolean;108    };109  };110}111 112type ManifestSource =113  | { manifestUrl: string; manifest?: never }114  | { manifest: ModelManifestV2; manifestUrl?: never };115 116export type LoadModelParams = ManifestSource & {117  modelId: ModelTierId;118  backend: RequestedBackend;119  contextSize?: number;120  warmupSystemPrompt?: string;121  threads?: number;122  benchmarkTuning?: BenchmarkLoadTuning;123  tensorPlacement?: FutureTensorPlacement;124};125 126export interface GenerateParams {127  messages: EngineChatMessage[];128  maxTokens?: number;129  temperature?: number;130  topP?: number;131  topK?: number;132  minP?: number;133  seed?: number;134  tools?: EngineChatTool[];135  toolChoice?: 'none' | 'auto' | 'required';136  cachePrompt?: boolean;137  returnTokenIds?: boolean;138}139 140export interface LoadModelResult {141  modelId: ModelTierId;142  backend: RuntimeBackend;143  gate: GateDecision;144  shardUrls: string[];145  context: {146    size: number;147    trainingSize: number;148    layerCount: number;149    vocabularySize: number;150    batchSize: number;151    microBatchSize: number;152  };153  tuning: {154    scope: 'release-defaults' | 'benchmark';155    requested: {156      nBatch: number | null;157      nUbatch: number | null;158      flashMode: BenchmarkFlashMode;159      cacheTypeK: BenchmarkKvCacheType;160      cacheTypeV: BenchmarkKvCacheType;161      wasmFlavor: BenchmarkWasmFlavor;162    };163    observed: {164      nBatch: number;165      nUbatch: number;166      flashAttention: boolean | null;167      cacheTypeK: string | null;168      cacheTypeV: string | null;169      kvBufferBytes: number | null;170      wasmFlavor: 'jspi' | 'compat';171      wasmSha256: string;172      compatWorkerSha256: string | null;173    };174    applied: boolean;175  };176  chatTemplate: {177    bytes: number;178    hasThinkMarker: boolean;179    hasToolCallMarker: boolean;180    hasToolResponseMarker: boolean;181  };182  backendReport: BackendReport;183}184 185export interface EngineTokenLogprob {186  id: number;187  logprob: number;188}189 190export interface EngineSampledTokenTraceEntry {191  selected: EngineTokenLogprob;192  topCandidates: EngineTokenLogprob[];193}194 195export interface GenerateResult {196  text: string;197  reasoningText: string;198  toolCallError?: string;199  tokenIds: number[] | null;200  tokenTrace: EngineSampledTokenTraceEntry[] | null;201  tokenTraceAccounting?: {202    usageCompletionTokens: number | null;203    tracedTokens: number;204    delta: number | null;205  } | null;206  finishReason: 'stop' | 'length' | 'tool_calls' | 'content_filter' | null;207  toolCalls: EngineChatToolCall[];208  usage: {209    promptTokens: number;210    completionTokens: number;211    totalTokens: number;212  } | null;213  timings: {214    promptTokensPerSecond: number;215    predictedTokensPerSecond: number;216  } | null;217}218 219export interface ScoreSequenceParams {220  promptTokenIds: number[];221  referenceTokenIds: number[];222  topK: 5;223}224 225export interface EngineTeacherForcedScoreEntry {226  index: number;227  selectedReference: EngineTokenLogprob;228  naturalTop1: EngineTokenLogprob;229  topCandidates: EngineTokenLogprob[];230  referenceRankInTopCandidatesZeroBased: number | null;231  top1Top2Margin: number;232}233 234export interface ScoreSequenceResult {235  method: {236    promptMode: 'raw-token-id-prefix';237    maxTokensPerStep: 1;238    temperature: 0;239    topK: 1;240    reportedTopLogprobs: 5;241    logitBias: 1_000;242    cachePromptFirst: false;243    cachePromptSubsequent: true;244  };245  entries: EngineTeacherForcedScoreEntry[];246  summary: {247    tokenCount: 1_024;248    meanNll: number;249    perplexity: number;250  };251}252 253export interface EmptyParams {254  readonly _empty?: never;255}256 257export interface EngineRequestMap {258  capabilities: EmptyParams;259  loadModel: LoadModelParams;260  generate: GenerateParams;261  scoreSequence: ScoreSequenceParams;262  abort: { targetRequestId: string };263  unload: EmptyParams;264  backendReport: EmptyParams;265  storageEstimate: EmptyParams;266  storagePersist: EmptyParams;267  storageClear: EmptyParams;268}269 270export interface EngineResponseMap {271  capabilities: EngineCapabilities;272  loadModel: LoadModelResult;273  generate: GenerateResult;274  scoreSequence: ScoreSequenceResult;275  abort: { targetRequestId: string; aborted: boolean };276  unload: { unloaded: true };277  backendReport: BackendReport;278  storageEstimate: StorageEstimate;279  storagePersist: StorageEstimate & { granted: boolean };280  storageClear: StorageEstimate & { cleared: true };281}282 283export type EngineMethod = keyof EngineRequestMap;284 285export type EngineRequest = {286  [Method in EngineMethod]: {287    type: 'request';288    requestId: string;289    method: Method;290    params: EngineRequestMap[Method];291  };292}[EngineMethod];293 294export interface EngineErrorPayload {295  code: string;296  message: string;297  details?: unknown;298}299 300export type ShardDownloadFailureKind = 'network' | 'abort' | 'verification';301 302export interface ShardDownloadFailureDetails {303  kind: 'shard-download-failure';304  failure: ShardDownloadFailureKind;305  causeCode: string;306  shardIndex: number;307  shardCount: number;308  shardPath: string;309  retryFromByteZero: true;310  partialDeleted: boolean;311}312 313export function isShardDownloadFailureDetails(314  value: unknown,315): value is ShardDownloadFailureDetails {316  if (typeof value !== 'object' || value === null) return false;317  const details = value as Partial<ShardDownloadFailureDetails>;318  return details.kind === 'shard-download-failure'319    && (details.failure === 'network' || details.failure === 'abort' || details.failure === 'verification')320    && typeof details.causeCode === 'string'321    && Number.isSafeInteger(details.shardIndex)322    && (details.shardIndex ?? -1) >= 0323    && Number.isSafeInteger(details.shardCount)324    && (details.shardCount ?? 0) > 0325    && (details.shardIndex ?? 0) < (details.shardCount ?? 0)326    && typeof details.shardPath === 'string'327    && details.shardPath.length > 0328    && details.retryFromByteZero === true329    && typeof details.partialDeleted === 'boolean';330}331 332export type EngineResponse = {333  [Method in EngineMethod]:334    | {335      type: 'response';336      requestId: string;337      method: Method;338      ok: true;339      result: EngineResponseMap[Method];340    }341    | {342      type: 'response';343      requestId: string;344      method: Method;345      ok: false;346      error: EngineErrorPayload;347    };348}[EngineMethod];349 350export type EngineShardProgressState =351  | 'queued'352  | 'cached'353  | 'downloading'354  | 'complete'355  | 'error';356 357export interface EngineShardProgress {358  index: number;359  path: string;360  loadedBytes: number;361  totalBytes: number;362  state: EngineShardProgressState;363}364 365export type EngineEvent =366  | {367    type: 'event';368    requestId: string;369    event: 'progress';370    phase: 'manifest' | 'cache' | 'download' | 'load' | 'warmup';371    loadedBytes: number;372    totalBytes: number;373    shardIndex: number | null;374    shardCount: number;375    shardPath: string | null;376    stageProgress: number | null;377    nativeStage: string | null;378    residentBytes: number | null;379    shards: EngineShardProgress[];380  }381  | {382    type: 'event';383    requestId: string;384    event: 'token';385    text: string;386    reasoningDelta?: string;387  }388  | {389    type: 'event';390    requestId: string;391    event: 'tool-call';392    index: number;393    id: string;394    name: string;395    argumentCharacters: number;396    argumentDelta: string;397  }398  | {399    type: 'event';400    requestId: string;401    event: 'generation';402    phase: 'preparing' | 'prefill' | 'decode';403    promptProcessed: number;404    promptTotal: number;405    promptCached: number;406    completionTokens: number;407    elapsedMs: number;408    promptTokensPerSecond: number;409    decodeTokensPerSecond: number;410  };411 412export type EngineWorkerMessage = EngineResponse | EngineEvent;413