OpenAI.hs

module LLM.Providers.OpenAI
  ( openAIGateway,
    openAIGatewayWith,
    openAIGatewayWithName,
    openAIProvider,
    openAIProviderWith,
    openAIProviderWithName,
    parseOpenAIResponse,
    parseOpenAIObjectResponse,
    parseOpenAIUsage,
    buildMessages,
    encodeTurn,
    encodeToolDef,
    parseOpenAIStream,
    openAIBuildBody,
    openAIBuildBodyPairs,
    authHeader,
  )
where

import Data.Aeson
  ( KeyValue ((.=)),
    Object,
    Value (Object, String),
    decodeStrict',
    encode,
    object,
    toJSON,
    withObject,
    (.!=),
    (.:),
    (.:?),
  )
import Data.Aeson.Types (Pair, Parser, parseEither, parseMaybe)
import Data.ByteString.Lazy qualified as BSL
import Control.Monad (when, zipWithM)
import Data.Either (lefts, rights)
import Data.Foldable (foldl', for_, forM_)
import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef)
import Data.Map.Strict qualified as Map
import Data.Maybe (catMaybes, fromMaybe, mapMaybe)
import Control.Applicative ((<|>))
import Data.Char (isSpace)
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding (decodeUtf8, encodeUtf8)
import LLM.Core.LLMProvider (LLMProvider (..), toGateway)
import LLM.Core.ProviderUtils (decodeObjectText, decodeToolArguments, handleStreamResponse, invalidToolArgumentsError, lenientConfig, normalizeSchemaOpenAI, readResponseBody, unexpectedResponse)
import LLM.Core.SSE (SSEEvent (sseData), malformedEvent, readSSEEvents, streamFailure)
import LLM.Core.Types
  ( commonImageMediaTypes,
    projectReasoning,
    ChatRequest
      ( reqConversation,
        reqMaxTokens,
        reqModel,
        reqSystem,
        reqTemperature,
        reqThinking,
        reqTools
      ),
    ContentPart (..),
    LLMError (EmptyResponse, HttpError),
    LLMGateway,
    LLMResult,
    LLMTextResult,
    MessageEncodeOptions (..),
    PartBody (..),
    StreamEvent (..),
    ThinkingContent (..),
    ThinkingMode (..),
    ToolCall (..),
    ToolDef (toolDescription, toolName, toolParameters),
    ToolResult (trCallId, trContent),
    Turn (..),
    ImageSource (..),
    defaultMessageEncodeOptions,
    mkChatResponse,
    mkToolCall,
    stripForeignOpaque,
    textPart,
    thinkingPart,
    toolCallPart,
  )
import LLM.Core.Usage (Usage (..))
import Network.HTTP.Client qualified as HC
import Network.HTTP.Req
  ( Option,
    POST (POST),
    ReqBodyJson (ReqBodyJson),
    Url,
    header,
    https,
    reqBr,
    runReq,
    (/:),
  )

openAIProviderName :: Text
openAIProviderName = "openai"

-- | Create an OpenAI client for api.openai.com. Takes the API key as a parameter.
openAIGateway :: Text -> LLMGateway
openAIGateway apiKey = toGateway $ openAIProvider apiKey

-- | Create an OpenAI-compatible client with a custom base URL.
openAIGatewayWith :: Url scheme -> Option scheme -> Text -> LLMGateway
openAIGatewayWith = openAIGatewayWithName "openai"

-- | Create an OpenAI-compatible client with a custom provider name and base URL.
openAIGatewayWithName :: Text -> Url scheme -> Option scheme -> Text -> LLMGateway
openAIGatewayWithName name baseUrl baseOpts apiKey = toGateway (openAIProviderWithName name baseUrl baseOpts apiKey)

-- | Create an OpenAI provider. Takes the API key as a parameter.
openAIProvider :: Text -> LLMProvider
openAIProvider = openAIProviderWith (https "api.openai.com") mempty

openAIProviderWith :: Url scheme -> Option scheme -> Text -> LLMProvider
openAIProviderWith = openAIProviderWithName "openai"

