cactus-react-native 1.0.1 → 1.1.0

package/src/hooks/useCactusSTT.ts

@@ -0,0 +1,285 @@
+import { useCallback, useEffect, useState, useRef } from 'react';
+import { CactusSTT } from '../classes/CactusSTT';
+import { CactusFileSystem } from '../native';
+import { getErrorMessage } from '../utils/error';
+import type {
+  CactusSTTParams,
+  CactusSTTTranscribeResult,
+  CactusSTTTranscribeParams,
+  CactusSTTDownloadParams,
+  CactusSTTAudioEmbedParams,
+  CactusSTTAudioEmbedResult,
+} from '../types/CactusSTT';
+import type { CactusModel } from '../types/CactusModel';
+
+export const useCactusSTT = ({
+  model = 'whisper-small',
+  contextSize = 2048,
+}: CactusSTTParams = {}) => {
+  const [cactusSTT, setCactusSTT] = useState(
+    () => new CactusSTT({ model, contextSize })
+  );
+
+  // State
+  const [response, setResponse] = useState('');
+  const [isGenerating, setIsGenerating] = useState(false);
+  const [isInitializing, setIsInitializing] = useState(false);
+  const [isDownloaded, setIsDownloaded] = useState(false);
+  const [isDownloading, setIsDownloading] = useState(false);
+  const [downloadProgress, setDownloadProgress] = useState(0);
+  const [error, setError] = useState<string | null>(null);
+
+  const currentModelRef = useRef(model);
+  const currentDownloadIdRef = useRef(0);
+
+  useEffect(() => {
+    currentModelRef.current = model;
+  }, [model]);
+
+  useEffect(() => {
+    setCactusSTT(new CactusSTT({ model, contextSize }));
+
+    setResponse('');
+    setIsGenerating(false);
+    setIsInitializing(false);
+    setIsDownloaded(false);
+    setIsDownloading(false);
+    setDownloadProgress(0);
+    setError(null);
+
+    let mounted = true;
+    CactusFileSystem.modelExists(model)
+      .then((exists) => {
+        if (!mounted) {
+          return;
+        }
+        setIsDownloaded(exists);
+      })
+      .catch((e) => {
+        if (!mounted) {
+          return;
+        }
+        setIsDownloaded(false);
+        setError(getErrorMessage(e));
+      });
+
+    return () => {
+      mounted = false;
+    };
+  }, [model, contextSize]);
+
+  useEffect(() => {
+    return () => {
+      cactusSTT.destroy().catch(() => {});
+    };
+  }, [cactusSTT]);
+
+  const download = useCallback(
+    async ({ onProgress }: CactusSTTDownloadParams = {}) => {
+      if (isDownloading) {
+        const message = 'CactusSTT is already downloading';
+        setError(message);
+        throw new Error(message);
+      }
+
+      setError(null);
+
+      if (isDownloaded) {
+        return;
+      }
+
+      const thisModel = currentModelRef.current;
+      const thisDownloadId = ++currentDownloadIdRef.current;
+
+      setDownloadProgress(0);
+      setIsDownloading(true);
+      try {
+        await cactusSTT.download({
+          onProgress: (progress) => {
+            if (
+              currentModelRef.current !== thisModel ||
+              currentDownloadIdRef.current !== thisDownloadId
+            ) {
+              return;
+            }
+
+            setDownloadProgress(progress);
+            onProgress?.(progress);
+          },
+        });
+
+        if (
+          currentModelRef.current !== thisModel ||
+          currentDownloadIdRef.current !== thisDownloadId
+        ) {
+          return;
+        }
+
+        setIsDownloaded(true);
+      } catch (e) {
+        if (
+          currentModelRef.current !== thisModel ||
+          currentDownloadIdRef.current !== thisDownloadId
+        ) {
+          return;
+        }
+
+        setError(getErrorMessage(e));
+        throw e;
+      } finally {
+        if (
+          currentModelRef.current !== thisModel ||
+          currentDownloadIdRef.current !== thisDownloadId
+        ) {
+          return;
+        }
+
+        setIsDownloading(false);
+        setDownloadProgress(0);
+      }
+    },
+    [cactusSTT, isDownloading, isDownloaded]
+  );
+
+  const init = useCallback(async () => {
+    if (isInitializing) {
+      const message = 'CactusSTT is already initializing';
+      setError(message);
+      throw new Error(message);
+    }
+
+    setError(null);
+    setIsInitializing(true);
+    try {
+      await cactusSTT.init();
+    } catch (e) {
+      setError(getErrorMessage(e));
+      throw e;
+    } finally {
+      setIsInitializing(false);
+    }
+  }, [cactusSTT, isInitializing]);
+
+  const transcribe = useCallback(
+    async ({
+      audioFilePath,
+      prompt,
+      options,
+      onToken,
+    }: CactusSTTTranscribeParams): Promise<CactusSTTTranscribeResult> => {
+      if (isGenerating) {
+        const message = 'CactusSTT is already generating';
+        setError(message);
+        throw new Error(message);
+      }
+
+      setError(null);
+      setResponse('');
+      setIsGenerating(true);
+      try {
+        return await cactusSTT.transcribe({
+          audioFilePath,
+          prompt,
+          options,
+          onToken: (token) => {
+            setResponse((prev) => prev + token);
+            onToken?.(token);
+          },
+        });
+      } catch (e) {
+        setError(getErrorMessage(e));
+        throw e;
+      } finally {
+        setIsGenerating(false);
+      }
+    },
+    [cactusSTT, isGenerating]
+  );
+
+  const audioEmbed = useCallback(
+    async ({
+      audioPath,
+    }: CactusSTTAudioEmbedParams): Promise<CactusSTTAudioEmbedResult> => {
+      if (isGenerating) {
+        const message = 'CactusSTT is already generating';
+        setError(message);
+        throw new Error(message);
+      }
+
+      setError(null);
+      setIsGenerating(true);
+      try {
+        return await cactusSTT.audioEmbed({ audioPath });
+      } catch (e) {
+        setError(getErrorMessage(e));
+        throw e;
+      } finally {
+        setIsGenerating(false);
+      }
+    },
+    [cactusSTT, isGenerating]
+  );
+
+  const stop = useCallback(async () => {
+    setError(null);
+    try {
+      await cactusSTT.stop();
+    } catch (e) {
+      setError(getErrorMessage(e));
+      throw e;
+    }
+  }, [cactusSTT]);
+
+  const reset = useCallback(async () => {
+    setError(null);
+    try {
+      await cactusSTT.reset();
+    } catch (e) {
+      setError(getErrorMessage(e));
+      throw e;
+    } finally {
+      setResponse('');
+    }
+  }, [cactusSTT]);
+
+  const destroy = useCallback(async () => {
+    setError(null);
+    try {
+      await cactusSTT.destroy();
+    } catch (e) {
+      setError(getErrorMessage(e));
+      throw e;
+    } finally {
+      setResponse('');
+    }
+  }, [cactusSTT]);
+
+  const getModels = useCallback(async (): Promise<CactusModel[]> => {
+    setError(null);
+    try {
+      return await cactusSTT.getModels();
+    } catch (e) {
+      setError(getErrorMessage(e));
+      throw e;
+    }
+  }, [cactusSTT]);
+
+  return {
+    response,
+    isGenerating,
+    isInitializing,
+    isDownloaded,
+    isDownloading,
+    downloadProgress,
+    error,
+
+    download,
+    init,
+    transcribe,
+    audioEmbed,
+    reset,
+    stop,
+    destroy,
+    getModels,
+  };
+};

