Fix native interface for LLM. Also small functionality extension for tool calling.

Author: pweglik

android/src/main/java/com/swmansion/rnexecutorch/LLM.kt

@@ -7,7 +7,7 @@ import org.pytorch.executorch.extension.llm.LlmCallback
 import org.pytorch.executorch.extension.llm.LlmModule
 
 class LLM(
-        reactContext: ReactApplicationContext,
+  reactContext: ReactApplicationContext,
 ) : NativeLLMSpec(reactContext), LlmCallback {
   private var llmModule: LlmModule? = null
 
@@ -26,9 +26,9 @@ class LLM(
   }
 
   override fun loadLLM(
-          modelSource: String,
-          tokenizerSource: String,
-          promise: Promise,
+    modelSource: String,
+    tokenizerSource: String,
+    promise: Promise,
   ) {
     try {
       llmModule = LlmModule(modelSource, tokenizerSource, 0.7f)
@@ -39,21 +39,21 @@ class LLM(
   }
 
   override fun runInference(
-          input: String,
-          promise: Promise,
+    input: String,
+    promise: Promise,
   ) {
     Thread {
-              llmModule!!.generate(input, this)
-              promise.resolve("Inference completed successfully")
-            }
-            .start()
+      llmModule!!.generate(input, this)
+      promise.resolve("Inference completed successfully")
+    }
+      .start()
   }
 
   override fun interrupt() {
     llmModule!!.stop()
   }
 
-  override fun deleteModule() {
+  override fun releaseResources() {
     llmModule = null
   }
 

src/controllers/LLMController.ts

@@ -138,7 +138,7 @@ export class LLMController {
   public delete() {
     this.onToken?.remove();
     this.onToken = null;
-    this.nativeModule.deleteModule();
+    this.nativeModule.releaseResources();
   }
 
   public async runInference(input: string) {
@@ -184,7 +184,7 @@ export class LLMController {
       return;
     }
 
-    if (!this.toolsConfig) {
+    if (!this.toolsConfig || this.toolsConfig.displayToolCalls) {
       this.responseCallback(
         this._response.replace(this.tokenizerConfig.eos_token, '')
       );

src/native/NativeLLM.ts

@@ -6,7 +6,7 @@ export interface Spec extends TurboModule {
   loadLLM(modelSource: string, tokenizerSource: string): Promise<string>;
   runInference(input: string): Promise<string>;
   interrupt(): void;
-  deleteModule(): void;
+  releaseResources(): void;
 
   readonly onToken: EventEmitter<string>;
 }

src/types/llm.ts

@@ -6,7 +6,7 @@ export interface LLMType {
   downloadProgress: number;
   error: string | null;
   runInference: (input: string) => Promise<void>;
-  sendMessage: (message: string, tools?: LLMTool[]) => Promise<void>;
+  sendMessage: (message: string) => Promise<void>;
   interrupt: () => void;
 }
 
@@ -35,6 +35,7 @@ export interface ChatConfig {
 export interface ToolsConfig {
   tools: LLMTool[];
   executeToolCallback: (call: ToolCall) => Promise<string | null>;
+  displayToolCalls?: boolean;
 }
 
 export const SPECIAL_TOKENS = [