Context.hs

-- | The 'Context' record — the part of a request that defines the
-- conversation: the optional system prompt, the message vector, and
-- the declared tools the model may invoke.
--
-- The per-call knobs — max tokens, temperature, API key — live on
-- 'Baikai.Options.Options' instead, so a conversation and the settings
-- it is dispatched with are separate values.
--
-- The @tools@ field is on the context because the same tool set applies
-- to every turn, and so is 'appendToolResult', which builds the
-- follow-up request after the model invoked one or more tools. The
-- helper lives here rather than in 'Baikai.Tool' so
-- that 'Baikai.Tool' can stay imports-light (the 'Tool' type is
-- referenced by the @tools@ field, so 'Baikai.Tool' cannot itself
-- depend on 'Context').
module Baikai.Context
  ( Context,
    systemPrompt,
    messages,
    tools,
    emptyContext,
    refuseAssistantDataHistory,
    contextOf,
    systemUser,
    addUser,
    addMessage,
    addResponse,
    appendToolResult,
    appendToolResultText,
  )
where

import Baikai.Content (AssistantContent (..), ToolCall (..), isCutOffToolCall)
import Baikai.Error (BaikaiError, unsupportedFeature)
import Baikai.Message (Message (..), ToolResult, toolResultErrorText, toolResultFromCallNow, toolResultText, user)
import Baikai.Response (Response (..), responseError, responseMessage)
import Baikai.Tool (Tool)
import Control.Applicative ((<|>))
import Control.Lens ((&), (.~), (^.))
import Data.Aeson (ToJSON)
import Data.Generics.Labels ()
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Vector (Vector)
import Data.Vector qualified as V
import GHC.Generics (Generic)

data Context = Context
  { systemPrompt :: !(Maybe Text),
    messages :: !(Vector Message),
    tools :: !(Vector Tool)
  }
  deriving stock (Eq, Show, Generic)
  deriving anyclass (ToJSON)

emptyContext :: Context
emptyContext =
  Context
    { systemPrompt = Nothing,
      messages = V.empty,
      tools = V.empty
    }

instance Semigroup Context where
  a <> b =
    Context
      { systemPrompt = systemPrompt a <|> systemPrompt b,
        messages = messages a <> messages b,
        tools = tools a <> tools b
      }

instance Monoid Context where
  mempty = emptyContext

-- | Build a context from an existing list of messages.
contextOf :: [Message] -> Context
contextOf msgs =
  emptyContext {messages = V.fromList msgs}

-- | Build a context with a system prompt and one user message.
systemUser :: Text -> Text -> Context
systemUser sys prompt =
  emptyContext
    { systemPrompt = Just sys,
      messages = V.singleton (user prompt)
    }

-- | Append a one-text-block user message.
addUser :: Text -> Context -> Context
addUser t = addMessage (user t)

-- | Append a message to the conversation.
addMessage :: Message -> Context -> Context
addMessage msg ctx =
  ctx {messages = V.snoc (messages ctx) msg}

-- | Append a provider response as an assistant message.
addResponse :: Response -> Context -> Context
addResponse resp = addMessage (responseMessage resp)

-- | Execute every 'AssistantToolCall' in @resp@'s message via the
-- caller-supplied dispatcher, then append the assistant message and
-- one 'Baikai.Message.ToolResultMessage' per call to @ctx@. The
-- returned 'Context' is ready to drive the follow-up request that
-- gives the model the tool results.
--
-- Calls are dispatched one at a time, in the order they appear, and the
-- dispatcher returns a rich 'ToolResult' carrying text blocks, image
-- blocks, and an error flag. Any timeout or sandboxing lives in the
-- dispatcher.
--
-- An __error-shaped response__ (one whose 'Baikai.Response.responseError'
-- is 'Just') appends nothing and dispatches nothing: the context comes
-- back unchanged. A failed call has no assistant turn worth replaying
-- and no tool calls to answer, and appending its empty message would put
-- a turn into the transcript that the model never took.
-- 'Baikai.Provider.Registry.runToolLoop' has always stopped on such a
-- response; the documented direct round trip in @docs\/user\/tools.md@
-- reaches here instead, and now behaves the same way.
--
-- A tool call cut off by the output cap
-- ('Baikai.Content.isCutOffToolCall') is __never dispatched__: its
-- arguments are the raw text the model got as far as sending, not a
-- request it finished making. It still gets a
-- 'Baikai.Message.ToolResultMessage', with @isError = True@ explaining
-- why, because a caller driving the exchange by hand expects one result
-- per call and must not silently lose the turn.
-- 'Baikai.Provider.Registry.runToolLoop' stops on such a response
-- instead of reaching here.
appendToolResult ::
  Context ->
  Response ->
  (ToolCall -> IO ToolResult) ->
  IO Context
appendToolResult ctx resp _dispatcher
  | Just _ <- responseError resp = pure ctx
appendToolResult ctx resp dispatcher = do
  let respPayload = resp ^. #message
      respMsg = responseMessage resp
      toolCalls = [tc | AssistantToolCall tc <- V.toList (respPayload ^. #content)]
  results <-
    traverse
      ( \tc -> do
          result <-
            if isCutOffToolCall tc
              then pure cutOffToolResult
              else dispatcher tc
          toolResultFromCallNow tc result
      )
      toolCalls
  pure $
    ctx
      & #messages
        .~ ( (ctx ^. #messages)
               <> V.singleton respMsg
               <> V.fromList results
           )

-- | What 'appendToolResult' reports instead of dispatching a call the
-- model never finished asking for.
cutOffToolResult :: ToolResult
cutOffToolResult =
  toolResultErrorText
    "tool call arguments were cut off by the output limit; the call was not dispatched — raise maxTokens and retry"

-- | Text-only convenience wrapper for the common case where every
-- tool call returns one successful text block.
appendToolResultText ::
  Context ->
  Response ->
  (ToolCall -> IO Text) ->
  IO Context
appendToolResultText ctx resp dispatcher =
  appendToolResult ctx resp (fmap toolResultText . dispatcher)

-- | Generative transports cannot replay structured answer evidence losslessly.
refuseAssistantDataHistory :: Context -> Either BaikaiError ()
refuseAssistantDataHistory ctx = case [ "messages[" <> Text.pack (show i) <> "].content[" <> Text.pack (show j) <> "]"
                                      | (i, AssistantMessage p) <- zip [0 :: Int ..] (V.toList (messages ctx)),
                                        (j, AssistantData _) <- zip [0 :: Int ..] (V.toList (p ^. #content))
                                      ] of
  [] -> Right ()
  path : _ -> Left (unsupportedFeature path "This transport cannot replay AssistantData evidence")