OpenAISpec.hs

{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}

module LLM.OpenAISpec (spec) where

import Autodocodec qualified as AC
import Autodocodec.Schema (jsonSchemaVia)
import Data.Aeson
  ( Value (Array, Object, String),
    eitherDecodeFileStrict',
    object,
    toJSON,
    (.=),
  )
import Data.Aeson.Key qualified as K
import Data.Aeson.KeyMap qualified as KM
import Data.Maybe (fromMaybe, mapMaybe)
import Data.Text (Text)
import Data.Vector qualified as V
import Data.ByteString qualified as BS
import Data.IORef (newIORef, readIORef, writeIORef)
import Data.Text qualified as T
import LLM.Core.Types
  ( ChatResponse (respText),
    LLMError (..),
    ProviderOpaque (..),
    ThinkingContent (..),
    ToolCall (tcArguments, tcId, tcName),
    Turn (..),
    cacheEphemeral,
    defaultMessageEncodeOptions,
    imageBase64Part,
    imageUrlPart,
    mkToolCall,
    textPart,
    thinkingPart,
    toolCallPart,
  )
import LLM.Core.ProviderUtils (normalizeSchemaOpenAI, stripBoundsAndComments)
import LLM.Core.Usage (Usage (..), mkUsage)
import LLM.Core.Utils (getToolCalls, hasToolCalls, isRetryable)
import LLM.OptionalFields (Node, Optionals)
import LLM.Providers.OpenAI (encodeTurn, parseOpenAIResponse, parseOpenAIStream, parseOpenAIUsage)
import Test.Hspec
  ( Spec,
    describe,
    expectationFailure,
    it,
    shouldBe,
    shouldContain,
    shouldSatisfy,
  )

