diff --git a/Sources/NextcloudKit/Models/Assistant/v2/TaskTypes.swift b/Sources/NextcloudKit/Models/Assistant/v2/TaskTypes.swift index b9201f55..2fa7301f 100644 --- a/Sources/NextcloudKit/Models/Assistant/v2/TaskTypes.swift +++ b/Sources/NextcloudKit/Models/Assistant/v2/TaskTypes.swift @@ -29,6 +29,9 @@ public struct TaskTypes: Codable { } public struct TaskTypeData: Codable { + public typealias TaskInputShape = [String: Shape] + public typealias TaskOutputShape = [String: Shape] + public var id: String? public let name: String? public let description: String? @@ -42,21 +45,21 @@ public struct TaskTypeData: Codable { self.inputShape = inputShape self.outputShape = outputShape } -} - -public struct TaskInputShape: Codable { - public let input: Shape? - public init(input: Shape?) { - self.input = input + public func isChat() -> Bool { + id == "core:text2text:chat" } -} -public struct TaskOutputShape: Codable { - public let output: Shape? + public func isTranslate() -> Bool { + id?.contains("translate") == true + } - public init(output: Shape?) { - self.output = output + public func isSingleTextInputOutput(supportedTaskType: String = "Text") -> Bool { + guard let inputShape, let outputShape else { return false } + return inputShape.count == 1 && + outputShape.count == 1 && + inputShape.values.first?.type == supportedTaskType && + outputShape.values.first?.type == supportedTaskType } } diff --git a/Sources/NextcloudKit/NextcloudKit+AssistantV2.swift b/Sources/NextcloudKit/NextcloudKit+AssistantV2.swift index da1d3162..a9aebc21 100644 --- a/Sources/NextcloudKit/NextcloudKit+AssistantV2.swift +++ b/Sources/NextcloudKit/NextcloudKit+AssistantV2.swift @@ -50,12 +50,11 @@ public extension NextcloudKit { if let result = try? decoder.decode(OCSTaskTypesResponse.self, from: data) { var types = result.ocs.data.types.map { (key, value) -> TaskTypeData in var taskType = value - taskType.id = key + taskType.id = taskType.id ?? key return taskType } types = types - .filter { $0.inputShape?.input?.type == supportedTaskType && $0.outputShape?.output?.type == supportedTaskType } - .sorted { ($0.id ?? "") < ($1.id ?? "") } + .filter { $0.isSingleTextInputOutput(supportedTaskType: supportedTaskType) || $0.isChat() || $0.isTranslate() } options.queue.async { continuation.resume(returning: (account: account, types: types, responseData: response, error: .success)) }