Model.hs

-- | A 'Model' is the data record callers hand to baikai when
-- dispatching a call. It carries everything baikai needs to know up
-- front to talk to a provider: the 'Api' tag (used to look up the
-- registered handler), the provider name, the base URL, the
-- per-million-token pricing rates, the context window and max output
-- cap, default per-call headers, and a per-API 'Compat' record
-- ('CompatNone' lets the provider auto-detect the record from the base
-- URL; see "Baikai.Compat").
--
-- Use 'modelId' to read the selected upstream model identifier, or
-- 'mkModel' to build a dispatchable record from the three
-- discriminators.
module Baikai.Model
  ( -- * Model
    Model,
    modelId,
    name,
    api,
    provider,
    baseUrl,
    reasoning,
    input,
    cost,
    fastModeCost,
    pricingPolicy,
    contextWindow,
    maxOutputTokens,
    headers,
    compat,
    emptyModel,
    mkModel,

    -- * Cost rates
    ModelCost (..),
    zeroModelCost,
    PricingPolicy (..),
    InputPriceTier (..),
    validatePricingPolicy,

    -- * Capabilities
    InputModality (..),
    OutputKind (..),
    outputKind,
    isJudgmentOnly,

    -- * Compatibility shim
    Compat (..),
    openaiCompletionsCompatFor,
    openaiResponsesCompatFor,
    anthropicMessagesCompatFor,
  )
where

import Baikai.Api (Api (..), renderApi)
import Baikai.Auth qualified as Auth
import Baikai.Compat
  ( AnthropicMessagesCompat,
    OpenAICompletionsCompat,
    OpenAIResponsesCompat,
    autoDetectAnthropicMessages,
    autoDetectOpenAICompletions,
    defaultOpenAIResponsesCompat,
  )
import Baikai.Header (HeaderName)
import Control.Monad (unless)
import Data.Aeson
  ( FromJSON (parseJSON),
    ToJSON (toEncoding, toJSON),
    defaultOptions,
    genericToEncoding,
    genericToJSON,
    withObject,
    (.!=),
    (.:),
    (.:?),
  )
import Data.Aeson qualified as Aeson
import Data.Aeson.KeyMap qualified as KM
import Data.List (nub, sort)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Text (Text)
import GHC.Generics (Generic)
import Numeric.Natural (Natural)

-- | What kinds of input a model accepts: which typed content blocks
-- ("Baikai.Content") the chosen 'Model' may be given.
data InputModality
  = InputText
  | InputImage
  | InputData
  deriving stock (Eq, Ord, Show, Generic)
  deriving anyclass (FromJSON, ToJSON)

-- | Catalog-owned output semantics. Legacy model JSON is generative.
data OutputKind = GenerativeOutput | JudgmentOutput
  deriving stock (Eq, Show, Generic)

instance ToJSON OutputKind where
  toJSON GenerativeOutput = Aeson.String "generative"
  toJSON JudgmentOutput = Aeson.String "judgment"

instance FromJSON OutputKind where
  parseJSON = Aeson.withText "OutputKind" $ \t -> case t of
    "generative" -> pure GenerativeOutput
    "judgment" -> pure JudgmentOutput
    _ -> fail "unknown output kind"

isJudgmentOnly :: Model -> Bool
isJudgmentOnly m = outputKind m == JudgmentOutput

-- | Per-million-token prices in USD as exact 'Rational' values.
data ModelCost = ModelCost
  { inputCost :: !Rational,
    outputCost :: !Rational,
    cacheReadCost :: !Rational,
    cacheWriteCost :: !Rational
  }
  deriving stock (Eq, Show, Generic)
  deriving anyclass (FromJSON, ToJSON)

-- | An exclusive input-context threshold. Its complete rate record
-- applies to every token category in the call once total input exceeds it.
-- Its long cache-write rate is the absolute per-million price of a
-- long-duration cache write in such a call. JSON without that key decodes
-- to 'Nothing'.
data InputPriceTier = InputPriceTier
  { inputAbove :: !Natural,
    rates :: !ModelCost,
    longCacheWriteCost :: !(Maybe Rational)
  }
  deriving stock (Eq, Show, Generic)
  deriving anyclass (ToJSON)

instance FromJSON InputPriceTier where
  parseJSON = withObject "InputPriceTier" $ \o ->
    InputPriceTier <$> o .: "inputAbove" <*> o .: "rates" <*> o .:? "longCacheWriteCost"

-- | Optional catalog policy layered over Model.cost. Long cache-write
-- pricing is an absolute per-million rate, selected only for a shaped
-- long-duration request. The policy-level rate applies while no tier is
-- crossed; a crossed tier supplies its own. A policy with a long rate
-- must therefore give every tier one.
data PricingPolicy = PricingPolicy
  { inputTiers :: ![InputPriceTier],
    longCacheWriteCost :: !(Maybe Rational)
  }
  deriving stock (Eq, Show, Generic)
  deriving anyclass (ToJSON)