openAIProviderWithName :: Text -> Url scheme -> Option scheme -> Text -> LLMProvider
openAIProviderWithName name baseUrl baseOpts apiKey =
  LLMProvider
    { providerName = name,
      providerImageMediaTypes = commonImageMediaTypes,
      providerImageUrlSupport = True,
      buildBody = openAIBuildBody,
      sendRequest = sendRequest,
      openStream = \body k ->
        runReq lenientConfig $ do
          let url = baseUrl /: "v1" /: "chat" /: "completions"
              opts = baseOpts <> authHeader apiKey
          reqBr POST url (ReqBodyJson body) opts $ \resp ->
            handleStreamResponse resp k,
      parseStream = const parseOpenAIStream,
      parseResponse = pure . parseOpenAIResponse,
      buildObjectBody = \r schema ->
        object $
          openAIBuildBodyPairs False r
            <> [ "response_format"
                   .= object
                     [ "type" .= ("json_schema" :: Text),
                       "json_schema"
                         .= object
                           [ "name" .= ("response" :: Text),
                             "schema" .= normalizeSchemaOpenAI schema,
                             "strict" .= True
                           ]
                     ]
               ],
      sendObjectRequest = sendRequest,
      parseObjectResponse = pure . parseOpenAIObjectResponse
    }
  where
    sendRequest body =
      runReq lenientConfig $ do
        let url = baseUrl /: "v1" /: "chat" /: "completions"
            opts = baseOpts <> authHeader apiKey
        reqBr POST url (ReqBodyJson body) opts readResponseBody

authHeader :: Text -> Option scheme
authHeader apiKey
  | T.null apiKey = mempty
  | otherwise = header "Authorization" ("Bearer " <> encodeUtf8 apiKey)

-- Request body

openAIBuildBody :: Bool -> ChatRequest -> Value
openAIBuildBody stream r = object $ openAIBuildBodyPairs stream r

openAIBuildBodyPairs :: Bool -> ChatRequest -> [Pair]
openAIBuildBodyPairs stream r =
  [ "model" .= r.reqModel,
    "max_completion_tokens" .= r.reqMaxTokens,
    "messages" .= buildMessages defaultMessageEncodeOptions r
  ]
    ++ ["temperature" .= t | Just t <- [r.reqTemperature]]
    ++ effortPairs r
    ++ ["tools" .= map encodeToolDef r.reqTools | not (null r.reqTools)]
    ++ ["stream" .= True | stream]
    ++ ["stream_options" .= object ["include_usage" .= True] | stream]

-- | Chat Completions reasoning_effort, when the catalog requests thinking.
effortPairs :: ChatRequest -> [Pair]
effortPairs r =
  case r.reqThinking of
    Just tm
      | tm.tmEnabled,
        Just e <- tm.tmEffort ->
          ["reasoning_effort" .= e]
    _ -> []

buildMessages :: MessageEncodeOptions -> ChatRequest -> [Value]
buildMessages opts r =
  maybe [] (\sys -> [object ["role" .= ("system" :: Text), "content" .= sys]]) r.reqSystem
    ++ concatMap (encodeTurn opts) r.reqConversation

encodeTurn :: MessageEncodeOptions -> Turn -> [Value]
encodeTurn _ (UserMessage parts) =
  [ object
      [ "role" .= ("user" :: Text),
        "content" .= encodeUserContent parts
      ]
  ]
encodeTurn opts (AssistantMessage parts) =
  let cleaned = map (stripForeignOpaque openAIProviderName) parts
      -- Cache hints are ignored on the OpenAI wire; content is unchanged.
      text = T.concat [t | ContentPart (TextPart t) _ <- cleaned]
      mReasoning = projectReasoning cleaned
      calls = [tc | ContentPart (ToolCallPart tc) _ <- cleaned]
      requireReasoning
        | opts.meoRequireReasoning = Just . fromMaybe ""
        | otherwise = id
   in [ object $
          ["role" .= ("assistant" :: Text)]
            -- Without tool calls, send content even when empty: OpenAI
            -- rejects an assistant message with neither (e.g. a Claude turn
            -- that held only thinking).
            ++ ["content" .= text | not (T.null text) || null calls]
            ++ [ "reasoning_content" .= rc
                 | opts.meoIncludeReasoning,
                   Just rc <- [requireReasoning mReasoning]
               ]
            ++ ["tool_calls" .= map encodeToolCall calls | not (null calls)]
      ]
