diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2026 Nadeem Bitar + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Nadeem Bitar nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/baikai-openai.cabal b/baikai-openai.cabal new file mode 100644 --- /dev/null +++ b/baikai-openai.cabal @@ -0,0 +1,74 @@ +cabal-version: 3.4 +name: baikai-openai +version: 0.1.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. + +category: AI +license: BSD-3-Clause +license-file: LICENSE +author: Nadeem Bitar +maintainer: nadeem@gmail.com +copyright: (c) 2026 Nadeem Bitar +build-type: Simple + +common common-options + ghc-options: + -Wall -Wcompat -Widentities -Wincomplete-uni-patterns + -Wincomplete-record-updates -Wredundant-constraints + -fhide-source-paths -Wmissing-export-lists -Wpartial-fields + -Wmissing-deriving-strategies + + default-language: GHC2024 + default-extensions: + DeriveAnyClass + DuplicateRecordFields + OverloadedLabels + OverloadedStrings + +library + import: common-options + hs-source-dirs: src + exposed-modules: + Baikai.Provider.OpenAI.Api + Baikai.Provider.OpenAI.Cli + Baikai.Provider.OpenAI.Interactive + + build-depends: + , aeson + , baikai ^>=0.1.0 + , base >=4.20 && <5 + , base64-bytestring + , bytestring + , containers + , generic-lens + , http-client + , http-client-tls + , lens ^>=5.3 + , openai + , process + , servant-client + , streamly >=0.11 && <0.13 + , streamly-core >=0.3 && <0.5 + , text ^>=2.1 + , time + , vector + +test-suite baikai-openai-test + import: common-options + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + build-depends: + , baikai ^>=0.1.0 + , baikai-openai + , base >=4.20 && <5 + , bytestring + , generic-lens + , lens ^>=5.3 + , streamly-core >=0.3 && <0.5 + , tasty + , tasty-hunit + , text ^>=2.1 + , vector diff --git a/src/Baikai/Provider/OpenAI/Api.hs b/src/Baikai/Provider/OpenAI/Api.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Provider/OpenAI/Api.hs @@ -0,0 +1,913 @@ +{-# LANGUAGE LambdaCase #-} + +-- | Provider wrapping the @openai@ package's Chat Completions API. +-- +-- Call 'register' once (typically from @main@) to install the +-- 'Baikai.Api.OpenAIChatCompletions' handler into the baikai +-- provider registry. After registration, any 'Baikai.Model.Model' +-- whose 'Baikai.Api.api' tag is 'OpenAIChatCompletions' dispatches +-- through this handler. +-- +-- The handler resolves 'Baikai.Options.apiKey' when present, falling +-- back to the @OPENAI_API_KEY@ env var via +-- 'Baikai.Auth.resolveApiKey'. +-- +-- 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. +-- +-- The synchronous 'complete' field is derived via +-- 'streamingComplete', so callers that drain the stream get the +-- same fully-assembled 'Response' they had before. +module Baikai.Provider.OpenAI.Api + ( register, + registerWithRegistry, + openaiChatStream, + ) +where + +import Baikai.Api (Api (..)) +import Baikai.Auth qualified as Auth +import Baikai.Compat + ( OpenAICompletionsCompat (..), + ThinkingFormat (..), + ) +import Baikai.Content qualified as Content +import Baikai.Context (Context (..)) +import Baikai.Cost (_Cost) +import Baikai.Cost.Pricing qualified as Pricing +import Baikai.Message qualified as Msg +import Baikai.Model (Model, openaiCompletionsCompatFor) +import Baikai.Options (Options (..)) +import Baikai.Provider.Registry + ( ApiProvider (..), + ProviderRegistry, + globalProviderRegistry, + registerApiProviderWith, + ) +import Baikai.StopReason qualified as Stop +import Baikai.Stream (streamingComplete) +import Baikai.Stream.Event + ( AssistantMessageEvent (..), + BlockEndPayload (..), + DeltaPayload (..), + IndexPayload (..), + StartPayload (..), + TerminalPayload (..), + ToolCallEndPayload (..), + ) +import Baikai.ThinkingLevel + ( ThinkingLevel (..), + ) +import Baikai.Tool qualified as Tool +import Baikai.Usage qualified as Usage +import Control.Concurrent (forkIO) +import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan) +import Control.Exception (SomeException, displayException, 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.Maybe (fromMaybe) +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as Text +import Data.Time.Clock (UTCTime, getCurrentTime) +import Data.Vector (Vector) +import Data.Vector qualified as Vector +import GHC.Generics (Generic) +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.Tool qualified as OpenAITool +import OpenAI.V1.ToolCall qualified as ToolCall +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 + +-- | 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 + } + +-- | Streaming producer for the OpenAI Chat Completions API. +-- +-- Forks one worker thread per call that drives +-- 'OpenAI.createChatCompletionStream' (the raw 'Aeson.Value' +-- variant, not the typed one — see module docs for why). The +-- worker pushes raw chunk values onto a 'Chan' terminated by +-- 'Nothing'; the consumer translates each chunk into zero or more +-- baikai 'AssistantMessageEvent' values and terminates with exactly +-- one 'EventDone' or 'EventError'. +openaiChatStream :: + Model -> Context -> Options -> Stream IO AssistantMessageEvent +openaiChatStream m ctx opts = + Stream.concatEffect $ do + setup <- prepareCall m ctx opts + case setup of + Left err -> pure (Stream.fromEffect (immediateError err)) + Right call -> do + ch <- newChan :: IO (Chan (Maybe RawChunk)) + tref <- newIORef False + _ <- forkIO (worker call ch) + startTime <- getCurrentTime + let initialState = + ProducerState + { chan = ch, + pending = [EventStart StartPayload {partial = skeletonStart m startTime}], + assembler = emptyAssembler m startTime, + finished = False, + terminalRef = tref + } + pure (Stream.unfoldrM step initialState) + +skeletonStart :: Model -> UTCTime -> Msg.Message +skeletonStart _m start = + Msg.AssistantMessage + Msg.AssistantPayload + { Msg.content = Vector.empty, + Msg.usage = Usage._Usage, + Msg.stopReason = Stop.Stop, + Msg.errorMessage = Nothing, + Msg.timestamp = start + } + +-- | Per-call prepared values. +data OpenAICall = OpenAICall + { methods :: !OpenAI.Methods, + request :: !Chat.CreateChatCompletion + } + deriving stock (Generic) + +prepareCall :: Model -> Context -> Options -> IO (Either Text OpenAICall) +prepareCall m ctx opts = case mapRequest m ctx opts of + Left e -> pure (Left 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 + } + } + 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 +-- ignored. Missing fields are 'Nothing' (we tolerate partial +-- tool-call deltas). +data RawChunk = RawChunk + { contentDelta :: !(Maybe Text), + finishReason :: !(Maybe Text), + toolDeltas :: ![RawToolDelta], + usage :: !(Maybe RawUsage), + error :: !(Maybe Text) + } + deriving stock (Show, Generic) + +data RawToolDelta = RawToolDelta + { index :: !Int, + id_ :: !(Maybe Text), + name :: !(Maybe Text), + args :: !(Maybe Text) + } + deriving stock (Show, Generic) + +data RawUsage = RawUsage + { inputTokens :: !Natural, + outputTokens :: !Natural, + cacheReadTokens :: !Natural, + reasoningTokens :: !(Maybe Natural) + } + deriving stock (Show, Generic) + +worker :: + OpenAICall -> Chan (Maybe RawChunk) -> IO () +worker call ch = do + let OpenAI.Methods {OpenAI.createChatCompletionStream = stream'} = call ^. #methods + 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) + case r of + Right () -> pure () + Left e -> + writeChan ch (Just (errorChunk (Text.pack (displayException 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 + choices <- o .:? "choices" + let firstChoice :: Maybe Aeson.Object + firstChoice = case choices of + Just (Aeson.Array a) + | Vector.length a > 0 -> + case Vector.head a of + Aeson.Object obj -> Just obj + _ -> Nothing + _ -> Nothing + (contentDelta, finishR, toolDeltas) <- case firstChoice of + Nothing -> pure (Nothing, Nothing, []) + Just ch -> do + finish <- ch .:? "finish_reason" + delta <- ch .:? "delta" + case delta of + Nothing -> pure (Nothing, finish, []) + Just (Aeson.Object dObj) -> do + cd <- dObj .:? "content" + tc <- dObj .:? "tool_calls" + let tds = parseToolCallDeltas tc + pure (cd, finish, tds) + _ -> pure (Nothing, finish, []) + usageM <- o .:? "usage" + let ru = case usageM of + Just (Aeson.Object uObj) -> parseUsage uObj + _ -> Nothing + pure + RawChunk + { contentDelta = contentDelta, + finishReason = finishR, + toolDeltas = toolDeltas, + usage = ru, + error = Nothing + } + +parseToolCallDeltas :: Maybe Value -> [RawToolDelta] +parseToolCallDeltas = \case + Just (Aeson.Array v) -> Vector.toList (Vector.mapMaybe oneDelta v) + _ -> [] + where + oneDelta :: Value -> Maybe RawToolDelta + oneDelta = \case + Aeson.Object o -> + let funcObj :: Maybe Aeson.Object + funcObj = case lookupField "function" o of + Just (Aeson.Object f) -> Just f + _ -> Nothing + getName = funcObj >>= lookupText "name" + getArgs = funcObj >>= lookupText "arguments" + in Just + RawToolDelta + { index = maybe 0 fromInt (lookupField "index" o), + id_ = lookupText "id" o, + name = getName, + args = getArgs + } + _ -> Nothing + +parseUsage :: Aeson.Object -> Maybe RawUsage +parseUsage o = + case Aeson.parseEither pUsage o of + Right u -> Just u + Left _ -> Nothing + where + pUsage obj = do + i <- obj .:? "prompt_tokens" + out <- obj .:? "completion_tokens" + ptd <- obj .:? "prompt_tokens_details" + ctd <- obj .:? "completion_tokens_details" + let cached = case ptd of + Just (Aeson.Object p) -> case lookupField "cached_tokens" p of + Just (Aeson.Number n) -> truncate n + _ -> 0 :: Natural + _ -> 0 + reasoning = case ctd of + Just (Aeson.Object c) -> case lookupField "reasoning_tokens" c of + Just (Aeson.Number n) -> Just (truncate n) + _ -> Nothing + _ -> Nothing + pure + RawUsage + { inputTokens = fromMaybe 0 i, + outputTokens = fromMaybe 0 out, + cacheReadTokens = cached, + reasoningTokens = reasoning + } + +lookupField :: Text -> Aeson.Object -> Maybe Value +lookupField k = KeyMap.lookup (AesonKey.fromText k) + +-- Pull a Text-valued field out of an Aeson object; tolerates +-- absent or non-Text values by returning 'Nothing'. +lookupText :: Text -> Aeson.Object -> Maybe Text +lookupText k o = case lookupField k o of + Just (Aeson.String t) -> Just t + _ -> Nothing + +fromInt :: Value -> Int +fromInt = \case + Aeson.Number n -> truncate n + _ -> 0 + +-- ============================================================ +-- Streamly state machine +-- ============================================================ + +data ProducerState = ProducerState + { chan :: !(Chan (Maybe RawChunk)), + pending :: ![AssistantMessageEvent], + assembler :: !Assembler, + finished :: !Bool, + terminalRef :: !(IORef Bool) + } + deriving stock (Generic) + +step :: ProducerState -> IO (Maybe (AssistantMessageEvent, ProducerState)) +step s + | (e : rest) <- s ^. #pending = do + writeTerminal s e + pure + ( Just + ( e, + s + & #pending .~ rest + & #finished .~ (s ^. #finished || terminal e) + ) + ) + | s ^. #finished = pure Nothing + | otherwise = do + mRaw <- readChan (s ^. #chan) + case mRaw of + Nothing -> do + alreadyTerminal <- readIORef (s ^. #terminalRef) + if alreadyTerminal + then pure Nothing + else do + now <- getCurrentTime + let (events, ass') = closeOpenStream now (s ^. #assembler) + case events of + [] -> pure Nothing + (e : rest) -> do + writeTerminal s e + pure + ( Just + ( e, + s + & #pending .~ rest + & #assembler .~ ass' + & #finished .~ True + ) + ) + Just raw -> do + now <- getCurrentTime + let (events, ass') = translate raw (s ^. #assembler) now + case events of + [] -> step (s & #assembler .~ ass') + (e : rest) -> do + writeTerminal s e + pure + ( Just + ( e, + s + & #pending .~ rest + & #assembler .~ ass' + & #finished .~ (s ^. #finished || terminal e) + ) + ) + +writeTerminal :: ProducerState -> AssistantMessageEvent -> IO () +writeTerminal s ev + | terminal ev = writeIORef (s ^. #terminalRef) True + | otherwise = pure () + +terminal :: AssistantMessageEvent -> Bool +terminal = \case + EventDone {} -> True + EventError {} -> True + _ -> False + +-- ============================================================ +-- Translation +-- ============================================================ + +-- | Translation state across one streaming call. +data Assembler = Assembler + { model :: !Model, + start :: !UTCTime, + -- | 'Just i' when a text block at baikai contentIndex @i@ is + -- currently open; 'Nothing' when no text block is open. + textOpen :: !(Maybe Int), + textAccum :: !Text, + textEverOpened :: !Bool, + -- | Maps OpenAI's per-call tool-call index to baikai's + -- 'contentIndex'. + toolIndexMap :: !(IntMap Int), + -- | baikai contentIndex → (id, name). + toolMeta :: !(IntMap (Text, Text)), + -- | baikai contentIndex → accumulated arguments JSON. + toolArgs :: !(IntMap Text), + closed :: !(IntMap Content.AssistantContent), + nextContentIndex :: !Int, + usage :: !Usage.Usage, + stopReason :: !Stop.StopReason, + -- | 'True' once a chunk carrying @finish_reason@ has been + -- observed. The terminal 'EventDone' fires on channel close so + -- the post-@finish_reason@ usage chunk (when @include_usage@ is + -- enabled) has a chance to land. + finishSeen :: !Bool, + errorMsg :: !(Maybe Text) + } + deriving stock (Generic) + +emptyAssembler :: Model -> UTCTime -> Assembler +emptyAssembler m s = + Assembler + { model = m, + start = s, + textOpen = Nothing, + textAccum = Text.empty, + textEverOpened = False, + toolIndexMap = IntMap.empty, + toolMeta = IntMap.empty, + toolArgs = IntMap.empty, + closed = IntMap.empty, + nextContentIndex = 0, + usage = Usage._Usage, + stopReason = Stop.Stop, + finishSeen = False, + errorMsg = Nothing + } + +translate :: + RawChunk -> + Assembler -> + UTCTime -> + ([AssistantMessageEvent], Assembler) +translate chunk ass now + | Just errMsg <- chunk ^. #error = + let msg = finalMessage ass now (Just errMsg) Stop.ErrorReason + in ( [EventError TerminalPayload {reason = Stop.ErrorReason, message = msg}], + 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 + -- 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) + +applyContentDelta :: + Maybe Text -> Assembler -> ([AssistantMessageEvent], Assembler) +applyContentDelta Nothing ass = ([], ass) +applyContentDelta (Just "") ass = ([], ass) +applyContentDelta (Just d) ass = + case ass ^. #textOpen of + Just i -> + ( [TextDelta DeltaPayload {contentIndex = i, delta = d}], + ass & #textAccum %~ (<> d) + ) + Nothing -> + let i = ass ^. #nextContentIndex + in ( [TextStart IndexPayload {contentIndex = i}, TextDelta DeltaPayload {contentIndex = i, delta = d}], + ass + & #textOpen .~ Just i + & #textAccum .~ d + & #textEverOpened .~ True + & #nextContentIndex .~ (i + 1) + ) + +applyToolDeltas :: + [RawToolDelta] -> Assembler -> ([AssistantMessageEvent], Assembler) +applyToolDeltas deltas ass = foldl' apply ([], ass) deltas + where + apply (acc, a) d = + let (events, a') = applyOneToolDelta d a + in (acc <> events, a') + +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) + Nothing -> + let i = ass ^. #nextContentIndex + ass' = + ass + & #toolIndexMap %~ IntMap.insert openaiIdx i + & #toolMeta %~ IntMap.insert i ("", "") + & #toolArgs %~ IntMap.insert i Text.empty + & #nextContentIndex .~ (i + 1) + in (i, ass', True) + -- Update metadata (id/name first delta only). + ass2 = + ass1 + & #toolMeta + %~ IntMap.adjust + ( \(existingId, existingName) -> + ( maybe existingId (\x -> if Text.null existingId then x else existingId) (d ^. #id_), + maybe existingName (\x -> if Text.null existingName then x else existingName) (d ^. #name) + ) + ) + baikaiIdx + -- Append args if present. + argsDelta = fromMaybe "" (d ^. #args) + ass3 = ass2 & #toolArgs %~ IntMap.adjust (<> argsDelta) baikaiIdx + events0 = if opened then [ToolCallStart IndexPayload {contentIndex = baikaiIdx}] else [] + events1 = + if Text.null argsDelta + then events0 + else events0 <> [ToolCallDelta DeltaPayload {contentIndex = baikaiIdx, delta = argsDelta}] + in (events1, ass3) + +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' + +-- | 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) + +-- | Close the open text block, if any, by emitting a 'TextEnd' and +-- storing the assembled content in 'closed'. +closeOpenText :: Assembler -> ([AssistantMessageEvent], Assembler) +closeOpenText ass = case ass ^. #textOpen of + Nothing -> ([], ass) + Just i -> + let body = ass ^. #textAccum + block = Content.AssistantText (Content.TextContent body) + in ( [TextEnd BlockEndPayload {contentIndex = i, content = body}], + ass + & #textOpen .~ Nothing + & #textAccum .~ Text.empty + & #closed %~ IntMap.insert i block + ) + +-- | Close every open tool call by emitting 'ToolCallEnd' (with the +-- fully parsed 'ToolCall') in index order. +closeOpenTools :: Assembler -> ([AssistantMessageEvent], Assembler) +closeOpenTools ass = + let openTools = IntMap.toAscList (ass ^. #toolArgs) + (events, ass') = foldl' closeOne ([], ass) openTools + in (events, ass') + where + closeOne (acc, a) (i, argsText) = + let (tid, tn) = fromMaybe ("", "") (IntMap.lookup i (a ^. #toolMeta)) + decoded :: Value + decoded = case Aeson.eitherDecodeStrict (Text.encodeUtf8 argsText) of + Right v -> v + Left _ -> Aeson.Object mempty + tc = + Content.ToolCall + { Content.id_ = tid, + Content.name = tn, + Content.arguments = decoded + } + block = Content.AssistantToolCall tc + in ( acc <> [ToolCallEnd ToolCallEndPayload {contentIndex = i, toolCall = tc}], + a + & #closed %~ IntMap.insert i block + & #toolArgs %~ IntMap.delete i + & #toolMeta %~ IntMap.delete i + ) + +closeOpenStream :: + UTCTime -> Assembler -> ([AssistantMessageEvent], Assembler) +closeOpenStream now ass + | ass ^. #finishSeen = + -- Channel closed cleanly after finish_reason. Emit + -- EventDone with the accumulated content + usage. + let reason = ass ^. #stopReason + msg = finalMessage ass now Nothing reason + in ([EventDone TerminalPayload {reason = reason, message = msg}], ass) + | otherwise = + -- Channel closed without a finish_reason. Force-close any + -- still-open blocks and emit EventError with the accumulated + -- content. + let (closeText, ass1) = closeOpenText ass + (closeTools, ass2) = closeOpenTools ass1 + reason = Stop.ErrorReason + msg = + finalMessage + ass2 + now + (Just "openai stream ended without finish_reason") + reason + errEv = EventError TerminalPayload {reason = reason, message = msg} + in (closeText <> closeTools <> [errEv], ass2) + +finalMessage :: + Assembler -> UTCTime -> Maybe Text -> Stop.StopReason -> Msg.Message +finalMessage ass now errMsg sr = + let blocks = blocksInOrder ass + m = ass ^. #model + usageBare = ass ^. #usage + computed = Pricing.computeCost m usageBare + usage' = usageBare & #cost .~ computed + in Msg.AssistantMessage + Msg.AssistantPayload + { Msg.content = blocks, + Msg.usage = usage', + Msg.stopReason = sr, + Msg.errorMessage = errMsg, + Msg.timestamp = 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 + now <- getCurrentTime + let msg = + Msg.AssistantMessage + Msg.AssistantPayload + { Msg.content = Vector.empty, + Msg.usage = Usage._Usage, + Msg.stopReason = Stop.ErrorReason, + Msg.errorMessage = Just errText, + Msg.timestamp = now + } + pure (EventError TerminalPayload {reason = Stop.ErrorReason, message = msg}) + +-- ============================================================ +-- 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) + 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 + } + +-- | 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 + ) + +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" + +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 diff --git a/src/Baikai/Provider/OpenAI/Cli.hs b/src/Baikai/Provider/OpenAI/Cli.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Provider/OpenAI/Cli.hs @@ -0,0 +1,183 @@ +-- | Provider that drives the @codex exec --json@ non-interactive CLI +-- as a subprocess. +-- +-- Call 'register' once (typically from @main@) to install the +-- 'Baikai.Api.OpenAICompletionsCli' handler with default config. +-- 'registerWith' accepts a caller-supplied 'CodexCliConfig'. +module Baikai.Provider.OpenAI.Cli + ( CodexCliConfig (..), + defaultCodexCliConfig, + register, + registerWith, + registerWithRegistry, + registerWithRegistryAndConfig, + ) +where + +import Baikai.Api (Api (..)) +import Baikai.Content (AssistantContent (..), TextContent (..)) +import Baikai.Context (Context) +import Baikai.Error (BaikaiError (..)) +import Baikai.Message (AssistantPayload (..)) +import Baikai.Model (Model) +import Baikai.Options (Options) +import Baikai.Provider.Cli.Internal qualified as Internal +import Baikai.Provider.Registry + ( ApiProvider (..), + ProviderRegistry, + globalProviderRegistry, + 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 Control.Lens ((^.)) +import Data.ByteString qualified as BS +import Data.Generics.Labels () +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.Process qualified as P + +-- | Configuration for the @codex exec --json@ subprocess. +data CodexCliConfig = CodexCliConfig + { executable :: !FilePath, + extraArgs :: !(Vector Text), + workingDir :: !(Maybe FilePath), + skipGitRepoCheck :: !Bool, + ephemeral :: !Bool + } + deriving stock (Eq, Show, Generic) + +defaultCodexCliConfig :: CodexCliConfig +defaultCodexCliConfig = + CodexCliConfig + { executable = "codex", + extraArgs = mempty, + workingDir = Nothing, + skipGitRepoCheck = True, + ephemeral = True + } + +-- | Install the Codex CLI handler with 'defaultCodexCliConfig'. +register :: IO () +register = registerWith defaultCodexCliConfig + +-- | Install the Codex CLI handler with a caller-supplied config. +-- +-- The Codex binary runs in batch mode. 'stream' wraps the batch +-- output in a synthetic one-shot event stream +-- (@EventStart, TextStart 0, TextDelta 0 body, TextEnd 0, EventDone@) +-- emitted after the subprocess exits. 'complete' stays on the +-- direct batch path so it preserves 'Response.latencyMs' rather than +-- recomputing it from synthetic event timestamps. EP-3's Decision +-- Log records the deviation from "complete = streamingComplete . +-- stream". +registerWith :: CodexCliConfig -> IO () +registerWith = registerWithRegistryAndConfig globalProviderRegistry + +-- | Install the Codex CLI handler with 'defaultCodexCliConfig' into an explicit +-- registry. +registerWithRegistry :: ProviderRegistry -> IO () +registerWithRegistry reg = registerWithRegistryAndConfig reg defaultCodexCliConfig + +-- | Install the Codex CLI handler with a caller-supplied config into an +-- explicit registry. +registerWithRegistryAndConfig :: ProviderRegistry -> CodexCliConfig -> IO () +registerWithRegistryAndConfig reg cfg = + registerApiProviderWith + reg + ApiProvider + { apiTag = OpenAICompletionsCli, + stream = liftCompleteToStream (runCodexCli cfg), + complete = runCodexCli cfg + } + +modelArgs :: Model -> [String] +modelArgs m = case Text.strip (m ^. #modelId) of + "" -> [] + mid -> ["--model", Text.unpack mid] + +handleStream :: Handle -> Stream IO BS.ByteString +handleStream h = Stream.unfoldrM step () + where + step _ = do + chunk <- BS.hGetSome h 4096 + if BS.null chunk + then pure Nothing + else pure (Just (chunk, ())) + +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] + procSpec = + (P.proc (cfg ^. #executable) baseArgs) + { 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 + +consume :: + UTCTime -> + Model -> + (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 + } + +millisBetween :: UTCTime -> UTCTime -> Integer +millisBetween a b = round (realToFrac (diffUTCTime b a) * (1000 :: Double)) diff --git a/src/Baikai/Provider/OpenAI/Interactive.hs b/src/Baikai/Provider/OpenAI/Interactive.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Provider/OpenAI/Interactive.hs @@ -0,0 +1,122 @@ +-- | Launch real interactive Codex sessions from Baikai's +-- provider-neutral interactive request type. +-- +-- This module is intentionally separate from +-- "Baikai.Provider.OpenAI.Cli": that module drives @codex exec@ as a +-- batch completion provider, while this module starts the interactive +-- terminal UI and returns only after the CLI exits. +module Baikai.Provider.OpenAI.Interactive + ( CodexInteractiveConfig (..), + defaultCodexInteractiveConfig, + codexInteractiveCommand, + codexInteractivePrompt, + launchCodexInteractive, + ) +where + +import Baikai.Interactive + ( CodexApprovalPolicy, + CodexSandboxMode, + InteractiveLaunchRequest (..), + InteractiveLaunchResult, + InteractiveProvider (..), + InteractiveSafety (..), + renderCodexApprovalPolicy, + renderCodexSandboxMode, + _InteractiveLaunchResult, + ) +import Baikai.Prelude +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) + } + deriving stock (Eq, Show, Generic) + +defaultCodexInteractiveConfig :: CodexInteractiveConfig +defaultCodexInteractiveConfig = + CodexInteractiveConfig + { executable = "codex", + extraArgs = mempty + } + +-- | Render the executable and arguments for an interactive Codex +-- launch. The final positional argument is the initial prompt. +codexInteractiveCommand :: + CodexInteractiveConfig -> InteractiveLaunchRequest -> (FilePath, [String]) +codexInteractiveCommand cfg req = + ( cfg ^. #executable, + modelArgs req + <> workingDirArgs req + <> extraDirArgs req + <> safetyArgs req + <> fmap Text.unpack (Vector.toList (cfg ^. #extraArgs)) + <> fmap Text.unpack (req ^. #extraArgs) + <> [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 + ] + +-- | Launch Codex with inherited stdin, stdout, and stderr so the +-- local CLI owns the interactive terminal experience. +launchCodexInteractive :: + CodexInteractiveConfig -> InteractiveLaunchRequest -> IO InteractiveLaunchResult +launchCodexInteractive cfg req = do + let (exe, args) = codexInteractiveCommand cfg req + spec = + (P.proc exe args) + { P.std_in = P.Inherit, + P.std_out = P.Inherit, + P.std_err = P.Inherit, + P.cwd = req ^. #workingDir + } + (_, _, _, ph) <- P.createProcess spec + code <- P.waitForProcess ph + pure (_InteractiveLaunchResult InteractiveCodex code) + +modelArgs :: InteractiveLaunchRequest -> [String] +modelArgs req = case Text.strip <$> req ^. #model of + Nothing -> [] + Just "" -> [] + Just mid -> ["--model", Text.unpack mid] + +workingDirArgs :: InteractiveLaunchRequest -> [String] +workingDirArgs req = case req ^. #workingDir of + Nothing -> [] + Just dir -> ["--cd", dir] + +extraDirArgs :: InteractiveLaunchRequest -> [String] +extraDirArgs req = + concatMap (\dir -> ["--add-dir", dir]) (req ^. #extraDirs) + +safetyArgs :: InteractiveLaunchRequest -> [String] +safetyArgs req = case req ^. #safety of + CodexSandbox sandbox approval -> codexSafetyArgs sandbox approval + DefaultSafety -> [] + ClaudeAllowedTools _ -> [] + +codexSafetyArgs :: CodexSandboxMode -> CodexApprovalPolicy -> [String] +codexSafetyArgs sandbox approval = + [ "--sandbox", + Text.unpack (renderCodexSandboxMode sandbox), + "--ask-for-approval", + Text.unpack (renderCodexApprovalPolicy approval) + ] diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,108 @@ +module Main (main) where + +import Baikai +import Baikai.Provider.OpenAI.Api +import Baikai.Provider.OpenAI.Interactive +import Control.Lens ((&), (.~), (^.)) +import Data.ByteString.Char8 qualified as BS8 +import Data.Generics.Labels () +import Data.Text qualified as Text +import Data.Vector qualified as Vector +import Streamly.Data.Stream qualified as Stream +import Test.Tasty (TestTree, defaultMain, testGroup) +import Test.Tasty.HUnit (assertBool, testCase, (@?=)) + +main :: IO () +main = + defaultMain $ + testGroup + "Baikai.Provider.OpenAI.Interactive" + [ commandRenderingTest, + promptRenderingTest, + compatDetectionTest, + rejectsImageToolResultsTest + ] + +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"] + } + req = + (_InteractiveLaunchRequest "inspect the repo") + & #systemPrompt .~ Just "Be precise." + & #model .~ Just "gpt-5-codex" + & #workingDir .~ Just "/work/project" + & #extraDirs .~ ["/work/shared", "/work/docs"] + & #safety .~ CodexSandbox CodexWorkspaceWrite CodexApprovalOnRequest + & #extraArgs .~ ["--search"] + codexInteractiveCommand cfg req + @?= ( "/bin/codex", + [ "--model", + "gpt-5-codex", + "--cd", + "/work/project", + "--add-dir", + "/work/shared", + "--add-dir", + "/work/docs", + "--sandbox", + "workspace-write", + "--ask-for-approval", + "on-request", + "--no-alt-screen", + "--search", + "System instructions:\nBe precise.\n\nUser request:\ninspect the repo" + ] + ) + +promptRenderingTest :: TestTree +promptRenderingTest = + testCase "omits the system-instruction wrapper when no system prompt is present" $ do + codexInteractivePrompt (_InteractiveLaunchRequest "hello") @?= "hello" + +compatDetectionTest :: TestTree +compatDetectionTest = + testCase "OpenAI-compatible hosts auto-detect request-shaping compat flags" $ do + let model = + _Model + & #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 + & #modelId .~ "gpt-test" + & #api .~ OpenAIChatCompletions + & #provider .~ "openai" + image = ImageContent {imageData = BS8.pack "png-bytes", mimeType = "image/png"} + ctx = + _Context + & #messages + .~ Vector.singleton + ( ToolResultMessage + ToolResultPayload + { toolCallId = "call_1", + toolName = "render", + content = Vector.singleton (ToolResultImage image), + isError = False, + timestamp = read "2026-06-05 00:00:00 UTC" + } + ) + events <- Stream.toList (openaiChatStream model ctx _Options) + 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)