instance FromJSON PricingPolicy where
  parseJSON = withObject "PricingPolicy" $ \o -> do
    p <- PricingPolicy <$> o .:? "inputTiers" .!= [] <*> o .:? "longCacheWriteCost"
    either (fail . show) (const (pure p)) (validatePricingPolicy p)

validatePricingPolicy :: PricingPolicy -> Either Text ()
validatePricingPolicy (PricingPolicy tiers policyLong) = do
  let thresholds = map inputAbove tiers
      validRates r = all (>= 0) [inputCost r, outputCost r, cacheReadCost r, cacheWriteCost r]
      tierLongs = [long | InputPriceTier {longCacheWriteCost = long} <- tiers]
  unless (thresholds == sort (nub thresholds)) (Left "Pricing thresholds must be strictly increasing")
  unless (all (validRates . rates) tiers && all (maybe True (>= 0)) (policyLong : tierLongs)) (Left "Pricing rates must be nonnegative")
  unless (policyLong == Nothing || Nothing `notElem` tierLongs) (Left "Every input tier must state its long cache-write rate when the policy has one")

-- | Per-API compatibility shim. 'CompatNone' tells the provider to
-- pick a sensible record by inspecting 'baseUrl'; the two real
-- constructors carry an explicit per-host record that overrides the
-- auto-detection. The records themselves live in 'Baikai.Compat' so
-- this module does not pull in every compat field as a transitive
-- dependency.
data Compat
  = CompatNone
  | CompatOpenAIResponses !OpenAIResponsesCompat
  | CompatOpenAICompletions !OpenAICompletionsCompat
  | CompatAnthropicMessages !AnthropicMessagesCompat
  deriving stock (Eq, Show, Generic)
  deriving anyclass (FromJSON, ToJSON)

-- | Project the 'OpenAICompletionsCompat' that applies to a 'Model':
-- the explicit one if 'compat' is 'CompatOpenAICompletions', otherwise
-- the result of inspecting 'baseUrl' via 'autoDetectOpenAICompletions'.
openaiCompletionsCompatFor :: Model -> OpenAICompletionsCompat
openaiCompletionsCompatFor m = case compat m of
  CompatOpenAICompletions c -> c
  _ -> autoDetectOpenAICompletions (baseUrl m)

-- | Explicit Responses facts or the native Responses defaults.
openaiResponsesCompatFor :: Model -> OpenAIResponsesCompat
openaiResponsesCompatFor m = case compat m of
  CompatOpenAIResponses c -> c
  _ -> defaultOpenAIResponsesCompat

-- | Project the 'AnthropicMessagesCompat' that applies to a 'Model':
-- the explicit one if 'compat' is 'CompatAnthropicMessages',
-- otherwise the result of inspecting 'baseUrl' via
-- 'autoDetectAnthropicMessages'.
--
-- An explicit record always wins. 'CompatNone' means host
-- auto-detection alone: the budget thinking style and sampling
-- parameters supported, which is what every generation before Opus 4.7
-- and every known compatible host accepts. The model id is never
-- consulted, because a generation's wire quirks are a fact of the
-- catalog record, not of its id — every Anthropic model in
-- "Baikai.Models.Generated" carries an explicit
-- 'CompatAnthropicMessages'. A hand-rolled model naming an
-- adaptive-era id must set its own record or start from the catalog
-- value.
anthropicMessagesCompatFor :: Model -> AnthropicMessagesCompat
anthropicMessagesCompatFor m = case compat m of
  CompatAnthropicMessages c -> c
  _ -> autoDetectAnthropicMessages (baseUrl m)

-- | The data record baikai dispatches on.
data Model = Model
  { modelId :: !Text,
    name :: !Text,
    api :: !Api,
    provider :: !Text,
    baseUrl :: !Text,
    reasoning :: !Bool,
    input :: ![InputModality],
    outputKind :: !OutputKind,
    cost :: !ModelCost,
    -- | Optional premium speed rates, in USD per million tokens.
    fastModeCost :: !(Maybe ModelCost),
    pricingPolicy :: !(Maybe PricingPolicy),
    contextWindow :: !Natural,
    -- | The provider's cap on output tokens for this model, or @0@
    -- when it is unknown (a hand-rolled model built from
    -- 'emptyModel', or a catalog entry upstream published no limit
    -- for). @0@ is not a request for zero output: the OpenAI adapter
    -- omits the cap entirely, and the Anthropic adapter — whose API
    -- requires the field and rejects @0@ — sends
    -- @Baikai.Provider.Claude.Internal.Request.uncappedMaxTokensFloor@
    -- instead. An explicit 'Baikai.Options.maxTokens' always wins,
    -- including an explicit @Just 0@.
    maxOutputTokens :: !Natural,
    headers :: !(Map HeaderName Text),
    compat :: !Compat
  }
  deriving stock (Eq, Generic)

