oai-swift/oAI/Providers/AIProvider.swift

//
//  AIProvider.swift
//  oAI
//
//  Protocol for AI provider implementations
//
// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0
// Copyright (C) 2026 Rune Olsen
//
// This file is part of oAI.
//
// oAI is licensed under the PolyForm Noncommercial License 1.0.0.
// You may use, study, modify, and share it for any noncommercial
// purpose. Commercial use — including selling oAI or any part of
// it, standalone or bundled into another product or service —
// requires a separate commercial license from the copyright holder.
//
// See the LICENSE file or
// <https://polyformproject.org/licenses/noncommercial/1.0.0> for
// the full license text. For commercial licensing, contact Rune
// Olsen via <https://oai.pm>.


import Foundation

// MARK: - Provider Protocol

protocol AIProvider {
    var name: String { get }
    var capabilities: ProviderCapabilities { get }
    
    func listModels() async throws -> [ModelInfo]
    func getModel(_ id: String) async throws -> ModelInfo?
    func chat(request: ChatRequest) async throws -> ChatResponse
    func streamChat(request: ChatRequest) -> AsyncThrowingStream<StreamChunk, Error>
    func getCredits() async throws -> Credits?

    /// Chat completion with pre-encoded messages for the MCP tool call loop.
    func chatWithToolMessages(model: String, messages: [[String: Any]], tools: [Tool]?, maxTokens: Int?, temperature: Double?) async throws -> ChatResponse
}

// MARK: - Provider Capabilities

struct ProviderCapabilities: Codable {
    let supportsStreaming: Bool
    let supportsVision: Bool
    let supportsTools: Bool
    let supportsOnlineSearch: Bool
    let maxContextLength: Int?
    
    static let `default` = ProviderCapabilities(
        supportsStreaming: true,
        supportsVision: false,
        supportsTools: false,
        supportsOnlineSearch: false,
        maxContextLength: nil
    )
}

// MARK: - Chat Request

struct ReasoningConfig: Sendable {
    let effort: String   // "high", "medium", "low", "minimal"
    let exclude: Bool    // true = use reasoning internally, hide from response
}

struct ChatRequest {
    let messages: [Message]
    let model: String
    let stream: Bool
    let maxTokens: Int?
    let temperature: Double?
    let topP: Double?
    let systemPrompt: String?
    let tools: [Tool]?
    let onlineMode: Bool
    let imageGeneration: Bool
    let reasoning: ReasoningConfig?

    init(
        messages: [Message],
        model: String,
        stream: Bool = true,
        maxTokens: Int? = nil,
        temperature: Double? = nil,
        topP: Double? = nil,
        systemPrompt: String? = nil,
        tools: [Tool]? = nil,
        onlineMode: Bool = false,
        imageGeneration: Bool = false,
        reasoning: ReasoningConfig? = nil
    ) {
        self.messages = messages
        self.model = model
        self.stream = stream
        self.maxTokens = maxTokens
        self.temperature = temperature
        self.topP = topP
        self.systemPrompt = systemPrompt
        self.tools = tools
        self.onlineMode = onlineMode
        self.imageGeneration = imageGeneration
        self.reasoning = reasoning
    }
}

// MARK: - Chat Response

struct ToolCallInfo {
    let id: String
    let type: String
    let functionName: String
    let arguments: String
}

struct ChatResponse: Codable {
    let id: String
    let model: String
    let content: String
    let role: String
    let finishReason: String?
    let usage: Usage?
    let created: Date
    let toolCalls: [ToolCallInfo]?
    let generatedImages: [Data]?

    struct Usage: Codable {
        let promptTokens: Int
        let completionTokens: Int
        let totalTokens: Int
        let cacheCreationInputTokens: Int?
        let cacheReadInputTokens: Int?
        /// Direct USD cost returned by the images API (bypasses token-based calculation).
        let rawCostUSD: Double?

        init(promptTokens: Int, completionTokens: Int, totalTokens: Int, cacheCreationInputTokens: Int? = nil, cacheReadInputTokens: Int? = nil, rawCostUSD: Double? = nil) {
            self.promptTokens = promptTokens
            self.completionTokens = completionTokens
            self.totalTokens = totalTokens
            self.cacheCreationInputTokens = cacheCreationInputTokens
            self.cacheReadInputTokens = cacheReadInputTokens
            self.rawCostUSD = rawCostUSD
        }

        // rawCostUSD is set programmatically, never decoded from API responses
        enum CodingKeys: String, CodingKey {
            case promptTokens = "prompt_tokens"
            case completionTokens = "completion_tokens"
            case totalTokens = "total_tokens"
            case cacheCreationInputTokens = "cache_creation_input_tokens"
            case cacheReadInputTokens = "cache_read_input_tokens"
        }