package/src/index.tsx

@@ -1,8 +1,10 @@
 // Classes
 export { CactusLM } from './classes/CactusLM';
+export { CactusSTT } from './classes/CactusSTT';
 
 // Hooks
 export { useCactusLM } from './hooks/useCactusLM';
+export { useCactusSTT } from './hooks/useCactusSTT';
 
 // Types
 export type { CactusModel } from './types/CactusModel';
@@ -10,14 +12,24 @@ export type {
   CactusLMParams,
   CactusLMDownloadParams,
   Message,
-  Options,
+  CompleteOptions,
   Tool,
   CactusLMCompleteParams,
   CactusLMCompleteResult,
   CactusLMEmbedParams,
   CactusLMEmbedResult,
-  CactusLMGetModelsParams,
+  CactusLMImageEmbedParams,
+  CactusLMImageEmbedResult,
 } from './types/CactusLM';
+export type {
+  CactusSTTParams,
+  CactusSTTDownloadParams,
+  TranscribeOptions,
+  CactusSTTTranscribeParams,
+  CactusSTTTranscribeResult,
+  CactusSTTAudioEmbedParams,
+  CactusSTTAudioEmbedResult,
+} from './types/CactusSTT';
 
 // Config
 export { CactusConfig } from './config/CactusConfig';