spec :: Spec
spec = describe "OpenAI" $ do
  describe "parseOpenAIResponse" $ do
    it "parses a text response" $ do
      Right val <- eitherDecodeFileStrict' "test/fixtures/openai-text.json"
      case parseOpenAIResponse val of
        Right resp -> do
          resp.respText `shouldBe` "Hello! How can I help you today?"
          hasToolCalls resp `shouldBe` False
        Left err -> expectationFailure $ "Parse failed: " <> show err

    it "parses a tool_calls response" $ do
      Right val <- eitherDecodeFileStrict' "test/fixtures/openai-tool-use.json"
      case parseOpenAIResponse val of
        Right resp -> do
          hasToolCalls resp `shouldBe` True
          let [tc] = getToolCalls resp
          tc.tcName `shouldBe` "get_weather"
          tc.tcId `shouldBe` "call_abc123"
        Left err -> expectationFailure $ "Parse failed: " <> show err

    it "returns a ParseError when finish_reason is length and tool arguments are cut off" $ do
      let resp =
            object
              [ "choices"
                  .= [ object
                         [ "finish_reason" .= ("length" :: Text),
                           "message"
                             .= object
                               [ "role" .= ("assistant" :: Text),
                                 "tool_calls"
                                   .= [ object
                                          [ "id" .= ("call_1" :: Text),
                                            "type" .= ("function" :: Text),
                                            "function"
                                              .= object
                                                [ "name" .= ("get_weather" :: Text),
                                                  "arguments" .= ("{\"location\": \"Lon" :: Text)
                                                ]
                                          ]
                                      ]
                               ]
                         ]
                     ]
              ]
      case parseOpenAIResponse resp of
        Left (ParseError msg) -> T.unpack msg `shouldContain` "length"
        other -> expectationFailure ("expected ParseError, got " <> show other)

  describe "parseOpenAIStream errors and truncation" $ do
    it "returns a retryable 429 for a mid-stream rate limit error" $ do
      reader <- fixtureReader "test/fixtures/openai-midstream-error.sse"
      result <- parseOpenAIStream reader (\_ -> pure ())
      case result of
        Left err@(HttpError 429 body) -> do
          T.unpack body `shouldContain` "Rate limit"
          isRetryable err `shouldBe` True
        other -> expectationFailure ("expected HttpError 429, got " <> show other)

    it "returns a ParseError instead of String args when finish_reason is length" $ do
      reader <- fixtureReader "test/fixtures/openai-length-tool.sse"
      result <- parseOpenAIStream reader (\_ -> pure ())
      case result of
        Left (ParseError msg) -> T.unpack msg `shouldContain` "length"
        other -> expectationFailure ("expected ParseError, got " <> show other)

  describe "parseOpenAIStream multiple tool calls" $ do
    it "parses every tool call in a chunk, in index order" $ do
      reader <- fixtureReader "test/fixtures/openai-multi-tool-chunk.sse"
      result <- parseOpenAIStream reader (\_ -> pure ())
      case result of
        Right resp ->
          [(tc.tcId, tc.tcName, tc.tcArguments) | tc <- getToolCalls resp]
            `shouldBe` [ ("call_a", "get_weather", object ["location" .= ("London" :: Text)]),
                         ("call_b", "get_time", object ["zone" .= ("CET" :: Text)])
                       ]
        Left err -> expectationFailure ("unexpected error: " <> show err)

    it "parses Ollama-style complete tool calls without index fields" $ do
      reader <- fixtureReader "test/fixtures/ollama-multi-tool-chunk.sse"
      result <- parseOpenAIStream reader (\_ -> pure ())
      case result of
        Right resp ->
          [(tc.tcId, tc.tcArguments) | tc <- getToolCalls resp]
            `shouldBe` [ ("call_x", object ["location" .= ("London" :: Text)]),
                         ("call_y", object ["location" .= ("Paris" :: Text)])
                       ]
        Left err -> expectationFailure ("unexpected error: " <> show err)

  describe "assistant message encoding" $ do
    it "sends empty content for a turn with only foreign thinking" $ do
      let redacted =
            ThinkingContent
              { thinkingText = Nothing,
                thinkingOpaque =
                  Just
                    ProviderOpaque
                      { poProvider = "claude",
                        poModel = Just "claude-haiku-4-5",
                        poPayload = object ["type" .= ("redacted_thinking" :: Text), "data" .= ("xyz" :: Text)]
                      }
              }
          msg = head (encodeTurn defaultMessageEncodeOptions (AssistantMessage [thinkingPart redacted]))
      lookupText "content" msg `shouldBe` Just ""
      lookupKey "tool_calls" msg `shouldBe` Nothing

    it "omits content for a tool-call-only turn" $ do
      let turn = AssistantMessage [toolCallPart (mkToolCall "c1" "get_weather" (object []))]
          msg = head (encodeTurn defaultMessageEncodeOptions turn)
      lookupKey "content" msg `shouldBe` Nothing
      lookupKey "tool_calls" msg `shouldSatisfy` (/= Nothing)

  describe "image request encoding" $ do
    it "encodes URL and base64 data-URL image parts" $ do
      Right b64 <- pure $ imageBase64Part "image/png" "aGVsbG8="
      let turn =
            UserMessage
              [ textPart "What is in this image?",
                imageUrlPart "https://example.com/boardwalk.jpg",
                b64
              ]
          content = userContent (head (encodeTurn defaultMessageEncodeOptions turn))
      contentTypes content `shouldBe` ["text", "image_url", "image_url"]
      case content of
        [Object textO, Object urlO, Object b64O] -> do
          lookupText "text" (Object textO) `shouldBe` Just "What is in this image?"
          nestedText ["image_url", "url"] (Object urlO)
            `shouldBe` Just "https://example.com/boardwalk.jpg"
          nestedText ["image_url", "url"] (Object b64O)
            `shouldBe` Just "data:image/png;base64,aGVsbG8="
        _ -> expectationFailure "expected text + two image_url parts"

  describe "cache hint encoding" $ do
    it "omits cache hints from the wire while keeping content identical" $ do
      Right b64 <- pure $ imageBase64Part "image/png" "aGVsbG8="
      let plain =
            UserMessage
              [ textPart "What is in this image?",
                imageUrlPart "https://example.com/boardwalk.jpg",
                b64
              ]
          hinted =
            UserMessage
              [ cacheEphemeral (textPart "What is in this image?"),
                imageUrlPart "https://example.com/boardwalk.jpg",
                cacheEphemeral b64
              ]
          plainMsg = head (encodeTurn defaultMessageEncodeOptions plain)
          hintedMsg = head (encodeTurn defaultMessageEncodeOptions hinted)
      hintedMsg `shouldBe` plainMsg
      case userContent hintedMsg of
        parts -> do
          contentTypes parts `shouldBe` ["text", "image_url", "image_url"]
          mapM_ (\p -> lookupKey "cache_control" p `shouldBe` Nothing) parts

  describe "normalizeSchemaOpenAI optional properties (strict mode)" $ do
    let schema = strictSchema (AC.codec :: AC.JSONCodec Optionals)
        prop k = schemaAt ["properties", k] schema

    it "lists every property as required" $
      schemaAt ["required"] schema
        `shouldBe` Just
          (toJSON (["def", "list", "named", "nested", "omit", "opt", "or_null", "plain", "req_maybe", "union"] :: [Text]))

    it "makes optionalField, optionalFieldWithDefault and optionalFieldWithOmittedDefault nullable" $ do
      prop "opt" `shouldBe` Just (nullable (typed "string"))
      prop "def" `shouldBe` Just (nullable (typed "integer"))
      prop "omit" `shouldBe` Just (nullable (typed "integer"))

    it "makes an optional nested object nullable, inline or by $ref" $ do
      fmap (const ()) (schemaAt ["properties", "nested", "anyOf", "0", "properties"] schema) `shouldBe` Just ()
      schemaAt ["properties", "nested", "anyOf", "1"] schema `shouldBe` Just (typed "null")
      schemaAt ["properties", "named", "anyOf", "0", "$ref"] schema `shouldBe` Just "#/$defs/NamedInner"
      schemaAt ["properties", "named", "anyOf", "1"] schema `shouldBe` Just (typed "null")

    it "leaves nullable and required properties as they are" $ do
      let nullOrString = object ["anyOf" .= [typed "null", typed "string"]]
      prop "req_maybe" `shouldBe` Just nullOrString
      prop "or_null" `shouldBe` Just nullOrString
      prop "plain" `shouldBe` Just (typed "string")
      schemaAt ["properties", "list", "type"] schema `shouldBe` Just "array"

    it "does the same inside nested objects, array items, union branches and definitions" $ do
      let inner path = do
            schemaAt (path <> ["properties", "a"]) schema `shouldBe` Just (nullable (typed "string"))
            schemaAt (path <> ["properties", "b"]) schema `shouldBe` Just (nullable (typed "integer"))
            schemaAt (path <> ["required"]) schema `shouldBe` Just (toJSON (["a", "b"] :: [Text]))
      inner ["properties", "nested", "anyOf", "0"]
      inner ["properties", "list", "items"]
      inner ["properties", "union", "anyOf", "0", "properties", "Left"]
      inner ["$defs", "NamedInner"]

    it "turns the oneOf of an Either into anyOf" $ do
      schemaAt ["properties", "union", "oneOf"] schema `shouldBe` Nothing
      schemaAt ["properties", "union", "anyOf", "0", "required"] schema `shouldBe` Just (toJSON (["Left"] :: [Text]))
      schemaAt ["properties", "union", "anyOf", "1", "required"] schema `shouldBe` Just (toJSON (["Right"] :: [Text]))

    it "changes only schema keywords, not property names or literal values" $ do
      let oneOfSchema = object ["oneOf" .= [typed "string", typed "integer"]]
          literal = object ["type" .= ("object" :: Text), "oneOf" .= [typed "string"]]
          input =
            object
              [ "type" .= ("object" :: Text),
                "properties" .= object ["oneOf" .= oneOfSchema],
                "required" .= (["oneOf"] :: [Text]),
                "default" .= object ["oneOf" .= literal]
              ]
      normalizeSchemaOpenAI input
        `shouldBe` object
          [ "type" .= ("object" :: Text),
            "properties" .= object ["oneOf" .= object ["anyOf" .= [typed "string", typed "integer"]]],
            "required" .= (["oneOf"] :: [Text]),
            "default" .= object ["oneOf" .= literal],
            "additionalProperties" .= False
          ]

    it "moves the $defs of a named codec to the root" $ do
      schemaAt ["properties", "named", "anyOf", "0"] schema
        `shouldBe` Just (object ["$ref" .= ("#/$defs/NamedInner" :: Text)])
      fmap (const ()) (schemaAt ["$defs", "NamedInner", "properties"] schema) `shouldBe` Just ()

    it "keeps $defs that are at the root already, and a property called $defs" $ do
      let node = strictSchema (AC.codec :: AC.JSONCodec Node)
      schemaAt ["$ref"] node `shouldBe` Just "#/$defs/Node"
      fmap (const ()) (schemaAt ["$defs", "Node", "properties"] node) `shouldBe` Just ()
      let input =
            object
              [ "type" .= ("object" :: Text),
                "properties" .= object ["$defs" .= typed "string"],
                "required" .= (["$defs"] :: [Text])
              ]
      normalizeSchemaOpenAI input
        `shouldBe` object
          [ "type" .= ("object" :: Text),
            "properties" .= object ["$defs" .= typed "string"],
            "required" .= (["$defs"] :: [Text]),
            "additionalProperties" .= False
          ]

    it "makes a recursive optional property nullable" $
      schemaAt ["$defs", "Node", "properties", "next"] (strictSchema (AC.codec :: AC.JSONCodec Node))
        `shouldBe` Just (nullable (object ["$ref" .= ("#/$defs/Node" :: Text)]))

  describe "parseOpenAIUsage" $ do
    it "extracts token counts" $ do
      Right val <- eitherDecodeFileStrict' "test/fixtures/openai-text.json"
      parseOpenAIUsage val `shouldBe` Just (mkUsage 15 9)

    it "reports cached_tokens without double-counting prompt_tokens" $ do
      let val =
            object
              [ "usage"
                  .= object
                    [ "prompt_tokens" .= (2006 :: Int),
                      "completion_tokens" .= (300 :: Int),
                      "prompt_tokens_details"
                        .= object
                          [ "cached_tokens" .= (1920 :: Int),
                            "cache_write_tokens" .= (40 :: Int)
                          ]
                    ]
              ]
      parseOpenAIUsage val
        `shouldBe` Just
          Usage
            { usageInputTokens = 2006,
              usageOutputTokens = 300,
              usageCacheReadTokens = 1920,
              usageCacheCreationTokens = 40,
              usageTotalCost = 0
            }

    it "reads zero cached_tokens from recorded conversation fixtures" $ do
      Right val <- eitherDecodeFileStrict' "test/fixtures/openai-conversation-generated.json"
      case val of
        Array arr ->
          case V.toList arr of
            (Object first : _) ->
              case KM.lookup "response" first of
                Just resp ->
                  case parseOpenAIUsage resp of
                    Just u -> do
                      u.usageCacheReadTokens `shouldBe` 0
                      u.usageInputTokens `shouldSatisfy` (> 0)
                    Nothing -> expectationFailure "failed to parse usage"
                _ -> expectationFailure "missing response"
            _ -> expectationFailure "empty conversation array"
        _ -> expectationFailure "expected conversation array"

