Usage.hs

-- | Token-usage accounting for a single provider call.
--
-- The prompt-side fields are provider-normalized into disjoint token
-- classes. 'inputTokens' counts only non-cached prompt tokens and
-- excludes both cache counters. 'cacheReadTokens' counts prompt tokens
-- served from a provider-side cache; 'cacheWriteTokens' counts prompt
-- tokens stored for future cache reads. Providers whose wire format is
-- inclusive, such as OpenAI Chat Completions where @prompt_tokens@
-- includes @prompt_tokens_details.cached_tokens@, must normalize by
-- subtraction when constructing 'Usage'.
--
-- 'totalTokens' is the sum of the billed token classes:
-- 'inputTokens' + 'outputTokens' + 'cacheReadTokens' +
-- 'cacheWriteTokens'. Provider mappings compute it from the normalized
-- parts rather than trusting a wire total. 'reasoningTokens', when
-- present, is an informational subset of 'outputTokens'; it is already
-- counted in 'outputTokens' and 'totalTokens' and is not billed
-- separately.
--
-- 'cost' is always populated — providers without pricing data fill it
-- with 'zeroCost' (zero across all rates) rather than a 'Nothing' that
-- every cost-reading caller would have to handle. 'Baikai.Cost.Pricing.computeCost'
-- depends on the token classes being disjoint so each class is billed
-- exactly once.
module Baikai.Usage (Usage (..), UsageAvailability (..), UsageCategory (..), BillingFact (..), observeBilling, zeroUsage, sumUsage) where

import Baikai.Cost (Cost, zeroCost)
import Data.Aeson (FromJSON (parseJSON), Options (constructorTagModifier, fieldLabelModifier), ToJSON (toJSON), camelTo2, defaultOptions, genericToJSON, (.!=), (.:), (.:?))
import Data.Aeson qualified as Aeson
import Data.Aeson.KeyMap qualified as KeyMap
import Data.Maybe (fromMaybe)
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import GHC.Generics (Generic)
import Numeric.Natural (Natural)

-- | Billed categories whose omission affects the local calculation.
data UsageCategory = InputUsage | OutputUsage | CacheReadUsage | CacheWriteUsage
  deriving stock (Eq, Ord, Show, Generic)

-- | Provider observations, independent of the requested tier and local rates.
data BillingFact = BillingServiceTier Text | BillingSpeed Text | BillingServerToolUse
  deriving stock (Eq, Ord, Show, Generic)

instance FromJSON BillingFact where parseJSON = Aeson.genericParseJSON usageOptions

instance ToJSON BillingFact where toJSON = genericToJSON usageOptions

-- | Provider facts, independent of local prices. Missing categories are not
-- observed zeroes; inconsistent counters cannot support an exact calculation.
data UsageAvailability = UsageAvailability
  { missingCategories :: !(Set UsageCategory),
    inconsistent :: !Bool,
    billingFacts :: !(Set BillingFact)
  }
  deriving stock (Eq, Show, Generic)

instance FromJSON UsageCategory where parseJSON = Aeson.genericParseJSON usageOptions

instance ToJSON UsageCategory where toJSON = genericToJSON usageOptions

instance FromJSON UsageAvailability where
  parseJSON = Aeson.withObject "UsageAvailability" $ \o -> UsageAvailability <$> o .: "missing_categories" <*> o .: "inconsistent" <*> o .:? "billing_facts" .!= Set.empty

instance ToJSON UsageAvailability where
  toJSON facts = case genericToJSON usageOptions facts of
    Aeson.Object o | Set.null (billingFacts facts) -> Aeson.Object (KeyMap.delete "billing_facts" o)
    value -> value

instance Semigroup UsageAvailability where
  a <> b = UsageAvailability (missingCategories a <> missingCategories b) (inconsistent a || inconsistent b) (billingFacts a <> billingFacts b)