instance FromJSON Model where
  parseJSON = Aeson.withObject "Model" $ \o ->
    Aeson.genericParseJSON Aeson.defaultOptions (Aeson.Object (if KM.member "outputKind" o then o else KM.insert "outputKind" (Aeson.String "generative") o))

-- | Rendered field by field rather than derived, so that the value of a
-- credential-carrying header prints as 'Auth.redactedMarker'. A 'Model'
-- is the record most likely to reach a log: it is embedded in every
-- 'Baikai.Response.Response', and the guides tell people to @print@
-- one.
--
-- The format is exactly what @deriving stock Show@ produces — the same
-- record syntax, field order and precedence — because the point is to
-- redact one value, not to invent a rendering. A test in
-- @baikai\/test\/Main.hs@ walks the 'Generic' representation and asserts
-- that every field name appears here, so a field added later cannot
-- silently vanish from 'show'.
--
-- 'Eq' is untouched, and so is 'FromJSON': the field itself still holds
-- what the caller put there and the header is still sent. The one lossy
-- path is a JSON round trip — 'toJSON' writes the marker, so decoding
-- the result gives a 'Model' whose credential header /is/ the marker.
-- That is deliberate; a serialised 'Model' is exactly the thing that
-- should not carry a key.
instance Show Model where
  showsPrec d m =
    showParen (d >= 11) $
      showString "Model {"
        . field "modelId" (modelId m)
        . next "name" (name m)
        . next "api" (api m)
        . next "provider" (provider m)
        . next "baseUrl" (baseUrl m)
        . next "reasoning" (reasoning m)
        . next "input" (input m)
        . next "outputKind" (outputKind m)
        . next "cost" (cost m)
        . next "fastModeCost" (fastModeCost m)
        . next "pricingPolicy" (pricingPolicy m)
        . next "contextWindow" (contextWindow m)
        . next "maxOutputTokens" (maxOutputTokens m)
        . next "headers" (Auth.redactHeaderValues (headers m))
        . next "compat" (compat m)
        . showChar '}'
    where
      field label v = showString label . showString " = " . showsPrec 0 v
      next label v = showString ", " . field label v

-- | Encoded through the 'Generic' representation of a copy whose
-- credential headers have been replaced, so the output is byte-identical
-- to the derived instance's for every model that carries none, and there
-- is no recursion back into this instance.
instance ToJSON Model where
  toJSON = genericToJSON defaultOptions . redactModel
  toEncoding = genericToEncoding defaultOptions . redactModel

redactModel :: Model -> Model
redactModel m = m {headers = Auth.redactHeaderValues (headers m)}

-- | A zero 'ModelCost' across all rates. Useful as a default for
-- models without published pricing (CLI providers, custom hosts).
zeroModelCost :: ModelCost
zeroModelCost =
  ModelCost
    { inputCost = 0,
      outputCost = 0,
      cacheReadCost = 0,
      cacheWriteCost = 0
    }

-- | A blank 'Model'. Useful as a record-update base for hand-rolled
-- 'Model' values in tests and one-shot scripts.
--
-- Its @api@ is @Custom ""@, which no handler can be registered under
-- meaningfully: dispatching a model that still carries it fails with
-- @No provider registered for API: \<blank Custom tag …\>@. Set @api@
-- (and @modelId@) before calling anything.
emptyModel :: Model
emptyModel =
  Model
    { modelId = "",
      name = "",
      api = Custom "",
      provider = "",
      baseUrl = "",
      reasoning = False,
      input = [InputText],
      outputKind = GenerativeOutput,
      cost = zeroModelCost,
      fastModeCost = Nothing,
      pricingPolicy = Nothing,
      contextWindow = 0,
      maxOutputTokens = 0,
      headers = Map.empty,
      compat = CompatNone
    }

-- | Build a dispatchable 'Model' from its three discriminators: the
-- 'Api' tag used for handler lookup, the upstream model id, and the
-- base URL. 'name' defaults to the model id, 'provider' defaults to
-- 'renderApi' of the tag, and all other fields come from 'emptyModel'.
mkModel :: Api -> Text -> Text -> Model
mkModel apiTag modelId_ baseUrl_ =
  emptyModel
    { modelId = modelId_,
      name = modelId_,
      api = apiTag,
      provider = renderApi apiTag,
      baseUrl = baseUrl_
    }