-- | The schema 'genObject' sends for a codec, as OpenAI receives it.
strictSchema :: AC.JSONCodec a -> Value
strictSchema c = normalizeSchemaOpenAI (stripBoundsAndComments (toJSON (jsonSchemaVia c)))

-- | The sub-schema at a path of keys; a number indexes an array.
schemaAt :: [Text] -> Value -> Maybe Value
schemaAt [] v = Just v
schemaAt (k : rest) (Object o) = KM.lookup (K.fromText k) o >>= schemaAt rest
schemaAt (k : rest) (Array a) = case reads (T.unpack k) of
  [(i, "")] -> a V.!? i >>= schemaAt rest
  _ -> Nothing
schemaAt _ _ = Nothing

nullable :: Value -> Value
nullable s = object ["anyOf" .= [s, object ["type" .= ("null" :: Text)]]]

typed :: Text -> Value
typed t = object ["type" .= t]

userContent :: Value -> [Value]
userContent (Object o) =
  case KM.lookup "content" o of
    Just (Array arr) -> V.toList arr
    _ -> []
userContent _ = []

contentTypes :: [Value] -> [Text]
contentTypes = mapMaybe typ
  where
    typ (Object o) = case KM.lookup "type" o of
      Just (String t) -> Just t
      _ -> Nothing
    typ _ = Nothing

lookupText :: Text -> Value -> Maybe Text
lookupText key (Object o) =
  case KM.lookup (K.fromText key) o of
    Just (String t) -> Just t
    _ -> Nothing
lookupText _ _ = Nothing

lookupKey :: Text -> Value -> Maybe Value
lookupKey key (Object o) = KM.lookup (K.fromText key) o
lookupKey _ _ = Nothing

nestedText :: [Text] -> Value -> Maybe Text
nestedText [key] v = lookupText key v
nestedText (key : rest) (Object o) =
  KM.lookup (K.fromText key) o >>= nestedText rest
nestedText _ _ = Nothing

-- | Yield the whole fixture once, then EOF.
fixtureReader :: FilePath -> IO (IO BS.ByteString)
fixtureReader path = do
  bs <- BS.readFile path
  ref <- newIORef (Just bs)
  pure $ do
    m <- readIORef ref
    writeIORef ref Nothing
    pure (fromMaybe BS.empty m)