package/src/native/Cactus.ts

@@ -1,28 +1,56 @@
 import { NitroModules } from 'react-native-nitro-modules';
 import type { Cactus as CactusSpec } from '../specs/Cactus.nitro';
+import { CactusImage } from './CactusImage';
 import type {
   CactusLMCompleteResult,
   Message,
-  Options,
+  CompleteOptions,
   Tool,
 } from '../types/CactusLM';
+import type {
+  CactusSTTTranscribeResult,
+  TranscribeOptions,
+} from '../types/CactusSTT';
 
 export class Cactus {
   private readonly hybridCactus =
     NitroModules.createHybridObject<CactusSpec>('Cactus');
 
-  public init(modelPath: string, contextSize: number): Promise<void> {
-    return this.hybridCactus.init(modelPath, contextSize);
+  public init(
+    modelPath: string,
+    contextSize: number,
+    corpusDir?: string
+  ): Promise<void> {
+    return this.hybridCactus.init(modelPath, contextSize, corpusDir);
   }
 
   public async complete(
     messages: Message[],
     responseBufferSize: number,
-    options?: Options,
-    tools?: Tool[],
+    options?: CompleteOptions,
+    tools?: { type: 'function'; function: Tool }[],
     callback?: (token: string, tokenId: number) => void
   ): Promise<CactusLMCompleteResult> {
-    const messagesJson = JSON.stringify(messages);
+    const messagesInternal: Message[] = [];
+    for (const message of messages) {
+      if (!message.images) {
+        messagesInternal.push(message);
+        continue;
+      }
+      const resizedImages: string[] = [];
+      for (const imagePath of message.images) {
+        const resizedImage = await CactusImage.resize(
+          imagePath.replace('file://', ''),
+          128,
+          128,
+          1
+        );
+        resizedImages.push(resizedImage);
+      }
+      messagesInternal.push({ ...message, images: resizedImages });
+    }
+
+    const messagesJson = JSON.stringify(messagesInternal);
     const optionsJson = options
       ? JSON.stringify({
           temperature: options.temperature,
@@ -61,10 +89,76 @@ export class Cactus {
     }
   }
 
+  public async transcribe(
+    audioFilePath: string,
+    prompt: string,
+    responseBufferSize: number,
+    options?: TranscribeOptions,
+    callback?: (token: string, tokenId: number) => void
+  ): Promise<CactusSTTTranscribeResult> {
+    const optionsJson = options
+      ? JSON.stringify({
+          temperature: options.temperature,
+          top_p: options.topP,
+          top_k: options.topK,
+          max_tokens: options.maxTokens,
+          stop_sequences: options.stopSequences,
+        })
+      : undefined;
+
+    const response = await this.hybridCactus.transcribe(
+      audioFilePath.replace('file://', ''),
+      prompt,
+      responseBufferSize,
+      optionsJson,
+      callback
+    );
+
+    try {
+      const parsed = JSON.parse(response);
+
+      return {
+        success: parsed.success,
+        response: parsed.response,
+        timeToFirstTokenMs: parsed.time_to_first_token_ms,
+        totalTimeMs: parsed.total_time_ms,
+        tokensPerSecond: parsed.tokens_per_second,
+        prefillTokens: parsed.prefill_tokens,
+        decodeTokens: parsed.decode_tokens,
+        totalTokens: parsed.total_tokens,
+      };
+    } catch {
+      throw new Error('Unable to parse transcription response');
+    }
+  }
+
   public embed(text: string, embeddingBufferSize: number): Promise<number[]> {
     return this.hybridCactus.embed(text, embeddingBufferSize);
   }
 
+  public async imageEmbed(
+    imagePath: string,
+    embeddingBufferSize: number
+  ): Promise<number[]> {
+    const resizedImage = await CactusImage.resize(
+      imagePath.replace('file://', ''),
+      128,
+      128,
+      1
+    );
+    return this.hybridCactus.imageEmbed(resizedImage, embeddingBufferSize);
+  }
+
+  public audioEmbed(
+    audioPath: string,
+    embeddingBufferSize: number
+  ): Promise<number[]> {
+    return this.hybridCactus.audioEmbed(
+      audioPath.replace('file://', ''),
+      embeddingBufferSize
+    );
+  }
+
   public reset(): Promise<void> {
     return this.hybridCactus.reset();
   }

package/src/native/CactusFileSystem.ts

@@ -35,10 +35,10 @@ export class CactusFileSystem {
 
   public static downloadModel(
     model: string,
+    url: string,
     onProgress?: (progress: number) => void
   ): Promise<void> {
-    const from = `https://vlqqczxwyaodtcdmdmlw.supabase.co/storage/v1/object/public/cactus-models/${model}.zip`;
-    return this.hybridCactusFileSystem.downloadModel(model, from, onProgress);
+    return this.hybridCactusFileSystem.downloadModel(model, url, onProgress);
   }
 
   public static deleteModel(model: string): Promise<void> {

package/src/native/CactusImage.ts

@@ -0,0 +1,20 @@
+import { NitroModules } from 'react-native-nitro-modules';
+import type { CactusImage as CactusImageSpec } from '../specs/CactusImage.nitro';
+
+export class CactusImage {
+  private static readonly hybridCactusImage =
+    NitroModules.createHybridObject<CactusImageSpec>('CactusImage');
+
+  public static base64(path: string): Promise<string> {
+    return this.hybridCactusImage.base64(path);
+  }
+
+  public static resize(
+    path: string,
+    height: number,
+    width: number,
+    quality: number
+  ): Promise<string> {
+    return this.hybridCactusImage.resize(path, height, width, quality);
+  }
+}

package/src/native/index.ts

@@ -2,4 +2,5 @@ export { Cactus } from './Cactus';
 export { CactusCrypto } from './CactusCrypto';
 export { CactusDeviceInfo } from './CactusDeviceInfo';
 export { CactusFileSystem } from './CactusFileSystem';
+export { CactusImage } from './CactusImage';
 export { CactusUtil } from './CactusUtil';

package/src/specs/Cactus.nitro.ts

@@ -1,7 +1,11 @@
 import type { HybridObject } from 'react-native-nitro-modules';
 
 export interface Cactus extends HybridObject<{ ios: 'c++'; android: 'c++' }> {
-  init(modelPath: string, contextSize: number): Promise<void>;
+  init(
+    modelPath: string,
+    contextSize: number,
+    corpusDir?: string
+  ): Promise<void>;
   complete(
     messagesJson: string,
     responseBufferSize: number,
@@ -9,7 +13,16 @@ export interface Cactus extends HybridObject<{ ios: 'c++'; android: 'c++' }> {
     toolsJson?: string,
     callback?: (token: string, tokenId: number) => void
   ): Promise<string>;
+  transcribe(
+    audioFilePath: string,
+    prompt: string,
+    responseBufferSize: number,
+    optionsJson?: string,
+    callback?: (token: string, tokenId: number) => void
+  ): Promise<string>;
   embed(text: string, embeddingBufferSize: number): Promise<number[]>;
+  imageEmbed(imagePath: string, embeddingBufferSize: number): Promise<number[]>;
+  audioEmbed(audioPath: string, embeddingBufferSize: number): Promise<number[]>;
   reset(): Promise<void>;
   stop(): Promise<void>;
   destroy(): Promise<void>;

package/src/specs/CactusImage.nitro.ts

@@ -0,0 +1,12 @@
+import type { HybridObject } from 'react-native-nitro-modules';
+
+export interface CactusImage
+  extends HybridObject<{ ios: 'swift'; android: 'kotlin' }> {
+  base64(path: string): Promise<string>;
+  resize(
+    path: string,
+    height: number,
+    width: number,
+    quality: number
+  ): Promise<string>;
+}

package/src/telemetry/Telemetry.ts

@@ -8,6 +8,7 @@ import {
 import { CactusConfig } from '../config/CactusConfig';
 import { packageVersion } from '../constants/packageVersion';
 import type { CactusLMCompleteResult } from '../types/CactusLM';
+import type { CactusSTTTranscribeResult } from '../types/CactusSTT';
 
 export interface LogRecord {
   // Framework
@@ -15,7 +16,13 @@ export interface LogRecord {
   framework_version: string;
 
   // Event
-  event_type: 'init' | 'completion' | 'embedding';
+  event_type:
+    | 'init'
+    | 'completion'
+    | 'transcription'
+    | 'embedding'
+    | 'image_embedding'
+    | 'audio_embedding';
   model: string;
   success: boolean;
   message?: string;
@@ -41,10 +48,17 @@ export class Telemetry {
   private static readonly logBufferPaths = {
     init: 'logs/init.json',
     completion: 'logs/completion.json',
+    transcription: 'logs/transcription.json',
     embedding: 'logs/embedding.json',
+    image_embedding: 'logs/image_embedding.json',
+    audio_embedding: 'logs/audio_embedding.json',
   };
 
   private static async handleLog(logRecord: LogRecord) {
+    if (!this.isInitialized()) {
+      return;
+    }
+
     if (!CactusConfig.isTelemetryEnabled) {
       return;
     }
@@ -78,6 +92,10 @@ export class Telemetry {
   }
 
   public static async init(cactusTelemetryToken?: string): Promise<void> {
+    if (this.isInitialized()) {
+      return;
+    }
+
     if (!CactusConfig.isTelemetryEnabled) {
       return;
     }
@@ -139,6 +157,29 @@ export class Telemetry {
     });
   }
 
+  public static logTranscribe(
+    model: string,
+    success: boolean,
+    message?: string,
+    result?: CactusSTTTranscribeResult
+  ): Promise<void> {
+    return this.handleLog({
+      framework: 'react-native',
+      framework_version: packageVersion,
+      event_type: 'transcription',
+      model,
+      success,
+      message,
+      telemetry_token: this.cactusTelemetryToken,
+      project_id: this.projectId,
+      device_id: this.deviceId,
+      tokens: result?.totalTokens,
+      response_time: result?.totalTimeMs,
+      ttft: result?.timeToFirstTokenMs,
+      tps: result?.tokensPerSecond,
+    });
+  }
+
   public static logEmbedding(
     model: string,
     success: boolean,
@@ -156,4 +197,40 @@ export class Telemetry {
       device_id: this.deviceId,
     });
   }
+
+  public static logImageEmbedding(
+    model: string,
+    success: boolean,
+    message?: string
+  ): Promise<void> {
+    return this.handleLog({
+      framework: 'react-native',
+      framework_version: packageVersion,
+      event_type: 'image_embedding',
+      model,
+      success,
+      message,
+      telemetry_token: this.cactusTelemetryToken,
+      project_id: this.projectId,
+      device_id: this.deviceId,
+    });
+  }
+
+  public static logAudioEmbedding(
+    model: string,
+    success: boolean,
+    message?: string
+  ): Promise<void> {
+    return this.handleLog({
+      framework: 'react-native',
+      framework_version: packageVersion,
+      event_type: 'audio_embedding',
+      model,
+      success,
+      message,
+      telemetry_token: this.cactusTelemetryToken,
+      project_id: this.projectId,
+      device_id: this.deviceId,
+    });
+  }
 }