markbaseengine/Sources/MarkBaseServer/EmbeddingsAPI.swift

import Foundation

// ─────────────────────────────────────────────────────────────
// Embeddings API Models
// OpenAI Compatible Format
// ─────────────────────────────────────────────────────────────

/// Embeddings request
public struct EmbeddingsRequest: Codable {
    public let model: String
    public let input: InputType
    public let encoding_format: String?
    
    public enum InputType: Codable, Sendable {
        case string(String)
        case strings([String])
        case tokens([Int])
        case tokensList([[Int]])
        
        public init(from decoder: Decoder) throws {
            let container = try decoder.singleValueContainer()
            
            if let string = try? container.decode(String.self) {
                self = .string(string)
            } else if let strings = try? container.decode([String].self) {
                self = .strings(strings)
            } else if let tokens = try? container.decode([Int].self) {
                self = .tokens(tokens)
            } else if let tokensList = try? container.decode([[Int]].self) {
                self = .tokensList(tokensList)
            } else {
                throw DecodingError.dataCorruptedError(
                    in: container,
                    debugDescription: "Invalid input type"
                )
            }
        }
        
        public func encode(to encoder: Encoder) throws {
            var container = encoder.singleValueContainer()
            
            switch self {
            case .string(let string):
                try container.encode(string)
            case .strings(let strings):
                try container.encode(strings)
            case .tokens(let tokens):
                try container.encode(tokens)
            case .tokensList(let tokensList):
                try container.encode(tokensList)
            }
        }
    }
    
    public init(
        model: String,
        input: InputType,
        encoding_format: String? = nil
    ) {
        self.model = model
        self.input = input
        self.encoding_format = encoding_format
    }
}

/// Embeddings response
public struct EmbeddingsResponse: Codable, Sendable {
    public let object: String
    public let data: [EmbeddingData]
    public let model: String
    public let usage: EmbeddingUsage
    
    public init(
        object: String = "list",
        data: [EmbeddingData],
        model: String,
        usage: EmbeddingUsage
    ) {
        self.object = object
        self.data = data
        self.model = model
        self.usage = usage
    }
}

/// Single embedding data
public struct EmbeddingData: Codable, Sendable {
    public let object: String
    public let index: Int
    public let embedding: [Float]
    
    public init(
        object: String = "embedding",
        index: Int,
        embedding: [Float]
    ) {
        self.object = object
        self.index = index
        self.embedding = embedding
    }
}

/// Embedding usage
public struct EmbeddingUsage: Codable, Sendable {
    public let prompt_tokens: Int
    public let total_tokens: Int
    
    public init(prompt_tokens: Int, total_tokens: Int) {
        self.prompt_tokens = prompt_tokens
        self.total_tokens = total_tokens
    }
}