refactor: move update plantuml method in AgentExecutor module

bsorrentino · Mar 17, 2024 · ddceebd · ddceebd
1 parent 9ba4b49
commit ddceebd
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 23 deletions.
diff --git a/AIAgent/Sources/AIAgent/AgentExecutor.swift b/AIAgent/Sources/AIAgent/AgentExecutor.swift
@@ -263,7 +263,9 @@ func routeDiagramTranslation( state: AgentExecutorState ) async throws -> String
     /* @objc optional */ func progress(_ message: String) -> Void
 }
 
-public func agentExecutor<T:AgentExecutorDelegate>( openAI: OpenAI, imageUrl: String, delegate:T ) async throws -> String? {
+public func runTranslateDrawingToPlantUML<T:AgentExecutorDelegate>( openAI: OpenAI, 
+                                                                    imageUrl: String,
+                                                                    delegate:T ) async throws -> String? {
 
     let workflow = GraphState { AgentExecutorState() }
 
@@ -302,3 +304,34 @@ public func agentExecutor<T:AgentExecutorDelegate>( openAI: OpenAI, imageUrl: St
 }
 
 
+public func updatePlantUML( openAI: OpenAI, 
+                            withModel model: Model,
+                            input: String,
+                            withInstruction instruction: String ) async throws -> String? {
+    let query = ChatQuery(
+        model: model,
+        messages: [
+            .init(role: .system, content:
+                            """
+                            You are my plantUML assistant.
+                            You must answer exclusively with diagram syntax.
+                            """),
+            .init( role: .assistant, content: input ),
+            .init( role: .user, content: instruction )
+        ],
+        temperature: 0.0,
+        topP: 1.0
+    )
+
+    let chat = try await openAI.chats(query: query)
+
+    let result = chat.choices[0].message.content
+
+    if case .string(let content) = result {
+
+        return content
+    }
+
+    return nil
+
+}
diff --git a/PlantUML/OpenAIObservableService.swift b/PlantUML/OpenAIObservableService.swift
@@ -97,7 +97,7 @@ class OpenAIObservableService : ObservableObject {
     }
 
     @MainActor
-    func query( input: String, instruction: String ) async -> String? {
+    func updatePlantUMLDiagram( input: String, instruction: String ) async -> String? {
 
         guard let openAI /*, let  openAIModel */, case .Ready = status else {
             return nil
@@ -107,26 +107,10 @@ class OpenAIObservableService : ObservableObject {
 
         do {
 
-            let query = ChatQuery(
-                model: openAIModel,
-                messages: [
-                    .init(role: .system, content:
-                                    """
-                                    You are my plantUML assistant.
-                                    You must answer exclusively with diagram syntax.
-                                    """),
-                    .init( role: .assistant, content: input ),
-                    .init( role: .user, content: instruction )
-                ],
-                temperature: 0.0,
-                topP: 1.0
-            )
-
-            let chat = try await openAI.chats(query: query)
-
-            let result = chat.choices[0].message.content
-
-            if case .string(let content) = result {
+            if let content = try await updatePlantUML(openAI: openAI,
+                                                      withModel: openAIModel,
+                                                      input: input,
+                                                      withInstruction: instruction) {
 
                 status = .Ready
 
@@ -167,7 +151,7 @@ extension OpenAIObservableService {
 
         do {
 
-            if let content = try await agentExecutor( openAI: openAI, imageUrl: imageUrl, delegate:delegate) {
+            if let content = try await runTranslateDrawingToPlantUML( openAI: openAI, imageUrl: imageUrl, delegate:delegate) {
 
                 status = .Ready