Fixed marking of response generation completed control flow in LlmInf…

…erence+Session.
google-ai-edge · Aug 26, 2024 · a4c6a6a · a4c6a6a
1 parent 3d650fd
commit a4c6a6a
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/mediapipe/tasks/ios/genai/core/sources/LlmSessionRunner.swift b/mediapipe/tasks/ios/genai/core/sources/LlmSessionRunner.swift
@@ -88,7 +88,7 @@ final class LlmSessionRunner {
   func predictAsync(
     progress: @escaping (_ partialResult: [String]?, _ error: Error?) -> Void,
     completion: @escaping (() -> Void)
-  ) throws {
+  ) {
     let callbackInfo = CallbackInfo(progress: progress, completion: completion)
     let callbackContext = UnsafeMutableRawPointer(Unmanaged.passRetained(callbackInfo).toOpaque())
 

diff --git a/mediapipe/tasks/ios/genai/inference/sources/LlmInference+Session.swift b/mediapipe/tasks/ios/genai/inference/sources/LlmInference+Session.swift
@@ -118,11 +118,13 @@ extension LlmInference {
       /// TODO: If simultaneous response generations on multiple sessions or the same session
       /// are allowed to happen it leads to a crash. Investigate if this can be handled by C++.
       try llmInference.shouldContinueWithResponseGeneration()
-
+
+      defer {
+        llmInference.markResponseGenerationCompleted()
+      }
+
       let tokens = try llmSessionRunner.predict()
 
-      llmInference.markResponseGenerationCompleted()
-
       guard let humanReadableLlmResponse = Session.humanReadableString(llmResponses: tokens)
       else {
         throw GenAiInferenceError.invalidResponse
@@ -159,7 +161,7 @@ extension LlmInference {
       /// Used to make a decision about whitespace stripping.
       var receivedFirstToken = true
 
-      try llmSessionRunner.predictAsync(
+      llmSessionRunner.predictAsync(
         progress: { partialResponseStrings, error in
           guard let responseStrings = partialResponseStrings,
             let humanReadableLlmResponse = Session.humanReadableString(