        init(from decoder: Decoder) throws {
            let c = try decoder.container(keyedBy: CodingKeys.self)
            promptTokens = try c.decode(Int.self, forKey: .promptTokens)
            completionTokens = try c.decode(Int.self, forKey: .completionTokens)
            totalTokens = try c.decode(Int.self, forKey: .totalTokens)
            cacheCreationInputTokens = try c.decodeIfPresent(Int.self, forKey: .cacheCreationInputTokens)
            cacheReadInputTokens = try c.decodeIfPresent(Int.self, forKey: .cacheReadInputTokens)
            rawCostUSD = nil
        }
    }

    // Custom Codable since ToolCallInfo/generatedImages are not from API directly
    enum CodingKeys: String, CodingKey {
        case id, model, content, role, finishReason, usage, created
    }

    init(id: String, model: String, content: String, role: String, finishReason: String?, usage: Usage?, created: Date, toolCalls: [ToolCallInfo]? = nil, generatedImages: [Data]? = nil) {
        self.id = id
        self.model = model
        self.content = content
        self.role = role
        self.finishReason = finishReason
        self.usage = usage
        self.created = created
        self.toolCalls = toolCalls
        self.generatedImages = generatedImages
    }

    init(from decoder: Decoder) throws {
        let container = try decoder.container(keyedBy: CodingKeys.self)
        id = try container.decode(String.self, forKey: .id)
        model = try container.decode(String.self, forKey: .model)
        content = try container.decode(String.self, forKey: .content)
        role = try container.decode(String.self, forKey: .role)
        finishReason = try container.decodeIfPresent(String.self, forKey: .finishReason)
        usage = try container.decodeIfPresent(Usage.self, forKey: .usage)
        created = try container.decode(Date.self, forKey: .created)
        toolCalls = nil
        generatedImages = nil
    }
}

// MARK: - Stream Chunk

struct StreamChunk {
    let id: String
    let model: String
    let delta: Delta
    let finishReason: String?
    let usage: ChatResponse.Usage?

    struct Delta {
        let content: String?
        let role: String?
        let images: [Data]?
        let thinking: String?   // reasoning/thinking tokens (nil if excluded or not supported)

        init(content: String?, role: String?, images: [Data]? = nil, thinking: String? = nil) {
            self.content = content
            self.role = role
            self.images = images
            self.thinking = thinking
        }
    }

    var deltaContent: String? {
        delta.content
    }
}

// MARK: - Tool Definition

struct Tool: Codable {
    let type: String
    let function: Function
    
    struct Function: Codable {
        let name: String
        let description: String
        let parameters: Parameters
        
        struct Parameters: Codable {
            let type: String
            let properties: [String: Property]
            let required: [String]?
            
            struct Property: Codable {
                let type: String
                let description: String
                let `enum`: [String]?
                let items: Items?

                /// Item schema for `type: "array"` properties (e.g. an array of strings).
                struct Items: Codable {
                    let type: String
                }

                init(type: String, description: String, enum: [String]? = nil, items: Items? = nil) {
                    self.type = type
                    self.description = description
                    self.enum = `enum`
                    self.items = items
                }
            }
        }
    }
}

// MARK: - Credits

struct Credits: Codable {
    let balance: Double
    let currency: String
    let usage: Double?
    let limit: Double?
    
    var balanceDisplay: String {
        String(format: "$%.2f", balance)
    }
    
    var usageDisplay: String? {
        guard let usage = usage else { return nil }
        return String(format: "$%.2f", usage)
    }
}

// MARK: - Provider Errors

enum ProviderError: LocalizedError {
    case invalidAPIKey
    case networkError(Error)
    case invalidResponse
    case rateLimitExceeded
    case modelNotFound(String)
    case insufficientCredits
    case timeout
    case unknown(String)
    
    var errorDescription: String? {
        switch self {
        case .invalidAPIKey:
            return "Invalid API key. Please check your settings."
        case .networkError(let error):
            return "Network error: \(error.localizedDescription)"
        case .invalidResponse:
            return "Received invalid response from API"
        case .rateLimitExceeded:
            return "Rate limit exceeded. Please try again later."
        case .modelNotFound(let model):
            return "Model '\(model)' not found"
        case .insufficientCredits:
            return "Insufficient credits"
        case .timeout:
            return "Request timed out"
        case .unknown(let message):
            return "Unknown error: \(message)"
        }
    }
}