encodeTurn _ (ToolTurn results) =
  map encodeToolResult results

-- | Single text stays a string (recorded-fixture compatible); mixed or image
-- content uses the OpenAI multimodal content-part array.
-- Cache hints are omitted from the wire JSON.
encodeUserContent :: [ContentPart] -> Value
encodeUserContent [ContentPart (TextPart t) _] = String t
encodeUserContent parts = toJSON (mapMaybe encodeUserPart parts)
  where
    encodeUserPart (ContentPart (TextPart t) _) =
      Just $ object ["type" .= ("text" :: Text), "text" .= t]
    encodeUserPart (ContentPart (ImagePart src) _) =
      Just $ encodeImagePart src
    encodeUserPart _ = Nothing

encodeImagePart :: ImageSource -> Value
encodeImagePart (ImageUrl url) =
  object
    [ "type" .= ("image_url" :: Text),
      "image_url" .= object ["url" .= url]
    ]
encodeImagePart (ImageBase64 mediaType data_) =
  object
    [ "type" .= ("image_url" :: Text),
      "image_url"
        .= object
          [ "url" .= ("data:" <> mediaType <> ";base64," <> data_)
          ]
    ]

encodeToolDef :: ToolDef -> Value
encodeToolDef td =
  object
    [ "type" .= ("function" :: Text),
      "function"
        .= object
          [ "name" .= td.toolName,
            "description" .= td.toolDescription,
            "parameters" .= td.toolParameters
          ]
    ]

encodeToolCall :: ToolCall -> Value
encodeToolCall tc =
  object
    [ "id" .= tc.tcId,
      "type" .= ("function" :: Text),
      "function"
        .= object
          [ "name" .= tc.tcName,
            "arguments" .= decodeUtf8 (BSL.toStrict (encode tc.tcArguments))
          ]
    ]

encodeToolResult :: ToolResult -> Value
encodeToolResult tr =
  object
    [ "role" .= ("tool" :: Text),
      "tool_call_id" .= tr.trCallId,
      "content" .= tr.trContent
    ]

-- Response parsing

parseOpenAIResponse :: Value -> LLMTextResult
parseOpenAIResponse v = case parseMaybe go v of
  Nothing -> Left EmptyResponse
  Just (finishReason, parts) ->
    case lefts parts of
      name : _ -> Left (invalidToolArgumentsError name (truncationReason finishReason))
      []
        | null parts -> Left EmptyResponse
        | otherwise -> Right (mkChatResponse (rights parts) (parseOpenAIUsage v))
  where
    go :: Value -> Parser (Maybe Text, [Either Text ContentPart])
    go = withObject "OpenAIResponse" $ \o -> do
      (choice : _) <- o .: "choices" :: Parser [Value]
      withObject
        "choice"
        ( \co -> do
            finishReason <- co .:? "finish_reason"
            msg <- co .: "message"
            parts <- withObject "message" parseMessage msg
            pure (finishReason, parts)
        )
        choice

    -- Tool calls with invalid JSON arguments come back as 'Left' tool name.
    parseMessage :: Object -> Parser [Either Text ContentPart]
    parseMessage mo = do
      mReasoning <- mo .:? "reasoning_content" :: Parser (Maybe Text)
      contentBlocks <- do
        mc <- mo .:? "content" :: Parser (Maybe Text)
        pure [textPart t | Just t <- [mc], not (T.null t)]
      toolBlocks <- do
        tcs <- mo .:? "tool_calls" .!= [] :: Parser [Value]
        mapM parseToolCall tcs
      let thinkingBlocks =
            [ thinkingPart (ThinkingContent (Just rc) Nothing)
              | Just rc <- [mReasoning],
                not (T.null rc)
            ]
      pure (map Right (thinkingBlocks ++ contentBlocks) ++ toolBlocks)

    parseToolCall :: Value -> Parser (Either Text ContentPart)
    parseToolCall = withObject "tool_call" $ \tc -> do
      cid <- tc .: "id"
      fn <- tc .: "function"
      withObject
        "function"
        ( \f -> do
            name <- f .: "name"
            argsStr <- f .: "arguments" :: Parser Text
            pure $ case decodeToolArguments argsStr of
              Just args -> Right (toolCallPart (mkToolCall cid name args))
              Nothing -> Left name
        )
        fn

