ChatGPTFunction.swift

import Foundation

public enum ChatGPTFunctionCallPhase {
    case detected
    case processing(argumentsJsonString: String)
    case ended(argumentsJsonString: String, result: ChatGPTFunctionResult)
    case error(argumentsJsonString: String, result: Error)
}

public enum ChatGPTFunctionResultUserReadableContent: Sendable {
    public struct ListItem: Sendable {
        public enum Detail: Sendable {
            case link(URL)
            case text(String)
        }
        
        public var title: String
        public var description: String?
        public var detail: Detail?
        
        public init(title: String, description: String? = nil, detail: Detail? = nil) {
            self.title = title
            self.description = description
            self.detail = detail
        }
    }
    
    case text(String)
    case list([ListItem])
    case searchResult([ListItem], queries: [String])
}

public protocol ChatGPTFunctionResult {
    var botReadableContent: String { get }
    var userReadableContent: ChatGPTFunctionResultUserReadableContent { get }
}

extension String: ChatGPTFunctionResult {
    public var botReadableContent: String { self }
    public var userReadableContent: ChatGPTFunctionResultUserReadableContent {
        .text(self)
    }
}

public struct NoChatGPTFunctionArguments: Decodable {}

public protocol ChatGPTFunction {
    typealias NoArguments = NoChatGPTFunctionArguments
    associatedtype Arguments: Decodable
    associatedtype Result: ChatGPTFunctionResult
    typealias ReportProgress = @Sendable (String) async -> Void

    /// The name of this function.
    /// May contain a-z, A-Z, 0-9, and underscores, with a maximum length of 64 characters.
    var name: String { get }
    /// A short description telling the bot when it should use this function.
    var description: String { get }
    /// The arguments schema that the function take in [JSON schema](https://json-schema.org).
    var argumentSchema: JSONSchemaValue { get }
    /// Prepare to call the function
    func prepare(reportProgress: @escaping ReportProgress) async
    /// Call the function with the given arguments.
    func call(arguments: Arguments, reportProgress: @escaping ReportProgress) async throws
        -> Result
}

public extension ChatGPTFunction {
    /// Call the function with the given arguments in JSON.
    func call(
        argumentsJsonString: String,
        reportProgress: @escaping ReportProgress
    ) async throws -> Result {
        let arguments = try await {
            do {
                return try JSONDecoder()
                    .decode(Arguments.self, from: argumentsJsonString.data(using: .utf8) ?? Data())
            } catch {
                await reportProgress(
                    "Error: Failed to decode arguments. \(error.localizedDescription)"
                )
                throw error
            }
        }()
        return try await call(arguments: arguments, reportProgress: reportProgress)
    }
}

public extension ChatGPTFunction where Arguments == NoArguments {
    var argumentSchema: JSONSchemaValue {
        [.type: "object", .properties: [:]]
    }
}

/// This kind of function is only used to get a structured output from the bot.
public protocol ChatGPTArgumentsCollectingFunction: ChatGPTFunction where Result == String {}

public extension ChatGPTArgumentsCollectingFunction {
    func prepare(reportProgress: @escaping ReportProgress = { _ in }) async {
        assertionFailure("This function is only used to get a structured output from the bot.")
    }

    func call(
        arguments: Arguments,
        reportProgress: @escaping ReportProgress = { _ in }
    ) async throws -> Result {
        assertionFailure("This function is only used to get a structured output from the bot.")
        return ""
    }

    func call(
        argumentsJsonString: String,
        reportProgress: @escaping ReportProgress
    ) async throws -> Result {
        assertionFailure("This function is only used to get a structured output from the bot.")
        return ""
    }
}

public struct ChatGPTFunctionSchema: Codable, Equatable, Sendable {
    public var name: String
    public var description: String
    public var parameters: JSONSchemaValue

    public init(name: String, description: String, parameters: JSONSchemaValue) {
        self.name = name
        self.description = description
        self.parameters = parameters
    }
}