diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,43 @@ +# Changelog + +All notable changes to baikai are recorded here. + +The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and +this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +## [Unreleased] + +## [baikai 0.1.0.0] - 2026-06-04 + +### Added + +- Initial release: unified Haskell interface for working with multiple AI + providers. Core modules including `Baikai`, `Baikai.Prelude`, `Baikai.Api`, + `Baikai.Provider`, `Baikai.Provider.Registry`, `Baikai.Response`, + `Baikai.Stream`, `Baikai.Tool`, `Baikai.Trace`, and the cost/usage modules. +- Depends on released `streamly` (`>=0.11 && <0.13`) and `streamly-core` + (`>=0.3 && <0.5`) from Hackage, so all dependencies resolve from Hackage. + +## [baikai-claude 0.1.0.0] - 2026-06-04 + +### Added + +- Initial release: Anthropic Claude providers for the baikai abstraction, + wrapping the `claude` package for both the Anthropic API and the `claude -p` + CLI (`Baikai.Provider.Claude.Api`, `.Cli`, `.Interactive`). + +## [baikai-openai 0.1.0.0] - 2026-06-04 + +### Added + +- Initial release: OpenAI providers for the baikai abstraction, wrapping the + `openai` package for OpenAI's Chat Completions API + (`Baikai.Provider.OpenAI.Api`, `.Cli`, `.Interactive`). + +## [baikai-trace-otel 0.1.0.0] - 2026-06-04 + +### Added + +- Initial release: OpenTelemetry `TraceSink` adapter for baikai + (`Baikai.Trace.Sink.OpenTelemetry`), emitting one OTel span per provider call + with GenAI semantic-convention attributes plus baikai cost and latency. 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.cabal b/baikai.cabal new file mode 100644 --- /dev/null +++ b/baikai.cabal @@ -0,0 +1,129 @@ +cabal-version: 3.4 +name: baikai +version: 0.1.0.0 +synopsis: Unified Haskell interface for multiple AI providers +description: + baikai provides a unified, provider-agnostic Haskell interface for working + with multiple AI providers. It abstracts chat and streaming completions, tool + use, model catalogs, cost accounting, and tracing behind a single API, with + concrete provider implementations supplied by companion packages such as + @baikai-claude@ and @baikai-openai@. + +category: AI +license: BSD-3-Clause +license-file: LICENSE +author: Nadeem Bitar +maintainer: nadeem@gmail.com +copyright: (c) 2026 Nadeem Bitar +build-type: Simple +extra-doc-files: CHANGELOG.md + +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 + Baikai.AgentAssets + Baikai.Api + Baikai.Auth + Baikai.CacheRetention + Baikai.Compat + Baikai.Content + Baikai.Context + Baikai.Cost + Baikai.Cost.Log + Baikai.Cost.Pricing + Baikai.Error + Baikai.Interactive + Baikai.Message + Baikai.Model + Baikai.Models.Generated + Baikai.Options + Baikai.Prelude + Baikai.Provider + Baikai.Provider.Cli.Internal + Baikai.Provider.Registry + Baikai.Response + Baikai.StopReason + Baikai.Stream + Baikai.Stream.Event + Baikai.ThinkingLevel + Baikai.Tool + Baikai.Trace + Baikai.Trace.Event + Baikai.Trace.Sink + Baikai.Usage + + build-depends: + , aeson + , base >=4.20 && <5 + , base64-bytestring + , bytestring + , containers + , generic-lens + , lens ^>=5.3 + , scientific + , streamly >=0.11 && <0.13 + , streamly-core >=0.3 && <0.5 + , text ^>=2.1 + , time + , unliftio-core + , vector + +executable baikai-gen-models + import: common-options + hs-source-dirs: gen + main-is: GenModels.hs + build-depends: + , aeson + , baikai + , base >=4.20 && <5 + , bytestring + , containers + , directory + , filepath + , scientific + , text ^>=2.1 + +test-suite baikai-test + import: common-options + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + other-modules: + AgentAssetsSpec + CatalogSpec + CostSpec + InteractiveSpec + TraceSpec + + build-tool-depends: baikai:baikai-gen-models + build-depends: + , aeson + , baikai + , base + , bytestring + , directory + , filepath + , process + , stm + , streamly-core >=0.3 && <0.5 + , tasty + , tasty-hunit + , temporary + , text + , vector diff --git a/gen/GenModels.hs b/gen/GenModels.hs new file mode 100644 --- /dev/null +++ b/gen/GenModels.hs @@ -0,0 +1,524 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedRecordDot #-} +{-# LANGUAGE OverloadedStrings #-} + +-- | Generator executable that reads JSON catalog files from +-- @baikai\/data\/models\/@ and emits a single +-- @baikai\/src\/Baikai\/Models\/Generated.hs@ module containing one +-- fully populated 'Baikai.Model.Model' value per (provider, model id) +-- pair. +-- +-- Run as @cabal run baikai-gen-models@ from anywhere inside the repo +-- tree. The default input and output paths are resolved relative to +-- the directory containing @baikai.cabal@, located by walking up from +-- the current working directory (and one step into a @baikai@ subdir +-- if needed), so the natural invocation from the repo root works +-- without any flags. Options: +-- +-- * @--models-dir PATH@ — directory containing catalog @.json@ files. +-- Default: @\\/data\/models@. +-- * @--out PATH@ — output Haskell source path. +-- Default: @\\/src\/Baikai\/Models\/Generated.hs@. +-- +-- Explicit paths are interpreted relative to the caller's CWD (or +-- absolute), not the package directory. +-- +-- The generator is intentionally deterministic: entries are sorted by +-- their generated Haskell identifier, the output formatting is fixed, +-- and re-running with unchanged inputs produces byte-identical output +-- (enforced by @CatalogSpec@ in @baikai\/test\/@). +module Main (main) where + +import Baikai.Api (Api (..), parseApi) +import Baikai.Compat + ( AnthropicMessagesCompat (..), + CacheControlFormat (..), + MaxTokensField (..), + OpenAICompletionsCompat (..), + ThinkingFormat (..), + defaultAnthropicMessagesCompat, + defaultOpenAICompletionsCompat, + ) +import Baikai.Model (InputModality (..)) +import Control.Monad (forM) +import Data.Aeson (FromJSON (..), (.!=), (.:), (.:?)) +import Data.Aeson qualified as Aeson +import Data.Aeson.Types (Parser, typeMismatch) +import Data.ByteString.Lazy qualified as BSL +import Data.List (sort, sortOn) +import Data.Ratio (denominator, numerator) +import Data.Scientific (Scientific) +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.IO qualified as TIO +import Numeric.Natural (Natural) +import System.Directory (doesFileExist, getCurrentDirectory, listDirectory) +import System.Environment (getArgs) +import System.Exit (die) +import System.FilePath (takeDirectory, takeExtension, ()) + +main :: IO () +main = do + args <- getArgs + (mModelsDir, mOutputPath) <- parseArgs args + (modelsDir, outputPath) <- resolveDefaults mModelsDir mOutputPath + jsonFiles <- listCatalogFiles modelsDir + catalogs <- forM jsonFiles $ \f -> do + raw <- BSL.readFile (modelsDir f) + case Aeson.eitherDecode raw of + Left err -> die $ f <> ": " <> err + Right c -> pure c + let allEntries = concatMap flattenEntries catalogs + sorted = sortOn fst allEntries + rendered = renderModule sorted + TIO.writeFile outputPath rendered + putStrLn $ + "Wrote " + <> outputPath + <> " (" + <> show (length sorted) + <> " enabled models)" + +-- | Parse command-line arguments. Recognises @--models-dir@ and +-- @--out@; everything else is a hard error. 'Nothing' for either +-- field means \"use the package-dir-anchored default\". +parseArgs :: [String] -> IO (Maybe FilePath, Maybe FilePath) +parseArgs = go Nothing Nothing + where + go d o [] = pure (d, o) + go _ o ("--models-dir" : v : rest) = go (Just v) o rest + go d _ ("--out" : v : rest) = go d (Just v) rest + go _ _ (a : _) = die $ "baikai-gen-models: unknown argument " <> show a + +-- | Resolve any unset path defaults against the @baikai@ package +-- directory so the generator works from anywhere inside the repo. +-- Explicit overrides are returned unchanged (CWD-relative or +-- absolute). If both paths are explicit no lookup happens. +resolveDefaults :: Maybe FilePath -> Maybe FilePath -> IO (FilePath, FilePath) +resolveDefaults (Just m) (Just o) = pure (m, o) +resolveDefaults mModelsDir mOutputPath = do + mPkg <- locatePackageDir + case mPkg of + Nothing -> + die + "baikai-gen-models: could not locate baikai.cabal in any \ + \ancestor directory; pass --models-dir and --out explicitly" + Just pkg -> + pure + ( maybe (pkg "data/models") id mModelsDir, + maybe (pkg "src/Baikai/Models/Generated.hs") id mOutputPath + ) + +-- | Locate the directory containing @baikai.cabal@ by walking up from +-- the current working directory. Also checks for a @baikai@ subdir at +-- each level so the executable works when invoked from the repo root +-- (where @baikai.cabal@ lives at @baikai\/baikai.cabal@). +locatePackageDir :: IO (Maybe FilePath) +locatePackageDir = getCurrentDirectory >>= walk + where + walk dir = do + here <- doesFileExist (dir "baikai.cabal") + if here + then pure (Just dir) + else do + sub <- doesFileExist (dir "baikai" "baikai.cabal") + if sub + then pure (Just (dir "baikai")) + else + let parent = takeDirectory dir + in if parent == dir then pure Nothing else walk parent + +-- | List the @.json@ files in @dir@, sorted alphabetically so the +-- output is deterministic across filesystems. +listCatalogFiles :: FilePath -> IO [FilePath] +listCatalogFiles dir = do + entries <- listDirectory dir + pure $ sort [f | f <- entries, takeExtension f == ".json"] + +-- * Catalog data types ----------------------------------------------- + +-- | One catalog file, in source form. +data CatalogFile = CatalogFile + { provider :: !Text, + baseUrl :: !Text, + api :: !Api, + compat :: !CatalogCompat, + models :: ![ModelEntry] + } + +instance FromJSON CatalogFile where + parseJSON = Aeson.withObject "CatalogFile" $ \o -> + CatalogFile + <$> o .: "provider" + <*> o .: "baseUrl" + <*> (parseApi <$> o .: "api") + <*> o .: "compat" + <*> o .: "models" + +-- | A compat directive in the catalog. @"auto"@ defers to EP-5's +-- @baseUrl@-driven auto-detection (rendered as 'CompatNone'). The two +-- structured constructors carry a full override record. +data CatalogCompat + = CatalogCompatAuto + | CatalogCompatOpenAI !OpenAICompletionsCompat + | CatalogCompatAnthropic !AnthropicMessagesCompat + deriving stock (Show) + +instance FromJSON CatalogCompat where + parseJSON = \case + Aeson.String "auto" -> pure CatalogCompatAuto + Aeson.Object o -> do + kind <- o .: "kind" :: Parser Text + case kind of + "openai-completions" -> + CatalogCompatOpenAI <$> parseOpenAICompat o + "anthropic-messages" -> + CatalogCompatAnthropic <$> parseAnthropicCompat o + _ -> + fail $ "CatalogCompat: unknown kind " <> show kind + v -> typeMismatch "CatalogCompat (expected \"auto\" or {\"kind\": ...})" v + +parseOpenAICompat :: Aeson.Object -> Parser OpenAICompletionsCompat +parseOpenAICompat o = do + let d = defaultOpenAICompletionsCompat + mtf <- optionalField o "maxTokensField" parseMaxTokensField (maxTokensField d) + sdr <- o .:? "supportsDeveloperRole" .!= supportsDeveloperRole d + sst <- o .:? "supportsStrictMode" .!= supportsStrictMode d + rtat <- o .:? "requiresThinkingAsText" .!= requiresThinkingAsText d + tf <- optionalField o "thinkingFormat" parseThinkingFormat (thinkingFormat d) + ccf <- optionalMaybeField o "cacheControlFormat" parseCacheControlFormat (cacheControlFormat d) + sus <- o .:? "supportsUsageInStreaming" .!= d.supportsUsageInStreaming + slcr <- o .:? "supportsLongCacheRetention" .!= d.supportsLongCacheRetention + pure + OpenAICompletionsCompat + { maxTokensField = mtf, + supportsDeveloperRole = sdr, + supportsStrictMode = sst, + requiresThinkingAsText = rtat, + thinkingFormat = tf, + cacheControlFormat = ccf, + supportsUsageInStreaming = sus, + supportsLongCacheRetention = slcr + } + +parseAnthropicCompat :: Aeson.Object -> Parser AnthropicMessagesCompat +parseAnthropicCompat o = do + let d = defaultAnthropicMessagesCompat + slcr <- o .:? "supportsLongCacheRetention" .!= d.supportsLongCacheRetention + scot <- o .:? "supportsCacheControlOnTools" .!= d.supportsCacheControlOnTools + seti <- o .:? "supportsEagerToolInputStreaming" .!= d.supportsEagerToolInputStreaming + ssah <- o .:? "sendSessionAffinityHeaders" .!= d.sendSessionAffinityHeaders + pure + AnthropicMessagesCompat + { supportsLongCacheRetention = slcr, + supportsCacheControlOnTools = scot, + supportsEagerToolInputStreaming = seti, + sendSessionAffinityHeaders = ssah + } + +parseMaxTokensField :: Text -> Parser MaxTokensField +parseMaxTokensField = \case + "max_tokens" -> pure MaxTokensField + "max_completion_tokens" -> pure MaxCompletionTokensField + t -> fail $ "unknown maxTokensField: " <> Text.unpack t + +parseThinkingFormat :: Text -> Parser ThinkingFormat +parseThinkingFormat = \case + "openai" -> pure ThinkingFormatOpenAI + "openrouter" -> pure ThinkingFormatOpenRouter + "deepseek" -> pure ThinkingFormatDeepseek + "together" -> pure ThinkingFormatTogether + "zai" -> pure ThinkingFormatZai + "qwen" -> pure ThinkingFormatQwen + "none" -> pure ThinkingFormatNone + t -> fail $ "unknown thinkingFormat: " <> Text.unpack t + +parseCacheControlFormat :: Text -> Parser CacheControlFormat +parseCacheControlFormat = \case + "anthropic" -> pure CacheControlFormatAnthropic + t -> fail $ "unknown cacheControlFormat: " <> Text.unpack t + +-- | Look up an optional field; missing returns @def@. +optionalField :: Aeson.Object -> Aeson.Key -> (Text -> Parser a) -> a -> Parser a +optionalField o key parser def = do + mv <- o .:? key + case mv of + Nothing -> pure def + Just v -> parser v + +-- | Look up an optional @Maybe@-typed field. JSON @null@ or absence +-- yields @def@; otherwise the value is parsed through @parser@ and +-- wrapped in 'Just'. +optionalMaybeField :: + Aeson.Object -> + Aeson.Key -> + (Text -> Parser a) -> + Maybe a -> + Parser (Maybe a) +optionalMaybeField o key parser def = do + mv <- o .:? key + case mv of + Nothing -> pure def + Just Aeson.Null -> pure Nothing + Just (Aeson.String t) -> Just <$> parser t + Just v -> typeMismatch "string or null" v + +-- | One model row in a catalog file. +data ModelEntry = ModelEntry + { entryId :: !Text, + entryName :: !Text, + entryReasoning :: !Bool, + entryInput :: ![InputModality], + entryCost :: !CostEntry, + entryContextWindow :: !Natural, + entryMaxOutputTokens :: !Natural, + entryEnabled :: !Bool, + entryCompatOverride :: !(Maybe CatalogCompat) + } + +instance FromJSON ModelEntry where + parseJSON = Aeson.withObject "ModelEntry" $ \o -> + ModelEntry + <$> o .: "id" + <*> o .: "name" + <*> o .:? "reasoning" .!= False + <*> (o .: "input" >>= traverse parseInputModality) + <*> o .: "cost" + <*> o .: "contextWindow" + <*> o .: "maxOutputTokens" + <*> o .:? "enabled" .!= True + <*> o .:? "compat" + +parseInputModality :: Text -> Parser InputModality +parseInputModality = \case + "text" -> pure InputText + "image" -> pure InputImage + t -> fail $ "unknown input modality: " <> Text.unpack t + +-- | Per-million-token cost rates. Parsed as 'Scientific' so the +-- round-trip into 'Rational' produces a small, canonical fraction +-- (@1.5@ → @3 % 2@) rather than an arbitrary IEEE-754 approximation. +data CostEntry = CostEntry + { costInput :: !Scientific, + costOutput :: !Scientific, + costCacheRead :: !Scientific, + costCacheWrite :: !Scientific + } + +instance FromJSON CostEntry where + parseJSON = Aeson.withObject "CostEntry" $ \o -> + CostEntry + <$> o .: "input" + <*> o .: "output" + <*> o .:? "cacheRead" .!= 0 + <*> o .:? "cacheWrite" .!= 0 + +-- * Flattening --------------------------------------------------------- + +-- | One generated Haskell identifier plus the 'Model'-shaped record +-- it should be rendered as. Keeping these together lets the renderer +-- stay a pure transformation over a flat list. +data GeneratedEntry = GeneratedEntry + { genIdent :: !Text, + genModelId :: !Text, + genName :: !Text, + genApi :: !Api, + genProvider :: !Text, + genBaseUrl :: !Text, + genReasoning :: !Bool, + genInput :: ![InputModality], + genCost :: !CostEntry, + genContextWindow :: !Natural, + genMaxOutputTokens :: !Natural, + genCompat :: !CatalogCompat + } + +flattenEntries :: CatalogFile -> [(Text, GeneratedEntry)] +flattenEntries c = + [ (genIdent g, g) + | m <- models c, + entryEnabled m, + let g = + GeneratedEntry + { genIdent = + sanitizeIdentifier (provider c <> "_" <> entryId m), + genModelId = entryId m, + genName = entryName m, + genApi = api c, + genProvider = provider c, + genBaseUrl = baseUrl c, + genReasoning = entryReasoning m, + genInput = entryInput m, + genCost = entryCost m, + genContextWindow = entryContextWindow m, + genMaxOutputTokens = entryMaxOutputTokens m, + genCompat = + case entryCompatOverride m of + Just c' -> c' + Nothing -> compat c + } + ] + +-- | Replace any non-identifier character with @_@. Haskell allows +-- letters, digits, underscore, and apostrophe; everything else +-- (slash, dash, dot, colon, …) becomes an underscore. +sanitizeIdentifier :: Text -> Text +sanitizeIdentifier = Text.map replace + where + replace c + | (c >= 'a' && c <= 'z') + || (c >= 'A' && c <= 'Z') + || (c >= '0' && c <= '9') + || c == '_' + || c == '\'' = + c + | otherwise = '_' + +-- * Source rendering -------------------------------------------------- + +renderModule :: [(Text, GeneratedEntry)] -> Text +renderModule entries = + Text.unlines + ( header + ++ concatMap (renderEntry . snd) entries + ) + where + header = + [ "-- AUTO-GENERATED by baikai-gen-models. Do not edit by hand.", + "-- Regenerate with: cabal run baikai-gen-models", + "{-# LANGUAGE OverloadedStrings #-}", + "{-# OPTIONS_GHC -Wno-missing-export-lists #-}", + "{-# OPTIONS_GHC -Wno-unused-imports #-}", + "", + "module Baikai.Models.Generated where", + "", + "import Baikai.Api (Api (..))", + "import Baikai.Compat", + " ( AnthropicMessagesCompat (..)", + " , CacheControlFormat (..)", + " , MaxTokensField (..)", + " , OpenAICompletionsCompat (..)", + " , ThinkingFormat (..)", + " )", + "import Baikai.Model", + " ( Compat (..)", + " , InputModality (..)", + " , Model (..)", + " , ModelCost (..)", + " )", + "import Data.Map.Strict qualified as Map", + "import Data.Ratio ((%))", + "" + ] + +renderEntry :: GeneratedEntry -> [Text] +renderEntry g = + [ genIdent g <> " :: Model", + genIdent g <> " =", + " Model", + " { modelId = " <> renderText (genModelId g), + " , name = " <> renderText (genName g), + " , api = " <> renderApiCtor (genApi g), + " , provider = " <> renderText (genProvider g), + " , baseUrl = " <> renderText (genBaseUrl g), + " , reasoning = " <> renderBool (genReasoning g), + " , input = " <> renderInputList (genInput g), + " , cost = " <> renderCost (genCost g), + " , contextWindow = " <> Text.pack (show (genContextWindow g)), + " , maxOutputTokens = " <> Text.pack (show (genMaxOutputTokens g)), + " , headers = Map.empty", + " , compat = " <> renderCompat (genCompat g), + " }", + "" + ] + +renderText :: Text -> Text +renderText t = + "\"" + <> Text.replace "\"" "\\\"" (Text.replace "\\" "\\\\" t) + <> "\"" + +renderBool :: Bool -> Text +renderBool True = "True" +renderBool False = "False" + +renderInputList :: [InputModality] -> Text +renderInputList ms = + "[" <> Text.intercalate ", " (map renderInput ms) <> "]" + where + renderInput InputText = "InputText" + renderInput InputImage = "InputImage" + +renderApiCtor :: Api -> Text +renderApiCtor = \case + OpenAIChatCompletions -> "OpenAIChatCompletions" + AnthropicMessages -> "AnthropicMessages" + OpenAICompletionsCli -> "OpenAICompletionsCli" + AnthropicMessagesCli -> "AnthropicMessagesCli" + Custom t -> "Custom " <> renderText t + +renderCost :: CostEntry -> Text +renderCost c = + Text.intercalate + "\n" + [ "ModelCost", + " { inputCost = " <> renderRational (toRational (costInput c)), + " , outputCost = " <> renderRational (toRational (costOutput c)), + " , cacheReadCost = " <> renderRational (toRational (costCacheRead c)), + " , cacheWriteCost = " <> renderRational (toRational (costCacheWrite c)), + " }" + ] + +renderRational :: Rational -> Text +renderRational r = + Text.pack (show (numerator r)) <> " % " <> Text.pack (show (denominator r)) + +renderCompat :: CatalogCompat -> Text +renderCompat = \case + CatalogCompatAuto -> "CompatNone" + CatalogCompatOpenAI c -> + Text.intercalate + "\n" + [ "CompatOpenAICompletions", + " OpenAICompletionsCompat", + " { maxTokensField = " <> renderMaxTokensField c.maxTokensField, + " , supportsDeveloperRole = " <> renderBool c.supportsDeveloperRole, + " , supportsStrictMode = " <> renderBool c.supportsStrictMode, + " , requiresThinkingAsText = " <> renderBool c.requiresThinkingAsText, + " , thinkingFormat = " <> renderThinkingFormat c.thinkingFormat, + " , cacheControlFormat = " <> renderMaybeCacheControl c.cacheControlFormat, + " , supportsUsageInStreaming = " <> renderBool c.supportsUsageInStreaming, + " , supportsLongCacheRetention = " <> renderBool c.supportsLongCacheRetention, + " }" + ] + CatalogCompatAnthropic c -> + Text.intercalate + "\n" + [ "CompatAnthropicMessages", + " AnthropicMessagesCompat", + " { supportsLongCacheRetention = " <> renderBool c.supportsLongCacheRetention, + " , supportsCacheControlOnTools = " <> renderBool c.supportsCacheControlOnTools, + " , supportsEagerToolInputStreaming = " <> renderBool c.supportsEagerToolInputStreaming, + " , sendSessionAffinityHeaders = " <> renderBool c.sendSessionAffinityHeaders, + " }" + ] + +renderMaxTokensField :: MaxTokensField -> Text +renderMaxTokensField = \case + MaxCompletionTokensField -> "MaxCompletionTokensField" + MaxTokensField -> "MaxTokensField" + +renderThinkingFormat :: ThinkingFormat -> Text +renderThinkingFormat = \case + ThinkingFormatOpenAI -> "ThinkingFormatOpenAI" + ThinkingFormatOpenRouter -> "ThinkingFormatOpenRouter" + ThinkingFormatDeepseek -> "ThinkingFormatDeepseek" + ThinkingFormatTogether -> "ThinkingFormatTogether" + ThinkingFormatZai -> "ThinkingFormatZai" + ThinkingFormatQwen -> "ThinkingFormatQwen" + ThinkingFormatNone -> "ThinkingFormatNone" + +renderMaybeCacheControl :: Maybe CacheControlFormat -> Text +renderMaybeCacheControl = \case + Nothing -> "Nothing" + Just CacheControlFormatAnthropic -> "Just CacheControlFormatAnthropic" diff --git a/src/Baikai.hs b/src/Baikai.hs new file mode 100644 --- /dev/null +++ b/src/Baikai.hs @@ -0,0 +1,57 @@ +-- | Public surface of the baikai library. Importing this module +-- gives a downstream consumer everything needed to declare a +-- 'Model', build a 'Context' and 'Options', and dispatch a call +-- through the registered handler. +module Baikai + ( -- * Types + module Baikai.AgentAssets, + module Baikai.Api, + module Baikai.Auth, + module Baikai.Model, + module Baikai.Content, + module Baikai.StopReason, + module Baikai.Message, + module Baikai.Tool, + module Baikai.Context, + module Baikai.Options, + module Baikai.Response, + module Baikai.Usage, + module Baikai.Cost, + module Baikai.Error, + module Baikai.Interactive, + + -- * Per-API compat shims and call-time options + module Baikai.Compat, + module Baikai.CacheRetention, + module Baikai.ThinkingLevel, + + -- * Provider registry + module Baikai.Provider, + + -- * Streaming + module Baikai.Stream, + module Baikai.Stream.Event, + ) +where + +import Baikai.AgentAssets +import Baikai.Api +import Baikai.Auth +import Baikai.CacheRetention +import Baikai.Compat +import Baikai.Content +import Baikai.Context +import Baikai.Cost +import Baikai.Error +import Baikai.Interactive +import Baikai.Message +import Baikai.Model +import Baikai.Options +import Baikai.Provider +import Baikai.Response +import Baikai.StopReason +import Baikai.Stream +import Baikai.Stream.Event +import Baikai.ThinkingLevel +import Baikai.Tool +import Baikai.Usage diff --git a/src/Baikai/AgentAssets.hs b/src/Baikai/AgentAssets.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/AgentAssets.hs @@ -0,0 +1,145 @@ +-- | Provider-native filesystem layouts for local agent assets. +-- +-- Baikai only describes where providers discover assets. Callers own +-- cloning, copying, updating, status reporting, and uninstalling kit +-- content. +module Baikai.AgentAssets + ( AgentAssetProvider, + AgentAssetScope, + AgentAssetKind (..), + AgentAssetFormat (..), + AgentAssetLayout (..), + CodexCustomAgent (..), + skillAsset, + customAgentAsset, + agentAssetLayout, + skillTargetPath, + agentTargetPath, + agentAssetFormat, + codexCustomAgentToml, + ) +where + +import Baikai.Interactive + ( InteractiveProvider (..), + InteractiveScope (..), + ) +import Baikai.Prelude +import Data.Text qualified as Text + +-- | Asset helpers use the same provider vocabulary as interactive +-- launchers: Claude Code and Codex are the local provider families. +type AgentAssetProvider = InteractiveProvider + +-- | User scope means a home-directory discovery path. Project scope +-- means a path relative to the working project. +type AgentAssetScope = InteractiveScope + +data AgentAssetKind + = SkillAsset + | CustomAgentAsset + deriving stock (Eq, Ord, Show, Generic) + +data AgentAssetFormat + = DirectoryAsset + | MarkdownFile + | TomlFile + deriving stock (Eq, Ord, Show, Generic) + +data AgentAssetLayout = AgentAssetLayout + { provider :: !AgentAssetProvider, + scope :: !AgentAssetScope, + kind :: !AgentAssetKind, + format :: !AgentAssetFormat, + path :: !FilePath + } + deriving stock (Eq, Show, Generic) + +data CodexCustomAgent = CodexCustomAgent + { name :: !Text, + description :: !Text, + developerInstructions :: !Text + } + deriving stock (Eq, Show, Generic) + +skillAsset :: + AgentAssetProvider -> AgentAssetScope -> FilePath -> AgentAssetLayout +skillAsset provider scope name = + agentAssetLayout provider scope SkillAsset name + +customAgentAsset :: + AgentAssetProvider -> AgentAssetScope -> FilePath -> AgentAssetLayout +customAgentAsset provider scope name = + agentAssetLayout provider scope CustomAgentAsset name + +agentAssetLayout :: + AgentAssetProvider -> + AgentAssetScope -> + AgentAssetKind -> + FilePath -> + AgentAssetLayout +agentAssetLayout provider scope kind name = + AgentAssetLayout + { provider, + scope, + kind, + format = agentAssetFormat provider kind, + path = case kind of + SkillAsset -> skillTargetPath provider scope name + CustomAgentAsset -> agentTargetPath provider scope name + } + +skillTargetPath :: AgentAssetProvider -> AgentAssetScope -> FilePath -> FilePath +skillTargetPath InteractiveClaude InteractiveProjectScope name = + joinPath [".claude", "skills", name] +skillTargetPath InteractiveClaude InteractiveUserScope name = + joinPath ["$HOME", ".claude", "skills", name] +skillTargetPath InteractiveCodex InteractiveProjectScope name = + joinPath [".agents", "skills", name] +skillTargetPath InteractiveCodex InteractiveUserScope name = + joinPath ["$HOME", ".agents", "skills", name] + +agentTargetPath :: AgentAssetProvider -> AgentAssetScope -> FilePath -> FilePath +agentTargetPath InteractiveClaude InteractiveProjectScope name = + joinPath [".claude", "agents", name <> ".md"] +agentTargetPath InteractiveClaude InteractiveUserScope name = + joinPath ["$HOME", ".claude", "agents", name <> ".md"] +agentTargetPath InteractiveCodex InteractiveProjectScope name = + joinPath [".codex", "agents", name <> ".toml"] +agentTargetPath InteractiveCodex InteractiveUserScope name = + joinPath ["$HOME", ".codex", "agents", name <> ".toml"] + +agentAssetFormat :: AgentAssetProvider -> AgentAssetKind -> AgentAssetFormat +agentAssetFormat _ SkillAsset = DirectoryAsset +agentAssetFormat InteractiveClaude CustomAgentAsset = MarkdownFile +agentAssetFormat InteractiveCodex CustomAgentAsset = TomlFile + +-- | Render the minimal TOML shape Codex custom agents expect. +codexCustomAgentToml :: CodexCustomAgent -> Text +codexCustomAgentToml agent = + Text.unlines + [ "name = " <> tomlString (agent ^. #name), + "description = " <> tomlString (agent ^. #description), + "developer_instructions = " <> tomlMultilineString (agent ^. #developerInstructions) + ] + +joinPath :: [FilePath] -> FilePath +joinPath [] = "" +joinPath (first : rest) = foldl' appendSegment first rest + where + appendSegment acc segment = acc <> "/" <> segment + +tomlString :: Text -> Text +tomlString t = + "\"" <> Text.concatMap escape t <> "\"" + where + escape '"' = "\\\"" + escape '\\' = "\\\\" + escape '\n' = "\\n" + escape '\r' = "\\r" + escape '\t' = "\\t" + escape c = Text.singleton c + +tomlMultilineString :: Text -> Text +tomlMultilineString t = + "\"\"\"\n" <> Text.replace "\"\"\"" "\\\"\\\"\\\"" t <> "\n\"\"\"" diff --git a/src/Baikai/Api.hs b/src/Baikai/Api.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Api.hs @@ -0,0 +1,59 @@ +{-# LANGUAGE LambdaCase #-} + +-- | The 'Api' tag — a typed discriminator that identifies which +-- upstream API a 'Baikai.Model.Model' speaks. +-- +-- 'Api' is a closed sum with an open 'Custom' escape hatch. The +-- closed constructors enable exhaustiveness checking inside the +-- library when a new built-in API is added; the 'Custom' constructor +-- preserves the open-world property so third-party callers can +-- register handlers under any tag without modifying baikai itself. +-- +-- The wire form is a single kebab-cased string +-- (@anthropic-messages@, @openai-chat-completions@, …). Round-trips +-- between 'Api' and 'Text' go through 'renderApi' and 'parseApi'; +-- any unknown tag becomes 'Custom !Text'. +module Baikai.Api + ( Api (..), + renderApi, + parseApi, + ) +where + +import Data.Aeson (FromJSON (parseJSON), ToJSON (toJSON), withText) +import Data.Text (Text) +import GHC.Generics (Generic) + +-- | The supported upstream API surfaces, plus an open escape hatch. +data Api + = OpenAIChatCompletions + | AnthropicMessages + | OpenAICompletionsCli + | AnthropicMessagesCli + | Custom !Text + deriving stock (Eq, Ord, Show, Generic) + +-- | Render an 'Api' tag as its canonical kebab-cased wire string. +renderApi :: Api -> Text +renderApi = \case + OpenAIChatCompletions -> "openai-chat-completions" + AnthropicMessages -> "anthropic-messages" + OpenAICompletionsCli -> "openai-completions-cli" + AnthropicMessagesCli -> "anthropic-messages-cli" + Custom t -> t + +-- | Parse a wire string into an 'Api' tag. Unknown strings become +-- 'Custom' values so callers can use the same tag space. +parseApi :: Text -> Api +parseApi = \case + "openai-chat-completions" -> OpenAIChatCompletions + "anthropic-messages" -> AnthropicMessages + "openai-completions-cli" -> OpenAICompletionsCli + "anthropic-messages-cli" -> AnthropicMessagesCli + t -> Custom t + +instance ToJSON Api where + toJSON = toJSON . renderApi + +instance FromJSON Api where + parseJSON = withText "Api" (pure . parseApi) diff --git a/src/Baikai/Auth.hs b/src/Baikai/Auth.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Auth.hs @@ -0,0 +1,59 @@ +{-# LANGUAGE LambdaCase #-} + +-- | API key sourcing for provider constructors. +-- +-- Providers accept an 'ApiKeySource' rather than a raw 'Text' so test code can +-- supply a literal token and production code can defer to an environment variable. +-- The lookup happens lazily inside 'resolveApiKey'; constructing an 'ApiKeyEnv' +-- value does not read the environment. +module Baikai.Auth + ( ApiKeySource (..), + renderApiKeySourceForDebug, + resolveApiKey, + ) +where + +import Baikai.Error (BaikaiError (..)) +import Control.Exception (throwIO) +import Control.Monad.IO.Class (MonadIO, liftIO) +import Data.Aeson (ToJSON (toJSON), object, (.=)) +import Data.Text (Text) +import Data.Text qualified as Text +import System.Environment qualified as Environment + +data ApiKeySource + = ApiKeyLiteral !Text + | ApiKeyEnv !String + deriving stock (Eq) + +instance Show ApiKeySource where + show = Text.unpack . renderApiKeySourceForDebug + +instance ToJSON ApiKeySource where + toJSON (ApiKeyLiteral _) = + object + [ "source" .= ("literal-redacted" :: Text) + ] + toJSON (ApiKeyEnv name) = + object + [ "source" .= ("env" :: Text), + "name" .= name + ] + +-- | Render a credential source for logs, test failures, and debugging without +-- exposing literal secret material. +renderApiKeySourceForDebug :: ApiKeySource -> Text +renderApiKeySourceForDebug (ApiKeyLiteral _) = "ApiKeyLiteral " +renderApiKeySourceForDebug (ApiKeyEnv name) = + "ApiKeyEnv " <> Text.pack (show name) + +-- | Resolve a key source to a plain 'Text'. Throws 'ProviderError' (a +-- constructor of 'BaikaiError') if 'ApiKeyEnv' is used and the named variable +-- is unset. +resolveApiKey :: (MonadIO m) => ApiKeySource -> m Text +resolveApiKey (ApiKeyLiteral t) = pure t +resolveApiKey (ApiKeyEnv name) = + liftIO $ + Environment.lookupEnv name >>= \case + Just v -> pure (Text.pack v) + Nothing -> throwIO (ProviderError ("env var " <> Text.pack name <> " is not set")) diff --git a/src/Baikai/CacheRetention.hs b/src/Baikai/CacheRetention.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/CacheRetention.hs @@ -0,0 +1,27 @@ +-- | Provider-agnostic prompt-cache retention preference. +-- +-- Each provider maps the value to its own primitive: Anthropic's +-- 'long' becomes @cache_control.ttl: "1h"@, 'short' becomes the +-- ephemeral marker with no TTL; OpenAI Responses API's 'long' would +-- become 24h. Hosts that do not advertise prompt caching ignore the +-- preference. +module Baikai.CacheRetention + ( CacheRetention (..), + ) +where + +import Data.Aeson (FromJSON, ToJSON) +import GHC.Generics (Generic) + +-- | The three buckets pi-mono settled on. +data CacheRetention + = -- | Do not request prompt caching at all. + CacheRetentionNone + | -- | Provider-default ephemeral retention (Anthropic: 5 minutes). + CacheRetentionShort + | -- | Long-retention bucket (Anthropic: @ttl: "1h"@; OpenAI + -- Responses: 24h). Downgrades to short on hosts that report + -- 'supportsLongCacheRetention' as 'False'. + CacheRetentionLong + deriving stock (Eq, Show, Generic) + deriving anyclass (FromJSON, ToJSON) diff --git a/src/Baikai/Compat.hs b/src/Baikai/Compat.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Compat.hs @@ -0,0 +1,237 @@ +{-# LANGUAGE LambdaCase #-} + +-- | Per-API compatibility records — feature flags that capture how a +-- specific host implements an OpenAI- or Anthropic-style endpoint. +-- The shape mirrors pi-mono's @OpenAICompletionsCompat@ and +-- @AnthropicMessagesCompat@ records. +-- +-- The driving observation is that two hosts can share an API tag +-- (both DeepSeek and OpenAI implement the Chat Completions wire +-- protocol) yet differ in small details: where to put the max-output +-- token cap, whether @strict: true@ is accepted on tool definitions, +-- which JSON key carries reasoning effort, whether thinking is +-- delivered in-band as @\@ tags, etc. Carrying those flags +-- on a per-'Baikai.Model.Model' record means the same provider +-- handler can serve both hosts without any per-host code branching. +-- +-- These record constructors are intentionally public in baikai 0.1. +-- They are the supported escape hatch for hand-rolled models and +-- generated catalog overrides when a host speaks a familiar protocol +-- but rejects or requires a small request-shaping difference. New +-- fields or enum constructors may be added in later minor versions as +-- new provider quirks are discovered; callers that want fewer upgrade +-- edits should start from the @default*@ values and use record updates +-- rather than constructing every field from scratch. +-- +-- Auto-detection from a 'Baikai.Model.Model' @baseUrl@ provides +-- reasonable defaults so callers rarely need to spell out a full +-- compat record. +module Baikai.Compat + ( -- * OpenAI Chat Completions compat + OpenAICompletionsCompat (..), + defaultOpenAICompletionsCompat, + MaxTokensField (..), + ThinkingFormat (..), + CacheControlFormat (..), + + -- * Anthropic Messages compat + AnthropicMessagesCompat (..), + defaultAnthropicMessagesCompat, + + -- * Auto-detection from baseUrl + autoDetectOpenAICompletions, + autoDetectAnthropicMessages, + ) +where + +import Data.Aeson (FromJSON, ToJSON) +import Data.Text (Text) +import Data.Text qualified as Text +import GHC.Generics (Generic) + +-- | Where the OpenAI-compatible host expects the max-output-tokens +-- cap. OpenAI's o-series models require @max_completion_tokens@; many +-- compatible hosts (DeepSeek, some OpenRouter routes) keep the +-- pre-o-series @max_tokens@ name. +data MaxTokensField + = MaxCompletionTokensField + | MaxTokensField + deriving stock (Eq, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | How the OpenAI-compatible host receives reasoning-effort +-- preferences. Each constructor names a wire shape we know about; new +-- shapes land as new constructors. +data ThinkingFormat + = -- | OpenAI-native: top-level @reasoning_effort: "minimal" | "low" + -- | "medium" | "high"@. + ThinkingFormatOpenAI + | -- | OpenRouter: nested @reasoning: { effort: "..." }@. + ThinkingFormatOpenRouter + | -- | DeepSeek: top-level @thinking: { type: "enabled" }@ plus + -- @reasoning_effort@. + ThinkingFormatDeepseek + | -- | Together AI: nested @reasoning: { enabled: true }@ plus + -- @reasoning_effort@. + ThinkingFormatTogether + | -- | Z.ai / Qwen: top-level @enable_thinking: true@. + ThinkingFormatZai + | -- | Qwen chat-template: top-level @enable_thinking: true@. + ThinkingFormatQwen + | -- | Host does not expose reasoning controls; the option is + -- silently dropped. + ThinkingFormatNone + deriving stock (Eq, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | The wire shape an OpenAI-compatible host accepts for prompt-cache +-- markers. Only Anthropic's @cache_control@ markers are known; the +-- field is 'Nothing' on hosts that do not advertise prompt caching. +data CacheControlFormat = CacheControlFormatAnthropic + deriving stock (Eq, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | Feature flags for one OpenAI-compatible host. +-- +-- Every flag has a default that reproduces OpenAI's own behaviour — +-- @defaultOpenAICompletionsCompat@. Hosts that diverge from OpenAI +-- override the differing fields only. +data OpenAICompletionsCompat = OpenAICompletionsCompat + { -- | Where the max-output-tokens cap goes in the request body. + maxTokensField :: !MaxTokensField, + -- | Whether the host accepts the @developer@ message role + -- (OpenAI's o-series) or only @system@. + supportsDeveloperRole :: !Bool, + -- | Whether the host accepts @strict: true@ on function tool + -- definitions. Set 'False' to drop the field on hosts that + -- reject it. + supportsStrictMode :: !Bool, + -- | Whether the host smuggles thinking into the assistant text as + -- @\...\@ markers. When 'True', the + -- provider's stream transformer extracts those markers into + -- typed thinking deltas on the way out. (The transformer itself + -- is wired in by EP-3; EP-5 only flips the switch.) + requiresThinkingAsText :: !Bool, + -- | The wire shape the host accepts for reasoning-effort + -- preferences. + thinkingFormat :: !ThinkingFormat, + -- | Whether the host accepts Anthropic-style @cache_control@ + -- markers (some OpenRouter routes pass them through to an + -- Anthropic backend). + cacheControlFormat :: !(Maybe CacheControlFormat), + -- | Whether the host emits a final @usage@ chunk in streaming + -- responses. OpenAI does; older OpenAI-compatible servers may + -- not. Currently advisory. + supportsUsageInStreaming :: !Bool, + -- | Whether the host honours long (1h) cache TTLs through the + -- Anthropic-style cache_control marker. Currently only relevant + -- when 'cacheControlFormat' is 'Just CacheControlFormatAnthropic'. + supportsLongCacheRetention :: !Bool + } + deriving stock (Eq, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | OpenAI's own host: every flag at the OpenAI default. +defaultOpenAICompletionsCompat :: OpenAICompletionsCompat +defaultOpenAICompletionsCompat = + OpenAICompletionsCompat + { maxTokensField = MaxCompletionTokensField, + supportsDeveloperRole = True, + supportsStrictMode = True, + requiresThinkingAsText = False, + thinkingFormat = ThinkingFormatOpenAI, + cacheControlFormat = Nothing, + supportsUsageInStreaming = True, + supportsLongCacheRetention = True + } + +-- | Feature flags for one Anthropic Messages-compatible host. +data AnthropicMessagesCompat = AnthropicMessagesCompat + { -- | Whether the host honours Anthropic's + -- @cache_control.ttl: "1h"@ long-retention marker. When 'False', + -- long-retention preferences silently downgrade to ephemeral. + supportsLongCacheRetention :: !Bool, + -- | Whether tool definitions accept @cache_control@ markers. + -- Anthropic's own host does; some compatible hosts do not. + supportsCacheControlOnTools :: !Bool, + -- | Whether the host supports per-tool eager input streaming + -- (Anthropic's @fine-grained-tool-streaming@ beta). Currently + -- advisory; the SDK does not expose the field directly. + supportsEagerToolInputStreaming :: !Bool, + -- | Whether to add session-affinity headers on every request + -- (Fireworks-style routing). + sendSessionAffinityHeaders :: !Bool + } + deriving stock (Eq, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | Anthropic's own host: every flag at its default. +defaultAnthropicMessagesCompat :: AnthropicMessagesCompat +defaultAnthropicMessagesCompat = + AnthropicMessagesCompat + { supportsLongCacheRetention = True, + supportsCacheControlOnTools = True, + supportsEagerToolInputStreaming = True, + sendSessionAffinityHeaders = False + } + +-- | Pick a sensible compat record for an unknown OpenAI-compatible +-- host based on its @baseUrl@. Falls back to +-- 'defaultOpenAICompletionsCompat' for hosts the table does not +-- recognise. +autoDetectOpenAICompletions :: Text -> OpenAICompletionsCompat +autoDetectOpenAICompletions url + | hasInfix "api.openai.com" = defaultOpenAICompletionsCompat + | hasInfix "api.deepseek.com" = + defaultOpenAICompletionsCompat + { thinkingFormat = ThinkingFormatDeepseek, + requiresThinkingAsText = True, + maxTokensField = MaxTokensField, + supportsStrictMode = False, + supportsDeveloperRole = False + } + | hasInfix "openrouter.ai" = + defaultOpenAICompletionsCompat + { thinkingFormat = ThinkingFormatOpenRouter, + supportsStrictMode = False, + supportsDeveloperRole = False, + cacheControlFormat = Just CacheControlFormatAnthropic + } + | hasInfix "together.xyz" || hasInfix "api.together.ai" = + defaultOpenAICompletionsCompat + { thinkingFormat = ThinkingFormatTogether, + supportsStrictMode = False, + supportsDeveloperRole = False + } + | hasInfix "z.ai" = + defaultOpenAICompletionsCompat + { thinkingFormat = ThinkingFormatZai, + supportsStrictMode = False + } + | hasInfix "dashscope" || hasInfix "qwen" = + defaultOpenAICompletionsCompat + { thinkingFormat = ThinkingFormatQwen, + supportsStrictMode = False + } + | otherwise = defaultOpenAICompletionsCompat + where + hasInfix :: Text -> Bool + hasInfix needle = needle `Text.isInfixOf` url + +-- | Pick a sensible compat record for an unknown +-- Anthropic-compatible host based on its @baseUrl@. Falls back to +-- 'defaultAnthropicMessagesCompat' for hosts the table does not +-- recognise. +autoDetectAnthropicMessages :: Text -> AnthropicMessagesCompat +autoDetectAnthropicMessages url + | hasInfix "api.anthropic.com" = defaultAnthropicMessagesCompat + | hasInfix "fireworks.ai" = + defaultAnthropicMessagesCompat + { supportsCacheControlOnTools = False, + sendSessionAffinityHeaders = True, + supportsLongCacheRetention = False + } + | otherwise = defaultAnthropicMessagesCompat + where + hasInfix :: Text -> Bool + hasInfix needle = needle `Text.isInfixOf` url diff --git a/src/Baikai/Content.hs b/src/Baikai/Content.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Content.hs @@ -0,0 +1,212 @@ +{-# LANGUAGE LambdaCase #-} + +-- | Typed content blocks for user, assistant, and tool-result messages. +-- +-- A message no longer carries a single opaque blob of text. Instead it +-- holds a vector of typed blocks. For user input the blocks can be text +-- or an inline base64-encoded image. For assistant output a block can be +-- text, a thinking trace (for reasoning-capable models), or a tool +-- invocation. For tool-result messages (a caller-supplied reply to a +-- model-issued tool call) blocks can be text or image. +-- +-- EP-1 introduces the types; EP-3 streams them, and EP-4 wires tool +-- round-tripping through the providers. Image content is restricted to +-- inline base64 with an explicit @mimeType@: the caller is responsible +-- for the (small, reversible) work of base64-encoding bytes once, and +-- every provider can consume the same shape without a URL-fetch path +-- that varies in failure modes. +module Baikai.Content + ( -- * Block primitives + TextContent (..), + ThinkingContent (..), + ToolCall (..), + ImageContent (..), + + -- * Per-role block sums + UserContent (..), + AssistantContent (..), + ToolResultContent (..), + + -- * Smart defaults + _TextContent, + _ThinkingContent, + _ToolCall, + _ImageContent, + ) +where + +import Data.Aeson + ( FromJSON (parseJSON), + Options (..), + SumEncoding (..), + ToJSON (toJSON), + Value (Null), + camelTo2, + defaultOptions, + genericParseJSON, + genericToJSON, + object, + withObject, + (.:), + (.=), + ) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.ByteString.Base64 qualified as Base64 +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as Text +import GHC.Generics (Generic) + +-- | A plain-text block. The wire form is @{"text": "..."}@. +newtype TextContent = TextContent + { text :: Text + } + deriving stock (Eq, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | A reasoning trace block produced by thinking-capable models. The +-- @signature@ field is an opaque continuation token returned by the +-- provider; it must be threaded back into the next request unchanged for +-- the model to resume from the same state. @redacted@ is 'True' when the +-- provider hid the underlying content (Anthropic flags this when its +-- safety system removes a thinking block from the wire response). +data ThinkingContent = ThinkingContent + { thinking :: !Text, + signature :: !(Maybe Text), + redacted :: !Bool + } + deriving stock (Eq, Show, Generic) + +-- | A model-issued tool invocation. @id_@ has a trailing underscore in +-- Haskell to dodge a clash with @Prelude.id@; the JSON encoding strips +-- it back to @id@. +data ToolCall = ToolCall + { id_ :: !Text, + name :: !Text, + arguments :: !Value + } + deriving stock (Eq, Show, Generic) + +-- | An inline image block. Bytes are stored decoded; the JSON encoding +-- emits base64 under @data@ and the @mimeType@ camel-snakes to +-- @mime_type@. +data ImageContent = ImageContent + { imageData :: !ByteString, + mimeType :: !Text + } + deriving stock (Eq, Show, Generic) + +-- | What a caller can put into a 'Baikai.Message.UserMessage'. +data UserContent + = UserText !TextContent + | UserImage !ImageContent + deriving stock (Eq, Show, Generic) + +-- | What a provider can put into a 'Baikai.Message.AssistantMessage'. +data AssistantContent + = AssistantText !TextContent + | AssistantThinking !ThinkingContent + | AssistantToolCall !ToolCall + deriving stock (Eq, Show, Generic) + +-- | What a caller can put into a 'Baikai.Message.ToolResultMessage'. +data ToolResultContent + = ToolResultText !TextContent + | ToolResultImage !ImageContent + deriving stock (Eq, Show, Generic) + +_TextContent :: TextContent +_TextContent = TextContent {text = Text.empty} + +_ThinkingContent :: ThinkingContent +_ThinkingContent = + ThinkingContent + { thinking = Text.empty, + signature = Nothing, + redacted = False + } + +_ToolCall :: ToolCall +_ToolCall = + ToolCall + { id_ = Text.empty, + name = Text.empty, + arguments = Null + } + +_ImageContent :: ImageContent +_ImageContent = + ImageContent + { imageData = BS.empty, + mimeType = Text.empty + } + +-- Aeson plumbing --------------------------------------------------------- + +snakeOptions :: Options +snakeOptions = defaultOptions {fieldLabelModifier = camelTo2 '_'} + +instance FromJSON ThinkingContent where + parseJSON = genericParseJSON snakeOptions + +instance ToJSON ThinkingContent where + toJSON = genericToJSON snakeOptions + +-- Strip the trailing underscore on @id_@ so the wire form is @id@; the +-- other fields keep their natural names. +toolCallOptions :: Options +toolCallOptions = + defaultOptions + { fieldLabelModifier = \case + "id_" -> "id" + s -> camelTo2 '_' s + } + +instance FromJSON ToolCall where + parseJSON = genericParseJSON toolCallOptions + +instance ToJSON ToolCall where + toJSON = genericToJSON toolCallOptions + +-- ImageContent is encoded manually so the bytes round-trip through +-- base64 on the wire under @data@ while staying decoded in Haskell. +instance ToJSON ImageContent where + toJSON c = + object + [ "data" .= Text.decodeUtf8 (Base64.encode (imageData c)), + "mime_type" .= mimeType c + ] + +instance FromJSON ImageContent where + parseJSON = withObject "ImageContent" $ \o -> do + encoded <- o .: "data" + mt <- o .: "mime_type" + case Base64.decode (Text.encodeUtf8 encoded) of + Left err -> fail ("ImageContent.data is not base64: " <> err) + Right raw -> pure ImageContent {imageData = raw, mimeType = mt} + +contentSumOptions :: Options +contentSumOptions = + defaultOptions + { sumEncoding = TaggedObject {tagFieldName = "type", contentsFieldName = "data"}, + constructorTagModifier = camelTo2 '_' + } + +instance FromJSON UserContent where + parseJSON = genericParseJSON contentSumOptions + +instance ToJSON UserContent where + toJSON = genericToJSON contentSumOptions + +instance FromJSON AssistantContent where + parseJSON = genericParseJSON contentSumOptions + +instance ToJSON AssistantContent where + toJSON = genericToJSON contentSumOptions + +instance FromJSON ToolResultContent where + parseJSON = genericParseJSON contentSumOptions + +instance ToJSON ToolResultContent where + toJSON = genericToJSON contentSumOptions diff --git a/src/Baikai/Context.hs b/src/Baikai/Context.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Context.hs @@ -0,0 +1,94 @@ +-- | The 'Context' record — the part of a request that defines the +-- conversation: the optional system prompt, the message vector, and +-- the declared tools the model may invoke. +-- +-- 'Context' replaces the prior 'Baikai.Request.Request' record's +-- conversation-related fields. The per-call knobs that previously +-- lived alongside the messages (max tokens, temperature, API key) +-- now live on 'Baikai.Options.Options' instead. +-- +-- EP-4 adds the @tools@ field and the 'appendToolResult' helper +-- that builds the follow-up request after the model invoked one or +-- more tools. The helper lives here rather than in 'Baikai.Tool' so +-- that 'Baikai.Tool' can stay imports-light (the 'Tool' type is +-- referenced by the @tools@ field, so 'Baikai.Tool' cannot itself +-- depend on 'Context'). +module Baikai.Context + ( Context (..), + _Context, + appendToolResult, + appendToolResultText, + ) +where + +import Baikai.Content (AssistantContent (..), ToolCall (..)) +import Baikai.Message (Message (..), ToolResult, toolResultFromCallNow, toolResultText) +import Baikai.Response (Response (..), responseMessage) +import Baikai.Tool (Tool) +import Control.Lens ((&), (.~), (^.)) +import Data.Aeson (ToJSON) +import Data.Generics.Labels () +import Data.Text (Text) +import Data.Vector (Vector) +import Data.Vector qualified as V +import GHC.Generics (Generic) + +data Context = Context + { systemPrompt :: !(Maybe Text), + messages :: !(Vector Message), + tools :: !(Vector Tool) + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +_Context :: Context +_Context = + Context + { systemPrompt = Nothing, + messages = V.empty, + tools = V.empty + } + +-- | Execute every 'AssistantToolCall' in @resp@'s message via the +-- caller-supplied dispatcher, then append the assistant message and +-- one 'Baikai.Message.ToolResultMessage' per call to @ctx@. The +-- returned 'Context' is ready to drive the follow-up request that +-- gives the model the tool results. +-- +-- The dispatcher receives one 'ToolCall' at a time and returns a rich +-- 'ToolResult' carrying text blocks, image blocks, and an error flag. +-- Any error handling (timeouts, sandboxing, multi-call concurrency) +-- lives in the dispatcher. +appendToolResult :: + Context -> + Response -> + (ToolCall -> IO ToolResult) -> + IO Context +appendToolResult ctx resp dispatcher = do + let respPayload = resp ^. #message + respMsg = responseMessage resp + toolCalls = [tc | AssistantToolCall tc <- V.toList (respPayload ^. #content)] + results <- + traverse + ( \tc -> do + result <- dispatcher tc + toolResultFromCallNow tc result + ) + toolCalls + pure $ + ctx + & #messages + .~ ( (ctx ^. #messages) + <> V.singleton respMsg + <> V.fromList results + ) + +-- | Text-only convenience wrapper for the common case where every +-- tool call returns one successful text block. +appendToolResultText :: + Context -> + Response -> + (ToolCall -> IO Text) -> + IO Context +appendToolResultText ctx resp dispatcher = + appendToolResult ctx resp (fmap toolResultText . dispatcher) diff --git a/src/Baikai/Cost.hs b/src/Baikai/Cost.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Cost.hs @@ -0,0 +1,60 @@ +module Baikai.Cost + ( Cost (..), + CostBreakdown (..), + _Cost, + _CostBreakdown, + usdAsScientific, + ) +where + +import Data.Aeson (ToJSON (toJSON), object, (.=)) +import Data.Scientific (Scientific, fromRationalRepetendUnlimited) +import GHC.Generics (Generic) + +data CostBreakdown = CostBreakdown + { inputUsd :: !Rational, + outputUsd :: !Rational, + cachedInputUsd :: !Rational, + cachedWriteUsd :: !Rational + } + deriving stock (Eq, Show, Generic) + +data Cost = Cost + { usd :: !Rational, + breakdown :: !CostBreakdown + } + deriving stock (Eq, Show, Generic) + +_CostBreakdown :: CostBreakdown +_CostBreakdown = + CostBreakdown + { inputUsd = 0, + outputUsd = 0, + cachedInputUsd = 0, + cachedWriteUsd = 0 + } + +_Cost :: Cost +_Cost = Cost {usd = 0, breakdown = _CostBreakdown} + +instance ToJSON CostBreakdown where + toJSON cb = + object + [ "input_usd" .= ratToSci (inputUsd cb), + "output_usd" .= ratToSci (outputUsd cb), + "cached_input_usd" .= ratToSci (cachedInputUsd cb), + "cached_write_usd" .= ratToSci (cachedWriteUsd cb) + ] + +instance ToJSON Cost where + toJSON c = + object + [ "usd" .= ratToSci (usd c), + "breakdown" .= breakdown c + ] + +usdAsScientific :: Cost -> Scientific +usdAsScientific = ratToSci . usd + +ratToSci :: Rational -> Scientific +ratToSci = fst . fromRationalRepetendUnlimited diff --git a/src/Baikai/Cost/Log.hs b/src/Baikai/Cost/Log.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Cost/Log.hs @@ -0,0 +1,222 @@ +-- | Opt-in JSONL call log for AI API calls. +-- +-- Each open handle owns a 'Chan' and a worker thread. 'appendEntry' +-- is a cheap channel push that returns immediately; the worker +-- drains the channel to disk through a streamly fold. Disk latency +-- therefore never pads the apparent latency of 'completeRequest'. +-- +-- The usual pattern is 'withCallLog', which opens a handle, runs +-- the body, and flushes pending entries on the way out: +-- +-- > withCallLog (CallLogConfig "/tmp/baikai.jsonl" True) $ \h -> do +-- > _ <- runRequestWithLog h model context options +-- > pure () +module Baikai.Cost.Log + ( CallLogConfig (..), + CallLogEntry (..), + CallLogHandle, + openCallLog, + closeCallLog, + withCallLog, + appendEntry, + runRequestWithLog, + runRequestWithLogWith, + summarizeContext, + ) +where + +import Baikai.Content (TextContent (..), UserContent (..)) +import Baikai.Context (Context) +import Baikai.Cost (usdAsScientific) +import Baikai.Message + ( Message (..), + UserPayload (..), + ) +import Baikai.Model (Model) +import Baikai.Options (Options) +import Baikai.Provider.Registry + ( ProviderRegistry, + completeRequestWith, + globalProviderRegistry, + ) +import Baikai.Response (Response) +import Baikai.Usage (Usage) +import Baikai.Usage qualified as Usage +import Control.Concurrent (forkIO) +import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan) +import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar, takeMVar) +import Control.Exception (bracket) +import Control.Lens ((^.)) +import Control.Monad.IO.Class (MonadIO, liftIO) +import Control.Monad.IO.Unlift (MonadUnliftIO, withRunInIO) +import Data.Aeson (FromJSON, ToJSON) +import Data.Aeson qualified as Aeson +import Data.ByteString.Lazy qualified as BSL +import Data.Foldable (find) +import Data.Function ((&)) +import Data.Generics.Labels () +import Data.Scientific (Scientific) +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Time (UTCTime, getCurrentTime) +import Data.Vector qualified as Vector +import GHC.Generics (Generic) +import Numeric.Natural (Natural) +import Streamly.Data.Fold qualified as Fold +import Streamly.Data.Stream qualified as Stream +import System.IO (BufferMode (LineBuffering), IOMode (AppendMode), hSetBuffering, withFile) + +-- | Where (and whether) to write the call log. +data CallLogConfig = CallLogConfig + { path :: !FilePath, + enabled :: !Bool + } + deriving stock (Eq, Show, Generic) + +-- | One line of the JSONL call log. Wire shape preserved from EP-0: +-- @cachedInputTokens@ keeps its name so existing log readers keep +-- parsing. +data CallLogEntry = CallLogEntry + { timestamp :: !UTCTime, + provider :: !Text, + model :: !Text, + inputTokens :: !(Maybe Natural), + outputTokens :: !(Maybe Natural), + cachedInputTokens :: !(Maybe Natural), + reasoningTokens :: !(Maybe Natural), + usd :: !(Maybe Scientific), + latencyMs :: !Integer, + promptSummary :: !Text + } + deriving stock (Eq, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | Opaque handle for an open call log. +data CallLogHandle = CallLogHandle + { chan :: !(Chan (Maybe CallLogEntry)), + done :: !(MVar ()), + cfg :: !CallLogConfig + } + +-- | Open a handle. When @enabled = True@, fork the worker thread +-- that drains entries to @path@ in append mode. +openCallLog :: (MonadIO m) => CallLogConfig -> m CallLogHandle +openCallLog c = liftIO $ do + ch <- newChan + d <- newEmptyMVar + case enabled c of + False -> putMVar d () + True -> do + _ <- forkIO (worker (path c) ch d) + pure () + pure CallLogHandle {chan = ch, done = d, cfg = c} + +-- | Signal shutdown and block until the worker has drained every +-- pending entry to disk. +closeCallLog :: (MonadIO m) => CallLogHandle -> m () +closeCallLog h = liftIO $ do + case enabled (cfg h) of + True -> writeChan (chan h) Nothing + False -> pure () + takeMVar (done h) + +-- | Bracketed lifetime: open the handle, run the body, close +-- exactly once on every path (including exceptions). +withCallLog :: (MonadUnliftIO m) => CallLogConfig -> (CallLogHandle -> m a) -> m a +withCallLog c body = + withRunInIO $ \run -> + bracket (openCallLog c) closeCallLog (run . body) + +-- | Non-blocking enqueue. When the handle is disabled, returns +-- immediately without touching the channel. +appendEntry :: (MonadIO m) => CallLogHandle -> CallLogEntry -> m () +appendEntry h entry + | not (enabled (cfg h)) = pure () + | otherwise = liftIO (writeChan (chan h) (Just entry)) + +-- | Dispatch through the registry, then (if logging is enabled) +-- enqueue a single JSONL record summarizing the call. +runRequestWithLog :: + (MonadIO m) => + CallLogHandle -> + Model -> + Context -> + Options -> + m Response +runRequestWithLog = runRequestWithLogWith globalProviderRegistry + +-- | Dispatch through the selected registry, then (if logging is enabled) +-- enqueue a single JSONL record summarizing the call. +runRequestWithLogWith :: + (MonadIO m) => + ProviderRegistry -> + CallLogHandle -> + Model -> + Context -> + Options -> + m Response +runRequestWithLogWith reg h m ctx opts = do + resp <- liftIO (completeRequestWith reg m ctx opts) + now <- liftIO getCurrentTime + let u :: Usage + u = (resp ^. #message) ^. #usage + meaningfulCost = (Usage.cost u) ^. #usd > 0 + entry = + CallLogEntry + { timestamp = now, + provider = resp ^. #provider, + model = (resp ^. #model) ^. #modelId, + inputTokens = positive (Usage.inputTokens u), + outputTokens = positive (Usage.outputTokens u), + cachedInputTokens = positive (Usage.cacheReadTokens u), + reasoningTokens = Usage.reasoningTokens u, + usd = if meaningfulCost then Just (usdAsScientific (Usage.cost u)) else Nothing, + latencyMs = resp ^. #latencyMs, + promptSummary = summarizeContext ctx + } + appendEntry h entry + pure resp + +-- | 'Just n' when @n > 0@, otherwise 'Nothing'. Keeps zero counts +-- out of the JSONL when the call did not exercise that dimension. +positive :: Natural -> Maybe Natural +positive 0 = Nothing +positive n = Just n + +-- | Worker loop: pull 'Maybe CallLogEntry' off the channel, drain +-- through a streamly fold that writes each entry as one JSON line. +worker :: FilePath -> Chan (Maybe CallLogEntry) -> MVar () -> IO () +worker p ch d = + withFile p AppendMode $ \fh -> do + hSetBuffering fh LineBuffering + let step :: () -> IO (Maybe (CallLogEntry, ())) + step () = do + msg <- readChan ch + case msg of + Nothing -> pure Nothing + Just e -> pure (Just (e, ())) + Stream.unfoldrM step () + & Stream.fold (Fold.drainMapM (writeEntry fh)) + putMVar d () + where + writeEntry fh entry = + BSL.hPut fh (Aeson.encode entry <> "\n") + +-- | Concatenate the first user message's text blocks, truncated to +-- 200 characters. Returns empty when no user message is present. +summarizeContext :: Context -> Text +summarizeContext ctx = + case find isUser (Vector.toList (ctx ^. #messages)) of + Just (UserMessage UserPayload {content = cs}) -> Text.take 200 (concatUserText cs) + _ -> Text.empty + where + isUser UserMessage {} = True + isUser _ = False + +-- | Concatenate the text-block payloads of a 'UserContent' vector. +-- Image blocks contribute nothing. +concatUserText :: Vector.Vector UserContent -> Text +concatUserText = Text.concat . Vector.toList . Vector.mapMaybe pickText + where + pickText (UserText (TextContent t)) = Just t + pickText _ = Nothing diff --git a/src/Baikai/Cost/Pricing.hs b/src/Baikai/Cost/Pricing.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Cost/Pricing.hs @@ -0,0 +1,69 @@ +-- | Cost computation from a 'Baikai.Model.Model' and a 'Usage'. +-- +-- The previous map-based lookup (@Map Text PricingRate@) is gone. +-- Pricing rates live on 'Baikai.Model.Model.cost' directly, so the +-- computation collapses to a record-field access. Models without +-- published pricing carry a zero 'Baikai.Model.ModelCost' (the +-- default in '_Model'), producing a zero 'Cost'. +module Baikai.Cost.Pricing + ( computeCost, + attachCost, + ) +where + +import Baikai.Cost (Cost (..), CostBreakdown (..)) +import Baikai.Message (AssistantPayload (..)) +import Baikai.Model (Model, ModelCost (..)) +import Baikai.Prelude +import Baikai.Response (Response (..)) +import Baikai.Usage (Usage (..)) + +-- | Compute a 'Cost' from a model's per-million-token rates and a +-- 'Usage'. Returns a zero 'Cost' when the model carries zero rates, +-- which is the truthful signal for providers without published +-- pricing (CLI providers, custom hosts). +computeCost :: Model -> Usage -> Cost +computeCost m u = + let rates = m ^. #cost + inRate = inputCost rates + outRate = outputCost rates + crRate = cacheReadCost rates + cwRate = cacheWriteCost rates + inUsd = toRational (u ^. #inputTokens) * inRate / 1_000_000 + outUsd = toRational (u ^. #outputTokens) * outRate / 1_000_000 + cachedUsd = toRational (u ^. #cacheReadTokens) * crRate / 1_000_000 + cacheWriteUsd = toRational (u ^. #cacheWriteTokens) * cwRate / 1_000_000 + total = inUsd + outUsd + cachedUsd + cacheWriteUsd + in Cost + { usd = total, + breakdown = + CostBreakdown + { inputUsd = inUsd, + outputUsd = outUsd, + cachedInputUsd = cachedUsd, + cachedWriteUsd = cacheWriteUsd + } + } + +-- | Replace the assistant response payload's embedded 'Cost' with one +-- computed from the supplied model. +attachCost :: Model -> Response -> Response +attachCost m r = + let AssistantPayload + { usage = u, + content = c, + stopReason = sr, + errorMessage = em, + timestamp = ts + } = r ^. #message + computed = computeCost m u + u' = u & #cost .~ computed + msg' = + AssistantPayload + { content = c, + usage = u', + stopReason = sr, + errorMessage = em, + timestamp = ts + } + in r & #message .~ msg' diff --git a/src/Baikai/Error.hs b/src/Baikai/Error.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Error.hs @@ -0,0 +1,13 @@ +module Baikai.Error (BaikaiError (..)) where + +import Control.Exception (Exception) +import Data.Text (Text) +import GHC.Generics (Generic) + +data BaikaiError + = ProviderError !Text + | RequestInvalid !Text + | DecodeError !Text + | ProcessError !Int !Text + deriving stock (Eq, Show, Generic) + deriving anyclass (Exception) diff --git a/src/Baikai/Interactive.hs b/src/Baikai/Interactive.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Interactive.hs @@ -0,0 +1,119 @@ +-- | Provider-neutral types for launching local interactive agent +-- CLIs such as Claude Code and Codex. +-- +-- This module deliberately does not implement process spawning. +-- The core package owns the shared vocabulary, while vendor +-- packages own the command-line flags for their local CLI. +module Baikai.Interactive + ( InteractiveProvider (..), + InteractiveScope (..), + InteractiveLaunchRequest (..), + InteractiveSafety (..), + CodexSandboxMode (..), + CodexApprovalPolicy (..), + InteractiveLaunchResult (..), + _InteractiveLaunchRequest, + _InteractiveLaunchResult, + renderInteractiveProvider, + renderInteractiveScope, + renderCodexSandboxMode, + renderCodexApprovalPolicy, + ) +where + +import Baikai.Prelude +import System.Exit (ExitCode) + +-- | Local interactive provider families that Baikai knows how to +-- describe without depending on a vendor package. +data InteractiveProvider + = InteractiveClaude + | InteractiveCodex + deriving stock (Eq, Ord, Show, Generic) + +-- | Scope for provider-native assets and configuration. User scope +-- means the provider's home-directory location. Project scope means +-- a location under the working project. +data InteractiveScope + = InteractiveUserScope + | InteractiveProjectScope + deriving stock (Eq, Ord, Show, Generic) + +-- | Inputs common to local interactive agent launches. +data InteractiveLaunchRequest = InteractiveLaunchRequest + { systemPrompt :: !(Maybe Text), + userPrompt :: !Text, + model :: !(Maybe Text), + workingDir :: !(Maybe FilePath), + extraDirs :: ![FilePath], + safety :: !InteractiveSafety, + extraArgs :: ![Text] + } + deriving stock (Eq, Show, Generic) + +-- | Safety configuration expressed in the shared core vocabulary. +-- Vendor launchers translate the selected branch into their CLI's +-- concrete flags. +data InteractiveSafety + = DefaultSafety + | ClaudeAllowedTools [Text] + | CodexSandbox CodexSandboxMode CodexApprovalPolicy + deriving stock (Eq, Ord, Show, Generic) + +data CodexSandboxMode + = CodexReadOnly + | CodexWorkspaceWrite + | CodexDangerFullAccess + deriving stock (Eq, Ord, Show, Generic) + +data CodexApprovalPolicy + = CodexApprovalUntrusted + | CodexApprovalOnFailure + | CodexApprovalOnRequest + | CodexApprovalNever + deriving stock (Eq, Ord, Show, Generic) + +-- | Process-level outcome after the interactive CLI exits. +data InteractiveLaunchResult = InteractiveLaunchResult + { provider :: !InteractiveProvider, + exitCode :: !ExitCode + } + deriving stock (Eq, Show, Generic) + +_InteractiveLaunchRequest :: Text -> InteractiveLaunchRequest +_InteractiveLaunchRequest prompt = + InteractiveLaunchRequest + { systemPrompt = Nothing, + userPrompt = prompt, + model = Nothing, + workingDir = Nothing, + extraDirs = [], + safety = DefaultSafety, + extraArgs = [] + } + +_InteractiveLaunchResult :: InteractiveProvider -> ExitCode -> InteractiveLaunchResult +_InteractiveLaunchResult p code = + InteractiveLaunchResult + { provider = p, + exitCode = code + } + +renderInteractiveProvider :: InteractiveProvider -> Text +renderInteractiveProvider InteractiveClaude = "claude" +renderInteractiveProvider InteractiveCodex = "codex" + +renderInteractiveScope :: InteractiveScope -> Text +renderInteractiveScope InteractiveUserScope = "user" +renderInteractiveScope InteractiveProjectScope = "project" + +renderCodexSandboxMode :: CodexSandboxMode -> Text +renderCodexSandboxMode CodexReadOnly = "read-only" +renderCodexSandboxMode CodexWorkspaceWrite = "workspace-write" +renderCodexSandboxMode CodexDangerFullAccess = "danger-full-access" + +renderCodexApprovalPolicy :: CodexApprovalPolicy -> Text +renderCodexApprovalPolicy CodexApprovalUntrusted = "untrusted" +renderCodexApprovalPolicy CodexApprovalOnFailure = "on-failure" +renderCodexApprovalPolicy CodexApprovalOnRequest = "on-request" +renderCodexApprovalPolicy CodexApprovalNever = "never" diff --git a/src/Baikai/Message.hs b/src/Baikai/Message.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Message.hs @@ -0,0 +1,305 @@ +-- | The conversation message ADT. +-- +-- A 'Message' is one of three constructors: +-- +-- * 'UserMessage' — caller input, carrying a 'UserPayload' (its +-- 'content' is a vector of text and inline image blocks, plus a +-- creation timestamp). +-- * 'AssistantMessage' — model output, carrying an 'AssistantPayload' +-- (its 'content' holds 'AssistantText', 'AssistantThinking', and +-- 'AssistantToolCall' blocks, alongside the call's 'Usage' (which now +-- embeds 'Cost' in-place), the 'StopReason' the model reported, an +-- optional 'errorMessage', and a timestamp). +-- * 'ToolResultMessage' — caller-supplied output for a model-issued +-- tool call, carrying a 'ToolResultPayload' (the tool-call id it +-- answers, the tool's name, the result 'content' (text or image), an +-- 'isError' flag, and a timestamp). +-- +-- The 'system' constructor from prior versions is removed: system +-- prompts live on 'Baikai.Request.Request.systemPrompt'. The 'Role' +-- enum is also removed — pattern-match on the constructor instead. +-- +-- Each constructor wraps a dedicated single-constructor payload record +-- ('UserPayload', 'AssistantPayload', 'ToolResultPayload') rather than +-- carrying record fields directly on the sum. This keeps every field +-- selector total: a partial selector defined across a multi-constructor +-- sum (e.g. an @assistantContent@ that throws on a 'UserMessage') is a +-- latent runtime hazard that @-Wpartial-fields@ flags. Splitting the +-- payloads out makes each field live in a single-constructor record +-- where selection cannot fail, and lets every payload reuse the clean +-- shared field name 'content' (via @DuplicateRecordFields@) instead of +-- the @userContent@/@assistantContent@/@toolResultContent@ workaround a +-- single data declaration would have forced. +module Baikai.Message + ( Message (..), + UserPayload (..), + AssistantPayload (..), + ToolResultPayload (..), + ToolResult (..), + user, + userAt, + userNow, + userImage, + userImageAt, + userImageNow, + assistant, + assistantAt, + assistantNow, + toolResultMessage, + toolResultMessageAt, + toolResultMessageNow, + toolResultFromCall, + toolResultFromCallAt, + toolResultFromCallNow, + toolResult, + toolResultAt, + toolResultNow, + toolResultText, + toolResultErrorText, + toolResultImage, + toolResultBlocks, + ) +where + +import Baikai.Content + ( AssistantContent (..), + ImageContent, + TextContent (TextContent), + ToolCall (..), + ToolResultContent (..), + UserContent (..), + ) +import Baikai.StopReason (StopReason (..)) +import Baikai.Usage (Usage, _Usage) +import Data.Aeson (ToJSON) +import Data.Text (Text) +import Data.Time (UTCTime, getCurrentTime) +import Data.Vector (Vector) +import Data.Vector qualified as V +import GHC.Generics (Generic) + +data Message + = UserMessage UserPayload + | AssistantMessage AssistantPayload + | ToolResultMessage ToolResultPayload + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +-- | Caller input: a vector of text and inline image blocks plus a +-- creation timestamp. +data UserPayload = UserPayload + { content :: !(Vector UserContent), + timestamp :: !UTCTime + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +-- | Model output: the assistant content blocks ('AssistantText', +-- 'AssistantThinking', 'AssistantToolCall'), the call's 'Usage' (which +-- embeds 'Cost' in-place), the reported 'StopReason', an optional +-- 'errorMessage', and a creation timestamp. +data AssistantPayload = AssistantPayload + { content :: !(Vector AssistantContent), + usage :: !Usage, + stopReason :: !StopReason, + errorMessage :: !(Maybe Text), + timestamp :: !UTCTime + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +-- | Caller-supplied output for a model-issued tool call: the tool-call +-- id it answers, the tool's name, the result 'content' (text or +-- image), an 'isError' flag, and a creation timestamp. +data ToolResultPayload = ToolResultPayload + { toolCallId :: !Text, + toolName :: !Text, + content :: !(Vector ToolResultContent), + isError :: !Bool, + timestamp :: !UTCTime + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +-- | A caller-produced result for one model-issued tool call. +-- +-- The result can contain one or more text or image blocks, and +-- @isError@ tells the model whether the tool execution failed in a +-- recoverable way. Use the smart constructors below for the common +-- cases. +data ToolResult = ToolResult + { content :: !(Vector ToolResultContent), + isError :: !Bool + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +defaultTimestamp :: UTCTime +defaultTimestamp = read "2000-01-01 00:00:00 UTC" + +-- | Build a one-text-block user message with a deterministic fixture +-- timestamp. Prefer 'userAt' when the timestamp matters, or 'userNow' +-- when the timestamp should be sampled in 'IO'. +user :: Text -> Message +user = userAt defaultTimestamp + +-- | Build a one-text-block user message at an explicit timestamp. +userAt :: UTCTime -> Text -> Message +userAt ts t = + UserMessage + UserPayload + { content = V.singleton (UserText (TextContent t)), + timestamp = ts + } + +-- | Build a one-text-block user message using the current time. +userNow :: Text -> IO Message +userNow t = do + now <- getCurrentTime + pure (userAt now t) + +-- | Build a user message carrying one inline image alongside optional +-- preceding text, using a deterministic fixture timestamp. +userImage :: ImageContent -> Maybe Text -> Message +userImage = userImageAt defaultTimestamp + +-- | Build a user image message at an explicit timestamp. Bytes are +-- stored decoded; the JSON encoding base64-encodes them on the wire. +userImageAt :: UTCTime -> ImageContent -> Maybe Text -> Message +userImageAt ts img mPrefix = + let prefix = case mPrefix of + Nothing -> V.empty + Just t -> V.singleton (UserText (TextContent t)) + in UserMessage + UserPayload + { content = prefix <> V.singleton (UserImage img), + timestamp = ts + } + +-- | Build a user image message using the current time. +userImageNow :: ImageContent -> Maybe Text -> IO Message +userImageNow img mPrefix = do + now <- getCurrentTime + pure (userImageAt now img mPrefix) + +-- | Build a one-text-block assistant message with zero usage and a +-- @stop@ stop reason, using a deterministic fixture timestamp. +assistant :: Text -> Message +assistant = assistantAt defaultTimestamp + +-- | Build a one-text-block assistant message at an explicit timestamp. +assistantAt :: UTCTime -> Text -> Message +assistantAt ts t = + AssistantMessage + AssistantPayload + { content = V.singleton (AssistantText (TextContent t)), + usage = _Usage, + stopReason = Stop, + errorMessage = Nothing, + timestamp = ts + } + +-- | Build a one-text-block assistant message using the current time. +assistantNow :: Text -> IO Message +assistantNow t = do + now <- getCurrentTime + pure (assistantAt now t) + +-- | Build a tool-result message answering a prior 'AssistantToolCall' +-- using rich content blocks and a deterministic fixture timestamp. +toolResultMessage :: Text -> Text -> ToolResult -> Message +toolResultMessage callId name = + toolResultMessageAt defaultTimestamp callId name + +-- | Build a tool-result message at an explicit timestamp. +toolResultMessageAt :: UTCTime -> Text -> Text -> ToolResult -> Message +toolResultMessageAt ts callId name ToolResult {content = blocks, isError = err} = + ToolResultMessage + ToolResultPayload + { toolCallId = callId, + toolName = name, + content = blocks, + isError = err, + timestamp = ts + } + +-- | Build a tool-result message using the current time. +toolResultMessageNow :: Text -> Text -> ToolResult -> IO Message +toolResultMessageNow callId name result = do + now <- getCurrentTime + pure (toolResultMessageAt now callId name result) + +-- | Build a rich tool-result message from the original tool call. +toolResultFromCall :: ToolCall -> ToolResult -> Message +toolResultFromCall ToolCall {id_ = callId, name = toolName} = + toolResultMessage callId toolName + +-- | Build a rich tool-result message from the original tool call at +-- an explicit timestamp. +toolResultFromCallAt :: UTCTime -> ToolCall -> ToolResult -> Message +toolResultFromCallAt ts ToolCall {id_ = callId, name = toolName} = + toolResultMessageAt ts callId toolName + +-- | Build a rich tool-result message from the original tool call using +-- the current time. +toolResultFromCallNow :: ToolCall -> ToolResult -> IO Message +toolResultFromCallNow ToolCall {id_ = callId, name = toolName} = + toolResultMessageNow callId toolName + +-- | Build a one-text-block tool-result message. This is the +-- text-only compatibility shorthand; prefer 'toolResultMessage' +-- when the result has multiple blocks or image content. +toolResult :: Text -> Text -> Text -> Bool -> Message +toolResult callId name body err = + toolResultAt defaultTimestamp callId name body err + +-- | Build a one-text-block tool-result message at an explicit timestamp. +toolResultAt :: UTCTime -> Text -> Text -> Text -> Bool -> Message +toolResultAt ts callId name body err = + toolResultMessageAt + ts + callId + name + ( ToolResult + { content = V.singleton (ToolResultText (TextContent body)), + isError = err + } + ) + +-- | Build a one-text-block tool-result message using the current time. +toolResultNow :: Text -> Text -> Text -> Bool -> IO Message +toolResultNow callId name body err = do + now <- getCurrentTime + pure (toolResultAt now callId name body err) + +-- | Successful one-text-block tool result. +toolResultText :: Text -> ToolResult +toolResultText body = + ToolResult + { content = V.singleton (ToolResultText (TextContent body)), + isError = False + } + +-- | Error one-text-block tool result. +toolResultErrorText :: Text -> ToolResult +toolResultErrorText body = + ToolResult + { content = V.singleton (ToolResultText (TextContent body)), + isError = True + } + +-- | Successful one-image-block tool result. +toolResultImage :: ImageContent -> ToolResult +toolResultImage img = + ToolResult + { content = V.singleton (ToolResultImage img), + isError = False + } + +-- | Tool result with explicit content blocks and error flag. +toolResultBlocks :: Vector ToolResultContent -> Bool -> ToolResult +toolResultBlocks blocks err = + ToolResult + { content = blocks, + isError = err + } diff --git a/src/Baikai/Model.hs b/src/Baikai/Model.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Model.hs @@ -0,0 +1,146 @@ +-- | A 'Model' is the data record callers hand to baikai when +-- dispatching a call. It carries everything baikai needs to know up +-- front to talk to a provider: the 'Api' tag (used to look up the +-- registered handler), the provider name, the base URL, the +-- per-million-token pricing rates, the context window and max output +-- cap, default per-call headers, and a per-API 'Compat' record (a +-- placeholder until EP-5 populates the real shims). +-- +-- The previous newtype @Model = Model Text@ — a thin tag for the +-- model id — is retired. The 'unModel' accessor is preserved as a +-- convenience until callers fully migrate to 'modelId'. +module Baikai.Model + ( -- * Model + Model (..), + _Model, + unModel, + + -- * Cost rates + ModelCost (..), + _ModelCost, + + -- * Capabilities + InputModality (..), + + -- * Compatibility shim + Compat (..), + openaiCompletionsCompatFor, + anthropicMessagesCompatFor, + ) +where + +import Baikai.Api (Api (..)) +import Baikai.Compat + ( AnthropicMessagesCompat, + OpenAICompletionsCompat, + autoDetectAnthropicMessages, + autoDetectOpenAICompletions, + ) +import Data.Aeson (FromJSON, ToJSON) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Text (Text) +import GHC.Generics (Generic) +import Numeric.Natural (Natural) + +-- | What kinds of input a model accepts. EP-1 introduced the typed +-- content blocks; this field documents which of them the chosen +-- 'Model' is allowed to consume. +data InputModality + = InputText + | InputImage + deriving stock (Eq, Ord, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | Per-million-token prices in USD as exact 'Rational' values. +data ModelCost = ModelCost + { inputCost :: !Rational, + outputCost :: !Rational, + cacheReadCost :: !Rational, + cacheWriteCost :: !Rational + } + deriving stock (Eq, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | Per-API compatibility shim. 'CompatNone' tells the provider to +-- pick a sensible record by inspecting 'baseUrl'; the two real +-- constructors carry an explicit per-host record that overrides the +-- auto-detection. The records themselves live in 'Baikai.Compat' so +-- this module does not pull in every compat field as a transitive +-- dependency. +data Compat + = CompatNone + | CompatOpenAICompletions !OpenAICompletionsCompat + | CompatAnthropicMessages !AnthropicMessagesCompat + deriving stock (Eq, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | Project the 'OpenAICompletionsCompat' that applies to a 'Model': +-- the explicit one if 'compat' is 'CompatOpenAICompletions', otherwise +-- the result of inspecting 'baseUrl' via 'autoDetectOpenAICompletions'. +openaiCompletionsCompatFor :: Model -> OpenAICompletionsCompat +openaiCompletionsCompatFor m = case compat m of + CompatOpenAICompletions c -> c + _ -> autoDetectOpenAICompletions (baseUrl m) + +-- | Project the 'AnthropicMessagesCompat' that applies to a 'Model': +-- the explicit one if 'compat' is 'CompatAnthropicMessages', +-- otherwise the result of inspecting 'baseUrl' via +-- 'autoDetectAnthropicMessages'. +anthropicMessagesCompatFor :: Model -> AnthropicMessagesCompat +anthropicMessagesCompatFor m = case compat m of + CompatAnthropicMessages c -> c + _ -> autoDetectAnthropicMessages (baseUrl m) + +-- | The data record baikai dispatches on. +data Model = Model + { modelId :: !Text, + name :: !Text, + api :: !Api, + provider :: !Text, + baseUrl :: !Text, + reasoning :: !Bool, + input :: ![InputModality], + cost :: !ModelCost, + contextWindow :: !Natural, + maxOutputTokens :: !Natural, + headers :: !(Map Text Text), + compat :: !Compat + } + deriving stock (Eq, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | The model id as 'Text'. Kept for ergonomic access from call +-- sites that previously held a newtype-wrapped 'Text'. +unModel :: Model -> Text +unModel = modelId + +-- | A zero 'ModelCost' across all rates. Useful as a default for +-- models without published pricing (CLI providers, custom hosts). +_ModelCost :: ModelCost +_ModelCost = + ModelCost + { inputCost = 0, + outputCost = 0, + cacheReadCost = 0, + cacheWriteCost = 0 + } + +-- | A blank 'Model'. Useful as a record-update base for hand-rolled +-- 'Model' values in tests and one-shot scripts. +_Model :: Model +_Model = + Model + { modelId = "", + name = "", + api = Custom "", + provider = "", + baseUrl = "", + reasoning = False, + input = [InputText], + cost = _ModelCost, + contextWindow = 0, + maxOutputTokens = 0, + headers = Map.empty, + compat = CompatNone + } diff --git a/src/Baikai/Models/Generated.hs b/src/Baikai/Models/Generated.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Models/Generated.hs @@ -0,0 +1,289 @@ +-- AUTO-GENERATED by baikai-gen-models. Do not edit by hand. +-- Regenerate with: cabal run baikai-gen-models +{-# LANGUAGE OverloadedStrings #-} +{-# OPTIONS_GHC -Wno-missing-export-lists #-} +{-# OPTIONS_GHC -Wno-unused-imports #-} + +module Baikai.Models.Generated where + +import Baikai.Api (Api (..)) +import Baikai.Compat + ( AnthropicMessagesCompat (..) + , CacheControlFormat (..) + , MaxTokensField (..) + , OpenAICompletionsCompat (..) + , ThinkingFormat (..) + ) +import Baikai.Model + ( Compat (..) + , InputModality (..) + , Model (..) + , ModelCost (..) + ) +import Data.Map.Strict qualified as Map +import Data.Ratio ((%)) + +anthropic_claude_haiku_4_5 :: Model +anthropic_claude_haiku_4_5 = + Model + { modelId = "claude-haiku-4-5" + , name = "Claude Haiku 4.5" + , api = AnthropicMessages + , provider = "anthropic" + , baseUrl = "https://api.anthropic.com" + , reasoning = False + , input = [InputText, InputImage] + , cost = ModelCost + { inputCost = 4 % 5 + , outputCost = 4 % 1 + , cacheReadCost = 2 % 25 + , cacheWriteCost = 1 % 1 + } + , contextWindow = 200000 + , maxOutputTokens = 8192 + , headers = Map.empty + , compat = CompatNone + } + +anthropic_claude_haiku_4_5_20251001 :: Model +anthropic_claude_haiku_4_5_20251001 = + Model + { modelId = "claude-haiku-4-5-20251001" + , name = "Claude Haiku 4.5 (2025-10-01)" + , api = AnthropicMessages + , provider = "anthropic" + , baseUrl = "https://api.anthropic.com" + , reasoning = False + , input = [InputText, InputImage] + , cost = ModelCost + { inputCost = 4 % 5 + , outputCost = 4 % 1 + , cacheReadCost = 2 % 25 + , cacheWriteCost = 1 % 1 + } + , contextWindow = 200000 + , maxOutputTokens = 8192 + , headers = Map.empty + , compat = CompatNone + } + +anthropic_claude_opus_4_7 :: Model +anthropic_claude_opus_4_7 = + Model + { modelId = "claude-opus-4-7" + , name = "Claude Opus 4.7" + , api = AnthropicMessages + , provider = "anthropic" + , baseUrl = "https://api.anthropic.com" + , reasoning = True + , input = [InputText, InputImage] + , cost = ModelCost + { inputCost = 15 % 1 + , outputCost = 75 % 1 + , cacheReadCost = 3 % 2 + , cacheWriteCost = 75 % 4 + } + , contextWindow = 200000 + , maxOutputTokens = 8192 + , headers = Map.empty + , compat = CompatNone + } + +anthropic_claude_sonnet_4_6 :: Model +anthropic_claude_sonnet_4_6 = + Model + { modelId = "claude-sonnet-4-6" + , name = "Claude Sonnet 4.6" + , api = AnthropicMessages + , provider = "anthropic" + , baseUrl = "https://api.anthropic.com" + , reasoning = False + , input = [InputText, InputImage] + , cost = ModelCost + { inputCost = 3 % 1 + , outputCost = 15 % 1 + , cacheReadCost = 3 % 10 + , cacheWriteCost = 15 % 4 + } + , contextWindow = 200000 + , maxOutputTokens = 8192 + , headers = Map.empty + , compat = CompatNone + } + +deepseek_deepseek_chat :: Model +deepseek_deepseek_chat = + Model + { modelId = "deepseek-chat" + , name = "DeepSeek Chat" + , api = OpenAIChatCompletions + , provider = "deepseek" + , baseUrl = "https://api.deepseek.com" + , reasoning = False + , input = [InputText] + , cost = ModelCost + { inputCost = 27 % 100 + , outputCost = 11 % 10 + , cacheReadCost = 7 % 100 + , cacheWriteCost = 0 % 1 + } + , contextWindow = 64000 + , maxOutputTokens = 8192 + , headers = Map.empty + , compat = CompatNone + } + +deepseek_deepseek_reasoner :: Model +deepseek_deepseek_reasoner = + Model + { modelId = "deepseek-reasoner" + , name = "DeepSeek Reasoner" + , api = OpenAIChatCompletions + , provider = "deepseek" + , baseUrl = "https://api.deepseek.com" + , reasoning = True + , input = [InputText] + , cost = ModelCost + { inputCost = 11 % 20 + , outputCost = 219 % 100 + , cacheReadCost = 7 % 50 + , cacheWriteCost = 0 % 1 + } + , contextWindow = 64000 + , maxOutputTokens = 8192 + , headers = Map.empty + , compat = CompatNone + } + +openai_gpt_4o :: Model +openai_gpt_4o = + Model + { modelId = "gpt-4o" + , name = "GPT-4o" + , api = OpenAIChatCompletions + , provider = "openai" + , baseUrl = "https://api.openai.com" + , reasoning = False + , input = [InputText, InputImage] + , cost = ModelCost + { inputCost = 5 % 2 + , outputCost = 10 % 1 + , cacheReadCost = 5 % 4 + , cacheWriteCost = 0 % 1 + } + , contextWindow = 128000 + , maxOutputTokens = 16384 + , headers = Map.empty + , compat = CompatNone + } + +openai_gpt_4o_mini :: Model +openai_gpt_4o_mini = + Model + { modelId = "gpt-4o-mini" + , name = "GPT-4o mini" + , api = OpenAIChatCompletions + , provider = "openai" + , baseUrl = "https://api.openai.com" + , reasoning = False + , input = [InputText, InputImage] + , cost = ModelCost + { inputCost = 3 % 20 + , outputCost = 3 % 5 + , cacheReadCost = 3 % 40 + , cacheWriteCost = 0 % 1 + } + , contextWindow = 128000 + , maxOutputTokens = 16384 + , headers = Map.empty + , compat = CompatNone + } + +openai_o1 :: Model +openai_o1 = + Model + { modelId = "o1" + , name = "OpenAI o1" + , api = OpenAIChatCompletions + , provider = "openai" + , baseUrl = "https://api.openai.com" + , reasoning = True + , input = [InputText, InputImage] + , cost = ModelCost + { inputCost = 15 % 1 + , outputCost = 60 % 1 + , cacheReadCost = 15 % 2 + , cacheWriteCost = 0 % 1 + } + , contextWindow = 200000 + , maxOutputTokens = 100000 + , headers = Map.empty + , compat = CompatNone + } + +openai_o1_mini :: Model +openai_o1_mini = + Model + { modelId = "o1-mini" + , name = "OpenAI o1-mini" + , api = OpenAIChatCompletions + , provider = "openai" + , baseUrl = "https://api.openai.com" + , reasoning = True + , input = [InputText] + , cost = ModelCost + { inputCost = 3 % 1 + , outputCost = 12 % 1 + , cacheReadCost = 3 % 2 + , cacheWriteCost = 0 % 1 + } + , contextWindow = 128000 + , maxOutputTokens = 65536 + , headers = Map.empty + , compat = CompatNone + } + +openrouter_anthropic_claude_sonnet_4 :: Model +openrouter_anthropic_claude_sonnet_4 = + Model + { modelId = "anthropic/claude-sonnet-4" + , name = "Claude Sonnet 4 via OpenRouter" + , api = OpenAIChatCompletions + , provider = "openrouter" + , baseUrl = "https://openrouter.ai/api" + , reasoning = False + , input = [InputText, InputImage] + , cost = ModelCost + { inputCost = 3 % 1 + , outputCost = 15 % 1 + , cacheReadCost = 3 % 10 + , cacheWriteCost = 15 % 4 + } + , contextWindow = 200000 + , maxOutputTokens = 8192 + , headers = Map.empty + , compat = CompatNone + } + +openrouter_openai_gpt_4o_mini :: Model +openrouter_openai_gpt_4o_mini = + Model + { modelId = "openai/gpt-4o-mini" + , name = "GPT-4o mini via OpenRouter" + , api = OpenAIChatCompletions + , provider = "openrouter" + , baseUrl = "https://openrouter.ai/api" + , reasoning = False + , input = [InputText, InputImage] + , cost = ModelCost + { inputCost = 3 % 20 + , outputCost = 3 % 5 + , cacheReadCost = 3 % 40 + , cacheWriteCost = 0 % 1 + } + , contextWindow = 128000 + , maxOutputTokens = 16384 + , headers = Map.empty + , compat = CompatNone + } + diff --git a/src/Baikai/Options.hs b/src/Baikai/Options.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Options.hs @@ -0,0 +1,57 @@ +-- | The 'Options' record — the per-call knobs that vary between +-- requests on the same conversation. +-- +-- 'apiKey' is 'Nothing' by default; the registered handler then +-- consults a provider-specific env var +-- (@ANTHROPIC_API_KEY@ for Anthropic, @OPENAI_API_KEY@ for OpenAI). +-- 'maxTokens' defaults to 'Nothing'; the handler falls back to the +-- chosen model's 'Baikai.Model.maxOutputTokens'. +-- +-- EP-4 added @toolChoice@. EP-5 adds @cacheRetention@ and @thinking@ +-- (provider-agnostic preferences that each provider maps to its own +-- primitive — see 'Baikai.CacheRetention' and 'Baikai.ThinkingLevel' +-- for the mappings). +module Baikai.Options + ( Options (..), + _Options, + ) +where + +import Baikai.Auth (ApiKeySource) +import Baikai.CacheRetention (CacheRetention) +import Baikai.ThinkingLevel (ThinkingLevel) +import Baikai.Tool (ToolChoice) +import Data.Aeson (ToJSON, Value) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Data.Text (Text) +import GHC.Generics (Generic) +import Numeric.Natural (Natural) + +data Options = Options + { maxTokens :: !(Maybe Natural), + temperature :: !(Maybe Double), + apiKey :: !(Maybe ApiKeySource), + timeoutMs :: !(Maybe Int), + headers :: !(Map Text Text), + metadata :: !(Map Text Value), + toolChoice :: !(Maybe ToolChoice), + cacheRetention :: !(Maybe CacheRetention), + thinking :: !(Maybe ThinkingLevel) + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +_Options :: Options +_Options = + Options + { maxTokens = Nothing, + temperature = Nothing, + apiKey = Nothing, + timeoutMs = Nothing, + headers = Map.empty, + metadata = Map.empty, + toolChoice = Nothing, + cacheRetention = Nothing, + thinking = Nothing + } diff --git a/src/Baikai/Prelude.hs b/src/Baikai/Prelude.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Prelude.hs @@ -0,0 +1,46 @@ +-- | Project-wide prelude for baikai. Re-exports the lens and +-- generic-lens vocabulary every module in this project is expected to +-- use, plus IO lifting, the scalar types used across the library, the +-- `Generic` class, and the aeson surface for deriving and writing +-- JSON instances. +-- +-- `Data.Generics.Labels` is imported for its orphan `IsLabel` instance +-- only (no values are re-exported); the instance still propagates to +-- anything that imports this Prelude. +module Baikai.Prelude + ( -- * Lens vocabulary + module Control.Lens, + + -- * Generic-lens vocabulary + module Data.Generics.Product, + module Data.Generics.Sum, + + -- * IO lifting (`liftIO` re-exported as a method of `MonadIO`) + MonadIO (..), + + -- * Scalar types + Text, + Vector, + Natural, + + -- * Generics + Generic, + + -- * JSON + FromJSON (..), + ToJSON (..), + genericParseJSON, + genericToJSON, + ) +where + +import Control.Lens hiding (Context) +import Control.Monad.IO.Class (MonadIO (..)) +import Data.Aeson (FromJSON (..), ToJSON (..), genericParseJSON, genericToJSON) +import Data.Generics.Labels () +import Data.Generics.Product +import Data.Generics.Sum +import Data.Text (Text) +import Data.Vector (Vector) +import GHC.Generics (Generic) +import Numeric.Natural (Natural) diff --git a/src/Baikai/Provider.hs b/src/Baikai/Provider.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Provider.hs @@ -0,0 +1,37 @@ +-- | The provider surface. +-- +-- The prior 'Provider' typeclass and 'SomeProvider' existential are +-- removed. Dispatch now goes through 'Baikai.Provider.Registry': +-- the caller picks a 'Baikai.Model.Model' record and calls +-- 'completeRequest'; the registry looks up the right handler by the +-- model's 'Baikai.Api.Api' tag. +-- +-- This module re-exports the registry surface so the +-- @import Baikai.Provider@ habit still resolves the symbols a +-- caller cares about. +module Baikai.Provider + ( ApiProvider (..), + ProviderRegistry, + newProviderRegistry, + globalProviderRegistry, + registerApiProviderWith, + registerApiProvider, + lookupApiProviderWith, + lookupApiProvider, + completeRequestWith, + completeRequest, + ) +where + +import Baikai.Provider.Registry + ( ApiProvider (..), + ProviderRegistry, + completeRequest, + completeRequestWith, + globalProviderRegistry, + lookupApiProvider, + lookupApiProviderWith, + newProviderRegistry, + registerApiProvider, + registerApiProviderWith, + ) diff --git a/src/Baikai/Provider/Cli/Internal.hs b/src/Baikai/Provider/Cli/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Provider/Cli/Internal.hs @@ -0,0 +1,164 @@ +-- | Internal helpers shared by the CLI providers in @baikai-claude@ +-- and @baikai-openai@. +-- +-- Not re-exported from "Baikai"; vendor packages import this module +-- directly. Anything exported here is considered part of the core +-- library's interface to the CLI provider packages and should not be +-- relied on by application code. +module Baikai.Provider.Cli.Internal + ( renderPrompt, + maybeApply, + decodeUtf8Lenient, + extractAgentMessage, + parseCodexJsonlStream, + ) +where + +import Baikai.Content + ( AssistantContent (..), + TextContent (..), + ToolResultContent (..), + UserContent (..), + ) +import Baikai.Context (Context) +import Baikai.Message + ( AssistantPayload (..), + Message (..), + ToolResultPayload (..), + UserPayload (..), + ) +import Control.Lens ((^.)) +import Data.Aeson (Value) +import Data.Aeson qualified as Aeson +import Data.Aeson.KeyMap qualified as KeyMap +import Data.Aeson.Types (parseMaybe, (.:?)) +import Data.ByteString (ByteString) +import Data.ByteString qualified as BS +import Data.Function ((&)) +import Data.Generics.Labels () +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 Data.Vector (Vector) +import Data.Vector qualified as Vector +import Data.Word (Word8) +import Streamly.Data.Fold qualified as Fold +import Streamly.Data.Stream (Stream) +import Streamly.Data.Stream qualified as Stream +import Streamly.Data.Unfold qualified as Unfold + +-- | Flatten a 'Context'\'s messages into a single prompt string +-- suitable for a one-shot CLI invocation. +-- +-- A context whose only message is a 'UserMessage' with a single +-- 'UserText' block is returned verbatim. Anything else is joined +-- with @[role]:@ markers so the CLI still sees a coherent +-- transcript. Image content and tool calls are dropped (CLI +-- providers do not participate in tool use; image bytes cannot be +-- passed verbatim through a CLI). The masterplan documents CLI +-- providers as text-only and tool-incapable. +renderPrompt :: Context -> Text +renderPrompt ctx = + let msgs = Vector.toList (ctx ^. #messages) + in case msgs of + [UserMessage UserPayload {content = uc}] + | [UserText (TextContent t)] <- Vector.toList uc -> + t + _ -> Text.intercalate "\n" (fmap tag msgs) + where + tag :: Message -> Text + tag (UserMessage UserPayload {content = uc}) = + "[user]: " <> flattenUser uc + tag (AssistantMessage AssistantPayload {content = ac}) = + "[assistant]: " <> flattenAssistant ac + tag (ToolResultMessage ToolResultPayload {toolName = n, content = trc}) = + "[tool:" <> n <> "]: " <> flattenToolResult trc + + flattenUser :: Vector UserContent -> Text + flattenUser = Text.concat . Vector.toList . Vector.mapMaybe pickU + where + pickU (UserText (TextContent t)) = Just t + pickU _ = Nothing + + flattenAssistant :: Vector AssistantContent -> Text + flattenAssistant = Text.concat . Vector.toList . Vector.mapMaybe pickA + where + pickA (AssistantText (TextContent t)) = Just t + pickA _ = Nothing + + flattenToolResult :: Vector ToolResultContent -> Text + flattenToolResult = Text.concat . Vector.toList . Vector.mapMaybe pickT + where + pickT (ToolResultText (TextContent t)) = Just t + pickT _ = Nothing + +-- | @maybeApply ma f x@ applies @f a x@ when @ma = Just a@, +-- otherwise returns @x@. Useful for threading optional configuration +-- into a cradle pipeline. +maybeApply :: Maybe a -> (a -> b -> b) -> b -> b +maybeApply Nothing _ b = b +maybeApply (Just a) f b = f a b + +-- | Decode UTF-8 bytes leniently, replacing invalid sequences with +-- U+FFFD. Used for surfacing CLI stderr in +-- 'Baikai.Error.ProcessError'. +decodeUtf8Lenient :: ByteString -> Text +decodeUtf8Lenient = Text.decodeUtf8With Text.lenientDecode + +-- | Best-effort extractor for the assistant text inside a single +-- Codex @--json@ event. See the original implementation's +-- documentation for the schema variants accepted. +extractAgentMessage :: Value -> Maybe Text +extractAgentMessage = parseMaybe parser + where + parser = Aeson.withObject "codex-event" $ \o -> do + mItem <- o .:? "item" + case (mItem :: Maybe Value) of + Just (Aeson.Object io) -> do + ty <- io .:? "type" + case (ty :: Maybe Text) of + Just "agent_message" -> pickText io + _ -> tryMsg o + _ -> tryMsg o + + tryMsg o = do + mInner <- o .:? "msg" + case (mInner :: Maybe Value) of + Just (Aeson.Object io) -> do + ty <- io .:? "type" + case (ty :: Maybe Text) of + Just "agent_message" -> pickText io + _ -> fail "not an agent_message" + _ -> tryFlat o + + tryFlat o = do + ty <- o .:? "type" + case (ty :: Maybe Text) of + Just "agent_message" -> pickText o + _ -> fail "not an agent_message" + + pickText o = case KeyMap.lookup "message" o of + Just (Aeson.String t) -> pure t + _ -> case KeyMap.lookup "text" o of + Just (Aeson.String t) -> pure t + _ -> fail "no payload" + +-- | Consume a stream of stdout bytes from @codex exec --json@, +-- split on newlines, decode each line as JSON, filter to +-- @agent_message@ events, and return the concatenation of their +-- payloads. +parseCodexJsonlStream :: Stream IO ByteString -> IO Text +parseCodexJsonlStream chunks = do + let bytes :: Stream IO Word8 + bytes = Stream.unfoldEach Unfold.fromList (fmap BS.unpack chunks) + lineFold = Fold.takeEndBy_ (== newlineByte) (Fold.foldl' BS.snoc BS.empty) + msgs <- + Stream.foldMany lineFold bytes + & Stream.mapMaybe Aeson.decodeStrict + & Stream.mapMaybe extractAgentMessage + & Stream.fold Fold.toList + pure (Text.concat msgs) + +newlineByte :: Word8 +newlineByte = 0x0a diff --git a/src/Baikai/Provider/Registry.hs b/src/Baikai/Provider/Registry.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Provider/Registry.hs @@ -0,0 +1,103 @@ +-- | The provider registry — the dispatch surface that replaces the +-- prior 'Baikai.Provider' typeclass and 'SomeProvider' existential. +-- +-- An 'ApiProvider' is the per-API handler. EP-3 promotes 'stream' to +-- the primary method: every handler exposes a streaming producer +-- that emits 'AssistantMessageEvent' values, and 'complete' is the +-- synchronous draining wrapper (typically @streamingComplete . stream@). +-- Callers can use an explicit 'ProviderRegistry' handle to isolate handler +-- sets, or use the global convenience registry for simple scripts. +-- +-- 'completeRequest' looks the handler up by the 'Model'\'s 'Api' +-- tag and dispatches; it throws 'Baikai.Error.ProviderError' when +-- no handler is registered for that tag. +module Baikai.Provider.Registry + ( ApiProvider (..), + ProviderRegistry, + newProviderRegistry, + globalProviderRegistry, + registerApiProviderWith, + registerApiProvider, + lookupApiProviderWith, + lookupApiProvider, + completeRequestWith, + completeRequest, + ) +where + +import Baikai.Api (Api, renderApi) +import Baikai.Context (Context) +import Baikai.Error (BaikaiError (..)) +import Baikai.Model (Model (..)) +import Baikai.Options (Options) +import Baikai.Response (Response) +import Baikai.Stream.Event (AssistantMessageEvent) +import Control.Exception (throwIO) +import Data.IORef (IORef, atomicModifyIORef', newIORef, readIORef) +import Data.Map.Strict (Map) +import Data.Map.Strict qualified as Map +import Streamly.Data.Stream (Stream) +import System.IO.Unsafe (unsafePerformIO) + +-- | A per-API handler. 'stream' is the primary streaming +-- entry point; 'complete' is the synchronous draining wrapper, +-- typically @streamingComplete . stream@ from "Baikai.Stream". +data ApiProvider = ApiProvider + { apiTag :: !Api, + stream :: !(Model -> Context -> Options -> Stream IO AssistantMessageEvent), + complete :: !(Model -> Context -> Options -> IO Response) + } + +-- | A mutable provider registry handle. Each handle owns its own handler map, +-- so tests and applications can maintain isolated provider sets in one process. +newtype ProviderRegistry = ProviderRegistry + { registryRef :: IORef (Map Api ApiProvider) + } + +-- | Construct an empty provider registry. +newProviderRegistry :: IO ProviderRegistry +newProviderRegistry = ProviderRegistry <$> newIORef Map.empty + +-- | Process-global handler registry used by the legacy convenience functions. +-- The 'unsafePerformIO' + 'NOINLINE' pattern mirrors 'Baikai.Trace'\'s +-- @eventCounter@; this is a single shared handle for the lifetime of the +-- process. +globalProviderRegistry :: ProviderRegistry +globalProviderRegistry = unsafePerformIO newProviderRegistry +{-# NOINLINE globalProviderRegistry #-} + +-- | Install (or replace) a handler. Idempotent for the same 'Api' +-- tag — calling 'registerApiProviderWith' twice for the same tag keeps only +-- the second handler. +registerApiProviderWith :: ProviderRegistry -> ApiProvider -> IO () +registerApiProviderWith reg p = + atomicModifyIORef' (registryRef reg) $ \m -> (Map.insert (apiTag p) p m, ()) + +-- | Install (or replace) a handler in the process-global registry. +registerApiProvider :: ApiProvider -> IO () +registerApiProvider = registerApiProviderWith globalProviderRegistry + +-- | Look up the handler registered for an 'Api' tag. +lookupApiProviderWith :: ProviderRegistry -> Api -> IO (Maybe ApiProvider) +lookupApiProviderWith reg tag = Map.lookup tag <$> readIORef (registryRef reg) + +-- | Look up the handler registered for an 'Api' tag in the process-global +-- registry. +lookupApiProvider :: Api -> IO (Maybe ApiProvider) +lookupApiProvider = lookupApiProviderWith globalProviderRegistry + +-- | Dispatch a synchronous request through the registered handler +-- for the model's 'Api' tag. Throws 'ProviderError' when no handler +-- is registered for that tag. +completeRequestWith :: ProviderRegistry -> Model -> Context -> Options -> IO Response +completeRequestWith reg m ctx opts = do + mProvider <- lookupApiProviderWith reg (api m) + case mProvider of + Just p -> complete p m ctx opts + Nothing -> + throwIO $ + ProviderError ("No provider registered for API: " <> renderApi (api m)) + +-- | Dispatch a synchronous request through the process-global registry. +completeRequest :: Model -> Context -> Options -> IO Response +completeRequest = completeRequestWith globalProviderRegistry diff --git a/src/Baikai/Response.hs b/src/Baikai/Response.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Response.hs @@ -0,0 +1,70 @@ +-- | The provider response envelope. +-- +-- A 'Response' wraps the model's assistant payload ('message') with +-- metadata baikai owns: the originating 'Api' tag (now typed via +-- 'Baikai.Api'), the input 'Model' echoed back, the provider name, +-- an optional 'responseId', and the measured 'latencyMs'. The +-- 'message' carries only assistant content blocks, 'Usage' (with +-- 'Cost' embedded), 'StopReason', optional error text, and timestamp. +-- +-- 'flattenAssistantBlocks' is a convenience accessor that flattens +-- the message's content blocks. +module Baikai.Response + ( Response (..), + _Response, + responseMessage, + flattenAssistantBlocks, + ) +where + +import Baikai.Api (Api (..)) +import Baikai.Content (AssistantContent) +import Baikai.Message (AssistantPayload (..), Message (..)) +import Baikai.Model (Model, _Model) +import Baikai.StopReason (StopReason (..)) +import Baikai.Usage (_Usage) +import Data.Text (Text) +import Data.Time (UTCTime) +import Data.Vector (Vector) +import Data.Vector qualified as V +import GHC.Generics (Generic) + +data Response = Response + { -- | Assistant-only payload returned by the provider. + message :: !AssistantPayload, + model :: !Model, + api :: !Api, + provider :: !Text, + responseId :: !(Maybe Text), + latencyMs :: !Integer + } + deriving stock (Eq, Show, Generic) + +-- | A blank assistant turn at epoch start. Useful as a fixture base +-- for tests and as the default in error paths where no message was +-- received. +_Response :: Response +_Response = + Response + { message = + AssistantPayload + { content = V.empty, + usage = _Usage, + stopReason = Stop, + errorMessage = Nothing, + timestamp = read "2000-01-01 00:00:00 UTC" :: UTCTime + }, + model = _Model, + api = Custom "", + provider = "", + responseId = Nothing, + latencyMs = 0 + } + +-- | Wrap the response payload as a conversation 'AssistantMessage'. +responseMessage :: Response -> Message +responseMessage r = AssistantMessage (message r) + +-- | Pull the response's assistant content blocks out. +flattenAssistantBlocks :: Response -> Vector AssistantContent +flattenAssistantBlocks Response {message = AssistantPayload {content = c}} = c diff --git a/src/Baikai/StopReason.hs b/src/Baikai/StopReason.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/StopReason.hs @@ -0,0 +1,54 @@ +{-# LANGUAGE LambdaCase #-} + +-- | Why the model stopped generating. +-- +-- A small closed sum populated on every 'Baikai.Message.AssistantMessage'. +-- API providers map their per-vendor stop reason onto this enum; CLI +-- providers populate @Stop@ on success and @ErrorReason@ when the +-- subprocess reports an error. +-- +-- Constructor encoding on the wire is snake-case: @"stop"@, @"length"@, +-- @"tool_use"@, @"error"@, @"aborted"@. @ErrorReason@ is renamed to +-- @"error"@ so the Haskell name does not clash with @Prelude.Either.Left@ +-- callers and the wire shape stays terse. +module Baikai.StopReason (StopReason (..)) where + +import Data.Aeson + ( FromJSON (parseJSON), + Options (..), + ToJSON (toJSON), + defaultOptions, + genericParseJSON, + genericToJSON, + ) +import Data.Char (toLower) +import GHC.Generics (Generic) + +data StopReason + = Stop + | Length + | ToolUse + | ErrorReason + | Aborted + deriving stock (Eq, Show, Generic) + +stopReasonOptions :: Options +stopReasonOptions = + defaultOptions + { constructorTagModifier = \case + "ErrorReason" -> "error" + s -> camelToSnake s + } + where + camelToSnake :: String -> String + camelToSnake [] = [] + camelToSnake (c : cs) = toLower c : go cs + where + go [] = [] + go (x : xs) + | x `elem` ['A' .. 'Z'] = '_' : toLower x : go xs + | otherwise = x : go xs + +instance FromJSON StopReason where parseJSON = genericParseJSON stopReasonOptions + +instance ToJSON StopReason where toJSON = genericToJSON stopReasonOptions diff --git a/src/Baikai/Stream.hs b/src/Baikai/Stream.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Stream.hs @@ -0,0 +1,379 @@ +{-# LANGUAGE LambdaCase #-} + +-- | The streaming entry point. +-- +-- 'streamRequest' looks up the per-API handler registered for a +-- model's 'Api' tag and returns its event stream. When no handler is +-- registered the function returns a one-event error stream (a single +-- 'EventError' with @stopReason = ErrorReason@) so a caller iterating +-- the stream always gets at least one terminal event. +-- +-- 'streamingComplete' folds an event stream into the synchronous +-- 'Response' shape. Vendor providers register both the @stream@ +-- handler and a @complete@ handler derived as +-- @streamingComplete . stream@, so a caller that does not care about +-- deltas keeps the original single-shot ergonomics with no new +-- ceremony. +module Baikai.Stream + ( streamRequest, + streamRequestWith, + streamingComplete, + reassembleResponse, + liftCompleteToStream, + ) +where + +import Baikai.Api (renderApi) +import Baikai.Content + ( AssistantContent (..), + TextContent (..), + ThinkingContent (..), + ) +import Baikai.Content qualified as Content +import Baikai.Context (Context) +import Baikai.Message (AssistantPayload (..), Message (AssistantMessage)) +import Baikai.Message qualified as Msg +import Baikai.Model (Model) +import Baikai.Options (Options) +import Baikai.Provider.Registry + ( ApiProvider (..), + ProviderRegistry, + globalProviderRegistry, + lookupApiProviderWith, + ) +import Baikai.Response (Response (..), responseMessage) +import Baikai.StopReason (StopReason (..)) +import Baikai.Stream.Event + ( AssistantMessageEvent (..), + BlockEndPayload (..), + DeltaPayload (..), + IndexPayload (..), + StartPayload (..), + TerminalPayload (..), + ToolCallEndPayload (..), + ) +import Baikai.Usage (_Usage) +import Control.Exception qualified +import Control.Lens ((%~), (&), (.~), (^.)) +import Data.Aeson qualified as Aeson +import Data.ByteString.Lazy qualified as BSL +import Data.Generics.Labels () +import Data.IntMap.Strict (IntMap) +import Data.IntMap.Strict qualified as IntMap +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.Encoding qualified as Text +import Data.Time (UTCTime, getCurrentTime) +import Data.Time.Clock qualified +import Data.Vector (Vector) +import Data.Vector qualified as Vector +import GHC.Generics (Generic) +import Streamly.Data.Fold (Fold) +import Streamly.Data.Fold qualified as Fold +import Streamly.Data.Stream (Stream) +import Streamly.Data.Stream qualified as Stream + +-- | Dispatch a streaming call through the registered handler for the +-- model's 'Api' tag. Returns a one-event error stream when no handler +-- is registered for that tag — the caller always sees at least one +-- terminal event. +streamRequest :: Model -> Context -> Options -> Stream IO AssistantMessageEvent +streamRequest = streamRequestWith globalProviderRegistry + +-- | Dispatch a streaming call through the selected provider registry. +-- Returns a one-event error stream when no handler is registered for that tag. +streamRequestWith :: + ProviderRegistry -> + Model -> + Context -> + Options -> + Stream IO AssistantMessageEvent +streamRequestWith reg m ctx opts = + Stream.concatEffect $ do + mProvider <- lookupApiProviderWith reg (m ^. #api) + case mProvider of + Just p -> pure (stream p m ctx opts) + Nothing -> pure (Stream.fromEffect (noProviderEvent m)) + +-- | Drain an event stream into a 'Response' by folding events into +-- the assembled message. +streamingComplete :: + (Model -> Context -> Options -> Stream IO AssistantMessageEvent) -> + Model -> + Context -> + Options -> + IO Response +streamingComplete f m ctx opts = + Stream.fold (reassembleResponse m) (f m ctx opts) + +-- | The fold used to reassemble a 'Response' from an event stream. +-- Exposed for the trace bridge so it does not duplicate the +-- reassembly logic. +reassembleResponse :: Model -> Fold IO AssistantMessageEvent Response +reassembleResponse m = + Fold.rmapM finalizeState (Fold.foldl' step (initialState m)) + +-- | One frozen-in-time view of the partial assembly. +data ReassemblyState = ReassemblyState + { model :: !Model, + -- | 'Just' once 'EventStart' has been observed. + skeleton :: !(Maybe Message), + -- | Captured from the EventStart message's timestamp. + startTime :: !(Maybe UTCTime), + -- | Content blocks closed so far, keyed by @contentIndex@. + blocks :: !(IntMap AssistantContent), + -- | Open text accumulators, indexed by @contentIndex@. + textBuf :: !(IntMap Text), + thinkBuf :: !(IntMap Text), + toolArgsBuf :: !(IntMap Text), + -- | 'Just (reason, msg)' once 'EventDone' or 'EventError' fires. + terminal :: !(Maybe (StopReason, Message)) + } + deriving stock (Show, Generic) + +initialState :: Model -> ReassemblyState +initialState m = + ReassemblyState + { model = m, + skeleton = Nothing, + startTime = Nothing, + blocks = IntMap.empty, + textBuf = IntMap.empty, + thinkBuf = IntMap.empty, + toolArgsBuf = IntMap.empty, + terminal = Nothing + } + +step :: ReassemblyState -> AssistantMessageEvent -> ReassemblyState +step s = \case + EventStart StartPayload {partial = sk} -> + let t = case sk of + AssistantMessage AssistantPayload {Msg.timestamp = ts} -> Just ts + _ -> Nothing + in s & #skeleton .~ Just sk & #startTime .~ t + TextStart IndexPayload {contentIndex = i} -> + s & #textBuf %~ IntMap.insert i Text.empty + TextDelta DeltaPayload {contentIndex = i, delta = d} -> + s & #textBuf %~ IntMap.insertWith (\new old -> old <> new) i d + TextEnd BlockEndPayload {contentIndex = i, content = body} -> + s + & #blocks %~ IntMap.insert i (AssistantText (TextContent body)) + & #textBuf %~ IntMap.delete i + ThinkingStart IndexPayload {contentIndex = i} -> + s & #thinkBuf %~ IntMap.insert i Text.empty + ThinkingDelta DeltaPayload {contentIndex = i, delta = d} -> + s & #thinkBuf %~ IntMap.insertWith (\new old -> old <> new) i d + ThinkingEnd BlockEndPayload {contentIndex = i, content = body} -> + s + & #blocks + %~ IntMap.insert + i + ( AssistantThinking + ThinkingContent {thinking = body, signature = Nothing, redacted = False} + ) + & #thinkBuf %~ IntMap.delete i + ToolCallStart IndexPayload {contentIndex = i} -> + s & #toolArgsBuf %~ IntMap.insert i Text.empty + ToolCallDelta DeltaPayload {contentIndex = i, delta = d} -> + s & #toolArgsBuf %~ IntMap.insertWith (\new old -> old <> new) i d + ToolCallEnd ToolCallEndPayload {contentIndex = i, toolCall = tc} -> + s + & #blocks %~ IntMap.insert i (AssistantToolCall tc) + & #toolArgsBuf %~ IntMap.delete i + EventDone TerminalPayload {reason = r, message = msg} -> + s & #terminal .~ Just (r, msg) + EventError TerminalPayload {reason = r, message = msg} -> + s & #terminal .~ Just (r, msg) + +finalizeState :: ReassemblyState -> IO Response +finalizeState s = do + now <- getCurrentTime + let m = s ^. #model + assembled = assembleBlocks s + (terminalMsg, terminalReason) = case s ^. #terminal of + Just (r, msg) -> (msg, r) + Nothing -> (synthesizeTerminal now assembled, Stop) + message' = overrideBlocksAndReason terminalReason terminalMsg assembled now + endTs = message' ^. #timestamp + latency = case s ^. #startTime of + Just startTs -> + round (realToFrac (Data.Time.Clock.diffUTCTime endTs startTs) * (1000 :: Double)) + Nothing -> 0 + pure + Response + { message = message', + model = m, + api = m ^. #api, + provider = m ^. #provider, + responseId = Nothing, + latencyMs = latency + } + +-- | Project the closed content blocks in 'contentIndex' order, plus +-- any still-open text/thinking accumulators flushed as best-effort +-- content (recovery path; providers should not exercise this). +assembleBlocks :: ReassemblyState -> Vector AssistantContent +assembleBlocks s = + let closed = IntMap.elems (s ^. #blocks) + danglingText = + [ AssistantText (TextContent t) + | (_, t) <- IntMap.toAscList (s ^. #textBuf), + not (Text.null t) + ] + danglingThink = + [ AssistantThinking + ThinkingContent {thinking = t, signature = Nothing, redacted = False} + | (_, t) <- IntMap.toAscList (s ^. #thinkBuf), + not (Text.null t) + ] + in Vector.fromList (closed <> danglingText <> danglingThink) + +-- | Replace the content vector and stop reason of an assistant +-- message, preserving usage, errorMessage, and timestamp. +overrideBlocksAndReason :: + StopReason -> Message -> Vector AssistantContent -> UTCTime -> AssistantPayload +overrideBlocksAndReason sr msg blocks fallbackTs = case msg of + AssistantMessage AssistantPayload {usage = u, errorMessage = em, timestamp = ts} -> + AssistantPayload + { Msg.content = blocks, + Msg.usage = u, + Msg.stopReason = sr, + Msg.errorMessage = em, + Msg.timestamp = ts + } + _ -> + AssistantPayload + { Msg.content = blocks, + Msg.usage = _Usage, + Msg.stopReason = sr, + Msg.errorMessage = Just "stream terminated with a non-assistant message", + Msg.timestamp = fallbackTs + } + +-- | Build a placeholder 'AssistantMessage' for streams that ended +-- without a terminal event (a producer bug; the recovery path). +synthesizeTerminal :: UTCTime -> Vector AssistantContent -> Message +synthesizeTerminal now blocks = + AssistantMessage + AssistantPayload + { Msg.content = blocks, + Msg.usage = _Usage, + Msg.stopReason = Stop, + Msg.errorMessage = Just "stream ended without terminal event", + Msg.timestamp = now + } + +-- | Wrap a synchronous @complete@ handler in a one-shot event +-- stream. Used by vendor providers (notably the CLI providers and +-- as a stepping stone for the API providers before they grow native +-- streaming producers) so a 'Baikai.Provider.Registry.ApiProvider' +-- can always populate its 'stream' field. The emitted event +-- sequence is: +-- +-- * 'EventStart' carrying the response's message skeleton (empty +-- content vector; api/provider/model id only). +-- * For each closed content block in the resolved 'Response', a +-- matching @_Start@ / @_Delta@ / @_End@ trio. Text blocks emit +-- one 'TextDelta' with the whole text; tool calls emit one +-- 'ToolCallDelta' with the JSON-encoded arguments. +-- * 'EventDone' carrying the fully assembled assistant message. +-- +-- An exception from @complete@ becomes a single-'EventError' stream +-- with @stopReason = ErrorReason@ and 'errorMessage' set from the +-- exception's display. +liftCompleteToStream :: + (Model -> Context -> Options -> IO Response) -> + Model -> + Context -> + Options -> + Stream IO AssistantMessageEvent +liftCompleteToStream f m ctx opts = + Stream.concatEffect $ do + startTs <- getCurrentTime + er <- tryAny (f m ctx opts) + case er of + Right resp -> pure (Stream.fromList (eventsFor startTs resp)) + Left e -> pure (Stream.fromEffect (errorEvent e)) + +-- | A typed 'try' over any synchronous exception. +tryAny :: IO a -> IO (Either Control.Exception.SomeException a) +tryAny = Control.Exception.try + +-- | Build the synthetic event list for a fully resolved 'Response'. +-- The 'EventStart' carries the supplied @startTs@ on its message +-- skeleton so 'reassembleResponse' can recover 'latencyMs' from the +-- start/end timestamps. +eventsFor :: UTCTime -> Response -> [AssistantMessageEvent] +eventsFor startTs resp = + let payload = resp ^. #message + msg = responseMessage resp + skeleton = + AssistantMessage + AssistantPayload + { Msg.content = Vector.empty, + Msg.usage = payload ^. #usage, + Msg.stopReason = payload ^. #stopReason, + Msg.errorMessage = payload ^. #errorMessage, + Msg.timestamp = startTs + } + blocks = Vector.toList (payload ^. #content) + blockEvents = + concat + [ blockEvent i b + | (i, b) <- zip [0 ..] blocks + ] + reason = payload ^. #stopReason + in [EventStart StartPayload {partial = skeleton}] + <> blockEvents + <> [EventDone TerminalPayload {reason = reason, message = msg}] + +blockEvent :: Int -> AssistantContent -> [AssistantMessageEvent] +blockEvent i = \case + AssistantText (TextContent t) -> + [ TextStart IndexPayload {contentIndex = i}, + TextDelta DeltaPayload {contentIndex = i, delta = t}, + TextEnd BlockEndPayload {contentIndex = i, content = t} + ] + AssistantThinking ThinkingContent {thinking = t} -> + [ ThinkingStart IndexPayload {contentIndex = i}, + ThinkingDelta DeltaPayload {contentIndex = i, delta = t}, + ThinkingEnd BlockEndPayload {contentIndex = i, content = t} + ] + AssistantToolCall tc -> + let argsText = + Text.decodeUtf8 (BSL.toStrict (Aeson.encode (Content.arguments tc))) + in [ ToolCallStart IndexPayload {contentIndex = i}, + ToolCallDelta DeltaPayload {contentIndex = i, delta = argsText}, + ToolCallEnd ToolCallEndPayload {contentIndex = i, toolCall = tc} + ] + +errorEvent :: Control.Exception.SomeException -> IO AssistantMessageEvent +errorEvent e = do + now <- getCurrentTime + let msg = + AssistantMessage + AssistantPayload + { Msg.content = Vector.empty, + Msg.usage = _Usage, + Msg.stopReason = ErrorReason, + Msg.errorMessage = Just (Text.pack (Control.Exception.displayException e)), + Msg.timestamp = now + } + pure (EventError TerminalPayload {reason = ErrorReason, message = msg}) + +-- | The single 'EventError' value used when no provider is +-- registered for the model's API tag. +noProviderEvent :: Model -> IO AssistantMessageEvent +noProviderEvent m = do + now <- getCurrentTime + let msg = + AssistantMessage + AssistantPayload + { Msg.content = Vector.empty, + Msg.usage = _Usage, + Msg.stopReason = ErrorReason, + Msg.errorMessage = + Just ("No provider registered for API: " <> renderApi (m ^. #api)), + Msg.timestamp = now + } + pure (EventError TerminalPayload {reason = ErrorReason, message = msg}) diff --git a/src/Baikai/Stream/Event.hs b/src/Baikai/Stream/Event.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Stream/Event.hs @@ -0,0 +1,162 @@ +{-# LANGUAGE LambdaCase #-} + +-- | The streaming event algebra: 'AssistantMessageEvent'. +-- +-- A provider call exposes its progress as a 'Streamly.Data.Stream.Stream +-- IO AssistantMessageEvent'. The stream begins with a single +-- 'EventStart' carrying an empty 'AssistantMessage' skeleton (api, +-- provider, model id), interleaves per-content-block lifecycle events +-- (@_Start@ / @_Delta@ / @_End@) keyed by 'contentIndex', and +-- terminates with exactly one 'EventDone' (success) or 'EventError' +-- (any failure that bubbled out of the producer). The terminal event +-- carries the fully assembled 'AssistantMessage' so a consumer that +-- only pattern-matches on the terminal event still gets a correct +-- response without folding deltas. +-- +-- The algebra is closed and shared by every provider in baikai 0.1. +-- Adding a new variant is a breaking change to baikai's public +-- surface. Consumers that prefer source resilience over exhaustiveness +-- can include a final wildcard branch in their pattern match, but +-- providers MUST NOT emit untyped provider-specific events through +-- this algebra. Providers MUST emit @_Start@, then zero or more +-- @_Delta@, then @_End@ for each content block in increasing +-- @contentIndex@ order; the reassembler defends against missing +-- @_End@ events but no built-in provider should rely on the recovery +-- path. +module Baikai.Stream.Event + ( AssistantMessageEvent (..), + StartPayload (..), + IndexPayload (..), + DeltaPayload (..), + BlockEndPayload (..), + ToolCallEndPayload (..), + TerminalPayload (..), + isTerminal, + ) +where + +import Baikai.Content (ToolCall) +import Baikai.Message (Message) +import Baikai.StopReason (StopReason) +import Data.Aeson (ToJSON) +import Data.Text (Text) +import GHC.Generics (Generic) + +-- | One observable step of an assistant turn. +-- +-- @contentIndex@ identifies which content block a per-block event +-- refers to. Indices are non-negative and strictly increasing in the +-- order blocks are opened; gaps are permitted (an OpenAI stream that +-- emits a text block at index 0 then a tool call at index 2 is +-- legal, as long as the @_Start@ of index 2 follows the @_End@ of +-- index 0). +-- Each constructor wraps a dedicated single-constructor payload record +-- rather than carrying record fields directly on the sum. A field +-- selector defined across this multi-constructor algebra would be +-- partial (e.g. a @delta@ that throws on 'EventStart'); @-Wpartial-fields@ +-- flags that latent runtime hazard. Same-shaped constructors share one +-- payload type ('IndexPayload', 'DeltaPayload', 'BlockEndPayload', +-- 'TerminalPayload'), so the payloads also collapse the repetition the +-- flat encoding spread across twelve constructors. +data AssistantMessageEvent + = -- | The first event in every stream. The payload's 'partial' is an + -- 'AssistantMessage' with empty content; downstream consumers that + -- care only about the message skeleton (api, provider, model id) + -- can read it here. + EventStart StartPayload + | -- | A text content block is about to receive deltas. + TextStart IndexPayload + | -- | A chunk of text appended to the text block at @contentIndex@. + TextDelta DeltaPayload + | -- | The text block at @contentIndex@ is closed; the payload's + -- 'content' is the concatenation of every preceding 'TextDelta' + -- for the same index. + TextEnd BlockEndPayload + | -- | A thinking content block is about to receive deltas. + ThinkingStart IndexPayload + | -- | A chunk of reasoning text appended to the thinking block. + ThinkingDelta DeltaPayload + | -- | The thinking block at @contentIndex@ is closed. + ThinkingEnd BlockEndPayload + | -- | A tool-call content block is about to receive argument-JSON + -- chunks. The tool's @id_@ and @name@ become known on this event + -- when the upstream API delivers them up front (Anthropic), or on + -- subsequent deltas otherwise (OpenAI). + ToolCallStart IndexPayload + | -- | A chunk of the tool call's arguments JSON. Concatenating + -- every 'ToolCallDelta' for a given @contentIndex@ yields a + -- syntactically valid JSON value. + ToolCallDelta DeltaPayload + | -- | The tool call at @contentIndex@ is closed; the payload's + -- 'toolCall' is the fully parsed call ('id_', 'name', and decoded + -- 'arguments'). + ToolCallEnd ToolCallEndPayload + | -- | The stream's terminal success event. The payload's 'message' + -- is the fully assembled 'AssistantMessage' carrying all content + -- blocks, the final 'Usage', and the resolved 'StopReason'. + EventDone TerminalPayload + | -- | The stream's terminal failure event. The payload's 'message' + -- is an 'AssistantMessage' carrying whatever content blocks were + -- already closed before the failure, plus a populated + -- 'errorMessage' and @stopReason = ErrorReason@ or + -- @stopReason = Aborted@. + EventError TerminalPayload + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +-- | Payload of 'EventStart': the message skeleton observed up front. +newtype StartPayload = StartPayload + { partial :: Message + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +-- | Payload of the @_Start@ events: which content block is opening. +newtype IndexPayload = IndexPayload + { contentIndex :: Int + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +-- | Payload of the @_Delta@ events: a chunk appended to the block at +-- @contentIndex@. +data DeltaPayload = DeltaPayload + { contentIndex :: !Int, + delta :: !Text + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +-- | Payload of the text/thinking @_End@ events: the closed block's +-- fully concatenated 'content'. +data BlockEndPayload = BlockEndPayload + { contentIndex :: !Int, + content :: !Text + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +-- | Payload of 'ToolCallEnd': the fully parsed 'toolCall'. +data ToolCallEndPayload = ToolCallEndPayload + { contentIndex :: !Int, + toolCall :: !ToolCall + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +-- | Payload of the terminal events 'EventDone' and 'EventError': the +-- resolved 'StopReason' and the assembled 'message'. +data TerminalPayload = TerminalPayload + { reason :: !StopReason, + message :: !Message + } + deriving stock (Eq, Show, Generic) + deriving anyclass (ToJSON) + +-- | 'True' when the event terminates the stream — exactly one +-- 'EventDone' or 'EventError' is emitted per call. +isTerminal :: AssistantMessageEvent -> Bool +isTerminal = \case + EventDone {} -> True + EventError {} -> True + _ -> False diff --git a/src/Baikai/ThinkingLevel.hs b/src/Baikai/ThinkingLevel.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/ThinkingLevel.hs @@ -0,0 +1,48 @@ +{-# LANGUAGE LambdaCase #-} + +-- | Provider-agnostic reasoning-effort preference. +-- +-- Each provider maps the value to its own primitive: OpenAI uses +-- @reasoning_effort: "minimal" | "low" | "medium" | "high"@; +-- Anthropic uses a token budget on its @thinking@ field. +-- 'thinkingTokenBudget' provides the recommended budget for the +-- token-based providers. +module Baikai.ThinkingLevel + ( ThinkingLevel (..), + renderThinkingLevel, + thinkingTokenBudget, + ) +where + +import Data.Aeson (FromJSON, ToJSON) +import Data.Text (Text) +import GHC.Generics (Generic) +import Numeric.Natural (Natural) + +-- | The four buckets pi-mono settled on. Coarse enough that callers +-- do not need to learn each provider's TTL conventions, fine enough +-- to influence behaviour in a meaningful way. +data ThinkingLevel + = ThinkingMinimal + | ThinkingLow + | ThinkingMedium + | ThinkingHigh + deriving stock (Eq, Ord, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | The wire string OpenAI's @reasoning_effort@ field expects. +renderThinkingLevel :: ThinkingLevel -> Text +renderThinkingLevel = \case + ThinkingMinimal -> "minimal" + ThinkingLow -> "low" + ThinkingMedium -> "medium" + ThinkingHigh -> "high" + +-- | Recommended token budget for providers that take an explicit +-- count (Anthropic's @thinking.budget_tokens@). +thinkingTokenBudget :: ThinkingLevel -> Natural +thinkingTokenBudget = \case + ThinkingMinimal -> 1024 + ThinkingLow -> 2048 + ThinkingMedium -> 8192 + ThinkingHigh -> 16384 diff --git a/src/Baikai/Tool.hs b/src/Baikai/Tool.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Tool.hs @@ -0,0 +1,85 @@ +-- | The 'Tool' type and tool-choice control. +-- +-- A 'Tool' is a name, a human-readable description, and a JSON +-- Schema describing the parameters. The schema is carried as a +-- 'Data.Aeson.Value' rather than a typed schema GADT so callers can +-- hand-write or generate the schema however they like; the +-- provider-side encoders forward it to the upstream API verbatim. +-- +-- 'ToolChoice' tells the model how aggressively to use the tools. +-- The default in 'Baikai.Options.Options' is 'Nothing' which lets +-- the provider apply its own default (typically @auto@ for +-- Anthropic and OpenAI). +-- +-- The conversation-level helper for the common case — execute the +-- model's tool calls and append the results to the next request — +-- lives in 'Baikai.Context.appendToolResult' and +-- 'Baikai.Context.appendToolResultText' to avoid an import cycle +-- between this module (which 'Baikai.Context' imports for the @tools@ +-- field type) and 'Baikai.Context' itself. +module Baikai.Tool + ( Tool (..), + ToolChoice (..), + _Tool, + ) +where + +import Data.Aeson + ( FromJSON (..), + Options (..), + SumEncoding (..), + ToJSON (..), + Value (..), + camelTo2, + defaultOptions, + genericParseJSON, + genericToJSON, + ) +import Data.Text (Text) +import Data.Text qualified as Text +import GHC.Generics (Generic) + +-- | A caller-declared tool. @parameters@ holds a JSON Schema; the +-- provider-side encoders pass it through unchanged. +data Tool = Tool + { name :: !Text, + description :: !Text, + parameters :: !Value + } + deriving stock (Eq, Show, Generic) + deriving anyclass (FromJSON, ToJSON) + +-- | How the model should pick between the registered tools. +-- +-- * 'ToolChoiceAuto' — model decides (the default at most providers). +-- * 'ToolChoiceNone' — disable tool calling for this request. +-- * 'ToolChoiceRequired' — must call some tool. +-- * 'ToolChoiceSpecific' — must call this exact tool by name. +data ToolChoice + = ToolChoiceAuto + | ToolChoiceNone + | ToolChoiceRequired + | ToolChoiceSpecific !Text + deriving stock (Eq, Show, Generic) + +toolChoiceOptions :: Options +toolChoiceOptions = + defaultOptions + { sumEncoding = TaggedObject {tagFieldName = "type", contentsFieldName = "name"}, + constructorTagModifier = camelTo2 '_' . drop (Text.length "ToolChoice") + } + +instance FromJSON ToolChoice where + parseJSON = genericParseJSON toolChoiceOptions + +instance ToJSON ToolChoice where + toJSON = genericToJSON toolChoiceOptions + +-- | An empty tool — useful as a base for record updates. +_Tool :: Tool +_Tool = + Tool + { name = Text.empty, + description = Text.empty, + parameters = Null + } diff --git a/src/Baikai/Trace.hs b/src/Baikai/Trace.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Trace.hs @@ -0,0 +1,351 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +-- | The 'withTrace' wrapper and supporting helpers. +-- +-- After EP-3, the trace bridge is stream-shaped at the core: +-- 'withTraceStream' returns a 'Stream IO AssistantMessageEvent' +-- that side-effects 'CallStarted' / 'CallFinished' / 'CallFailed' +-- events to a user-supplied 'TraceSink' as the stream's lifecycle +-- unfolds. 'withTrace' is the synchronous draining wrapper: +-- @withTrace sink m ctx opts = Stream.fold (reassembleResponse m) +-- (withTraceStream sink m ctx opts)@. +-- +-- Per-call plumbing: open a 'Chan' of @Maybe TraceEvent@, fork a +-- worker thread that drains the channel through the sink's fold, +-- push 'CallStarted' eagerly (before the first +-- 'AssistantMessageEvent' is emitted), then watch for the stream's +-- terminal event ('EventDone' or 'EventError') and push the +-- matching 'CallFinished' / 'CallFailed' before yielding the +-- terminal event to the consumer. Cleanup ('Nothing' sentinel on +-- the channel + 'takeMVar' on the worker) is idempotent and runs +-- through 'Stream.finallyIO' so an early-aborting consumer never +-- leaks the worker. +module Baikai.Trace + ( -- * Re-exports + TraceEvent (..), + TraceSink (..), + + -- * Wrappers + withTrace, + withTraceWith, + withTraceStream, + withTraceStreamWith, + runRequestWith, + runRequestWithRegistry, + + -- * Helpers + newEventId, + summarizeContext, + ) +where + +import Baikai.Context (Context) +import Baikai.Cost (usdAsScientific) +import Baikai.Cost qualified as Cost +import Baikai.Cost.Log + ( CallLogEntry (..), + CallLogHandle, + appendEntry, + summarizeContext, + ) +import Baikai.Message (AssistantPayload (..), Message (..)) +import Baikai.Model (Model) +import Baikai.Options (Options) +import Baikai.Prelude +import Baikai.Provider.Registry (ProviderRegistry, globalProviderRegistry) +import Baikai.Response (Response) +import Baikai.Stream (reassembleResponse, streamRequestWith) +import Baikai.Stream.Event (AssistantMessageEvent (..), TerminalPayload (..)) +import Baikai.Trace.Event (TraceEvent (..)) +import Baikai.Trace.Sink (TraceSink (..)) +import Baikai.Usage (Usage) +import Baikai.Usage qualified as Usage +import Control.Concurrent (forkIO) +import Control.Concurrent.Chan (Chan, newChan, readChan, writeChan) +import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar, takeMVar) +-- Control.Exception no longer needed; producer errors flow through +-- the stream as 'EventError'. +import Control.Monad (unless) +import Control.Monad.IO.Unlift (MonadUnliftIO, withRunInIO) +import Data.Bits (unsafeShiftL, (.&.), (.|.)) +import Data.IORef (IORef, atomicModifyIORef', newIORef) +import Data.Maybe (fromMaybe) +import Data.Text qualified as Text +import Data.Time (UTCTime, diffUTCTime, getCurrentTime) +import Data.Time.Clock.POSIX (getPOSIXTime) +import Numeric (showHex) +import Streamly.Data.Stream (Stream) +import Streamly.Data.Stream qualified as Stream +import System.IO.Unsafe (unsafePerformIO) + +-- ============================================================ +-- Stream-shaped trace bridge +-- ============================================================ + +-- | Decorate a streaming provider call with structured tracing +-- events emitted to the supplied 'TraceSink'. +-- +-- The returned stream yields the same 'AssistantMessageEvent's as +-- 'Baikai.Stream.streamRequest'. As a side effect: one 'CallStarted' +-- event is pushed to the sink before the first +-- 'AssistantMessageEvent' is observed, and one 'CallFinished' (on +-- 'EventDone') or 'CallFailed' (on 'EventError') is pushed before +-- the terminal event is yielded to the consumer. Intermediate delta +-- events are not traced — emitting one trace per token would +-- explode trace volume. +withTraceStream :: + TraceSink -> + Model -> + Context -> + Options -> + Stream IO AssistantMessageEvent +withTraceStream = withTraceStreamWith globalProviderRegistry + +-- | Decorate a streaming provider call through an explicit registry handle. +withTraceStreamWith :: + ProviderRegistry -> + TraceSink -> + Model -> + Context -> + Options -> + Stream IO AssistantMessageEvent +withTraceStreamWith reg (TraceSink sinkFold) m ctx opts = + Stream.concatEffect $ do + state <- newTraceState + let c = state ^. #chan + d = state ^. #done + _ <- + forkIO $ do + let stepDrain () = do + msg <- readChan c + pure (fmap (\e -> (e, ())) msg) + Stream.unfoldrM stepDrain () + & Stream.fold sinkFold + putMVar d () + eid <- newEventId + start <- getCurrentTime + writeChan c $ + Just + CallStarted + { eventId = eid, + timestamp = start, + provider = m ^. #provider, + model = m ^. #modelId, + maxTokens = resolvedMaxTokens m opts, + promptSummary = summarizeContext ctx + } + pure $ + Stream.finallyIO + (cleanupTrace state) + (Stream.mapM (traceEvent state eid start m) (streamRequestWith reg m ctx opts)) + +-- | Synchronous trace wrapper. Drains 'withTraceStream' into a +-- 'Response' through 'reassembleResponse'. +-- +-- Unlike the EP-2 'withTrace' (which re-threw the producer's +-- exception), this implementation never throws for producer-side +-- failures: errors flow through the stream as a terminal +-- 'EventError' and the drained 'Response' carries +-- @stopReason = ErrorReason@ plus 'errorMessage'. The masterplan's +-- Vision & Scope section commits to "partial output is always +-- recoverable" and the plan's Decision Log records that producer +-- failures must surface as response data, not exceptions. +-- Downstream-of-the-fold exceptions (e.g. an 'appendEntry' that +-- fails) still propagate unchanged. +withTrace :: + (MonadUnliftIO m) => + TraceSink -> Model -> Context -> Options -> m Response +withTrace = withTraceWith globalProviderRegistry + +-- | Synchronous trace wrapper that dispatches through an explicit provider +-- registry handle. +withTraceWith :: + (MonadUnliftIO m) => + ProviderRegistry -> + TraceSink -> + Model -> + Context -> + Options -> + m Response +withTraceWith reg sink model ctx opts = + withRunInIO $ \_ -> + Stream.fold + (reassembleResponse model) + (withTraceStreamWith reg sink model ctx opts) + +-- ============================================================ +-- Per-call trace state +-- ============================================================ + +data TraceState = TraceState + { chan :: !(Chan (Maybe TraceEvent)), + done :: !(MVar ()), + closed :: !(IORef Bool) + } + deriving stock (Generic) + +newTraceState :: IO TraceState +newTraceState = do + c <- newChan + d <- newEmptyMVar + r <- newIORef False + pure TraceState {chan = c, done = d, closed = r} + +cleanupTrace :: TraceState -> IO () +cleanupTrace s = do + alreadyClosed <- + atomicModifyIORef' (s ^. #closed) (\b -> (True, b)) + unless alreadyClosed $ do + writeChan (s ^. #chan) Nothing + takeMVar (s ^. #done) + +traceEvent :: + TraceState -> + Text -> + UTCTime -> + Model -> + AssistantMessageEvent -> + IO AssistantMessageEvent +traceEvent state eid start m ev = do + case ev of + EventDone TerminalPayload {message = msg} -> do + now <- getCurrentTime + let latency = millisBetween start now + mu = assistantUsageFromMsg msg + meaningfulCost = maybe False (\u -> usdRat (Usage.cost u) > 0) mu + finished = + CallFinished + { eventId = eid, + timestamp = now, + provider = m ^. #provider, + model = m ^. #modelId, + latencyMs = latency, + inputTokens = fmap Usage.inputTokens mu, + outputTokens = fmap Usage.outputTokens mu, + usd = + if meaningfulCost + then fmap (usdAsScientific . Usage.cost) mu + else Nothing + } + writeChan (state ^. #chan) (Just finished) + cleanupTrace state + EventError TerminalPayload {message = msg} -> do + now <- getCurrentTime + let latency = millisBetween start now + errMsg = case msg of + AssistantMessage AssistantPayload {errorMessage = Just t} -> t + _ -> "stream terminated with EventError" + failed = + CallFailed + { eventId = eid, + timestamp = now, + provider = m ^. #provider, + model = m ^. #modelId, + latencyMs = latency, + errorMessage = errMsg + } + writeChan (state ^. #chan) (Just failed) + cleanupTrace state + _ -> pure () + pure ev + +-- ============================================================ +-- Cost-log convenience wrapper +-- ============================================================ + +-- | Combine tracing with call-log persistence in one call. Traces +-- the streaming call, drains it into a 'Response', then appends a +-- 'CallLogEntry' to the given handle. +runRequestWith :: + (MonadUnliftIO m) => + TraceSink -> + CallLogHandle -> + Model -> + Context -> + Options -> + m Response +runRequestWith = runRequestWithRegistry globalProviderRegistry + +-- | Combine tracing with call-log persistence while dispatching through an +-- explicit provider registry handle. +runRequestWithRegistry :: + (MonadUnliftIO m) => + ProviderRegistry -> + TraceSink -> + CallLogHandle -> + Model -> + Context -> + Options -> + m Response +runRequestWithRegistry reg sink h m ctx opts = do + resp <- withTraceWith reg sink m ctx opts + now <- liftIO getCurrentTime + let mu = assistantUsage resp + meaningfulCost = maybe False (\u -> usdRat (Usage.cost u) > 0) mu + entry = + CallLogEntry + { timestamp = now, + provider = m ^. #provider, + model = m ^. #modelId, + inputTokens = mu >>= positiveNat . Usage.inputTokens, + outputTokens = mu >>= positiveNat . Usage.outputTokens, + cachedInputTokens = mu >>= positiveNat . Usage.cacheReadTokens, + reasoningTokens = mu >>= Usage.reasoningTokens, + usd = + if meaningfulCost + then fmap (usdAsScientific . Usage.cost) mu + else Nothing, + latencyMs = resp ^. #latencyMs, + promptSummary = summarizeContext ctx + } + appendEntry h entry + pure resp + +-- ============================================================ +-- Internal helpers +-- ============================================================ + +-- | Project the assistant turn's 'Usage' out of a response. +assistantUsage :: Response -> Maybe Usage +assistantUsage resp = Just ((resp ^. #message) ^. #usage) + +assistantUsageFromMsg :: Message -> Maybe Usage +assistantUsageFromMsg = \case + AssistantMessage AssistantPayload {usage = u} -> Just u + _ -> Nothing + +usdRat :: Cost.Cost -> Rational +usdRat = Cost.usd + +positiveNat :: Natural -> Maybe Natural +positiveNat 0 = Nothing +positiveNat n = Just n + +resolvedMaxTokens :: Model -> Options -> Natural +resolvedMaxTokens m opts = fromMaybe (m ^. #maxOutputTokens) (opts ^. #maxTokens) + +millisBetween :: UTCTime -> UTCTime -> Integer +millisBetween a b = round (realToFrac (diffUTCTime b a) * (1000 :: Double)) + +-- ============================================================ +-- Event id +-- ============================================================ + +newEventId :: IO Text +newEventId = do + n <- atomicModifyIORef' eventCounter (\k -> (k + 1, k)) + let raw :: Word + raw = (eventBase .&. 0xFFFF) `unsafeShiftL` 16 .|. (n .&. 0xFFFF) + hex = showHex raw "" + padded = replicate (8 - length hex) '0' <> hex + pure (Text.pack padded) + +eventCounter :: IORef Word +eventCounter = unsafePerformIO (newIORef 0) +{-# NOINLINE eventCounter #-} + +eventBase :: Word +eventBase = unsafePerformIO $ do + t <- getPOSIXTime + pure (fromIntegral (floor t :: Integer)) +{-# NOINLINE eventBase #-} diff --git a/src/Baikai/Trace/Event.hs b/src/Baikai/Trace/Event.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Trace/Event.hs @@ -0,0 +1,95 @@ +{-# OPTIONS_GHC -Wno-partial-fields #-} + +-- | The 'TraceEvent' sum and its JSON encoding. +-- +-- A trace event is one of three discriminated cases: 'CallStarted' fires +-- when a provider call begins, 'CallFinished' when it returns a response, +-- and 'CallFailed' when it throws. The 'sumEncoding' tag field is @kind@, +-- so a JSON-Lines stream of these can be filtered with +-- @jq 'select(.kind == "call_finished")'@. +module Baikai.Trace.Event + ( TraceEvent (..), + traceEventOptions, + ) +where + +import Data.Aeson + ( FromJSON (parseJSON), + Options (..), + SumEncoding (..), + ToJSON (..), + defaultOptions, + genericParseJSON, + genericToEncoding, + genericToJSON, + ) +import Data.Char (toLower) +import Data.Scientific (Scientific) +import Data.Text (Text) +import Data.Time (UTCTime) +import GHC.Generics (Generic) +import Numeric.Natural (Natural) + +-- | One observable event from a provider call. +-- +-- Every event carries an 'eventId' that correlates the @started@ event +-- with its matching @finished@ or @failed@ event within a single process +-- run. Token counts and dollar cost are 'Maybe' because subscription-based +-- providers (the CLIs) do not report them; 'omitNothingFields' keeps the +-- absent fields out of the rendered JSON. +data TraceEvent + = CallStarted + { eventId :: !Text, + timestamp :: !UTCTime, + provider :: !Text, + model :: !Text, + maxTokens :: !Natural, + promptSummary :: !Text + } + | CallFinished + { eventId :: !Text, + timestamp :: !UTCTime, + provider :: !Text, + model :: !Text, + latencyMs :: !Integer, + inputTokens :: !(Maybe Natural), + outputTokens :: !(Maybe Natural), + usd :: !(Maybe Scientific) + } + | CallFailed + { eventId :: !Text, + timestamp :: !UTCTime, + provider :: !Text, + model :: !Text, + latencyMs :: !Integer, + errorMessage :: !Text + } + deriving stock (Eq, Show, Generic) + +-- | Aeson options shared by 'ToJSON' and 'FromJSON' instances. +-- +-- * Sum encoding: @{"kind":"","data":{...}}@. +-- * Constructor tags: snake-case (@call_started@, @call_finished@, +-- @call_failed@). +-- * Field labels: kept as-is (camelCase). +-- * Nothing fields are dropped from the encoded JSON. +traceEventOptions :: Options +traceEventOptions = + defaultOptions + { sumEncoding = TaggedObject {tagFieldName = "kind", contentsFieldName = "data"}, + constructorTagModifier = dropWhile (== '_') . camelToSnake, + omitNothingFields = True + } + where + camelToSnake :: String -> String + camelToSnake [] = [] + camelToSnake (c : cs) + | c `elem` ['A' .. 'Z'] = '_' : toLower c : camelToSnake cs + | otherwise = c : camelToSnake cs + +instance ToJSON TraceEvent where + toJSON = genericToJSON traceEventOptions + toEncoding = genericToEncoding traceEventOptions + +instance FromJSON TraceEvent where + parseJSON = genericParseJSON traceEventOptions diff --git a/src/Baikai/Trace/Sink.hs b/src/Baikai/Trace/Sink.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Trace/Sink.hs @@ -0,0 +1,100 @@ +{-# LANGUAGE LambdaCase #-} + +-- | The 'TraceSink' newtype and four built-in sinks. +-- +-- A 'TraceSink' is a thin wrapper around a streamly @'Fold' IO 'TraceEvent' ()@. +-- The fold shape is deliberate: 'Streamly.Data.Fold' exports composition +-- combinators like 'Fold.tee' (fan to two folds), 'Fold.filter' (drop inputs +-- failing a predicate), and 'Fold.lmap' (project each input), so future +-- sinks (OpenTelemetry, redaction, projection) plug in without an adapter. +module Baikai.Trace.Sink + ( TraceSink (..), + silent, + stdoutSink, + fileSink, + multiSink, + renderHuman, + ) +where + +import Baikai.Trace.Event (TraceEvent (..)) +import Data.Aeson qualified as Aeson +import Data.ByteString.Lazy qualified as BSL +import Data.Text (Text) +import Data.Text qualified as Text +import Data.Text.IO qualified as Text.IO +import Data.Time (defaultTimeLocale, formatTime) +import Streamly.Data.Fold (Fold) +import Streamly.Data.Fold qualified as Fold +import System.IO (IOMode (AppendMode), withFile) + +-- | A trace sink is a streamly fold over 'TraceEvent' values. Folds +-- compose: 'Fold.tee' fans events to two sinks, 'Fold.filter' drops +-- events that fail a predicate, 'Fold.lmap' projects each event before +-- feeding the inner fold. +newtype TraceSink = TraceSink + { runSink :: Fold IO TraceEvent () + } + +-- | Consume events without effect. Useful in tests. +silent :: TraceSink +silent = TraceSink Fold.drain + +-- | Print each event to stdout using 'renderHuman'. +stdoutSink :: TraceSink +stdoutSink = TraceSink (Fold.drainMapM (Text.IO.putStrLn . renderHuman)) + +-- | Append each event as one JSON-encoded line. Open-per-write is +-- intentional: crash safety beats throughput for trace events. A +-- kept-open-handle variant is a future enhancement. +fileSink :: FilePath -> IO TraceSink +fileSink path = + pure $ TraceSink $ Fold.drainMapM $ \e -> + withFile path AppendMode $ \h -> + BSL.hPut h (Aeson.encode e <> "\n") + +-- | Fan every event out to every sink in the list. Implemented by folding +-- 'Fold.tee' across the input list; 'Fold.tee' runs both folds on each +-- input and returns the pair of their accumulators, which we discard. +multiSink :: [TraceSink] -> TraceSink +multiSink sinks = + TraceSink (foldr step Fold.drain sinks) + where + step (TraceSink f) acc = fmap (const ()) (Fold.tee f acc) + +-- | Format an event as a single human-readable line. +renderHuman :: TraceEvent -> Text +renderHuman = \case + CallStarted {timestamp, provider, model, maxTokens, promptSummary} -> + Text.unwords + [ "[" <> fmtTime timestamp <> "]", + provider, + model, + "START", + "max=" <> tshow maxTokens, + Text.take 80 promptSummary + ] + CallFinished {timestamp, provider, model, latencyMs, inputTokens, outputTokens, usd} -> + Text.unwords + [ "[" <> fmtTime timestamp <> "]", + provider, + model, + "->", + tshow latencyMs <> "ms", + maybe "" (\n -> "in=" <> tshow n) inputTokens, + maybe "" (\n -> "out=" <> tshow n) outputTokens, + maybe "(no-cost)" (\s -> "$" <> tshow s) usd + ] + CallFailed {timestamp, provider, model, latencyMs, errorMessage} -> + Text.unwords + [ "[" <> fmtTime timestamp <> "]", + provider, + model, + "FAILED", + tshow latencyMs <> "ms:", + errorMessage + ] + where + tshow :: (Show a) => a -> Text + tshow x = Text.pack (show x) + fmtTime t = Text.pack (formatTime defaultTimeLocale "%Y-%m-%dT%H:%M:%SZ" t) diff --git a/src/Baikai/Usage.hs b/src/Baikai/Usage.hs new file mode 100644 --- /dev/null +++ b/src/Baikai/Usage.hs @@ -0,0 +1,48 @@ +-- | Token-usage accounting for a single provider call. +-- +-- The shape splits cache-read tokens (the prompt fragment served from a +-- prior cache write) from cache-write tokens (the prompt fragment the +-- provider stored for future calls), so callers can reason about both +-- billing dimensions. 'cost' is always populated — providers without +-- pricing data fill it with '_Cost' (zero across all rates) rather than +-- a 'Nothing' that every cost-reading caller would have to handle. +module Baikai.Usage (Usage (..), _Usage) where + +import Baikai.Cost (Cost, _Cost) +import Data.Aeson + ( Options (fieldLabelModifier), + ToJSON (toJSON), + camelTo2, + defaultOptions, + genericToJSON, + ) +import GHC.Generics (Generic) +import Numeric.Natural (Natural) + +data Usage = Usage + { inputTokens :: !Natural, + outputTokens :: !Natural, + cacheReadTokens :: !Natural, + cacheWriteTokens :: !Natural, + reasoningTokens :: !(Maybe Natural), + totalTokens :: !Natural, + cost :: !Cost + } + deriving stock (Eq, Show, Generic) + +usageOptions :: Options +usageOptions = defaultOptions {fieldLabelModifier = camelTo2 '_'} + +instance ToJSON Usage where toJSON = genericToJSON usageOptions + +_Usage :: Usage +_Usage = + Usage + { inputTokens = 0, + outputTokens = 0, + cacheReadTokens = 0, + cacheWriteTokens = 0, + reasoningTokens = Nothing, + totalTokens = 0, + cost = _Cost + } diff --git a/test/AgentAssetsSpec.hs b/test/AgentAssetsSpec.hs new file mode 100644 --- /dev/null +++ b/test/AgentAssetsSpec.hs @@ -0,0 +1,78 @@ +module AgentAssetsSpec (tests) where + +import Baikai.AgentAssets +import Baikai.Interactive +import Data.Text qualified as Text +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Baikai.AgentAssets" + [ pathTests, + layoutTests, + codexTomlTest + ] + +pathTests :: TestTree +pathTests = + testGroup + "target paths" + [ testCase "Claude Code project paths" $ do + skillTargetPath InteractiveClaude InteractiveProjectScope "example" + @?= ".claude/skills/example" + agentTargetPath InteractiveClaude InteractiveProjectScope "example" + @?= ".claude/agents/example.md", + testCase "Claude Code user paths" $ do + skillTargetPath InteractiveClaude InteractiveUserScope "example" + @?= "$HOME/.claude/skills/example" + agentTargetPath InteractiveClaude InteractiveUserScope "example" + @?= "$HOME/.claude/agents/example.md", + testCase "Codex project paths use .agents for skills and .codex for agents" $ do + skillTargetPath InteractiveCodex InteractiveProjectScope "example" + @?= ".agents/skills/example" + agentTargetPath InteractiveCodex InteractiveProjectScope "example" + @?= ".codex/agents/example.toml", + testCase "Codex user paths use home discovery roots" $ do + skillTargetPath InteractiveCodex InteractiveUserScope "example" + @?= "$HOME/.agents/skills/example" + agentTargetPath InteractiveCodex InteractiveUserScope "example" + @?= "$HOME/.codex/agents/example.toml" + ] + +layoutTests :: TestTree +layoutTests = + testGroup + "layout metadata" + [ testCase "skills are directory assets" $ do + agentAssetFormat InteractiveClaude SkillAsset @?= DirectoryAsset + agentAssetFormat InteractiveCodex SkillAsset @?= DirectoryAsset, + testCase "custom agents use provider-native file formats" $ do + agentAssetFormat InteractiveClaude CustomAgentAsset @?= MarkdownFile + agentAssetFormat InteractiveCodex CustomAgentAsset @?= TomlFile, + testCase "layout carries provider, scope, kind, format, and path" $ do + customAgentAsset InteractiveCodex InteractiveProjectScope "reviewer" + @?= AgentAssetLayout + { provider = InteractiveCodex, + scope = InteractiveProjectScope, + kind = CustomAgentAsset, + format = TomlFile, + path = ".codex/agents/reviewer.toml" + } + ] + +codexTomlTest :: TestTree +codexTomlTest = + testCase "Codex custom-agent TOML escapes strings and preserves instructions" $ do + codexCustomAgentToml + CodexCustomAgent + { name = "repo\"reviewer", + description = "Reviews\tchanges", + developerInstructions = "Read first.\nAvoid triple quotes: \"\"\"" + } + @?= Text.unlines + [ "name = \"repo\\\"reviewer\"", + "description = \"Reviews\\tchanges\"", + "developer_instructions = \"\"\"\nRead first.\nAvoid triple quotes: \\\"\\\"\\\"\n\"\"\"" + ] diff --git a/test/CatalogSpec.hs b/test/CatalogSpec.hs new file mode 100644 --- /dev/null +++ b/test/CatalogSpec.hs @@ -0,0 +1,43 @@ +-- | Regression test that locks down the contract between the JSON +-- catalog files under @baikai\/data\/models\/@ and the auto-generated +-- @baikai\/src\/Baikai\/Models\/Generated.hs@ module. +-- +-- The test invokes the @baikai-gen-models@ executable (made available +-- on @PATH@ by the @build-tool-depends@ entry in @baikai.cabal@) with +-- @--out@ pointing at a fresh temp file, then asserts the regenerated +-- output is byte-identical to the committed module. If the test fails +-- the remediation is always: +-- +-- @ +-- cabal run baikai-gen-models +-- git add baikai\/src\/Baikai\/Models\/Generated.hs +-- @ +-- +-- A failure means either someone edited @Baikai.Models.Generated@ by +-- hand (the file says \"do not edit\" at the top) or a @data\/models@ +-- JSON file changed without a paired regeneration. +module CatalogSpec (tests) where + +import Data.ByteString qualified as BS +import System.IO.Temp (withSystemTempDirectory) +import System.Process (callProcess) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertEqual, testCase) + +tests :: TestTree +tests = + testGroup + "Baikai.Models.Generated" + [ testCase "regenerating from data/models produces no diff" $ + withSystemTempDirectory "baikai-catalog-spec" $ \tmpDir -> do + let regenPath = tmpDir <> "/Generated.hs" + committedPath = "src/Baikai/Models/Generated.hs" + callProcess "baikai-gen-models" ["--out", regenPath] + committed <- BS.readFile committedPath + regenerated <- BS.readFile regenPath + assertEqual + "Generated.hs is out of sync with data/models/*.json.\n\ + \Run `cabal run baikai-gen-models` and commit the result." + committed + regenerated + ] diff --git a/test/CostSpec.hs b/test/CostSpec.hs new file mode 100644 --- /dev/null +++ b/test/CostSpec.hs @@ -0,0 +1,225 @@ +module CostSpec (tests) where + +import Baikai.Api (Api (..)) +import Baikai.Content (AssistantContent (..), TextContent (..)) +import Baikai.Context (Context (..), _Context) +import Baikai.Cost qualified as Cost +import Baikai.Cost.Log + ( CallLogConfig (..), + CallLogEntry, + runRequestWithLog, + withCallLog, + ) +import Baikai.Cost.Pricing (attachCost, computeCost) +import Baikai.Message (AssistantPayload (..), user) +import Baikai.Model (Model (..), ModelCost (..), _Model) +import Baikai.Options (Options, _Options) +import Baikai.Prelude +import Baikai.Provider + ( ApiProvider (..), + registerApiProvider, + ) +import Baikai.Response (Response (..), flattenAssistantBlocks) +import Baikai.StopReason (StopReason (..)) +import Baikai.Stream (liftCompleteToStream) +import Baikai.Usage (Usage, _Usage) +import Data.Aeson qualified as Aeson +import Data.ByteString.Lazy.Char8 qualified as BSL +import Data.List.NonEmpty (NonEmpty ((:|)), nonEmpty) +import Data.Maybe (fromJust, isJust) +import Data.Vector qualified as V +import System.Directory (getTemporaryDirectory, removeFile) +import System.FilePath (()) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Baikai.Cost" + [ computeTests, + attachCostTests, + callLogTests + ] + +-- Sample: 1000 input + 500 output tokens against a claude-haiku +-- pricing record (input $1/M, output $5/M) is exactly +-- 1000*(1/1_000_000) + 500*(5/1_000_000) = 7/2000 USD. +sampleUsage :: Usage +sampleUsage = + _Usage + & #inputTokens + .~ 1000 + & #outputTokens + .~ 500 + +-- Build a known-pricing model that matches the prior +-- claude-haiku-4-5-20251001 rates: input $1/M, output $5/M, +-- cache-read $0.10/M, cache-write $1.25/M. +knownModel :: Model +knownModel = + _Model + & #modelId + .~ "claude-haiku-4-5-20251001" + & #api + .~ Custom "test" + & #provider + .~ "anthropic" + & #cost + .~ ModelCost + { inputCost = 1, + outputCost = 5, + cacheReadCost = 1 / 10, + cacheWriteCost = 5 / 4 + } + +unknownModel :: Model +unknownModel = + _Model + & #modelId + .~ "totally-fake-model" + & #api + .~ Custom "test" + +computeTests :: TestTree +computeTests = + testGroup + "computeCost" + [ testCase "deterministic cost for the known model" $ + Cost.usd (computeCost knownModel sampleUsage) + @?= 7 / 2000, + testCase "zero cost for unknown models" $ + Cost.usd (computeCost unknownModel sampleUsage) + @?= 0, + testCase "cacheReadTokens contribute when present" $ do + let u :: Usage + u = + _Usage + & #cacheReadTokens + .~ 1000 + Cost.usd (computeCost knownModel u) @?= 1 / 10000, + testCase "cacheWriteTokens contribute against the known model" $ do + let u :: Usage + u = + _Usage + & #cacheWriteTokens + .~ 1000 + Cost.usd (computeCost knownModel u) @?= 1 / 800 + ] + +attachCostTests :: TestTree +attachCostTests = + testGroup + "attachCost" + [ testCase "fills the cost field on known models" $ + attachedUsd knownModel @?= 7 / 2000, + testCase "leaves cost zero on unknown models" $ + attachedUsd unknownModel @?= 0, + testCase "leaves the response's content alone" $ do + let resp = attachCost knownModel (mkResp knownModel) + flattenAssistantBlocks resp + @?= V.singleton (AssistantText (TextContent "hi")) + ] + where + attachedUsd m = + let AssistantPayload {usage = u} = attachCost m (mkResp m) ^. #message + in Cost.usd (u ^. #cost) + mkResp m = + Response + { message = + AssistantPayload + { content = V.singleton (AssistantText (TextContent "hi")), + usage = sampleUsage, + stopReason = Stop, + errorMessage = Nothing, + timestamp = read "2026-05-14 00:00:00 UTC" + }, + model = m, + api = Custom "test", + provider = "claude-api", + responseId = Nothing, + latencyMs = 100 + } + +-- Register a handler under a private API tag that returns a canned +-- response. Used by the call-log tests below. +cannedApi :: Api +cannedApi = Custom "baikai-cost-canned" + +cannedHaiku :: Response +cannedHaiku = + let u = sampleUsage & #cost .~ computeCost knownModel sampleUsage + in Response + { message = + AssistantPayload + { content = V.singleton (AssistantText (TextContent "ok")), + usage = u, + stopReason = Stop, + errorMessage = Nothing, + timestamp = read "2026-05-14 00:00:00 UTC" + }, + model = knownModel & #api .~ cannedApi, + api = cannedApi, + provider = "canned", + responseId = Nothing, + latencyMs = 7 + } + +registerCanned :: Response -> IO () +registerCanned resp = + let handler _m _ctx _opts = pure resp + in registerApiProvider + ApiProvider + { apiTag = cannedApi, + stream = liftCompleteToStream handler, + complete = handler + } + +cannedModel :: Model +cannedModel = knownModel & #api .~ cannedApi + +ctxHello :: Context +ctxHello = _Context & #messages .~ V.fromList [user "Hello world"] + +optsZero :: Options +optsZero = _Options + +callLogTests :: TestTree +callLogTests = + testGroup + "CallLog" + [ testCase "disabled handle skips disk I/O" $ do + registerCanned cannedHaiku + let cfg = CallLogConfig {path = "/dev/null", enabled = False} + withCallLog cfg $ \h -> do + resp <- runRequestWithLog h cannedModel ctxHello optsZero + flattenAssistantBlocks resp + @?= V.singleton (AssistantText (TextContent "ok")), + testCase "enabled handle writes one JSONL record per call" $ do + registerCanned cannedHaiku + tmp <- getTemporaryDirectory + let path' = tmp "baikai-cost-test.jsonl" + writeFile path' "" + let cfg = CallLogConfig {path = path', enabled = True} + withCallLog cfg $ \h -> do + _ <- runRequestWithLog h cannedModel ctxHello optsZero + pure () + raw <- BSL.readFile path' + firstLine <- case nonEmpty (BSL.lines raw) of + Nothing -> fail "expected one JSONL line, got an empty file" + Just (l :| rest) -> do + length rest @?= 0 + pure l + let mEntry :: Maybe CallLogEntry + mEntry = Aeson.decode firstLine + isJust mEntry @?= True + let entry = fromJust mEntry + entry ^. #provider @?= "canned" + entry ^. #model @?= "claude-haiku-4-5-20251001" + entry ^. #inputTokens @?= Just 1000 + entry ^. #outputTokens @?= Just 500 + entry ^. #latencyMs @?= 7 + entry ^. #promptSummary @?= "Hello world" + isJust (entry ^. #usd) @?= True + removeFile path' + ] diff --git a/test/InteractiveSpec.hs b/test/InteractiveSpec.hs new file mode 100644 --- /dev/null +++ b/test/InteractiveSpec.hs @@ -0,0 +1,55 @@ +module InteractiveSpec (tests) where + +import Baikai.Interactive +import Baikai.Prelude +import System.Exit (ExitCode (..)) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Baikai.Interactive" + [ requestDefaultTest, + providerRenderingTest, + codexSafetyRenderingTest, + resultConstructorTest + ] + +requestDefaultTest :: TestTree +requestDefaultTest = + testCase "_InteractiveLaunchRequest keeps optional launch settings empty" $ do + let req = _InteractiveLaunchRequest "start here" + req ^. #systemPrompt @?= Nothing + req ^. #userPrompt @?= "start here" + req ^. #model @?= Nothing + req ^. #workingDir @?= Nothing + req ^. #extraDirs @?= [] + req ^. #safety @?= DefaultSafety + req ^. #extraArgs @?= [] + +providerRenderingTest :: TestTree +providerRenderingTest = + testCase "provider and scope values render to stable names" $ do + renderInteractiveProvider InteractiveClaude @?= "claude" + renderInteractiveProvider InteractiveCodex @?= "codex" + renderInteractiveScope InteractiveUserScope @?= "user" + renderInteractiveScope InteractiveProjectScope @?= "project" + +codexSafetyRenderingTest :: TestTree +codexSafetyRenderingTest = + testCase "codex sandbox and approval values render to CLI-ready names" $ do + renderCodexSandboxMode CodexReadOnly @?= "read-only" + renderCodexSandboxMode CodexWorkspaceWrite @?= "workspace-write" + renderCodexSandboxMode CodexDangerFullAccess @?= "danger-full-access" + renderCodexApprovalPolicy CodexApprovalUntrusted @?= "untrusted" + renderCodexApprovalPolicy CodexApprovalOnFailure @?= "on-failure" + renderCodexApprovalPolicy CodexApprovalOnRequest @?= "on-request" + renderCodexApprovalPolicy CodexApprovalNever @?= "never" + +resultConstructorTest :: TestTree +resultConstructorTest = + testCase "_InteractiveLaunchResult records provider identity and process status" $ do + let result = _InteractiveLaunchResult InteractiveCodex ExitSuccess + result ^. #provider @?= InteractiveCodex + result ^. #exitCode @?= ExitSuccess diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,260 @@ +module Main (main) where + +import AgentAssetsSpec qualified +import Baikai +import Baikai.Prelude +import CatalogSpec qualified +import CostSpec qualified +import Data.Aeson qualified as Aeson +import Data.ByteString.Char8 qualified as BS8 +import Data.ByteString.Lazy.Char8 qualified as LBS8 +import Data.Text qualified as Text +import Data.Vector qualified as V +import InteractiveSpec qualified +import Streamly.Data.Stream qualified as Stream +import Test.Tasty (TestTree, defaultMain, testGroup) +import Test.Tasty.HUnit (assertBool, testCase, (@?=)) +import TraceSpec qualified + +-- | Ground the test provider on a 'Custom' API tag so it does not +-- clash with the real Anthropic/OpenAI handlers if a future test +-- registers them in the same process. +testApi :: Api +testApi = Custom "baikai-test" + +testModel :: Model +testModel = + _Model + & #modelId + .~ "test-model" + & #name + .~ "Test Model" + & #api + .~ testApi + & #provider + .~ "test" + +-- | Install a handler that returns a fixed assistant message for +-- the 'testApi' tag. Idempotent: re-registering the same tag +-- overwrites. +registerTestHandler :: Text -> IO () +registerTestHandler canned = + registerApiProvider (testProvider "test" canned) + +testProvider :: Text -> Text -> ApiProvider +testProvider providerName canned = + let handler m _ctx _opts = + pure $ + _Response + & #message + .~ AssistantPayload + { content = V.singleton (AssistantText (TextContent canned)), + usage = _Usage, + stopReason = Stop, + errorMessage = Nothing, + timestamp = read "2026-06-05 00:00:00 UTC" + } + & #model + .~ m + & #api + .~ testApi + & #provider + .~ providerName + in ApiProvider + { apiTag = testApi, + stream = liftCompleteToStream handler, + complete = handler + } + +main :: IO () +main = do + registerTestHandler "hello from the test provider" + defaultMain $ + testGroup + "baikai" + [ tests, + AgentAssetsSpec.tests, + CatalogSpec.tests, + CostSpec.tests, + InteractiveSpec.tests, + TraceSpec.tests + ] + +tests :: TestTree +tests = + testGroup + "baikai EP-2" + [ testCase "_Context defaults are zero-y" $ do + _Context ^. #systemPrompt @?= Nothing + V.length (_Context ^. #messages) @?= 0, + testCase "_Options defaults are zero-y" $ do + _Options ^. #maxTokens @?= Nothing + _Options ^. #temperature @?= Nothing + _Options ^. #apiKey @?= Nothing, + testCase "Options Show redacts literal API keys" $ do + let secret = "sk-baikai-secret-never-print" + opts = _Options & #apiKey .~ Just (ApiKeyLiteral secret) + assertBool + "show opts must not contain the raw API key" + (not (secret `Text.isInfixOf` Text.pack (show opts))), + testCase "Options JSON redacts literal API keys" $ do + let secret = "sk-baikai-secret-never-print" + opts = _Options & #apiKey .~ Just (ApiKeyLiteral secret) + assertBool + "Aeson.encode opts must not contain the raw API key" + (not (secret `Text.isInfixOf` Text.pack (LBS8.unpack (Aeson.encode opts)))), + testCase "completeRequest dispatches through the registered handler" $ do + let ctx = _Context & #messages .~ V.fromList [user "ping"] + resp <- completeRequest testModel ctx _Options + flattenAssistantBlocks resp + @?= V.singleton (AssistantText (TextContent "hello from the test provider")) + (resp ^. #model) ^. #modelId @?= "test-model" + resp ^. #provider @?= "test", + testCase "explicit registries isolate providers for the same API" $ do + regA <- newProviderRegistry + regB <- newProviderRegistry + registerApiProviderWith regA (testProvider "provider-a" "hello from A") + registerApiProviderWith regB (testProvider "provider-b" "hello from B") + let ctx = _Context & #messages .~ V.fromList [user "ping"] + respA <- completeRequestWith regA testModel ctx _Options + respB <- completeRequestWith regB testModel ctx _Options + flattenAssistantBlocks respA + @?= V.singleton (AssistantText (TextContent "hello from A")) + flattenAssistantBlocks respB + @?= V.singleton (AssistantText (TextContent "hello from B")) + respA ^. #provider @?= "provider-a" + respB ^. #provider @?= "provider-b", + testCase "streamRequestWith dispatches through an explicit registry" $ do + reg <- newProviderRegistry + registerApiProviderWith reg (testProvider "stream-provider" "hello from stream") + let ctx = _Context & #messages .~ V.fromList [user "ping"] + events <- Stream.toList (streamRequestWith reg testModel ctx _Options) + resp <- Stream.fold (reassembleResponse testModel) (Stream.fromList events) + flattenAssistantBlocks resp + @?= V.singleton (AssistantText (TextContent "hello from stream")), + testCase "OpenAI compat auto-detection drives provider request policy" $ do + let deepseek = + _Model + & #api + .~ OpenAIChatCompletions + & #baseUrl + .~ "https://api.deepseek.com" + compat = openaiCompletionsCompatFor deepseek + compat ^. #thinkingFormat @?= ThinkingFormatDeepseek + compat ^. #maxTokensField @?= MaxTokensField + compat ^. #supportsStrictMode @?= False + compat ^. #supportsDeveloperRole @?= False, + testCase "explicit OpenAI compat overrides baseUrl auto-detection" $ do + let explicit = + defaultOpenAICompletionsCompat + { supportsStrictMode = False, + thinkingFormat = ThinkingFormatNone + } + model = + _Model + & #api + .~ OpenAIChatCompletions + & #baseUrl + .~ "https://api.openai.com" + & #compat + .~ CompatOpenAICompletions explicit + compat = openaiCompletionsCompatFor model + compat ^. #supportsStrictMode @?= False + compat ^. #thinkingFormat @?= ThinkingFormatNone, + testCase "Anthropic compat auto-detection drives cache request policy" $ do + let fireworks = + _Model + & #api + .~ AnthropicMessages + & #baseUrl + .~ "https://api.fireworks.ai/inference/v1" + compat = anthropicMessagesCompatFor fireworks + compat ^. #supportsCacheControlOnTools @?= False + compat ^. #sendSessionAffinityHeaders @?= True + compat ^. #supportsLongCacheRetention @?= False, + testCase "user smart constructor produces a UserMessage" $ do + let ts = read "2026-06-05 01:02:03 UTC" + case userAt ts "hello" of + UserMessage UserPayload {content = uc, timestamp = actualTs} -> do + uc @?= V.singleton (UserText (TextContent "hello")) + actualTs @?= ts + _ -> error "expected UserMessage", + testCase "assistant smart constructor produces an AssistantMessage" $ do + let ts = read "2026-06-05 01:02:03 UTC" + case assistantAt ts "world" of + AssistantMessage AssistantPayload {content = ac, stopReason = sr, timestamp = actualTs} -> do + ac @?= V.singleton (AssistantText (TextContent "world")) + sr @?= Stop + actualTs @?= ts + _ -> error "expected AssistantMessage", + testCase "effectful user constructor produces a UserMessage in IO" $ do + msg <- userNow "hello now" + case msg of + UserMessage UserPayload {content = uc} -> + uc @?= V.singleton (UserText (TextContent "hello now")) + _ -> error "expected UserMessage", + testCase "appendToolResult carries text, image, and error payloads" $ do + let textCall = + _ToolCall + { id_ = "call_text", + name = "text_tool", + arguments = Aeson.object [] + } + imageCall = + _ToolCall + { id_ = "call_image", + name = "image_tool", + arguments = Aeson.object [] + } + errorCall = + _ToolCall + { id_ = "call_error", + name = "error_tool", + arguments = Aeson.object [] + } + assistantTurn = + AssistantMessage + AssistantPayload + { content = + V.fromList + [ AssistantToolCall textCall, + AssistantToolCall imageCall, + AssistantToolCall errorCall + ], + usage = _Usage, + stopReason = ToolUse, + errorMessage = Nothing, + timestamp = read "2026-06-05 00:00:00 UTC" + } + resp = _Response & #message .~ assistantPayload + assistantPayload = case assistantTurn of + AssistantMessage p -> p + _ -> error "expected assistant fixture" + ctx0 = _Context & #messages .~ V.singleton (user "use tools") + image = ImageContent {imageData = BS8.pack "png-bytes", mimeType = "image/png"} + dispatcher tc = case tc ^. #name of + "text_tool" -> pure (toolResultText "text result") + "image_tool" -> pure (toolResultImage image) + "error_tool" -> pure (toolResultErrorText "tool failed") + other -> error ("unexpected tool: " <> Text.unpack other) + ctx1 <- appendToolResult ctx0 resp dispatcher + case V.toList (ctx1 ^. #messages) of + [_, assistantMsg, ToolResultMessage textPayload, ToolResultMessage imagePayload, ToolResultMessage errorPayload] -> do + assistantMsg @?= assistantTurn + case textPayload of + ToolResultPayload {toolCallId = callId, content = blocks, isError = err} -> do + callId @?= "call_text" + blocks @?= V.singleton (ToolResultText (TextContent "text result")) + err @?= False + case imagePayload of + ToolResultPayload {toolCallId = callId, content = blocks, isError = err} -> do + callId @?= "call_image" + blocks @?= V.singleton (ToolResultImage image) + err @?= False + case errorPayload of + ToolResultPayload {toolCallId = callId, content = blocks, isError = err} -> do + callId @?= "call_error" + blocks @?= V.singleton (ToolResultText (TextContent "tool failed")) + err @?= True + msgs -> error ("unexpected context messages: " <> show msgs) + ] diff --git a/test/TraceSpec.hs b/test/TraceSpec.hs new file mode 100644 --- /dev/null +++ b/test/TraceSpec.hs @@ -0,0 +1,161 @@ +module TraceSpec (tests) where + +import Baikai.Api (Api (..)) +import Baikai.Content (AssistantContent (..), TextContent (..)) +import Baikai.Context (Context (..), _Context) +import Baikai.Error (BaikaiError (..)) +import Baikai.Message (AssistantPayload (..), user) +import Baikai.Model (Model (..), _Model) +import Baikai.Options (Options, _Options) +import Baikai.Prelude +import Baikai.Provider (ApiProvider (..), registerApiProvider) +import Baikai.Response (Response (..)) +import Baikai.StopReason (StopReason (..)) +import Baikai.Stream (liftCompleteToStream) +import Baikai.Trace (withTrace) +import Baikai.Trace.Event (TraceEvent (..)) +import Baikai.Trace.Sink (TraceSink (..), silent) +import Baikai.Usage (_Usage) +import Control.Concurrent.STM (TVar, atomically, modifyTVar', newTVarIO, readTVarIO) +import Control.Exception (throwIO) +import Data.Text qualified as Text +import Data.Vector qualified as V +import Streamly.Data.Fold qualified as Fold +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +tests :: TestTree +tests = + testGroup + "Baikai.Trace" + [ silentTest, + memoryFinishTest, + memoryFailTest + ] + +-- | Each test uses its own private 'Api' tag so tasty's parallel +-- test scheduler cannot race the (process-global) registry between +-- tests. +stubModel :: Api -> Model +stubModel a = + _Model + & #modelId + .~ "stub-1" + & #api + .~ a + & #provider + .~ "stub.trace" + & #maxOutputTokens + .~ 16 + +stubContext :: Context +stubContext = _Context & #messages .~ V.fromList [user "hello"] + +stubOptions :: Options +stubOptions = _Options & #maxTokens .~ Just 16 + +stubResponse :: Api -> Response +stubResponse a = + Response + { message = + AssistantPayload + { content = V.singleton (AssistantText (TextContent "hi")), + usage = _Usage, + stopReason = Stop, + errorMessage = Nothing, + timestamp = read "2026-05-14 00:00:00 UTC" + }, + model = stubModel a, + api = a, + provider = "stub.trace", + responseId = Nothing, + latencyMs = 0 + } + +registerOk :: Api -> IO () +registerOk a = + let handler _m _ctx _opts = pure (stubResponse a) + in registerApiProvider + ApiProvider + { apiTag = a, + stream = liftCompleteToStream handler, + complete = handler + } + +registerFail :: Api -> BaikaiError -> IO () +registerFail a e = + let handler _m _ctx _opts = throwIO e + in registerApiProvider + ApiProvider + { apiTag = a, + stream = liftCompleteToStream handler, + complete = handler + } + +memorySink :: IO (TVar [TraceEvent], TraceSink) +memorySink = do + ref <- newTVarIO [] + let step () e = atomically (modifyTVar' ref (e :)) + sink = TraceSink (Fold.foldlM' step (pure ())) + pure (ref, sink) + +silentTest :: TestTree +silentTest = + testGroup + "silent sink" + [ testCase "returns the response on success" $ do + let a = Custom "baikai-trace-silent-ok" + registerOk a + _ <- withTrace silent (stubModel a) stubContext stubOptions + pure (), + testCase "encodes failure as ErrorReason in the response" $ do + let a = Custom "baikai-trace-silent-fail" + registerFail a (ProviderError "boom") + resp <- withTrace silent (stubModel a) stubContext stubOptions + let AssistantPayload {stopReason = sr, errorMessage = em} = resp ^. #message + sr @?= ErrorReason + assertBool + ("expected errorMessage to mention boom, got: " <> show em) + (maybe False ("boom" `Text.isInfixOf`) em) + ] + +memoryFinishTest :: TestTree +memoryFinishTest = + testCase "memory sink records CallStarted then CallFinished" $ do + let a = Custom "baikai-trace-memory-ok" + registerOk a + (ref, sink) <- memorySink + _ <- withTrace sink (stubModel a) stubContext stubOptions + rev <- readTVarIO ref + let events = reverse rev + length events @?= 2 + case events of + [s@CallStarted {}, f@CallFinished {}] -> do + (s ^. #eventId :: Text) @?= (f ^. #eventId :: Text) + (s ^. #provider :: Text) @?= "stub.trace" + (f ^. #provider :: Text) @?= "stub.trace" + (s ^. #model :: Text) @?= "stub-1" + (f ^. #model :: Text) @?= "stub-1" + _ -> assertFailure ("unexpected event sequence: " <> show events) + +memoryFailTest :: TestTree +memoryFailTest = + testCase "memory sink records CallStarted then CallFailed on stream error" $ do + let a = Custom "baikai-trace-memory-fail" + registerFail a (ProviderError "stub-failure") + (ref, sink) <- memorySink + resp <- withTrace sink (stubModel a) stubContext stubOptions + -- The producer-side failure surfaces as an ErrorReason on the + -- response (no throw) and as CallFailed on the trace sink. + let AssistantPayload {stopReason = sr} = resp ^. #message + sr @?= ErrorReason + rev <- readTVarIO ref + let events = reverse rev + length events @?= 2 + case events of + [s@CallStarted {}, f@CallFailed {errorMessage = msg}] -> do + (s ^. #eventId :: Text) @?= (f ^. #eventId :: Text) + assertBool + ("expected error to mention stub-failure, got: " <> show msg) + ("stub-failure" `Text.isInfixOf` msg) + _ -> assertFailure ("unexpected event sequence: " <> show events)