-- | The finish reason, if it says the response hit the token limit.
truncationReason :: Maybe Text -> Maybe Text
truncationReason r = if r == Just "length" then r else Nothing

-- | Parse a structured-output response from an OpenAI-compatible chat
-- completions API: the JSON text in the first choice's message content.
--
-- No choice, or a missing or null content, is 'EmptyResponse'; so is
-- content that is empty once code fences are stripped. Content that is not
-- valid JSON, and a body without the expected fields, are 'ParseError's
-- that quote the decoder's message and the start of the text or body.
parseOpenAIObjectResponse :: Value -> LLMResult (Value, Maybe Usage)
parseOpenAIObjectResponse v = case parseEither content v of
  Left err -> Left (unexpectedResponse err v)
  Right Nothing -> Left EmptyResponse
  Right (Just text) -> (\obj -> (obj, parseOpenAIUsage v)) <$> decodeObjectText text
  where
    content :: Value -> Parser (Maybe Text)
    content = withObject "OpenAIObjectResponse" $ \o -> do
      choices <- o .: "choices" :: Parser [Value]
      case choices of
        [] -> pure Nothing
        choice : _ ->
          withObject "choice" (\co -> co .: "message" >>= withObject "message" (.:? "content")) choice

parseOpenAIUsage :: Value -> Maybe Usage
parseOpenAIUsage = parseMaybe $ withObject "OpenAIResponse" $ \o -> do
  u <- o .: "usage"
  withObject "usage" parseOpenAIUsageObject u

-- | Shared OpenAI Chat Completions / DeepSeek usage shape.
--
-- @prompt_tokens@ is already the total input (including cache hits). Cache
-- reads come from DeepSeek's @prompt_cache_hit_tokens@ or OpenAI's
-- @prompt_tokens_details.cached_tokens@. Optional @cache_write_tokens@ map to
-- creation counters.
parseOpenAIUsageObject :: Object -> Parser Usage
parseOpenAIUsageObject uo = do
  prompt <- uo .: "prompt_tokens"
  completion <- uo .: "completion_tokens"
  cacheReadDeepSeek <- uo .:? "prompt_cache_hit_tokens"
  details <- uo .:? "prompt_tokens_details" :: Parser (Maybe Value)
  (cacheReadDetails, cacheWrite) <- case details of
    Nothing -> pure (Nothing, Nothing)
    Just d ->
      withObject
        "prompt_tokens_details"
        ( \dto ->
            (,)
              <$> dto .:? "cached_tokens"
              <*> dto .:? "cache_write_tokens"
        )
        d
  let cacheRead = fromMaybe 0 (cacheReadDeepSeek <|> cacheReadDetails)
      cacheCreate = fromMaybe 0 cacheWrite
  pure
    Usage
      { usageInputTokens = prompt,
        usageOutputTokens = completion,
        usageCacheReadTokens = cacheRead,
        usageCacheCreationTokens = cacheCreate,
        usageTotalCost = 0
      }

-- Streaming