-- | Provider-normalized token usage for one model call.
--
-- The prompt-side classes are disjoint: 'inputTokens' excludes
-- 'cacheReadTokens' and 'cacheWriteTokens'. 'totalTokens' is the sum
-- of all billed token classes, and 'reasoningTokens' is an optional
-- breakdown already included in 'outputTokens'.
data Usage = Usage
  { -- | Non-cached prompt tokens. Excludes 'cacheReadTokens' and
    -- 'cacheWriteTokens'.
    inputTokens :: !Natural,
    -- | Output tokens billed at the model's output rate. Any
    -- 'reasoningTokens' are already included here.
    outputTokens :: !Natural,
    -- | Prompt tokens served from a provider-side cache.
    cacheReadTokens :: !Natural,
    -- | Prompt tokens written into a provider-side cache for future
    -- reads.
    cacheWriteTokens :: !Natural,
    -- | Optional reasoning/thinking-token breakdown. These tokens are
    -- an informational subset of 'outputTokens', not an additional
    -- billed class.
    reasoningTokens :: !(Maybe Natural),
    -- | Sum of the billed token classes:
    -- 'inputTokens' + 'outputTokens' + 'cacheReadTokens' +
    -- 'cacheWriteTokens'.
    totalTokens :: !Natural,
    -- | Availability of provider billing facts. Nothing is the legacy/manual
    -- representation; normalized API responses always carry an annotation.
    availability :: !(Maybe UsageAvailability),
    -- | Computed cost. Incomplete usage or prices carry estimation reasons.
    cost :: !Cost
  }
  deriving stock (Eq, Show, Generic)

usageOptions :: Options
usageOptions = defaultOptions {fieldLabelModifier = camelTo2 '_', constructorTagModifier = camelTo2 '_'}

instance ToJSON Usage where
  toJSON u = case genericToJSON usageOptions u of
    Aeson.Object o | Nothing <- availability u -> Aeson.Object (KeyMap.delete "availability" o)
    value -> value

-- | Empty usage with every count and cost set to zero.
zeroUsage :: Usage
zeroUsage =
  Usage
    { inputTokens = 0,
      outputTokens = 0,
      cacheReadTokens = 0,
      cacheWriteTokens = 0,
      reasoningTokens = Nothing,
      totalTokens = 0,
      availability = Nothing,
      cost = zeroCost
    }

-- | Combine two optional reasoning-token counts. Presence wins: an
-- absent side counts as zero, but the result is only 'Nothing' when
-- both sides are 'Nothing', so a non-reasoning call summed with a
-- reasoning call keeps the reasoning total.
combineReasoning :: Maybe Natural -> Maybe Natural -> Maybe Natural
combineReasoning Nothing Nothing = Nothing
combineReasoning a b = Just (fromMaybe 0 a + fromMaybe 0 b)

instance Semigroup Usage where
  a <> b =
    Usage
      { inputTokens = inputTokens a + inputTokens b,
        outputTokens = outputTokens a + outputTokens b,
        cacheReadTokens = cacheReadTokens a + cacheReadTokens b,
        cacheWriteTokens = cacheWriteTokens a + cacheWriteTokens b,
        reasoningTokens = combineReasoning (reasoningTokens a) (reasoningTokens b),
        totalTokens = totalTokens a + totalTokens b,
        availability = availability a <> availability b,
        cost = cost a <> cost b
      }

instance Monoid Usage where
  mempty = zeroUsage

-- | Total a collection of per-call usages into one.
sumUsage :: (Foldable f) => f Usage -> Usage
sumUsage = foldl' (<>) mempty

-- | Add actual response observations without overwriting missing-count facts.
observeBilling :: [BillingFact] -> Usage -> Usage
observeBilling [] u = u
observeBilling facts u =
  let previous = fromMaybe (UsageAvailability Set.empty False Set.empty) (availability u)
   in u {availability = Just previous {billingFacts = billingFacts previous <> Set.fromList facts}}