Add MathChatPlugin

Author: intitni

Core/Sources/ChatPlugins/MathChatPlugin/MathChatPlugin.swift

@@ -0,0 +1,63 @@
+import ChatPlugin
+import Environment
+import Foundation
+import OpenAIService
+
+/// Use Python to solve math problems.
+public actor MathChatPlugin: ChatPlugin {
+    public static var command: String { "math" }
+    public nonisolated var name: String { "Math" }
+
+    let chatGPTService: any ChatGPTServiceType
+    var isCancelled = false
+    weak var delegate: ChatPluginDelegate?
+
+    public init(inside chatGPTService: any ChatGPTServiceType, delegate: ChatPluginDelegate) {
+        self.chatGPTService = chatGPTService
+        self.delegate = delegate
+    }
+
+    public func send(content: String, originalMessage: String) async {
+        delegate?.pluginDidStart(self)
+        delegate?.pluginDidStartResponding(self)
+
+        let id = "\(Self.command)-\(UUID().uuidString)"
+        var reply = ChatMessage(id: id, role: .assistant, content: "Calculating...")
+
+        await chatGPTService.mutateHistory { history in
+            history.append(.init(role: .user, content: originalMessage, summary: content))
+            history.append(reply)
+        }
+
+        do {
+            let result = try await solveMathProblem(content)
+            await chatGPTService.mutateHistory { history in
+                if history.last?.id == id {
+                    history.removeLast()
+                }
+                reply.content = result
+                history.append(reply)
+            }
+        } catch {
+            await chatGPTService.mutateHistory { history in
+                if history.last?.id == id {
+                    history.removeLast()
+                }
+                reply.content = error.localizedDescription
+                history.append(reply)
+            }
+        }
+
+        delegate?.pluginDidEndResponding(self)
+        delegate?.pluginDidEnd(self)
+    }
+
+    public func cancel() async {
+        isCancelled = true
+    }
+
+    public func stopResponding() async {
+        isCancelled = true
+    }
+}
+

Core/Sources/ChatPlugins/MathChatPlugin/SolveMathProblem.swift

@@ -0,0 +1,34 @@
+import Foundation
+import LangChain
+import PythonKit
+
+func solveMathProblem(_ problem: String) async throws -> String {
+    #if DEBUG
+    let verbose = true
+    #else
+    let verbose = false
+    #endif
+
+    struct E: Error, LocalizedError {
+        var errorDescription: String? {
+            "Failed to parse answer."
+        }
+    }
+
+    let task = Task {
+        try withReadableThrowingPython {
+            let langchain = try Python.attemptImport("langchain")
+            let LLMMathChain = langchain.LLMMathChain
+            let llm = try LangChainChatModel.DynamicChatOpenAI(temperature: 0)
+            let llmMath = LLMMathChain.from_llm(llm, verbose: verbose)
+            let result = try llmMath.run.throwing.dynamicallyCall(withArguments: problem)
+            let answer = String(result)
+            if let answer { return answer }
+
+            throw E()
+        }
+    }
+
+    return try await task.value
+}
+

Core/Sources/ChatService/AllPlugins.swift

@@ -0,0 +1,8 @@
+import ChatPlugin
+import MathChatPlugin
+
+let allPlugins: [ChatPlugin.Type] = [
+    TerminalChatPlugin.self,
+    AITerminalChatPlugin.self,
+    MathChatPlugin.self,
+]

Core/Sources/ChatService/ChatPluginController.swift

@@ -7,8 +7,8 @@ final class ChatPluginController {
     let chatGPTService: any ChatGPTServiceType
     let plugins: [String: ChatPlugin.Type]
     var runningPlugin: ChatPlugin?
-
-    init(chatGPTService: any ChatGPTServiceType, plugins: ChatPlugin.Type...) {
+    
+    init(chatGPTService: any ChatGPTServiceType, plugins: [ChatPlugin.Type]) {
         self.chatGPTService = chatGPTService
         var all = [String: ChatPlugin.Type]()
         for plugin in plugins {
@@ -17,6 +17,10 @@ final class ChatPluginController {
         self.plugins = all
     }
 
+    convenience init(chatGPTService: any ChatGPTServiceType, plugins: ChatPlugin.Type...) {
+        self.init(chatGPTService: chatGPTService, plugins: plugins)
+    }
+
     /// Handle the message in a plugin if required. Return false if no plugin handles the message.
     func handleContent(_ content: String) async throws -> Bool {
         // look for the prefix of content, see if there is something like /command.

Core/Sources/ChatService/ChatService.swift

@@ -25,12 +25,7 @@ public final class ChatService: ObservableObject {
 
     public init<T: ChatGPTServiceType>(chatGPTService: T) {
         self.chatGPTService = chatGPTService
-        pluginController = ChatPluginController(
-            chatGPTService: chatGPTService,
-            plugins:
-            TerminalChatPlugin.self,
-            AITerminalChatPlugin.self
-        )
+        pluginController = ChatPluginController(chatGPTService: chatGPTService, plugins: allPlugins)
         contextController = DynamicContextController(
             chatGPTService: chatGPTService,
             contextCollectors: ActiveDocumentChatContextCollector()