Add StructuredOutputChatModelChain

Author: intitni

Tool/Sources/LangChain/Agent.swift

@@ -96,8 +96,6 @@ public protocol Agent {
     associatedtype Output: AgentOutputParsable
     associatedtype ScratchPadContent: Equatable
     var chatModelChain: ChatModelChain<AgentInput<Input, ScratchPadContent>> { get }
-    var observationPrefix: String { get }
-    var llmPrefix: String { get }
 
     func validateTools(tools: [AgentTool]) throws
     func constructScratchpad(intermediateSteps: [AgentAction]) -> AgentScratchPad<ScratchPadContent>

Tool/Sources/LangChain/Agents/FunctionCallingChatAgent.swift

@@ -6,6 +6,15 @@ public protocol CallbackEvent {
 }
 
 public struct CallbackEvents {
+    public struct UnTypedEvent: CallbackEvent {
+        public var info: String
+        init(info: String) {
+            self.info = info
+        }
+    }
+    
+    public var untyped: UnTypedEvent.Type { UnTypedEvent.self }
+    
     private init() {}
 }
 
@@ -52,6 +61,12 @@ public struct CallbackManager {
             observer.handler(info)
         }
     }
+    
+    public func send(_ string: String) {
+        for case let observer as Observer<CallbackEvents.UnTypedEvent> in observers {
+            observer.handler(string)
+        }
+    }
 }
 
 public extension [CallbackManager] {
@@ -65,5 +80,9 @@ public extension [CallbackManager] {
     ) {
         for cb in self { cb.send(keyPath, info) }
     }
+    
+    func send(_ event: String) {
+        for cb in self { cb.send(event) }
+    }
 }
 

Tool/Sources/LangChain/Callback.swift

@@ -0,0 +1,106 @@
+import Foundation
+import Logger
+import OpenAIService
+
+/// This is an agent used to get a structured output.
+public class StructuredOutputChatModelChain<Output: Decodable>: Chain {
+    public struct EndFunction: ChatGPTFunction {
+        public typealias Argument = Output
+        public typealias Result = String
+        public var name: String { "saveFinalAnswer" }
+        public var description: String { "Save the final answer when it's ready" }
+        public let argumentSchema: JSONSchemaValue
+        public var reportProgress: (String) async -> Void = { _ in }
+        public func prepare() async {}
+        public func call(arguments: Argument) async throws -> Result { "" }
+        public init(argumentSchema: JSONSchemaValue) {
+            self.argumentSchema = argumentSchema
+        }
+    }
+
+    struct FunctionProvider: ChatGPTFunctionProvider {
+        var endFunction: EndFunction
+        var functions: [any ChatGPTFunction] {
+            [endFunction]
+        }
+
+        var functionCallStrategy: FunctionCallStrategy? {
+            .name(endFunction.name)
+        }
+    }
+
+    public typealias Input = String
+    public let chatModelChain: ChatModelChain<String>
+    var functionProvider: FunctionProvider
+
+    public init(
+        configuration: ChatGPTConfiguration = UserPreferenceChatGPTConfiguration(),
+        tools: [AgentTool] = [],
+        endFunction: EndFunction,
+        extraSystemPrompt: String = ""
+    ) {
+        functionProvider = .init(
+            endFunction: endFunction
+        )
+        chatModelChain = .init(
+            chatModel: OpenAIChat(
+                configuration: configuration.overriding {
+                    $0.runFunctionsAutomatically = false
+                },
+                memory: nil,
+                functionProvider: functionProvider,
+                stream: false
+            ),
+            stops: ["Observation:"],
+            promptTemplate: { input in
+                [
+                    .init(
+                        role: .system,
+                        content: """
+                        You are a helpful assistant
+                        Generate a final answer to my query as concisely, helpfully and accurately as possible.
+                        You don't ask me for additional information.
+                        \(extraSystemPrompt)
+                        """
+                    ),
+                    .init(role: .user, content: input),
+                ]
+            }
+        )
+    }
+
+    public func callLogic(
+        _ input: String,
+        callbackManagers: [CallbackManager]
+    ) async throws -> Output? {
+        let output = try await chatModelChain.call(input, callbackManagers: callbackManagers)
+        return await parseOutput(output)
+    }
+
+    public func parseOutput(_ output: Output?) -> String {
+        return String(describing: output)
+    }
+
+    public func parseOutput(_ message: ChatMessage) async -> Output? {
+        if let functionCall = message.functionCall {
+            if let function = functionProvider.functions.first(where: {
+                $0.name == functionCall.name
+            }) {
+                if function.name == functionProvider.endFunction.name {
+                    do {
+                        let result = try JSONDecoder().decode(
+                            Output.self,
+                            from: functionCall.arguments.data(using: .utf8) ?? Data()
+                        )
+                        return result
+                    } catch {
+                        return nil
+                    }
+                }
+            }
+        }
+
+        return nil
+    }
+}
+