Content.hs

{-# LANGUAGE LambdaCase #-}

-- | Typed content blocks for user, assistant, and tool-result messages.
--
-- A message no longer carries a single opaque blob of text. Instead it
-- holds a vector of typed blocks. For user input the blocks can be text
-- or an inline base64-encoded image. For assistant output a block can be
-- text, a thinking trace (for reasoning-capable models), or a tool
-- invocation. For tool-result messages (a caller-supplied reply to a
-- model-issued tool call) blocks can be text or image.
--
-- EP-1 introduces the types; EP-3 streams them, and EP-4 wires tool
-- round-tripping through the providers. Image content is restricted to
-- inline base64 with an explicit @mimeType@: the caller is responsible
-- for the (small, reversible) work of base64-encoding bytes once, and
-- every provider can consume the same shape without a URL-fetch path
-- that varies in failure modes.
module Baikai.Content
  ( -- * Block primitives
    TextContent (..),
    ThinkingContent (..),
    ToolCall (..),
    ImageContent (..),

    -- * Per-role block sums
    UserContent (..),
    AssistantContent (..),
    ToolResultContent (..),

    -- * Smart defaults
    emptyTextContent,
    emptyThinkingContent,
    emptyToolCall,
    emptyImageContent,
    _TextContent,
    _ThinkingContent,
    _ToolCall,
    _ImageContent,
  )
where

import Data.Aeson
  ( FromJSON (parseJSON),
    Options (..),
    SumEncoding (..),
    ToJSON (toJSON),
    Value (Null),
    camelTo2,
    defaultOptions,
    genericParseJSON,
    genericToJSON,
    object,
    withObject,
    (.:),
    (.=),
  )
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Base64 qualified as Base64
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.Encoding qualified as Text
import GHC.Generics (Generic)

-- | A plain-text block. The wire form is @{"text": "..."}@.
newtype TextContent = TextContent
  { text :: Text
  }
  deriving stock (Eq, Show, Generic)
  deriving anyclass (FromJSON, ToJSON)

-- | A reasoning trace block produced by thinking-capable models. The
-- @signature@ field is an opaque continuation token returned by the
-- provider; it must be threaded back into the next request unchanged for
-- the model to resume from the same state. @redacted@ is 'True' when the
-- provider hid the underlying content (Anthropic flags this when its
-- safety system removes a thinking block from the wire response). When
-- @redacted@ is 'True', @thinking@ holds the provider's opaque encrypted
-- payload verbatim; callers must not display or edit it.
data ThinkingContent = ThinkingContent
  { thinking :: !Text,
    signature :: !(Maybe Text),
    redacted :: !Bool
  }
  deriving stock (Eq, Show, Generic)

-- | A model-issued tool invocation. @id_@ has a trailing underscore in
-- Haskell to dodge a clash with @Prelude.id@; the JSON encoding strips
-- it back to @id@.
data ToolCall = ToolCall
  { id_ :: !Text,
    name :: !Text,
    arguments :: !Value
  }
  deriving stock (Eq, Show, Generic)

-- | An inline image block. Bytes are stored decoded; the JSON encoding
-- emits base64 under @data@ and the @mimeType@ camel-snakes to
-- @mime_type@.
data ImageContent = ImageContent
  { imageData :: !ByteString,
    mimeType :: !Text
  }
  deriving stock (Eq, Show, Generic)

-- | What a caller can put into a 'Baikai.Message.UserMessage'.
data UserContent
  = UserText !TextContent
  | UserImage !ImageContent
  deriving stock (Eq, Show, Generic)

-- | What a provider can put into a 'Baikai.Message.AssistantMessage'.
data AssistantContent
  = AssistantText !TextContent
  | AssistantThinking !ThinkingContent
  | AssistantToolCall !ToolCall
  deriving stock (Eq, Show, Generic)

-- | What a caller can put into a 'Baikai.Message.ToolResultMessage'.
data ToolResultContent
  = ToolResultText !TextContent
  | ToolResultImage !ImageContent
  deriving stock (Eq, Show, Generic)

emptyTextContent :: TextContent
emptyTextContent = TextContent {text = Text.empty}

emptyThinkingContent :: ThinkingContent
emptyThinkingContent =
  ThinkingContent
    { thinking = Text.empty,
      signature = Nothing,
      redacted = False
    }

emptyToolCall :: ToolCall
emptyToolCall =
  ToolCall
    { id_ = Text.empty,
      name = Text.empty,
      arguments = Null
    }

emptyImageContent :: ImageContent
emptyImageContent =
  ImageContent
    { imageData = BS.empty,
      mimeType = Text.empty
    }

-- Aeson plumbing ---------------------------------------------------------

snakeOptions :: Options
snakeOptions = defaultOptions {fieldLabelModifier = camelTo2 '_'}

instance FromJSON ThinkingContent where
  parseJSON = genericParseJSON snakeOptions

instance ToJSON ThinkingContent where
  toJSON = genericToJSON snakeOptions

-- Strip the trailing underscore on @id_@ so the wire form is @id@; the
-- other fields keep their natural names.
toolCallOptions :: Options
toolCallOptions =
  defaultOptions
    { fieldLabelModifier = \case
        "id_" -> "id"
        s -> camelTo2 '_' s
    }

instance FromJSON ToolCall where
  parseJSON = genericParseJSON toolCallOptions

instance ToJSON ToolCall where
  toJSON = genericToJSON toolCallOptions

-- ImageContent is encoded manually so the bytes round-trip through
-- base64 on the wire under @data@ while staying decoded in Haskell.
instance ToJSON ImageContent where
  toJSON c =
    object
      [ "data" .= Text.decodeUtf8 (Base64.encode (imageData c)),
        "mime_type" .= mimeType c
      ]

instance FromJSON ImageContent where
  parseJSON = withObject "ImageContent" $ \o -> do
    encoded <- o .: "data"
    mt <- o .: "mime_type"
    case Base64.decode (Text.encodeUtf8 encoded) of
      Left err -> fail ("ImageContent.data is not base64: " <> err)
      Right raw -> pure ImageContent {imageData = raw, mimeType = mt}

contentSumOptions :: Options
contentSumOptions =
  defaultOptions
    { sumEncoding = TaggedObject {tagFieldName = "type", contentsFieldName = "data"},
      constructorTagModifier = camelTo2 '_'
    }

instance FromJSON UserContent where
  parseJSON = genericParseJSON contentSumOptions

instance ToJSON UserContent where
  toJSON = genericToJSON contentSumOptions

instance FromJSON AssistantContent where
  parseJSON = genericParseJSON contentSumOptions

instance ToJSON AssistantContent where
  toJSON = genericToJSON contentSumOptions

instance FromJSON ToolResultContent where
  parseJSON = genericParseJSON contentSumOptions

instance ToJSON ToolResultContent where
  toJSON = genericToJSON contentSumOptions

{-# DEPRECATED _TextContent "Use emptyTextContent instead." #-}
_TextContent :: TextContent
_TextContent = emptyTextContent

{-# DEPRECATED _ThinkingContent "Use emptyThinkingContent instead." #-}
_ThinkingContent :: ThinkingContent
_ThinkingContent = emptyThinkingContent

{-# DEPRECATED _ToolCall "Use emptyToolCall instead." #-}
_ToolCall :: ToolCall
_ToolCall = emptyToolCall

{-# DEPRECATED _ImageContent "Use emptyImageContent instead." #-}
_ImageContent :: ImageContent
_ImageContent = emptyImageContent