diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,75 @@ ## [Unreleased] +## [0.2.0.0] - 2026-10-06 + +### Changed + +- **Breaking:** conversation turns use ordered content parts. + `UserMessage` / `AssistantMessage` hold `[ContentPart]`; `UserTurn` is a + bidirectional pattern synonym for a single text part. Replace + `AssistantTurn text reasoning calls` with `assistantTurn` (canonical order) + or `AssistantMessage respContent` for authoritative replay. +- **Breaking:** `ChatResponse.respContent` is now `[ContentPart]` (authoritative). + `respText` / `respReasoning` remain convenience projections. +- **Breaking:** `ToolCall.tcProviderMeta` is `Maybe ProviderOpaque` (tagged + provider + optional model + opaque payload), not bare `Value`. +- **Breaking:** removed `ContentBlock`; use `ContentPart` / `PartBody`. +- **Breaking:** `Turn` / `ToolCall` JSON shapes changed (role/content based). + Persisted histories must be migrated; old and new shapes are not mixed. +- Claude adapter: thinking request mapping (`thinking.type=enabled` + + `budget_tokens` from catalog effort), ordered thinking/text/tool_use parse + and stream, signed-block replay around tool rounds. Temperature is omitted + when thinking is enabled (Anthropic incompatibility). +- Gemini adapter: catalog thinking maps to `thinkingConfig`; tool-call thought + signatures use `ProviderOpaque` and still replay only for a matching model. +- OpenAI / DeepSeek / Ollama adapters encode and parse ordered parts; DeepSeek + reasoning continues via `reasoning_content` on ordered `ThinkingPart`s. +- Foreign opaque thinking / tool metadata is stripped at encode time on + provider fallback (history unchanged). +- **Breaking:** `ModelConfig` gains `mcCapabilities` (`ModelCapabilities`). + Manual record construction must set it (use `defaultModelCapabilities`). +- Example catalog: `capabilities` on vision/thinking models; removed + `temperature` from `haiku_4_5` (Claude thinking temperature footgun). + +- **Breaking:** `Usage` adds `usageCacheReadTokens` / `usageCacheCreationTokens`. + `usageInputTokens` is total input including cache read/write tokens (providers + normalize wire semantics so totals are not double-counted). Prefer `mkUsage`. +- **Breaking:** `PricingInfo` adds optional `pricePerMillionCacheRead` / + `pricePerMillionCacheWrite`; absent rates fall back to the ordinary input rate. + Prefer `defaultPricingInfo` for catalogs without cache rates. +- Cost estimation: ordinary input × input rate + cache read × cache-read rate + + cache creation × cache-write rate + output × output rate. +- **Breaking:** `ContentPart` gains `partCacheHint :: Maybe CacheHint`. Manual + construction and pattern matches must account for the new field; prefer + `textPart` / `cacheEphemeral`. `UserTurn` still matches only a single + unannotated text part. + +### Added + +- Image input: `ImageSource`, `ImagePart`, `imageUrlPart`, `imageBase64Part` / + `mkImageBase64` with MIME and base64 validation. Encoded for Claude, Gemini, + and OpenAI Chat Completions (DeepSeek/Ollama reuse the OpenAI shape when + `capabilities.vision` is declared). +- Catalog `capabilities` object (`thinking`, `vision`, `promptCaching`; missing + flags default to `false`). Carried on `ModelCatalogItem` / `ModelConfig`. +- Fallback candidates are validated before I/O: images require `vision`; + enabled thinking requires `thinking`. Unsupported capability yields + `UnsupportedCapability` and continues the fallback chain. +- `ProviderOpaque`, `ThinkingContent`, `ContentPart`, `PartBody`, `textPart`, + `thinkingPart`, `toolCallPart`, `mkChatResponse`, `projectText`, + `projectReasoning`, `turnToolCalls`, `validateTurn`, `stripForeignOpaque`. +- Claude thinking fixture and unit tests for ordered replay / foreign opaque + omission; Gemini signature matching tests. +- Image request-shape tests (Claude/Gemini/OpenAI) and vision fallback tests. +- `mkUsage`, `defaultPricingInfo`, `usageOrdinaryInputTokens`; provider parsers + report cache read/creation counters (Claude, OpenAI, DeepSeek, Gemini). +- Optional Claude prompt-cache breakpoints: `CacheHint` / `CacheEphemeral` on + `ContentPart.partCacheHint`, helpers `withCacheHint` / `cacheEphemeral`. + Claude emits wire `cache_control: {type: ephemeral}` (default TTL) on marked + parts without reordering; other providers ignore hints and leave content + unchanged. Hints round-trip in conversation JSON and survive agent tool rounds. + ## [0.1.1.0] - 2026-08-23 ### Changed @@ -89,3 +158,4 @@ [0.1.0.2]: https://github.com/aische/llm-simple/compare/v0.1.0.1...v0.1.0.2 [0.1.0.1]: https://github.com/aische/llm-simple/compare/v0.1.0.0...v0.1.0.1 [0.1.0.0]: https://github.com/aische/llm-simple/releases/tag/v0.1.0.0 +[0.2.0.0]: https://github.com/aische/llm-simple/compare/v0.1.1.0...v0.2.0.0 diff --git a/Readme.md b/Readme.md --- a/Readme.md +++ b/Readme.md @@ -4,7 +4,7 @@ The API stays small on purpose: JSON model and provider catalogs, bundled filesystem tools, and a straight line from `loadModelOrThrow` to `generateText`. Under that surface you still get multi-provider gateways, fallbacks, streaming, structured output, and sandboxed workspace tools — enough to prototype and compare ideas without assembling the pieces yourself. -**Status:** early 0.1.x release. APIs may change. +**Status:** early 0.2.x release. APIs may change. ## Features @@ -61,7 +61,7 @@ | `modelConfigName` | Name used in code to look up this config | | `providerName` | Provider key defined in `providers.json` | | `modelName` | Provider-specific model identifier | -| `pricing` | `pricePerMillionInput` / `pricePerMillionOutput` for usage cost tracking | +| `pricing` | `pricePerMillionInput` / `pricePerMillionOutput` for usage cost tracking; optional `pricePerMillionCacheRead` / `pricePerMillionCacheWrite` (fall back to input rate) | | `maxTokens` | Max tokens per request | | `temperature` | Sampling temperature (optional) | | `requestTimeout` | Request timeout in ms (optional) | @@ -69,9 +69,50 @@ | `retryCount` | Number of retries on failure | | `jitterBackoff` | Backoff jitter in ms between retries | | `thinking` | Extended thinking effort level (optional, provider-dependent) | +| `capabilities` | Optional object: `thinking`, `vision`, `promptCaching` (default `false`) | A provider is only available if its API key is set (except Ollama, which is always available). +Capability flags are declarations used before each fallback candidate is +called: image parts require `vision`; enabled thinking requires `thinking`. +Unsupported candidates fail with `UnsupportedCapability` and the next +fallback is tried. Cache hints may be ignored by providers that do not +support prompt caching (they do not change prompt semantics). + +### Image input + +User messages may mix text and images: + +```haskell +import LLM (imageUrlPart, imageBase64Part, textPart, UserMessage) + +let Right png = imageBase64Part "image/png" "" + msg = UserMessage [imageUrlPart "https://example.com/photo.jpg", png, textPart "describe"] +``` + +`imageBase64Part` validates MIME type (`image/jpeg|png|gif|webp|heic|heif`) +and base64 payload locally. Use a model with `"capabilities": {"vision": true}`. + +### Prompt cache breakpoints + +Mark stable content with `cacheEphemeral` (or `withCacheHint CacheEphemeral`). +Claude serializes that as `cache_control: {"type":"ephemeral"}` on the content +block (provider default TTL). Other providers omit the marker and send the same +content. Hints persist in conversation JSON and through agent tool rounds. + +```haskell +import LLM (cacheEphemeral, textPart, UserMessage) + +let msg = + UserMessage + [ cacheEphemeral (textPart "large stable context…"), + textPart "question about the context" + ] +``` + +`UserTurn "text"` still matches only a single unannotated text part; use +`UserMessage` when parts carry cache hints. + ### Provider catalog Providers are defined in `providers.json` (bundled with the package as a @@ -127,7 +168,7 @@ import Control.Exception (SomeException, catch) import Heptapod (generate) import LLM.Agent (Agent (..), RuntimeArgs (..), generateText, noEventObserver) -import LLM.Core (Turn (UserTurn)) +import LLM.Core (pattern UserTurn) import LLM.Generate (ModelWithFallbacks (..), Hooks (..), llmHooks, noHooks) import LLM.Load (fsTools, loadModelsOrThrow) diff --git a/app/Main.hs b/app/Main.hs --- a/app/Main.hs +++ b/app/Main.hs @@ -12,7 +12,7 @@ noEventObserver, ) import LLM.Core - ( Turn (UserTurn), + ( pattern UserTurn, ) import LLM.Generate ( ModelWithFallbacks (..), diff --git a/llm-simple.cabal b/llm-simple.cabal --- a/llm-simple.cabal +++ b/llm-simple.cabal @@ -2,7 +2,7 @@ name: llm-simple -version: 0.1.1.0 +version: 0.2.0.0 synopsis: Multi-provider LLM library with agent tool loops and filesystem tools @@ -16,7 +16,7 @@ loops, structured output via Autodocodec, JSON model- and provider-catalog loading, and workspace-scoped filesystem tools. . - Status: early 0.1.x release — APIs may change. + Status: early 0.2.x release — APIs may change. license: BSD-3-Clause @@ -65,6 +65,7 @@ OverloadedRecordDot DuplicateRecordFields NoFieldSelectors + PatternSynonyms StrictData GADTs TypeFamilies @@ -245,5 +246,7 @@ heptapod >=1.1 && <1.2, llm-simple, mtl >=2.3 && <2.4, - retry >=0.9 && <0.10 + retry >=0.9 && <0.10, + scientific >=0.3 && <0.4, + vector >=0.13 && <0.14 build-tool-depends: hspec-discover:hspec-discover diff --git a/model-catalog.json b/model-catalog.json --- a/model-catalog.json +++ b/model-catalog.json @@ -12,7 +12,10 @@ "requestTimeout": 10000, "throttleDelay": 1000, "retryCount": 3, - "jitterBackoff": 1000 + "jitterBackoff": 1000, + "capabilities": { + "vision": true + } }, { "modelConfigName": "gpt_5_6_terra", @@ -27,7 +30,10 @@ "requestTimeout": 10000, "throttleDelay": 1000, "retryCount": 3, - "jitterBackoff": 1000 + "jitterBackoff": 1000, + "capabilities": { + "vision": true + } }, { "modelConfigName": "llama_3_2", @@ -72,7 +78,11 @@ "requestTimeout": 10000, "throttleDelay": 1000, "retryCount": 3, - "jitterBackoff": 1000 + "jitterBackoff": 1000, + "capabilities": { + "thinking": true, + "vision": true + } }, { "modelConfigName": "haiku_4_5", @@ -83,11 +93,15 @@ "pricePerMillionOutput": 5 }, "maxTokens": 4096, - "temperature": 0.5, "requestTimeout": 30000, "throttleDelay": 5000, "retryCount": 3, - "jitterBackoff": 5000 + "jitterBackoff": 5000, + "capabilities": { + "thinking": true, + "vision": true, + "promptCaching": true + } }, { "modelConfigName": "deepseek4flash", @@ -102,6 +116,9 @@ "requestTimeout": 20000, "throttleDelay": 3000, "retryCount": 3, - "jitterBackoff": 1000 + "jitterBackoff": 1000, + "capabilities": { + "thinking": true + } } ] diff --git a/src/LLM.hs b/src/LLM.hs --- a/src/LLM.hs +++ b/src/LLM.hs @@ -28,7 +28,7 @@ -- agContextWindow = Nothing } -- rt <- mkRuntime -- your RuntimeArgs -- result <- generateText agent (ModelWithFallbacks model []) tools rt --- [UserTurn \"hello\"] +-- [UserTurn \"hello\"] -- pattern synonym -- ... -- @ -- @@ -51,19 +51,37 @@ module LLM ( -- * Core conversation types Turn (..), + pattern UserTurn, + assistantTurn, + ContentPart (..), + PartBody (..), + CacheHint (..), + ImageSource (..), + ThinkingContent (..), + ProviderOpaque (..), + textPart, + thinkingPart, + toolCallPart, + imageUrlPart, + imageBase64Part, + withCacheHint, + cacheEphemeral, ToolCall (..), ToolResult (..), ToolDef (..), TypedTool (..), ChatResponse (..), - ContentBlock (..), + mkChatResponse, LLMGateway (..), LLMError (..), Usage (..), PricingInfo (..), emptyUsage, + mkUsage, addUsage, estimateCost, + usageOrdinaryInputTokens, + defaultPricingInfo, ThinkingMode (..), LLMHooks (..), AbortSignal, @@ -74,6 +92,8 @@ streamTextWithFallbacks, genObject, genObjectUntyped, + ModelCapabilities (..), + defaultModelCapabilities, ModelConfig (..), ModelWithFallbacks (..), GenRequest (..), @@ -135,11 +155,16 @@ import LLM.Core ( AbortSignal, ChatResponse (..), - ContentBlock (..), + CacheHint (..), + ContentPart (..), + ImageSource (..), LLMError (..), LLMGateway (..), LLMHooks (..), + PartBody (..), PricingInfo (..), + ProviderOpaque (..), + ThinkingContent (..), ThinkingMode (..), ToolCall (..), ToolDef (..), @@ -148,9 +173,22 @@ TypedTool (..), Usage (..), addUsage, + assistantTurn, + cacheEphemeral, + defaultPricingInfo, emptyUsage, estimateCost, getToolCalls, + imageBase64Part, + imageUrlPart, + mkChatResponse, + mkUsage, + pattern UserTurn, + textPart, + thinkingPart, + toolCallPart, + usageOrdinaryInputTokens, + withCacheHint, ) import LLM.Generate ( GenRequest (..), @@ -160,12 +198,14 @@ GenerateTextResult (..), GeneratableObject, Hooks (..), + ModelCapabilities (..), ModelConfig (..), ModelWithFallbacks (..), StreamChunk (..), RoundTextRole (..), debugHooks, defaultDebugHooks, + defaultModelCapabilities, genObject, genObjectUntyped, generateTextWithFallbacks, diff --git a/src/LLM/Agent/Generate.hs b/src/LLM/Agent/Generate.hs --- a/src/LLM/Agent/Generate.hs +++ b/src/LLM/Agent/Generate.hs @@ -124,32 +124,32 @@ toolCalls = getToolCalls resp roundUsage = fromMaybe emptyUsage resp.respUsage newUsage = currentUsage <> roundUsage + -- Authoritative ordered parts from the provider response. + assistantMsg = AssistantMessage resp.respContent case toolCalls of [] -> do - let finalTurn = AssistantTurn txt resp.respReasoning [] - finalTurnsAcc = newTurnsAcc ++ [finalTurn] + let finalTurnsAcc = newTurnsAcc ++ [assistantMsg] successResult = GenerateTextResult rt.rtGenerationId finalTurnsAcc txt newUsage - emitEvent rt (MessageFinalized finalTurn) + emitEvent rt (MessageFinalized assistantMsg) emitEvent rt (GenerationFinished successResult) pure $ Right successResult _ -> do - let assistantTurn = AssistantTurn txt resp.respReasoning toolCalls - toolContext = createToolContext agent currentTurns newUsage rt + let toolContext = createToolContext agent currentTurns newUsage rt tools = getResolvedTools id agent toolMap rt - emitEvent rt (MessageCreated assistantTurn) + emitEvent rt (MessageCreated assistantMsg) emitEvent rt (ToolRoundStarted loopCount) toolResultsE <- executeToolsWithAbort rt.rtAbortSignal rt.rtHooks toolContext tools toolCalls case toolResultsE of Left err -> do - let errResult = GenerateErrorResult err (newTurnsAcc ++ [assistantTurn]) newUsage + let errResult = GenerateErrorResult err (newTurnsAcc ++ [assistantMsg]) newUsage emitEvent rt (GenerationFailed err errResult) pure $ Left errResult Right toolResults -> do let toolTurn = ToolTurn toolResults emitEvent rt (MessageCreated toolTurn) emitEvent rt (ToolRoundFinished loopCount) - let turnsToAdd = [assistantTurn, toolTurn] + let turnsToAdd = [assistantMsg, toolTurn] go (currentTurns ++ turnsToAdd) (newTurnsAcc ++ turnsToAdd) newUsage (loopCount + 1) diff --git a/src/LLM/Agent/ToolUtils.hs b/src/LLM/Agent/ToolUtils.hs --- a/src/LLM/Agent/ToolUtils.hs +++ b/src/LLM/Agent/ToolUtils.hs @@ -151,7 +151,7 @@ | idx < 0 = 0 | remaining <= 0 = idx + 1 | otherwise = case conv !! idx of - UserTurn _ -> go (idx - 1) (remaining - 1) + UserMessage _ -> go (idx - 1) (remaining - 1) _ -> go (idx - 1) remaining -- | Build a 'GenRequest' from agent configuration and runtime state. diff --git a/src/LLM/Agent/Tools/HistoryTool.hs b/src/LLM/Agent/Tools/HistoryTool.hs --- a/src/LLM/Agent/Tools/HistoryTool.hs +++ b/src/LLM/Agent/Tools/HistoryTool.hs @@ -6,7 +6,15 @@ import Data.Text qualified as T import GHC.Generics (Generic) import LLM.Agent.Types (ToolContext (..)) -import LLM.Core.Types (ToolCall (tcName), ToolResult (trContent, trName), Turn (..), TypedTool (..)) +import LLM.Core.Types + ( ToolCall (tcName), + ToolResult (trContent, trName), + Turn (..), + TypedTool (..), + projectReasoning, + projectText, + turnToolCalls, + ) newtype HistoryToolArgs = HistoryToolArgs { _historyChunk :: Int @@ -53,7 +61,7 @@ countUserTurns = length . filter isUserTurn isUserTurn :: Turn -> Bool -isUserTurn (UserTurn _) = True +isUserTurn (UserMessage _) = True isUserTurn _ = False -- | Split a conversation into pages of @n@ user messages each, working @@ -88,7 +96,7 @@ | idx < 0 = 0 | remaining <= 0 = idx + 1 | otherwise = case conv !! idx of - UserTurn _ -> go (idx - 1) (remaining - 1) + UserMessage _ -> go (idx - 1) (remaining - 1) _ -> go (idx - 1) remaining -- | Extract a slice [start, end) from a list. @@ -100,14 +108,18 @@ formatChunk = T.intercalate "\n" . map formatTurn formatTurn :: Turn -> Text -formatTurn (UserTurn t) = "[User] " <> t -formatTurn (AssistantTurn t mReasoning calls) = - "[Assistant] " - <> t - <> maybe "" (\r -> " [reasoning: " <> T.take 200 r <> "]") mReasoning - <> if null calls - then "" - else " [called: " <> T.intercalate ", " (map (\x -> x.tcName) calls) <> "]" +formatTurn (UserMessage parts) = + "[User] " <> projectText parts +formatTurn (AssistantMessage parts) = + let t = projectText parts + mReasoning = projectReasoning parts + calls = turnToolCalls parts + in "[Assistant] " + <> t + <> maybe "" (\r -> " [reasoning: " <> T.take 200 r <> "]") mReasoning + <> if null calls + then "" + else " [called: " <> T.intercalate ", " (map (\x -> x.tcName) calls) <> "]" formatTurn (ToolTurn results) = "[Tool results] " <> T.intercalate ", " [r.trName <> ": " <> T.take 200 r.trContent | r <- results] diff --git a/src/LLM/Core.hs b/src/LLM/Core.hs --- a/src/LLM/Core.hs +++ b/src/LLM/Core.hs @@ -3,17 +3,39 @@ ( LLMGateway (..), ChatRequest (..), ChatResponse (..), + mkChatResponse, LLMTextResult, LLMObjectResult, LLMResult, LLMError (..), Turn (..), + pattern UserTurn, + ContentPart (..), + PartBody (..), + CacheHint (..), + ImageSource (..), + ThinkingContent (..), + ProviderOpaque (..), + textPart, + thinkingPart, + toolCallPart, + imageUrlPart, + imageBase64Part, + withCacheHint, + cacheEphemeral, + mkImageBase64, + supportedImageMediaTypes, + projectText, + projectReasoning, + turnToolCalls, + conversationHasImages, + coalesceAdjacentTextParts, + validateTurn, ToolCall (..), ToolResult (..), LLMHooks (..), TypedTool (..), ToolDef (..), - ContentBlock (..), StreamEvent (..), ThinkingMode (..), MessageEncodeOptions (..), @@ -23,8 +45,11 @@ Usage (..), PricingInfo (..), emptyUsage, + mkUsage, addUsage, estimateCost, + usageOrdinaryInputTokens, + defaultPricingInfo, AbortSignal, newAbortSignal, abort, diff --git a/src/LLM/Core/Types.hs b/src/LLM/Core/Types.hs --- a/src/LLM/Core/Types.hs +++ b/src/LLM/Core/Types.hs @@ -1,21 +1,54 @@ +{-# LANGUAGE PatternSynonyms #-} + module LLM.Core.Types - ( Turn (..), + ( -- * Conversation turns + Turn (..), + pattern UserTurn, assistantTurn, - ContentBlock (..), + ContentPart (..), + PartBody (..), + CacheHint (..), + ImageSource (..), + ThinkingContent (..), + ProviderOpaque (..), + textPart, + thinkingPart, + toolCallPart, + imageUrlPart, + imageBase64Part, + withCacheHint, + cacheEphemeral, + mkImageBase64, + supportedImageMediaTypes, + projectText, + projectReasoning, + turnToolCalls, + conversationHasImages, + coalesceAdjacentTextParts, + validateTurn, + opaqueForProvider, + stripForeignOpaque, + + -- * Chat ChatRequest (..), ChatResponse (..), + mkChatResponse, LLMError (..), LLMTextResult, LLMObjectResult, LLMResult, + + -- * Tools ToolDef (..), ToolCall (..), mkToolCall, - LLMGateway (..), ToolResult (..), - StreamEvent (..), TypedTool (..), + + -- * Provider surface + LLMGateway (..), LLMHooks (..), + StreamEvent (..), ThinkingMode (..), MessageEncodeOptions (..), defaultMessageEncodeOptions, @@ -23,8 +56,22 @@ ) where -import Data.Aeson (FromJSON, ToJSON, Value) +import Data.Aeson + ( FromJSON (..), + ToJSON (..), + Value (..), + object, + withObject, + withText, + (.:), + (.:?), + (.=), + ) +import Data.Aeson.Types (Parser) +import Data.Char (isSpace, isAsciiUpper, isAsciiLower, isDigit) +import Data.Maybe (mapMaybe) import Data.Text (Text) +import Data.Text qualified as T import GHC.Generics (Generic) import LLM.Core.Usage (Usage) @@ -60,10 +107,10 @@ onLLMResponseError :: Text -> Text -> IO () } --- | DeepSeek thinking mode configuration. +-- | Thinking / reasoning mode configuration shared across providers. data ThinkingMode = ThinkingMode { tmEnabled :: Bool, - tmEffort :: Maybe Text -- e.g. @high@ or @max@ + tmEffort :: Maybe Text -- e.g. @high@, @max@, or a provider-specific token budget } deriving (Show, Eq) @@ -79,16 +126,362 @@ deepSeekMessageEncodeOptions :: MessageEncodeOptions deepSeekMessageEncodeOptions = MessageEncodeOptions {meoIncludeReasoning = True} --- | A single turn in a conversation +-- | Opaque, provider-owned payload that must round-trip for replay. +-- +-- Used for Claude thinking signatures, Gemini thought signatures, and any +-- similar provider-bound state. Consumers must treat 'poPayload' as opaque. +data ProviderOpaque = ProviderOpaque + { poProvider :: Text, + poModel :: Maybe Text, + poPayload :: Value + } + deriving (Show, Eq, Generic) + +instance ToJSON ProviderOpaque where + toJSON po = + object $ + [ "provider" .= po.poProvider, + "payload" .= po.poPayload + ] + ++ ["model" .= m | Just m <- [po.poModel]] + +instance FromJSON ProviderOpaque where + parseJSON = withObject "ProviderOpaque" $ \o -> + ProviderOpaque + <$> o .: "provider" + <*> o .:? "model" + <*> o .: "payload" + +-- | Displayable and/or opaque thinking content for one assistant part. +data ThinkingContent = ThinkingContent + { thinkingText :: Maybe Text, + thinkingOpaque :: Maybe ProviderOpaque + } + deriving (Show, Eq, Generic, FromJSON, ToJSON) + +-- | Image input by HTTPS URL or base64 payload. +data ImageSource + = ImageUrl Text + | ImageBase64 {imageMediaType :: Text, imageData :: Text} + deriving (Show, Eq, Generic) + +instance ToJSON ImageSource where + toJSON (ImageUrl url) = + object ["type" .= ("url" :: Text), "url" .= url] + toJSON (ImageBase64 mediaType data_) = + object + [ "type" .= ("base64" :: Text), + "media_type" .= mediaType, + "data" .= data_ + ] + +instance FromJSON ImageSource where + parseJSON = withObject "ImageSource" $ \o -> do + typ <- o .: "type" :: Parser Text + case typ of + "url" -> ImageUrl <$> o .: "url" + "base64" -> + ImageBase64 + <$> o .: "media_type" + <*> o .: "data" + _ -> fail $ "Unknown image source type: " <> T.unpack typ + +-- | MIME types accepted by 'mkImageBase64' / 'imageBase64Part'. +supportedImageMediaTypes :: [Text] +supportedImageMediaTypes = + [ "image/jpeg", + "image/png", + "image/gif", + "image/webp", + "image/heic", + "image/heif" + ] + +-- | Validate MIME type and base64 payload for an inline image. +mkImageBase64 :: Text -> Text -> Either Text ImageSource +mkImageBase64 mediaType rawData + | mediaType `notElem` supportedImageMediaTypes = + Left $ + "unsupported image media type: " + <> mediaType + <> "; expected one of: " + <> T.intercalate ", " supportedImageMediaTypes + | T.null cleaned = + Left "image base64 data must not be empty" + | not (isBase64Text cleaned) = + Left "image data is not valid base64" + | otherwise = + Right $ ImageBase64 mediaType cleaned + where + cleaned = T.filter (not . isSpace) rawData + +isBase64Text :: Text -> Bool +isBase64Text t = + let n = T.length t + in n > 0 + && n `mod` 4 == 0 + && T.all isBase64Char t + where + isBase64Char c = + isAsciiUpper c + || isAsciiLower c + || isDigit c + || c == '+' + || c == '/' + || c == '=' + +-- | Body of one ordered content part. +data PartBody + = TextPart Text + | ImagePart ImageSource + | ThinkingPart ThinkingContent + | ToolCallPart ToolCall + deriving (Show, Eq, Generic) + +instance ToJSON PartBody where + toJSON (TextPart t) = object ["type" .= ("text" :: Text), "text" .= t] + toJSON (ImagePart src) = object ["type" .= ("image" :: Text), "image" .= src] + toJSON (ThinkingPart tc) = + object $ + ["type" .= ("thinking" :: Text)] + ++ ["text" .= t | Just t <- [tc.thinkingText]] + ++ ["opaque" .= o | Just o <- [tc.thinkingOpaque]] + toJSON (ToolCallPart tc) = + object + [ "type" .= ("tool_call" :: Text), + "tool_call" .= tc + ] + +instance FromJSON PartBody where + parseJSON = withObject "PartBody" $ \o -> do + typ <- o .: "type" :: Parser Text + case typ of + "text" -> TextPart <$> o .: "text" + "image" -> ImagePart <$> o .: "image" + "thinking" -> do + mText <- o .:? "text" + mOpaque <- o .:? "opaque" + pure $ ThinkingPart (ThinkingContent mText mOpaque) + "tool_call" -> ToolCallPart <$> o .: "tool_call" + _ -> fail $ "Unknown part type: " <> T.unpack typ + +-- | Portable prompt-cache breakpoint intent. +-- +-- Unsupported providers ignore the hint without dropping the underlying +-- content. Claude serializes 'CacheEphemeral' as wire @cache_control@. +data CacheHint = CacheEphemeral + deriving (Show, Eq, Generic) + +instance ToJSON CacheHint where + toJSON CacheEphemeral = String "ephemeral" + +instance FromJSON CacheHint where + parseJSON = withText "CacheHint" $ \t -> + case t of + "ephemeral" -> pure CacheEphemeral + _ -> fail $ "Unknown cache hint: " <> T.unpack t + +-- | One ordered content part in a user or assistant message. +data ContentPart = ContentPart + { partBody :: PartBody, + -- | Optional cache breakpoint. Portable intent; see 'CacheHint'. + partCacheHint :: Maybe CacheHint + } + deriving (Show, Eq, Generic) + +instance ToJSON ContentPart where + toJSON (ContentPart body mHint) = + object $ + ("partBody" .= body) : ["partCacheHint" .= h | Just h <- [mHint]] + +instance FromJSON ContentPart where + parseJSON = withObject "ContentPart" $ \o -> + ContentPart + <$> o .: "partBody" + <*> o .:? "partCacheHint" + +-- | Content part with no cache hint. +mkPart :: PartBody -> ContentPart +mkPart body = ContentPart body Nothing + +textPart :: Text -> ContentPart +textPart t = mkPart (TextPart t) + +thinkingPart :: ThinkingContent -> ContentPart +thinkingPart tc = mkPart (ThinkingPart tc) + +toolCallPart :: ToolCall -> ContentPart +toolCallPart tc = mkPart (ToolCallPart tc) + +-- | User image part from a publicly reachable URL. +imageUrlPart :: Text -> ContentPart +imageUrlPart url = mkPart (ImagePart (ImageUrl url)) + +-- | User image part from base64 data; validates MIME type and payload. +imageBase64Part :: Text -> Text -> Either Text ContentPart +imageBase64Part mediaType data_ = + mkPart . ImagePart <$> mkImageBase64 mediaType data_ + +-- | Attach a cache hint to a content part. +withCacheHint :: CacheHint -> ContentPart -> ContentPart +withCacheHint hint cp = cp {partCacheHint = Just hint} + +-- | Mark a content part with Claude's default ephemeral cache breakpoint. +cacheEphemeral :: ContentPart -> ContentPart +cacheEphemeral = withCacheHint CacheEphemeral + +-- | A single turn in a conversation. +-- +-- Prefer 'UserTurn' / 'assistantTurn' for simple text construction. Match +-- 'UserMessage' / 'AssistantMessage' when inspecting arbitrary ordered parts. data Turn - = UserTurn Text - | AssistantTurn Text (Maybe Text) [ToolCall] -- content, reasoning_content, tool calls + = UserMessage [ContentPart] + | AssistantMessage [ContentPart] | ToolTurn [ToolResult] - deriving (Show, Eq, Generic, FromJSON, ToJSON) + deriving (Show, Eq, Generic) +instance ToJSON Turn where + toJSON (UserMessage parts) = + object ["role" .= ("user" :: Text), "content" .= parts] + toJSON (AssistantMessage parts) = + object ["role" .= ("assistant" :: Text), "content" .= parts] + toJSON (ToolTurn results) = + object ["role" .= ("tool" :: Text), "results" .= results] + +instance FromJSON Turn where + parseJSON = withObject "Turn" $ \o -> do + role <- o .: "role" :: Parser Text + case role of + "user" -> UserMessage <$> o .: "content" + "assistant" -> AssistantMessage <$> o .: "content" + "tool" -> ToolTurn <$> o .: "results" + _ -> fail $ "Unknown turn role: " <> T.unpack role + +-- | Bidirectional pattern for a single unannotated user text part. +-- +-- Matches only when there is no cache hint. Use 'UserMessage' for annotated +-- or multi-part content. +pattern UserTurn :: Text -> Turn +pattern UserTurn text = UserMessage [ContentPart (TextPart text) Nothing] + +{-# COMPLETE UserMessage, AssistantMessage, ToolTurn #-} + +-- | Migration helper: emit thinking (if any), then text, then tool calls. +-- +-- Does not recover arbitrary provider block order; use 'AssistantMessage' +-- with 'respContent' when replaying authoritative ordered parts. assistantTurn :: Text -> Maybe Text -> [ToolCall] -> Turn -assistantTurn = AssistantTurn +assistantTurn text mReasoning calls = + AssistantMessage $ + [ thinkingPart (ThinkingContent (Just r) Nothing) + | Just r <- [mReasoning], + not (T.null r) + ] + ++ [textPart text | not (T.null text)] + ++ map toolCallPart calls +-- | Concatenate text parts in order. +projectText :: [ContentPart] -> Text +projectText = T.concat . mapMaybe go + where + go (ContentPart (TextPart t) _) = Just t + go _ = Nothing + +-- | First non-empty thinking text, if any. +projectReasoning :: [ContentPart] -> Maybe Text +projectReasoning = go + where + go [] = Nothing + go (ContentPart (ThinkingPart tc) _ : rest) = + case tc.thinkingText of + Just t | not (T.null t) -> Just t + _ -> go rest + go (_ : rest) = go rest + +-- | Tool calls in part order. +turnToolCalls :: [ContentPart] -> [ToolCall] +turnToolCalls = mapMaybe go + where + go (ContentPart (ToolCallPart tc) _) = Just tc + go _ = Nothing + +-- | Merge runs of adjacent text parts (e.g. streamed token deltas). +-- +-- Only coalesces when both parts share the same cache hint, so breakpoints +-- are not lost. +coalesceAdjacentTextParts :: [ContentPart] -> [ContentPart] +coalesceAdjacentTextParts = go + where + go [] = [] + go (ContentPart (TextPart t1) h1 : ContentPart (TextPart t2) h2 : rest) + | h1 == h2 = + go (ContentPart (TextPart (t1 <> t2)) h1 : rest) + go (p : rest) = p : go rest + +-- | Whether any turn in the conversation contains an image part. +conversationHasImages :: [Turn] -> Bool +conversationHasImages = any turnHasImage + where + turnHasImage (UserMessage parts) = any isImagePart parts + turnHasImage (AssistantMessage parts) = any isImagePart parts + turnHasImage (ToolTurn _) = False + isImagePart (ContentPart (ImagePart _) _) = True + isImagePart _ = False + +-- | Validate role/part combinations for a turn. +-- +-- User messages may contain text and image parts; assistant messages may +-- contain text, thinking, and tool-call parts. Returns 'Left' with an error +-- message when the combination is invalid. Cache hints do not affect validity. +validateTurn :: Turn -> Either Text () +validateTurn (UserMessage parts) = + mapM_ userPart parts + where + userPart (ContentPart (TextPart _) _) = Right () + userPart (ContentPart (ImagePart _) _) = Right () + userPart (ContentPart (ThinkingPart _) _) = + Left "user messages may not contain thinking parts" + userPart (ContentPart (ToolCallPart _) _) = + Left "user messages may not contain tool-call parts" +validateTurn (AssistantMessage parts) = + mapM_ assistantPart parts + where + assistantPart (ContentPart (TextPart _) _) = Right () + assistantPart (ContentPart (ThinkingPart _) _) = Right () + assistantPart (ContentPart (ToolCallPart _) _) = Right () + assistantPart (ContentPart (ImagePart _) _) = + Left "assistant messages may not contain image parts" +validateTurn (ToolTurn _) = Right () + +-- | Keep opaque metadata only when it belongs to @provider@. +opaqueForProvider :: Text -> Maybe ProviderOpaque -> Maybe ProviderOpaque +opaqueForProvider provider (Just o) + | o.poProvider == provider = Just o +opaqueForProvider _ _ = Nothing + +-- | Drop foreign opaque state from thinking parts and tool-call metadata. +-- +-- Used by provider encoders during fallback so stored history is not mutated. +-- Cache hints are preserved. +stripForeignOpaque :: Text -> ContentPart -> ContentPart +stripForeignOpaque provider (ContentPart (ThinkingPart tc) hint) = + ContentPart + ( ThinkingPart + tc + { thinkingOpaque = opaqueForProvider provider tc.thinkingOpaque + } + ) + hint +stripForeignOpaque provider (ContentPart (ToolCallPart tc) hint) = + ContentPart + ( ToolCallPart + tc + { tcProviderMeta = opaqueForProvider provider tc.tcProviderMeta + } + ) + hint +stripForeignOpaque _ p = p + -- | A tool definition sent to the model data ToolDef = ToolDef { toolName :: Text, @@ -113,21 +506,34 @@ -- | A tool invocation returned by the model. -- --- 'tcProviderMeta' is an opaque, provider-owned JSON bag that must round-trip --- back to the provider on subsequent requests when this tool call is replayed --- as part of the conversation history. It exists because some providers --- (currently Gemini 2.5 thinking models, with their @thoughtSignature@) bind --- internal state to a specific tool-call part and reject the request if it --- isn't echoed back verbatim. Each provider is responsible for the shape of --- its own metadata; consumers should treat it as opaque. +-- 'tcProviderMeta' is opaque, provider-owned state that must round-trip back +-- to the same provider (and usually the same model) when this tool call is +-- replayed. Gemini thought signatures are the current example. data ToolCall = ToolCall { tcId :: Text, -- provider-specific call id tcName :: Text, tcArguments :: Value, - tcProviderMeta :: Maybe Value + tcProviderMeta :: Maybe ProviderOpaque } - deriving (Show, Eq, Generic, FromJSON, ToJSON) + deriving (Show, Eq, Generic) +instance ToJSON ToolCall where + toJSON tc = + object $ + [ "id" .= tc.tcId, + "name" .= tc.tcName, + "arguments" .= tc.tcArguments + ] + ++ ["provider_meta" .= m | Just m <- [tc.tcProviderMeta]] + +instance FromJSON ToolCall where + parseJSON = withObject "ToolCall" $ \o -> + ToolCall + <$> o .: "id" + <*> o .: "name" + <*> o .: "arguments" + <*> o .:? "provider_meta" + -- | Smart constructor for a 'ToolCall' with no provider metadata. Use this -- everywhere except when a provider parser is attaching its own metadata. mkToolCall :: Text -> Text -> Value -> ToolCall @@ -150,6 +556,8 @@ | EmptyResponse -- valid JSON, but no content in it | ToolLoopExceeded Int -- hit the max tool rounds limit | Aborted -- user cancelled the request + | -- | Model lacks a required catalog capability (e.g. vision) for this request. + UnsupportedCapability Text deriving (Show, Eq, Generic, ToJSON, FromJSON) -- | A request to an LLM provider @@ -164,20 +572,30 @@ } deriving (Show, Eq) --- | A content block in a response — either text or a tool call -data ContentBlock - = TextBlock Text - | ToolCallBlock ToolCall - deriving (Show, Eq) - --- | A response from an LLM provider +-- | A response from an LLM provider. +-- +-- 'respContent' is authoritative ordered content. 'respText' and +-- 'respReasoning' are convenience projections and may be lossy. data ChatResponse = ChatResponse { respText :: Text, - respContent :: [ContentBlock], + respContent :: [ContentPart], respUsage :: Maybe Usage, respReasoning :: Maybe Text } deriving (Show, Eq) + +-- | Build a 'ChatResponse' with text/reasoning projections derived from parts. +-- +-- Adjacent text parts are coalesced so streamed deltas become one part. +mkChatResponse :: [ContentPart] -> Maybe Usage -> ChatResponse +mkChatResponse parts usage = + let coalesced = coalesceAdjacentTextParts parts + in ChatResponse + { respText = projectText coalesced, + respContent = coalesced, + respUsage = usage, + respReasoning = projectReasoning coalesced + } -- | Events emitted during streaming data StreamEvent diff --git a/src/LLM/Core/Usage.hs b/src/LLM/Core/Usage.hs --- a/src/LLM/Core/Usage.hs +++ b/src/LLM/Core/Usage.hs @@ -2,22 +2,58 @@ ( Usage (..), PricingInfo (..), emptyUsage, + mkUsage, addUsage, estimateCost, + usageOrdinaryInputTokens, + defaultPricingInfo, ) where -import Data.Aeson (FromJSON, ToJSON) +import Data.Aeson (FromJSON (..), ToJSON (..), Value, object, withObject, (.!=), (.:), (.:?), (.=)) +import Data.Aeson.Types (Object, Parser) +import Data.Maybe (fromMaybe) import GHC.Generics (Generic) --- | Token usage from a single API call +-- | Token usage from a single API call. +-- +-- 'usageInputTokens' is the total input token count, including tokens that were +-- cache reads or cache writes. Provider parsers normalize differing wire +-- semantics so this total is never double-counted. +-- +-- Cost estimation treats +-- @usageInputTokens - usageCacheReadTokens - usageCacheCreationTokens@ as +-- ordinary (uncached) input. data Usage = Usage { usageInputTokens :: !Int, usageOutputTokens :: !Int, + usageCacheReadTokens :: !Int, + usageCacheCreationTokens :: !Int, usageTotalCost :: !Double } - deriving (Show, Eq, Generic, ToJSON, FromJSON) + deriving (Show, Eq, Generic) +instance ToJSON Usage where + toJSON :: Usage -> Value + toJSON u = + object + [ "usageInputTokens" .= u.usageInputTokens, + "usageOutputTokens" .= u.usageOutputTokens, + "usageCacheReadTokens" .= u.usageCacheReadTokens, + "usageCacheCreationTokens" .= u.usageCacheCreationTokens, + "usageTotalCost" .= u.usageTotalCost + ] + +instance FromJSON Usage where + parseJSON :: Value -> Parser Usage + parseJSON = withObject "Usage" $ \o -> + Usage + <$> o .: "usageInputTokens" + <*> o .: "usageOutputTokens" + <*> o .:? "usageCacheReadTokens" .!= 0 + <*> o .:? "usageCacheCreationTokens" .!= 0 + <*> o .:? "usageTotalCost" .!= 0 + instance Semigroup Usage where (<>) :: Usage -> Usage -> Usage (<>) = addUsage @@ -27,24 +63,90 @@ mempty = emptyUsage emptyUsage :: Usage -emptyUsage = Usage {usageInputTokens = 0, usageOutputTokens = 0, usageTotalCost = 0} +emptyUsage = + Usage + { usageInputTokens = 0, + usageOutputTokens = 0, + usageCacheReadTokens = 0, + usageCacheCreationTokens = 0, + usageTotalCost = 0 + } +-- | Construct usage with zero cache counters and zero cost. +mkUsage :: Int -> Int -> Usage +mkUsage input output = + Usage + { usageInputTokens = input, + usageOutputTokens = output, + usageCacheReadTokens = 0, + usageCacheCreationTokens = 0, + usageTotalCost = 0 + } + addUsage :: Usage -> Usage -> Usage addUsage a b = Usage { usageInputTokens = a.usageInputTokens + b.usageInputTokens, usageOutputTokens = a.usageOutputTokens + b.usageOutputTokens, + usageCacheReadTokens = a.usageCacheReadTokens + b.usageCacheReadTokens, + usageCacheCreationTokens = a.usageCacheCreationTokens + b.usageCacheCreationTokens, usageTotalCost = a.usageTotalCost + b.usageTotalCost } --- | Pricing in dollars per million tokens +-- | Uncached input tokens used for ordinary input pricing. +usageOrdinaryInputTokens :: Usage -> Int +usageOrdinaryInputTokens u = + max 0 (u.usageInputTokens - u.usageCacheReadTokens - u.usageCacheCreationTokens) + +-- | Pricing in dollars per million tokens. +-- +-- Optional cache rates fall back to 'pricePerMillionInput' when absent. data PricingInfo = PricingInfo { pricePerMillionInput :: Double, - pricePerMillionOutput :: Double + pricePerMillionOutput :: Double, + pricePerMillionCacheRead :: Maybe Double, + pricePerMillionCacheWrite :: Maybe Double } - deriving (Eq, Ord, Show, Generic, FromJSON, ToJSON) + deriving (Eq, Ord, Show, Generic) +instance ToJSON PricingInfo where + toJSON :: PricingInfo -> Value + toJSON p = + object $ + [ "pricePerMillionInput" .= p.pricePerMillionInput, + "pricePerMillionOutput" .= p.pricePerMillionOutput + ] + ++ ["pricePerMillionCacheRead" .= r | Just r <- [p.pricePerMillionCacheRead]] + ++ ["pricePerMillionCacheWrite" .= w | Just w <- [p.pricePerMillionCacheWrite]] + +instance FromJSON PricingInfo where + parseJSON :: Value -> Parser PricingInfo + parseJSON = withObject "PricingInfo" parsePricingInfo + +parsePricingInfo :: Object -> Parser PricingInfo +parsePricingInfo o = + PricingInfo + <$> o .: "pricePerMillionInput" + <*> o .: "pricePerMillionOutput" + <*> o .:? "pricePerMillionCacheRead" + <*> o .:? "pricePerMillionCacheWrite" + +-- | Pricing with only ordinary input/output rates (no cache-specific rates). +defaultPricingInfo :: Double -> Double -> PricingInfo +defaultPricingInfo input output = + PricingInfo + { pricePerMillionInput = input, + pricePerMillionOutput = output, + pricePerMillionCacheRead = Nothing, + pricePerMillionCacheWrite = Nothing + } + estimateCost :: PricingInfo -> Usage -> Double estimateCost p u = - fromIntegral u.usageInputTokens * p.pricePerMillionInput / 1_000_000 - + fromIntegral u.usageOutputTokens * p.pricePerMillionOutput / 1_000_000 + let ordinary = usageOrdinaryInputTokens u + cacheReadRate = fromMaybe p.pricePerMillionInput p.pricePerMillionCacheRead + cacheWriteRate = fromMaybe p.pricePerMillionInput p.pricePerMillionCacheWrite + in fromIntegral ordinary * p.pricePerMillionInput / 1_000_000 + + fromIntegral u.usageCacheReadTokens * cacheReadRate / 1_000_000 + + fromIntegral u.usageCacheCreationTokens * cacheWriteRate / 1_000_000 + + fromIntegral u.usageOutputTokens * p.pricePerMillionOutput / 1_000_000 diff --git a/src/LLM/Core/Utils.hs b/src/LLM/Core/Utils.hs --- a/src/LLM/Core/Utils.hs +++ b/src/LLM/Core/Utils.hs @@ -21,11 +21,19 @@ import Data.Text qualified as T import LLM.Core.Types ( ChatResponse (..), - ContentBlock (..), + ContentPart (..), + ImageSource (..), LLMError (..), LLMResult, + PartBody (..), + ProviderOpaque (..), + ThinkingContent (..), ToolCall (..), ToolResult (..), + mkChatResponse, + projectReasoning, + thinkingPart, + turnToolCalls, ) import LLM.Core.Usage (Usage (..)) import System.Timeout (timeout) @@ -40,10 +48,7 @@ -- | Extract tool calls from a response getToolCalls :: ChatResponse -> [ToolCall] -getToolCalls r = concatMap go r.respContent - where - go (ToolCallBlock tc) = [tc] - go _ = [] +getToolCalls r = turnToolCalls r.respContent -- | Whether an error is worth retrying isRetryable :: LLMError -> Bool @@ -82,56 +87,120 @@ streamResponseJson r = object [ "text" .= r.respText, - "content" .= map blockToJson r.respContent, + "content" .= map partToJson r.respContent, "usage" .= fmap usageToJson r.respUsage, "reasoning" .= r.respReasoning ] where - blockToJson (TextBlock t) = object ["type" .= ("text" :: Text), "text" .= t] - blockToJson (ToolCallBlock tc) = + partToJson (ContentPart (TextPart t) hint) = object $ + ["type" .= ("text" :: Text), "text" .= t] + ++ cacheHintPair hint + partToJson (ContentPart (ImagePart src) hint) = + object $ + ["type" .= ("image" :: Text), "image" .= imageToJson src] + ++ cacheHintPair hint + partToJson (ContentPart (ThinkingPart tc) hint) = + object $ + ["type" .= ("thinking" :: Text)] + ++ ["text" .= t | Just t <- [tc.thinkingText]] + ++ ["opaque" .= opaqueToJson o | Just o <- [tc.thinkingOpaque]] + ++ cacheHintPair hint + partToJson (ContentPart (ToolCallPart tc) hint) = + object $ [ "type" .= ("tool_call" :: Text), "id" .= tc.tcId, "name" .= tc.tcName, "arguments" .= tc.tcArguments ] - ++ ["provider_meta" .= m | Just m <- [tc.tcProviderMeta]] + ++ ["provider_meta" .= opaqueToJson m | Just m <- [tc.tcProviderMeta]] + ++ cacheHintPair hint + cacheHintPair hint = + ["cache_hint" .= h | Just h <- [hint]] + imageToJson (ImageUrl url) = + object ["type" .= ("url" :: Text), "url" .= url] + imageToJson (ImageBase64 mediaType data_) = + object + [ "type" .= ("base64" :: Text), + "media_type" .= mediaType, + "data" .= data_ + ] + opaqueToJson o = + object $ + [ "provider" .= o.poProvider, + "payload" .= o.poPayload + ] + ++ ["model" .= m | Just m <- [o.poModel]] usageToJson u = object [ "input_tokens" .= u.usageInputTokens, - "output_tokens" .= u.usageOutputTokens + "output_tokens" .= u.usageOutputTokens, + "cache_read_tokens" .= u.usageCacheReadTokens, + "cache_creation_tokens" .= u.usageCacheCreationTokens ] parseChatResponse :: Value -> Parser ChatResponse parseChatResponse = AE.withObject "ChatResponse" $ \v -> do - text <- v AE..: "text" - content <- v AE..: "content" >>= mapM parseContentBlock + content <- v AE..: "content" >>= mapM parseContentPart usage <- v AE..:? "usage" >>= mapM parseUsage - reasoning <- v AE..:? "reasoning" - pure - ChatResponse - { respText = text, - respContent = content, - respUsage = usage, - respReasoning = reasoning - } + -- Synthetic stream summaries store reasoning beside content blocks. + mReasoning <- v AE..:? "reasoning" + let contentWithReasoning = + case (projectReasoning content, mReasoning) of + (Nothing, Just rc) + | not (T.null rc) -> + thinkingPart (ThinkingContent (Just rc) Nothing) : content + _ -> content + pure $ mkChatResponse contentWithReasoning usage where - parseContentBlock = AE.withObject "ContentBlock" $ \o -> do + parseContentPart = AE.withObject "ContentPart" $ \o -> do t <- o AE..: "type" - case (t :: Text) of - "text" -> TextBlock <$> o AE..: "text" + mHint <- o AE..:? "cache_hint" + body <- case (t :: Text) of + "text" -> TextPart <$> o AE..: "text" + "image" -> ImagePart <$> (o AE..: "image" >>= parseImageSource) + "thinking" -> do + mText <- o AE..:? "text" + mOpaque <- o AE..:? "opaque" >>= mapM parseOpaque + pure $ ThinkingPart (ThinkingContent mText mOpaque) "tool_call" -> do tcId <- o AE..: "id" tcName <- o AE..: "name" tcArgs <- o AE..: "arguments" - tcMeta <- o AE..:? "provider_meta" - pure $ ToolCallBlock $ ToolCall tcId tcName tcArgs tcMeta - _ -> fail "Unknown content block type" + tcMeta <- o AE..:? "provider_meta" >>= mapM parseOpaque + pure $ ToolCallPart (ToolCall tcId tcName tcArgs tcMeta) + _ -> fail "Unknown content part type" + pure $ ContentPart body mHint + parseImageSource = AE.withObject "ImageSource" $ \o -> do + typ <- o AE..: "type" :: Parser Text + case typ of + "url" -> ImageUrl <$> o AE..: "url" + "base64" -> + ImageBase64 + <$> o AE..: "media_type" + <*> o AE..: "data" + _ -> fail "Unknown image source type" + + parseOpaque = AE.withObject "ProviderOpaque" $ \o -> + ProviderOpaque + <$> o AE..: "provider" + <*> o AE..:? "model" + <*> o AE..: "payload" + parseUsage = AE.withObject "Usage" $ \o -> do input <- o AE..: "input_tokens" output <- o AE..: "output_tokens" - pure $ Usage input output 0.0 + cacheRead <- fromMaybe 0 <$> o AE..:? "cache_read_tokens" + cacheCreate <- fromMaybe 0 <$> o AE..:? "cache_creation_tokens" + pure $ + Usage + { usageInputTokens = input, + usageOutputTokens = output, + usageCacheReadTokens = cacheRead, + usageCacheCreationTokens = cacheCreate, + usageTotalCost = 0.0 + } printValue :: Value -> IO () printValue val = L8.putStrLn (encode val) diff --git a/src/LLM/Generate.hs b/src/LLM/Generate.hs --- a/src/LLM/Generate.hs +++ b/src/LLM/Generate.hs @@ -13,6 +13,8 @@ streamTextWithFallbacks, genObject, genObjectUntyped, + ModelCapabilities (..), + defaultModelCapabilities, ModelConfig (..), ModelWithFallbacks (..), GenRequest (..), diff --git a/src/LLM/Generate/GenerateUtils.hs b/src/LLM/Generate/GenerateUtils.hs --- a/src/LLM/Generate/GenerateUtils.hs +++ b/src/LLM/Generate/GenerateUtils.hs @@ -4,6 +4,7 @@ usageWithModelCost, callWithRetryTimeout, withModelFallbacks, + validateModelCapabilities, llmHooks, ) where @@ -16,12 +17,16 @@ LLMError (..), LLMGateway (gwName), LLMHooks (..), + ThinkingMode (..), + Turn, + conversationHasImages, ) import LLM.Core.Usage (Usage (..), estimateCost) import LLM.Core.Utils (withRetry, withTimeout) import LLM.Generate.Logger (Hooks (..), LogLevel (..)) import LLM.Generate.ModelConfig - ( ModelConfig (..), + ( ModelCapabilities (..), + ModelConfig (..), ModelWithFallbacks (..), mfwToModelConfigs, modelRetryPolicy, @@ -47,6 +52,37 @@ usageWithModelCost :: ModelConfig -> Usage -> Usage usageWithModelCost mc u = u {usageTotalCost = estimateCost mc.mcPricing u} +-- | Check that a candidate model declares every capability the request needs. +-- +-- Unsupported vision is a candidate failure (do not drop images). Thinking +-- configuration requires a declared thinking capability. Cache hints are not +-- validated here; unsupported providers ignore them. +validateModelCapabilities :: ModelConfig -> [Turn] -> Either LLMError () +validateModelCapabilities mc turns = do + whenNeedsVision + whenNeedsThinking + where + caps = mc.mcCapabilities + whenNeedsVision + | conversationHasImages turns, + not caps.capVision = + Left $ + UnsupportedCapability $ + "model " + <> mc.mcModel + <> " does not support vision" + | otherwise = Right () + whenNeedsThinking = + case mc.mcThinking of + Just ThinkingMode {tmEnabled = True} + | not caps.capThinking -> + Left $ + UnsupportedCapability $ + "model " + <> mc.mcModel + <> " does not support thinking" + _ -> Right () + callWithRetryTimeout :: GenRequest -> ModelConfig -> @@ -73,16 +109,24 @@ maybe GErrAllModelsFailed GErrLLM mLast loop (mc : rest) _ = do gr.grHooks.onLog Info (formatTryingModel mc) - r <- invokePerModel mc - case r of - Left Aborted -> pure $ Left GErrAborted + case validateModelCapabilities mc gr.grMessages of Left err -> case rest of [] -> pure $ Left (GErrLLM err) _ -> do gr.grHooks.onLog Warn (formatModelFallback mc err) loop rest (Just err) - Right a -> pure $ Right a + Right () -> do + r <- invokePerModel mc + case r of + Left Aborted -> pure $ Left GErrAborted + Left err -> + case rest of + [] -> pure $ Left (GErrLLM err) + _ -> do + gr.grHooks.onLog Warn (formatModelFallback mc err) + loop rest (Just err) + Right a -> pure $ Right a formatTryingModel :: ModelConfig -> Text formatTryingModel mc = diff --git a/src/LLM/Generate/ModelConfig.hs b/src/LLM/Generate/ModelConfig.hs --- a/src/LLM/Generate/ModelConfig.hs +++ b/src/LLM/Generate/ModelConfig.hs @@ -1,5 +1,7 @@ module LLM.Generate.ModelConfig - ( ModelConfig (..), + ( ModelCapabilities (..), + defaultModelCapabilities, + ModelConfig (..), ModelWithFallbacks (..), mfwToModelConfigs, modelRetryPolicy, @@ -7,10 +9,46 @@ where import Control.Retry (RetryPolicyM, fullJitterBackoff, limitRetries) +import Data.Aeson (FromJSON (..), ToJSON (..), object, withObject, (.!=), (.:?), (.=)) import Data.Text (Text) +import GHC.Generics (Generic) import LLM.Core.Types (LLMGateway, ThinkingMode) import LLM.Core.Usage (PricingInfo) +-- | Declared model capabilities from the catalog. +-- +-- Missing flags decode as 'False'. These are declarations, not automatic +-- API discovery; request validation consults them before I/O. +data ModelCapabilities = ModelCapabilities + { capThinking :: Bool, + capVision :: Bool, + capPromptCaching :: Bool + } + deriving (Show, Eq, Ord, Generic) + +defaultModelCapabilities :: ModelCapabilities +defaultModelCapabilities = + ModelCapabilities + { capThinking = False, + capVision = False, + capPromptCaching = False + } + +instance ToJSON ModelCapabilities where + toJSON caps = + object + [ "thinking" .= caps.capThinking, + "vision" .= caps.capVision, + "promptCaching" .= caps.capPromptCaching + ] + +instance FromJSON ModelCapabilities where + parseJSON = withObject "capabilities" $ \o -> + ModelCapabilities + <$> o .:? "thinking" .!= False + <*> o .:? "vision" .!= False + <*> o .:? "promptCaching" .!= False + -- | Provider connection and per-model tuning parameters. -- -- Usually constructed via 'LLM.Load.loadModelOrThrow' from a JSON catalog entry, @@ -28,6 +66,8 @@ mcTemperature :: Maybe Double, -- | Thinking / reasoning mode (DeepSeek, etc.), when supported. mcThinking :: Maybe ThinkingMode, + -- | Declared capabilities (vision, thinking, prompt caching). + mcCapabilities :: ModelCapabilities, -- | Whole-request timeout in milliseconds ('Nothing' = no timeout). mcRequestTimeout :: Maybe Int, -- | Delay in milliseconds before each API call (rate limiting). diff --git a/src/LLM/Load.hs b/src/LLM/Load.hs --- a/src/LLM/Load.hs +++ b/src/LLM/Load.hs @@ -10,6 +10,8 @@ -- "pricePerMillionInput": 0.0, -- "pricePerMillionOutput": 0.0 -- }, +-- -- optional: pricePerMillionCacheRead / pricePerMillionCacheWrite +-- -- (fall back to pricePerMillionInput when absent) -- "maxTokens": 1024, -- "temperature": 0.5, -- "requestTimeout": 10000, diff --git a/src/LLM/Load/LoadModels.hs b/src/LLM/Load/LoadModels.hs --- a/src/LLM/Load/LoadModels.hs +++ b/src/LLM/Load/LoadModels.hs @@ -18,8 +18,12 @@ import Data.Map (Map) import Data.Map qualified as Map import Data.Text (Text) +import Data.Maybe (fromMaybe) import LLM.Core.Types (ThinkingMode (..)) -import LLM.Generate.ModelConfig (ModelConfig (..)) +import LLM.Generate.ModelConfig + ( ModelConfig (..), + defaultModelCapabilities, + ) import LLM.Load.LoadGateways (GatewayMap, loadGatewaysFromCatalogOrThrow) import LLM.Load.ModelCatalog (ModelCatalogItem (..), loadModelCatalog) import LLM.Load.ProviderCatalog (loadProviderCatalogForModelCatalog) @@ -45,6 +49,7 @@ mcMaxTokens = item.maxTokens, mcTemperature = item.temperature, mcThinking = fmap (\x -> ThinkingMode {tmEnabled = True, tmEffort = Just x}) item.thinking, + mcCapabilities = fromMaybe defaultModelCapabilities item.capabilities, mcRequestTimeout = item.requestTimeout, mcThrottleDelay = item.throttleDelay, mcRetryCount = item.retryCount, diff --git a/src/LLM/Load/ModelCatalog.hs b/src/LLM/Load/ModelCatalog.hs --- a/src/LLM/Load/ModelCatalog.hs +++ b/src/LLM/Load/ModelCatalog.hs @@ -8,6 +8,7 @@ import Data.Text qualified as T import GHC.Generics (Generic) import LLM.Core.Usage (PricingInfo) +import LLM.Generate.ModelConfig (ModelCapabilities) import LLM.Load.Types (LoadConfigError (..)) import LLM.Load.Utils (decodeJsonFile) @@ -20,6 +21,8 @@ temperature :: Maybe Double, requestTimeout :: Maybe Int, thinking :: Maybe Text, + -- | Optional capability flags; missing object or keys default to false. + capabilities :: Maybe ModelCapabilities, throttleDelay :: Maybe Int, retryCount :: Int, jitterBackoff :: Int diff --git a/src/LLM/Providers/Claude.hs b/src/LLM/Providers/Claude.hs --- a/src/LLM/Providers/Claude.hs +++ b/src/LLM/Providers/Claude.hs @@ -4,21 +4,32 @@ claudeProvider, claudeProviderWith, parseClaudeResponse, + parseClaudeStream, parseClaudeUsage, + claudeBuildBody, + encodeTurn, + effortToBudgetTokens, ) where import Data.Aeson ( KeyValue ((.=)), - Value (String), + Value (Object, String), decodeStrict', encode, object, + toJSON, withObject, + (.!=), (.:), + (.:?), ) -import Data.Aeson.Types (Parser, parseMaybe) +import Data.Aeson.KeyMap qualified as KM +import Data.Aeson.Types (Object, Pair, Parser, parseMaybe) +import Data.ByteString qualified as BS +import Data.Foldable (for_) import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef) +import Data.Maybe (fromMaybe, mapMaybe) import Data.Text (Text) import Data.Text qualified as T import Data.Text.Encoding (encodeUtf8) @@ -28,26 +39,38 @@ import LLM.Core.ProviderUtils (handleStreamResponse, lenientConfig, stripJsonFences) import LLM.Core.SSE (SSEEvent (sseData, sseEvent), readSSEEvents) import LLM.Core.Types - ( ChatRequest + ( CacheHint (..), + ChatRequest ( reqConversation, reqMaxTokens, reqModel, reqSystem, reqTemperature, + reqThinking, reqTools ), - ChatResponse (ChatResponse), - ContentBlock (..), + ContentPart (..), + ImageSource (..), LLMError (EmptyResponse), LLMGateway, LLMResult, LLMTextResult, + PartBody (..), + ProviderOpaque (..), StreamEvent (..), + ThinkingContent (..), + ThinkingMode (..), ToolCall (..), - mkToolCall, ToolDef (toolDescription, toolName, toolParameters), ToolResult (trCallId, trContent), Turn (..), + mkChatResponse, + mkToolCall, + stripForeignOpaque, + textPart, + thinkingPart, + toolCallPart, + pattern UserTurn, ) import LLM.Core.Usage (Usage (..), emptyUsage) import Network.HTTP.Client qualified as HC @@ -67,6 +90,9 @@ (/:), ) +claudeProviderName :: Text +claudeProviderName = "claude" + -- | Create a LLMGateway for the Claude provider at api.anthropic.com. claudeGateway :: Text -> LLMGateway claudeGateway apiKey = toGateway $ claudeProvider apiKey @@ -85,15 +111,16 @@ claudeProviderWith :: Url scheme -> Option scheme -> Text -> LLMProvider claudeProviderWith baseUrl baseOpts apiKey = LLMProvider - { providerName = "claude", + { providerName = claudeProviderName, buildBody = claudeBuildBody, sendRequest = sendRequest, sendStreamRequest = \body callback -> runReq lenientConfig $ do let url = baseUrl /: "v1" /: "messages" opts = baseOpts <> claudeAuthOpts apiKey + modelHint = fromMaybe "" (parseMaybe (withObject "body" (.: "model")) body) reqBr POST url (ReqBodyJson body) opts $ \resp -> - handleStreamResponse resp (`parseClaudeStream` callback), + handleStreamResponse resp (\br -> parseClaudeStream modelHint (HC.brRead br) callback), parseResponse = pure . parseClaudeResponse, buildObjectBody = \r schema -> let schemaText = TL.toStrict . decodeUtf8 $ encode schema @@ -111,87 +138,424 @@ resp <- req POST url (ReqBodyJson body) jsonResponse opts pure (responseStatusCode resp, responseBody resp) --- Internal helpers - claudeAuthOpts :: Text -> Option scheme claudeAuthOpts apiKey = header "x-api-key" (encodeUtf8 apiKey) <> header "anthropic-version" "2023-06-01" -parseClaudeStream :: HC.BodyReader -> (StreamEvent -> IO ()) -> IO LLMTextResult -parseClaudeStream reader callback = do - blocksRef <- newIORef ([] :: [ContentBlock]) +-- | Map catalog effort labels to Claude extended-thinking @budget_tokens@. +effortToBudgetTokens :: Text -> Int +effortToBudgetTokens = \case + "low" -> 1024 + "medium" -> 4096 + "high" -> 10000 + "max" -> 16000 + other -> + case reads (T.unpack other) of + [(n, "")] | n >= 1024 -> n + _ -> 10000 + +claudeBuildBody :: Bool -> ChatRequest -> Value +claudeBuildBody stream r = + object $ + [ "model" .= r.reqModel, + "max_tokens" .= r.reqMaxTokens, + "messages" .= concatMap (encodeTurn r.reqModel) r.reqConversation + ] + ++ ["system" .= sys | Just sys <- [r.reqSystem]] + ++ temperaturePairs r + ++ thinkingPairs r + ++ ["tools" .= map encodeToolDef r.reqTools | not (null r.reqTools)] + ++ ["stream" .= True | stream] + +-- | Anthropic rejects non-default temperature while thinking is active. +-- Omit temperature whenever thinking is enabled. +temperaturePairs :: ChatRequest -> [Pair] +temperaturePairs r = + case r.reqThinking of + Just tm | tm.tmEnabled -> [] + _ -> ["temperature" .= t | Just t <- [r.reqTemperature]] + +thinkingPairs :: ChatRequest -> [Pair] +thinkingPairs r = + case r.reqThinking of + Just tm + | tm.tmEnabled -> + let budget = maybe 10000 effortToBudgetTokens tm.tmEffort + in [ "thinking" + .= object + [ "type" .= ("enabled" :: Text), + "budget_tokens" .= budget + ] + ] + Just tm + | not tm.tmEnabled -> + ["thinking" .= object ["type" .= ("disabled" :: Text)]] + _ -> [] + +encodeTurn :: Text -> Turn -> [Value] +encodeTurn _ (UserMessage parts) = + [ object + [ "role" .= ("user" :: Text), + "content" .= encodeUserContent parts + ] + ] +encodeTurn currentModel (AssistantMessage parts) = + [ object + [ "role" .= ("assistant" :: Text), + "content" .= mapMaybe (encodeAssistantPart currentModel . stripForeignOpaque claudeProviderName) parts + ] + ] +encodeTurn _ (ToolTurn results) = + [ object + [ "role" .= ("user" :: Text), + "content" .= map encodeToolResult results + ] + ] + +encodeUserContent :: [ContentPart] -> Value +-- Bare string only when a single unannotated text part (fixture-compatible). +encodeUserContent [ContentPart (TextPart t) Nothing] = String t +encodeUserContent parts = toJSON (mapMaybe encodeUserPart parts) + where + encodeUserPart (ContentPart (TextPart t) hint) = + Just $ withCacheControl hint $ object ["type" .= ("text" :: Text), "text" .= t] + encodeUserPart (ContentPart (ImagePart src) hint) = + Just $ withCacheControl hint $ encodeImageBlock src + encodeUserPart _ = Nothing + +encodeImageBlock :: ImageSource -> Value +encodeImageBlock (ImageUrl url) = + object + [ "type" .= ("image" :: Text), + "source" + .= object + [ "type" .= ("url" :: Text), + "url" .= url + ] + ] +encodeImageBlock (ImageBase64 mediaType data_) = + object + [ "type" .= ("image" :: Text), + "source" + .= object + [ "type" .= ("base64" :: Text), + "media_type" .= mediaType, + "data" .= data_ + ] + ] + +-- | Attach Anthropic @cache_control@ for the default ephemeral TTL. +withCacheControl :: Maybe CacheHint -> Value -> Value +withCacheControl (Just CacheEphemeral) (Object o) = + Object $ KM.insert "cache_control" (object ["type" .= ("ephemeral" :: Text)]) o +withCacheControl _ v = v + +encodeAssistantPart :: Text -> ContentPart -> Maybe Value +encodeAssistantPart currentModel (ContentPart (ThinkingPart tc) hint) = + case opaqueForClaude currentModel tc.thinkingOpaque of + Just payload -> Just $ withCacheControl hint payload + Nothing -> + -- Without a Claude signature we must not invent a thinking block. + Nothing +encodeAssistantPart _ (ContentPart (TextPart t) hint) + | T.null t = Nothing + | otherwise = + Just $ + withCacheControl hint $ + object ["type" .= ("text" :: Text), "text" .= t] +encodeAssistantPart _ (ContentPart (ToolCallPart tc) hint) = + Just $ withCacheControl hint $ encodeToolUseBlock tc +encodeAssistantPart _ (ContentPart (ImagePart _) _) = Nothing + +opaqueForClaude :: Text -> Maybe ProviderOpaque -> Maybe Value +opaqueForClaude _ (Just o) + | o.poProvider == claudeProviderName = Just o.poPayload +opaqueForClaude _ _ = Nothing + +encodeToolDef :: ToolDef -> Value +encodeToolDef td = + object + [ "name" .= td.toolName, + "description" .= td.toolDescription, + "input_schema" .= td.toolParameters + ] + +encodeToolUseBlock :: ToolCall -> Value +encodeToolUseBlock tc = + object + [ "type" .= ("tool_use" :: Text), + "id" .= tc.tcId, + "name" .= tc.tcName, + "input" .= tc.tcArguments + ] + +encodeToolResult :: ToolResult -> Value +encodeToolResult tr = + object + [ "type" .= ("tool_result" :: Text), + "tool_use_id" .= tr.trCallId, + "content" .= tr.trContent + ] + +parseClaudeResponse :: Value -> LLMTextResult +parseClaudeResponse v = case parseMaybe (go modelVer) v of + Nothing -> Left EmptyResponse + Just parts -> + case parts of + [] -> Left EmptyResponse + _ -> Right (mkChatResponse parts (parseClaudeUsage v)) + where + modelVer = parseMaybe (withObject "ClaudeResponse" (.: "model")) v + + go :: Maybe Text -> Value -> Parser [ContentPart] + go mv = withObject "ClaudeResponse" $ \o -> do + content <- o .: "content" :: Parser [Value] + mapM (parseBlock mv) content + +parseBlock :: Maybe Text -> Value -> Parser ContentPart +parseBlock mv = withObject "content_block" $ \o -> do + typ <- o .: "type" :: Parser Text + case typ of + "text" -> textPart <$> o .: "text" + "thinking" -> do + thinkingTxt <- o .:? "thinking" :: Parser (Maybe Text) + signature <- o .:? "signature" :: Parser (Maybe Text) + let payload = + object $ + ["type" .= ("thinking" :: Text)] + ++ ["thinking" .= t | Just t <- [thinkingTxt]] + ++ ["signature" .= s | Just s <- [signature]] + opaque = + ProviderOpaque + { poProvider = claudeProviderName, + poModel = mv, + poPayload = payload + } + mText = case thinkingTxt of + Just t | not (T.null t) -> Just t + _ -> Nothing + pure $ thinkingPart (ThinkingContent mText (Just opaque)) + "redacted_thinking" -> do + data_ <- o .: "data" :: Parser Text + let payload = + object + [ "type" .= ("redacted_thinking" :: Text), + "data" .= data_ + ] + opaque = + ProviderOpaque + { poProvider = claudeProviderName, + poModel = mv, + poPayload = payload + } + pure $ thinkingPart (ThinkingContent Nothing (Just opaque)) + "tool_use" -> do + cid <- o .: "id" + name <- o .: "name" + args <- o .: "input" + pure $ toolCallPart (mkToolCall cid name args) + _ -> fail $ "Unknown content block type: " <> T.unpack typ + +parseClaudeUsage :: Value -> Maybe Usage +parseClaudeUsage = parseMaybe $ withObject "ClaudeResponse" $ \o -> do + u <- o .: "usage" + withObject "usage" parseClaudeUsageObject u + +parseClaudeUsageObject :: Object -> Parser Usage +parseClaudeUsageObject uo = do + -- Anthropic: input_tokens is uncached-only; total = uncached + read + creation. + uncached <- uo .: "input_tokens" + output <- uo .: "output_tokens" + cacheRead <- uo .:? "cache_read_input_tokens" .!= 0 + cacheCreate <- uo .:? "cache_creation_input_tokens" .!= 0 + pure + Usage + { usageInputTokens = uncached + cacheRead + cacheCreate, + usageOutputTokens = output, + usageCacheReadTokens = cacheRead, + usageCacheCreationTokens = cacheCreate, + usageTotalCost = 0 + } + +parseClaudeObjectResponse :: Value -> IO (LLMResult (Value, Maybe Usage)) +parseClaudeObjectResponse v = case parseMaybe go v of + Nothing -> pure $ Left EmptyResponse + Just text -> case decodeStrict' (encodeUtf8 (stripJsonFences text)) of + Nothing -> pure $ Left EmptyResponse + Just obj -> pure $ Right (obj, parseClaudeUsage v) + where + go :: Value -> Parser Text + go = withObject "ClaudeResponse" $ \o -> do + content <- o .: "content" :: Parser [Value] + texts <- mapMaybeM textOf content + case texts of + (t : _) -> pure t + _ -> fail "No content" + textOf = parseMaybe $ withObject "content_block" $ \o -> do + typ <- o .: "type" :: Parser Text + case typ of + "text" -> o .: "text" + _ -> fail "not text" + mapMaybeM f xs = pure (mapMaybe f xs) + +-- Streaming ------------------------------------------------------------------ + +data StreamBlock + = StreamText Text + | StreamThinking Text (Maybe Text) -- text, signature + | StreamRedacted Text -- data + | StreamTool Text Text Value -- id, name, args + +parseClaudeStream :: Text -> IO BS.ByteString -> (StreamEvent -> IO ()) -> IO LLMTextResult +parseClaudeStream modelHint readChunk callback = do + blocksRef <- newIORef ([] :: [StreamBlock]) usageRef <- newIORef emptyUsage - -- For accumulating tool_use input JSON across deltas - toolAccRef <- newIORef (Nothing :: Maybe (Text, Text, Text)) -- (id, name, json_so_far) - readSSEEvents (HC.brRead reader) $ \sse -> do + toolAccRef <- newIORef (Nothing :: Maybe (Text, Text, Text)) + thinkingAccRef <- newIORef (Nothing :: Maybe (Text, Maybe Text)) -- text, signature + redactedRef <- newIORef (Nothing :: Maybe Text) + modelRef <- newIORef modelHint + readSSEEvents readChunk $ \sse -> do case sse.sseEvent of Just "message_start" -> - -- Extract input token count from message.usage case decodeStrict' (encodeUtf8 sse.sseData) of - Just v -> case parseMaybe parseMessageStartUsage v of - Just inputToks -> modifyIORef' usageRef $ \u -> u {usageInputTokens = inputToks} - Nothing -> pure () + Just v -> do + case parseMaybe parseMessageStartUsage v of + Just startUsage -> modifyIORef' usageRef $ \u -> + u + { usageInputTokens = startUsage.usageInputTokens, + usageCacheReadTokens = startUsage.usageCacheReadTokens, + usageCacheCreationTokens = startUsage.usageCacheCreationTokens + } + Nothing -> pure () + for_ + (parseMaybe parseMessageStartModel v) + (writeIORef modelRef) Nothing -> pure () Just "content_block_start" -> case decodeStrict' (encodeUtf8 sse.sseData) of Just v -> case parseMaybe parseContentBlockStart v of - Just (cid, name) -> writeIORef toolAccRef (Just (cid, name, "")) - Nothing -> pure () -- text block start, nothing to do + Just (StartTool cid name) -> writeIORef toolAccRef (Just (cid, name, "")) + Just StartThinking -> writeIORef thinkingAccRef (Just ("", Nothing)) + Just (StartRedacted data_) -> writeIORef redactedRef (Just data_) + Just StartText -> pure () + Nothing -> pure () Nothing -> pure () Just "content_block_delta" -> case decodeStrict' (encodeUtf8 sse.sseData) of Just v -> do - -- Try text delta case parseMaybe parseTextDelta v of Just txt -> do - modifyIORef' blocksRef (TextBlock txt :) + modifyIORef' blocksRef (StreamText txt :) callback (StreamDelta txt) Nothing -> pure () - -- Try tool input delta + case parseMaybe parseThinkingDelta v of + Just txt -> do + modifyIORef' thinkingAccRef $ fmap (\(acc, sig) -> (acc <> txt, sig)) + callback (StreamReasoningDelta txt) + Nothing -> pure () + case parseMaybe parseSignatureDelta v of + Just sig -> + modifyIORef' thinkingAccRef $ fmap (\(acc, _) -> (acc, Just sig)) + Nothing -> pure () case parseMaybe parseInputJsonDelta v of Just fragment -> modifyIORef' toolAccRef $ fmap (\(cid, name, acc) -> (cid, name, acc <> fragment)) Nothing -> pure () Nothing -> pure () Just "content_block_stop" -> do + mThinking <- readIORef thinkingAccRef + case mThinking of + Just (txt, mSig) -> do + modifyIORef' blocksRef (StreamThinking txt mSig :) + writeIORef thinkingAccRef Nothing + Nothing -> pure () + mRedacted <- readIORef redactedRef + case mRedacted of + Just data_ -> do + modifyIORef' blocksRef (StreamRedacted data_ :) + writeIORef redactedRef Nothing + Nothing -> pure () mTool <- readIORef toolAccRef case mTool of - Just (cid, name, jsonStr) | not (T.null jsonStr) -> do + Just (cid, name, jsonStr) -> do let args = case decodeStrict' (encodeUtf8 jsonStr) of - Just v -> v + Just a -> a Nothing -> String jsonStr tc = mkToolCall cid name args - modifyIORef' blocksRef (ToolCallBlock tc :) + modifyIORef' blocksRef (StreamTool cid name args :) callback (StreamToolCall tc) writeIORef toolAccRef Nothing - _ -> writeIORef toolAccRef Nothing + Nothing -> pure () Just "message_delta" -> case decodeStrict' (encodeUtf8 sse.sseData) of Just v -> case parseMaybe parseMessageDeltaUsage v of Just outputToks -> modifyIORef' usageRef $ \u -> u {usageOutputTokens = outputToks} Nothing -> pure () Nothing -> pure () - _ -> pure () -- message_stop, ping, etc. - blocks <- reverse <$> readIORef blocksRef + _ -> pure () + rawBlocks <- reverse <$> readIORef blocksRef usage <- readIORef usageRef - let text = T.concat [t | TextBlock t <- blocks] - if null blocks + model <- readIORef modelRef + let parts = map (streamBlockToPart model) rawBlocks + if null parts then pure $ Left EmptyResponse - else pure $ Right (ChatResponse text blocks (Just usage) Nothing) + else pure $ Right (mkChatResponse parts (Just usage)) --- Parsers for streaming events -parseMessageStartUsage :: Value -> Parser Int +streamBlockToPart :: Text -> StreamBlock -> ContentPart +streamBlockToPart _ (StreamText t) = textPart t +streamBlockToPart model (StreamThinking txt mSig) = + let payload = + object $ + ["type" .= ("thinking" :: Text)] + ++ ["thinking" .= txt | not (T.null txt)] + ++ ["signature" .= s | Just s <- [mSig]] + opaque = + ProviderOpaque + { poProvider = claudeProviderName, + poModel = if T.null model then Nothing else Just model, + poPayload = payload + } + mText = if T.null txt then Nothing else Just txt + in thinkingPart (ThinkingContent mText (Just opaque)) +streamBlockToPart model (StreamRedacted data_) = + let payload = + object + [ "type" .= ("redacted_thinking" :: Text), + "data" .= data_ + ] + opaque = + ProviderOpaque + { poProvider = claudeProviderName, + poModel = if T.null model then Nothing else Just model, + poPayload = payload + } + in thinkingPart (ThinkingContent Nothing (Just opaque)) +streamBlockToPart _model (StreamTool cid name args) = + toolCallPart (mkToolCall cid name args) + +data BlockStart + = StartText + | StartThinking + | StartRedacted Text + | StartTool Text Text + +parseMessageStartUsage :: Value -> Parser Usage parseMessageStartUsage = withObject "message_start" $ \o -> do msg <- o .: "message" - withObject "message" (\mo -> do u <- mo .: "usage"; withObject "usage" (.: "input_tokens") u) msg + withObject "message" (\mo -> do u <- mo .: "usage"; withObject "usage" parseClaudeUsageObject u) msg +parseMessageStartModel :: Value -> Parser Text +parseMessageStartModel = withObject "message_start" $ \o -> do + msg <- o .: "message" + withObject "message" (.: "model") msg + parseMessageDeltaUsage :: Value -> Parser Int parseMessageDeltaUsage = withObject "message_delta" $ \o -> do u <- o .: "usage" withObject "usage" (.: "output_tokens") u -parseContentBlockStart :: Value -> Parser (Text, Text) +parseContentBlockStart :: Value -> Parser BlockStart parseContentBlockStart = withObject "content_block_start" $ \o -> do cb <- o .: "content_block" withObject @@ -199,8 +563,11 @@ ( \cbo -> do typ <- cbo .: "type" :: Parser Text case typ of - "tool_use" -> (,) <$> cbo .: "id" <*> cbo .: "name" - _ -> fail "not tool_use" + "tool_use" -> StartTool <$> cbo .: "id" <*> cbo .: "name" + "thinking" -> pure StartThinking + "redacted_thinking" -> StartRedacted <$> (cbo .:? "data" .!= "") + "text" -> pure StartText + _ -> fail "unknown block start" ) cb @@ -217,6 +584,32 @@ ) d +parseThinkingDelta :: Value -> Parser Text +parseThinkingDelta = withObject "delta_event" $ \o -> do + d <- o .: "delta" + withObject + "delta" + ( \d' -> do + typ <- d' .: "type" :: Parser Text + case typ of + "thinking_delta" -> d' .: "thinking" + _ -> fail "not thinking_delta" + ) + d + +parseSignatureDelta :: Value -> Parser Text +parseSignatureDelta = withObject "delta_event" $ \o -> do + d <- o .: "delta" + withObject + "delta" + ( \d' -> do + typ <- d' .: "type" :: Parser Text + case typ of + "signature_delta" -> d' .: "signature" + _ -> fail "not signature_delta" + ) + d + parseInputJsonDelta :: Value -> Parser Text parseInputJsonDelta = withObject "delta_event" $ \o -> do d <- o .: "delta" @@ -229,108 +622,3 @@ _ -> fail "not input_json_delta" ) d - -claudeBuildBody :: Bool -> ChatRequest -> Value -claudeBuildBody stream r = - object $ - [ "model" .= r.reqModel, - "max_tokens" .= r.reqMaxTokens, - "messages" .= concatMap encodeTurn r.reqConversation - ] - ++ ["system" .= sys | Just sys <- [r.reqSystem]] - ++ ["temperature" .= t | Just t <- [r.reqTemperature]] - ++ ["tools" .= map encodeToolDef r.reqTools | not (null r.reqTools)] - ++ ["stream" .= True | stream] - -encodeTurn :: Turn -> [Value] -encodeTurn (UserTurn content) = - [ object - [ "role" .= ("user" :: Text), - "content" .= content - ] - ] -encodeTurn (AssistantTurn text _mReasoning calls) = - [ object - [ "role" .= ("assistant" :: Text), - "content" .= (textBlocks ++ toolBlocks) - ] - ] - where - textBlocks = [object ["type" .= ("text" :: Text), "text" .= text] | not (T.null text)] - toolBlocks = map encodeToolUseBlock calls -encodeTurn (ToolTurn results) = - [ object - [ "role" .= ("user" :: Text), - "content" .= map encodeToolResult results - ] - ] - -encodeToolDef :: ToolDef -> Value -encodeToolDef td = - object - [ "name" .= td.toolName, - "description" .= td.toolDescription, - "input_schema" .= td.toolParameters - ] - -encodeToolUseBlock :: ToolCall -> Value -encodeToolUseBlock tc = - object - [ "type" .= ("tool_use" :: Text), - "id" .= tc.tcId, - "name" .= tc.tcName, - "input" .= tc.tcArguments - ] - -encodeToolResult :: ToolResult -> Value -encodeToolResult tr = - object - [ "type" .= ("tool_result" :: Text), - "tool_use_id" .= tr.trCallId, - "content" .= tr.trContent - ] - -parseClaudeResponse :: Value -> LLMTextResult -parseClaudeResponse v = case parseMaybe go v of - Nothing -> Left EmptyResponse - Just blocks -> case blocks of - [] -> Left EmptyResponse - _ -> - let text = T.concat [t | TextBlock t <- blocks] - in Right (ChatResponse text blocks (parseClaudeUsage v) Nothing) - where - go :: Value -> Parser [ContentBlock] - go = withObject "ClaudeResponse" $ \o -> do - content <- o .: "content" :: Parser [Value] - mapM parseBlock content - - parseBlock :: Value -> Parser ContentBlock - parseBlock = withObject "content_block" $ \o -> do - typ <- o .: "type" :: Parser Text - case typ of - "text" -> TextBlock <$> o .: "text" - "tool_use" -> do - cid <- o .: "id" - name <- o .: "name" - args <- o .: "input" - pure $ ToolCallBlock (mkToolCall cid name args) - _ -> fail $ "Unknown content block type: " <> T.unpack typ - -parseClaudeUsage :: Value -> Maybe Usage -parseClaudeUsage = parseMaybe $ withObject "ClaudeResponse" $ \o -> do - u <- o .: "usage" - withObject "usage" (\uo -> Usage <$> uo .: "input_tokens" <*> uo .: "output_tokens" <*> pure 0) u - -parseClaudeObjectResponse :: Value -> IO (LLMResult (Value, Maybe Usage)) -parseClaudeObjectResponse v = case parseMaybe go v of - Nothing -> pure $ Left EmptyResponse - Just text -> case decodeStrict' (encodeUtf8 (stripJsonFences text)) of - Nothing -> pure $ Left EmptyResponse - Just obj -> pure $ Right (obj, parseClaudeUsage v) - where - go :: Value -> Parser Text - go = withObject "ClaudeResponse" $ \o -> do - content <- o .: "content" :: Parser [Value] - case content of - (block : _) -> withObject "content_block" (.: "text") block - _ -> fail "No content" diff --git a/src/LLM/Providers/DeepSeek.hs b/src/LLM/Providers/DeepSeek.hs --- a/src/LLM/Providers/DeepSeek.hs +++ b/src/LLM/Providers/DeepSeek.hs @@ -35,8 +35,8 @@ LLMError (EmptyResponse), LLMGateway, ThinkingMode (..), - Turn (UserTurn), deepSeekMessageEncodeOptions, + pattern UserTurn, ) import LLM.Providers.OpenAI ( authHeader, diff --git a/src/LLM/Providers/Gemini.hs b/src/LLM/Providers/Gemini.hs --- a/src/LLM/Providers/Gemini.hs +++ b/src/LLM/Providers/Gemini.hs @@ -5,12 +5,15 @@ geminiProviderWith, parseGeminiResponse, parseGeminiUsage, + encodeTurn, + signatureForModel, ) where import Control.Applicative ((<|>)) import Data.Aeson ( KeyValue ((.=)), + Object, Value (Object, String), decodeStrict', object, @@ -22,7 +25,7 @@ import Data.Aeson.KeyMap qualified as KM import Data.Aeson.Types (Pair, Parser, parseMaybe) import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef) -import Data.Maybe (fromMaybe) +import Data.Maybe (fromMaybe, mapMaybe) import Data.Text (Text) import Data.Text qualified as T import Data.Text.Encoding (encodeUtf8) @@ -37,20 +40,30 @@ reqModel, reqSystem, reqTemperature, + reqThinking, reqTools ), - ChatResponse (ChatResponse), - ContentBlock (..), + ContentPart (..), LLMError (EmptyResponse), LLMGateway, LLMObjectResult, LLMTextResult, + PartBody (..), + ProviderOpaque (..), StreamEvent (..), + ThinkingContent (..), + ThinkingMode (..), ToolCall (..), ToolDef (toolDescription, toolName, toolParameters), ToolResult (trContent, trName), Turn (..), + ImageSource (..), + mkChatResponse, mkToolCall, + stripForeignOpaque, + textPart, + thinkingPart, + toolCallPart, ) import LLM.Core.Usage (Usage (..)) import Network.HTTP.Client qualified as HC @@ -71,6 +84,9 @@ (=:), ) +geminiProviderName :: Text +geminiProviderName = "gemini" + -- | Create a LLMGateway for the Gemini provider at generativelanguage.googleapis.com. geminiGateway :: Text -> LLMGateway geminiGateway apiKey = toGateway (geminiProvider apiKey) @@ -89,7 +105,7 @@ geminiProviderWith :: Url scheme -> Option scheme -> Text -> LLMProvider geminiProviderWith baseUrl baseOpts apiKey = LLMProvider - { providerName = "gemini", + { providerName = geminiProviderName, buildBody = const geminiBuildBody, sendRequest = sendRequest, sendStreamRequest = \body callback -> @@ -115,6 +131,7 @@ "responseSchema" .= schema ] ++ ["temperature" .= t | Just t <- [r.reqTemperature]] + ++ thinkingConfigPairs r ) ] ++ [ "system_instruction" .= object ["parts" .= [object ["text" .= sys]]] @@ -129,8 +146,6 @@ where sendRequest body = runReq lenientConfig $ do - -- For non-streaming we need the model name from the body to construct the URL. - -- We extract it from the request body JSON since the LLMProvider only passes Value. let model = extractModel body url = baseUrl @@ -144,85 +159,67 @@ geminiAuthOpts :: Text -> Option scheme geminiAuthOpts apiKey = header "x-goog-api-key" (encodeUtf8 apiKey) --- | Extract model name stashed in the request body by geminiBuildBody. extractModel :: Value -> Text extractModel v = fromMaybe "gemini-2.0-flash" (parseMaybe (withObject "body" (.: "_model")) v) --- | Remove the internal '_model' field before sending to the API. stripModel :: Value -> Value stripModel (Object o) = Object (KM.delete "_model" o) stripModel v = v parseGeminiStream :: HC.BodyReader -> (StreamEvent -> IO ()) -> IO LLMTextResult parseGeminiStream reader callback = do - blocksRef <- newIORef ([] :: [ContentBlock]) + partsRef <- newIORef ([] :: [ContentPart]) usageRef <- newIORef Nothing readSSEEvents (HC.brRead reader) $ \sse -> do case decodeStrict' (encodeUtf8 sse.sseData) of Nothing -> pure () Just v -> do - -- modelVersion is needed so we can later guard against replaying - -- thought signatures into a different model than the one that - -- produced them. let modelVer = parseMaybe parseModelVersion v case parseMaybe (parseChunkParts modelVer) v of Just parts -> do - newBlocks <- mapM (assignToolId callback) parts - modifyIORef' blocksRef (++ newBlocks) + newParts <- mapM (assignToolId callback) parts + modifyIORef' partsRef (++ newParts) Nothing -> pure () case parseMaybe parseUsageMetadata v of Just u -> writeIORef usageRef (Just u) Nothing -> pure () - blocks <- readIORef blocksRef + parts <- readIORef partsRef usage <- readIORef usageRef - let text = T.concat [t | TextBlock t <- blocks] - if null blocks + if null parts then pure $ Left EmptyResponse - else pure $ Right (ChatResponse text blocks usage Nothing) + else pure $ Right (mkChatResponse parts usage) where - assignToolId :: (StreamEvent -> IO ()) -> ContentBlock -> IO ContentBlock - assignToolId cb (TextBlock t) = do + assignToolId :: (StreamEvent -> IO ()) -> ContentPart -> IO ContentPart + assignToolId cb (ContentPart (TextPart t) hint) = do cb (StreamDelta t) - pure (TextBlock t) - assignToolId cb (ToolCallBlock tc) = do + pure (textPart t) {partCacheHint = hint} + assignToolId cb (ContentPart (ThinkingPart tc) hint) = do + case tc.thinkingText of + Just t | not (T.null t) -> cb (StreamReasoningDelta t) + _ -> pure () + pure (thinkingPart tc) {partCacheHint = hint} + assignToolId cb (ContentPart (ToolCallPart tc) hint) = do tc' <- normalizeToolCallId tc cb (StreamToolCall tc') - pure (ToolCallBlock tc') + pure (toolCallPart tc') {partCacheHint = hint} + assignToolId _ (ContentPart (ImagePart src) hint) = + pure (ContentPart (ImagePart src) hint) - parseChunkParts :: Maybe Text -> Value -> Parser [ContentBlock] + parseChunkParts :: Maybe Text -> Value -> Parser [ContentPart] parseChunkParts modelVer = withObject "GeminiChunk" $ \o -> do (cand : _) <- o .: "candidates" :: Parser [Value] withObject "candidate" ( \co -> do cont <- co .: "content" - withObject "content" (\cco -> cco .: "parts" >>= mapM (parsePartBlock modelVer)) cont + withObject "content" (\cco -> cco .: "parts" >>= mapM (parsePart modelVer)) cont ) cand - parsePartBlock :: Maybe Text -> Value -> Parser ContentBlock - parsePartBlock modelVer = withObject "part" $ \o -> do - mSig <- o .:? "thoughtSignature" :: Parser (Maybe Text) - let tryText = TextBlock <$> (o .: "text") - tryFunctionCall = do - fc <- o .: "functionCall" - withObject - "functionCall" - ( \fco -> do - name <- fco .: "name" - args <- fco .:? "args" .!= object [] - pure $ ToolCallBlock (attachGeminiMeta modelVer mSig (mkToolCall name name args)) - ) - fc - tryText <|> tryFunctionCall - parseUsageMetadata :: Value -> Parser Usage parseUsageMetadata = withObject "GeminiChunk" $ \o -> do u <- o .: "usageMetadata" - withObject - "usageMetadata" - (\uo -> Usage <$> uo .: "promptTokenCount" <*> uo .: "candidatesTokenCount" <*> pure 0) - u + withObject "usageMetadata" parseGeminiUsageObject u geminiBuildBody :: ChatRequest -> Value geminiBuildBody r = object $ geminiBuildBodyPairs r @@ -240,25 +237,28 @@ | not (null r.reqTools) ] --- | Encode a turn for Gemini. The current request model is threaded through --- so that thought signatures captured from a previous response can be --- replayed only when the receiving model matches the one that emitted them. +-- | Encode a turn for Gemini. Foreign opaque thinking / tool metadata is omitted +-- at encode time so fallbacks never replay Claude (or other) state. encodeTurn :: Text -> Turn -> [Value] -encodeTurn _ (UserTurn content) = +encodeTurn _ (UserMessage parts) = [ object [ "role" .= ("user" :: Text), - "parts" .= [object ["text" .= content]] + "parts" .= mapMaybe encodeUserPart parts ] ] -encodeTurn currentModel (AssistantTurn text _mReasoning calls) = + where + -- Cache hints are ignored on the Gemini wire; content is unchanged. + encodeUserPart (ContentPart (TextPart t) _) = Just $ object ["text" .= t] + encodeUserPart (ContentPart (ImagePart src) _) = Just $ encodeImagePart src + encodeUserPart _ = Nothing +encodeTurn currentModel (AssistantMessage parts) = [ object [ "role" .= ("model" :: Text), - "parts" .= (textParts ++ callParts) + "parts" .= mapMaybe (encodeAssistantPart currentModel) cleaned ] ] where - textParts = [object ["text" .= text] | not (T.null text)] - callParts = map (encodeFunctionCall currentModel) calls + cleaned = map (stripForeignOpaque geminiProviderName) parts encodeTurn _ (ToolTurn results) = [ object [ "role" .= ("user" :: Text), @@ -266,6 +266,61 @@ ] ] +encodeImagePart :: ImageSource -> Value +encodeImagePart (ImageUrl url) = + object + [ "fileData" + .= object + [ "mimeType" .= guessImageMimeFromUrl url, + "fileUri" .= url + ] + ] +encodeImagePart (ImageBase64 mediaType data_) = + object + [ "inlineData" + .= object + [ "mimeType" .= mediaType, + "data" .= data_ + ] + ] + +-- | Best-effort MIME guess for Gemini fileData URL parts. +guessImageMimeFromUrl :: Text -> Text +guessImageMimeFromUrl url = + let path = T.toLower $ T.takeWhile (/= '?') url + in case () of + _ + | ".png" `T.isSuffixOf` path -> "image/png" + | ".gif" `T.isSuffixOf` path -> "image/gif" + | ".webp" `T.isSuffixOf` path -> "image/webp" + | ".heic" `T.isSuffixOf` path -> "image/heic" + | ".heif" `T.isSuffixOf` path -> "image/heif" + | otherwise -> "image/jpeg" + +encodeAssistantPart :: Text -> ContentPart -> Maybe Value +encodeAssistantPart _ (ContentPart (TextPart t) _) + | T.null t = Nothing + | otherwise = Just $ object ["text" .= t] +encodeAssistantPart currentModel (ContentPart (ThinkingPart tc) _) = + case tc.thinkingOpaque of + Just o + | o.poProvider == geminiProviderName, + modelOk currentModel o.poModel -> + Just o.poPayload + _ -> + case tc.thinkingText of + Just t + | not (T.null t) -> + Just $ object ["text" .= t, "thought" .= True] + _ -> Nothing +encodeAssistantPart currentModel (ContentPart (ToolCallPart tc) _) = + Just $ encodeFunctionCall currentModel tc +encodeAssistantPart _ (ContentPart (ImagePart _) _) = Nothing + +modelOk :: Text -> Maybe Text -> Bool +modelOk _ Nothing = True +modelOk current (Just m) = modelsMatch current m + encodeToolDef :: ToolDef -> Value encodeToolDef td = object @@ -274,15 +329,6 @@ "parameters" .= td.toolParameters ] --- | Encode an assistant tool call back to a Gemini @functionCall@ part. --- --- Gemini 2.5 thinking models attach an opaque 'thoughtSignature' to each --- function-call part. That signature must be replayed verbatim on subsequent --- requests against the same model, or the API rejects the request with a --- @function call is missing a thought_signature@ error. We stored the --- signature plus the emitting model name in 'tcProviderMeta' at parse time; --- here we re-emit it only when 'currentModel' matches, because signatures --- are bound to the model that produced them. encodeFunctionCall :: Text -> ToolCall -> Value encodeFunctionCall currentModel tc = object $ @@ -304,37 +350,78 @@ ] ] --- | Generate a unique call ID for Gemini tool calls (which lack native IDs) normalizeToolCallId :: ToolCall -> IO ToolCall normalizeToolCallId tc = do u <- newUnique let callId = "call_" <> T.pack (show (hashUnique u)) pure tc {tcId = callId} -normalizeBlock :: ContentBlock -> IO ContentBlock -normalizeBlock (ToolCallBlock tc) = ToolCallBlock <$> normalizeToolCallId tc -normalizeBlock b = pure b +normalizePart :: ContentPart -> IO ContentPart +normalizePart (ContentPart (ToolCallPart tc) hint) = + (\tc' -> (toolCallPart tc') {partCacheHint = hint}) <$> normalizeToolCallId tc +normalizePart p = pure p genConfig :: ChatRequest -> Value genConfig r = object $ ("maxOutputTokens" .= r.reqMaxTokens) : ["temperature" .= t | Just t <- [r.reqTemperature]] + ++ thinkingConfigPairs r +-- | Map 'ThinkingMode' to Gemini @thinkingConfig@. +-- +-- Level strings (@low@/@medium@/@high@/@minimal@) become @thinkingLevel@ +-- (Gemini 3). Numeric effort becomes @thinkingBudget@ (Gemini 2.5). Enabled +-- without effort requests dynamic budget (@-1@). +thinkingConfigPairs :: ChatRequest -> [Pair] +thinkingConfigPairs r = + case r.reqThinking of + Nothing -> [] + Just tm + | not tm.tmEnabled -> + ["thinkingConfig" .= object ["thinkingBudget" .= (0 :: Int)]] + | Just e <- tm.tmEffort, + e `elem` ["minimal", "low", "medium", "high"] -> + [ "thinkingConfig" + .= object + [ "thinkingLevel" .= e, + "includeThoughts" .= True + ] + ] + | Just e <- tm.tmEffort, + Just n <- readMaybeInt e -> + [ "thinkingConfig" + .= object + [ "thinkingBudget" .= n, + "includeThoughts" .= True + ] + ] + | otherwise -> + [ "thinkingConfig" + .= object + [ "thinkingBudget" .= (-1 :: Int), + "includeThoughts" .= True + ] + ] + +readMaybeInt :: Text -> Maybe Int +readMaybeInt t = + case reads (T.unpack t) of + [(n, "")] -> Just n + _ -> Nothing + parseGeminiResponse :: Value -> IO LLMTextResult parseGeminiResponse v = case parseMaybe (go modelVer) v of Nothing -> pure $ Left EmptyResponse - Just blocks -> do - blocks' <- mapM normalizeBlock blocks - case blocks' of + Just parts -> do + parts' <- mapM normalizePart parts + case parts' of [] -> pure $ Left EmptyResponse - _ -> - let text = T.concat [t | TextBlock t <- blocks'] - in pure $ Right (ChatResponse text blocks' (parseGeminiUsage v) Nothing) + _ -> pure $ Right (mkChatResponse parts' (parseGeminiUsage v)) where modelVer = parseMaybe parseModelVersion v - go :: Maybe Text -> Value -> Parser [ContentBlock] + go :: Maybe Text -> Value -> Parser [ContentPart] go mv = withObject "GeminiResponse" $ \o -> do (cand : _) <- o .: "candidates" :: Parser [Value] withObject @@ -344,72 +431,100 @@ withObject "content" ( \cco -> do - parts <- cco .: "parts" :: Parser [Value] - mapM (parsePart mv) parts + ps <- cco .: "parts" :: Parser [Value] + mapM (parsePart mv) ps ) cont ) cand - parsePart :: Maybe Text -> Value -> Parser ContentBlock - parsePart mv = withObject "part" $ \o -> do - mSig <- o .:? "thoughtSignature" :: Parser (Maybe Text) - let tryText = TextBlock <$> (o .: "text") - tryFunctionCall = do - fc <- o .: "functionCall" - withObject - "functionCall" - ( \fco -> do - name <- fco .: "name" - args <- fco .:? "args" .!= object [] - -- Gemini doesn't provide a call id; use the function name. - -- normalizeBlock replaces it with a unique id later. - pure $ ToolCallBlock (attachGeminiMeta mv mSig (mkToolCall name name args)) - ) - fc - tryText <|> tryFunctionCall +parsePart :: Maybe Text -> Value -> Parser ContentPart +parsePart mv = withObject "part" $ \o -> do + mSig <- o .:? "thoughtSignature" :: Parser (Maybe Text) + mThought <- o .:? "thought" :: Parser (Maybe Bool) + let tryText = do + t <- o .: "text" + if mThought == Just True + then do + let opaque = + case mSig of + Nothing -> Nothing + Just sig -> + Just + ProviderOpaque + { poProvider = geminiProviderName, + poModel = mv, + poPayload = + object $ + ["text" .= t, "thought" .= True] + ++ ["thoughtSignature" .= sig] + } + pure $ thinkingPart (ThinkingContent (Just t) opaque) + else + -- Plain text; preserve signature as opaque thinking sidecar only when + -- required — Gemini 2.5 may put the signature on the first part. + pure $ textPart t + tryFunctionCall = do + fc <- o .: "functionCall" + withObject + "functionCall" + ( \fco -> do + name <- fco .: "name" + args <- fco .:? "args" .!= object [] + pure $ toolCallPart (attachGeminiMeta mv mSig (mkToolCall name name args)) + ) + fc + tryFunctionCall <|> tryText --- | Extract @modelVersion@ from a Gemini response or chunk. This is the --- canonical model name (e.g. @gemini-2.5-pro-001@) used as the binding key --- for thought signatures. parseModelVersion :: Value -> Parser Text parseModelVersion = withObject "GeminiResponse" (.: "modelVersion") --- | Build the provider-metadata bag stored on a 'ToolCall' for Gemini. --- Returns the original tool call unchanged when there's no signature to --- preserve, so we don't pollute logs with empty metadata. attachGeminiMeta :: Maybe Text -> Maybe Text -> ToolCall -> ToolCall attachGeminiMeta _ Nothing tc = tc attachGeminiMeta mModel (Just sig) tc = - tc {tcProviderMeta = Just (object (("thoughtSignature" .= sig) : modelField))} - where - modelField = ["model" .= m | Just m <- [mModel]] + tc + { tcProviderMeta = + Just + ProviderOpaque + { poProvider = geminiProviderName, + poModel = mModel, + poPayload = object ["thoughtSignature" .= sig] + } + } -- | Pull a thought signature out of 'tcProviderMeta' iff it was emitted by --- the model we're currently calling. Cross-model replay is unsafe — Gemini --- treats the signature as an opaque per-model token and will reject it. -signatureForModel :: Text -> Maybe Value -> Maybe Text -signatureForModel currentModel (Just (Object o)) = do - String sig <- KM.lookup "thoughtSignature" o - case KM.lookup "model" o of - Just (String m) | not (modelsMatch currentModel m) -> Nothing - _ -> Just sig +-- the model we're currently calling. +signatureForModel :: Text -> Maybe ProviderOpaque -> Maybe Text +signatureForModel currentModel (Just o) + | o.poProvider == geminiProviderName, + modelOk currentModel o.poModel = + case o.poPayload of + Object m | Just (String sig) <- KM.lookup "thoughtSignature" m -> Just sig + _ -> Nothing signatureForModel _ _ = Nothing --- | A signature emitted by @gemini-2.5-pro-001@ is safe to replay against --- @gemini-2.5-pro@ (and vice-versa): the response carries the resolved --- version string while requests typically use an alias. We accept either --- direction being a prefix of the other. modelsMatch :: Text -> Text -> Bool modelsMatch a b = a == b || T.isPrefixOf a b || T.isPrefixOf b a parseGeminiUsage :: Value -> Maybe Usage parseGeminiUsage = parseMaybe $ withObject "GeminiResponse" $ \o -> do u <- o .: "usageMetadata" - withObject - "usageMetadata" - (\uo -> Usage <$> uo .: "promptTokenCount" <*> uo .: "candidatesTokenCount" <*> pure 0) - u + withObject "usageMetadata" parseGeminiUsageObject u + +-- | Gemini @promptTokenCount@ already includes cached content tokens. +parseGeminiUsageObject :: Object -> Parser Usage +parseGeminiUsageObject uo = do + prompt <- uo .: "promptTokenCount" + candidates <- uo .: "candidatesTokenCount" + cacheRead <- uo .:? "cachedContentTokenCount" .!= 0 + pure + Usage + { usageInputTokens = prompt, + usageOutputTokens = candidates, + usageCacheReadTokens = cacheRead, + usageCacheCreationTokens = 0, + usageTotalCost = 0 + } parseGeminiObjectResponse :: Value -> IO LLMObjectResult parseGeminiObjectResponse v = case parseMaybe go v of diff --git a/src/LLM/Providers/OpenAI.hs b/src/LLM/Providers/OpenAI.hs --- a/src/LLM/Providers/OpenAI.hs +++ b/src/LLM/Providers/OpenAI.hs @@ -24,6 +24,7 @@ decodeStrict', encode, object, + toJSON, withObject, (.!=), (.:), @@ -33,7 +34,8 @@ import Data.ByteString.Lazy qualified as BSL import Data.Foldable (forM_) import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef) -import Data.Maybe (fromMaybe, isNothing) +import Data.Maybe (fromMaybe, listToMaybe, mapMaybe) +import Control.Applicative ((<|>)) import Data.Text (Text) import Data.Text qualified as T import Data.Text.Encoding (decodeUtf8, encodeUtf8) @@ -47,21 +49,30 @@ reqModel, reqSystem, reqTemperature, + reqThinking, reqTools ), - ChatResponse (ChatResponse, respContent, respReasoning, respText, respUsage), - ContentBlock (..), + ContentPart (..), LLMError (EmptyResponse), LLMGateway, LLMTextResult, MessageEncodeOptions (..), + PartBody (..), StreamEvent (..), + ThinkingContent (..), + ThinkingMode (..), ToolCall (..), - mkToolCall, ToolDef (toolDescription, toolName, toolParameters), ToolResult (trCallId, trContent), Turn (..), + ImageSource (..), defaultMessageEncodeOptions, + mkChatResponse, + mkToolCall, + stripForeignOpaque, + textPart, + thinkingPart, + toolCallPart, ) import LLM.Core.Usage (Usage (..)) import Network.HTTP.Client qualified as HC @@ -81,6 +92,9 @@ (/:), ) +openAIProviderName :: Text +openAIProviderName = "openai" + -- | Create an OpenAI client for api.openai.com. Takes the API key as a parameter. openAIGateway :: Text -> LLMGateway openAIGateway apiKey = toGateway $ openAIProvider apiKey @@ -164,36 +178,85 @@ "messages" .= buildMessages defaultMessageEncodeOptions r ] ++ ["temperature" .= t | Just t <- [r.reqTemperature]] + ++ effortPairs r ++ ["tools" .= map encodeToolDef r.reqTools | not (null r.reqTools)] ++ ["stream" .= True | stream] ++ ["stream_options" .= object ["include_usage" .= True] | stream] +-- | Chat Completions reasoning_effort, when the catalog requests thinking. +effortPairs :: ChatRequest -> [Pair] +effortPairs r = + case r.reqThinking of + Just tm + | tm.tmEnabled, + Just e <- tm.tmEffort -> + ["reasoning_effort" .= e] + _ -> [] + buildMessages :: MessageEncodeOptions -> ChatRequest -> [Value] buildMessages opts r = maybe [] (\sys -> [object ["role" .= ("system" :: Text), "content" .= sys]]) r.reqSystem ++ concatMap (encodeTurn opts) r.reqConversation encodeTurn :: MessageEncodeOptions -> Turn -> [Value] -encodeTurn _ (UserTurn content) = +encodeTurn _ (UserMessage parts) = [ object [ "role" .= ("user" :: Text), - "content" .= content + "content" .= encodeUserContent parts ] ] -encodeTurn opts (AssistantTurn text mReasoning calls) = - [ object $ - ["role" .= ("assistant" :: Text)] - ++ ["content" .= text | not (T.null text)] - ++ [ "reasoning_content" .= rc - | opts.meoIncludeReasoning, - Just rc <- [mReasoning], - not (T.null rc) - ] - ++ ["tool_calls" .= map encodeToolCall calls | not (null calls)] - ] +encodeTurn opts (AssistantMessage parts) = + let cleaned = map (stripForeignOpaque openAIProviderName) parts + -- Cache hints are ignored on the OpenAI wire; content is unchanged. + text = T.concat [t | ContentPart (TextPart t) _ <- cleaned] + mReasoning = + listToMaybe + [ t + | ContentPart (ThinkingPart tc) _ <- cleaned, + Just t <- [tc.thinkingText], + not (T.null t) + ] + calls = [tc | ContentPart (ToolCallPart tc) _ <- cleaned] + in [ object $ + ["role" .= ("assistant" :: Text)] + ++ ["content" .= text | not (T.null text)] + ++ [ "reasoning_content" .= rc + | opts.meoIncludeReasoning, + Just rc <- [mReasoning] + ] + ++ ["tool_calls" .= map encodeToolCall calls | not (null calls)] + ] encodeTurn _ (ToolTurn results) = map encodeToolResult results +-- | Single text stays a string (recorded-fixture compatible); mixed or image +-- content uses the OpenAI multimodal content-part array. +-- Cache hints are omitted from the wire JSON. +encodeUserContent :: [ContentPart] -> Value +encodeUserContent [ContentPart (TextPart t) _] = String t +encodeUserContent parts = toJSON (mapMaybe encodeUserPart parts) + where + encodeUserPart (ContentPart (TextPart t) _) = + Just $ object ["type" .= ("text" :: Text), "text" .= t] + encodeUserPart (ContentPart (ImagePart src) _) = + Just $ encodeImagePart src + encodeUserPart _ = Nothing + +encodeImagePart :: ImageSource -> Value +encodeImagePart (ImageUrl url) = + object + [ "type" .= ("image_url" :: Text), + "image_url" .= object ["url" .= url] + ] +encodeImagePart (ImageBase64 mediaType data_) = + object + [ "type" .= ("image_url" :: Text), + "image_url" + .= object + [ "url" .= ("data:" <> mediaType <> ";base64," <> data_) + ] + ] + encodeToolDef :: ToolDef -> Value encodeToolDef td = object @@ -231,20 +294,12 @@ parseOpenAIResponse :: Value -> LLMTextResult parseOpenAIResponse v = case parseMaybe go v of Nothing -> Left EmptyResponse - Just (mReasoning, blocks) -> - if null blocks && isNothing mReasoning + Just parts -> + if null parts then Left EmptyResponse - else - let text = T.concat [t | TextBlock t <- blocks] - in Right - ChatResponse - { respText = text, - respContent = blocks, - respUsage = parseOpenAIUsage v, - respReasoning = mReasoning - } + else Right (mkChatResponse parts (parseOpenAIUsage v)) where - go :: Value -> Parser (Maybe Text, [ContentBlock]) + go :: Value -> Parser [ContentPart] go = withObject "OpenAIResponse" $ \o -> do (choice : _) <- o .: "choices" :: Parser [Value] withObject @@ -255,18 +310,23 @@ ) choice - parseMessage :: Object -> Parser (Maybe Text, [ContentBlock]) + parseMessage :: Object -> Parser [ContentPart] parseMessage mo = do mReasoning <- mo .:? "reasoning_content" :: Parser (Maybe Text) contentBlocks <- do mc <- mo .:? "content" :: Parser (Maybe Text) - pure [TextBlock t | Just t <- [mc], not (T.null t)] + pure [textPart t | Just t <- [mc], not (T.null t)] toolBlocks <- do tcs <- mo .:? "tool_calls" .!= [] :: Parser [Value] mapM parseToolCall tcs - pure (mReasoning, contentBlocks ++ toolBlocks) + let thinkingBlocks = + [ thinkingPart (ThinkingContent (Just rc) Nothing) + | Just rc <- [mReasoning], + not (T.null rc) + ] + pure (thinkingBlocks ++ contentBlocks ++ toolBlocks) - parseToolCall :: Value -> Parser ContentBlock + parseToolCall :: Value -> Parser ContentPart parseToolCall = withObject "tool_call" $ \tc -> do cid <- tc .: "id" fn <- tc .: "function" @@ -278,23 +338,56 @@ let args = case decodeStrict' (encodeUtf8 argsStr) of Just v' -> v' Nothing -> String argsStr - pure $ ToolCallBlock (mkToolCall cid name args) + pure $ toolCallPart (mkToolCall cid name args) ) fn parseOpenAIUsage :: Value -> Maybe Usage parseOpenAIUsage = parseMaybe $ withObject "OpenAIResponse" $ \o -> do u <- o .: "usage" - withObject "usage" (\uo -> Usage <$> uo .: "prompt_tokens" <*> uo .: "completion_tokens" <*> pure 0) u + withObject "usage" parseOpenAIUsageObject u +-- | Shared OpenAI Chat Completions / DeepSeek usage shape. +-- +-- @prompt_tokens@ is already the total input (including cache hits). Cache +-- reads come from DeepSeek's @prompt_cache_hit_tokens@ or OpenAI's +-- @prompt_tokens_details.cached_tokens@. Optional @cache_write_tokens@ map to +-- creation counters. +parseOpenAIUsageObject :: Object -> Parser Usage +parseOpenAIUsageObject uo = do + prompt <- uo .: "prompt_tokens" + completion <- uo .: "completion_tokens" + cacheReadDeepSeek <- uo .:? "prompt_cache_hit_tokens" + details <- uo .:? "prompt_tokens_details" :: Parser (Maybe Value) + (cacheReadDetails, cacheWrite) <- case details of + Nothing -> pure (Nothing, Nothing) + Just d -> + withObject + "prompt_tokens_details" + ( \dto -> + (,) + <$> dto .:? "cached_tokens" + <*> dto .:? "cache_write_tokens" + ) + d + let cacheRead = fromMaybe 0 (cacheReadDeepSeek <|> cacheReadDetails) + cacheCreate = fromMaybe 0 cacheWrite + pure + Usage + { usageInputTokens = prompt, + usageOutputTokens = completion, + usageCacheReadTokens = cacheRead, + usageCacheCreationTokens = cacheCreate, + usageTotalCost = 0 + } + -- Streaming parseOpenAIStream :: HC.BodyReader -> (StreamEvent -> IO ()) -> IO LLMTextResult parseOpenAIStream reader callback = do - blocksRef <- newIORef ([] :: [ContentBlock]) + blocksRef <- newIORef ([] :: [ContentPart]) reasoningRef <- newIORef Nothing usageRef <- newIORef Nothing - -- Track in-flight tool calls: index -> (id, name, accumulated args) toolAccRef <- newIORef ([] :: [(Int, Text, Text, Text)]) readSSEEvents (HC.brRead reader) $ \sse -> do let raw = sse.sseData @@ -303,33 +396,27 @@ else case decodeStrict' (encodeUtf8 raw) of Nothing -> pure () Just v -> do - -- Reasoning deltas case parseMaybe parseStreamReasoningDelta v of Just (Just txt) | not (T.null txt) -> do modifyIORef' reasoningRef (Just . maybe txt (<> txt)) callback (StreamReasoningDelta txt) _ -> pure () - -- Text deltas case parseMaybe parseStreamTextDelta v of Just txt | not (T.null txt) -> do - modifyIORef' blocksRef (TextBlock txt :) + modifyIORef' blocksRef (textPart txt :) callback (StreamDelta txt) _ -> pure () - -- Tool call deltas case parseMaybe parseStreamToolDelta v of Just (idx, mId, mName, argChunk) -> do modifyIORef' toolAccRef $ \acc -> case lookup idx [(i, (i, cid, n, a)) | (i, cid, n, a) <- acc] of Nothing -> - -- New tool call let cid = fromMaybe "" mId n = fromMaybe "" mName in acc ++ [(idx, cid, n, argChunk)] Just (_, cid, n, a) -> - -- Accumulate arguments [(if i == idx then (i, cid, n, a <> argChunk) else entry) | entry@(i, _, _, _) <- acc] Nothing -> pure () - -- Finish reason: emit accumulated tool calls case parseMaybe parseFinishReason v of Just "tool_calls" -> do tools <- readIORef toolAccRef @@ -338,38 +425,33 @@ Just a -> a Nothing -> String argsStr tc = mkToolCall cid name args - modifyIORef' blocksRef (ToolCallBlock tc :) + modifyIORef' blocksRef (toolCallPart tc :) callback (StreamToolCall tc) writeIORef toolAccRef [] _ -> pure () - -- Usage (in the final chunk when stream_options.include_usage is set) case parseMaybe parseStreamUsage v of Just u -> writeIORef usageRef (Just u) Nothing -> pure () - -- Flush any remaining tool calls tools <- readIORef toolAccRef forM_ tools $ \(_, cid, name, argsStr) -> do let args = case decodeStrict' (encodeUtf8 argsStr) of Just a -> a Nothing -> String argsStr tc = mkToolCall cid name args - modifyIORef' blocksRef (ToolCallBlock tc :) + modifyIORef' blocksRef (toolCallPart tc :) callback (StreamToolCall tc) - blocks <- reverse <$> readIORef blocksRef + textBlocks <- reverse <$> readIORef blocksRef mReasoning <- readIORef reasoningRef usage <- readIORef usageRef - let text = T.concat [t | TextBlock t <- blocks] - if null blocks && isNothing mReasoning + let thinkingBlocks = + [ thinkingPart (ThinkingContent (Just rc) Nothing) + | Just rc <- [mReasoning], + not (T.null rc) + ] + parts = thinkingBlocks ++ textBlocks + if null parts then pure $ Left EmptyResponse - else - pure $ - Right - ChatResponse - { respText = text, - respContent = blocks, - respUsage = usage, - respReasoning = mReasoning - } + else pure $ Right (mkChatResponse parts usage) parseStreamReasoningDelta :: Value -> Parser (Maybe Text) parseStreamReasoningDelta = withObject "chunk" $ \o -> do @@ -421,4 +503,4 @@ parseStreamUsage :: Value -> Parser Usage parseStreamUsage = withObject "chunk" $ \o -> do u <- o .: "usage" - withObject "usage" (\uo -> Usage <$> uo .: "prompt_tokens" <*> uo .: "completion_tokens" <*> pure 0) u + withObject "usage" parseOpenAIUsageObject u diff --git a/test/LLM/ChatSpec.hs b/test/LLM/ChatSpec.hs --- a/test/LLM/ChatSpec.hs +++ b/test/LLM/ChatSpec.hs @@ -3,11 +3,12 @@ module LLM.ChatSpec (spec) where import Data.Aeson (object, (.=)) +import Data.IORef (IORef, modifyIORef', newIORef, readIORef) import Data.Map qualified as Map import Data.Text (Text) import Heptapod (generate) import LLM.Agent.Events (noEventObserver) -import LLM.Agent.Generate (generateText) +import LLM.Agent.Generate (generateText, streamText) import LLM.Agent.Types ( Agent (..), RuntimeArgs (..), @@ -16,21 +17,28 @@ ) import LLM.Core.Abort (AbortSignal, abort, newAbortSignal) import LLM.Core.Types - ( ChatRequest (..), + ( CacheHint (..), + ChatRequest (..), ChatResponse (..), - ContentBlock (TextBlock, ToolCallBlock), + ContentPart (..), + PartBody (..), + textPart, toolCallPart, pattern UserTurn, + Turn (..), + cacheEphemeral, + imageUrlPart, LLMError (..), LLMGateway (..), LLMHooks (..), ToolDef (ToolDef, toolDescription, toolName, toolParameters, toolReadonly), - Turn (..), mkToolCall, ) -import LLM.Core.Usage (PricingInfo (..), Usage (Usage)) +import LLM.Core.Usage (PricingInfo (..), defaultPricingInfo, mkUsage) import LLM.Generate.Logger (noHooks) import LLM.Generate.ModelConfig - ( ModelConfig (..), + ( ModelCapabilities (..), + ModelConfig (..), ModelWithFallbacks (..), + defaultModelCapabilities, ) import LLM.Generate.Types ( GenerateError (..), @@ -43,6 +51,7 @@ expectationFailure, it, shouldBe, + shouldSatisfy, ) -- | A mock gateway that returns a fixed response @@ -72,10 +81,10 @@ { gwName = "mock-tool", gwGenerateText = \_ req -> if any isToolTurn req.reqConversation - then pure $ Right (ChatResponse "The weather is sunny." [TextBlock "The weather is sunny."] (Just (Usage 80 15 0)) Nothing) + then pure $ Right (ChatResponse "The weather is sunny." [textPart "The weather is sunny."] (Just (mkUsage 80 15)) Nothing) else let tc = mkToolCall "call_1" "get_weather" (object ["location" .= ("London" :: Text)]) - in pure $ Right (ChatResponse "" [ToolCallBlock tc] (Just (Usage 50 10 0)) Nothing), + in pure $ Right (ChatResponse "" [toolCallPart tc] (Just (mkUsage 50 10)) Nothing), gwStreamText = \_ _ _ -> pure $ Right (ChatResponse "" [] Nothing Nothing), gwGenerateObject = \_ _ _ -> pure $ Right (object [], Nothing) } @@ -84,7 +93,7 @@ isToolTurn _ = False zeroPricing :: PricingInfo -zeroPricing = PricingInfo 0 0 +zeroPricing = defaultPricingInfo 0 0 -- | Wrap a gateway in a ModelConfig with test defaults mockModel :: LLMGateway -> ModelConfig @@ -96,6 +105,7 @@ mcMaxTokens = 1024, mcTemperature = Nothing, mcThinking = Nothing, + mcCapabilities = defaultModelCapabilities, mcRequestTimeout = Nothing, mcThrottleDelay = Nothing, mcRetryCount = 0, @@ -154,19 +164,64 @@ rt <- mkRuntime mSig generateText agent models toolMap rt turns +runStreamGenerate :: + Agent -> + ModelWithFallbacks -> + ToolMap Text -> + Maybe AbortSignal -> + [Turn] -> + IO (Either GenerateErrorResult GenerateTextResult) +runStreamGenerate agent models toolMap mSig turns = do + rt <- mkRuntime mSig + streamText (\_ -> pure ()) agent models toolMap rt turns + +-- | Gateway that records each request conversation, then behaves like 'mockToolGateway'. +capturingToolGateway :: IORef [[Turn]] -> LLMGateway +capturingToolGateway ref = + let respond req = do + modifyIORef' ref (req.reqConversation :) + if any isToolTurn req.reqConversation + then pure $ Right (ChatResponse "The weather is sunny." [textPart "The weather is sunny."] (Just (mkUsage 80 15)) Nothing) + else + let tc = mkToolCall "call_1" "get_weather" (object ["location" .= ("London" :: Text)]) + in pure $ Right (ChatResponse "" [toolCallPart tc] (Just (mkUsage 50 10)) Nothing) + in LLMGateway + { gwName = "mock-tool-capture", + gwGenerateText = \_ -> respond, + gwStreamText = \_ req _onEvent -> respond req, + gwGenerateObject = \_ _ _ -> pure $ Right (object [], Nothing) + } + where + isToolTurn (ToolTurn _) = True + isToolTurn _ = False + +hasCachedUserPrefix :: [Turn] -> Bool +hasCachedUserPrefix = + any + ( \case + UserMessage parts -> + any + ( \case + ContentPart (TextPart _) (Just CacheEphemeral) -> True + _ -> False + ) + parts + _ -> False + ) + spec :: Spec spec = describe "Chat" $ do let toolMap = Map.fromList [("get_weather", weatherTool)] describe "generateText" $ do it "returns text for a simple response" $ do - let gw = mockGateway (ChatResponse "Hi there!" [TextBlock "Hi there!"] (Just (Usage 10 5 0)) Nothing) + let gw = mockGateway (ChatResponse "Hi there!" [textPart "Hi there!"] (Just (mkUsage 10 5)) Nothing) models = ModelWithFallbacks (mockModel gw) [] result <- runGenerate defaultAgent models toolMap Nothing [UserTurn "hello"] case result of Right r -> do r.gtrText `shouldBe` "Hi there!" - length r.gtrNewMessages `shouldBe` 1 -- AssistantTurn - r.gtrUsage `shouldBe` Usage 10 5 0 + length r.gtrNewMessages `shouldBe` 1 -- assistantTurn + r.gtrUsage `shouldBe` mkUsage 10 5 Left err -> expectationFailure $ show err it "propagates errors" $ do @@ -184,18 +239,37 @@ case result of Right r -> do r.gtrText `shouldBe` "The weather is sunny." - -- AssistantTurn(tool call) + ToolTurn + AssistantTurn(final) + -- assistantTurn(tool call) + ToolTurn + assistantTurn(final) length r.gtrNewMessages `shouldBe` 3 - r.gtrUsage `shouldBe` Usage 130 25 0 -- 50+80 input, 10+15 output + r.gtrUsage `shouldBe` mkUsage 130 25 -- 50+80 input, 10+15 output Left err -> expectationFailure $ show err + it "preserves cache hints on user parts across non-streaming tool rounds" $ do + ref <- newIORef [] + let agent = defaultAgent {agTools = ["get_weather"]} + models = ModelWithFallbacks (mockModel (capturingToolGateway ref)) [] + msgs = + [ UserMessage + [ cacheEphemeral (textPart "static system-like context"), + textPart "weather in london?" + ] + ] + result <- runGenerate agent models toolMap Nothing msgs + case result of + Right r -> do + r.gtrText `shouldBe` "The weather is sunny." + convs <- reverse <$> readIORef ref + length convs `shouldBe` 2 + mapM_ (`shouldSatisfy` hasCachedUserPrefix) convs + Left err -> expectationFailure $ show err + it "respects maxToolRounds" $ do let infiniteToolGateway = LLMGateway { gwName = "mock-infinite", gwGenerateText = \_ _ -> let tc = mkToolCall "call_1" "get_weather" (object []) - in pure $ Right (ChatResponse "" [ToolCallBlock tc] Nothing Nothing), + in pure $ Right (ChatResponse "" [toolCallPart tc] Nothing Nothing), gwStreamText = \_ _ _ -> pure $ Right (ChatResponse "" [] Nothing Nothing), gwGenerateObject = \_ _ _ -> pure $ Right (object [], Nothing) } @@ -208,7 +282,7 @@ it "falls back to next model on retryable error" $ do let failGw = mockErrorGateway (HttpError 503 "service unavailable") - okGw = mockGateway (ChatResponse "Fallback worked!" [TextBlock "Fallback worked!"] (Just (Usage 10 5 0)) Nothing) + okGw = mockGateway (ChatResponse "Fallback worked!" [textPart "Fallback worked!"] (Just (mkUsage 10 5)) Nothing) models = ModelWithFallbacks (mockModel failGw) [mockModel okGw] result <- runGenerate defaultAgent models toolMap Nothing [UserTurn "hello"] case result of @@ -217,13 +291,44 @@ it "falls back on non-retryable error too" $ do let failGw = mockErrorGateway (HttpError 400 "bad request") - okGw = mockGateway (ChatResponse "Fallback worked!" [TextBlock "Fallback worked!"] (Just (Usage 10 5 0)) Nothing) + okGw = mockGateway (ChatResponse "Fallback worked!" [textPart "Fallback worked!"] (Just (mkUsage 10 5)) Nothing) models = ModelWithFallbacks (mockModel failGw) [mockModel okGw] result <- runGenerate defaultAgent models toolMap Nothing [UserTurn "hello"] case result of Right r -> r.gtrText `shouldBe` "Fallback worked!" Left err -> expectationFailure $ "Expected fallback success, got: " <> show err + it "skips a non-vision model and succeeds with a vision-capable fallback" $ do + let boomGw = + LLMGateway + { gwName = "should-not-run", + gwGenerateText = \_ _ -> pure (Left (HttpError 500 "should not be called")), + gwStreamText = \_ _ _ -> pure (Left (HttpError 500 "should not be called")), + gwGenerateObject = \_ _ _ -> pure (Left (HttpError 500 "should not be called")) + } + okGw = mockGateway (ChatResponse "I see a cat." [textPart "I see a cat."] (Just (mkUsage 10 5)) Nothing) + noVision = mockModel boomGw + withVision = + (mockModel okGw) + { mcCapabilities = defaultModelCapabilities {capVision = True} + } + models = ModelWithFallbacks noVision [withVision] + msgs = [UserMessage [imageUrlPart "https://example.com/cat.png", textPart "what is this?"]] + result <- runGenerate defaultAgent models toolMap Nothing msgs + case result of + Right r -> r.gtrText `shouldBe` "I see a cat." + Left err -> expectationFailure $ "Expected vision fallback success, got: " <> show err + + it "returns UnsupportedCapability when no vision-capable model remains" $ do + let noVision1 = mockModel (mockErrorGateway (HttpError 500 "unused")) + noVision2 = mockModel (mockErrorGateway (HttpError 500 "unused")) + models = ModelWithFallbacks noVision1 [noVision2] + msgs = [UserMessage [imageUrlPart "https://example.com/cat.png", textPart "what?"]] + result <- runGenerate defaultAgent models toolMap Nothing msgs + case result of + Left GenerateErrorResult {gerError = GErrLLM (UnsupportedCapability _)} -> pure () + other -> expectationFailure $ "Expected UnsupportedCapability, got: " <> show other + it "returns error from last model when all fail" $ do let failGw1 = mockErrorGateway (HttpError 503 "service unavailable") failGw2 = mockErrorGateway (HttpError 400 "bad request") @@ -234,7 +339,7 @@ other -> expectationFailure $ "Expected HttpError 400 from last model, got: " <> show other it "returns Aborted when signal is fired before the call" $ do - let gw = mockGateway (ChatResponse "Hi!" [TextBlock "Hi!"] Nothing Nothing) + let gw = mockGateway (ChatResponse "Hi!" [textPart "Hi!"] Nothing Nothing) models = ModelWithFallbacks (mockModel gw) [] sig <- newAbortSignal abort sig @@ -264,7 +369,7 @@ gwGenerateText = \_ _ -> let tc1 = mkToolCall "c1" "slow" (object []) tc2 = mkToolCall "c2" "slow" (object []) - in pure $ Right (ChatResponse "" [ToolCallBlock tc1, ToolCallBlock tc2] Nothing Nothing), + in pure $ Right (ChatResponse "" [toolCallPart tc1, toolCallPart tc2] Nothing Nothing), gwStreamText = \_ _ _ -> pure $ Right (ChatResponse "" [] Nothing Nothing), gwGenerateObject = \_ _ _ -> pure $ Right (object [], Nothing) } @@ -277,8 +382,8 @@ other -> expectationFailure $ "Expected GErrAborted during tools, got: " <> show other it "does not fall back on Aborted" $ do - let gw = mockGateway (ChatResponse "Hi!" [TextBlock "Hi!"] Nothing Nothing) - okGw = mockGateway (ChatResponse "Fallback" [TextBlock "Fallback"] Nothing Nothing) + let gw = mockGateway (ChatResponse "Hi!" [textPart "Hi!"] Nothing Nothing) + okGw = mockGateway (ChatResponse "Fallback" [textPart "Fallback"] Nothing Nothing) models = ModelWithFallbacks (mockModel gw) [mockModel okGw] sig <- newAbortSignal abort sig @@ -286,3 +391,23 @@ case result of Left GenerateErrorResult {gerError = GErrAborted} -> pure () other -> expectationFailure $ "Expected GErrAborted (no fallback), got: " <> show other + + describe "streamText" $ do + it "preserves cache hints on user parts across streaming tool rounds" $ do + ref <- newIORef [] + let agent = defaultAgent {agTools = ["get_weather"]} + models = ModelWithFallbacks (mockModel (capturingToolGateway ref)) [] + msgs = + [ UserMessage + [ cacheEphemeral (textPart "static system-like context"), + textPart "weather in london?" + ] + ] + result <- runStreamGenerate agent models toolMap Nothing msgs + case result of + Right r -> do + r.gtrText `shouldBe` "The weather is sunny." + convs <- reverse <$> readIORef ref + length convs `shouldBe` 2 + mapM_ (`shouldSatisfy` hasCachedUserPrefix) convs + Left err -> expectationFailure $ show err diff --git a/test/LLM/ClaudeSpec.hs b/test/LLM/ClaudeSpec.hs --- a/test/LLM/ClaudeSpec.hs +++ b/test/LLM/ClaudeSpec.hs @@ -2,20 +2,58 @@ module LLM.ClaudeSpec (spec) where -import Data.Aeson (eitherDecodeFileStrict') +import Data.Aeson + ( Value (Array, Number, Object, String), + eitherDecodeFileStrict', + object, + (.=), + ) +import Data.Aeson.Key qualified as K +import Data.Aeson.KeyMap qualified as KM +import Data.ByteString qualified as BS +import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef) +import Data.Maybe (mapMaybe) +import Data.Scientific (toBoundedInteger, toRealFloat) +import Data.Text (Text) +import Data.Vector qualified as V import LLM.Core.Types - ( ChatResponse (respText), - ToolCall (tcId, tcName), + ( ChatRequest (..), + ChatResponse (respContent, respReasoning, respText), + ContentPart (..), + PartBody (..), + ProviderOpaque (..), + StreamEvent (..), + ThinkingContent (..), + ThinkingMode (..), + ToolCall (..), + ToolResult (..), + Turn (..), + mkToolCall, + textPart, + thinkingPart, + toolCallPart, + imageUrlPart, + imageBase64Part, + cacheEphemeral, + pattern UserTurn, ) -import LLM.Core.Usage (Usage (Usage)) +import LLM.Core.Usage (Usage (..), mkUsage) import LLM.Core.Utils (getToolCalls, hasToolCalls) -import LLM.Providers.Claude (parseClaudeResponse, parseClaudeUsage) +import LLM.Providers.Claude + ( claudeBuildBody, + encodeTurn, + effortToBudgetTokens, + parseClaudeResponse, + parseClaudeStream, + parseClaudeUsage, + ) import Test.Hspec ( Spec, describe, expectationFailure, it, shouldBe, + shouldContain, ) spec :: Spec @@ -40,11 +78,337 @@ tc.tcId `shouldBe` "toolu_01A09q90qw90lq917835lq9" Left err -> expectationFailure $ "Parse failed: " <> show err + it "preserves thinking -> text -> tool_use order and opaque signature" $ do + Right val <- eitherDecodeFileStrict' "test/fixtures/claude-thinking-tool-use.json" + case parseClaudeResponse val of + Right resp -> do + resp.respReasoning `shouldBe` Just "I should call the weather tool." + resp.respText `shouldBe` "Checking the weather." + case resp.respContent of + [ ContentPart (ThinkingPart tc) Nothing, + ContentPart (TextPart "Checking the weather.") Nothing, + ContentPart (ToolCallPart tool) Nothing + ] -> do + tc.thinkingText `shouldBe` Just "I should call the weather tool." + case tc.thinkingOpaque of + Just o -> do + o.poProvider `shouldBe` "claude" + o.poModel `shouldBe` Just "claude-haiku-4-5-20251001" + lookupText "signature" o.poPayload `shouldBe` Just "sig_thinking_abc123" + Nothing -> expectationFailure "expected thinking opaque" + tool.tcId `shouldBe` "toolu_weather_1" + other -> expectationFailure $ "unexpected parts: " <> show other + Left err -> expectationFailure $ "Parse failed: " <> show err + + describe "parseClaudeStream" $ do + it "preserves thinking -> text -> tool_use order and opaque signature" $ do + sse <- BS.readFile "test/fixtures/claude-thinking-tool-use.sse" + eventsRef <- newIORef ([] :: [StreamEvent]) + reader <- mkBodyReader sse + result <- parseClaudeStream "claude-haiku-4-5-20251001" reader $ \ev -> + modifyIORef' eventsRef (ev :) + case result of + Right resp -> do + resp.respReasoning `shouldBe` Just "I should call the weather tool." + resp.respText `shouldBe` "Checking the weather." + case resp.respContent of + [ ContentPart (ThinkingPart tc) Nothing, + ContentPart (TextPart "Checking the weather.") Nothing, + ContentPart (ToolCallPart tool) Nothing + ] -> do + tc.thinkingText `shouldBe` Just "I should call the weather tool." + case tc.thinkingOpaque of + Just o -> do + o.poProvider `shouldBe` "claude" + o.poModel `shouldBe` Just "claude-haiku-4-5-20251001" + lookupText "signature" o.poPayload `shouldBe` Just "sig_thinking_abc123" + Nothing -> expectationFailure "expected thinking opaque" + tool.tcId `shouldBe` "toolu_weather_1" + tool.tcName `shouldBe` "get_weather" + let encoded = + encodeTurn + "claude-haiku-4-5-20251001" + (AssistantMessage resp.respContent) + content = messageContent (head encoded) + contentTypes content `shouldBe` ["thinking", "text", "tool_use"] + case content of + (Object o : _) -> + lookupText "signature" (Object o) `shouldBe` Just "sig_thinking_abc123" + _ -> expectationFailure "expected thinking object" + let withResult = + encodeTurn + "claude-haiku-4-5-20251001" + ( ToolTurn + [ ToolResult + { trCallId = "toolu_weather_1", + trName = "get_weather", + trContent = "sunny" + } + ] + ) + length withResult `shouldBe` 1 + other -> expectationFailure $ "unexpected parts: " <> show other + events <- reverse <$> readIORef eventsRef + events + `shouldContain` [ StreamReasoningDelta "I should call the weather tool.", + StreamDelta "Checking the weather." + ] + case [tc | StreamToolCall tc <- events] of + [tc] -> tc.tcId `shouldBe` "toolu_weather_1" + other -> expectationFailure $ "expected one StreamToolCall, got: " <> show other + Left err -> expectationFailure $ "Stream parse failed: " <> show err + + describe "encodeTurn replay" $ do + it "replays signed thinking blocks before tool_use for the same model" $ do + Right val <- eitherDecodeFileStrict' "test/fixtures/claude-thinking-tool-use.json" + case parseClaudeResponse val of + Right resp -> do + let encoded = encodeTurn "claude-haiku-4-5-20251001" (AssistantMessage resp.respContent) + content = messageContent (head encoded) + contentTypes content `shouldBe` ["thinking", "text", "tool_use"] + case content of + (Object o : _) -> + lookupText "signature" (Object o) `shouldBe` Just "sig_thinking_abc123" + _ -> expectationFailure "expected thinking object" + let withResult = + encodeTurn + "claude-haiku-4-5-20251001" + ( ToolTurn + [ ToolResult + { trCallId = "toolu_weather_1", + trName = "get_weather", + trContent = "sunny" + } + ] + ) + length withResult `shouldBe` 1 + Left err -> expectationFailure $ show err + + it "omits foreign thinking opaque and foreign tool meta on encode" $ do + let foreignThinking = + thinkingPart + ThinkingContent + { thinkingText = Just "gemini thoughts", + thinkingOpaque = + Just + ProviderOpaque + { poProvider = "gemini", + poModel = Just "gemini-2.5-flash", + poPayload = object ["thoughtSignature" .= ("sig" :: Text)] + } + } + foreignTool = + let base = mkToolCall "c1" "get_weather" (object []) + in toolCallPart + base + { tcProviderMeta = + Just + ProviderOpaque + { poProvider = "gemini", + poModel = Just "gemini-2.5-flash", + poPayload = object ["thoughtSignature" .= ("sig" :: Text)] + } + } + turn = + AssistantMessage + [ foreignThinking, + textPart "hello", + foreignTool + ] + content = messageContent (head (encodeTurn "claude-haiku-4-5-20251001" turn)) + contentTypes content `shouldBe` ["text", "tool_use"] + case content of + [Object textO, Object toolO] -> do + lookupText "text" (Object textO) `shouldBe` Just "hello" + lookupText "name" (Object toolO) `shouldBe` Just "get_weather" + KM.lookup "thoughtSignature" toolO `shouldBe` Nothing + _ -> expectationFailure "expected text + tool_use only" + + describe "thinking request mapping" $ do + it "maps effort to budget_tokens and omits temperature when thinking is on" $ do + let req = + ChatRequest + { reqModel = "claude-haiku-4-5-20251001", + reqConversation = [UserTurn "hi"], + reqSystem = Nothing, + reqMaxTokens = 4096, + reqTemperature = Just 0.5, + reqTools = [], + reqThinking = Just ThinkingMode {tmEnabled = True, tmEffort = Just "high"} + } + body = claudeBuildBody False req + nestedText ["thinking", "type"] body `shouldBe` Just "enabled" + nestedInt ["thinking", "budget_tokens"] body `shouldBe` Just (effortToBudgetTokens "high") + lookupKey "temperature" body `shouldBe` Nothing + + it "keeps temperature when thinking is off" $ do + let req = + ChatRequest + { reqModel = "claude-haiku-4-5-20251001", + reqConversation = [UserTurn "hi"], + reqSystem = Nothing, + reqMaxTokens = 4096, + reqTemperature = Just 0.5, + reqTools = [], + reqThinking = Nothing + } + body = claudeBuildBody False req + lookupNumber "temperature" body `shouldBe` Just 0.5 + + describe "image request encoding" $ do + it "encodes URL and base64 image parts" $ do + Right b64 <- pure $ imageBase64Part "image/png" "aGVsbG8=" + let turn = + UserMessage + [ imageUrlPart "https://example.com/cat.jpg", + b64, + textPart "describe" + ] + content = messageContent (head (encodeTurn "claude-haiku-4-5-20251001" turn)) + contentTypes content `shouldBe` ["image", "image", "text"] + case content of + [Object urlO, Object b64O, Object textO] -> do + nestedText ["source", "type"] (Object urlO) `shouldBe` Just "url" + nestedText ["source", "url"] (Object urlO) + `shouldBe` Just "https://example.com/cat.jpg" + nestedText ["source", "type"] (Object b64O) `shouldBe` Just "base64" + nestedText ["source", "media_type"] (Object b64O) `shouldBe` Just "image/png" + nestedText ["source", "data"] (Object b64O) `shouldBe` Just "aGVsbG8=" + lookupText "text" (Object textO) `shouldBe` Just "describe" + _ -> expectationFailure "expected image/image/text blocks" + + describe "cache_control breakpoints" $ do + it "places cache_control on the marked block without reordering" $ do + Right b64 <- pure $ imageBase64Part "image/png" "aGVsbG8=" + let turn = + UserMessage + [ imageUrlPart "https://example.com/cat.jpg", + cacheEphemeral b64, + cacheEphemeral (textPart "describe") + ] + content = messageContent (head (encodeTurn "claude-haiku-4-5-20251001" turn)) + contentTypes content `shouldBe` ["image", "image", "text"] + case content of + [Object urlO, Object b64O, Object textO] -> do + KM.lookup "cache_control" urlO `shouldBe` Nothing + nestedText ["cache_control", "type"] (Object b64O) `shouldBe` Just "ephemeral" + nestedText ["cache_control", "type"] (Object textO) `shouldBe` Just "ephemeral" + nestedText ["source", "data"] (Object b64O) `shouldBe` Just "aGVsbG8=" + lookupText "text" (Object textO) `shouldBe` Just "describe" + _ -> expectationFailure "expected image/image/text blocks" + + it "uses a content-block array for a single cached text part" $ do + let turn = UserMessage [cacheEphemeral (textPart "static context")] + msg = head (encodeTurn "claude-haiku-4-5-20251001" turn) + case lookupKey "content" msg of + Just (Array arr) -> do + V.length arr `shouldBe` 1 + nestedText ["cache_control", "type"] (V.head arr) `shouldBe` Just "ephemeral" + lookupText "text" (V.head arr) `shouldBe` Just "static context" + Just (String _) -> + expectationFailure "cached single text must not encode as a bare string" + _ -> expectationFailure "expected content array" + describe "parseClaudeUsage" $ do it "extracts token counts" $ do Right val <- eitherDecodeFileStrict' "test/fixtures/claude-text.json" - parseClaudeUsage val `shouldBe` Just (Usage 25 10 0) + parseClaudeUsage val `shouldBe` Just (mkUsage 25 10) it "extracts token counts from tool_use response" $ do Right val <- eitherDecodeFileStrict' "test/fixtures/claude-tool-use.json" - parseClaudeUsage val `shouldBe` Just (Usage 50 35 0) + parseClaudeUsage val `shouldBe` Just (mkUsage 50 35) + + it "normalizes cache read/creation into total input without double-counting" $ do + let val = + object + [ "usage" + .= object + [ "input_tokens" .= (50 :: Int), + "output_tokens" .= (20 :: Int), + "cache_read_input_tokens" .= (100_000 :: Int), + "cache_creation_input_tokens" .= (1_200 :: Int) + ] + ] + parseClaudeUsage val + `shouldBe` Just + Usage + { usageInputTokens = 101_250, + usageOutputTokens = 20, + usageCacheReadTokens = 100_000, + usageCacheCreationTokens = 1_200, + usageTotalCost = 0 + } + + it "treats missing or zero cache fields as zero on recorded fixtures" $ do + Right val <- eitherDecodeFileStrict' "test/fixtures/claude-conversation-generated.json" + case val of + Array arr -> + case V.toList arr of + (Object first : _) -> + case KM.lookup "response" first of + Just resp -> + parseClaudeUsage resp + `shouldBe` Just (mkUsage 623 54) + _ -> expectationFailure "missing response" + _ -> expectationFailure "empty conversation array" + _ -> expectationFailure "expected conversation array" + +-- | Yield the full SSE payload once, then empty chunks (EOF). +mkBodyReader :: BS.ByteString -> IO (IO BS.ByteString) +mkBodyReader bs = do + ref <- newIORef (Just bs) + pure $ do + m <- readIORef ref + case m of + Just chunk -> do + writeIORef ref Nothing + pure chunk + Nothing -> pure BS.empty + +messageContent :: Value -> [Value] +messageContent (Object o) = + case KM.lookup "content" o of + Just (Array a) -> V.toList a + _ -> [] +messageContent _ = [] + +contentTypes :: [Value] -> [Text] +contentTypes = mapMaybe typ + where + typ (Object o) = case KM.lookup "type" o of + Just (String t) -> Just t + _ -> Nothing + typ _ = Nothing + +lookupText :: Text -> Value -> Maybe Text +lookupText key (Object o) = + KM.lookup (K.fromText key) o >>= \case + String t -> Just t + _ -> Nothing +lookupText _ _ = Nothing + +lookupNumber :: Text -> Value -> Maybe Double +lookupNumber key (Object o) = + KM.lookup (K.fromText key) o >>= \case + Number sci -> Just (toRealFloat sci) + _ -> Nothing +lookupNumber _ _ = Nothing + +lookupKey :: Text -> Value -> Maybe Value +lookupKey key (Object o) = KM.lookup (K.fromText key) o +lookupKey _ _ = Nothing + +nestedText :: [Text] -> Value -> Maybe Text +nestedText [key] v = lookupText key v +nestedText (key : rest) (Object o) = + KM.lookup (K.fromText key) o >>= nestedText rest +nestedText _ _ = Nothing + +nestedInt :: [Text] -> Value -> Maybe Int +nestedInt [key] (Object o) = + KM.lookup (K.fromText key) o >>= \case + Number sci -> toBoundedInteger sci + _ -> Nothing +nestedInt (key : rest) (Object o) = + KM.lookup (K.fromText key) o >>= nestedInt rest +nestedInt _ _ = Nothing diff --git a/test/LLM/DeepSeekSpec.hs b/test/LLM/DeepSeekSpec.hs --- a/test/LLM/DeepSeekSpec.hs +++ b/test/LLM/DeepSeekSpec.hs @@ -1,28 +1,36 @@ module LLM.DeepSeekSpec (spec) where -import Data.Aeson (Value (Object, String), object, (.=)) +import Data.Aeson (Value (Array, Object, String), eitherDecodeFileStrict', object, (.=)) import Data.Aeson.Key qualified as K import Data.Aeson.KeyMap qualified as KM import Data.Text (Text) +import Data.Vector qualified as V import LLM.Core.Types ( ChatRequest (..), - ChatResponse (respReasoning, respText), + ChatResponse (respContent, respReasoning, respText), + ContentPart (..), + PartBody (..), + ThinkingContent (..), ThinkingMode (..), Turn (..), + assistantTurn, + cacheEphemeral, deepSeekMessageEncodeOptions, defaultMessageEncodeOptions, mkToolCall, + textPart, ) +import LLM.Core.Usage (Usage (..)) import LLM.Providers.DeepSeek (deepSeekBuildBodyPairs) -import LLM.Providers.OpenAI (encodeTurn, parseOpenAIResponse) -import Test.Hspec (Spec, describe, it, shouldBe) +import LLM.Providers.OpenAI (encodeTurn, parseOpenAIResponse, parseOpenAIUsage) +import Test.Hspec (Spec, describe, expectationFailure, it, shouldBe) spec :: Spec spec = describe "DeepSeek thinking mode" $ do describe "message encoding" $ do it "includes reasoning_content when replaying assistant tool turns" $ do let turn = - AssistantTurn + assistantTurn "Let me check." (Just "I should call the weather tool.") [mkToolCall "call_1" "get_weather" (object ["location" .= ("London" :: Text)])] @@ -30,10 +38,18 @@ lookupText "reasoning_content" msg `shouldBe` Just "I should call the weather tool." it "omits reasoning_content for OpenAI-compatible default encoding" $ do - let turn = AssistantTurn "Hello" (Just "thinking") [] + let turn = assistantTurn "Hello" (Just "thinking") [] msg = head (encodeTurn defaultMessageEncodeOptions turn) lookupText "reasoning_content" msg `shouldBe` Nothing + it "omits cache hints from the wire while keeping content identical" $ do + let plain = UserMessage [textPart "static context"] + hinted = UserMessage [cacheEphemeral (textPart "static context")] + plainMsg = head (encodeTurn deepSeekMessageEncodeOptions plain) + hintedMsg = head (encodeTurn deepSeekMessageEncodeOptions hinted) + hintedMsg `shouldBe` plainMsg + lookupText "content" hintedMsg `shouldBe` Just "static context" + describe "request body" $ do it "disables thinking by default" $ do let body = object (deepSeekBuildBodyPairs False sampleRequest) @@ -75,14 +91,41 @@ Right resp -> do resp.respReasoning `shouldBe` Just "Let me think." resp.respText `shouldBe` "The answer is 42." + case resp.respContent of + [ ContentPart (ThinkingPart (ThinkingContent (Just "Let me think.") Nothing)) Nothing, + ContentPart (TextPart "The answer is 42.") Nothing + ] -> pure () + other -> fail $ "unexpected ordered parts: " <> show other Left err -> fail $ show err + describe "parseOpenAIUsage (DeepSeek cache fields)" $ do + it "maps prompt_cache_hit_tokens without double-counting prompt_tokens" $ do + Right val <- eitherDecodeFileStrict' "test/fixtures/deepseek-conversation-generated.json" + case val of + Array arr -> + case V.toList arr of + (Object first : _) -> + case KM.lookup "response" first of + Just resp -> + parseOpenAIUsage resp + `shouldBe` Just + Usage + { usageInputTokens = 336, + usageOutputTokens = 68, + usageCacheReadTokens = 256, + usageCacheCreationTokens = 0, + usageTotalCost = 0 + } + _ -> expectationFailure "missing response" + _ -> expectationFailure "empty conversation array" + _ -> expectationFailure "expected conversation array" + sampleRequest :: ChatRequest sampleRequest = ChatRequest { reqModel = "deepseek-v4-pro", reqConversation = - [ AssistantTurn "Hi" (Just "CoT") [mkToolCall "c1" "get_date" (object [])], + [ assistantTurn "Hi" (Just "CoT") [mkToolCall "c1" "get_date" (object [])], ToolTurn [] ], reqSystem = Nothing, diff --git a/test/LLM/GeminiSpec.hs b/test/LLM/GeminiSpec.hs --- a/test/LLM/GeminiSpec.hs +++ b/test/LLM/GeminiSpec.hs @@ -2,11 +2,27 @@ module LLM.GeminiSpec (spec) where -import Data.Aeson (eitherDecodeFileStrict') -import LLM.Core.Types (ChatResponse (respText), ToolCall (tcName)) -import LLM.Core.Usage (Usage (Usage)) +import Data.Aeson (Value (Array, Object, String), eitherDecodeFileStrict', object, (.=)) +import Data.Aeson.KeyMap qualified as KM +import Data.Text (Text) +import Data.Vector qualified as V +import LLM.Core.Types + ( ChatResponse (respText), + ProviderOpaque (..), + ThinkingContent (..), + ToolCall (..), + Turn (..), + mkToolCall, + textPart, + thinkingPart, + toolCallPart, + imageUrlPart, + imageBase64Part, + cacheEphemeral, + ) +import LLM.Core.Usage (Usage (..), mkUsage) import LLM.Core.Utils (getToolCalls, hasToolCalls) -import LLM.Providers.Gemini (parseGeminiResponse, parseGeminiUsage) +import LLM.Providers.Gemini (encodeTurn, parseGeminiResponse, parseGeminiUsage, signatureForModel) import Test.Hspec ( Spec, describe, @@ -37,7 +53,157 @@ tc.tcName `shouldBe` "get_weather" Left err -> expectationFailure $ "Parse failed: " <> show err + describe "thought signature round-trip" $ do + it "replays signature only for a matching model" $ do + let tc = + (mkToolCall "call_1" "get_weather" (object ["location" .= ("London" :: Text)])) + { tcProviderMeta = + Just + ProviderOpaque + { poProvider = "gemini", + poModel = Just "gemini-2.5-pro-001", + poPayload = object ["thoughtSignature" .= ("sig-xyz" :: Text)] + } + } + turn = AssistantMessage [toolCallPart tc] + matching = head (encodeTurn "gemini-2.5-pro" turn) + mismatch = head (encodeTurn "gemini-2.5-flash" turn) + hasThoughtSignature matching `shouldBe` True + hasThoughtSignature mismatch `shouldBe` False + signatureForModel "gemini-2.5-pro" tc.tcProviderMeta `shouldBe` Just "sig-xyz" + signatureForModel "gemini-2.5-flash" tc.tcProviderMeta `shouldBe` Nothing + + it "omits foreign Claude thinking opaque while keeping text and tool calls" $ do + let turn = + AssistantMessage + [ thinkingPart + ThinkingContent + { thinkingText = Just "claude thoughts", + thinkingOpaque = + Just + ProviderOpaque + { poProvider = "claude", + poModel = Just "claude-haiku-4-5-20251001", + poPayload = + object + [ "type" .= ("thinking" :: Text), + "signature" .= ("sig" :: Text) + ] + } + }, + textPart "hello", + toolCallPart (mkToolCall "c1" "get_weather" (object [])) + ] + encoded = head (encodeTurn "gemini-2.5-flash" turn) + parts = assistantParts encoded + -- Foreign opaque dropped; plain thinking text may remain as thought part, + -- plus text and functionCall => at least text + functionCall. + length parts `shouldBe` 3 + hasFunctionCall encoded `shouldBe` True + hasThoughtSignature encoded `shouldBe` False + + describe "image request encoding" $ do + it "encodes URL fileData and base64 inlineData parts" $ do + Right b64 <- pure $ imageBase64Part "image/jpeg" "aGVsbG8=" + let turn = + UserMessage + [ imageUrlPart "https://example.com/photo.png", + b64, + textPart "caption" + ] + parts = assistantParts (head (encodeTurn "gemini-3.1-flash-lite" turn)) + length parts `shouldBe` 3 + case parts of + [Object urlP, Object b64P, Object textP] -> do + KM.member "fileData" urlP `shouldBe` True + KM.member "inlineData" b64P `shouldBe` True + KM.lookup "text" textP `shouldBe` Just (String "caption") + case KM.lookup "fileData" urlP of + Just (Object fd) -> do + KM.lookup "fileUri" fd + `shouldBe` Just (String "https://example.com/photo.png") + KM.lookup "mimeType" fd `shouldBe` Just (String "image/png") + _ -> expectationFailure "expected fileData" + case KM.lookup "inlineData" b64P of + Just (Object idata) -> do + KM.lookup "mimeType" idata `shouldBe` Just (String "image/jpeg") + KM.lookup "data" idata `shouldBe` Just (String "aGVsbG8=") + _ -> expectationFailure "expected inlineData" + _ -> expectationFailure "expected three parts" + + describe "cache hint encoding" $ do + it "omits cache hints from the wire while keeping content identical" $ do + Right b64 <- pure $ imageBase64Part "image/jpeg" "aGVsbG8=" + let plain = + UserMessage + [ imageUrlPart "https://example.com/photo.png", + b64, + textPart "caption" + ] + hinted = + UserMessage + [ cacheEphemeral (imageUrlPart "https://example.com/photo.png"), + b64, + cacheEphemeral (textPart "caption") + ] + plainMsg = head (encodeTurn "gemini-3.1-flash-lite" plain) + hintedMsg = head (encodeTurn "gemini-3.1-flash-lite" hinted) + hintedMsg `shouldBe` plainMsg + describe "parseGeminiUsage" $ do it "extracts token counts" $ do Right val <- eitherDecodeFileStrict' "test/fixtures/gemini-text.json" - parseGeminiUsage val `shouldBe` Just (Usage 20 8 0) + parseGeminiUsage val `shouldBe` Just (mkUsage 20 8) + + it "reports cachedContentTokenCount without double-counting promptTokenCount" $ do + let val = + object + [ "usageMetadata" + .= object + [ "promptTokenCount" .= (1000 :: Int), + "candidatesTokenCount" .= (50 :: Int), + "cachedContentTokenCount" .= (800 :: Int) + ] + ] + parseGeminiUsage val + `shouldBe` Just + Usage + { usageInputTokens = 1000, + usageOutputTokens = 50, + usageCacheReadTokens = 800, + usageCacheCreationTokens = 0, + usageTotalCost = 0 + } + +hasThoughtSignature :: Value -> Bool +hasThoughtSignature (Object o) = + case KM.lookup "parts" o of + Just (Array arr) -> + any + ( \case + Object p -> KM.member "thoughtSignature" p + _ -> False + ) + arr + _ -> False +hasThoughtSignature _ = False + +hasFunctionCall :: Value -> Bool +hasFunctionCall (Object o) = + case KM.lookup "parts" o of + Just (Array arr) -> + any + ( \case + Object p -> KM.member "functionCall" p + _ -> False + ) + arr + _ -> False +hasFunctionCall _ = False + +assistantParts :: Value -> [Value] +assistantParts (Object o) = + case KM.lookup "parts" o of + Just (Array arr) -> V.toList arr + _ -> [] +assistantParts _ = [] diff --git a/test/LLM/GenerateObjectSpec.hs b/test/LLM/GenerateObjectSpec.hs --- a/test/LLM/GenerateObjectSpec.hs +++ b/test/LLM/GenerateObjectSpec.hs @@ -10,12 +10,13 @@ import LLM.Agent.GenerateObject (generateObject, generateObjectUntyped) import LLM.Agent.Types (Agent (..), RuntimeArgs (..)) import LLM.Core.Abort (AbortSignal, abort, newAbortSignal) -import LLM.Core.Types (ChatRequest (..), ChatResponse (..), LLMError (..), LLMGateway (..), LLMHooks (..), Turn (..)) -import LLM.Core.Usage (PricingInfo (..), Usage (..)) +import LLM.Core.Types (ChatRequest (..), ChatResponse (..), LLMError (..), LLMGateway (..), LLMHooks (..), assistantTurn, pattern UserTurn) +import LLM.Core.Usage (PricingInfo (..), Usage (..), defaultPricingInfo, mkUsage) import LLM.Generate.Logger (noHooks) import LLM.Generate.ModelConfig ( ModelConfig (..), ModelWithFallbacks (..), + defaultModelCapabilities, ) import LLM.Generate.Types (GenerateError (..), GenerateErrorResult (..)) import LLM.WeatherTool (WeatherToolArgs (..)) @@ -32,13 +33,13 @@ spec = describe "GenerateObject" $ do describe "generateObjectUntyped" $ do it "returns the provider JSON and usage" $ do - let gw = objectGateway (object ["location" .= ("Paris" :: Text)]) (Usage 12 3 0) + let gw = objectGateway (object ["location" .= ("Paris" :: Text)]) (mkUsage 12 3) models = ModelWithFallbacks (mockModel gw) [] result <- runUntyped models (object ["type" .= ("object" :: Text)]) case result of Right (value, usage) -> do value `shouldBe` object ["location" .= ("Paris" :: Text)] - usage `shouldBe` Usage 12 3 0 + usage `shouldBe` mkUsage 12 3 Left err -> expectationFailure $ show err it "propagates provider errors" $ do @@ -52,7 +53,7 @@ it "returns Aborted when the signal is already set" $ do sig <- newAbortSignal abort sig - let gw = objectGateway (object []) (Usage 0 0 0) + let gw = objectGateway (object []) (mkUsage 0 0) models = ModelWithFallbacks (mockModel gw) [] rt <- mkRuntime (Just sig) result <- generateObjectUntyped defaultAgent models rt [UserTurn "go"] (object ["type" .= ("object" :: Text)]) @@ -62,7 +63,7 @@ it "falls back to the next model on failure" $ do let failGw = errorObjectGateway (HttpError 500 "boom") - okGw = objectGateway (object ["ok" .= True]) (Usage 1 1 0) + okGw = objectGateway (object ["ok" .= True]) (mkUsage 1 1) models = ModelWithFallbacks (mockModel failGw) [mockModel okGw] result <- runUntyped models (object ["type" .= ("object" :: Text)]) case result of @@ -71,17 +72,17 @@ describe "generateObject" $ do it "decodes a typed object from provider JSON" $ do - let gw = objectGateway (object ["location" .= ("London" :: Text)]) (Usage 4 2 0) + let gw = objectGateway (object ["location" .= ("London" :: Text)]) (mkUsage 4 2) models = ModelWithFallbacks (mockModel gw) [] result <- runTyped models case result of Right (WeatherToolArgs loc, usage) -> do loc `shouldBe` "London" - usage `shouldBe` Usage 4 2 0 + usage `shouldBe` mkUsage 4 2 Left err -> expectationFailure $ show err it "reports a parse error when provider JSON does not match the codec" $ do - let gw = objectGateway (object ["wrong" .= (1 :: Int)]) (Usage 0 0 0) + let gw = objectGateway (object ["wrong" .= (1 :: Int)]) (mkUsage 0 0) models = ModelWithFallbacks (mockModel gw) [] result <- runTyped models case result of @@ -91,12 +92,12 @@ it "does not advertise tools even when agContextWindow would inject get_history" $ do captured <- newIORef Nothing - let gw = capturingObjectGateway captured (object ["location" .= ("Paris" :: Text)]) (Usage 1 0 0) + let gw = capturingObjectGateway captured (object ["location" .= ("Paris" :: Text)]) (mkUsage 1 0) models = ModelWithFallbacks (mockModel gw) [] agent = defaultAgent {agContextWindow = Just 1} conv = [ UserTurn "first", - AssistantTurn "a1" Nothing [], + assistantTurn "a1" Nothing [], UserTurn "second" ] rt <- mkRuntime Nothing @@ -145,6 +146,7 @@ mcMaxTokens = 256, mcTemperature = Nothing, mcThinking = Nothing, + mcCapabilities = defaultModelCapabilities, mcRequestTimeout = Nothing, mcThrottleDelay = Nothing, mcRetryCount = 0, @@ -152,7 +154,7 @@ } zeroPricing :: PricingInfo -zeroPricing = PricingInfo 0 0 +zeroPricing = defaultPricingInfo 0 0 defaultAgent :: Agent defaultAgent = diff --git a/test/LLM/GenericConversationTest.hs b/test/LLM/GenericConversationTest.hs --- a/test/LLM/GenericConversationTest.hs +++ b/test/LLM/GenericConversationTest.hs @@ -9,9 +9,14 @@ import LLM.Agent.Types (Agent (..), RuntimeArgs (..)) import LLM.Core.LLMProvider (LLMProvider, toGateway) import LLM.Core.Types (LLMGateway, ThinkingMode (..)) -import LLM.Core.Usage (PricingInfo (..)) +import LLM.Core.Usage (defaultPricingInfo) import LLM.Generate.Logger (noHooks) -import LLM.Generate.ModelConfig (ModelConfig (..), ModelWithFallbacks (ModelWithFallbacks)) +import LLM.Generate.ModelConfig + ( ModelCapabilities (..), + ModelConfig (..), + ModelWithFallbacks (ModelWithFallbacks), + defaultModelCapabilities, + ) import LLM.TestKit ( loadRecordedConversation, mockProvider, @@ -35,10 +40,14 @@ ModelConfig { mcGateway = provider, mcModel = T.pack opts.modelName, - mcPricing = PricingInfo {pricePerMillionInput = 0.0, pricePerMillionOutput = 0.0}, + mcPricing = defaultPricingInfo 0.0 0.0, mcMaxTokens = 1024, mcTemperature = Nothing, mcThinking = opts.specThinking, + mcCapabilities = + defaultModelCapabilities + { capThinking = maybe False (.tmEnabled) opts.specThinking + }, mcRequestTimeout = Nothing, mcThrottleDelay = Nothing, mcRetryCount = 3, diff --git a/test/LLM/HistoryToolSpec.hs b/test/LLM/HistoryToolSpec.hs --- a/test/LLM/HistoryToolSpec.hs +++ b/test/LLM/HistoryToolSpec.hs @@ -9,7 +9,7 @@ import LLM.Agent.ToolUtils (toTool, windowOffset) import LLM.Agent.Tools.HistoryTool (historyToolTyped) import LLM.Agent.Types (Agent (..), RuntimeArgs (..), Tool (..), ToolContext (..)) -import LLM.Core.Types (Turn (..)) +import LLM.Core.Types (Turn (..), assistantTurn, pattern UserTurn) import LLM.Generate.GenerateUtils (llmHooks) import LLM.Generate.Logger (noHooks) import Test.Hspec @@ -58,8 +58,8 @@ it "returns the full hidden prefix when the visible window has no user turns" $ do let conv = [ UserTurn "hidden question", - AssistantTurn "hidden answer" Nothing [], - AssistantTurn "visible assistant only" Nothing [] + assistantTurn "hidden answer" Nothing [], + assistantTurn "visible assistant only" Nothing [] ] -- Offset past both user+assistant hidden turns; visible slice is -- assistant-only so countUserTurns == 0. @@ -71,11 +71,11 @@ sampleConversation :: [Turn] sampleConversation = [ UserTurn "first question", - AssistantTurn "first answer" Nothing [], + assistantTurn "first answer" Nothing [], UserTurn "second question", - AssistantTurn "second answer" Nothing [], + assistantTurn "second answer" Nothing [], UserTurn "third question", - AssistantTurn "third answer" Nothing [] + assistantTurn "third answer" Nothing [] ] noWindowAgent :: Agent diff --git a/test/LLM/LoadSpec.hs b/test/LLM/LoadSpec.hs --- a/test/LLM/LoadSpec.hs +++ b/test/LLM/LoadSpec.hs @@ -3,14 +3,16 @@ import Control.Exception (IOException, SomeException, catch, displayException, fromException, try) import Control.Monad.Except (runExceptT) import Data.Aeson (Value (Array), eitherDecodeFileStrict') +import Data.List (find) import Data.Map qualified as Map import Data.Maybe (isJust) import Data.Text (Text) import Data.Text qualified as T import LLM.Core.Types (LLMGateway (..)) -import LLM.Generate.ModelConfig (ModelConfig (..)) +import LLM.Generate.ModelConfig (ModelCapabilities (..), ModelConfig (..), defaultModelCapabilities) import LLM.Load.LoadGateways (buildGateway, loadGateways, loadGatewaysFromCatalog, parseProviderBaseUrl) import LLM.Load.LoadModels (loadModelOrThrow, loadModelOrThrow_, loadModelsOrThrow) +import LLM.Load.ModelCatalog (ModelCatalogItem (..)) import LLM.Load.ProviderCatalog ( ProviderCatalogItem (..), ProviderProtocol (..), @@ -63,6 +65,26 @@ it "loads a known model config" $ do cfg <- loadModelOrThrow ollamaCatalog "llama_3_2" cfg.mcModel `shouldBe` "llama3.2:latest" + cfg.mcCapabilities.capVision `shouldBe` False + + it "loads catalog capabilities when present" $ do + result <- eitherDecodeFileStrict' "./model-catalog.json" + case result of + Right (items :: [ModelCatalogItem]) -> do + let Just haiku = find (\i -> i.modelConfigName == "haiku_4_5") items + haiku.capabilities + `shouldBe` Just + ( ModelCapabilities + { capThinking = True, + capVision = True, + capPromptCaching = True + } + ) + Left err -> expectationFailure err + + it "loads catalogs without capabilities (defaults false)" $ do + cfg <- loadModelOrThrow ollamaCatalog "mistral" + cfg.mcCapabilities `shouldBe` defaultModelCapabilities it "throws when a model config is missing" $ do result <- try @LoadConfigError (loadModelsOrThrow ollamaCatalog ("llama_3_2" :: Text, "missing_model" :: Text)) diff --git a/test/LLM/OpenAISpec.hs b/test/LLM/OpenAISpec.hs --- a/test/LLM/OpenAISpec.hs +++ b/test/LLM/OpenAISpec.hs @@ -2,20 +2,37 @@ module LLM.OpenAISpec (spec) where -import Data.Aeson (eitherDecodeFileStrict') +import Data.Aeson + ( Value (Array, Object, String), + eitherDecodeFileStrict', + object, + (.=), + ) +import Data.Aeson.Key qualified as K +import Data.Aeson.KeyMap qualified as KM +import Data.Maybe (mapMaybe) +import Data.Text (Text) +import Data.Vector qualified as V import LLM.Core.Types ( ChatResponse (respText), ToolCall (tcId, tcName), + Turn (..), + cacheEphemeral, + defaultMessageEncodeOptions, + imageBase64Part, + imageUrlPart, + textPart, ) -import LLM.Core.Usage (Usage (Usage)) +import LLM.Core.Usage (Usage (..), mkUsage) import LLM.Core.Utils (getToolCalls, hasToolCalls) -import LLM.Providers.OpenAI (parseOpenAIResponse, parseOpenAIUsage) +import LLM.Providers.OpenAI (encodeTurn, parseOpenAIResponse, parseOpenAIUsage) import Test.Hspec ( Spec, describe, expectationFailure, it, shouldBe, + shouldSatisfy, ) spec :: Spec @@ -39,7 +56,130 @@ tc.tcId `shouldBe` "call_abc123" Left err -> expectationFailure $ "Parse failed: " <> show err + describe "image request encoding" $ do + it "encodes URL and base64 data-URL image parts" $ do + Right b64 <- pure $ imageBase64Part "image/png" "aGVsbG8=" + let turn = + UserMessage + [ textPart "What is in this image?", + imageUrlPart "https://example.com/boardwalk.jpg", + b64 + ] + content = userContent (head (encodeTurn defaultMessageEncodeOptions turn)) + contentTypes content `shouldBe` ["text", "image_url", "image_url"] + case content of + [Object textO, Object urlO, Object b64O] -> do + lookupText "text" (Object textO) `shouldBe` Just "What is in this image?" + nestedText ["image_url", "url"] (Object urlO) + `shouldBe` Just "https://example.com/boardwalk.jpg" + nestedText ["image_url", "url"] (Object b64O) + `shouldBe` Just "data:image/png;base64,aGVsbG8=" + _ -> expectationFailure "expected text + two image_url parts" + + describe "cache hint encoding" $ do + it "omits cache hints from the wire while keeping content identical" $ do + Right b64 <- pure $ imageBase64Part "image/png" "aGVsbG8=" + let plain = + UserMessage + [ textPart "What is in this image?", + imageUrlPart "https://example.com/boardwalk.jpg", + b64 + ] + hinted = + UserMessage + [ cacheEphemeral (textPart "What is in this image?"), + imageUrlPart "https://example.com/boardwalk.jpg", + cacheEphemeral b64 + ] + plainMsg = head (encodeTurn defaultMessageEncodeOptions plain) + hintedMsg = head (encodeTurn defaultMessageEncodeOptions hinted) + hintedMsg `shouldBe` plainMsg + case userContent hintedMsg of + parts -> do + contentTypes parts `shouldBe` ["text", "image_url", "image_url"] + mapM_ (\p -> lookupKey "cache_control" p `shouldBe` Nothing) parts + + it "Ollama (shared OpenAI encoder) also omits cache hints" $ do + -- Ollama requests use the same encodeTurn + defaultMessageEncodeOptions path. + let plain = UserMessage [textPart "hello"] + hinted = UserMessage [cacheEphemeral (textPart "hello")] + head (encodeTurn defaultMessageEncodeOptions hinted) + `shouldBe` head (encodeTurn defaultMessageEncodeOptions plain) + describe "parseOpenAIUsage" $ do it "extracts token counts" $ do Right val <- eitherDecodeFileStrict' "test/fixtures/openai-text.json" - parseOpenAIUsage val `shouldBe` Just (Usage 15 9 0) + parseOpenAIUsage val `shouldBe` Just (mkUsage 15 9) + + it "reports cached_tokens without double-counting prompt_tokens" $ do + let val = + object + [ "usage" + .= object + [ "prompt_tokens" .= (2006 :: Int), + "completion_tokens" .= (300 :: Int), + "prompt_tokens_details" + .= object + [ "cached_tokens" .= (1920 :: Int), + "cache_write_tokens" .= (40 :: Int) + ] + ] + ] + parseOpenAIUsage val + `shouldBe` Just + Usage + { usageInputTokens = 2006, + usageOutputTokens = 300, + usageCacheReadTokens = 1920, + usageCacheCreationTokens = 40, + usageTotalCost = 0 + } + + it "reads zero cached_tokens from recorded conversation fixtures" $ do + Right val <- eitherDecodeFileStrict' "test/fixtures/openai-conversation-generated.json" + case val of + Array arr -> + case V.toList arr of + (Object first : _) -> + case KM.lookup "response" first of + Just resp -> + case parseOpenAIUsage resp of + Just u -> do + u.usageCacheReadTokens `shouldBe` 0 + u.usageInputTokens `shouldSatisfy` (> 0) + Nothing -> expectationFailure "failed to parse usage" + _ -> expectationFailure "missing response" + _ -> expectationFailure "empty conversation array" + _ -> expectationFailure "expected conversation array" + +userContent :: Value -> [Value] +userContent (Object o) = + case KM.lookup "content" o of + Just (Array arr) -> V.toList arr + _ -> [] +userContent _ = [] + +contentTypes :: [Value] -> [Text] +contentTypes = mapMaybe typ + where + typ (Object o) = case KM.lookup "type" o of + Just (String t) -> Just t + _ -> Nothing + typ _ = Nothing + +lookupText :: Text -> Value -> Maybe Text +lookupText key (Object o) = + case KM.lookup (K.fromText key) o of + Just (String t) -> Just t + _ -> Nothing +lookupText _ _ = Nothing + +lookupKey :: Text -> Value -> Maybe Value +lookupKey key (Object o) = KM.lookup (K.fromText key) o +lookupKey _ _ = Nothing + +nestedText :: [Text] -> Value -> Maybe Text +nestedText [key] v = lookupText key v +nestedText (key : rest) (Object o) = + KM.lookup (K.fromText key) o >>= nestedText rest +nestedText _ _ = Nothing diff --git a/test/LLM/StreamingSpec.hs b/test/LLM/StreamingSpec.hs --- a/test/LLM/StreamingSpec.hs +++ b/test/LLM/StreamingSpec.hs @@ -7,17 +7,20 @@ import Data.Text (Text) import LLM.Core.Types ( ChatResponse (..), - ContentBlock (TextBlock, ToolCallBlock), LLMGateway (..), StreamEvent (..), Turn (..), + assistantTurn, mkToolCall, + pattern UserTurn, + textPart, + toolCallPart, ) -import LLM.Core.Usage (PricingInfo (..), Usage (..)) +import LLM.Core.Usage (PricingInfo (..), defaultPricingInfo, mkUsage) import LLM.Generate.Generate (streamTextLLM) import LLM.Generate.GenerateUtils (llmHooks) import LLM.Generate.Logger (noHooks) -import LLM.Generate.ModelConfig (ModelConfig (..)) +import LLM.Generate.ModelConfig (ModelConfig (..), defaultModelCapabilities) import LLM.Generate.Types ( GenRequest (..), RoundTextRole (..), @@ -32,7 +35,7 @@ let gw = streamGateway [StreamDelta "hello"] - (ChatResponse "hello" [TextBlock "hello"] (Just (Usage 1 1 0)) Nothing) + (ChatResponse "hello" [textPart "hello"] (Just (mkUsage 1 1)) Nothing) chunks <- runStream gw [] reverse chunks `shouldBe` [ TextDelta "hello", @@ -40,11 +43,11 @@ ] it "routes text to answer after a tool turn" $ do - let prior = [UserTurn "q", AssistantTurn "" Nothing [mkToolCall "1" "t" (object [])], ToolTurn []] + let prior = [UserTurn "q", assistantTurn "" Nothing [mkToolCall "1" "t" (object [])], ToolTurn []] gw = streamGateway [StreamDelta "done"] - (ChatResponse "done" [TextBlock "done"] (Just (Usage 1 1 0)) Nothing) + (ChatResponse "done" [textPart "done"] (Just (mkUsage 1 1)) Nothing) chunks <- runStream gw prior reverse chunks `shouldBe` [AnswerDelta "done"] @@ -53,7 +56,7 @@ gw = streamGateway [StreamDelta "searching", StreamToolCall tc] - (ChatResponse "" [ToolCallBlock tc] (Just (Usage 2 0 0)) Nothing) + (ChatResponse "" [toolCallPart tc] (Just (mkUsage 2 0)) Nothing) chunks <- runStream gw [UserTurn "find x"] reverse chunks `shouldBe` [ TextDelta "searching", @@ -65,7 +68,7 @@ let gw = streamGateway [StreamReasoningDelta "think"] - (ChatResponse "ok" [TextBlock "ok"] Nothing Nothing) + (ChatResponse "ok" [textPart "ok"] Nothing Nothing) chunks <- runStream gw [] reverse chunks `shouldBe` [ReasoningDelta "think", RoundTextRoleCommitted AnswerRole] @@ -99,6 +102,7 @@ mcMaxTokens = 256, mcTemperature = Nothing, mcThinking = Nothing, + mcCapabilities = defaultModelCapabilities, mcRequestTimeout = Nothing, mcThrottleDelay = Nothing, mcRetryCount = 0, @@ -108,4 +112,4 @@ readIORef ref zeroPricing :: PricingInfo -zeroPricing = PricingInfo 0 0 +zeroPricing = defaultPricingInfo 0 0 diff --git a/test/LLM/TestKit.hs b/test/LLM/TestKit.hs --- a/test/LLM/TestKit.hs +++ b/test/LLM/TestKit.hs @@ -32,12 +32,17 @@ import LLM.Agent.ToolUtils (toTool) import LLM.Agent.Types (Agent (..), RuntimeArgs (..), ToolMap) import LLM.Core.LLMProvider (LLMProvider (..)) -import LLM.Core.Types (LLMGateway, ThinkingMode (..), Turn (..)) -import LLM.Core.Usage (PricingInfo (..), addUsage, emptyUsage) +import LLM.Core.Types (LLMGateway, ThinkingMode (..), Turn (..), pattern UserTurn) +import LLM.Core.Usage (addUsage, defaultPricingInfo, emptyUsage) import LLM.Core.Utils (parseChatResponse) import LLM.Generate.GenerateUtils (llmHooks) import LLM.Generate.Logger (Hooks (..), noHooks) -import LLM.Generate.ModelConfig (ModelConfig (..), ModelWithFallbacks (..)) +import LLM.Generate.ModelConfig + ( ModelCapabilities (..), + ModelConfig (..), + ModelWithFallbacks (..), + defaultModelCapabilities, + ) import LLM.Generate.Types (GenerateErrorResult (..), GenerateTextResult (..)) import LLM.WeatherTool (weatherToolTyped) import System.Directory (createDirectoryIfMissing) @@ -150,10 +155,14 @@ ModelConfig { mcGateway = gateway, mcModel = modelName, - mcPricing = PricingInfo {pricePerMillionInput = 0.0, pricePerMillionOutput = 0.0}, + mcPricing = defaultPricingInfo 0.0 0.0, mcMaxTokens = 1024, mcTemperature = Nothing, mcThinking = thinking, + mcCapabilities = + defaultModelCapabilities + { capThinking = maybe False (.tmEnabled) thinking + }, mcRequestTimeout = Nothing, mcThrottleDelay = Nothing, mcRetryCount = 3, diff --git a/test/LLM/TypesSpec.hs b/test/LLM/TypesSpec.hs --- a/test/LLM/TypesSpec.hs +++ b/test/LLM/TypesSpec.hs @@ -1,25 +1,48 @@ module LLM.TypesSpec (spec) where -import Data.Aeson (object, (.=)) +import Data.Aeson (eitherDecode, encode, object, (.=)) +import Data.Text qualified as T import LLM.Core.Types - ( ChatResponse (ChatResponse), - ContentBlock (TextBlock, ToolCallBlock), + ( CacheHint (..), + ChatResponse (ChatResponse), + ContentPart (..), LLMError (EmptyResponse, HttpError, NetworkError), + PartBody (..), + ThinkingContent (..), + Turn (..), + cacheEphemeral, + imageBase64Part, + imageUrlPart, mkToolCall, + pattern UserTurn, + textPart, + thinkingPart, + toolCallPart, + validateTurn, ) import LLM.Core.Usage ( PricingInfo (..), Usage (..), addUsage, + defaultPricingInfo, emptyUsage, estimateCost, + mkUsage, + usageOrdinaryInputTokens, ) import LLM.Core.Utils ( getToolCalls, hasToolCalls, isRetryable, ) -import Test.Hspec (Spec, describe, it, shouldBe) +import Test.Hspec + ( Spec, + describe, + expectationFailure, + it, + shouldBe, + shouldSatisfy, + ) spec :: Spec spec = describe "Types" $ do @@ -27,39 +50,184 @@ it "emptyUsage has zero tokens" $ do emptyUsage.usageInputTokens `shouldBe` 0 emptyUsage.usageOutputTokens `shouldBe` 0 + emptyUsage.usageCacheReadTokens `shouldBe` 0 + emptyUsage.usageCacheCreationTokens `shouldBe` 0 - it "addUsage sums token counts" $ do - let u1 = Usage 10 20 0 - u2 = Usage 30 40 0 - addUsage u1 u2 `shouldBe` Usage 40 60 0 + it "addUsage sums token counts including cache counters" $ do + let u1 = + Usage + { usageInputTokens = 10, + usageOutputTokens = 20, + usageCacheReadTokens = 3, + usageCacheCreationTokens = 2, + usageTotalCost = 0.25 + } + u2 = + Usage + { usageInputTokens = 30, + usageOutputTokens = 40, + usageCacheReadTokens = 4, + usageCacheCreationTokens = 1, + usageTotalCost = 0.5 + } + addUsage u1 u2 + `shouldBe` Usage + { usageInputTokens = 40, + usageOutputTokens = 60, + usageCacheReadTokens = 7, + usageCacheCreationTokens = 3, + usageTotalCost = 0.75 + } it "addUsage is associative" $ do - let u1 = Usage 1 2 0 - u2 = Usage 3 4 0 - u3 = Usage 5 6 0 + let u1 = mkUsage 1 2 + u2 = mkUsage 3 4 + u3 = mkUsage 5 6 addUsage (addUsage u1 u2) u3 `shouldBe` addUsage u1 (addUsage u2 u3) + it "Semigroup matches addUsage and Monoid uses emptyUsage" $ do + let u1 = mkUsage 10 1 + u2 = mkUsage 5 2 + (u1 <> u2) `shouldBe` addUsage u1 u2 + (mempty :: Usage) `shouldBe` emptyUsage + + it "JSON round-trips with cache fields" $ do + let u = + Usage + { usageInputTokens = 100, + usageOutputTokens = 20, + usageCacheReadTokens = 40, + usageCacheCreationTokens = 10, + usageTotalCost = 0.5 + } + eitherDecode (encode u) `shouldBe` Right u + + it "JSON decode defaults missing cache fields to zero" $ do + let json = "{\"usageInputTokens\":12,\"usageOutputTokens\":3}" + eitherDecode json + `shouldBe` Right (mkUsage 12 3) + describe "estimateCost" $ do it "calculates cost in dollars from per-million pricing" $ do - let pricing = PricingInfo {pricePerMillionInput = 1.0, pricePerMillionOutput = 5.0} - usage = Usage 1_000_000 1_000_000 0 + let pricing = defaultPricingInfo 1.0 5.0 + usage = mkUsage 1_000_000 1_000_000 estimateCost pricing usage `shouldBe` 6.0 it "returns 0 for zero usage" $ do - let pricing = PricingInfo {pricePerMillionInput = 1.0, pricePerMillionOutput = 5.0} + let pricing = defaultPricingInfo 1.0 5.0 estimateCost pricing emptyUsage `shouldBe` 0.0 + it "prices cache read/write separately when rates are set" $ do + let pricing = + PricingInfo + { pricePerMillionInput = 3.0, + pricePerMillionOutput = 15.0, + pricePerMillionCacheRead = Just 0.3, + pricePerMillionCacheWrite = Just 3.75 + } + usage = + Usage + { usageInputTokens = 1_000_000, + usageOutputTokens = 0, + usageCacheReadTokens = 400_000, + usageCacheCreationTokens = 100_000, + usageTotalCost = 0 + } + -- ordinary 500k * 3 + read 400k * 0.3 + write 100k * 3.75 = 1.5 + 0.12 + 0.375 + estimateCost pricing usage `shouldBe` 1.995 + usageOrdinaryInputTokens usage `shouldBe` 500_000 + + it "falls back to input rate when cache rates are absent" $ do + let pricing = defaultPricingInfo 2.0 0.0 + usage = + Usage + { usageInputTokens = 1_000_000, + usageOutputTokens = 0, + usageCacheReadTokens = 250_000, + usageCacheCreationTokens = 250_000, + usageTotalCost = 0 + } + estimateCost pricing usage `shouldBe` 2.0 + + describe "PricingInfo JSON" $ do + it "decodes catalogs without cache rates" $ do + let json = "{\"pricePerMillionInput\":1.0,\"pricePerMillionOutput\":5.0}" + eitherDecode json `shouldBe` Right (defaultPricingInfo 1.0 5.0) + + it "round-trips optional cache rates" $ do + let pricing = + PricingInfo + { pricePerMillionInput = 3.0, + pricePerMillionOutput = 15.0, + pricePerMillionCacheRead = Just 0.3, + pricePerMillionCacheWrite = Just 3.75 + } + eitherDecode (encode pricing) `shouldBe` Right pricing + describe "hasToolCalls / getToolCalls" $ do it "returns False for text-only response" $ do - let resp = ChatResponse "hello" [TextBlock "hello"] Nothing Nothing + let resp = ChatResponse "hello" [textPart "hello"] Nothing Nothing hasToolCalls resp `shouldBe` False getToolCalls resp `shouldBe` [] it "returns True when tool calls present" $ do let tc = mkToolCall "id1" "get_weather" (object ["location" .= ("London" :: String)]) - resp = ChatResponse "" [ToolCallBlock tc] Nothing Nothing + resp = ChatResponse "" [toolCallPart tc] Nothing Nothing hasToolCalls resp `shouldBe` True getToolCalls resp `shouldBe` [tc] + + describe "cache hints" $ do + it "round-trips CacheEphemeral on ContentPart / Turn JSON" $ do + let part = cacheEphemeral (textPart "cached prefix") + turn = UserMessage [part, textPart "question"] + eitherDecode (encode part) + `shouldBe` Right (ContentPart (TextPart "cached prefix") (Just CacheEphemeral)) + eitherDecode (encode turn) `shouldBe` Right turn + + it "UserTurn does not match a cache-annotated single text part" $ do + let annotated = UserMessage [cacheEphemeral (textPart "hi")] + case annotated of + UserTurn _ -> expectationFailure "UserTurn should not match annotated text" + UserMessage [ContentPart (TextPart "hi") (Just CacheEphemeral)] -> pure () + other -> expectationFailure $ "unexpected: " <> show other + + describe "UserTurn / validateTurn" $ do + it "UserTurn constructs a single text user message" $ do + UserTurn "hi" `shouldBe` UserMessage [textPart "hi"] + + it "rejects thinking parts on user messages" $ do + validateTurn + ( UserMessage + [thinkingPart (ThinkingContent (Just "nope") Nothing)] + ) + `shouldBe` Left "user messages may not contain thinking parts" + + it "allows image parts on user messages" $ do + validateTurn + ( UserMessage + [imageUrlPart "https://example.com/a.png", textPart "look"] + ) + `shouldBe` Right () + + it "rejects image parts on assistant messages" $ do + validateTurn + (AssistantMessage [imageUrlPart "https://example.com/a.png"]) + `shouldBe` Left "assistant messages may not contain image parts" + + describe "imageBase64Part" $ do + it "accepts valid png base64" $ do + imageBase64Part "image/png" "aGVsbG8=" `shouldSatisfy` either (const False) (const True) + + it "rejects unsupported MIME types" $ do + case imageBase64Part "image/svg+xml" "aGVsbG8=" of + Left msg -> msg `shouldSatisfy` T.isInfixOf "unsupported image media type" + Right _ -> fail "expected MIME failure" + + it "rejects empty base64" $ do + imageBase64Part "image/png" " " `shouldBe` Left "image base64 data must not be empty" + + it "rejects invalid base64 characters" $ do + imageBase64Part "image/png" "!!!!" `shouldBe` Left "image data is not valid base64" describe "isRetryable" $ do it "retries on 429" $ do