oai-swift/oAI/Models/ModelCategory.swift

//
//  ModelCategory.swift
//  Confab
//
//  Category tags for AI models, inferred from model name/id/description.
//
// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0
// Copyright (C) 2026 Rune Olsen
//
// This file is part of Confab.
//
// Confab 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 Confab 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 SwiftUI

enum ModelCategory: String, CaseIterable, Codable, Sendable {
    case programming = "Programming"
    case math        = "Math"
    case medical     = "Medical"
    case translation = "Translation"
    case roleplay    = "Roleplay"
    case creative    = "Creative"
    case science     = "Science"
    case finance     = "Finance"
    case legal       = "Legal"

    var color: Color {
        switch self {
        case .programming: return .blue
        case .math:        return .orange
        case .medical:     return .red
        case .translation: return .teal
        case .roleplay:    return .pink
        case .creative:    return .purple
        case .science:     return .green
        case .finance:     return Color(red: 0.75, green: 0.60, blue: 0.0)
        case .legal:       return Color(red: 0.55, green: 0.40, blue: 0.20)
        }
    }

    var systemImage: String {
        switch self {
        case .programming: return "chevron.left.forwardslash.chevron.right"
        case .math:        return "function"
        case .medical:     return "cross.fill"
        case .translation: return "globe"
        case .roleplay:    return "theatermasks.fill"
        case .creative:    return "pencil.and.outline"
        case .science:     return "atom"
        case .finance:     return "chart.line.uptrend.xyaxis"
        case .legal:       return "building.columns.fill"
        }
    }

    // MARK: - Category Inference

    /// Infer categories from a model's name, id, and description.
    static func infer(name: String, id: String, description: String?) -> [ModelCategory] {
        let nameId = (name + " " + id).lowercased()
        let desc   = description?.lowercased() ?? ""
        return allCases.filter { $0.matches(nameId: nameId, desc: desc) }
    }

    private func matches(nameId: String, desc: String) -> Bool {
        if nameKeywords.contains(where: { nameId.contains($0) }) { return true }
        if desc.count > 40 && descKeywords.contains(where: { desc.contains($0) }) { return true }
        return false
    }

    // Patterns matched against lowercased "name + id" string
    private var nameKeywords: [String] {
        switch self {
        case .programming:
            return ["code", "coder", "codex", "codellama", "starcoder", "phind",
                    "codestral", "opencoder", "swe-", "devin-", "wizard-code",
                    "replit-code", "qwen-coder", "deepseek-coder", "devstral",
                    "granite-code", "yi-coder", "artigenz", "wavecoder",
                    "programming", "software-", "cursor-"]
        case .math:
            return ["math", "mathem", "numina", "minerva-math", "wizard-math",
                    "deepseek-math", "qwen-math", "numinamath", "mathstral",
                    "qwq", "internlm-math", "mammoth", "mathcoder", "orion-math",
                    "abel-", "metamath"]
        case .medical:
            return ["medical", "meditron", "med42", "medllama", "biomed",
                    "health-llm", "biosage", "clinicalbert", "pubmedbert",
                    "clinical", "llama-med", "openbiomed", "pmc-llama",
                    "doctorglm", "biolm", "biomistral", "medalpaca",
                    "medpalm", "pharmallm", "mimic"]
        case .translation:
            return ["nllb", "madlad", "-aya-", "seamless", "tower-instruct",
                    "alma-", "bayling", "opus-mt", "m2m-100", "mbart",
                    "translate", "multilingual-", "xglm", "madlad-400"]
        case .roleplay:
            return ["roleplay", "role-play", "mytho", "capybara", "cinematika",
                    "manticore", "weaver-", "noromaid", "airoboros", "toppy",
                    "dolphin", "hermes", "openhermes", "psyfighter",
                    "bluemoon", "midnight", "remm", "rose-20b"]
        case .creative:
            return ["creative-writing", "story-writer", "storyllm", "fimbulvetr",
                    "rp-", "goliath", "lzlv", "mlewd"]
        case .science:
            return ["scibert", "biogpt", "galactica", "science-llm", "scillm",
                    "darwin-", "newton-", "eureka-", "sci-"]
        case .finance:
            return ["fingpt", "finma", "finance-llm", "financellm", "pixiu",
                    "flang-", "alphafin", "bloom-finance", "finbert",
                    "stockgpt", "traderllm"]
        case .legal:
            return ["lawbench", "legalbench", "legalbert", "lawgpt", "legal-llm",
                    "lawyerllm", "legalai", "chatlaw", "jurisllm"]
        }
    }

    // Phrases matched against lowercased description (only when description > 40 chars)
    private var descKeywords: [String] {
        switch self {
        case .programming:
            return ["code generation", "designed for coding", "built for code",
                    "coding-focused", "programming assistant", "specialized for code",
                    "optimized for coding", "coding tasks", "software engineering",
                    "for developers", "code completion", "software development",
                    "coding and", "and coding", "writing code"]
        case .math:
            return ["mathematical reasoning", "math competition", "math olympiad",
                    "designed for math", "theorem proving", "quantitative reasoning",
                    "numerical problem", "math, code", "math and code",
                    "mathematics and", "advanced math", "math tasks",
                    "solving math", "competition math", "math problems"]
        case .medical:
            return ["medical knowledge", "clinical reasoning", "healthcare",
                    "biomedical research", "medical question", "trained on medical",
                    "medical domain", "medical literature", "clinical decision",
                    "health information", "medical text", "medical imaging"]
        case .translation:
            return ["machine translation", "language translation",
                    "multilingual translation", "cross-lingual",
                    "translation tasks", "translation between",
                    "natural language translation"]
        case .roleplay:
            return ["designed for roleplay", "roleplay scenarios",
                    "creative roleplay", "interactive roleplay", "character roleplay",
                    "role-playing", "roleplaying", "uncensored", "nsfw",
                    "adult content", "creative fiction", "interactive story"]
        case .creative:
            return ["creative writing", "storytelling", "narrative generation",
                    "fiction writing", "prose generation", "story writing",
                    "creative text", "story generation", "write stories"]
        case .science:
            return ["scientific literature", "scientific research",
                    "chemistry tasks", "biology research", "physics",
                    "scientific reasoning", "science tasks", "stem tasks",
                    "scientific knowledge"]
        case .finance:
            return ["financial analysis", "quantitative finance",
                    "financial modeling", "market analysis", "financial reasoning",
                    "investment analysis", "economic analysis", "trading"]
        case .legal:
            return ["legal document", "case law", "legal reasoning",
                    "legal text", "law and legal", "legal questions",
                    "legal analysis", "legal research", "contract analysis"]
        }
    }
}