feat: adds 'regenerateLast' feature to answer session

oramasearch · Jul 24, 2024 · 5fcfa73 · 5fcfa73
1 parent 58692e5
commit 5fcfa73
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 0 deletions.
diff --git a/Sources/oramacloud-client/answer-session.swift b/Sources/oramacloud-client/answer-session.swift
@@ -86,6 +86,12 @@ class AnswerSession<Doc: Codable> {
         let message: String
     }
 
+    enum MaybeStream {
+        case stream(AsyncThrowingStream<String, Error>)
+        case string(String)
+        case error(Error)
+    }
+
     private let eventEmitter = EventEmitter()
     private let endpointBaseURL = "https://answer.api.orama.com"
     private var abortController: Task<Void, Error>?
@@ -98,6 +104,7 @@ class AnswerSession<Doc: Codable> {
     private var messages: [AnswerParams<Doc>.Message]
     private var inferenceType: AnswerParams<Doc>.InferenceType
     private var state: [AnswerParams<Doc>.Interaction<Doc>] = []
+    private var lastInteractionParams: AnswerParams<Doc>.AskParams?
 
     init(params: AnswerParams<Doc>) {
         userContext = params.userContext ?? .string
@@ -142,9 +149,35 @@ class AnswerSession<Doc: Codable> {
         return response
     }
 
+    public func regenerateLast(stream: Bool = true) async throws -> MaybeStream {
+        if state.isEmpty || self.messages.isEmpty {
+            throw OramaClientError.runtimeError("No messages to regenerate")
+        }
+
+        let isLastMessageAssistant = self.messages.last?.role == .assistant
+
+        if !isLastMessageAssistant {
+            throw OramaClientError.runtimeError("Last message is not an assistant message")
+        }
+
+        if self.lastInteractionParams == nil {
+            throw OramaClientError.runtimeError("Cannot find last interaction params")
+        }
+
+        self.messages.removeLast()
+        self.state.removeLast()
+
+        if stream {
+            return .stream(try await fetchAnswer(params: self.lastInteractionParams!))
+        } else {
+            return .string(try await ask(params: self.lastInteractionParams!))
+        }
+    }
+
     private func fetchAnswer(params: AnswerParams<Doc>.AskParams) async throws -> AsyncThrowingStream<String, Error> {
         AsyncThrowingStream { continuation in
             let interactionId = Cuid.generateId()
+            self.lastInteractionParams = params
             self.abortController = Task {
                 do {
                     self.state.append(AnswerParams<Doc>.Interaction(

diff --git a/Sources/oramacloud-client/types.swift b/Sources/oramacloud-client/types.swift
@@ -77,3 +77,9 @@ struct SearchRequestPayload: Encodable {
         return jsonObject
     }
 }
+
+// ======================== ERRORS ========================
+
+enum OramaClientError: Error {
+    case runtimeError(String)
+}
diff --git a/Tests/oramacloud-clientTests/oramacloud_clientTests.swift b/Tests/oramacloud-clientTests/oramacloud_clientTests.swift
@@ -117,6 +117,31 @@ final class oramacloud_clientTests: XCTestCase {
         wait(for: [expectation], timeout: 60.0)
     }
 
+    func testE2ERegenerateLastAnswer() async throws {
+        let expectation = XCTestExpectation(description: "Can correctly regenerate the last answer")
+        var state: [AnswerParams<E2EDoc>.Interaction<E2EDoc>] = []
+
+        do {
+            _ = answerSession.on(event: .stateChange) {
+                state = $0 as! [AnswerParams<E2EDoc>.Interaction<E2EDoc>]
+            }
+
+            let _ = try await answerSession.ask(params: AnswerParams.AskParams(query: "german", userData: nil, related: nil))
+            let _ = try await answerSession.ask(params: AnswerParams.AskParams(query: "labrador", userData: nil, related: nil))
+            let _ = try await answerSession.regenerateLast(stream: false)
+
+            expectation.fulfill()
+            XCTAssertEqual(state.count, 2, "Should contain 2 interactions")
+            XCTAssertEqual(state.last!.query, "labrador", "Second query should be 'labrador'")
+
+            await fulfillment(of: [expectation], timeout: 120.0)
+
+        } catch {
+            XCTFail("Test failed with error: \(error)")
+            expectation.fulfill()
+        }
+    }
+
     func testE2EIndexManager() throws {
         struct DocumentStruct: Codable {
             let breed: String
@@ -181,5 +206,7 @@ extension oramacloud_clientTests {
         ("testE2EAnswerSession", testE2EAnswerSession),
         ("testOnMessageLoading", testOnMessageLoading),
         ("testAsyncE2EAnswerSession", testAsyncE2EAnswerSession),
+        ("testE2ERegenerateLastAnswer", testE2ERegenerateLastAnswer),
+        ("testE2EIndexManager", testE2EIndexManager)
     ]
 }