parseOpenAIStream :: HC.BodyReader -> (StreamEvent -> IO ()) -> IO LLMTextResult
parseOpenAIStream reader callback = do
  blocksRef <- newIORef ([] :: [ContentPart])
  -- Reasoning fragments, newest first; joined once at the end.
  reasoningRef <- newIORef ([] :: [Text])
  usageRef <- newIORef Nothing
  -- Tool calls being assembled, keyed by index: (id, name, argument
  -- fragments newest first).
  toolAccRef <- newIORef (Map.empty :: Map.Map Int (Text, Text, [Text]))
  errorRef <- newIORef (Nothing :: Maybe LLMError)
  finishReasonRef <- newIORef (Nothing :: Maybe Text)
  badToolRef <- newIORef (Nothing :: Maybe Text) -- first tool with invalid JSON args
  let flushTools = do
        tools <- readIORef toolAccRef
        forM_ (Map.elems tools) $ \(cid, name, fragments) ->
          case decodeToolArguments (T.concat (reverse fragments)) of
            Just args -> do
              let tc = mkToolCall cid name args
              modifyIORef' blocksRef (toolCallPart tc :)
              callback (StreamToolCall tc)
            Nothing -> modifyIORef' badToolRef (<|> Just name)
        writeIORef toolAccRef Map.empty
  -- Set by @[DONE]@ or a non-null @finish_reason@. Some compatible servers
  -- send only one of the two.
  doneRef <- newIORef False
  malformedRef <- newIORef (Nothing :: Maybe Text)
  badUtf8 <- readSSEEvents (HC.brRead reader) $ \sse -> do
    let raw = sse.sseData
    if raw == "[DONE]"
      then writeIORef doneRef True
      else case decodeStrict' (encodeUtf8 raw) of
        Nothing
          | T.all isSpace raw -> pure ()
          | otherwise -> modifyIORef' malformedRef (<|> Just (malformedEvent raw))
        Just v -> do
          for_ (parseMaybe parseStreamError v) $ \status ->
            modifyIORef' errorRef (<|> Just (HttpError status raw))
          case parseMaybe parseStreamReasoningDelta v of
            Just (Just txt) | not (T.null txt) -> do
              modifyIORef' reasoningRef (txt :)
              callback (StreamReasoningDelta txt)
            _ -> pure ()
          case parseMaybe parseStreamTextDelta v of
            Just txt | not (T.null txt) -> do
              modifyIORef' blocksRef (textPart txt :)
              callback (StreamDelta txt)
            _ -> pure ()
          for_ (parseMaybe parseStreamToolDeltas v) $ \deltas ->
            modifyIORef' toolAccRef (\acc -> foldl' addToolDelta acc deltas)
          case parseMaybe parseFinishReason v of
            Just reason -> do
              writeIORef finishReasonRef (Just reason)
              writeIORef doneRef True
              when (reason == "tool_calls") flushTools
            Nothing -> pure ()
          case parseMaybe parseStreamUsage v of
            Just u -> writeIORef usageRef (Just u)
            Nothing -> pure ()
  flushTools
  textBlocks <- reverse <$> readIORef blocksRef
  reasoning <- T.concat . reverse <$> readIORef reasoningRef
  usage <- readIORef usageRef
  let thinkingBlocks =
        [ thinkingPart (ThinkingContent (Just reasoning) Nothing)
          | not (T.null reasoning)
        ]
      parts = thinkingBlocks ++ textBlocks
  mError <- readIORef errorRef
  mBadTool <- readIORef badToolRef
  finishReason <- readIORef finishReasonRef
  done <- readIORef doneRef
  malformed <- (badUtf8 <|>) <$> readIORef malformedRef
  pure $ case (mError <|> streamFailure "[DONE] or finish_reason" done malformed, mBadTool) of
    (Just err, _) -> Left err
    (Nothing, Just name) -> Left (invalidToolArgumentsError name (truncationReason finishReason))
    _
      | null parts -> Left EmptyResponse
      | otherwise -> Right (mkChatResponse parts usage)

