refactor: refine AgentExecutor implementation

LangGraph/Sources/LangGraph/AgentExecutor.swift

@@ -1,6 +1,6 @@
 //
 //  File.swift
-//  
+//
 //
 //  Created by bsorrentino on 12/03/24.
 //
@@ -18,10 +18,11 @@ func loadPromptFromBundle( fileName: String ) throws -> String {
     guard let filepath = Bundle.module.path(forResource: fileName, ofType: "txt") else {
         throw _EX("prompt file \(fileName) not found!")
     }
-
+    
     return try String(contentsOfFile: filepath, encoding: .utf8)
 }
 
+
 struct DiagramParticipant : Codable {
     var name: String
     var shape: String
@@ -38,19 +39,36 @@ struct DiagramContainer: Codable {
     var children: [String] // destination
     var description: String
 }
+
+enum DiagramNLPDescription : Codable {
+    case string(String)
+    case array([String])
+
+    init(from decoder: Decoder) throws {
+        let container = try decoder.singleValueContainer()
+        if let string = try? container.decode(String.self) {
+            self = .string(string)
+        } else if let array = try? container.decode([String].self) {
+            self = .array(array)
+        } else {
+            throw _EX("Expected string or array of strings")
+        }
+    }
+}
+
 struct DiagramDescription : Codable {
     var type: String
     var title: String
     var participants: [DiagramParticipant]
     var relations: [DiagramRelation]
     var containers: [DiagramContainer]
-    var description: [String] // NLP description
+    var description: DiagramNLPDescription // NLP description
 }
 
 struct AgentExecutorState : AgentState {
 
     var data: [String : Any]
-
+    
     init() {
         data = [:]
     }
@@ -59,33 +77,29 @@ struct AgentExecutorState : AgentState {
         data = initState
     }
 
-    var openAI:OpenAI? {
-        data["openai"] as? OpenAI
-    }
-
     var diagramImageUrlOrData:String? {
         data["diagram_image_url_or_data"] as? String
     }
-
+    
     var diagramCode:String? {
         data["diagram_code"] as? String
     }
-
+    
     var diagram:DiagramDescription? {
         data["diagram"] as? DiagramDescription
     }
 }
 
 func diagramDescriptionOutputParse( _ content: String ) throws -> DiagramDescription {
 
-    let regex = #/```(json)?(?<code>.*[^`]{3})(```)?/#.dotMatchesNewlines()
+    let regex = #/```(json\n)?({)(?<code>.*)(}\n(```)?)/#.dotMatchesNewlines()
 
-    if let match = try regex.wholeMatch(in: content) {
+    if let match = try regex.firstMatch(in: content) {
 
 
         let decoder = JSONDecoder()
 
-        let code = match.code
+        let code = "{\(match.code)}"
 
         if let data = code.data(using: .utf8) {
 
@@ -98,19 +112,20 @@ func diagramDescriptionOutputParse( _ content: String ) throws -> DiagramDescrip
     else {
         throw _EX( "content doesn't match schema!")
     }
-        
+
 
 }
 
-func describeDiagramImage( state: AgentExecutorState ) async throws -> PartialAgentState {
+func describeDiagramImage<T:AgentExecutorDelegate>( state: AgentExecutorState,
+                                                    openAI:OpenAI,
+                                                    delegate:T ) async throws -> PartialAgentState {
 
-    guard let openAI = state.openAI else {
-        throw _EX("OpenAI not initialized!")
-    }
     guard let imageUrl = state.diagramImageUrlOrData else {
         throw _EX("diagramImageUrlOrData not initialized!")
     }
 
+    await delegate.progress("starting analyze\ndiagram 👀")
+
     let prompt = try loadPromptFromBundle(fileName: "describe_diagram_prompt")
 
     let query = ChatQuery(
@@ -123,33 +138,47 @@ func describeDiagramImage( state: AgentExecutorState ) async throws -> PartialAg
         ],
         maxTokens: 2000
     )
-
+        
     let chatResult = try await openAI.chats(query: query)
 
+    await delegate.progress( "diagram processed ✅")
+
     let result = chatResult.choices[0].message.content
-
+    
     if case .string(let content) = result {
-        return [ "diagram": try diagramDescriptionOutputParse( content ) ]
+        let diagram = try diagramDescriptionOutputParse( content )
+
+        await delegate.progress( "diagram type\n '\(diagram.type)'")
+
+        return [ "diagram": diagram ]
     }
 
     throw _EX("invalid content")
 }
 
-func translateSequenceDiagramDescriptionToPlantUML( state: AgentExecutorState ) async throws -> PartialAgentState {
+func translateSequenceDiagramDescriptionToPlantUML<T:AgentExecutorDelegate>( state: AgentExecutorState,
+                                                    openAI:OpenAI,
+                                                    delegate:T ) async throws -> PartialAgentState {
 
-    guard let openAI = state.openAI else {
-        throw _EX("OpenAI not initialized!")
-    }
     guard let diagram = state.diagram else {
         throw _EX("diagram not initialized!")
     }
 
+    await delegate.progress("starting translate diagram into sequence Diagram")
+
     var prompt = try loadPromptFromBundle(fileName: "sequence_diagram_prompt")
+
+    let description:String = switch(diagram.description) {
+                case .string(let string):
+                    string
+                case .array(let array ):
+                    array.joined(separator: "\n")
+                }
 
     prompt = prompt
         .replacingOccurrences(of: "{diagram_title}", with: diagram.title)
-        .replacingOccurrences(of: "{diagram_description}", with: diagram.description.joined(separator: "\n"))
-
+        .replacingOccurrences(of: "{diagram_description}", with: description)
+    
     let query = ChatQuery(
         model: .gpt3_5Turbo,
         messages: [
@@ -160,22 +189,61 @@ func translateSequenceDiagramDescriptionToPlantUML( state: AgentExecutorState )
         maxTokens: 2000
     )
 
-        let chatResult = try await openAI.chats(query: query)
-
-        let result = chatResult.choices[0].message.content
-
-        if case .string(let content) = result {
-            return [ "diagram_code": content ]
-
-        }
-
+    let chatResult = try await openAI.chats(query: query)
+
+    let result = chatResult.choices[0].message.content
+
+    guard case .string(let content) = result else  {
         throw _EX( "invalid result!" )
-
+    }
+
+    return [ "diagram_code": content ]
+
 }
 
-func translateGenericDiagramDescriptionToPlantUML( state: AgentExecutorState ) async throws -> PartialAgentState {
+func translateGenericDiagramDescriptionToPlantUML<T:AgentExecutorDelegate>( state: AgentExecutorState, 
+                                                                            openAI:OpenAI,
+                                                                            delegate:T ) async throws -> PartialAgentState {
+
+    guard let diagram = state.diagram else {
+        throw _EX("diagram not initialized!")
+    }
+
+    await delegate.progress("starting translate diagram into generic Diagram")
+
+    var prompt = try loadPromptFromBundle(fileName: "generic_diagram_prompt")
+
+    let encoder = JSONEncoder()
+
+    let data = try encoder.encode(diagram)
+
+    guard let content = String(data: data, encoding: .utf8) else {
+        throw _EX("diagram encoding error!")
+    }
+
+    prompt = prompt
+            .replacingOccurrences(of: "{diagram_description}", with: content)
+
+    let query = ChatQuery(
+        model: .gpt3_5Turbo,
+        messages: [
+            Chat(role: .user, content: [
+                ChatContent(text: prompt),
+            ])
+        ],
+        maxTokens: 2000
+    )
+
+    let chatResult = try await openAI.chats(query: query)
+
+    let result = chatResult.choices[0].message.content
+
+    guard case .string(let content) = result else {
+        throw _EX( "invalid result!" )
+    }
+
+    return [ "diagram_code": content ]
 
-    return [:]
 }
 
 func routeDiagramTranslation( state: AgentExecutorState ) async throws -> String {
@@ -190,14 +258,25 @@ func routeDiagramTranslation( state: AgentExecutorState ) async throws -> String
     }
 }
 
-public func agentExecutor( openAI: OpenAI, imageUrl: String ) async throws -> String? {
+@MainActor /* @objc */ public protocol AgentExecutorDelegate {
+
+    /* @objc optional */ func progress(_ message: String) -> Void
+}
+
+public func agentExecutor<T:AgentExecutorDelegate>( openAI: OpenAI, imageUrl: String, delegate:T ) async throws -> String? {
 
     let workflow = GraphState( stateType: AgentExecutorState.self )
 
-    try workflow.addNode("agent_describer", action: describeDiagramImage )
-    try workflow.addNode("agent_sequence_plantuml", action: translateSequenceDiagramDescriptionToPlantUML)
-    try workflow.addNode("agent_generic_plantuml", action: translateGenericDiagramDescriptionToPlantUML)
-
+    try workflow.addNode("agent_describer", action: { state in
+        try await describeDiagramImage(state: state, openAI: openAI, delegate: delegate)
+    })
+    try workflow.addNode("agent_sequence_plantuml", action: { state in
+        try await translateSequenceDiagramDescriptionToPlantUML( state: state, openAI:openAI, delegate:delegate )
+    })
+     try workflow.addNode("agent_generic_plantuml", action: { state in
+         try await translateGenericDiagramDescriptionToPlantUML( state: state, openAI:openAI, delegate:delegate )
+    })
+
     try workflow.addEdge(sourceId: "agent_sequence_plantuml", targetId: END)
     try workflow.addEdge(sourceId: "agent_generic_plantuml", targetId: END)
 
@@ -210,16 +289,16 @@ public func agentExecutor( openAI: OpenAI, imageUrl: String ) async throws -> St
         ]
     )
     workflow.setEntryPoint( "agent_describer")
-
+    
     let app = try workflow.compile()
-
+    
     let inputs:[String : Any] = [
         "openai": openAI,
         "diagram_image_url_or_data": imageUrl
-    ] 
-
+    ]
+    
     let response = try await app.invoke( inputs: inputs)
-
+    
     return response.diagramCode
 }