JSON.hs

module GenAILib.JSON
  ( customJSONOptions
  , flattenOptions
  , jsonToText
  )
  where

import Data.Aeson (Options, Value (Object), defaultOptions, encode,
  fieldLabelModifier, omitNothingFields)
import Data.Aeson.Key (fromText)
import Data.Aeson.KeyMap ((!?), delete, union)
import Data.String.Conv (toS)
import Data.Text.Lazy qualified as TL


stripInitialUnderscore :: String -> String
stripInitialUnderscore ('_' : fieldName) = fieldName
stripInitialUnderscore fieldName = fieldName


customJSONOptions :: Options
customJSONOptions = defaultOptions
  { fieldLabelModifier = stripInitialUnderscore
  , omitNothingFields = True
  }


-- | Most LLM back-ends expect the options like temperature to be at the
--   top-level of the object. This function will move the options Object up and
--   merge it with the other keys, removing the options Object as well.
--
--   Used internally by this library to construct LLM request bodies
--
-- This
--
-- > { "foo" : 1
-- > , "bar" : 5
-- > , "options" : { "bar" : 2, "baz" : 3 }
-- > }
--
-- Becomes this
--
-- > { "foo" : 1
-- > , "bar" : 2
-- > , "baz" : 3
-- > }
flattenOptions :: Value -> Value

flattenOptions (Object origKm) = Object newKm
  where
    optsKey = fromText "options"
    newKm = case origKm !? optsKey of
      (Just (Object optsKm)) -> delete optsKey origKm `union` optsKm
      _ -> origKm

flattenOptions nonObjValue = nonObjValue


-- | Convenience function to textify a JSON @Value@
jsonToText :: Value -> TL.Text
jsonToText = toS . encode