-- | A mid-stream @{"error": ...}@ chunk, mapped to an HTTP status so
-- 'LLM.Core.Utils.isRetryable' treats it like the same error sent before
-- streaming started. Uses a numeric @code@/@status@ when the server sends
-- one; otherwise 429 for rate limits, 503 for overload, else 500.
parseStreamError :: Value -> Parser Int
parseStreamError = withObject "chunk" $ \o -> do
  e <- o .: "error"
  case e of
    String msg -> pure (classify [msg])
    Object eo -> do
      mCode <- eo .:? "code" <|> pure Nothing
      mStatus <- eo .:? "status" <|> pure Nothing
      texts <- mapM (\k -> eo .:? k <|> pure Nothing) ["type", "code", "message"]
      pure $ case filter (\c -> c >= 400 && c < 600) (catMaybes [mCode, mStatus]) of
        c : _ -> c
        [] -> classify (catMaybes texts)
    _ -> fail "not an error"
  where
    classify ts
      | any (hasWord "rate") ts = 429
      | any (hasWord "overload") ts = 503
      | otherwise = 500
    hasWord w t = w `T.isInfixOf` T.toLower t

parseStreamReasoningDelta :: Value -> Parser (Maybe Text)
parseStreamReasoningDelta = withObject "chunk" $ \o -> do
  (c : _) <- o .: "choices" :: Parser [Value]
  withObject "choice" (\co -> do d <- co .: "delta"; withObject "delta" (.:? "reasoning_content") d) c

parseStreamTextDelta :: Value -> Parser Text
parseStreamTextDelta = withObject "chunk" $ \o -> do
  (c : _) <- o .: "choices" :: Parser [Value]
  withObject "choice" (\co -> do d <- co .: "delta"; withObject "delta" (.: "content") d) c

-- | One streamed tool-call fragment: index, id, name, argument text.
type ToolDelta = (Int, Maybe Text, Maybe Text, Text)

-- | Merge a fragment into the accumulated tool calls. A fragment whose id
-- differs from the call already at its index starts a new call (servers
-- that omit @index@ would otherwise merge distinct calls).
addToolDelta :: Map.Map Int (Text, Text, [Text]) -> ToolDelta -> Map.Map Int (Text, Text, [Text])
addToolDelta acc (idx, mId, mName, argChunk) =
  case Map.lookup idx acc of
    Just (cid, n, a)
      | maybe True (== cid) mId || T.null cid ->
          Map.insert idx (fromMaybe cid mId, fromMaybe n mName, argChunk : a) acc
      | otherwise -> Map.insert (nextIndex acc) new acc
    Nothing -> Map.insert idx new acc
  where
    new = (fromMaybe "" mId, fromMaybe "" mName, [argChunk])
    nextIndex m = maybe 0 ((+ 1) . fst) (Map.lookupMax m)

-- | All tool-call fragments in a chunk. Ollama may send several complete
-- calls in one chunk and may omit @index@; the position is used then.
parseStreamToolDeltas :: Value -> Parser [ToolDelta]
parseStreamToolDeltas = withObject "chunk" $ \o -> do
  (c : _) <- o .: "choices" :: Parser [Value]
  withObject
    "choice"
    ( \co -> do
        d <- co .: "delta"
        withObject
          "delta"
          ( \d' -> do
              tcs <- d' .: "tool_calls" :: Parser [Value]
              zipWithM parseOne [0 ..] tcs
          )
          d
    )
    c
  where
    parseOne pos = withObject "tool_call" $ \tco -> do
      idx <- tco .:? "index" .!= pos
      mId <- tco .:? "id"
      fn <- tco .:? "function" .!= object []
      withObject
        "function"
        ( \f -> do
            mName <- f .:? "name"
            args <- f .:? "arguments" .!= ""
            pure (idx, mId, mName, args)
        )
        fn

parseFinishReason :: Value -> Parser Text
parseFinishReason = withObject "chunk" $ \o -> do
  (c : _) <- o .: "choices" :: Parser [Value]
  withObject "choice" (.: "finish_reason") c

parseStreamUsage :: Value -> Parser Usage
parseStreamUsage = withObject "chunk" $ \o -> do
  u <- o .: "usage"
  withObject "usage" parseOpenAIUsageObject u