Better state handelling

StanfordBDHG · Feb 13, 2024 · b5d6ccb · b5d6ccb
1 parent 12a59ea
commit b5d6ccb
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 45 deletions.
diff --git a/LLMonFHIR/FHIR Display/MultipleResourcesChatView.swift b/LLMonFHIR/FHIR Display/MultipleResourcesChatView.swift
@@ -32,10 +32,13 @@ struct MultipleResourcesChatView: View {
                     let contextBinding = Binding { llm.context } set: { llm.context = $0 }
                     ChatView(
                         contextBinding,
-                        disableInput: multipleResourceInterpreter.viewState == .processing
+                        disableInput: llm.state.representation == .processing
                     )
+                        .viewStateAlert(state: llm.state)
                         .onChange(of: llm.context) {
-                            multipleResourceInterpreter.queryLLM()
+                            if llm.state == .ready {
+                                multipleResourceInterpreter.queryLLM()
+                            }
                         }
                 } else {
                     ChatView(
@@ -47,7 +50,6 @@ struct MultipleResourcesChatView: View {
                 .toolbar {
                     toolbar
                 }
-                .viewStateAlert(state: $multipleResourceInterpreter.viewState)
                 .onAppear {
                     multipleResourceInterpreter.queryLLM()
                 }
@@ -58,7 +60,7 @@ struct MultipleResourcesChatView: View {
 
     @MainActor @ToolbarContentBuilder private var toolbar: some ToolbarContent {
         ToolbarItem(placement: .cancellationAction) {
-            if multipleResourceInterpreter.viewState == .processing {
+            if multipleResourceInterpreter.llm?.state.representation == .processing {
                 ProgressView()
             } else {
                 Button("Close") {
@@ -92,7 +94,7 @@ struct MultipleResourcesChatView: View {
                         .accessibilityLabel(Text("Reset Chat"))
                 }
             )
-            .disabled(multipleResourceInterpreter.viewState == .processing)
+            .disabled(multipleResourceInterpreter.llm?.state.representation == .processing)
         }
     }
 }
diff --git a/LLMonFHIR/FHIR Interpretation/FHIRInterpretationModule.swift b/LLMonFHIR/FHIR Interpretation/FHIRInterpretationModule.swift
@@ -59,6 +59,7 @@ class FHIRInterpretationModule: Module {
                     systemPrompt: nil   // No system prompt as this will be determined later by the resource interpreter
                 )
             ) {
+                // FHIR interpretation function
                 FHIRInterpretationFunction(
                     fhirStore: self.fhirStore,
                     resourceSummary: self.resourceSummary,

diff --git a/LLMonFHIR/FHIR Interpretation/FHIRMultipleResourceInterpreter.swift b/LLMonFHIR/FHIR Interpretation/FHIRMultipleResourceInterpreter.swift
@@ -6,6 +6,7 @@
 // SPDX-License-Identifier: MIT
 //
 
+import os
 import Spezi
 import SpeziFHIR
 import SpeziFHIRInterpretation
@@ -23,13 +24,14 @@ private enum FHIRMultipleResourceInterpreterConstants {
 
 @Observable
 class FHIRMultipleResourceInterpreter {
+    static let logger = Logger(subsystem: "edu.stanford.bdhg", category: "LLMonFHIR")
+
     private let localStorage: LocalStorage
     private let llmRunner: LLMRunner
     private let llmSchema: any LLMSchema
     private let fhirStore: FHIRStore
     private let resourceSummary: FHIRResourceSummary
 
-    @MainActor var viewState: ViewState = .idle
     var llm: (any LLMSession)?
 
 
@@ -56,50 +58,38 @@ class FHIRMultipleResourceInterpreter {
 
     @MainActor
     func queryLLM() {
-        guard viewState == .idle, llm?.context.last?.role == .user || !(llm?.context.contains(where: { $0.role == .assistant }) ?? false) else {
+        guard llm?.context.last?.role == .user || !(llm?.context.contains(where: { $0.role == .assistant }) ?? false) else {
             return
         }
 
         Task {
-            do {
-                var llm: LLMSession
-                if let llmTemp = self.llm {
-                    llm = llmTemp
-                } else {
-                    llm = await llmRunner(with: llmSchema)
-                    self.llm = llm
-                }
-
-                viewState = .processing
-
-                if llm.context.isEmpty {
-                    llm.context.append(systemMessage: FHIRPrompt.interpretMultipleResources.prompt)
-                }
-
-                if let patient = fhirStore.patient {
-                    llm.context.append(systemMessage: patient.jsonDescription)
-                }
-
-                print("The Multiple Resource Interpreter has access to \(fhirStore.llmRelevantResources.count) resources.")
-
-                do {
-                    let stream = try await llm.generate()
-
-                    for try await token in stream {
-                        llm.context.append(assistantOutput: token)
-                    }
-                } catch let error as LLMError {
-                    llm.state = .error(error: error)
-                } catch {
-                    llm.state = .error(error: LLMDefaultError.unknown(error))
-                }
-
-                try localStorage.store(llm.context, storageKey: FHIRMultipleResourceInterpreterConstants.chat)
-
-                viewState = .idle
-            } catch {
-                viewState = .error(error.localizedDescription)
+            var llm: LLMSession
+            if let llmTemp = self.llm {
+                llm = llmTemp
+            } else {
+                llm = await llmRunner(with: llmSchema)
+                self.llm = llm
+            }
+
+            if llm.context.isEmpty {
+                llm.context.append(systemMessage: FHIRPrompt.interpretMultipleResources.prompt)
+            }
+
+            if let patient = fhirStore.patient {
+                llm.context.append(systemMessage: patient.jsonDescription)
+            }
+
+            Self.logger.debug("The Multiple Resource Interpreter has access to \(self.fhirStore.llmRelevantResources.count) resources.")
+
+            guard let stream = try? await llm.generate() else {
+                return
+            }
+
+            for try await token in stream {
+                llm.context.append(assistantOutput: token)
             }
+
+            try localStorage.store(llm.context, storageKey: FHIRMultipleResourceInterpreterConstants.chat)
         }
     }
 }