diff --git a/baikai-openai.cabal b/baikai-openai.cabal --- a/baikai-openai.cabal +++ b/baikai-openai.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: baikai-openai -version: 0.2.0.0 +version: 0.3.0.0 synopsis: OpenAI providers for the baikai abstraction description: Wraps the openai Haskell package as a Baikai Provider for OpenAI's Chat Completions API. @@ -33,12 +33,16 @@ exposed-modules: Baikai.Provider.OpenAI.Api Baikai.Provider.OpenAI.Cli - Baikai.Provider.OpenAI.ErrorClass Baikai.Provider.OpenAI.Interactive + Baikai.Provider.OpenAI.Internal.ErrorClass + Baikai.Provider.OpenAI.Internal.Request + Baikai.Provider.OpenAI.Shape + Baikai.Provider.OpenAI.Sse + Baikai.Provider.OpenAI.Transport build-depends: , aeson - , baikai ^>=0.2.0 + , baikai ^>=0.3.0 , base >=4.20 && <5 , base64-bytestring , bytestring @@ -63,16 +67,25 @@ type: exitcode-stdio-1.0 hs-source-dirs: test main-is: Main.hs - other-modules: ErrorClassSpec + other-modules: + ErrorClassSpec + ReasoningSpec + ShapeSpec + SseSpec + TransportSpec + build-depends: , aeson - , baikai ^>=0.2.0 + , baikai ^>=0.3.0 , baikai-openai , base >=4.20 && <5 , bytestring , case-insensitive , containers + , directory + , filepath , generic-lens + , http-client , http-types , lens ^>=5.3 , openai @@ -81,4 +94,5 @@ , tasty , tasty-hunit , text ^>=2.1 + , time , vector diff --git a/src/Baikai/Provider/OpenAI/Api.hs b/src/Baikai/Provider/OpenAI/Api.hs --- a/src/Baikai/Provider/OpenAI/Api.hs +++ b/src/Baikai/Provider/OpenAI/Api.hs @@ -9,19 +9,20 @@ -- through this handler. -- -- The handler resolves 'Baikai.Options.apiKey' when present, falling --- back to the @OPENAI_API_KEY@ env var via --- 'Baikai.Auth.resolveApiKey'. +-- back to the host-specific env var from +-- 'Baikai.Auth.defaultApiKeyEnvForBaseUrl'. Unknown hosts require an +-- explicit key source. -- -- EP-3 promotes streaming to the primary entry point. The handler -- exposes a 'streamly' 'Stream' of 'AssistantMessageEvent' values --- bridged from the upstream SDK's raw --- 'createChatCompletionStream' callback. We deliberately bypass --- the typed variant because the typed @ChatCompletionChunk@ --- requires @id@ + @function.name@ on every tool-call delta — fields --- that OpenAI omits on partial-argument continuation chunks — so a --- tool-using stream fails to parse end-to-end. Parsing the raw --- 'Aeson.Value' chunk manually lets us tolerate missing fields the --- way the upstream wire protocol intends. +-- bridged from a local SSE transport. Requests start as the SDK's +-- typed 'OpenAI.V1.Chat.Completions.CreateChatCompletion' value, then +-- 'Baikai.Provider.OpenAI.Shape.streamRequestBody' rewrites the raw +-- JSON body for OpenAI-compatible host quirks before +-- 'Baikai.Provider.OpenAI.Sse.openaiSseStreamValueWithHeaders' sends +-- it with cached transport settings and caller headers. Streaming +-- responses are parsed from raw 'Aeson.Value' chunks so partial +-- tool-call deltas may omit fields such as @id@ and @function.name@. -- -- The synchronous 'complete' field is derived via -- 'streamingComplete', so callers that drain the stream get the @@ -29,33 +30,49 @@ module Baikai.Provider.OpenAI.Api ( register, registerWithRegistry, + openaiChatProvider, openaiChatStream, - mapRequest, + RawChunk (..), + RawToolDelta (..), + parseChunk, + TagScanState (..), + _TagScanState, + scanThinkTags, + Assembler (..), + emptyAssembler, + translate, + closeOpenStream, + + -- * Usage mapping + + -- Exposed for tests; may move behind an .Internal namespace in a later plan. + RawUsage (..), + parseUsage, + rawUsageToUsage, ) where import Baikai.Api (Api (..)) -import Baikai.Auth qualified as Auth -import Baikai.Compat - ( OpenAICompletionsCompat (..), - ThinkingFormat (..), - ) +import Baikai.Compat (OpenAICompletionsCompat (requiresThinkingAsText)) import Baikai.Content qualified as Content import Baikai.Context (Context (..)) -import Baikai.Cost (_Cost) +import Baikai.Cost (zeroCost) import Baikai.Cost.Pricing qualified as Pricing -import Baikai.Error (BaikaiError, invalidRequest) +import Baikai.Error (BaikaiError, invalidRequest, providerError) import Baikai.Message qualified as Msg import Baikai.Model (Model, openaiCompletionsCompatFor) import Baikai.Options (Options (..)) -import Baikai.Provider.OpenAI.ErrorClass (classifyErrorText, classifyException) +import Baikai.Provider.OpenAI.Internal.ErrorClass (classifyException) +import Baikai.Provider.OpenAI.Internal.Request (mapRequest) +import Baikai.Provider.OpenAI.Shape (streamRequestBody) +import Baikai.Provider.OpenAI.Sse (openaiSseStreamValueWithHeaders) +import Baikai.Provider.OpenAI.Transport qualified as Transport import Baikai.Provider.Registry ( ApiProvider (..), ProviderRegistry, - globalProviderRegistry, + registerApiProvider, registerApiProviderWith, ) -import Baikai.ResponseFormat (ResponseFormat (..)) import Baikai.StopReason qualified as Stop import Baikai.Stream (streamingComplete) import Baikai.Stream.Event @@ -64,30 +81,28 @@ DeltaPayload (..), IndexPayload (..), StartPayload (..), + ThinkingEndPayload (..), ToolCallEndPayload (..), doneTerminal, errorTerminal, ) -import Baikai.ThinkingLevel - ( ThinkingLevel (..), - ) -import Baikai.Tool qualified as Tool import Baikai.Usage qualified as Usage +import Control.Applicative ((<|>)) import Control.Concurrent (forkIO) import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan) -import Control.Exception (SomeException, try) +import Control.Exception (SomeAsyncException (..), SomeException, fromException, throwIO, try) import Control.Lens ((%~), (&), (.~), (^.)) import Data.Aeson (Value (..), (.:?)) import Data.Aeson qualified as Aeson import Data.Aeson.Key qualified as AesonKey import Data.Aeson.KeyMap qualified as KeyMap import Data.Aeson.Types qualified as Aeson -import Data.ByteString.Base64 qualified as Base64 -import Data.ByteString.Lazy qualified as BSL import Data.Generics.Labels () import Data.IORef (IORef, newIORef, readIORef, writeIORef) import Data.IntMap.Strict (IntMap) import Data.IntMap.Strict qualified as IntMap +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map import Data.Maybe (fromMaybe) import Data.Text (Text) import Data.Text qualified as Text @@ -96,30 +111,34 @@ import Data.Vector (Vector) import Data.Vector qualified as Vector import GHC.Generics (Generic) +import Network.HTTP.Types.Header (RequestHeaders) import Numeric.Natural (Natural) -import OpenAI.V1 qualified as OpenAI -import OpenAI.V1.Chat.Completions qualified as Chat -import OpenAI.V1.Models qualified as OpenAIModels -import OpenAI.V1.ResponseFormat qualified as RF -import OpenAI.V1.Tool qualified as OpenAITool -import OpenAI.V1.ToolCall qualified as ToolCall +import Servant.Client qualified as Client import Streamly.Data.Stream (Stream) import Streamly.Data.Stream qualified as Stream -- | Install the OpenAI Chat Completions handler into the registry. register :: IO () -register = registerWithRegistry globalProviderRegistry +register = registerApiProvider openaiChatProvider +-- | First-class OpenAI Chat Completions provider value. Use with +-- 'registerApiProviderWith' or 'newProviderRegistryFrom' for explicit +-- registries. +openaiChatProvider :: ApiProvider +openaiChatProvider = + ApiProvider + { apiTag = OpenAIChatCompletions, + stream = openaiChatStream, + complete = streamingComplete openaiChatStream + } + -- | Install the OpenAI Chat Completions handler into an explicit registry. registerWithRegistry :: ProviderRegistry -> IO () registerWithRegistry reg = registerApiProviderWith reg - ApiProvider - { apiTag = OpenAIChatCompletions, - stream = openaiChatStream, - complete = streamingComplete openaiChatStream - } + openaiChatProvider +{-# DEPRECATED registerWithRegistry "use registerApiProviderWith reg openaiChatProvider" #-} -- | Streaming producer for the OpenAI Chat Completions API. -- @@ -134,23 +153,22 @@ Model -> Context -> Options -> Stream IO AssistantMessageEvent openaiChatStream m ctx opts = Stream.concatEffect $ do - setup <- prepareCall m ctx opts + setupResult <- trySync (prepareCall m ctx opts) + let setup = either (Left . exceptionToError) id setupResult case setup of - Left err -> pure (Stream.fromEffect (immediateError err)) + Left err -> Stream.fromList <$> immediateError err Right call -> do - ch <- newChan :: IO (Chan (Maybe RawChunk)) + ch <- newChan :: IO (Chan (Maybe (Either BaikaiError RawChunk))) tref <- newIORef False - eref <- newIORef Nothing - _ <- forkIO (worker call ch eref) + _ <- forkIO (worker call ch) startTime <- getCurrentTime let initialState = ProducerState { chan = ch, - pending = [EventStart StartPayload {partial = skeletonStart m startTime}], + pending = [EventStart StartPayload {partial = skeletonStart m startTime, responseId = Nothing}], assembler = emptyAssembler m startTime, finished = False, - terminalRef = tref, - errInfoRef = eref + terminalRef = tref } pure (Stream.unfoldrM step initialState) @@ -159,44 +177,42 @@ Msg.AssistantMessage Msg.AssistantPayload { Msg.content = Vector.empty, - Msg.usage = Usage._Usage, + Msg.usage = Usage.zeroUsage, Msg.stopReason = Stop.Stop, Msg.errorMessage = Nothing, - Msg.timestamp = start + Msg.timestamp = Just start } -- | Per-call prepared values. data OpenAICall = OpenAICall - { methods :: !OpenAI.Methods, - request :: !Chat.CreateChatCompletion + { clientEnv :: !Client.ClientEnv, + requestHeaders :: !RequestHeaders, + timeoutMs :: !(Maybe Int), + requestBody :: !Aeson.Value } deriving stock (Generic) -prepareCall :: Model -> Context -> Options -> IO (Either Text OpenAICall) +prepareCall :: Model -> Context -> Options -> IO (Either BaikaiError OpenAICall) prepareCall m ctx opts = case mapRequest m ctx opts of - Left e -> pure (Left e) + Left e -> pure (Left (invalidRequest e)) Right req -> do - key <- resolveKey opts let url = case m ^. #baseUrl of "" -> "https://api.openai.com" u -> u - env <- OpenAI.getClientEnv url - let mtds = OpenAI.makeMethods env key Nothing Nothing - req' = - req - { Chat.stream = Just True, - Chat.stream_options = - Just - Chat._ChatCompletionStreamOptions - { Chat.include_usage = Just True - } + key <- Transport.resolveKey url opts + env <- Transport.getClientEnvCached url + let compat = openaiCompletionsCompatFor m + body = streamRequestBody compat opts req + headers = Transport.requestHeaders key m opts + pure + ( Right + OpenAICall + { clientEnv = env, + requestHeaders = headers, + timeoutMs = opts ^. #timeoutMs, + requestBody = body } - pure (Right OpenAICall {methods = mtds, request = req'}) - -resolveKey :: Options -> IO Text -resolveKey opts = case opts ^. #apiKey of - Just source -> Auth.resolveApiKey source - Nothing -> Auth.resolveApiKey (Auth.ApiKeyEnv "OPENAI_API_KEY") + ) -- | A loose summary of one streamed chunk. The raw 'Aeson.Value' is -- pre-parsed into the fields we care about; unknown fields are @@ -204,15 +220,15 @@ -- tool-call deltas). data RawChunk = RawChunk { contentDelta :: !(Maybe Text), + reasoningDelta :: !(Maybe Text), finishReason :: !(Maybe Text), toolDeltas :: ![RawToolDelta], - usage :: !(Maybe RawUsage), - error :: !(Maybe Text) + usage :: !(Maybe RawUsage) } deriving stock (Show, Generic) data RawToolDelta = RawToolDelta - { index :: !Int, + { index :: !(Maybe Int), id_ :: !(Maybe Text), name :: !(Maybe Text), args :: !(Maybe Text) @@ -228,34 +244,22 @@ deriving stock (Show, Generic) worker :: - OpenAICall -> Chan (Maybe RawChunk) -> IORef (Maybe BaikaiError) -> IO () -worker call ch errInfoRef = do - let OpenAI.Methods {OpenAI.createChatCompletionStream = stream'} = call ^. #methods + OpenAICall -> Chan (Maybe (Either BaikaiError RawChunk)) -> IO () +worker call ch = do r <- - try @SomeException $ - stream' (call ^. #request) $ \case - Left errText -> writeChan ch (Just (errorChunk errText)) - Right val -> case parseChunk val of - Left err -> writeChan ch (Just (errorChunk (Text.pack err))) - Right chunk -> writeChan ch (Just chunk) + trySync $ + Transport.runWithTimeout (call ^. #timeoutMs) $ + openaiSseStreamValueWithHeaders (call ^. #clientEnv) (call ^. #requestHeaders) (call ^. #requestBody) $ \case + Left be -> writeChan ch (Just (Left be)) + Right val -> case parseChunk val of + Left err -> writeChan ch (Just (Left (providerError (Text.pack err)))) + Right chunk -> writeChan ch (Just (Right chunk)) case r of - Right () -> pure () - -- An HTTP-level exception carries an HTTP status; classify it and - -- stash the structured error for the end-of-stream recovery path - -- (see 'step' / 'closeOpenStream'). No terminal chunk is written. - Left e -> writeIORef errInfoRef (Just (classifyException e)) + Right Nothing -> pure () + Right (Just be) -> writeChan ch (Just (Left be)) + Left e -> writeChan ch (Just (Left (exceptionToError e))) writeChan ch Nothing -errorChunk :: Text -> RawChunk -errorChunk t = - RawChunk - { contentDelta = Nothing, - finishReason = Nothing, - toolDeltas = [], - usage = Nothing, - error = Just t - } - -- | Aeson parser tolerant of partial tool-call fields. parseChunk :: Value -> Either String RawChunk parseChunk = Aeson.parseEither $ Aeson.withObject "ChatCompletionChunk" $ \o -> do @@ -268,19 +272,20 @@ Aeson.Object obj -> Just obj _ -> Nothing _ -> Nothing - (contentDelta, finishR, toolDeltas) <- case firstChoice of - Nothing -> pure (Nothing, Nothing, []) + (contentDelta, reasoningDelta, finishR, toolDeltas) <- case firstChoice of + Nothing -> pure (Nothing, Nothing, Nothing, []) Just ch -> do finish <- ch .:? "finish_reason" delta <- ch .:? "delta" case delta of - Nothing -> pure (Nothing, finish, []) + Nothing -> parseMessageObject ch finish Just (Aeson.Object dObj) -> do cd <- dObj .:? "content" + let rd = reasoningText dObj tc <- dObj .:? "tool_calls" let tds = parseToolCallDeltas tc - pure (cd, finish, tds) - _ -> pure (Nothing, finish, []) + pure (cd, rd, finish, tds) + _ -> parseMessageObject ch finish usageM <- o .:? "usage" let ru = case usageM of Just (Aeson.Object uObj) -> parseUsage uObj @@ -288,12 +293,28 @@ pure RawChunk { contentDelta = contentDelta, + reasoningDelta = reasoningDelta, finishReason = finishR, toolDeltas = toolDeltas, - usage = ru, - error = Nothing + usage = ru } +parseMessageObject :: + Aeson.Object -> + Maybe Text -> + Aeson.Parser (Maybe Text, Maybe Text, Maybe Text, [RawToolDelta]) +parseMessageObject ch finish = do + msg <- ch .:? "message" + case msg of + Just (Aeson.Object mObj) -> do + cd <- mObj .:? "content" + pure (cd, reasoningText mObj, finish, []) + _ -> pure (Nothing, Nothing, finish, []) + +reasoningText :: Aeson.Object -> Maybe Text +reasoningText obj = + lookupText "reasoning_content" obj <|> lookupText "reasoning" obj + parseToolCallDeltas :: Maybe Value -> [RawToolDelta] parseToolCallDeltas = \case Just (Aeson.Array v) -> Vector.toList (Vector.mapMaybe oneDelta v) @@ -310,7 +331,7 @@ getArgs = funcObj >>= lookupText "arguments" in Just RawToolDelta - { index = maybe 0 fromInt (lookupField "index" o), + { index = fromInt <$> lookupField "index" o, id_ = lookupText "id" o, name = getName, args = getArgs @@ -366,14 +387,11 @@ -- ============================================================ data ProducerState = ProducerState - { chan :: !(Chan (Maybe RawChunk)), + { chan :: !(Chan (Maybe (Either BaikaiError RawChunk))), pending :: ![AssistantMessageEvent], assembler :: !Assembler, finished :: !Bool, - terminalRef :: !(IORef Bool), - -- | Set by the worker when an HTTP-level exception is caught, so the - -- end-of-stream recovery path can surface a categorised error. - errInfoRef :: !(IORef (Maybe BaikaiError)) + terminalRef :: !(IORef Bool) } deriving stock (Generic) @@ -399,8 +417,7 @@ then pure Nothing else do now <- getCurrentTime - mErr <- readIORef (s ^. #errInfoRef) - let (events, ass') = closeOpenStream now mErr (s ^. #assembler) + let (events, ass') = closeOpenStream now Nothing (s ^. #assembler) case events of [] -> pure Nothing (e : rest) -> do @@ -446,6 +463,99 @@ -- Translation -- ============================================================ +data TagMode + = TagVisible + | TagReasoning + deriving stock (Eq, Show, Generic) + +-- | Incremental scanner state for hosts that stream reasoning in +-- assistant text using @@ or @@ tags. +data TagScanState = TagScanState + { tagMode :: !TagMode, + tagPending :: !Text + } + deriving stock (Eq, Show, Generic) + +_TagScanState :: TagScanState +_TagScanState = + TagScanState + { tagMode = TagVisible, + tagPending = Text.empty + } + +-- | Split one text delta into reasoning fragments ('Left') and +-- visible text fragments ('Right'), preserving partial tag prefixes +-- across chunk boundaries. +scanThinkTags :: TagScanState -> Text -> (TagScanState, [Either Text Text]) +scanThinkTags st input = + let (mode', pending', parts) = go (tagMode st) (tagPending st <> input) [] + in (TagScanState {tagMode = mode', tagPending = pending'}, parts) + where + go mode txt acc = + case findTag mode txt of + Just (before, after, nextMode) -> + go nextMode after (appendPart mode before acc) + Nothing -> + let (emitNow, pending) = splitPending mode txt + in (mode, pending, appendPart mode emitNow acc) + + appendPart _ "" acc = acc + appendPart TagVisible t acc = acc <> [Right t] + appendPart TagReasoning t acc = acc <> [Left t] + +findTag :: TagMode -> Text -> Maybe (Text, Text, TagMode) +findTag mode txt = + case earliest markers of + Nothing -> Nothing + Just (idx, marker) -> + Just + ( Text.take idx txt, + Text.drop (idx + Text.length marker) txt, + nextMode + ) + where + (markers, nextMode) = case mode of + TagVisible -> (openingTags, TagReasoning) + TagReasoning -> (closingTags, TagVisible) + earliest = + foldr + ( \marker best -> + case Text.breakOn marker txt of + (_, "") -> best + (before, _) -> + let candidate = (Text.length before, marker) + in case best of + Nothing -> Just candidate + Just (oldIdx, _) | Text.length before < oldIdx -> Just candidate + _ -> best + ) + Nothing + +splitPending :: TagMode -> Text -> (Text, Text) +splitPending mode txt = + let suffix = longestTagPrefix (case mode of TagVisible -> openingTags; TagReasoning -> closingTags) txt + in (Text.dropEnd (Text.length suffix) txt, suffix) + +longestTagPrefix :: [Text] -> Text -> Text +longestTagPrefix markers txt = + foldr longer Text.empty candidates + where + candidates = + [ suffix + | n <- [1 .. Text.length txt], + let suffix = Text.takeEnd n txt, + any (suffix `Text.isPrefixOf`) markers + ] + longer a b + | Text.length a > Text.length b = a + | otherwise = b + +openingTags :: [Text] +openingTags = ["", ""] + +closingTags :: [Text] +closingTags = ["", ""] + -- | Translation state across one streaming call. data Assembler = Assembler { model :: !Model, @@ -455,9 +565,14 @@ textOpen :: !(Maybe Int), textAccum :: !Text, textEverOpened :: !Bool, + reasoningOpen :: !(Maybe Int), + reasoningAccum :: !Text, + tagScanState :: !TagScanState, -- | Maps OpenAI's per-call tool-call index to baikai's -- 'contentIndex'. toolIndexMap :: !(IntMap Int), + toolIdMap :: !(Map Text Int), + lastToolIdx :: !(Maybe Int), -- | baikai contentIndex → (id, name). toolMeta :: !(IntMap (Text, Text)), -- | baikai contentIndex → accumulated arguments JSON. @@ -471,7 +586,8 @@ -- the post-@finish_reason@ usage chunk (when @include_usage@ is -- enabled) has a chance to land. finishSeen :: !Bool, - errorMsg :: !(Maybe Text) + pendingError :: !(Maybe BaikaiError), + finishNote :: !(Maybe Text) } deriving stock (Generic) @@ -483,58 +599,108 @@ textOpen = Nothing, textAccum = Text.empty, textEverOpened = False, + reasoningOpen = Nothing, + reasoningAccum = Text.empty, + tagScanState = _TagScanState, toolIndexMap = IntMap.empty, + toolIdMap = Map.empty, + lastToolIdx = Nothing, toolMeta = IntMap.empty, toolArgs = IntMap.empty, closed = IntMap.empty, nextContentIndex = 0, - usage = Usage._Usage, + usage = Usage.zeroUsage, stopReason = Stop.Stop, finishSeen = False, - errorMsg = Nothing + pendingError = Nothing, + finishNote = Nothing } translate :: - RawChunk -> + Either BaikaiError RawChunk -> Assembler -> UTCTime -> ([AssistantMessageEvent], Assembler) translate chunk ass now - | Just errMsg <- chunk ^. #error = - let msg = finalMessage ass now (Just errMsg) Stop.ErrorReason - in ( [EventError (errorTerminal Stop.ErrorReason msg (classifyErrorText errMsg))], - ass & #errorMsg .~ Just errMsg - ) - | otherwise = - let -- 1. Apply content delta (open text block if needed). - (textEvents, ass1) = applyContentDelta (chunk ^. #contentDelta) ass - -- 2. Apply tool-call deltas. - (toolEvents, ass2) = applyToolDeltas (chunk ^. #toolDeltas) ass1 - -- 3. Apply usage chunk if present. - ass3 = applyUsage (chunk ^. #usage) ass2 - -- 4. If finish_reason is set, close any open text/tool + | Left be <- chunk = + let msg = finalMessage ass now (Just (be ^. #message)) Stop.ErrorReason + in ([EventError (errorTerminal Nothing Stop.ErrorReason msg be)], ass) + | Right raw <- chunk = + let -- 1. Apply field-based reasoning delta. + (reasoningEvents, ass1) = applyReasoningDelta (raw ^. #reasoningDelta) ass + -- 2. Apply content delta (open text block if needed). + (textEvents, ass2) = applyContentDelta (raw ^. #contentDelta) ass1 + -- 3. Apply tool-call deltas. + (toolEvents, ass3) = applyToolDeltas (raw ^. #toolDeltas) ass2 + -- 4. Apply usage chunk if present. + ass4 = applyUsage (raw ^. #usage) ass3 + -- 5. If finish_reason is set, close any open text/tool -- blocks and stash the reason. EventDone is deferred -- to channel close so the post-finish_reason usage -- chunk has a chance to land. - (closeEvents, ass4) = case chunk ^. #finishReason of - Just fr -> closeOnFinish fr ass3 - Nothing -> ([], ass3) - in (textEvents <> toolEvents <> closeEvents, ass4) + (closeEvents, ass5) = case raw ^. #finishReason of + Just fr -> closeOnFinish fr ass4 + Nothing -> ([], ass4) + in (reasoningEvents <> textEvents <> toolEvents <> closeEvents, ass5) +applyReasoningDelta :: + Maybe Text -> Assembler -> ([AssistantMessageEvent], Assembler) +applyReasoningDelta Nothing ass = ([], ass) +applyReasoningDelta (Just "") ass = ([], ass) +applyReasoningDelta (Just d) ass = + case ass ^. #reasoningOpen of + Just i -> + ( [ThinkingDelta DeltaPayload {contentIndex = i, delta = d}], + ass & #reasoningAccum %~ (<> d) + ) + Nothing -> + let i = ass ^. #nextContentIndex + in ( [ThinkingStart IndexPayload {contentIndex = i}, ThinkingDelta DeltaPayload {contentIndex = i, delta = d}], + ass + & #reasoningOpen .~ Just i + & #reasoningAccum .~ d + & #nextContentIndex .~ (i + 1) + ) + applyContentDelta :: Maybe Text -> Assembler -> ([AssistantMessageEvent], Assembler) applyContentDelta Nothing ass = ([], ass) applyContentDelta (Just "") ass = ([], ass) applyContentDelta (Just d) ass = + if requiresThinkingAsText (openaiCompletionsCompatFor (ass ^. #model)) + then + let (tagState', parts) = scanThinkTags (ass ^. #tagScanState) d + (events, ass') = foldl' applyTaggedPart ([], ass & #tagScanState .~ tagState') parts + in (events, ass') + else applyVisibleTextDelta d ass + +applyTaggedPart :: + ([AssistantMessageEvent], Assembler) -> + Either Text Text -> + ([AssistantMessageEvent], Assembler) +applyTaggedPart (acc, ass) = \case + Left reasoning -> + let (events, ass') = applyReasoningDelta (Just reasoning) ass + in (acc <> events, ass') + Right visible -> + let (events, ass') = applyVisibleTextDelta visible ass + in (acc <> events, ass') + +applyVisibleTextDelta :: + Text -> Assembler -> ([AssistantMessageEvent], Assembler) +applyVisibleTextDelta "" ass = ([], ass) +applyVisibleTextDelta d ass = case ass ^. #textOpen of Just i -> - ( [TextDelta DeltaPayload {contentIndex = i, delta = d}], - ass & #textAccum %~ (<> d) - ) + let (reasoningEvents, ass1) = closeOpenReasoning ass + in ( reasoningEvents <> [TextDelta DeltaPayload {contentIndex = i, delta = d}], + ass1 & #textAccum %~ (<> d) + ) Nothing -> - let i = ass ^. #nextContentIndex - in ( [TextStart IndexPayload {contentIndex = i}, TextDelta DeltaPayload {contentIndex = i, delta = d}], - ass + let (reasoningEvents, ass1) = closeOpenReasoning ass + i = ass1 ^. #nextContentIndex + in ( reasoningEvents <> [TextStart IndexPayload {contentIndex = i}, TextDelta DeltaPayload {contentIndex = i, delta = d}], + ass1 & #textOpen .~ Just i & #textAccum .~ d & #textEverOpened .~ True @@ -543,7 +709,11 @@ applyToolDeltas :: [RawToolDelta] -> Assembler -> ([AssistantMessageEvent], Assembler) -applyToolDeltas deltas ass = foldl' apply ([], ass) deltas +applyToolDeltas [] ass = ([], ass) +applyToolDeltas deltas ass = + let (reasoningEvents, ass0) = closeOpenReasoning ass + (toolEvents, ass') = foldl' apply ([], ass0) deltas + in (reasoningEvents <> toolEvents, ass') where apply (acc, a) d = let (events, a') = applyOneToolDelta d a @@ -552,14 +722,29 @@ applyOneToolDelta :: RawToolDelta -> Assembler -> ([AssistantMessageEvent], Assembler) applyOneToolDelta d ass = - let openaiIdx = d ^. #index - (baikaiIdx, ass1, opened) = case IntMap.lookup openaiIdx (ass ^. #toolIndexMap) of - Just i -> (i, ass, False) + let mOpenaiIdx = d ^. #index + mToolId = d ^. #id_ + byIndex = mOpenaiIdx >>= \idx -> IntMap.lookup idx (ass ^. #toolIndexMap) + byId = mToolId >>= \tid -> Map.lookup tid (ass ^. #toolIdMap) + byLast = + case (mOpenaiIdx, mToolId) of + (Nothing, Nothing) -> ass ^. #lastToolIdx + _ -> Nothing + (baikaiIdx, ass1, opened) = case byIndex <|> byId <|> byLast of + Just i -> + ( i, + ass + & #toolIdMap %~ maybe id (`Map.insert` i) mToolId + & #lastToolIdx .~ Just i, + False + ) Nothing -> let i = ass ^. #nextContentIndex ass' = ass - & #toolIndexMap %~ IntMap.insert openaiIdx i + & #toolIndexMap %~ maybe id (`IntMap.insert` i) mOpenaiIdx + & #toolIdMap %~ maybe id (`Map.insert` i) mToolId + & #lastToolIdx .~ Just i & #toolMeta %~ IntMap.insert i ("", "") & #toolArgs %~ IntMap.insert i Text.empty & #nextContentIndex .~ (i + 1) @@ -585,32 +770,87 @@ else events0 <> [ToolCallDelta DeltaPayload {contentIndex = baikaiIdx, delta = argsDelta}] in (events1, ass3) +-- | Normalize OpenAI's inclusive usage counters into baikai's +-- disjoint 'Usage.Usage' convention. OpenAI's @prompt_tokens@ +-- includes @prompt_tokens_details.cached_tokens@, so the cached count +-- is subtracted out of 'Usage.inputTokens'. The subtraction is clamped +-- at zero because 'Natural' subtraction throws on underflow and because +-- OpenAI-compatible hosts can report inconsistent counters. +-- 'Usage.totalTokens' is recomputed from the normalized parts; +-- 'Usage.reasoningTokens' is a subset of 'Usage.outputTokens' and is +-- not added to the total. OpenAI does not bill cache writes, so +-- 'Usage.cacheWriteTokens' is always zero. +rawUsageToUsage :: RawUsage -> Usage.Usage +rawUsageToUsage u = + let prompt = u ^. #inputTokens + cached = u ^. #cacheReadTokens + out = u ^. #outputTokens + nonCached = if cached >= prompt then 0 else prompt - cached + in Usage.Usage + { Usage.inputTokens = nonCached, + Usage.outputTokens = out, + Usage.cacheReadTokens = cached, + Usage.cacheWriteTokens = 0, + Usage.reasoningTokens = u ^. #reasoningTokens, + Usage.totalTokens = nonCached + out + cached, + Usage.cost = zeroCost + } + applyUsage :: Maybe RawUsage -> Assembler -> Assembler applyUsage Nothing ass = ass -applyUsage (Just u) ass = - let usage' = - Usage.Usage - { Usage.inputTokens = u ^. #inputTokens, - Usage.outputTokens = u ^. #outputTokens, - Usage.cacheReadTokens = u ^. #cacheReadTokens, - Usage.cacheWriteTokens = 0, - Usage.reasoningTokens = u ^. #reasoningTokens, - Usage.totalTokens = (u ^. #inputTokens) + (u ^. #outputTokens) + (u ^. #cacheReadTokens), - Usage.cost = _Cost - } - in ass & #usage .~ usage' +applyUsage (Just u) ass = ass & #usage .~ rawUsageToUsage u -- | Close all open content blocks and stash the resolved stop -- reason; defer 'EventDone' to channel close. closeOnFinish :: Text -> Assembler -> ([AssistantMessageEvent], Assembler) closeOnFinish finishReason ass = - let (closeText, ass1) = closeOpenText ass - (closeTools, ass2) = closeOpenTools ass1 - reason = mapFinishReason finishReason - ass3 = ass2 & #stopReason .~ reason & #finishSeen .~ True - in (closeText <> closeTools, ass3) + let (tagEvents, ass0) = flushTagScanPending ass + (closeReasoning, ass1) = closeOpenReasoning ass0 + (closeText, ass2) = closeOpenText ass1 + (closeTools, ass3) = closeOpenTools ass2 + (reason, note) = mapFinishReason finishReason + pending = + if reason == Stop.ErrorReason + then Just (providerError ("provider stopped the response: finish_reason=" <> finishReason)) + else Nothing + ass4 = + ass3 + & #stopReason .~ reason + & #finishSeen .~ True + & #pendingError .~ pending + & #finishNote .~ note + in (tagEvents <> closeReasoning <> closeText <> closeTools, ass4) +flushTagScanPending :: Assembler -> ([AssistantMessageEvent], Assembler) +flushTagScanPending ass = + let st = ass ^. #tagScanState + pending = tagPending st + ass0 = ass & #tagScanState .~ st {tagPending = Text.empty} + in case (tagMode st, pending) of + (_, "") -> ([], ass0) + (TagVisible, t) -> applyVisibleTextDelta t ass0 + (TagReasoning, t) -> applyReasoningDelta (Just t) ass0 + +closeOpenReasoning :: Assembler -> ([AssistantMessageEvent], Assembler) +closeOpenReasoning ass = case ass ^. #reasoningOpen of + Nothing -> ([], ass) + Just i -> + let body = ass ^. #reasoningAccum + thinkingContent = + Content.ThinkingContent + { Content.thinking = body, + Content.signature = Nothing, + Content.redacted = False + } + block = Content.AssistantThinking thinkingContent + in ( [ThinkingEnd ThinkingEndPayload {contentIndex = i, content = thinkingContent}], + ass + & #reasoningOpen .~ Nothing + & #reasoningAccum .~ Text.empty + & #closed %~ IntMap.insert i block + ) + -- | Close the open text block, if any, by emitting a 'TextEnd' and -- storing the assembled content in 'closed'. closeOpenText :: Assembler -> ([AssistantMessageEvent], Assembler) @@ -652,31 +892,43 @@ & #closed %~ IntMap.insert i block & #toolArgs %~ IntMap.delete i & #toolMeta %~ IntMap.delete i + & #toolIdMap %~ (if Text.null tid then id else Map.delete tid) + & #lastToolIdx .~ Nothing ) closeOpenStream :: UTCTime -> Maybe BaikaiError -> Assembler -> ([AssistantMessageEvent], Assembler) closeOpenStream now mErr ass | ass ^. #finishSeen = - -- Channel closed cleanly after finish_reason. Emit - -- EventDone with the accumulated content + usage. + -- Channel closed cleanly after finish_reason. let reason = ass ^. #stopReason - msg = finalMessage ass now Nothing reason - in ([EventDone (doneTerminal reason msg)], ass) + terminalErr = + (ass ^. #pendingError) + <|> if reason == Stop.ErrorReason + then Just (providerError "provider stopped the response with an error finish_reason") + else Nothing + msg = finalMessage ass now (fmap (^. #message) terminalErr) reason + terminalEvent = case terminalErr of + Just be -> EventError (errorTerminal Nothing reason msg be) + Nothing -> EventDone (doneTerminal Nothing reason msg) + in ([terminalEvent], ass) | otherwise = -- Channel closed without a finish_reason. Force-close any -- still-open blocks and emit EventError. When the worker stored a -- classified HTTP error ('Just be'), surface it structurally; -- otherwise report the unexpected end of stream. - let (closeText, ass1) = closeOpenText ass - (closeTools, ass2) = closeOpenTools ass1 + let (tagEvents, ass0) = flushTagScanPending ass + (closeReasoning, ass1) = closeOpenReasoning ass0 + (closeText, ass2) = closeOpenText ass1 + (closeTools, ass3) = closeOpenTools ass2 reason = Stop.ErrorReason errText = case mErr of Just be -> be ^. #message Nothing -> "openai stream ended without finish_reason" - msg = finalMessage ass2 now (Just errText) reason - errEv = EventError (errorTerminal reason msg mErr) - in (closeText <> closeTools <> [errEv], ass2) + msg = finalMessage ass3 now (Just errText) reason + errInfo = fromMaybe (providerError errText) mErr + errEv = EventError (errorTerminal Nothing reason msg errInfo) + in (tagEvents <> closeReasoning <> closeText <> closeTools <> [errEv], ass3) finalMessage :: Assembler -> UTCTime -> Maybe Text -> Stop.StopReason -> Msg.Message @@ -691,257 +943,51 @@ { Msg.content = blocks, Msg.usage = usage', Msg.stopReason = sr, - Msg.errorMessage = errMsg, - Msg.timestamp = now + Msg.errorMessage = errMsg <|> (ass ^. #finishNote), + Msg.timestamp = Just now } blocksInOrder :: Assembler -> Vector Content.AssistantContent blocksInOrder ass = Vector.fromList (IntMap.elems (ass ^. #closed)) --- | Immediate single-error stream emitted when the request itself --- could not be built (e.g. message mapping failed). -immediateError :: Text -> IO AssistantMessageEvent -immediateError errText = do +-- | Immediate error stream emitted when the request itself could not +-- be built (e.g. message mapping failed). +immediateError :: BaikaiError -> IO [AssistantMessageEvent] +immediateError err = do now <- getCurrentTime + let errText = err ^. #message let msg = Msg.AssistantMessage Msg.AssistantPayload { Msg.content = Vector.empty, - Msg.usage = Usage._Usage, + Msg.usage = Usage.zeroUsage, Msg.stopReason = Stop.ErrorReason, Msg.errorMessage = Just errText, - Msg.timestamp = now + Msg.timestamp = Just now } - pure (EventError (errorTerminal Stop.ErrorReason msg (Just (invalidRequest errText)))) - --- ============================================================ --- Request mapping (preserved from EP-2 with minor refactoring) --- ============================================================ - -mapRequest :: - Model -> Context -> Options -> Either Text Chat.CreateChatCompletion -mapRequest m ctx opts = do - body <- traverse mapMessage (Vector.toList (ctx ^. #messages)) - let compat = openaiCompletionsCompatFor m - prefix = case ctx ^. #systemPrompt of - Nothing -> [] - Just sp -> - [ Chat.System - { Chat.content = Vector.singleton Chat.Text {Chat.text = sp}, - Chat.name = Nothing - } - ] - mt = fromMaybe (m ^. #maxOutputTokens) (opts ^. #maxTokens) - toolsField = - if Vector.null (ctx ^. #tools) - then Nothing - else Just (Vector.map (mkOpenAITool compat) (ctx ^. #tools)) - toolChoiceField = fmap mkOpenAIToolChoice (opts ^. #toolChoice) - reasoningEffortField = - applyThinkingFormat compat (opts ^. #thinking) - responseFormatField = - fmap mkOpenAIResponseFormat (opts ^. #responseFormat) pure - Chat._CreateChatCompletion - { Chat.messages = Vector.fromList (prefix <> body), - Chat.model = OpenAIModels.Model (m ^. #modelId), - Chat.max_completion_tokens = Just mt, - Chat.temperature = opts ^. #temperature, - Chat.tools = toolsField, - Chat.tool_choice = toolChoiceField, - Chat.reasoning_effort = reasoningEffortField, - Chat.response_format = responseFormatField - } - --- | Map a baikai 'ResponseFormat' onto the upstream OpenAI --- 'RF.ResponseFormat'. 'JsonObject' becomes plain-JSON mode; --- 'JsonSchema' becomes a named, optionally-strict schema. The --- schema 'Value' is forwarded verbatim. -mkOpenAIResponseFormat :: ResponseFormat -> RF.ResponseFormat -mkOpenAIResponseFormat = \case - JsonObject -> RF.JSON_Object - JsonSchema {name = n, schema = s, strict = st} -> - RF.JSON_Schema - { RF.json_schema = - RF.JSONSchema - { RF.description = Nothing, - RF.name = n, - RF.schema = Just s, - RF.strict = Just st - } - } - --- | Map a 'Baikai.ThinkingLevel.ThinkingLevel' onto the OpenAI SDK's --- 'Chat.ReasoningEffort' enum. Returns 'Nothing' when the caller did --- not request a level, when the host's 'thinkingFormat' is --- 'ThinkingFormatNone', or when the host expects a non-OpenAI shape --- the SDK does not support natively. --- --- The non-OpenAI thinking formats (DeepSeek, OpenRouter, Together, --- Z.ai, Qwen) require additional top-level JSON keys the upstream --- @openai@ Haskell SDK does not expose. They are silently dropped on --- this revision; see the EP-5 Decision Log for the rationale and --- pointers to the workaround when one is needed. -applyThinkingFormat :: - OpenAICompletionsCompat -> - Maybe ThinkingLevel -> - Maybe Chat.ReasoningEffort -applyThinkingFormat _ Nothing = Nothing -applyThinkingFormat compat (Just lvl) = case thinkingFormat compat of - ThinkingFormatOpenAI -> Just (toReasoningEffort lvl) - _ -> Nothing - -toReasoningEffort :: ThinkingLevel -> Chat.ReasoningEffort -toReasoningEffort = \case - ThinkingMinimal -> Chat.ReasoningEffort_Minimal - ThinkingLow -> Chat.ReasoningEffort_Low - ThinkingMedium -> Chat.ReasoningEffort_Medium - ThinkingHigh -> Chat.ReasoningEffort_High - --- | Map a baikai 'Tool.Tool' into the upstream OpenAI 'Tool_Function' --- shape. The compat record's 'supportsStrictMode' flag controls --- whether the @strict@ field is sent ('True') or dropped ('False'); --- some OpenAI-compatible hosts reject strict-mode tools entirely. --- --- The default ('defaultOpenAICompletionsCompat') leaves strict --- unset, which OpenAI treats as the default-permissive behaviour; --- callers that want @strict: true@ on every tool can flip the field --- on their compat record (a future enhancement; currently we pass --- 'Nothing' even when 'supportsStrictMode' is 'True', matching the --- pre-EP-5 behaviour). -mkOpenAITool :: OpenAICompletionsCompat -> Tool.Tool -> OpenAITool.Tool -mkOpenAITool _compat t = - OpenAITool.Tool_Function - { OpenAITool.function = - OpenAITool.Function - { OpenAITool.name = Tool.name t, - OpenAITool.description = Just (Tool.description t), - OpenAITool.parameters = Just (Tool.parameters t), - OpenAITool.strict = Nothing - } - } - --- | Map a baikai 'Tool.ToolChoice' into the upstream OpenAI --- 'ToolChoice'. OpenAI accepts @none@, @auto@, @required@, and a --- specific function reference; the SDK's 'ToolChoiceTool' takes the --- whole 'OpenAITool.Tool' value so we synthesise a stub function --- tool carrying just the name (OpenAI ignores the schema in this --- position). -mkOpenAIToolChoice :: Tool.ToolChoice -> OpenAITool.ToolChoice -mkOpenAIToolChoice = \case - Tool.ToolChoiceAuto -> OpenAITool.ToolChoiceAuto - Tool.ToolChoiceNone -> OpenAITool.ToolChoiceNone - Tool.ToolChoiceRequired -> OpenAITool.ToolChoiceRequired - Tool.ToolChoiceSpecific n -> - OpenAITool.ToolChoiceTool - ( OpenAITool.Tool_Function - { OpenAITool.function = - OpenAITool.Function - { OpenAITool.name = n, - OpenAITool.description = Nothing, - OpenAITool.parameters = Nothing, - OpenAITool.strict = Nothing - } - } - ) - -mapMessage :: Msg.Message -> Either Text (Chat.Message (Vector Chat.Content)) -mapMessage = \case - Msg.UserMessage Msg.UserPayload {Msg.content = uc} -> - Right - Chat.User - { Chat.content = Vector.map userContentToPart uc, - Chat.name = Nothing - } - Msg.AssistantMessage Msg.AssistantPayload {Msg.content = ac} -> - let textBody = collectAssistantText ac - toolCalls = collectToolCalls ac - content - | Text.null textBody = Nothing - | otherwise = Just (Vector.singleton Chat.Text {Chat.text = textBody}) - toolCallVec - | Vector.null toolCalls = Nothing - | otherwise = Just toolCalls - in Right - Chat.Assistant - { Chat.assistant_content = content, - Chat.refusal = Nothing, - Chat.name = Nothing, - Chat.assistant_audio = Nothing, - Chat.tool_calls = toolCallVec - } - Msg.ToolResultMessage - Msg.ToolResultPayload - { Msg.toolCallId = tid, - Msg.content = trc, - Msg.isError = err - } -> - case collectToolResultText trc of - Left unsupported -> Left unsupported - Right body -> - let decorated = if err then "[error] " <> body else body - in Right - Chat.Tool - { Chat.content = Vector.singleton Chat.Text {Chat.text = decorated}, - Chat.tool_call_id = tid - } - -userContentToPart :: Content.UserContent -> Chat.Content -userContentToPart = \case - Content.UserText (Content.TextContent t) -> Chat.Text {Chat.text = t} - Content.UserImage img -> - let encoded = Text.decodeUtf8 (Base64.encode (Content.imageData img)) - uri = "data:" <> Content.mimeType img <> ";base64," <> encoded - in Chat.Image_URL - { Chat.image_url = Chat.ImageURL {Chat.url = uri, Chat.detail = Nothing} - } - -collectAssistantText :: Vector Content.AssistantContent -> Text -collectAssistantText = - Text.concat - . Vector.toList - . Vector.mapMaybe - ( \case - Content.AssistantText (Content.TextContent t) -> Just t - Content.AssistantThinking th -> - if Content.redacted th - then Nothing - else Just ("" <> Content.thinking th <> "") - Content.AssistantToolCall _ -> Nothing - ) + [ EventStart StartPayload {partial = msg, responseId = Nothing}, + EventError (errorTerminal Nothing Stop.ErrorReason msg err) + ] -collectToolCalls :: Vector Content.AssistantContent -> Vector ToolCall.ToolCall -collectToolCalls = - Vector.mapMaybe - ( \case - Content.AssistantToolCall tc -> - Just - ToolCall.ToolCall_Function - { ToolCall.id = Content.id_ tc, - ToolCall.function = - ToolCall.Function - { ToolCall.name = Content.name tc, - ToolCall.arguments = - Text.decodeUtf8 (BSL.toStrict (Aeson.encode (Content.arguments tc))) - } - } - _ -> Nothing - ) +mapFinishReason :: Text -> (Stop.StopReason, Maybe Text) +mapFinishReason r = case r of + "stop" -> (Stop.Stop, Nothing) + "length" -> (Stop.Length, Nothing) + "tool_calls" -> (Stop.ToolUse, Nothing) + "function_call" -> (Stop.ToolUse, Nothing) + "content_filter" -> (Stop.ErrorReason, Nothing) + _ -> (Stop.Stop, Just ("unrecognized finish_reason: " <> r)) -collectToolResultText :: Vector Content.ToolResultContent -> Either Text Text -collectToolResultText = - fmap (Text.concat . Vector.toList) . traverse oneBlock - where - oneBlock = \case - Content.ToolResultText (Content.TextContent t) -> Right t - Content.ToolResultImage _ -> - Left "OpenAI Chat Completions cannot encode ToolResultImage blocks in tool-result messages" +trySync :: IO a -> IO (Either SomeException a) +trySync action = do + r <- try action + case r of + Left e + | Just (SomeAsyncException _) <- (fromException e :: Maybe SomeAsyncException) -> + throwIO e + | otherwise -> pure (Left e) + Right a -> pure (Right a) -mapFinishReason :: Text -> Stop.StopReason -mapFinishReason r = case r of - "stop" -> Stop.Stop - "length" -> Stop.Length - "tool_calls" -> Stop.ToolUse - "function_call" -> Stop.ToolUse - "content_filter" -> Stop.ErrorReason - _ -> Stop.Stop +exceptionToError :: SomeException -> BaikaiError +exceptionToError e = fromMaybe (classifyException e) (fromException e) diff --git a/src/Baikai/Provider/OpenAI/Cli.hs b/src/Baikai/Provider/OpenAI/Cli.hs --- a/src/Baikai/Provider/OpenAI/Cli.hs +++ b/src/Baikai/Provider/OpenAI/Cli.hs @@ -5,8 +5,16 @@ -- 'Baikai.Api.OpenAICompletionsCli' handler with default config. -- 'registerWith' accepts a caller-supplied 'CodexCliConfig'. module Baikai.Provider.OpenAI.Cli - ( CodexCliConfig (..), + ( CodexCliConfig, + executable, + extraArgs, + workingDir, + skipGitRepoCheck, + ephemeral, + codexCliCommand, + codexCliPrompt, defaultCodexCliConfig, + codexCliProvider, register, registerWith, registerWithRegistry, @@ -17,7 +25,7 @@ import Baikai.Api (Api (..)) import Baikai.Content (AssistantContent (..), TextContent (..)) import Baikai.Context (Context) -import Baikai.Error (processError, providerError) +import Baikai.Error (BaikaiError, processError, providerError) import Baikai.Message (AssistantPayload (..)) import Baikai.Model (Model) import Baikai.Options (Options) @@ -25,33 +33,35 @@ import Baikai.Provider.Registry ( ApiProvider (..), ProviderRegistry, - globalProviderRegistry, + registerApiProvider, registerApiProviderWith, ) import Baikai.Response qualified as Resp import Baikai.StopReason (StopReason (..)) import Baikai.Stream (liftCompleteToStream) -import Baikai.Usage (_Usage) -import Control.Exception (bracket, throwIO) +import Baikai.Usage (zeroUsage) +import Control.Concurrent (forkIO) +import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar) +import Control.Exception (SomeAsyncException (..), SomeException, displayException, fromException, throwIO, try) import Control.Lens ((^.)) import Data.ByteString qualified as BS import Data.Generics.Labels () +import Data.Maybe (fromMaybe) import Data.Text (Text) import Data.Text qualified as Text import Data.Time.Clock (UTCTime, diffUTCTime, getCurrentTime) -import Data.Vector (Vector) import Data.Vector qualified as Vector import GHC.Generics (Generic) import Streamly.Data.Stream (Stream) import Streamly.Data.Stream qualified as Stream import System.Exit (ExitCode (..)) -import System.IO (Handle, hClose) +import System.IO (Handle) import System.Process qualified as P -- | Configuration for the @codex exec --json@ subprocess. data CodexCliConfig = CodexCliConfig { executable :: !FilePath, - extraArgs :: !(Vector Text), + extraArgs :: ![Text], workingDir :: !(Maybe FilePath), skipGitRepoCheck :: !Bool, ephemeral :: !Bool @@ -70,8 +80,17 @@ -- | Install the Codex CLI handler with 'defaultCodexCliConfig'. register :: IO () -register = registerWith defaultCodexCliConfig +register = registerApiProvider (codexCliProvider defaultCodexCliConfig) +-- | First-class Codex CLI provider value for a caller-supplied config. +codexCliProvider :: CodexCliConfig -> ApiProvider +codexCliProvider cfg = + ApiProvider + { apiTag = OpenAICompletionsCli, + stream = liftCompleteToStream (runCodexCli cfg), + complete = runCodexCli cfg + } + -- | Install the Codex CLI handler with a caller-supplied config. -- -- The Codex binary runs in batch mode. 'stream' wraps the batch @@ -83,12 +102,14 @@ -- Log records the deviation from "complete = streamingComplete . -- stream". registerWith :: CodexCliConfig -> IO () -registerWith = registerWithRegistryAndConfig globalProviderRegistry +registerWith cfg = registerApiProvider (codexCliProvider cfg) +{-# DEPRECATED registerWith "use registerApiProvider (codexCliProvider cfg)" #-} -- | Install the Codex CLI handler with 'defaultCodexCliConfig' into an explicit -- registry. registerWithRegistry :: ProviderRegistry -> IO () registerWithRegistry reg = registerWithRegistryAndConfig reg defaultCodexCliConfig +{-# DEPRECATED registerWithRegistry "use registerApiProviderWith reg (codexCliProvider defaultCodexCliConfig)" #-} -- | Install the Codex CLI handler with a caller-supplied config into an -- explicit registry. @@ -96,11 +117,8 @@ registerWithRegistryAndConfig reg cfg = registerApiProviderWith reg - ApiProvider - { apiTag = OpenAICompletionsCli, - stream = liftCompleteToStream (runCodexCli cfg), - complete = runCodexCli cfg - } + (codexCliProvider cfg) +{-# DEPRECATED registerWithRegistryAndConfig "use registerApiProviderWith reg (codexCliProvider cfg)" #-} modelArgs :: Model -> [String] modelArgs m = case Text.strip (m ^. #modelId) of @@ -116,69 +134,107 @@ then pure Nothing else pure (Just (chunk, ())) +-- | The full prompt text passed to @codex exec@: the flattened +-- conversation, wrapped with the system prompt when one is set. +-- @codex exec --help@ exposes no system-prompt flag, so the system +-- prompt travels in prompt text. +codexCliPrompt :: Context -> Text +codexCliPrompt ctx = + Internal.wrapSystemPrompt (ctx ^. #systemPrompt) (Internal.renderPrompt ctx) + +-- | Render the executable and arguments for a @codex exec --json@ +-- batch call. The prompt is preceded by @--@ so dash-leading prompts +-- cannot be parsed as options. +codexCliCommand :: CodexCliConfig -> Model -> Context -> (FilePath, [String]) +codexCliCommand cfg m ctx = + ( cfg ^. #executable, + ["exec"] + <> modelArgs m + <> ["--json"] + <> ["--skip-git-repo-check" | cfg ^. #skipGitRepoCheck] + <> ["--ephemeral" | cfg ^. #ephemeral] + <> fmap Text.unpack (cfg ^. #extraArgs) + <> ["--", Text.unpack (codexCliPrompt ctx)] + ) + runCodexCli :: CodexCliConfig -> Model -> Context -> Options -> IO Resp.Response runCodexCli cfg m ctx _opts = do - let prompt = Internal.renderPrompt ctx - baseArgs = - ["exec"] - <> modelArgs m - <> ["--json"] - <> ["--skip-git-repo-check" | cfg ^. #skipGitRepoCheck] - <> ["--ephemeral" | cfg ^. #ephemeral] - <> fmap Text.unpack (Vector.toList (cfg ^. #extraArgs)) - <> [Text.unpack prompt] + let (exe, args) = codexCliCommand cfg m ctx procSpec = - (P.proc (cfg ^. #executable) baseArgs) + (P.proc exe args) { P.std_in = P.NoStream, P.std_out = P.CreatePipe, P.std_err = P.CreatePipe, P.cwd = cfg ^. #workingDir } start <- getCurrentTime - bracket - (P.createProcess procSpec) - cleanup - (consume start m) - -cleanup :: (Maybe Handle, Maybe Handle, Maybe Handle, P.ProcessHandle) -> IO () -cleanup (_, mOut, mErr, ph) = do - maybe (pure ()) hClose mOut - maybe (pure ()) hClose mErr - P.terminateProcess ph + result <- trySync (P.withCreateProcess procSpec (consume start m)) + case result of + Right resp -> pure resp + Left ex -> do + end <- getCurrentTime + pure (Resp.errorResponse m end (millisBetween start end) (exceptionToError ex)) consume :: UTCTime -> Model -> - (Maybe Handle, Maybe Handle, Maybe Handle, P.ProcessHandle) -> + Maybe Handle -> + Maybe Handle -> + Maybe Handle -> + P.ProcessHandle -> IO Resp.Response -consume start m (_, mOut, mErr, ph) = do - hOut <- maybe (throwIO (providerError "codex: stdout handle missing")) pure mOut - hErr <- maybe (throwIO (providerError "codex: stderr handle missing")) pure mErr - body <- Internal.parseCodexJsonlStream (handleStream hOut) - errBytes <- BS.hGetContents hErr - exitCode <- P.waitForProcess ph - end <- getCurrentTime - case exitCode of - ExitFailure n -> throwIO (processError n (Internal.decodeUtf8Lenient errBytes)) - ExitSuccess -> - pure - Resp.Response - { Resp.message = - AssistantPayload - { content = - Vector.singleton (AssistantText (TextContent (Text.strip body))), - usage = _Usage, - stopReason = Stop, - errorMessage = Nothing, - timestamp = end - }, - Resp.model = m, - Resp.api = OpenAICompletionsCli, - Resp.provider = m ^. #provider, - Resp.responseId = Nothing, - Resp.latencyMs = millisBetween start end, - Resp.errorInfo = Nothing - } +consume start m _ mOut mErr ph = do + case (mOut, mErr) of + (Nothing, _) -> errorNow (providerError "codex: stdout handle missing") + (_, Nothing) -> errorNow (providerError "codex: stderr handle missing") + (Just hOut, Just hErr) -> do + errVar <- newEmptyMVar + _ <- + forkIO $ do + result <- try (BS.hGetContents hErr) :: IO (Either SomeException BS.ByteString) + putMVar errVar (either (const BS.empty) id result) + body <- Internal.parseCodexJsonlStream (handleStream hOut) + errBytes <- takeMVar errVar + exitCode <- P.waitForProcess ph + end <- getCurrentTime + case exitCode of + ExitFailure n -> pure (Resp.errorResponse m end (millisBetween start end) (processError n (Internal.decodeUtf8Lenient errBytes))) + ExitSuccess -> + pure + Resp.Response + { Resp.message = + AssistantPayload + { content = + Vector.singleton (AssistantText (TextContent (Text.strip body))), + usage = zeroUsage, + stopReason = Stop, + errorMessage = Nothing, + timestamp = Just end + }, + Resp.model = m, + Resp.api = OpenAICompletionsCli, + Resp.provider = m ^. #provider, + Resp.responseId = Nothing, + Resp.latencyMs = millisBetween start end, + Resp.errorInfo = Nothing + } + where + errorNow err = do + end <- getCurrentTime + pure (Resp.errorResponse m end (millisBetween start end) err) -millisBetween :: UTCTime -> UTCTime -> Integer +millisBetween :: UTCTime -> UTCTime -> Int millisBetween a b = round (realToFrac (diffUTCTime b a) * (1000 :: Double)) + +trySync :: IO a -> IO (Either SomeException a) +trySync action = do + r <- try action + case r of + Left e + | Just (SomeAsyncException _) <- (fromException e :: Maybe SomeAsyncException) -> + throwIO e + | otherwise -> pure (Left e) + Right a -> pure (Right a) + +exceptionToError :: SomeException -> BaikaiError +exceptionToError e = fromMaybe (providerError (Text.pack (displayException e))) (fromException e) diff --git a/src/Baikai/Provider/OpenAI/ErrorClass.hs b/src/Baikai/Provider/OpenAI/ErrorClass.hs deleted file mode 100644 --- a/src/Baikai/Provider/OpenAI/ErrorClass.hs +++ /dev/null @@ -1,104 +0,0 @@ --- | Map failures from the OpenAI SDK onto baikai's typed 'BaikaiError'. --- 'classifyException' handles an exception thrown by the --- @servant-client@ HTTP layer (carrying an HTTP status); --- 'classifyErrorText' handles an error that arrives mid-stream as a --- plain text message (the OpenAI streaming chunk's @error@ field). -module Baikai.Provider.OpenAI.ErrorClass - ( classifyException, - classifyErrorText, - -- | Exposed for testing the HTTP-status mapping without a live call. - responseToError, - ) -where - -import Baikai.Error - ( BaikaiError (..), - ErrorCategory (..), - bodyIndicatesOverflow, - classifyHttpStatusWithBody, - decodeError, - providerError, - ) -import Control.Exception (SomeException, displayException, fromException) -import Data.ByteString (ByteString) -import Data.ByteString.Lazy qualified as LBS -import Data.CaseInsensitive qualified as CI -import Data.Foldable (toList) -import Data.Sequence (Seq) -import Data.Text (Text) -import Data.Text qualified as Text -import Data.Text.Encoding qualified as Text -import Data.Text.Encoding.Error qualified as Text -import Network.HTTP.Types.Status (statusCode) -import Servant.Client (ClientError, ResponseF (..)) -import Servant.Client qualified as Servant -import Text.Read (readMaybe) - --- | Convert any exception caught from the OpenAI SDK into a categorised --- 'BaikaiError'. Recognises @servant-client@ 'ClientError'; anything --- else degrades to a generic provider error carrying the displayed --- exception text. -classifyException :: SomeException -> BaikaiError -classifyException ex = case fromException ex of - Just clientErr -> fromClientError clientErr - Nothing -> providerError (Text.pack (displayException ex)) - -fromClientError :: ClientError -> BaikaiError -fromClientError clientErr = case clientErr of - Servant.FailureResponse _req resp -> responseToError resp - Servant.DecodeFailure detail _ -> decodeError detail - Servant.UnsupportedContentType _ _ -> decodeError "unsupported content type in OpenAI response" - Servant.InvalidContentTypeHeader _ -> decodeError "invalid content-type header in OpenAI response" - Servant.ConnectionError exc -> - (providerError ("connection error: " <> Text.pack (displayException exc))) - { category = TransientError - } - -responseToError :: ResponseF LBS.ByteString -> BaikaiError -responseToError resp = - BaikaiError - { category = classifyHttpStatusWithBody status retryAfter body, - message = msg, - httpStatus = Just status, - retryAfterSeconds = retryAfter, - exitCode = Nothing - } - where - status = statusCode (responseStatusCode resp) - body = decodeLenient (LBS.toStrict (responseBody resp)) - retryAfter = parseRetryAfter (responseHeaders resp) - snippet = Text.take 300 (Text.strip body) - msg = - "HTTP " - <> Text.pack (show status) - <> (if Text.null snippet then "" else ": " <> snippet) - -parseRetryAfter :: Seq (CI.CI ByteString, ByteString) -> Maybe Int -parseRetryAfter headers = do - raw <- lookup (CI.mk "Retry-After") (toList headers) - readMaybe (Text.unpack (Text.strip (decodeLenient raw))) - -decodeLenient :: ByteString -> Text -decodeLenient = Text.decodeUtf8With Text.lenientDecode - --- | Best-effort classification of an OpenAI streamed error message, --- which arrives as plain text without an HTTP status. Returns 'Nothing' --- for empty text (the caller keeps the raw message and 'errorInfo' --- stays absent). -classifyErrorText :: Text -> Maybe BaikaiError -classifyErrorText t - | Text.null (Text.strip t) = Nothing - | otherwise = Just (providerError t) {category = cat} - where - lower = Text.toLower t - has needle = needle `Text.isInfixOf` lower - cat - | bodyIndicatesOverflow t = ContextOverflow - | has "rate limit" || has "rate_limit" = RateLimited - | has "overloaded" || has "server_error" || has "service unavailable" = TransientError - | has "insufficient_quota" - || has "invalid api key" - || has "incorrect api key" - || has "invalid_api_key" = - AuthError - | otherwise = OtherError diff --git a/src/Baikai/Provider/OpenAI/Interactive.hs b/src/Baikai/Provider/OpenAI/Interactive.hs --- a/src/Baikai/Provider/OpenAI/Interactive.hs +++ b/src/Baikai/Provider/OpenAI/Interactive.hs @@ -6,7 +6,9 @@ -- batch completion provider, while this module starts the interactive -- terminal UI and returns only after the CLI exits. module Baikai.Provider.OpenAI.Interactive - ( CodexInteractiveConfig (..), + ( CodexInteractiveConfig, + executable, + extraArgs, defaultCodexInteractiveConfig, codexInteractiveCommand, codexInteractivePrompt, @@ -17,24 +19,24 @@ import Baikai.Interactive ( CodexApprovalPolicy, CodexSandboxMode, - InteractiveLaunchRequest (..), + InteractiveLaunchRequest, InteractiveLaunchResult, InteractiveProvider (..), InteractiveSafety (..), + interactiveLaunchResult, renderCodexApprovalPolicy, renderCodexSandboxMode, - _InteractiveLaunchResult, ) import Baikai.Prelude +import Baikai.Provider.Cli.Internal qualified as Internal import Data.Generics.Labels () import Data.Text qualified as Text -import Data.Vector qualified as Vector import System.Process qualified as P -- | Configuration for the interactive @codex@ process. data CodexInteractiveConfig = CodexInteractiveConfig { executable :: !FilePath, - extraArgs :: !(Vector Text) + extraArgs :: ![Text] } deriving stock (Eq, Show, Generic) @@ -55,25 +57,17 @@ <> workingDirArgs req <> extraDirArgs req <> safetyArgs req - <> fmap Text.unpack (Vector.toList (cfg ^. #extraArgs)) + <> fmap Text.unpack (cfg ^. #extraArgs) <> fmap Text.unpack (req ^. #extraArgs) - <> [Text.unpack (codexInteractivePrompt req)] + <> ["--", Text.unpack (codexInteractivePrompt req)] ) -- | Codex does not currently expose a top-level interactive -- system-prompt flag. Preserve Baikai's request shape by placing the -- system prompt before the user prompt in the initial prompt text. codexInteractivePrompt :: InteractiveLaunchRequest -> Text -codexInteractivePrompt req = case Text.strip <$> req ^. #systemPrompt of - Nothing -> req ^. #userPrompt - Just "" -> req ^. #userPrompt - Just prompt -> - Text.concat - [ "System instructions:\n", - prompt, - "\n\nUser request:\n", - req ^. #userPrompt - ] +codexInteractivePrompt req = + Internal.wrapSystemPrompt (req ^. #systemPrompt) (req ^. #userPrompt) -- | Launch Codex with inherited stdin, stdout, and stderr so the -- local CLI owns the interactive terminal experience. @@ -88,12 +82,11 @@ P.std_err = P.Inherit, P.cwd = req ^. #workingDir } - (_, _, _, ph) <- P.createProcess spec - code <- P.waitForProcess ph - pure (_InteractiveLaunchResult InteractiveCodex code) + code <- P.withCreateProcess spec (\_ _ _ ph -> P.waitForProcess ph) + pure (interactiveLaunchResult InteractiveCodex code) modelArgs :: InteractiveLaunchRequest -> [String] -modelArgs req = case Text.strip <$> req ^. #model of +modelArgs req = case Text.strip <$> req ^. #modelId of Nothing -> [] Just "" -> [] Just mid -> ["--model", Text.unpack mid] diff --git a/src/Baikai/Provider/OpenAI/Internal/ErrorClass.hs b/src/Baikai/Provider/OpenAI/Internal/ErrorClass.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Provider/OpenAI/Internal/ErrorClass.hs @@ -0,0 +1,140 @@ +-- | Internal failure classification for the OpenAI provider. +-- +-- This module is exposed for provider tests and debugging, but it is +-- not part of baikai's PVP-stable application surface. Names, types, +-- and semantics here may change in minor releases. +-- +-- 'classifyException' handles an exception thrown by the +-- @servant-client@ HTTP layer; 'classifyErrorText' handles an error +-- that arrives mid-stream as a plain text message. +module Baikai.Provider.OpenAI.Internal.ErrorClass + ( classifyException, + classifyErrorText, + -- | Exposed for testing the HTTP-status mapping without a live call. + responseToError, + ) +where + +import Baikai.Error + ( BaikaiError (..), + ErrorCategory (..), + bodyIndicatesOverflow, + decodeError, + httpError, + invalidRequest, + parseRetryAfterSeconds, + providerError, + ) +import Control.Exception (SomeException, displayException, fromException) +import Data.ByteString (ByteString) +import Data.ByteString.Lazy qualified as LBS +import Data.CaseInsensitive qualified as CI +import Data.Foldable (toList) +import Data.Sequence (Seq) +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as Text +import Data.Text.Encoding.Error qualified as Text +import Network.HTTP.Client qualified as HTTP +import Network.HTTP.Types.Status (statusCode) +import Servant.Client (ClientError, ResponseF (..)) +import Servant.Client qualified as Servant +import Text.Read (readMaybe) + +-- | Convert any exception caught from the OpenAI SDK into a categorised +-- 'BaikaiError'. Recognises @servant-client@ 'ClientError' and raw +-- @http-client@ 'HttpException'; anything else degrades to a generic +-- provider error carrying the displayed exception text. +classifyException :: SomeException -> BaikaiError +classifyException ex + | Just clientErr <- fromException ex = fromClientError clientErr + | Just httpEx <- fromException ex = fromHttpException httpEx + | otherwise = providerError (Text.pack (displayException ex)) + +fromClientError :: ClientError -> BaikaiError +fromClientError clientErr = case clientErr of + Servant.FailureResponse _req resp -> responseToError resp + Servant.DecodeFailure detail _ -> decodeError detail + Servant.UnsupportedContentType _ _ -> decodeError "unsupported content type in OpenAI response" + Servant.InvalidContentTypeHeader _ -> decodeError "invalid content-type header in OpenAI response" + Servant.ConnectionError exc -> + (providerError ("connection error: " <> Text.pack (displayException exc))) + { category = TransientError + } + +responseToError :: ResponseF LBS.ByteString -> BaikaiError +responseToError resp = httpError status retryAfter body + where + status = statusCode (responseStatusCode resp) + body = decodeLenient (LBS.toStrict (responseBody resp)) + retryAfter = parseRetryAfter (responseHeaders resp) + +parseRetryAfter :: Seq (CI.CI ByteString, ByteString) -> Maybe Int +parseRetryAfter headers = do + raw <- lookup (CI.mk "Retry-After") (toList headers) + parseRetryAfterSeconds (decodeLenient raw) + +fromHttpException :: HTTP.HttpException -> BaikaiError +fromHttpException = \case + HTTP.InvalidUrlException url reason -> + invalidRequest (Text.pack (url <> ": " <> reason)) + HTTP.HttpExceptionRequest _ content -> fromHttpExceptionContent content + +fromHttpExceptionContent :: HTTP.HttpExceptionContent -> BaikaiError +fromHttpExceptionContent = \case + HTTP.StatusCodeException resp body -> + httpError + (statusCode (HTTP.responseStatus resp)) + (parseRetryAfterHttp (HTTP.responseHeaders resp)) + (decodeLenient body) + HTTP.ConnectionFailure e -> transient (Text.pack (displayException e)) + HTTP.ConnectionTimeout -> transient "connection timeout" + HTTP.ResponseTimeout -> transient "response timeout" + HTTP.ConnectionClosed -> transient "connection closed" + HTTP.NoResponseDataReceived -> transient "no response data received" + HTTP.IncompleteHeaders -> transient "incomplete response headers" + other -> providerError (Text.pack (show other)) + where + transient t = (providerError ("connection error: " <> t)) {category = TransientError} + +parseRetryAfterHttp :: [(CI.CI ByteString, ByteString)] -> Maybe Int +parseRetryAfterHttp headers = do + raw <- lookup (CI.mk "Retry-After") headers + parseRetryAfterSeconds (decodeLenient raw) + +decodeLenient :: ByteString -> Text +decodeLenient = Text.decodeUtf8With Text.lenientDecode + +-- | Best-effort classification of an OpenAI streamed error message, +-- which arrives as plain text without an HTTP status. Returns 'Nothing' +-- for empty text (the caller keeps the raw message and 'errorInfo' +-- stays absent). +classifyErrorText :: Text -> Maybe BaikaiError +classifyErrorText t + | Just e <- classifySdkHttpText t = Just e + | Text.null (Text.strip t) = Nothing + | otherwise = Just (providerError t) {category = cat} + where + lower = Text.toLower t + has needle = needle `Text.isInfixOf` lower + cat + | bodyIndicatesOverflow t = ContextOverflow + | has "rate limit" || has "rate_limit" = RateLimited + | has "overloaded" || has "server_error" || has "service unavailable" = TransientError + | has "insufficient_quota" + || has "invalid api key" + || has "incorrect api key" + || has "invalid_api_key" = + AuthError + | otherwise = OtherError + +classifySdkHttpText :: Text -> Maybe BaikaiError +classifySdkHttpText raw = do + rest <- Text.stripPrefix "HTTP error " raw + let (codeText, afterCode) = Text.breakOn " " rest + code <- readMaybe (Text.unpack codeText) + let body = case Text.breakOn ": " afterCode of + (_, sepBody) + | not (Text.null sepBody) -> Text.drop 2 sepBody + _ -> "" + pure (httpError code Nothing body) diff --git a/src/Baikai/Provider/OpenAI/Internal/Request.hs b/src/Baikai/Provider/OpenAI/Internal/Request.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Provider/OpenAI/Internal/Request.hs @@ -0,0 +1,265 @@ +{-# LANGUAGE LambdaCase #-} + +-- | Internal request mapping for the OpenAI Chat Completions provider. +-- +-- This module is exposed for provider tests and debugging but is not covered by PVP +-- compatibility guarantees. Import the public provider module for stable application code. +module Baikai.Provider.OpenAI.Internal.Request + ( mapRequest, + ) +where + +import Baikai.Compat + ( OpenAICompletionsCompat (..), + ThinkingFormat (..), + ) +import Baikai.Content qualified as Content +import Baikai.Context (Context (..)) +import Baikai.Message qualified as Msg +import Baikai.Model (Model, openaiCompletionsCompatFor) +import Baikai.Options (Options (..)) +import Baikai.ResponseFormat (ResponseFormat (..)) +import Baikai.ThinkingLevel (ThinkingLevel (..)) +import Baikai.Tool qualified as Tool +import Control.Lens ((^.)) +import Data.Aeson qualified as Aeson +import Data.ByteString.Base64 qualified as Base64 +import Data.ByteString.Lazy qualified as BSL +import Data.Generics.Labels () +import Data.Maybe (fromMaybe) +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as Text +import Data.Vector (Vector) +import Data.Vector qualified as Vector +import OpenAI.V1.Chat.Completions qualified as Chat +import OpenAI.V1.Models qualified as OpenAIModels +import OpenAI.V1.ResponseFormat qualified as RF +import OpenAI.V1.Tool qualified as OpenAITool +import OpenAI.V1.ToolCall qualified as ToolCall + +-- ============================================================ +-- Request mapping (preserved from EP-2 with minor refactoring) +-- ============================================================ + +mapRequest :: + Model -> Context -> Options -> Either Text Chat.CreateChatCompletion +mapRequest m ctx opts = do + body <- traverse mapMessage (Vector.toList (ctx ^. #messages)) + let compat = openaiCompletionsCompatFor m + prefix = case ctx ^. #systemPrompt of + Nothing -> [] + Just sp -> + [ Chat.System + { Chat.content = Vector.singleton Chat.Text {Chat.text = sp}, + Chat.name = Nothing + } + ] + mt = fromMaybe (m ^. #maxOutputTokens) (opts ^. #maxTokens) + maxTokensField = + if mt == 0 then Nothing else Just mt + toolsField = + if Vector.null (ctx ^. #tools) + then Nothing + else Just (Vector.map (mkOpenAITool compat) (ctx ^. #tools)) + toolChoiceField = fmap mkOpenAIToolChoice (opts ^. #toolChoice) + reasoningEffortField = + applyThinkingFormat compat (opts ^. #thinking) + responseFormatField = + fmap (mkOpenAIResponseFormat compat) (opts ^. #responseFormat) + pure + Chat._CreateChatCompletion + { Chat.messages = Vector.fromList (prefix <> body), + Chat.model = OpenAIModels.Model (m ^. #modelId), + Chat.max_completion_tokens = maxTokensField, + Chat.temperature = opts ^. #temperature, + Chat.top_p = opts ^. #topP, + Chat.stop = opts ^. #stopSequences, + Chat.seed = opts ^. #seed, + Chat.frequency_penalty = opts ^. #frequencyPenalty, + Chat.presence_penalty = opts ^. #presencePenalty, + Chat.tools = toolsField, + Chat.tool_choice = toolChoiceField, + Chat.reasoning_effort = reasoningEffortField, + Chat.response_format = responseFormatField + } + +-- | Map a baikai 'ResponseFormat' onto the upstream OpenAI +-- 'RF.ResponseFormat'. 'JsonObject' becomes plain-JSON mode; +-- 'JsonSchema' becomes a named, optionally-strict schema. The +-- schema 'Value' is forwarded verbatim. +mkOpenAIResponseFormat :: OpenAICompletionsCompat -> ResponseFormat -> RF.ResponseFormat +mkOpenAIResponseFormat _ JsonObject = RF.JSON_Object +mkOpenAIResponseFormat compat JsonSchema {name = n, schema = s, strict = st} = + RF.JSON_Schema + { RF.json_schema = + RF.JSONSchema + { RF.description = Nothing, + RF.name = n, + RF.schema = Just s, + RF.strict = if supportsStrictMode compat then Just st else Nothing + } + } + +-- | Map a 'Baikai.ThinkingLevel.ThinkingLevel' onto the OpenAI SDK's +-- 'Chat.ReasoningEffort' enum. Returns 'Nothing' when the caller did +-- not request a level, when the host's 'thinkingFormat' is +-- 'ThinkingFormatNone', or when the host expects a non-OpenAI shape +-- the SDK does not support natively. +-- +-- The non-OpenAI thinking formats (DeepSeek, OpenRouter, Together, +-- Z.ai, Qwen) require additional top-level JSON keys the upstream +-- @openai@ Haskell SDK does not expose. They are silently dropped on +-- this revision; see the EP-5 Decision Log for the rationale and +-- pointers to the workaround when one is needed. +applyThinkingFormat :: + OpenAICompletionsCompat -> + Maybe ThinkingLevel -> + Maybe Chat.ReasoningEffort +applyThinkingFormat _ Nothing = Nothing +applyThinkingFormat compat (Just lvl) = case thinkingFormat compat of + ThinkingFormatOpenAI -> Just (toReasoningEffort lvl) + _ -> Nothing + +toReasoningEffort :: ThinkingLevel -> Chat.ReasoningEffort +toReasoningEffort = \case + ThinkingMinimal -> Chat.ReasoningEffort_Minimal + ThinkingLow -> Chat.ReasoningEffort_Low + ThinkingMedium -> Chat.ReasoningEffort_Medium + ThinkingHigh -> Chat.ReasoningEffort_High + +-- | Map a baikai 'Tool.Tool' into the upstream OpenAI 'Tool_Function' +-- shape. The compat record's 'supportsStrictMode' flag controls +-- whether the @strict@ field is sent ('True') or dropped ('False'); +-- some OpenAI-compatible hosts reject strict-mode tools entirely. +-- +-- The default ('defaultOpenAICompletionsCompat') leaves strict +-- unset, which OpenAI treats as the default-permissive behaviour; +-- callers that want @strict: true@ on every tool can flip the field +-- on their compat record (a future enhancement; currently we pass +-- 'Nothing' even when 'supportsStrictMode' is 'True', matching the +-- pre-EP-5 behaviour). +mkOpenAITool :: OpenAICompletionsCompat -> Tool.Tool -> OpenAITool.Tool +mkOpenAITool _compat t = + OpenAITool.Tool_Function + { OpenAITool.function = + OpenAITool.Function + { OpenAITool.name = Tool.name t, + OpenAITool.description = Just (Tool.description t), + OpenAITool.parameters = Just (Tool.parameters t), + OpenAITool.strict = Nothing + } + } + +-- | Map a baikai 'Tool.ToolChoice' into the upstream OpenAI +-- 'ToolChoice'. OpenAI accepts @none@, @auto@, @required@, and a +-- specific function reference; the SDK's 'ToolChoiceTool' takes the +-- whole 'OpenAITool.Tool' value so we synthesise a stub function +-- tool carrying just the name (OpenAI ignores the schema in this +-- position). +mkOpenAIToolChoice :: Tool.ToolChoice -> OpenAITool.ToolChoice +mkOpenAIToolChoice = \case + Tool.ToolChoiceAuto -> OpenAITool.ToolChoiceAuto + Tool.ToolChoiceNone -> OpenAITool.ToolChoiceNone + Tool.ToolChoiceRequired -> OpenAITool.ToolChoiceRequired + Tool.ToolChoiceSpecific n -> + OpenAITool.ToolChoiceTool + ( OpenAITool.Tool_Function + { OpenAITool.function = + OpenAITool.Function + { OpenAITool.name = n, + OpenAITool.description = Nothing, + OpenAITool.parameters = Nothing, + OpenAITool.strict = Nothing + } + } + ) + +mapMessage :: Msg.Message -> Either Text (Chat.Message (Vector Chat.Content)) +mapMessage = \case + Msg.UserMessage Msg.UserPayload {Msg.content = uc} -> + Right + Chat.User + { Chat.content = Vector.map userContentToPart uc, + Chat.name = Nothing + } + Msg.AssistantMessage Msg.AssistantPayload {Msg.content = ac} -> + let textBody = collectAssistantText ac + toolCalls = collectToolCalls ac + content + | Text.null textBody = Nothing + | otherwise = Just (Vector.singleton Chat.Text {Chat.text = textBody}) + toolCallVec + | Vector.null toolCalls = Nothing + | otherwise = Just toolCalls + in Right + Chat.Assistant + { Chat.assistant_content = content, + Chat.refusal = Nothing, + Chat.name = Nothing, + Chat.assistant_audio = Nothing, + Chat.tool_calls = toolCallVec + } + Msg.ToolResultMessage + Msg.ToolResultPayload + { Msg.toolCallId = tid, + Msg.content = trc, + Msg.isError = err + } -> + case collectToolResultText trc of + Left unsupported -> Left unsupported + Right body -> + let decorated = if err then "[error] " <> body else body + in Right + Chat.Tool + { Chat.content = Vector.singleton Chat.Text {Chat.text = decorated}, + Chat.tool_call_id = tid + } + +userContentToPart :: Content.UserContent -> Chat.Content +userContentToPart = \case + Content.UserText (Content.TextContent t) -> Chat.Text {Chat.text = t} + Content.UserImage img -> + let encoded = Text.decodeUtf8 (Base64.encode (Content.imageData img)) + uri = "data:" <> Content.mimeType img <> ";base64," <> encoded + in Chat.Image_URL + { Chat.image_url = Chat.ImageURL {Chat.url = uri, Chat.detail = Nothing} + } + +collectAssistantText :: Vector Content.AssistantContent -> Text +collectAssistantText = + Text.concat + . Vector.toList + . Vector.mapMaybe + ( \case + Content.AssistantText (Content.TextContent t) -> Just t + Content.AssistantThinking _ -> Nothing + Content.AssistantToolCall _ -> Nothing + ) + +collectToolCalls :: Vector Content.AssistantContent -> Vector ToolCall.ToolCall +collectToolCalls = + Vector.mapMaybe + ( \case + Content.AssistantToolCall tc -> + Just + ToolCall.ToolCall_Function + { ToolCall.id = Content.id_ tc, + ToolCall.function = + ToolCall.Function + { ToolCall.name = Content.name tc, + ToolCall.arguments = + Text.decodeUtf8 (BSL.toStrict (Aeson.encode (Content.arguments tc))) + } + } + _ -> Nothing + ) + +collectToolResultText :: Vector Content.ToolResultContent -> Either Text Text +collectToolResultText = + fmap (Text.concat . Vector.toList) . traverse oneBlock + where + oneBlock = \case + Content.ToolResultText (Content.TextContent t) -> Right t + Content.ToolResultImage _ -> + Left "OpenAI Chat Completions cannot encode ToolResultImage blocks in tool-result messages" diff --git a/src/Baikai/Provider/OpenAI/Shape.hs b/src/Baikai/Provider/OpenAI/Shape.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Provider/OpenAI/Shape.hs @@ -0,0 +1,216 @@ +{-# LANGUAGE LambdaCase #-} + +-- | Pure request-body shaping for OpenAI-compatible Chat Completions hosts. +module Baikai.Provider.OpenAI.Shape + ( shapeRequestBody, + streamRequestBody, + renameMaxTokens, + dropUnsupportedStrict, + injectThinkingShape, + injectCacheControl, + ) +where + +import Baikai.CacheRetention (CacheRetention (..)) +import Baikai.Compat + ( CacheControlFormat (..), + MaxTokensField (..), + OpenAICompletionsCompat + ( cacheControlFormat, + maxTokensField, + supportsLongCacheRetention, + supportsStrictMode, + supportsUsageInStreaming, + thinkingFormat + ), + ThinkingFormat (..), + ) +import Baikai.Options (Options, cacheRetention, thinking) +import Baikai.ThinkingLevel (ThinkingLevel (..)) +import Data.Aeson (Value (..), (.=)) +import Data.Aeson qualified as Aeson +import Data.Aeson.Key qualified as AesonKey +import Data.Aeson.KeyMap (KeyMap) +import Data.Aeson.KeyMap qualified as KeyMap +import Data.Text (Text) +import Data.Vector (Vector) +import Data.Vector qualified as Vector +import OpenAI.V1.Chat.Completions qualified as Chat + +shapeRequestBody :: + OpenAICompletionsCompat -> Options -> Aeson.Value -> Aeson.Value +shapeRequestBody compat opts = + injectCacheControl compat opts + . injectThinkingShape compat opts + . dropUnsupportedStrict compat + . renameMaxTokens compat + +streamRequestBody :: + OpenAICompletionsCompat -> + Options -> + Chat.CreateChatCompletion -> + Aeson.Value +streamRequestBody compat opts req = + shapeRequestBody compat opts (Aeson.toJSON req') + where + req' = + req + { Chat.stream = Just True, + Chat.stream_options = + if supportsUsageInStreaming compat + then + Just + Chat._ChatCompletionStreamOptions + { Chat.include_usage = Just True + } + else Nothing + } + +renameMaxTokens :: OpenAICompletionsCompat -> Aeson.Value -> Aeson.Value +renameMaxTokens compat + | maxTokensField compat /= MaxTokensField = id + | otherwise = + mapObject $ \obj -> + case KeyMap.lookup (key "max_completion_tokens") obj of + Nothing -> obj + Just v -> + KeyMap.insert (key "max_tokens") v $ + KeyMap.delete (key "max_completion_tokens") obj + +dropUnsupportedStrict :: OpenAICompletionsCompat -> Aeson.Value -> Aeson.Value +dropUnsupportedStrict compat + | supportsStrictMode compat = id + | otherwise = + mapObject $ + adjustKey (key "response_format") $ + mapObject $ + adjustKey (key "json_schema") $ + mapObject (KeyMap.delete (key "strict")) + +injectThinkingShape :: OpenAICompletionsCompat -> Options -> Aeson.Value -> Aeson.Value +injectThinkingShape compat opts body = + case thinking opts of + Nothing -> body + Just lvl -> case thinkingFormat compat of + ThinkingFormatOpenAI -> body + ThinkingFormatNone -> body + ThinkingFormatOpenRouter -> + insertTop "reasoning" (Aeson.object ["effort" .= effort lvl]) body + ThinkingFormatDeepseek -> + insertTop "reasoning_effort" (String (effort lvl)) $ + insertTop "thinking" (Aeson.object ["type" .= ("enabled" :: Text)]) body + ThinkingFormatTogether -> + insertTop "reasoning_effort" (String (effort lvl)) $ + insertTop "reasoning" (Aeson.object ["enabled" .= True]) body + ThinkingFormatZai -> + insertTop "enable_thinking" (Bool True) body + ThinkingFormatQwen -> + insertTop "enable_thinking" (Bool True) body + +injectCacheControl :: OpenAICompletionsCompat -> Options -> Aeson.Value -> Aeson.Value +injectCacheControl compat opts body = + case (cacheControlFormat compat, cacheRetention opts) of + (Just CacheControlFormatAnthropic, Just retention) + | Just marker <- cacheControlMarker compat retention -> + mapObject + (adjustKey (key "messages") (shapeMessages marker)) + body + _ -> body + +cacheControlMarker :: OpenAICompletionsCompat -> CacheRetention -> Maybe Aeson.Value +cacheControlMarker compat = \case + CacheRetentionNone -> Nothing + CacheRetentionShort -> Just (Aeson.object ["type" .= ("ephemeral" :: Text)]) + CacheRetentionLong -> + if supportsLongCacheRetention compat + then + Just + ( Aeson.object + [ "type" .= ("ephemeral" :: Text), + "ttl" .= ("1h" :: Text) + ] + ) + else Just (Aeson.object ["type" .= ("ephemeral" :: Text)]) + +shapeMessages :: Aeson.Value -> Aeson.Value -> Aeson.Value +shapeMessages marker = \case + Array messages -> + let target = + findLastIndex isSystemMessage messages + `orElse` findLastIndex isUserMessage messages + in case target of + Nothing -> Array messages + Just i -> + Array (messages Vector.// [(i, shapeMessage marker (messages Vector.! i))]) + other -> other + +shapeMessage :: Aeson.Value -> Aeson.Value -> Aeson.Value +shapeMessage marker = + mapObject (adjustKey (key "content") (shapeContent marker)) + +shapeContent :: Aeson.Value -> Aeson.Value -> Aeson.Value +shapeContent marker = \case + Array parts + | not (Vector.null parts) -> + let i = Vector.length parts - 1 + in Array (parts Vector.// [(i, shapeContentPart marker (parts Vector.! i))]) + other -> other + +shapeContentPart :: Aeson.Value -> Aeson.Value -> Aeson.Value +shapeContentPart marker = + mapObject (KeyMap.insert (key "cache_control") marker) + +isSystemMessage :: Aeson.Value -> Bool +isSystemMessage = hasRole "system" + +isUserMessage :: Aeson.Value -> Bool +isUserMessage = hasRole "user" + +hasRole :: Text -> Aeson.Value -> Bool +hasRole role = \case + Object obj -> KeyMap.lookup (key "role") obj == Just (String role) + _ -> False + +findLastIndex :: (Aeson.Value -> Bool) -> Vector Aeson.Value -> Maybe Int +findLastIndex p = + fmap fst + . lastMay + . filter (p . snd) + . zip [0 ..] + . Vector.toList + +lastMay :: [a] -> Maybe a +lastMay [] = Nothing +lastMay xs = Just (last xs) + +orElse :: Maybe a -> Maybe a -> Maybe a +orElse (Just x) _ = Just x +orElse Nothing y = y + +insertTop :: Text -> Aeson.Value -> Aeson.Value -> Aeson.Value +insertTop k v = mapObject (KeyMap.insert (key k) v) + +mapObject :: (KeyMap Aeson.Value -> KeyMap Aeson.Value) -> Aeson.Value -> Aeson.Value +mapObject f = \case + Object obj -> Object (f obj) + other -> other + +adjustKey :: + AesonKey.Key -> + (Aeson.Value -> Aeson.Value) -> + KeyMap Aeson.Value -> + KeyMap Aeson.Value +adjustKey k f obj = + case KeyMap.lookup k obj of + Nothing -> obj + Just old -> KeyMap.insert k (f old) obj + +key :: Text -> AesonKey.Key +key = AesonKey.fromText + +effort :: ThinkingLevel -> Text +effort = \case + ThinkingMinimal -> "low" + ThinkingLow -> "low" + ThinkingMedium -> "medium" + ThinkingHigh -> "high" diff --git a/src/Baikai/Provider/OpenAI/Sse.hs b/src/Baikai/Provider/OpenAI/Sse.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Provider/OpenAI/Sse.hs @@ -0,0 +1,159 @@ +{-# LANGUAGE LambdaCase #-} + +-- | Local SSE transport for OpenAI Chat Completions streams. +module Baikai.Provider.OpenAI.Sse + ( openaiSseStream, + openaiSseStreamValue, + openaiSseStreamValueWithHeaders, + sseFromResponse, + ) +where + +import Baikai.Error (BaikaiError, decodeError, httpError, parseRetryAfterSeconds) +import Control.Monad (foldM, when) +import Data.Aeson qualified as Aeson +import Data.ByteString qualified as SBS +import Data.ByteString.Char8 qualified as S8 +import Data.CaseInsensitive qualified as CI +import Data.IORef qualified as IORef +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as Text +import Data.Text.Encoding.Error qualified as Text +import Network.HTTP.Client qualified as HTTP +import Network.HTTP.Types.Header (RequestHeaders) +import Network.HTTP.Types.Status qualified as Status +import OpenAI.V1.Chat.Completions qualified as Chat +import Servant.Client qualified as Client + +-- | POST the request to @/v1/chat/completions@ and feed decoded SSE +-- JSON payloads to the callback. A @data: [DONE]@ frame ends the +-- stream without producing a callback value. +openaiSseStream :: + Client.ClientEnv -> + Text -> + Chat.CreateChatCompletion -> + (Either BaikaiError Aeson.Value -> IO ()) -> + IO () +openaiSseStream env apiKey req = + openaiSseStreamValue env apiKey (Aeson.toJSON req) + +openaiSseStreamValue :: + Client.ClientEnv -> + Text -> + Aeson.Value -> + (Either BaikaiError Aeson.Value -> IO ()) -> + IO () +openaiSseStreamValue env apiKey = + openaiSseStreamValueWithHeaders + env + [ ("Authorization", Text.encodeUtf8 ("Bearer " <> apiKey)), + ("Accept", "text/event-stream"), + ("Content-Type", "application/json") + ] + +openaiSseStreamValueWithHeaders :: + Client.ClientEnv -> + RequestHeaders -> + Aeson.Value -> + (Either BaikaiError Aeson.Value -> IO ()) -> + IO () +openaiSseStreamValueWithHeaders env requestHeaders requestBody onEvent = do + let base = Client.baseUrl env + secure = case Client.baseUrlScheme base of + Client.Http -> False + Client.Https -> True + request = + HTTP.defaultRequest + { HTTP.secure = secure, + HTTP.host = S8.pack (Client.baseUrlHost base), + HTTP.port = Client.baseUrlPort base, + HTTP.method = "POST", + HTTP.path = S8.pack (normalizePath (Client.baseUrlPath base) <> "/v1/chat/completions"), + HTTP.requestHeaders = requestHeaders, + HTTP.requestBody = HTTP.RequestBodyLBS (Aeson.encode requestBody), + -- EP-8 wires Options.timeoutMs through this local transport. + HTTP.responseTimeout = HTTP.responseTimeoutNone + } + HTTP.withResponse request (Client.manager env) (`sseFromResponse` onEvent) + +sseFromResponse :: + HTTP.Response HTTP.BodyReader -> + (Either BaikaiError Aeson.Value -> IO ()) -> + IO () +sseFromResponse response onEvent = do + let st = HTTP.responseStatus response + if not (Status.statusIsSuccessful st) + then do + bodyChunks <- HTTP.brConsume (HTTP.responseBody response) + let bodyText = decodeLenient (SBS.concat bodyChunks) + retryAfter = + parseRetryAfterSeconds . decodeLenient + =<< lookup (CI.mk "Retry-After") (HTTP.responseHeaders response) + onEvent (Left (httpError (Status.statusCode st) retryAfter bodyText)) + else do + lineBufRef <- IORef.newIORef SBS.empty + eventBufRef <- IORef.newIORef ([] :: [SBS.ByteString]) + let flushEvent = do + es <- IORef.atomicModifyIORef' eventBufRef (\buf -> ([], reverse buf)) + case es of + [] -> pure False + _ -> do + let payload = S8.concat es + if payload == "[DONE]" + then pure True + else case Aeson.eitherDecodeStrict payload of + Left err -> onEvent (Left (decodeError (Text.pack err))) >> pure False + Right val -> onEvent (Right val) >> pure False + + handleLine line = + let l = stripCR line + in if S8.null l + then flushEvent + else + if "data:" `S8.isPrefixOf` l + then do + let d = S8.dropWhile (== ' ') (S8.drop 5 l) + IORef.modifyIORef' eventBufRef (d :) + pure False + else pure False + + loop = do + chunk <- HTTP.brRead (HTTP.responseBody response) + if SBS.null chunk + then do + pendingLine <- IORef.readIORef lineBufRef + when (not (SBS.null pendingLine)) $ do + _ <- handleLine pendingLine + IORef.writeIORef lineBufRef SBS.empty + _ <- flushEvent + pure () + else do + prev <- IORef.readIORef lineBufRef + let combined = prev <> chunk + ls = S8.split '\n' combined + case unsnoc ls of + Nothing -> loop + Just (completeLines, lastLine) -> do + IORef.writeIORef lineBufRef lastLine + stop <- foldM (\acc ln -> if acc then pure True else handleLine ln) False completeLines + if stop then pure () else loop + loop + +normalizePath :: String -> String +normalizePath = \case + "" -> "" + p@('/' : _) -> p + p -> '/' : p + +stripCR :: SBS.ByteString -> SBS.ByteString +stripCR bs = case S8.unsnoc bs of + Just (initBs, '\r') -> initBs + _ -> bs + +unsnoc :: [a] -> Maybe ([a], a) +unsnoc [] = Nothing +unsnoc xs = Just (init xs, last xs) + +decodeLenient :: SBS.ByteString -> Text +decodeLenient = Text.decodeUtf8With Text.lenientDecode diff --git a/src/Baikai/Provider/OpenAI/Transport.hs b/src/Baikai/Provider/OpenAI/Transport.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Provider/OpenAI/Transport.hs @@ -0,0 +1,104 @@ +module Baikai.Provider.OpenAI.Transport + ( getClientEnvCached, + cachedClientEnvCount, + requestHeaders, + resolveKey, + runWithTimeout, + ) +where + +import Baikai.Auth qualified as Auth +import Baikai.Error (BaikaiError (..), ErrorCategory (..), authError) +import Baikai.Model (Model (..)) +import Baikai.Options (Options (..)) +import Control.Concurrent.MVar (MVar, modifyMVar, newMVar) +import Control.Exception (throwIO) +import Control.Lens ((^.)) +import Data.CaseInsensitive qualified as CI +import Data.Generics.Labels () +import Data.Map.Strict qualified as Map +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as Text +import Network.HTTP.Client qualified as HTTP +import Network.HTTP.Client.TLS qualified as TLS +import Network.HTTP.Types.Header (RequestHeaders) +import Servant.Client qualified as Client +import System.IO.Unsafe (unsafePerformIO) +import System.Timeout qualified as Timeout + +getClientEnvCached :: Text -> IO Client.ClientEnv +getClientEnvCached baseUrl = + modifyMVar clientEnvCache $ \cache -> + case Map.lookup baseUrl cache of + Just env -> pure (cache, env) + Nothing -> do + env <- newClientEnv baseUrl + pure (Map.insert baseUrl env cache, env) + +cachedClientEnvCount :: IO Int +cachedClientEnvCount = + modifyMVar clientEnvCache $ \cache -> pure (cache, Map.size cache) + +requestHeaders :: Text -> Model -> Options -> RequestHeaders +requestHeaders apiKey m opts = + applyHeaderOverrides + [ ("Authorization", Text.encodeUtf8 ("Bearer " <> apiKey)), + ("Accept", "text/event-stream"), + ("Content-Type", "application/json") + ] + (Map.toList (m ^. #headers) <> Map.toList (opts ^. #headers)) + +resolveKey :: Text -> Options -> IO Text +resolveKey baseUrl opts = case opts ^. #apiKey of + Just source -> Auth.resolveApiKey source + Nothing -> case Auth.defaultApiKeyEnvForBaseUrl baseUrl of + Just name -> Auth.resolveApiKey (Auth.ApiKeyEnv name) + Nothing -> + throwIO $ + authError $ + "no default API key env is known for " <> baseUrl <> "; set Options.apiKey explicitly" + +runWithTimeout :: Maybe Int -> IO () -> IO (Maybe BaikaiError) +runWithTimeout Nothing action = action >> pure Nothing +runWithTimeout (Just ms) action = do + result <- Timeout.timeout (max 0 ms * 1000) action + pure $ case result of + Just () -> Nothing + Nothing -> Just (timeoutError ms) + +timeoutError :: Int -> BaikaiError +timeoutError ms = + BaikaiError + { category = TransientError, + message = "provider stream exceeded timeoutMs=" <> Text.pack (show ms), + httpStatus = Nothing, + retryAfterSeconds = Nothing, + exitCode = Nothing + } + +newClientEnv :: Text -> IO Client.ClientEnv +newClientEnv baseUrl = do + parsed <- Client.parseBaseUrl (Text.unpack baseUrl) + manager <- + TLS.newTlsManagerWith + TLS.tlsManagerSettings + { HTTP.managerResponseTimeout = HTTP.responseTimeoutNone + } + pure (Client.mkClientEnv manager parsed) + +applyHeaderOverrides :: + RequestHeaders -> + [(Text, Text)] -> + RequestHeaders +applyHeaderOverrides = + foldl addHeader + where + addHeader headers (name, value) = + let nameBytes = Text.encodeUtf8 name + ciName = CI.mk nameBytes + in (ciName, Text.encodeUtf8 value) : filter ((/= ciName) . fst) headers + +{-# NOINLINE clientEnvCache #-} +clientEnvCache :: MVar (Map.Map Text Client.ClientEnv) +clientEnvCache = unsafePerformIO (newMVar Map.empty) diff --git a/test/ErrorClassSpec.hs b/test/ErrorClassSpec.hs --- a/test/ErrorClassSpec.hs +++ b/test/ErrorClassSpec.hs @@ -1,7 +1,7 @@ module ErrorClassSpec (tests) where -import Baikai.Error (BaikaiError (..), ErrorCategory (..)) -import Baikai.Provider.OpenAI.ErrorClass +import Baikai.Error (BaikaiError (..), ErrorCategory (..), isRetryable) +import Baikai.Provider.OpenAI.Internal.ErrorClass ( classifyErrorText, classifyException, responseToError, @@ -12,6 +12,7 @@ import Data.CaseInsensitive qualified as CI import Data.Sequence qualified as Seq import Data.Text qualified as Text +import Network.HTTP.Client qualified as HTTP import Network.HTTP.Types.Status (mkStatus) import Network.HTTP.Types.Version (http11) import Servant.Client (ResponseF (..)) @@ -21,8 +22,9 @@ tests :: TestTree tests = testGroup - "Baikai.Provider.OpenAI.ErrorClass" + "Baikai.Provider.OpenAI.Internal.ErrorClass" [ httpStatusTests, + sdkTextTests, streamedErrorTests, fallbackTests ] @@ -77,9 +79,42 @@ fallbackTests = testGroup "classifyException fallback" - [ testCase "non-ClientError exception -> OtherError, text preserved" $ do + [ testCase "http-client connection failures are transient" $ do + let e = + classifyException $ + toException $ + HTTP.HttpExceptionRequest + HTTP.defaultRequest + (HTTP.ConnectionFailure (toException (userError "reset"))) + category e @?= TransientError + assertBool "connection failure is retryable" (isRetryable e), + testCase "http-client response timeouts are transient" $ do + let e = + classifyException $ + toException $ + HTTP.HttpExceptionRequest + HTTP.defaultRequest + HTTP.ResponseTimeout + category e @?= TransientError + assertBool "response timeout is retryable" (isRetryable e), + testCase "non-ClientError exception -> OtherError, text preserved" $ do let e = classifyException (toException (userError "weird failure")) category e @?= OtherError assertBool "message keeps the original text" $ "weird failure" `Text.isInfixOf` message e + ] + +sdkTextTests :: TestTree +sdkTextTests = + testGroup + "classifyErrorText (SDK HTTP text)" + [ testCase "429 SDK text -> RateLimited with status" $ do + let parsed = + classifyErrorText + "HTTP error 429 Too Many Requests: {\"error\":{\"message\":\"Rate limit reached...\",\"type\":\"tokens\"}}" + fmap category parsed @?= Just RateLimited + fmap httpStatus parsed @?= Just (Just 429), + testCase "401 SDK text -> AuthError" $ + fmap category (classifyErrorText "HTTP error 401 Unauthorized: {\"error\":{\"message\":\"bad key\"}}") + @?= Just AuthError ] diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -1,10 +1,24 @@ module Main (main) where import Baikai +import Baikai.Cost qualified as Cost +import Baikai.Cost.Pricing (computeCost) import Baikai.Provider.OpenAI.Api + ( RawChunk (..), + closeOpenStream, + emptyAssembler, + openaiChatStream, + parseUsage, + rawUsageToUsage, + translate, + ) +import Baikai.Provider.OpenAI.Cli qualified as CodexCli import Baikai.Provider.OpenAI.Interactive +import Baikai.Provider.OpenAI.Internal.Request (mapRequest) +import Control.Exception (bracket) import Control.Lens ((&), (.~), (^.)) import Data.Aeson qualified as Aeson +import Data.Aeson.Types qualified as AesonTypes import Data.ByteString.Char8 qualified as BS8 import Data.Generics.Labels () import Data.Text qualified as Text @@ -12,9 +26,17 @@ import ErrorClassSpec qualified import OpenAI.V1.Chat.Completions qualified as Chat import OpenAI.V1.ResponseFormat qualified as RF +import ReasoningSpec qualified +import ShapeSpec qualified +import SseSpec qualified import Streamly.Data.Stream qualified as Stream +import System.Directory (getPermissions, getTemporaryDirectory, setOwnerExecutable, setPermissions) +import System.Environment (lookupEnv, setEnv, unsetEnv) +import System.FilePath (()) +import System.Timeout (timeout) import Test.Tasty (TestTree, defaultMain, testGroup) -import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) +import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, testCase, (@?=)) +import TransportSpec qualified main :: IO () main = @@ -22,11 +44,23 @@ testGroup "Baikai.Provider.OpenAI" [ commandRenderingTest, + batchCommandRenderingTest, + batchSystemPromptTest, + stderrFloodTest, + usageMappingTests, promptRenderingTest, compatDetectionTest, rejectsImageToolResultsTest, + noKeyStreamTest, + codexMissingBinaryTest, + finishReasonTests, responseFormatMappingTest, - ErrorClassSpec.tests + optionsMappingTest, + ErrorClassSpec.tests, + ReasoningSpec.tests, + ShapeSpec.tests, + SseSpec.tests, + TransportSpec.tests ] -- | A 'JsonSchema' on 'Options.responseFormat' maps onto the @@ -37,7 +71,7 @@ responseFormatMappingTest = testCase "responseFormat JsonSchema maps onto OpenAI response_format" $ do let model = - _Model + emptyModel & #modelId .~ "gpt-4o-mini" & #api .~ OpenAIChatCompletions & #provider .~ "openai" @@ -52,9 +86,9 @@ "required" Aeson..= (["name", "age"] :: [Text.Text]), "additionalProperties" Aeson..= False ] - ctx = _Context + ctx = emptyContext opts = - _Options + emptyOptions & #responseFormat .~ Just (JsonSchema {name = "person", schema = personSchema, strict = True}) case mapRequest model ctx opts of @@ -67,18 +101,130 @@ RF.description js @?= Nothing other -> assertFailure ("expected JSON_Schema, got: " <> show other) +optionsMappingTest :: TestTree +optionsMappingTest = + testCase "sampling Options map onto OpenAI request fields" $ do + let model = + emptyModel + & #modelId .~ "gpt-4o-mini" + & #api .~ OpenAIChatCompletions + & #provider .~ "openai" + opts = + emptyOptions + & #topP .~ Just 0.9 + & #stopSequences .~ Just (Vector.fromList ["END", "STOP"]) + & #seed .~ Just 7 + & #frequencyPenalty .~ Just 0.2 + & #presencePenalty .~ Just 0.3 + case mapRequest model emptyContext opts of + Left e -> assertFailure ("mapRequest failed: " <> Text.unpack e) + Right req -> do + Chat.top_p req @?= Just 0.9 + Chat.stop req @?= Just (Vector.fromList ["END", "STOP"]) + Chat.seed req @?= Just 7 + Chat.frequency_penalty req @?= Just 0.2 + Chat.presence_penalty req @?= Just 0.3 + +usageMappingTests :: TestTree +usageMappingTests = + testGroup + "usage mapping" + [ testCase "cached prompt tokens map to disjoint fields" $ do + u <- normalizedUsage cachedUsagePayload + inputTokens u @?= 20 + cacheReadTokens u @?= 80 + outputTokens u @?= 50 + reasoningTokens u @?= Just 20 + cacheWriteTokens u @?= 0 + totalTokens u @?= 150, + testCase "computeCost bills each token class exactly once" $ do + u <- normalizedUsage cachedUsagePayload + let c = computeCost usageCostModel u + -- The double-billing bug produced 358 / 1000000 by charging + -- cached tokens at both the input and cache-read rates. + Cost.usd c @?= (139 / 500000 :: Rational) + Cost.inputUsd (Cost.breakdown c) @?= (20 / 1000000 :: Rational) + Cost.cachedInputUsd (Cost.breakdown c) @?= (8 / 1000000 :: Rational), + testCase "clamps when a compatible host over-reports cached tokens" $ do + u <- normalizedUsage overCachedUsagePayload + inputTokens u @?= 0 + cacheReadTokens u @?= 120 + outputTokens u @?= 50 + totalTokens u @?= 170, + testCase "no cache details means no cache tokens" $ do + u <- normalizedUsage uncachedUsagePayload + inputTokens u @?= 100 + cacheReadTokens u @?= 0 + outputTokens u @?= 50 + reasoningTokens u @?= Nothing + totalTokens u @?= 150 + ] + +cachedUsagePayload :: Aeson.Object +cachedUsagePayload = + usageObject + [ "prompt_tokens" Aeson..= (100 :: Int), + "completion_tokens" Aeson..= (50 :: Int), + "total_tokens" Aeson..= (150 :: Int), + "prompt_tokens_details" Aeson..= Aeson.object ["cached_tokens" Aeson..= (80 :: Int)], + "completion_tokens_details" Aeson..= Aeson.object ["reasoning_tokens" Aeson..= (20 :: Int)] + ] + +overCachedUsagePayload :: Aeson.Object +overCachedUsagePayload = + usageObject + [ "prompt_tokens" Aeson..= (100 :: Int), + "completion_tokens" Aeson..= (50 :: Int), + "total_tokens" Aeson..= (150 :: Int), + "prompt_tokens_details" Aeson..= Aeson.object ["cached_tokens" Aeson..= (120 :: Int)] + ] + +uncachedUsagePayload :: Aeson.Object +uncachedUsagePayload = + usageObject + [ "prompt_tokens" Aeson..= (100 :: Int), + "completion_tokens" Aeson..= (50 :: Int), + "total_tokens" Aeson..= (150 :: Int) + ] + +usageObject :: [AesonTypes.Pair] -> Aeson.Object +usageObject pairs = + case Aeson.object pairs of + Aeson.Object o -> o + _ -> error "unreachable: Aeson.object builds an Object" + +normalizedUsage :: Aeson.Object -> IO Usage +normalizedUsage payload = + case parseUsage payload of + Just raw -> pure (rawUsageToUsage raw) + Nothing -> assertFailure "expected usage payload to parse" + +usageCostModel :: Model +usageCostModel = + emptyModel + & #modelId .~ "gpt-test" + & #api .~ OpenAIChatCompletions + & #provider .~ "openai" + & #cost + .~ ModelCost + { inputCost = 1, + outputCost = 5, + cacheReadCost = 1 / 10, + cacheWriteCost = 5 / 4 + } + commandRenderingTest :: TestTree commandRenderingTest = testCase "renders model, working directory, extra dirs, sandbox, approval, and extra args" $ do let cfg = defaultCodexInteractiveConfig { executable = "/bin/codex", - extraArgs = Vector.fromList ["--no-alt-screen"] + extraArgs = ["--no-alt-screen"] } req = - (_InteractiveLaunchRequest "inspect the repo") + (interactiveLaunchRequest "inspect the repo") & #systemPrompt .~ Just "Be precise." - & #model .~ Just "gpt-5-codex" + & #modelId .~ Just "gpt-5-codex" & #workingDir .~ Just "/work/project" & #extraDirs .~ ["/work/shared", "/work/docs"] & #safety .~ CodexSandbox CodexWorkspaceWrite CodexApprovalOnRequest @@ -99,39 +245,108 @@ "on-request", "--no-alt-screen", "--search", + "--", "System instructions:\nBe precise.\n\nUser request:\ninspect the repo" ] ) +batchCommandRenderingTest :: TestTree +batchCommandRenderingTest = + testCase "codex exec argv terminates options before a dash-leading prompt" $ do + let model = + emptyModel + & #modelId .~ "" + & #api .~ OpenAICompletionsCli + & #provider .~ "openai" + ctx = emptyContext & #messages .~ Vector.singleton (user "-begin with a dash") + CodexCli.codexCliCommand CodexCli.defaultCodexCliConfig model ctx + @?= ( "codex", + [ "exec", + "--json", + "--skip-git-repo-check", + "--ephemeral", + "--", + "-begin with a dash" + ] + ) + +batchSystemPromptTest :: TestTree +batchSystemPromptTest = + testCase "codex exec argv carries system prompt in the prompt text" $ do + let model = + emptyModel + & #modelId .~ "" + & #api .~ OpenAICompletionsCli + & #provider .~ "openai" + ctx = + emptyContext + & #systemPrompt .~ Just "Be terse." + & #messages .~ Vector.singleton (user "ping") + CodexCli.codexCliCommand CodexCli.defaultCodexCliConfig model ctx + @?= ( "codex", + [ "exec", + "--json", + "--skip-git-repo-check", + "--ephemeral", + "--", + "System instructions:\nBe terse.\n\nUser request:\nping" + ] + ) + +stderrFloodTest :: TestTree +stderrFloodTest = + testCase "codex batch provider survives a 1MiB stderr flood without deadlock" $ do + dir <- getTemporaryDirectory + let script = dir "baikai-codex-stderr-flood.sh" + writeFile script $ + unlines + [ "#!/bin/sh", + "head -c 1048576 /dev/zero | tr '\\0' 'e' >&2", + "printf '{\"type\":\"agent_message\",\"message\":\"pong\"}\\n'" + ] + perms <- getPermissions script + setPermissions script (setOwnerExecutable True perms) + reg <- newProviderRegistry + registerApiProviderWith reg (CodexCli.codexCliProvider CodexCli.defaultCodexCliConfig {CodexCli.executable = script}) + let model = + emptyModel + & #modelId .~ "" + & #api .~ OpenAICompletionsCli + & #provider .~ "openai" + ctx = emptyContext & #messages .~ Vector.singleton (user "ping") + mResp <- timeout 30000000 (completeRequestWith reg model ctx emptyOptions) + case mResp of + Nothing -> assertFailure "deadlock: stderr was not drained concurrently" + Just resp -> assistantText resp @?= "pong" + promptRenderingTest :: TestTree promptRenderingTest = testCase "omits the system-instruction wrapper when no system prompt is present" $ do - codexInteractivePrompt (_InteractiveLaunchRequest "hello") @?= "hello" + codexInteractivePrompt (interactiveLaunchRequest "hello") @?= "hello" compatDetectionTest :: TestTree compatDetectionTest = testCase "OpenAI-compatible hosts auto-detect request-shaping compat flags" $ do let model = - _Model + emptyModel & #api .~ OpenAIChatCompletions & #baseUrl .~ "https://api.deepseek.com" compat = openaiCompletionsCompatFor model compat ^. #thinkingFormat @?= ThinkingFormatDeepseek compat ^. #maxTokensField @?= MaxTokensField compat ^. #supportsStrictMode @?= False - compat ^. #supportsDeveloperRole @?= False rejectsImageToolResultsTest :: TestTree rejectsImageToolResultsTest = testCase "OpenAI API mapping rejects image tool-result blocks instead of dropping them" $ do let model = - _Model + emptyModel & #modelId .~ "gpt-test" & #api .~ OpenAIChatCompletions & #provider .~ "openai" image = ImageContent {imageData = BS8.pack "png-bytes", mimeType = "image/png"} ctx = - _Context + emptyContext & #messages .~ Vector.singleton ( ToolResultMessage @@ -140,13 +355,118 @@ toolName = "render", content = Vector.singleton (ToolResultImage image), isError = False, - timestamp = read "2026-06-05 00:00:00 UTC" + timestamp = Just (read "2026-06-05 00:00:00 UTC") } ) - events <- Stream.toList (openaiChatStream model ctx _Options) + events <- Stream.toList (openaiChatStream model ctx emptyOptions) + assertErrorContract events case events of - [EventError TerminalPayload {message = AssistantMessage AssistantPayload {errorMessage = Just msg}}] -> - assertBool - ("expected ToolResultImage error, got: " <> Text.unpack msg) - ("ToolResultImage" `Text.isInfixOf` msg) - other -> error ("expected one EventError; got: " <> show other) + [ EventStart StartPayload {}, + EventError TerminalPayload {message = AssistantMessage AssistantPayload {errorMessage = Just msg}} + ] -> + assertBool + ("expected ToolResultImage error, got: " <> Text.unpack msg) + ("ToolResultImage" `Text.isInfixOf` msg) + other -> error ("expected EventStart then EventError; got: " <> show other) + +noKeyStreamTest :: TestTree +noKeyStreamTest = + testCase "missing OPENAI_API_KEY yields one terminal EventError" $ + withUnsetEnv "OPENAI_API_KEY" $ do + let model = + emptyModel + & #modelId .~ "gpt-test" + & #api .~ OpenAIChatCompletions + & #provider .~ "openai" + events <- Stream.toList (openaiChatStream model emptyContext emptyOptions) + assertErrorContract events + case last events of + EventError TerminalPayload {errorInfo = Just be} -> + be ^. #category @?= AuthError + other -> assertFailure ("expected terminal EventError with AuthError, got: " <> show other) + +codexMissingBinaryTest :: TestTree +codexMissingBinaryTest = + testCase "codex CLI missing binary returns an error-shaped Response" $ do + reg <- newProviderRegistry + registerApiProviderWith + reg + (CodexCli.codexCliProvider CodexCli.defaultCodexCliConfig {CodexCli.executable = "/nonexistent/codex-binary"}) + let model = + emptyModel + & #modelId .~ "" + & #api .~ OpenAICompletionsCli + & #provider .~ "openai" + ctx = emptyContext & #messages .~ Vector.singleton (user "ping") + resp <- completeRequestWith reg model ctx emptyOptions + case responseError resp of + Just be -> be ^. #category @?= OtherError + Nothing -> assertFailure "expected missing binary to be returned in-band" + +finishReasonTests :: TestTree +finishReasonTests = + testGroup + "finish_reason handling" + [ testCase "content_filter terminates as EventError" $ do + let (_events1, ass1) = + translate + (Right RawChunk {contentDelta = Just "partial", reasoningDelta = Nothing, finishReason = Nothing, toolDeltas = [], usage = Nothing}) + (emptyAssembler openaiTestModel (read "2026-06-05 00:00:00 UTC")) + (read "2026-06-05 00:00:01 UTC") + (events2, ass2) = + translate + (Right RawChunk {contentDelta = Nothing, reasoningDelta = Nothing, finishReason = Just "content_filter", toolDeltas = [], usage = Nothing}) + ass1 + (read "2026-06-05 00:00:02 UTC") + (events3, _) = closeOpenStream (read "2026-06-05 00:00:03 UTC") Nothing ass2 + let terminalEvents = events2 <> events3 + assertErrorContract terminalEvents + case last terminalEvents of + EventError TerminalPayload {errorInfo = Just be} -> do + be ^. #category @?= OtherError + assertBool "message mentions content_filter" ("content_filter" `Text.isInfixOf` (be ^. #message)) + other -> assertFailure ("expected EventError for content_filter, got: " <> show other), + testCase "unknown finish_reason is a successful diagnostic" $ do + let (_events, ass1) = + translate + (Right RawChunk {contentDelta = Nothing, reasoningDelta = Nothing, finishReason = Just "mystery", toolDeltas = [], usage = Nothing}) + (emptyAssembler openaiTestModel (read "2026-06-05 00:00:00 UTC")) + (read "2026-06-05 00:00:01 UTC") + (terminalEvents, _) = closeOpenStream (read "2026-06-05 00:00:02 UTC") Nothing ass1 + case terminalEvents of + [EventDone TerminalPayload {message = AssistantMessage AssistantPayload {stopReason = Stop, errorMessage = Just msg}}] -> + msg @?= "unrecognized finish_reason: mystery" + other -> assertFailure ("expected successful diagnostic EventDone, got: " <> show other) + ] + +openaiTestModel :: Model +openaiTestModel = + emptyModel + & #modelId .~ "gpt-test" + & #api .~ OpenAIChatCompletions + & #provider .~ "openai" + +withUnsetEnv :: String -> IO a -> IO a +withUnsetEnv name action = + bracket + (lookupEnv name) + restore + (const (unsetEnv name >> action)) + where + restore = maybe (unsetEnv name) (setEnv name) + +assertErrorContract :: [AssistantMessageEvent] -> Assertion +assertErrorContract events = do + let terminals = filter isTerminal events + length terminals @?= 1 + case terminals of + [EventError TerminalPayload {errorInfo = Nothing}] -> + assertFailure "terminal EventError omitted errorInfo" + _ -> pure () + +assistantText :: Response -> Text.Text +assistantText resp = + Text.concat + [ t + | AssistantText (TextContent t) <- Vector.toList (resp ^. #message ^. #content) + ] diff --git a/test/ReasoningSpec.hs b/test/ReasoningSpec.hs new file mode 100644 --- /dev/null +++ b/test/ReasoningSpec.hs @@ -0,0 +1,255 @@ +{-# LANGUAGE LambdaCase #-} + +module ReasoningSpec (tests) where + +import Baikai +import Baikai.Models.Generated +import Baikai.Provider.OpenAI.Api + ( RawChunk (..), + closeOpenStream, + emptyAssembler, + parseChunk, + scanThinkTags, + translate, + _TagScanState, + ) +import Baikai.Provider.OpenAI.Internal.Request (mapRequest) +import Control.Lens ((&), (.~)) +import Data.Aeson qualified as Aeson +import Data.Generics.Labels () +import Data.Text qualified as Text +import Data.Time.Clock (UTCTime) +import Data.Vector qualified as Vector +import OpenAI.V1.Chat.Completions qualified as Chat +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "ReasoningSpec" + [ parseReasoningTests, + assemblyTests, + tagScannerTests, + taggedTextCompatTest, + replayDropsThinkingTest + ] + +parseReasoningTests :: TestTree +parseReasoningTests = + testGroup + "parseChunk reasoning fields" + [ testCase "DeepSeek reasoning_content delta" $ + case parseChunk (chunkValue "reasoning_content" "because") of + Right RawChunk {reasoningDelta = Just r} -> r @?= "because" + other -> assertFailure ("unexpected parse result: " <> show other), + testCase "OpenRouter reasoning delta" $ + case parseChunk (chunkValue "reasoning" "because") of + Right RawChunk {reasoningDelta = Just r} -> r @?= "because" + other -> assertFailure ("unexpected parse result: " <> show other) + ] + +assemblyTests :: TestTree +assemblyTests = + testGroup + "reasoning assembly" + [ testCase "reasoning deltas close before visible content" $ do + let chunks = + [ emptyChunk {reasoningDelta = Just "because "}, + emptyChunk {reasoningDelta = Just "therefore"}, + emptyChunk {contentDelta = Just "answer "}, + emptyChunk {contentDelta = Just "done"}, + emptyChunk {finishReason = Just "stop"} + ] + events = runChunks deepseek_deepseek_reasoner chunks + eventShape events + @?= [ "ThinkingStart:0", + "ThinkingDelta:0:because ", + "ThinkingDelta:0:therefore", + "ThinkingEnd:0:because therefore", + "TextStart:1", + "TextDelta:1:answer ", + "TextDelta:1:done", + "TextEnd:1:answer done", + "EventDone" + ] + terminalContent events + @?= Vector.fromList + [ AssistantThinking + ThinkingContent + { thinking = "because therefore", + signature = Nothing, + redacted = False + }, + AssistantText (TextContent "answer done") + ], + testCase "whole message shape yields reasoning then text" $ do + let raw = + Aeson.object + [ "choices" + Aeson..= [ Aeson.object + [ "message" + Aeson..= Aeson.object + [ "reasoning_content" Aeson..= ("because" :: Text.Text), + "content" Aeson..= ("answer" :: Text.Text) + ], + "finish_reason" Aeson..= ("stop" :: Text.Text) + ] + ] + ] + chunk <- either (assertFailure . ("parse failed: " <>)) pure (parseChunk raw) + terminalContent (runChunks deepseek_deepseek_reasoner [chunk]) + @?= Vector.fromList + [ AssistantThinking + ThinkingContent + { thinking = "because", + signature = Nothing, + redacted = False + }, + AssistantText (TextContent "answer") + ] + ] + +tagScannerTests :: TestTree +tagScannerTests = + testGroup + "scanThinkTags" + [ testCase "split tags across deltas" $ do + let (st1, p1) = scanThinkTags _TagScanState "reasoninganswer" + p1 <> p2 <> p3 @?= [Left "reasoning", Right "answer"], + testCase "literal less-than text passes through" $ do + let (_st, parts) = scanThinkTags _TagScanState "2 < 3" + parts @?= [Right "2 < 3"] + ] + +taggedTextCompatTest :: TestTree +taggedTextCompatTest = + testCase "requiresThinkingAsText gates tag extraction" $ do + let tagged = "reasoninganswer" + deepseekEvents = + runChunks + deepseek_deepseek_reasoner + [ emptyChunk {contentDelta = Just tagged}, + emptyChunk {finishReason = Just "stop"} + ] + openaiEvents = + runChunks + openai_gpt_4o_mini + [ emptyChunk {contentDelta = Just tagged}, + emptyChunk {finishReason = Just "stop"} + ] + terminalContent deepseekEvents + @?= Vector.fromList + [ AssistantThinking + ThinkingContent + { thinking = "reasoning", + signature = Nothing, + redacted = False + }, + AssistantText (TextContent "answer") + ] + terminalContent openaiEvents + @?= Vector.singleton (AssistantText (TextContent tagged)) + +replayDropsThinkingTest :: TestTree +replayDropsThinkingTest = + testCase "OpenAI-compatible replay drops AssistantThinking blocks" $ do + let msg = + AssistantMessage + AssistantPayload + { content = + Vector.fromList + [ AssistantThinking + ThinkingContent + { thinking = "internal", + signature = Nothing, + redacted = False + }, + AssistantText (TextContent "visible") + ], + usage = zeroUsage, + stopReason = Stop, + errorMessage = Nothing, + timestamp = Just testTime + } + ctx = emptyContext & #messages .~ Vector.singleton msg + case mapRequest deepseek_deepseek_reasoner ctx emptyOptions of + Left e -> assertFailure ("mapRequest failed: " <> Text.unpack e) + Right req -> case Vector.toList (requestMessages req) of + [Chat.Assistant {Chat.assistant_content = Just parts}] -> + case Vector.toList parts of + [Chat.Text {Chat.text = body}] -> body @?= "visible" + other -> assertFailure ("unexpected assistant content parts: " <> show other) + other -> assertFailure ("unexpected mapped messages: " <> show other) + +chunkValue :: Text.Text -> Text.Text -> Aeson.Value +chunkValue key value = + Aeson.object + [ "choices" + Aeson..= [ Aeson.object + [ "delta" + Aeson..= case key of + "reasoning_content" -> + Aeson.object ["reasoning_content" Aeson..= value] + "reasoning" -> + Aeson.object ["reasoning" Aeson..= value] + _ -> + Aeson.object [] + ] + ] + ] + +emptyChunk :: RawChunk +emptyChunk = + RawChunk + { contentDelta = Nothing, + reasoningDelta = Nothing, + finishReason = Nothing, + toolDeltas = [], + usage = Nothing + } + +runChunks :: Model -> [RawChunk] -> [AssistantMessageEvent] +runChunks model chunks = + let (events, ass) = + foldl + ( \(acc, st) chunk -> + let (newEvents, st') = translate (Right chunk) st testTime + in (acc <> newEvents, st') + ) + ([], emptyAssembler model testTime) + chunks + (terminalEvents, _) = closeOpenStream testTime Nothing ass + in events <> terminalEvents + +eventShape :: [AssistantMessageEvent] -> [Text.Text] +eventShape = + fmap $ \case + ThinkingStart IndexPayload {contentIndex = i} -> "ThinkingStart:" <> tshow i + ThinkingDelta DeltaPayload {contentIndex = i, delta = d} -> "ThinkingDelta:" <> tshow i <> ":" <> d + ThinkingEnd ThinkingEndPayload {contentIndex = i, content = ThinkingContent {thinking = body}} -> + "ThinkingEnd:" <> tshow i <> ":" <> body + TextStart IndexPayload {contentIndex = i} -> "TextStart:" <> tshow i + TextDelta DeltaPayload {contentIndex = i, delta = d} -> "TextDelta:" <> tshow i <> ":" <> d + TextEnd BlockEndPayload {contentIndex = i, content = body} -> "TextEnd:" <> tshow i <> ":" <> body + EventDone {} -> "EventDone" + EventError {} -> "EventError" + _ -> "other" + +terminalContent :: [AssistantMessageEvent] -> Vector.Vector AssistantContent +terminalContent events = + case last events of + EventDone TerminalPayload {message = AssistantMessage AssistantPayload {content = blocks}} -> blocks + EventError TerminalPayload {message = AssistantMessage AssistantPayload {content = blocks}} -> blocks + _ -> Vector.empty + +requestMessages :: Chat.CreateChatCompletion -> Vector.Vector (Chat.Message (Vector.Vector Chat.Content)) +requestMessages Chat.CreateChatCompletion {Chat.messages = msgs} = msgs + +tshow :: (Show a) => a -> Text.Text +tshow = Text.pack . show + +testTime :: UTCTime +testTime = read "2026-07-03 12:00:00 UTC" diff --git a/test/ShapeSpec.hs b/test/ShapeSpec.hs new file mode 100644 --- /dev/null +++ b/test/ShapeSpec.hs @@ -0,0 +1,216 @@ +{-# LANGUAGE LambdaCase #-} + +module ShapeSpec (tests) where + +import Baikai +import Baikai.Content qualified as Content +import Baikai.Models.Generated qualified as Models +import Baikai.Provider.OpenAI.Api + ( RawChunk (..), + RawToolDelta (..), + closeOpenStream, + emptyAssembler, + translate, + ) +import Baikai.Provider.OpenAI.Internal.Request (mapRequest) +import Baikai.Provider.OpenAI.Shape (streamRequestBody) +import Control.Lens ((&), (.~)) +import Data.Aeson (Value (..), (.=)) +import Data.Aeson qualified as Aeson +import Data.Aeson.Key qualified as AesonKey +import Data.Aeson.KeyMap qualified as KeyMap +import Data.Generics.Labels () +import Data.Text qualified as Text +import Data.Time.Clock (UTCTime) +import Data.Vector qualified as Vector +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "ShapeSpec" + [ deepseekShapeTest, + openRouterCacheControlTest, + strictModeGateTest, + usageStreamingGateTest, + zeroCapOmissionTest, + indexlessToolDeltaTest + ] + +deepseekShapeTest :: TestTree +deepseekShapeTest = + testCase "DeepSeek request body uses max_tokens and reasoning shape" $ do + value <- + shapedBody + Models.deepseek_deepseek_chat + (emptyOptions & #thinking .~ Just ThinkingHigh) + emptyContext + lookupTop "max_completion_tokens" value @?= Nothing + lookupTop "max_tokens" value @?= Just (Number 8192) + lookupTop "thinking" value + @?= Just (Aeson.object ["type" .= ("enabled" :: Text.Text)]) + lookupTop "reasoning_effort" value @?= Just (String "high") + +openRouterCacheControlTest :: TestTree +openRouterCacheControlTest = + testCase "OpenRouter cache marker lands on the system content part with ttl" $ do + let ctx = + emptyContext + & #systemPrompt .~ Just "cache this prefix" + & #messages .~ Vector.singleton (user "answer") + opts = emptyOptions & #cacheRetention .~ Just CacheRetentionLong + value <- shapedBody Models.openrouter_openai_gpt_4o_mini opts ctx + systemCacheControl value + @?= Just + ( Aeson.object + [ "type" .= ("ephemeral" :: Text.Text), + "ttl" .= ("1h" :: Text.Text) + ] + ) + +strictModeGateTest :: TestTree +strictModeGateTest = + testCase "supportsStrictMode gates response_format json_schema strict" $ do + let schema = Aeson.object ["type" .= ("object" :: Text.Text)] + opts = + emptyOptions + & #responseFormat + .~ Just (JsonSchema {name = "shape", schema = schema, strict = True}) + value <- shapedBody Models.deepseek_deepseek_chat opts emptyContext + lookupPath ["response_format", "json_schema", "strict"] value + @?= Nothing + +usageStreamingGateTest :: TestTree +usageStreamingGateTest = + testCase "supportsUsageInStreaming gates stream_options" $ do + let compat = defaultOpenAICompletionsCompat {supportsUsageInStreaming = False} + model = + Models.openai_gpt_4o_mini + & #compat .~ CompatOpenAICompletions compat + value <- shapedBody model emptyOptions emptyContext + lookupTop "stream" value @?= Just (Bool True) + lookupTop "stream_options" value @?= Nothing + +zeroCapOmissionTest :: TestTree +zeroCapOmissionTest = + testCase "unknown zero maxOutputTokens omits max_completion_tokens" $ do + let model = + emptyModel + & #modelId .~ "custom" + & #api .~ OpenAIChatCompletions + & #provider .~ "custom" + & #maxOutputTokens .~ 0 + req <- either (assertFailure . Text.unpack) pure (mapRequest model emptyContext emptyOptions) + lookupTop "max_completion_tokens" (Aeson.toJSON req) @?= Nothing + +indexlessToolDeltaTest :: TestTree +indexlessToolDeltaTest = + testCase "id-bearing index-less tool deltas remain separate" $ do + let chunks = + [ emptyChunk + { toolDeltas = + [ RawToolDelta + { index = Nothing, + id_ = Just "call_a", + name = Just "first", + args = Just "{\"a\":" + }, + RawToolDelta + { index = Nothing, + id_ = Just "call_b", + name = Just "second", + args = Just "{\"b\":" + } + ] + }, + emptyChunk + { toolDeltas = + [ RawToolDelta + { index = Nothing, + id_ = Just "call_a", + name = Nothing, + args = Just "1}" + }, + RawToolDelta + { index = Nothing, + id_ = Just "call_b", + name = Nothing, + args = Just "2}" + } + ] + }, + emptyChunk {finishReason = Just "tool_calls"} + ] + events = runChunks chunks + toolCalls = + [ toolCall + | ToolCallEnd ToolCallEndPayload {toolCall = toolCall} <- events + ] + fmap Content.id_ toolCalls @?= ["call_a", "call_b"] + fmap Content.name toolCalls @?= ["first", "second"] + fmap Content.arguments toolCalls + @?= [ Aeson.object ["a" .= (1 :: Int)], + Aeson.object ["b" .= (2 :: Int)] + ] + +shapedBody :: Model -> Options -> Context -> IO Value +shapedBody model opts ctx = do + req <- either (assertFailure . Text.unpack) pure (mapRequest model ctx opts) + pure (streamRequestBody (openaiCompletionsCompatFor model) opts req) + +lookupTop :: Text.Text -> Value -> Maybe Value +lookupTop field = lookupPath [field] + +lookupPath :: [Text.Text] -> Value -> Maybe Value +lookupPath [] value = Just value +lookupPath (field : rest) (Object obj) = + KeyMap.lookup (AesonKey.fromText field) obj >>= lookupPath rest +lookupPath _ _ = Nothing + +systemCacheControl :: Value -> Maybe Value +systemCacheControl value = do + Array messages <- lookupTop "messages" value + systemMessage <- + firstMay + [ msg + | msg@(Object _) <- Vector.toList messages, + lookupPath ["role"] msg == Just (String "system") + ] + Array content <- lookupPath ["content"] systemMessage + contentPart <- lastMay (Vector.toList content) + lookupPath ["cache_control"] contentPart + +firstMay :: [a] -> Maybe a +firstMay [] = Nothing +firstMay (x : _) = Just x + +lastMay :: [a] -> Maybe a +lastMay [] = Nothing +lastMay xs = Just (last xs) + +emptyChunk :: RawChunk +emptyChunk = + RawChunk + { contentDelta = Nothing, + reasoningDelta = Nothing, + finishReason = Nothing, + toolDeltas = [], + usage = Nothing + } + +runChunks :: [RawChunk] -> [AssistantMessageEvent] +runChunks chunks = + let (events, ass) = + foldl + ( \(acc, st) chunk -> + let (newEvents, st') = translate (Right chunk) st testTime + in (acc <> newEvents, st') + ) + ([], emptyAssembler Models.openai_gpt_4o_mini testTime) + chunks + (terminalEvents, _) = closeOpenStream testTime Nothing ass + in events <> terminalEvents + +testTime :: UTCTime +testTime = read "2026-07-03 12:00:00 UTC" diff --git a/test/SseSpec.hs b/test/SseSpec.hs new file mode 100644 --- /dev/null +++ b/test/SseSpec.hs @@ -0,0 +1,58 @@ +module SseSpec (tests) where + +import Baikai.Error (ErrorCategory (..), category, httpStatus, retryAfterSeconds) +import Baikai.Provider.OpenAI.Sse (sseFromResponse) +import Data.Aeson qualified as Aeson +import Data.ByteString (ByteString) +import Data.CaseInsensitive qualified as CI +import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef) +import Network.HTTP.Client.Internal qualified as HTTP +import Network.HTTP.Types.Status (mkStatus) +import Network.HTTP.Types.Version (http11) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Baikai.Provider.OpenAI.Sse" + [ testCase "non-2xx response preserves Retry-After and status" $ do + eventsRef <- newIORef [] + resp <- mkResponse 429 [("Retry-After", "9")] ["{\"error\":{\"message\":\"rate limited\",\"type\":\"tokens\"}}"] + sseFromResponse resp (\ev -> modifyIORef' eventsRef (<> [ev])) + events <- readIORef eventsRef + case events of + [Left e] -> do + category e @?= RateLimited + retryAfterSeconds e @?= Just 9 + httpStatus e @?= Just 429 + other -> assertFailure ("expected one classified error, got: " <> show other), + testCase "[DONE] terminates without emitting a JSON event" $ do + eventsRef <- newIORef [] + resp <- mkResponse 200 [] ["data: {\"choices\":[]}\n\n", "data: [DONE]\n\n", "data: {\"ignored\":true}\n\n"] + sseFromResponse resp (\ev -> modifyIORef' eventsRef (<> [ev])) + events <- readIORef eventsRef + case events of + [Right (Aeson.Object _)] -> pure () + other -> assertFailure ("expected one JSON event before [DONE], got: " <> show other) + ] + +mkResponse :: Int -> [(ByteString, ByteString)] -> [ByteString] -> IO (HTTP.Response HTTP.BodyReader) +mkResponse status headers chunks = do + ref <- newIORef chunks + let bodyReader = do + remaining <- readIORef ref + case remaining of + [] -> pure "" + (x : xs) -> writeIORef ref xs >> pure x + pure + HTTP.Response + { HTTP.responseStatus = mkStatus status "", + HTTP.responseVersion = http11, + HTTP.responseHeaders = [(CI.mk k, v) | (k, v) <- headers], + HTTP.responseBody = bodyReader, + HTTP.responseCookieJar = HTTP.createCookieJar [], + HTTP.responseClose' = HTTP.ResponseClose (pure ()), + HTTP.responseOriginalRequest = HTTP.defaultRequest, + HTTP.responseEarlyHints = [] + } diff --git a/test/TransportSpec.hs b/test/TransportSpec.hs new file mode 100644 --- /dev/null +++ b/test/TransportSpec.hs @@ -0,0 +1,81 @@ +module TransportSpec (tests) where + +import Baikai +import Baikai.Provider.OpenAI.Transport qualified as Transport +import Control.Concurrent (threadDelay) +import Control.Exception (bracket, try) +import Control.Lens ((&), (.~), (^.)) +import Data.CaseInsensitive qualified as CI +import Data.Map.Strict qualified as Map +import Data.Text qualified as Text +import Data.Text.Encoding qualified as Text +import Network.HTTP.Types.Header (RequestHeaders) +import System.Environment (lookupEnv, setEnv, unsetEnv) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Baikai.Provider.OpenAI.Transport" + [ clientEnvCacheTest, + requestHeadersTest, + timeoutTest, + unknownHostKeyTest + ] + +clientEnvCacheTest :: TestTree +clientEnvCacheTest = + testCase "cached ClientEnv is allocated once for a base URL" $ do + let url = "https://cache-openai.test" + before <- Transport.cachedClientEnvCount + _ <- Transport.getClientEnvCached url + afterFirst <- Transport.cachedClientEnvCount + _ <- Transport.getClientEnvCached url + afterSecond <- Transport.cachedClientEnvCount + afterFirst @?= before + 1 + afterSecond @?= afterFirst + +requestHeadersTest :: TestTree +requestHeadersTest = + testCase "model and option headers reach the wire, options winning case-insensitively" $ do + let model = + emptyModel + & #headers .~ Map.fromList [("X-Trace", "model"), ("authorization", "model-auth")] + opts = + emptyOptions + & #headers .~ Map.fromList [("x-trace", "option"), ("Authorization", "option-auth")] + headers = Transport.requestHeaders "secret" model opts + header "X-Trace" headers @?= Just "option" + header "authorization" headers @?= Just "option-auth" + header "Accept" headers @?= Just "text/event-stream" + +timeoutTest :: TestTree +timeoutTest = + testCase "runWithTimeout classifies an elapsed whole-call timeout as transient" $ do + result <- Transport.runWithTimeout (Just 1) (threadDelay 100000) + case result of + Just be -> do + be ^. #category @?= TransientError + "timeoutMs=1" `Text.isInfixOf` (be ^. #message) @?= True + Nothing -> assertFailure "expected timeout error" + +unknownHostKeyTest :: TestTree +unknownHostKeyTest = + testCase "unknown hosts do not fall back to OPENAI_API_KEY" $ + withEnv "OPENAI_API_KEY" "openai-secret" $ do + result <- try (Transport.resolveKey "https://unknown.example" emptyOptions) :: IO (Either BaikaiError Text.Text) + case result of + Left be -> be ^. #category @?= AuthError + Right _ -> assertFailure "expected AuthError for unknown host" + +header :: Text.Text -> RequestHeaders -> Maybe Text.Text +header name headers = + Text.decodeUtf8 <$> lookup (CI.mk (Text.encodeUtf8 name)) headers + +withEnv :: String -> String -> IO a -> IO a +withEnv name value = + bracket + (lookupEnv name <* setEnv name value) + (maybe (unsetEnv name) (setEnv name)) + . const