diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,142 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.4.0.0] - 2026-10-24 + +### Fixed + +- Provider requests are no longer resent by req. req's default config + retried timeouts and 408/504 responses up to 5 times, resending (and + billing) the POST, before `withRetry` retried again. Its default 30 s + response timeout also failed any non-streaming call that took longer than + 30 s to generate. All providers now share one TLS manager without a + response timeout and a config that never retries; `mcRequestTimeout` is the + only response deadline. New `mkLenientConfig` (`LLM.Core.ProviderUtils`) + builds that config for a given manager. +- One missing API key no longer breaks loading of every model. + `loadModelOrThrow` / `loadModelsOrThrow` used to build every catalog entry + and fail on the first one whose provider had no key, with a misleading + "Provider not found". Now only the requested models must be available; for + one that is not, the error names the reason, e.g. "Model haiku_4_5 needs + provider claude, but env var CLAUDE_API_KEY is not set" (or the invalid + base URL). The returned `ModelConfigMap` / `GatewayMap` hold all currently + available models and gateways. "Model not found" errors now name the model. + New `resolveGateway`, `resolveGateways` and `ProviderUnavailable` + (`LLM.Load.LoadGateways`). +- File tools no longer corrupt non-UTF-8 files or depend on the locale. + `replace_in_file` and `multi_replace_in_file` decoded leniently and wrote + the whole file back, replacing every invalid byte with U+FFFD; they now + refuse files that are not valid UTF-8 and leave them unchanged. Writes + (`writefile` and both replace tools) always produce UTF-8; under `LANG=C` + they used to fail or mangle non-ASCII text. `grep` and `file_info` read + bytes and decode leniently, so `grep` no longer skips files it cannot + decode. `getSystemPrompt` reports a `file:` prompt that is not valid UTF-8 + as a `LoadSystemPromptError` naming the file. +- An abort during a tool round no longer leaves an unanswered tool call in + `gerPartialNewMessages`. The partial history used to end with the + assistant's tool calls and no `ToolTurn`, which Claude and OpenAI reject on + the next request. It now ends with a `ToolTurn` holding the real results of + the calls that ran and "Aborted before execution" for the rest. + **Breaking:** `executeToolsWithAbort` returns the new `ToolRoundOutcome` + (`ToolRoundCompleted` / `ToolRoundAborted`) instead of + `GenerateResult [ToolResult]`. +- Gemini structured output works with thinking enabled. The response's first + part was the thought summary, which the object parser tried to decode as + JSON (`EmptyResponse`). The parser now joins the text of all parts not + marked `"thought": true`, and object requests no longer ask for thought + summaries (`includeThoughts`). +- Ollama honours the max-token limit. Chat requests sent `num_predict` and + object requests `max_completion_tokens`; Ollama's OpenAI-compatible endpoint + ignores both, so `mcMaxTokens` had no effect. Both now send `max_tokens`. + Streamed Ollama responses now request usage (`stream_options.include_usage`) + as OpenAI's do. The recorded Ollama conversation fixtures were re-recorded. +- File tools handle symbolic links as documented. `sandboxPath` resolves + links, so the tools' `isSymlink` checks on its result never fired: + `remove_file` and `remove_directory` on a link deleted the link's target + inside the workspace, and `move_file` / `copy_file` moved or copied the + target. `remove_file` and `remove_directory` now remove the link itself + and leave its target unchanged (also for links pointing outside the + workspace); `move_file` and `copy_file` refuse a link source. New + `sandboxSymlink` (`LLM.Tools.FsConfig`). +- Claude and Gemini no longer send turns with empty content. After a + fallback, an assistant turn holding only foreign thinking (e.g. a DeepSeek + turn cut off during reasoning, or a Claude turn with only redacted + thinking) encoded to `"content": []` / `"parts": []`, which the API + rejects with a 400. Such turns, and user turns whose parts all encode to + nothing, are now left out; both APIs accept the consecutive same-role + turns this can leave. The OpenAI protocol was fixed in 0.3.0.0. +- Gemini thought signatures are only replayed to the model that produced + them. Model names were compared by prefix, so `gemini-2.5-flash` and + `gemini-2.5-flash-lite` counted as the same model and received each + other's signatures. Names now must be equal apart from a `models/` prefix + and a stable version suffix (`-001`); other calls get the documented + placeholder signature. +- Claude non-streaming responses with unknown content blocks (e.g. + `server_tool_use`, `web_search_tool_result`) no longer fail as + `EmptyResponse`. Unknown blocks are skipped, as streaming already did. +- Non-streaming HTTP errors keep their status and raw body. `HttpError` + bodies (and the `onLLMResponseError` text) were Haskell `show` output + (`Object (fromList [...])`); they are now the body as sent, as for + streaming. A non-JSON error response (e.g. an HTML 401/502 from a proxy) + made req throw, which became a retryable `NetworkError` without its + status; it is now `HttpError status body`. A 200 body that is not JSON is + a `ParseError`. Streaming error bodies that are not valid UTF-8 are + decoded leniently instead of throwing. **Breaking:** `sendRequest` and + `sendObjectRequest` in `LLMProvider` return the raw body + (`IO (Int, ByteString)`) instead of a decoded `Value`. +- HTTP 500, 502 and 504 are retried, like 429/503/529. They are usually + transient, and a Claude mid-stream `api_error` (mapped to 500) was never + retried. Each retry delay is now capped at 30 s; `fullJitterBackoff` grew + without bound for large `retryCount`s. `Retry-After` is still ignored; + the `mcJitterBackoff` docs and the README now say so. +- Reasoning is no longer cut to its first chunk. `respReasoning` (and + `projectReasoning`) returned only the first thinking part, so a streamed + Gemini response kept only its first thought chunk, and the OpenAI/DeepSeek + encoder replayed a truncated `reasoning_content`. Both now concatenate all + thinking text. `mkChatResponse` merges adjacent thinking parts that carry + no opaque payload; signed parts (e.g. Claude thinking blocks) stay one per + block. +- The generate layer honours `grAbortSignal`, as its docs promised. + `generateTextWithFallbacks` / `streamTextWithFallbacks` (and the + `Rotated` variants) never read it, so the `Aborted` path was dead. The + signal is now checked before each model in the fallback chain and before + each attempt (also after throttle and backoff sleeps), and a stream is cut + off at its next event; all yield `GErrAborted`. A non-streaming HTTP call + already in flight is not interrupted. **Breaking:** removed the unused + `LLMError` constructor `ToolLoopExceeded` (the agent loop reports + `GErrToolExceeded`). +- Documented what happens at `agMaxToolRounds`: the last round's tools run + and their results are in `gerPartialNewMessages`, but the model is not + called again, so it never sees them. +- Smaller file-tool gaps. `grep` on an explicit file path skipped the 1 MB + size cap that applies to files found by its directory walk; such a file is + now refused with an error. `read_file_paginated` threw on bytes that are + not valid UTF-8; it now shows them as U+FFFD, like the other read-only + tools. `find_files` backtracked exponentially on patterns with many `*`s + (e.g. `*a*a*a*a*a*a*a*a*b` took seconds on a 40-character name); matching + is now O(pattern × name), with the same results. Its walk no longer + recounts and reverses the match list at every entry. +- Gemini tool-call ids no longer repeat across processes. They were + `call_` from `Data.Unique`, which restarts at 1 in every process, so a + saved conversation reloaded in a new process got new calls with ids it + already held; after a fallback, Claude and OpenAI reject or mismatch + those. Ids are now `call_`. New direct dependency: + `random`. +- Gemini errors, blocks and finish reasons are no longer all + `EmptyResponse`. A stream that fails midway stays HTTP 200 and sends + `{"error": {"code": 503, ...}}` as a data event; that event was ignored, + so the call ended as a non-retryable `EmptyResponse` (or as a shortened + success). It is now `HttpError code body`, retried like the same error + before streaming, as for Claude and OpenAI. A response without content + names why Gemini stopped: a blocked prompt (`promptFeedback.blockReason`) + or a filter finish reason (`SAFETY`, `RECITATION`, `PROHIBITED_CONTENT`, + …) gives the new `ResponseBlocked`; `MAX_TOKENS` (common when thinking + uses the whole budget) and other finish reasons give a `ParseError` naming + the reason. A response with content is still returned when the token + limit cut it off. Structured output cut off by `MAX_TOKENS` is a + `ParseError` saying so. **Breaking:** new `LLMError` constructor + `ResponseBlocked`. + ## [0.3.0.0] - 2026-10-13 ### Added @@ -274,6 +410,7 @@ DeepSeek), single-shot generation with fallbacks, agent tool loops, structured output, JSON model catalog loading, and workspace-scoped filesystem tools. +[0.4.0.0]: https://github.com/aische/llm-simple/compare/v0.3.0.0...v0.4.0.0 [0.3.0.0]: https://github.com/aische/llm-simple/compare/v0.2.0.0...v0.3.0.0 [0.2.0.0]: https://github.com/aische/llm-simple/compare/v0.1.1.0...v0.2.0.0 [0.1.1.0]: https://github.com/aische/llm-simple/compare/v0.1.0.2...v0.1.1.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.3.x release. APIs may change. +**Status:** early 0.4.x release. APIs may change. ## Features @@ -34,7 +34,12 @@ DEEPSEEK_API_KEY=... ``` -Ollama needs no API key. The library reads keys from the **process environment**. +Ollama needs no API key. It is used through its OpenAI-compatible endpoint +(`/v1/chat/completions`); tested with Ollama 0.32.4. Streamed responses +request usage via `stream_options.include_usage`; an Ollama version that does +not support it streams without usage. + +The library reads keys from the **process environment**. Load `.env` yourself in `main` (the example executable does this), or call `loadGatewaysWithDotenv` when building gateways manually. @@ -66,8 +71,8 @@ | `temperature` | Sampling temperature (optional) | | `requestTimeout` | Timeout in ms per attempt (optional); for streaming, time until the first event; timeouts are retried | | `throttleDelay` | Delay between requests in ms (optional) | -| `retryCount` | Number of retries on failure | -| `jitterBackoff` | Backoff jitter in ms between retries | +| `retryCount` | Number of retries on retryable errors (429, 500, 502, 503, 504, 529, network errors, timeouts) | +| `jitterBackoff` | Base backoff in ms between retries; full jitter, exponential, each delay capped at 30 s. `Retry-After` is not honoured, so raise this for providers that often rate-limit you | | `thinking` | Extended thinking effort level (optional, provider-dependent) | | `capabilities` | Optional object: `thinking`, `vision`, `promptCaching` (default `false`) | 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.3.0.0 +version: 0.4.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.2.x release — APIs may change. + Status: pre-1.0 — APIs may change between major versions. license: BSD-3-Clause @@ -153,15 +153,17 @@ dotenv >=0.9 && <0.13, filepath >=1.4 && <1.6, http-client >=0.7 && <0.8, + http-client-tls >=0.3 && <0.5, http-types >=0.12 && <0.13, lens >=5.3 && <5.4, mtl >=2.3 && <2.4, + random >=1.2.1 && <1.4, req >=3.13.4 && <3.14, retry >=0.9 && <0.10, text >=1.2 && <2.2, time >=1.11 && <1.15, unordered-containers >=0.2 && < 0.3, - uuid-types >=1.0 && <1.1, + uuid-types >=1.0.4 && <1.1, vector >=0.13 && <0.14, executable llm-simple @@ -216,6 +218,7 @@ LLM.GeminiSpec LLM.GenericConversationTest LLM.RecordedConversationSpec + LLM.OllamaSpec LLM.OpenAISpec LLM.TypesSpec LLM.ChatSpec @@ -226,6 +229,8 @@ LLM.LoadSpec LLM.StreamingSpec LLM.HistoryToolSpec + LLM.HttpConfigSpec + LLM.LLMProviderSpec LLM.GenerateObjectSpec LLM.ToolUtilsSpec LLM.CrossProviderSpec @@ -245,9 +250,13 @@ bytestring >=0.11 && <0.13, hspec >=2.10 && <2.12, heptapod >=1.1 && <1.2, + http-client >=0.7 && <0.8, llm-simple, + network >=3.1 && <3.3, + req >=3.13.4 && <3.14, mtl >=2.3 && <2.4, retry >=0.9 && <0.10, scientific >=0.3 && <0.4, + uuid-types >=1.0.4 && <1.1, vector >=0.13 && <0.14 build-tool-depends: hspec-discover:hspec-discover 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 @@ -12,6 +12,7 @@ ( createGenRequest, createToolContext, executeToolsWithAbort, + ToolRoundOutcome (..), getResolvedTools, ) import LLM.Agent.Types @@ -26,7 +27,7 @@ Turn (..), ) import LLM.Core.Usage (Usage (..), emptyUsage) -import LLM.Core.Utils (getToolCalls) +import LLM.Core.Utils (getToolCalls, toolResult) import LLM.Generate.Generate ( generateTextWithFallbacksRotated, streamTextWithFallbacksRotated, @@ -148,11 +149,17 @@ toolResultsE <- executeToolsWithAbort rt.rtAbortSignal rt.rtHooks toolContext tools toolCalls case toolResultsE of - Left err -> do - let errResult = GenerateErrorResult err (newTurnsAcc ++ [assistantMsg]) newUsage - emitEvent rt (GenerationFailed err errResult) + ToolRoundAborted done notRun -> do + -- Answer every tool call so the history stays valid + -- for the next request: real results for the calls + -- that ran (they may have had side effects), a + -- placeholder for the rest. + let toolTurn = ToolTurn (done ++ map (`toolResult` "Aborted before execution") notRun) + errResult = GenerateErrorResult GErrAborted (newTurnsAcc ++ [assistantMsg, toolTurn]) newUsage + emitEvent rt (MessageCreated toolTurn) + emitEvent rt (GenerationFailed GErrAborted errResult) pure $ Left errResult - Right toolResults -> do + ToolRoundCompleted toolResults -> do let toolTurn = ToolTurn toolResults emitEvent rt (MessageCreated toolTurn) emitEvent rt (ToolRoundFinished loopCount) 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 @@ -2,6 +2,7 @@ ( executeTool, executeTools, executeToolsWithAbort, + ToolRoundOutcome (..), createToolContext, getSchema, toTool, @@ -41,11 +42,7 @@ import LLM.Core.Usage (Usage) import LLM.Core.Utils (toolResult, trySync) import LLM.Generate.Logger (Hooks (..)) -import LLM.Generate.Types - ( GenRequest (..), - GenerateError (..), - GenerateResult, - ) +import LLM.Generate.Types (GenRequest (..)) import LLM.Tools.FsConfig (SandboxViolation, formatSandboxViolation) -- | Execute a single tool call by looking it up in the tool list @@ -76,17 +73,25 @@ executeTools :: Hooks -> ToolContext -> [Tool Text] -> [ToolCall] -> IO [ToolResult] executeTools hooks ctx tools = mapM (executeTool hooks ctx tools) --- | Execute tool calls one at a time, checking the abort signal between each. --- Returns @Left Aborted@ if the signal fires before all calls finish. -executeToolsWithAbort :: Maybe AbortSignal -> Hooks -> ToolContext -> [Tool Text] -> [ToolCall] -> IO (GenerateResult [ToolResult]) -executeToolsWithAbort Nothing hooks ctx tools tcs = Right <$> executeTools hooks ctx tools tcs +-- | Outcome of a tool round run by 'executeToolsWithAbort'. +data ToolRoundOutcome + = -- | Every call ran. + ToolRoundCompleted [ToolResult] + | -- | The abort signal fired: results of the calls that ran, and the calls + -- that did not. + ToolRoundAborted [ToolResult] [ToolCall] + deriving (Show, Eq) + +-- | Execute tool calls one at a time, checking the abort signal before each. +executeToolsWithAbort :: Maybe AbortSignal -> Hooks -> ToolContext -> [Tool Text] -> [ToolCall] -> IO ToolRoundOutcome +executeToolsWithAbort Nothing hooks ctx tools tcs = ToolRoundCompleted <$> executeTools hooks ctx tools tcs executeToolsWithAbort (Just sig) hooks ctx tools tcs = go [] tcs where - go acc [] = pure (Right (reverse acc)) - go acc (tc : rest) = do + go acc [] = pure (ToolRoundCompleted (reverse acc)) + go acc remaining@(tc : rest) = do aborted <- isAborted sig if aborted - then pure (Left GErrAborted) + then pure (ToolRoundAborted (reverse acc) remaining) else do r <- executeTool hooks ctx tools tc go (r : acc) rest diff --git a/src/LLM/Agent/Types.hs b/src/LLM/Agent/Types.hs --- a/src/LLM/Agent/Types.hs +++ b/src/LLM/Agent/Types.hs @@ -34,6 +34,9 @@ -- Only these names are exposed; the map may contain additional tools. agTools :: [Text], -- | Maximum number of model→tool rounds before returning 'GErrToolExceeded'. + -- The tools of the last round run, and their results are in + -- 'gerPartialNewMessages', but no further model call is made, so the + -- model never sees them. agMaxToolRounds :: Int, -- | When 'Just', only the last @n@ user messages (and their follow-on -- turns) are sent to the provider. Older turns remain in 'ToolContext' and diff --git a/src/LLM/Core/LLMProvider.hs b/src/LLM/Core/LLMProvider.hs --- a/src/LLM/Core/LLMProvider.hs +++ b/src/LLM/Core/LLMProvider.hs @@ -7,16 +7,19 @@ where import Control.Exception (try) -import Data.Aeson (Value) +import Data.Aeson (Value, eitherDecodeStrict') import Data.ByteString (ByteString) import Data.Text (Text) import Data.Text qualified as T +import Data.Text.Encoding (decodeUtf8With) +import Data.Text.Encoding.Error (lenientDecode) import LLM.Core.Types ( ChatRequest (reqTools), - LLMError (HttpError, NetworkError), + LLMError (HttpError, NetworkError, ParseError), LLMGateway (..), LLMHooks (..), LLMObjectResult, + LLMResult, LLMTextResult, StreamEvent, ) @@ -32,8 +35,9 @@ providerImageUrlSupport :: Bool, -- | Build the JSON request body. Bool indicates whether streaming is requested. buildBody :: Bool -> ChatRequest -> Value, - -- | Make a non-streaming HTTP call, returning (status code, response JSON). - sendRequest :: Value -> IO (Int, Value), + -- | Make a non-streaming HTTP call, returning (status code, raw response + -- body). Must not throw on a non-2xx status. + sendRequest :: Value -> IO (Int, ByteString), -- | Make a streaming HTTP call. On a 200 response the continuation receives -- a reader over the raw body (empty chunk = EOF); other statuses become an -- 'HttpError' without calling it. @@ -45,8 +49,9 @@ parseResponse :: Value -> IO LLMTextResult, -- | Build the JSON request body for object generation. buildObjectBody :: ChatRequest -> Value -> Value, - -- | Make a non-streaming HTTP call for object generation, returning (status code, response JSON). - sendObjectRequest :: Value -> IO (Int, Value), + -- | Make a non-streaming HTTP call for object generation, returning + -- (status code, raw response body). + sendObjectRequest :: Value -> IO (Int, ByteString), -- | Parse a complete JSON response body for object generation. parseObjectResponse :: Value -> IO LLMObjectResult } @@ -80,13 +85,7 @@ Left e -> do hooks.onLLMResponseError p.providerName (T.pack (show (e :: HttpException))) pure $ Left $ NetworkError (T.pack (show (e :: HttpException))) - Right (status, respBody) -> do - hooks.onLLMResponse p.providerName respBody - if status == 200 - then p.parseResponse respBody - else do - hooks.onLLMResponseError p.providerName ("HTTP error: " <> T.pack (show status) <> " " <> T.pack (show respBody)) - pure $ Left $ HttpError status (T.pack $ show respBody) + Right resp -> handleResponse p hooks p.parseResponse resp -- | Generic object generation via the LLMProvider. genericGenerateObject :: LLMProvider -> LLMHooks -> Value -> ChatRequest -> IO LLMObjectResult @@ -98,13 +97,27 @@ Left e -> do hooks.onLLMResponseError p.providerName (T.pack (show (e :: HttpException))) pure $ Left $ NetworkError (T.pack (show (e :: HttpException))) - Right (status, respBody) -> do - hooks.onLLMResponse p.providerName respBody - if status == 200 - then p.parseObjectResponse respBody - else do - hooks.onLLMResponseError p.providerName ("HTTP error: " <> T.pack (show status) <> " " <> T.pack (show respBody)) - pure $ Left $ HttpError status (T.pack $ show respBody) + Right resp -> handleResponse p hooks p.parseObjectResponse resp + +-- | Turn a non-streaming (status, raw body) into a result. A non-200 status +-- is an 'HttpError' carrying the body as text, whether or not it is JSON +-- (an HTML 502 from a proxy keeps its status). A 200 body that is not JSON +-- is a 'ParseError'. 'onLLMResponse' sees every body that decodes as JSON. +handleResponse :: LLMProvider -> LLMHooks -> (Value -> IO (LLMResult a)) -> (Int, ByteString) -> IO (LLMResult a) +handleResponse p hooks parse (status, raw) = do + let decoded = eitherDecodeStrict' raw + either (const (pure ())) (hooks.onLLMResponse p.providerName) decoded + if status /= 200 + then do + let bodyText = decodeUtf8With lenientDecode raw + hooks.onLLMResponseError p.providerName ("HTTP error: " <> T.pack (show status) <> " " <> bodyText) + pure $ Left $ HttpError status bodyText + else case decoded of + Left err -> do + let msg = "response is not valid JSON: " <> T.pack err + hooks.onLLMResponseError p.providerName msg + pure $ Left $ ParseError msg + Right v -> parse v -- | Convert any LLMProvider instance into a LLMGateway. -- LLMHooks are not baked in — they are passed at call time. diff --git a/src/LLM/Core/ProviderUtils.hs b/src/LLM/Core/ProviderUtils.hs --- a/src/LLM/Core/ProviderUtils.hs +++ b/src/LLM/Core/ProviderUtils.hs @@ -1,5 +1,6 @@ module LLM.Core.ProviderUtils ( lenientConfig, + mkLenientConfig, readAll, handleStreamResponse, normalizeSchemaOpenAI, @@ -10,24 +11,49 @@ ) where +import Control.Retry (limitRetries) import Data.Aeson (ToJSON (toJSON), Value (..), decodeStrict', object) import Data.Aeson.KeyMap qualified as KM import Data.ByteString qualified as BS import Data.Text (Text) import Data.Text qualified as T -import Data.Text.Encoding (decodeUtf8, encodeUtf8) +import Data.Text.Encoding (decodeUtf8With, encodeUtf8) +import Data.Text.Encoding.Error (lenientDecode) import LLM.Core.Types (LLMError (HttpError, ParseError), LLMTextResult) import Network.HTTP.Client qualified as HC -import Network.HTTP.Req (HttpConfig, defaultHttpConfig, httpConfigCheckResponse) +import Network.HTTP.Client.TLS (tlsManagerSettings) +import Network.HTTP.Req (HttpConfig, defaultHttpConfig, httpConfigAltManager, httpConfigCheckResponse, httpConfigRetryPolicy) import Network.HTTP.Types.Status (statusCode) +import System.IO.Unsafe (unsafePerformIO) --- | Don't let req throw on non-2xx; we handle errors ourselves +-- | The HTTP config every provider uses: 'mkLenientConfig' with one shared +-- TLS manager that has no response timeout. lenientConfig :: HttpConfig -lenientConfig = +lenientConfig = mkLenientConfig sharedManager + +-- | HTTP config for provider requests using the given manager: +-- +-- * req does not throw on non-2xx; we handle errors ourselves. +-- * req never retries. A retry resends a paid POST, and retrying is the job +-- of 'LLM.Core.Utils.withRetry'. +mkLenientConfig :: HC.Manager -> HttpConfig +mkLenientConfig mgr = defaultHttpConfig - { httpConfigCheckResponse = \_ _ _ -> Nothing + { httpConfigCheckResponse = \_ _ _ -> Nothing, + httpConfigRetryPolicy = limitRetries 0, + httpConfigAltManager = Just mgr } +-- | Manager shared by all provider requests. A non-streaming call receives +-- its response headers only after generation finishes, so http-client's +-- default 30 s response timeout would cut off long calls. The only response +-- deadline is the one the application configures ('mcRequestTimeout'). +sharedManager :: HC.Manager +sharedManager = + unsafePerformIO $ + HC.newManager tlsManagerSettings {HC.managerResponseTimeout = HC.responseTimeoutNone} +{-# NOINLINE sharedManager #-} + -- | Accumulate all chunks from a BodyReader for error messages readAll :: HC.BodyReader -> IO [BS.ByteString] readAll br = do @@ -62,7 +88,9 @@ if status /= 200 then do chunks <- readAll (HC.responseBody resp) - pure $ Left $ HttpError status (decodeUtf8 (BS.concat chunks)) + -- Lenient: an error page need not be UTF-8, and a decode exception + -- here would escape the caller's 'HttpException' handler. + pure $ Left $ HttpError status (decodeUtf8With lenientDecode (BS.concat chunks)) else handler (HC.responseBody resp) -- | Recursively enforce OpenAI structured output constraints: 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 @@ -408,16 +408,13 @@ go (ContentPart (TextPart t) _) = Just t go _ = Nothing --- | First non-empty thinking text, if any. +-- | All non-empty thinking text concatenated in order (without separator, +-- as it was streamed), or 'Nothing' if there is none. 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 +projectReasoning parts = + case [t | ContentPart (ThinkingPart tc) _ <- parts, Just t <- [tc.thinkingText], not (T.null t)] of + [] -> Nothing + ts -> Just (T.concat ts) -- | Tool calls in part order. turnToolCalls :: [ContentPart] -> [ToolCall] @@ -439,6 +436,22 @@ go (ContentPart (TextPart (t1 <> t2)) h1 : rest) go (p : rest) = p : go rest +-- | Merge runs of adjacent thinking parts (e.g. streamed Gemini thought +-- chunks). Only parts without an opaque payload and with the same cache hint +-- are merged; opaque parts (e.g. Claude signed blocks) stay one per block. +coalesceAdjacentThinkingParts :: [ContentPart] -> [ContentPart] +coalesceAdjacentThinkingParts = go + where + go [] = [] + go + ( ContentPart (ThinkingPart (ThinkingContent t1 Nothing)) h1 + : ContentPart (ThinkingPart (ThinkingContent t2 Nothing)) h2 + : rest + ) + | h1 == h2 = + go (ContentPart (ThinkingPart (ThinkingContent (t1 <> t2) Nothing)) h1 : rest) + go (p : rest) = p : go rest + -- | Whether any turn in the conversation contains an image part. conversationHasImages :: [Turn] -> Bool conversationHasImages = any turnHasImage @@ -575,10 +588,12 @@ | TimeoutError -- request timed out | ParseError Text -- JSON we couldn't make sense of | EmptyResponse -- valid JSON, but no content in it - | ToolLoopExceeded Int -- hit the max tool rounds limit - | Aborted -- user cancelled the request + | Aborted -- the request's abort signal fired | -- | Model lacks a required catalog capability (e.g. vision) for this request. UnsupportedCapability Text + | -- | The provider blocked the prompt or the response (e.g. a safety or + -- recitation filter); the text names the reason. Not retried. + ResponseBlocked Text deriving (Show, Eq, Generic, ToJSON, FromJSON) -- | A request to an LLM provider @@ -607,10 +622,11 @@ -- | Build a 'ChatResponse' with text/reasoning projections derived from parts. -- --- Adjacent text parts are coalesced so streamed deltas become one part. +-- Adjacent text parts, and adjacent thinking parts without an opaque +-- payload, are coalesced so streamed deltas become one part. mkChatResponse :: [ContentPart] -> Maybe Usage -> ChatResponse mkChatResponse parts usage = - let coalesced = coalesceAdjacentTextParts parts + let coalesced = coalesceAdjacentThinkingParts (coalesceAdjacentTextParts parts) in ChatResponse { respText = projectText coalesced, respContent = coalesced, 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 @@ -70,7 +70,7 @@ -- | Whether an error is worth retrying isRetryable :: LLMError -> Bool -isRetryable (HttpError status _) = status `elem` [429, 503, 529] +isRetryable (HttpError status _) = status `elem` [429, 500, 502, 503, 504, 529] isRetryable (NetworkError _) = True isRetryable TimeoutError = True isRetryable _ = False diff --git a/src/LLM/Generate/Generate.hs b/src/LLM/Generate/Generate.hs --- a/src/LLM/Generate/Generate.hs +++ b/src/LLM/Generate/Generate.hs @@ -9,11 +9,14 @@ ) where +import Control.Exception (Exception, handle, throwIO) import Control.Monad (unless, when) import Data.IORef (IORef, newIORef, readIORef, writeIORef) import Data.Maybe (fromMaybe) +import LLM.Core.Abort (isAbortedMaybe) import LLM.Core.Types ( ChatResponse (..), + LLMError (Aborted), LLMGateway (..), LLMTextResult, StreamEvent (..), @@ -127,14 +130,26 @@ writeIORef emittedRef False ProviderStreamCallback {pscOnEvent, pscFinalize} <- mkProviderStreamCallback gr onChunk' + -- Cut the stream off at the first event after the abort signal + -- fired: the callback throws out of the provider's stream reader. + let onEvent e = do + aborted <- isAbortedMaybe gr.grAbortSignal + if aborted then throwIO StreamAborted else pscOnEvent e result <- - withFirstEventTimeout - mc.mcRequestTimeout - (mc.mcGateway.gwStreamText gr.grLLMHooks request) - pscOnEvent + handle (\StreamAborted -> pure (Left Aborted)) $ + withFirstEventTimeout + mc.mcRequestTimeout + (mc.mcGateway.gwStreamText gr.grLLMHooks request) + onEvent case result of Right _ -> pscFinalize >> pure result Left err -> pure (Left err) + +-- | Thrown from the stream callback once the abort signal has fired. +data StreamAborted = StreamAborted + deriving (Show) + +instance Exception StreamAborted data RoundPhase = Unclassified | Preamble | Answer deriving (Eq) 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 @@ -16,6 +16,7 @@ import Data.Text (Text) import Data.Text qualified as T import Data.List (nub) +import LLM.Core.Abort (isAbortedMaybe) import LLM.Core.Types ( ChatRequest (..), ContentPart (..), @@ -145,8 +146,16 @@ IO (Either LLMError a) callWithRetry gr mc attempt = maybeThrottle mc.mcThrottleDelay $ - withRetry (modelRetryPolicy mc) (gr.grHooks.onLog Warn) attempt + withRetry (modelRetryPolicy mc) (gr.grHooks.onLog Warn) (unlessAborted gr attempt) +-- | Run @act@ unless 'grAbortSignal' has fired, else return 'Aborted' +-- (which is not retried). Checked before every attempt, so also after +-- throttle and backoff sleeps. +unlessAborted :: GenRequest -> IO (Either LLMError a) -> IO (Either LLMError a) +unlessAborted gr act = do + aborted <- isAbortedMaybe gr.grAbortSignal + if aborted then pure (Left Aborted) else act + -- | 'callWithRetry' where each attempt is limited by 'mcRequestTimeout'. -- -- The timeout applies to each attempt separately, not to the whole @@ -207,7 +216,10 @@ pure $ Left $ maybe GErrAllModelsFailed GErrLLM mLast - loop failed (mc : rest) _ = do + loop failed (mc : rest) mLast = do + aborted <- isAbortedMaybe gr.grAbortSignal + if aborted then pure (Left GErrAborted) else tryModel failed mc rest mLast + tryModel failed mc rest _ = do gr.grHooks.onLog Info (formatTryingModel mc) case validateModelCapabilities mc gr.grMessages of Left err -> 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 @@ -8,7 +8,7 @@ ) where -import Control.Retry (RetryPolicyM, fullJitterBackoff, limitRetries) +import Control.Retry (RetryPolicyM, capDelay, fullJitterBackoff, limitRetries) import Data.Aeson (FromJSON (..), ToJSON (..), object, withObject, (.!=), (.:?), (.=)) import Data.Text (Text) import GHC.Generics (Generic) @@ -76,7 +76,9 @@ mcThrottleDelay :: Maybe Int, -- | Number of retries on retryable 'LLMError's before giving up on this model. mcRetryCount :: Int, - -- | Base backoff in milliseconds between retries (jittered). + -- | Base backoff in milliseconds between retries (jittered). The + -- server's @Retry-After@ header is not honoured, so for a provider that + -- often rate-limits (429), raise this (each delay is capped at 30 s). mcJitterBackoff :: Int } @@ -92,5 +94,9 @@ mfwToModelConfigs :: ModelWithFallbacks -> [ModelConfig] mfwToModelConfigs mwf = mwf.mwfModel : mwf.mwfFallbacks +-- | 'mcRetryCount' retries with full-jitter exponential backoff from +-- 'mcJitterBackoff', each delay capped at 30 s. @Retry-After@ is ignored. modelRetryPolicy :: ModelConfig -> RetryPolicyM IO -modelRetryPolicy mc = limitRetries mc.mcRetryCount <> fullJitterBackoff (mc.mcJitterBackoff * 1000) +modelRetryPolicy mc = + capDelay 30_000_000 $ + limitRetries mc.mcRetryCount <> fullJitterBackoff (mc.mcJitterBackoff * 1000) diff --git a/src/LLM/Generate/Types.hs b/src/LLM/Generate/Types.hs --- a/src/LLM/Generate/Types.hs +++ b/src/LLM/Generate/Types.hs @@ -35,7 +35,10 @@ grMessages :: [Turn], -- | Tool definitions included in the provider request. grTools :: [ToolDef], - -- | Cooperative cancellation checked before and during the call. + -- | Cooperative cancellation. Checked before each model in the fallback + -- chain and before each attempt (also after throttle and backoff + -- sleeps); a stream is cut off at its next event. A non-streaming HTTP + -- call already in flight is not interrupted. Yields 'GErrAborted'. grAbortSignal :: Maybe AbortSignal, -- | Provider wire hooks for raw request/response JSON. grLLMHooks :: LLMHooks, @@ -62,7 +65,9 @@ -- | Failure result preserving partial progress from the current run. data GenerateErrorResult = GenerateErrorResult { gerError :: GenerateError, - -- | Turns produced before the failure (may include an in-progress tool round). + -- | Turns produced before the failure. Safe to append to the history: + -- every tool call has a result. On an abort during a tool round, calls + -- that did not run get the result "Aborted before execution". gerPartialNewMessages :: [Turn], gerUsage :: Usage } @@ -72,7 +77,9 @@ data GenerateError = -- | Provider-level failure (HTTP, network, parse, etc.). GErrLLM LLMError - | -- | 'Agent.agMaxToolRounds' exhausted without a final text reply. + | -- | 'Agent.agMaxToolRounds' exhausted without a final text reply. The + -- last round's tool results are in 'gerPartialNewMessages' but were not + -- sent to the model. GErrToolExceeded | -- | Every model in the fallback chain failed. GErrAllModelsFailed diff --git a/src/LLM/Load/LoadGateways.hs b/src/LLM/Load/LoadGateways.hs --- a/src/LLM/Load/LoadGateways.hs +++ b/src/LLM/Load/LoadGateways.hs @@ -7,18 +7,17 @@ loadGatewaysFromCatalog, loadGatewaysFromCatalogOrThrow, parseProviderBaseUrl, + ProviderUnavailable (..), + describeProviderUnavailable, + resolveGateway, + resolveGateways, ) where import Configuration.Dotenv (defaultConfig, loadFile) import Control.Exception (SomeException, catch, throwIO) -import Control.Monad (forM) -import Control.Monad.Except (ExceptT (..), MonadError (throwError), runExceptT) -import Control.Monad.IO.Class (MonadIO (liftIO)) -import Data.Functor ((<&>)) import Data.Map qualified as Map import Data.ByteString.Char8 qualified as B -import Data.Maybe (catMaybes) import Data.Text (Text) import Data.Text qualified as T import LLM.Core.Types (LLMGateway) @@ -115,65 +114,51 @@ mbOverride <- lookupEnv (T.unpack envVar) pure $ maybe item.baseUrl T.pack mbOverride -tryBuildGateway :: ProviderCatalogItem -> IO (Maybe (Text, LLMGateway)) -tryBuildGateway item = - case item.apiKeyEnv of - Nothing -> buildWithKey "" - Just envVar -> - lookupEnv (T.unpack envVar) >>= (\case - Nothing -> pure Nothing - Just apiKey -> buildWithKey apiKey) . fmap T.pack - where - buildWithKey apiKey = do - baseUrlText <- resolveBaseUrl item - case parseProviderBaseUrl baseUrlText of - Left _ -> pure Nothing - Right baseUrl -> - pure $ - Just - ( item.providerName, - buildGateway item baseUrl apiKey - ) +-- | Why a provider in the catalog has no gateway. +data ProviderUnavailable + = -- | The provider's API key environment variable is not set. + MissingApiKey Text + | -- | The provider's base URL cannot be parsed (the parse error). + InvalidBaseUrl Text + deriving (Eq, Show) -tryBuildGatewayOrThrow :: ProviderCatalogItem -> ExceptT LoadConfigError IO (Maybe (Text, LLMGateway)) -tryBuildGatewayOrThrow item = - case item.apiKeyEnv of - Nothing -> buildWithKey "" - Just envVar -> - liftIO (lookupEnv (T.unpack envVar) <&> fmap T.pack) >>= \case - Nothing -> pure Nothing - Just apiKey -> buildWithKey apiKey - where - buildWithKey apiKey = do - baseUrlText <- liftIO $ resolveBaseUrl item - baseUrl <- case parseProviderBaseUrl baseUrlText of - Left err -> - throwError $ - LoadModelConfigError $ - "Invalid baseUrl for provider " - <> item.providerName - <> ": " - <> err - Right parsed -> pure parsed - pure $ - Just - ( item.providerName, - buildGateway item baseUrl apiKey - ) +-- | Describe why a provider is unavailable, e.g. "env var X is not set". +describeProviderUnavailable :: ProviderUnavailable -> Text +describeProviderUnavailable = \case + MissingApiKey envVar -> "env var " <> envVar <> " is not set" + InvalidBaseUrl err -> "it has an " <> err +-- | Build the gateway for a provider, or say why it is unavailable. Never +-- throws. +resolveGateway :: ProviderCatalogItem -> IO (Either ProviderUnavailable LLMGateway) +resolveGateway item = do + apiKey <- case item.apiKeyEnv of + Nothing -> pure (Right "") + Just envVar -> maybe (Left (MissingApiKey envVar)) (Right . T.pack) <$> lookupEnv (T.unpack envVar) + baseUrlText <- resolveBaseUrl item + pure $ do + key <- apiKey + baseUrl <- either (Left . InvalidBaseUrl) Right (parseProviderBaseUrl baseUrlText) + pure (buildGateway item baseUrl key) + +-- | Resolve every provider in the catalog. See 'resolveGateway'. +resolveGateways :: ProviderCatalogMap -> IO (Map.Map Text (Either ProviderUnavailable LLMGateway)) +resolveGateways = traverse resolveGateway + +-- | Gateways of all providers that are available. Providers without an API +-- key or with an invalid base URL are left out. loadGatewaysFromCatalog :: ProviderCatalogMap -> IO GatewayMap -loadGatewaysFromCatalog catalog = do - pairs <- forM (Map.elems catalog) tryBuildGateway - pure $ Map.fromList $ catMaybes pairs +loadGatewaysFromCatalog catalog = Map.mapMaybe (either (const Nothing) Just) <$> resolveGateways catalog +-- | Like 'loadGatewaysFromCatalog', but throws a 'LoadModelConfigError' if +-- any provider has an invalid base URL. loadGatewaysFromCatalogOrThrow :: ProviderCatalogMap -> IO GatewayMap loadGatewaysFromCatalogOrThrow catalog = do - result <- runExceptT $ do - pairs <- forM (Map.elems catalog) tryBuildGatewayOrThrow - pure $ Map.fromList $ catMaybes pairs - case result of - Left err -> throwIO err - Right gatewayMap -> pure gatewayMap + resolved <- resolveGateways catalog + case [(name, err) | (name, Left (InvalidBaseUrl err)) <- Map.toList resolved] of + (name, err) : _ -> + throwIO $ LoadModelConfigError $ "Invalid baseUrl for provider " <> name <> ": " <> err + [] -> pure $ Map.mapMaybe (either (const Nothing) Just) resolved -- | Build a gateway map from the bundled default provider catalog and the -- current process environment. Does not read a @.env@ file; set variables 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 @@ -12,64 +12,102 @@ import Control.Exception (throwIO) import Control.Lens (Each (each), mapMOf) -import Control.Monad (forM) -import Control.Monad.Except (ExceptT, MonadError (throwError), liftEither, runExceptT) +import Control.Monad.Except (ExceptT, liftEither, runExceptT) import Control.Monad.IO.Class (MonadIO (liftIO)) 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.Core.Types (LLMGateway, ThinkingMode (..)) import LLM.Generate.ModelConfig ( ModelConfig (..), defaultModelCapabilities, ) -import LLM.Load.LoadGateways (GatewayMap, loadGatewaysFromCatalogOrThrow) -import LLM.Load.ModelCatalog (ModelCatalogItem (..), loadModelCatalog) +import LLM.Load.LoadGateways (GatewayMap, ProviderUnavailable, describeProviderUnavailable, resolveGateways) +import LLM.Load.ModelCatalog (ModelCatalogItem (..), ModelCatalogMap, loadModelCatalog) import LLM.Load.ProviderCatalog (loadProviderCatalogForModelCatalog) import LLM.Load.Types (LoadConfigError (LoadModelConfigError)) +-- | Model configs by catalog name. type ModelConfigMap = Map Text ModelConfig -loadModelConfigMap :: FilePath -> ExceptT LoadConfigError IO (ModelConfigMap, GatewayMap) -loadModelConfigMap filePath = do +-- | Providers by name, each either its gateway or why it is unavailable. +type ResolvedProviders = Map Text (Either ProviderUnavailable LLMGateway) + +-- | Parse the model and provider catalogs and resolve every provider. Fails +-- only if a catalog cannot be read or parsed; unavailable providers are kept +-- with the reason. +loadCatalogs :: FilePath -> ExceptT LoadConfigError IO (ModelCatalogMap, ResolvedProviders) +loadCatalogs filePath = do providerCatalog <- loadProviderCatalogForModelCatalog filePath - gatewayMap <- liftIO $ loadGatewaysFromCatalogOrThrow providerCatalog - modelCatalogMap <- loadModelCatalog filePath - modelConfigs <- forM (Map.toList modelCatalogMap) $ \(modelConfigName, item) -> do - case Map.lookup item.providerName gatewayMap of - Nothing -> throwError $ LoadModelConfigError $ "Provider not found: " <> item.providerName - Just gateway -> - pure - ( modelConfigName, - ModelConfig - { mcGateway = gateway, - mcModel = item.modelName, - mcPricing = item.pricing, - 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, - mcJitterBackoff = item.jitterBackoff - } - ) - pure (Map.fromList modelConfigs, gatewayMap) + modelCatalog <- loadModelCatalog filePath + providers <- liftIO $ resolveGateways providerCatalog + pure (modelCatalog, providers) +-- | Load the catalogs next to @filePath@ and return all models and gateways +-- that are currently available. Models whose provider is unavailable (API key +-- env var not set, invalid base URL, unknown provider) are left out; this +-- never fails because of them. +loadModelConfigMap :: FilePath -> ExceptT LoadConfigError IO (ModelConfigMap, GatewayMap) +loadModelConfigMap filePath = available <$> loadCatalogs filePath + +available :: (ModelCatalogMap, ResolvedProviders) -> (ModelConfigMap, GatewayMap) +available (models, providers) = + ( Map.mapMaybe (either (const Nothing) Just . resolveModel providers) models, + Map.mapMaybe (either (const Nothing) Just) providers + ) + +-- | Look up a requested model. Fails with the reason when the model is not in +-- the catalog or its provider is unavailable. +lookupModel :: ModelCatalogMap -> ResolvedProviders -> Text -> Either LoadConfigError ModelConfig +lookupModel models providers name = + case Map.lookup name models of + Nothing -> Left $ LoadModelConfigError $ "Model not found: " <> name + Just item -> resolveModel providers item + +resolveModel :: ResolvedProviders -> ModelCatalogItem -> Either LoadConfigError ModelConfig +resolveModel providers item = + case Map.lookup item.providerName providers of + Nothing -> unavailable "is not in the provider catalog" + Just (Left reason) -> unavailable (describeProviderUnavailable reason) + Just (Right gateway) -> + Right + ModelConfig + { mcGateway = gateway, + mcModel = item.modelName, + mcPricing = item.pricing, + 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, + mcJitterBackoff = item.jitterBackoff + } + where + unavailable reason = + Left $ + LoadModelConfigError $ + "Model " <> item.modelConfigName <> " needs provider " <> item.providerName <> ", but " <> reason + +-- | Load the requested models. Also returns all available models and +-- gateways (see 'loadModelConfigMap'). loadModels :: (Each s t Text ModelConfig) => FilePath -> s -> ExceptT LoadConfigError IO (t, ModelConfigMap, GatewayMap) loadModels filePath s = do - (modelConfigMap, gatewayMap) <- loadModelConfigMap filePath - r <- liftEither $ maybe (Left $ LoadModelConfigError "Model not found") Right $ mapMOf each (`Map.lookup` modelConfigMap) s + catalogs@(models, providers) <- loadCatalogs filePath + r <- liftEither $ mapMOf each (lookupModel models providers) s + let (modelConfigMap, gatewayMap) = available catalogs pure (r, modelConfigMap, gatewayMap) +-- | Load one model. Also returns all available models and gateways (see +-- 'loadModelConfigMap'). loadModel :: FilePath -> Text -> ExceptT LoadConfigError IO (ModelConfig, ModelConfigMap, GatewayMap) loadModel filePath modelConfigName = do - (modelConfigMap, gatewayMap) <- loadModelConfigMap filePath - case Map.lookup modelConfigName modelConfigMap of - Nothing -> throwError $ LoadModelConfigError $ "Model not found: " <> modelConfigName - Just modelConfig -> pure (modelConfig, modelConfigMap, gatewayMap) + catalogs@(models, providers) <- loadCatalogs filePath + modelConfig <- liftEither $ lookupModel models providers modelConfigName + let (modelConfigMap, gatewayMap) = available catalogs + pure (modelConfig, modelConfigMap, gatewayMap) -- | Load models from a json file and throw an error if any model is not found -- @@ -91,14 +129,5 @@ (\(x, _, _) -> x) <$> loadModelOrThrow_ filePath modelConfigName loadModelOrThrow_ :: FilePath -> Text -> IO (ModelConfig, ModelConfigMap, GatewayMap) -loadModelOrThrow_ filePath modelConfigName = do - r <- runExceptT $ loadModelConfigMap filePath - case r of - Left err -> throwIO err - Right (modelConfigMap, gatewayMap) -> do - case Map.lookup modelConfigName modelConfigMap of - Nothing -> - throwIO $ - LoadModelConfigError $ - "Model not found: " <> modelConfigName - Just modelConfig -> pure (modelConfig, modelConfigMap, gatewayMap) +loadModelOrThrow_ filePath modelConfigName = + runExceptT (loadModel filePath modelConfigName) >>= either throwIO pure diff --git a/src/LLM/Load/Utils.hs b/src/LLM/Load/Utils.hs --- a/src/LLM/Load/Utils.hs +++ b/src/LLM/Load/Utils.hs @@ -4,9 +4,10 @@ import Control.Monad.Except (ExceptT (..), MonadError (throwError)) import Control.Monad.IO.Class (MonadIO (liftIO)) import Data.Aeson (FromJSON, eitherDecodeFileStrict) +import Data.ByteString qualified as BS import Data.Text (Text) import Data.Text qualified as T -import Data.Text.IO qualified as TIO +import Data.Text.Encoding (decodeUtf8') import LLM.Load.Types (LoadConfigError (..)) decodeJsonFile :: @@ -27,8 +28,10 @@ if T.isPrefixOf "file:" t then do let filePath = drop 5 (T.unpack t) - result <- liftIO (try (TIO.readFile filePath) :: IO (Either IOException Text)) + result <- liftIO (try (BS.readFile filePath) :: IO (Either IOException BS.ByteString)) case result of Left (e :: IOException) -> throwError (LoadSystemPromptError (T.pack (show e))) - Right content -> pure (Just content) + Right bytes -> case decodeUtf8' bytes of + Left _ -> throwError (LoadSystemPromptError (T.pack filePath <> " is not valid UTF-8")) + Right content -> pure (Just content) else pure $ Just t 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 @@ -16,7 +16,7 @@ import Control.Applicative ((<|>)) import Data.Aeson ( KeyValue ((.=)), - Value (Object, String), + Value (Array, Object, String), decodeStrict', encode, object, @@ -31,7 +31,7 @@ import Data.ByteString qualified as BS import Data.Foldable (for_) import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef) -import Data.Maybe (fromMaybe, mapMaybe) +import Data.Maybe (catMaybes, fromMaybe, mapMaybe) import Data.Text (Text) import Data.Text qualified as T import Data.Text.Encoding (encodeUtf8) @@ -84,7 +84,7 @@ Url, header, https, - jsonResponse, + bsResponse, req, reqBr, responseBody, @@ -157,7 +157,7 @@ runReq lenientConfig $ do let url = baseUrl /: "v1" /: "messages" opts = baseOpts <> claudeAuthOpts apiKey - resp <- req POST url (ReqBodyJson body) jsonResponse opts + resp <- req POST url (ReqBodyJson body) bsResponse opts pure (responseStatusCode resp, responseBody resp) claudeAuthOpts :: Text -> Option scheme @@ -245,19 +245,29 @@ ["thinking" .= object ["type" .= ("disabled" :: Text)]] _ -> [] +-- | Encode a turn for Claude. A user or assistant turn whose content encodes +-- to nothing (e.g. an assistant turn holding only foreign thinking) is +-- dropped: Claude rejects @"content": []@, and merges the consecutive +-- same-role turns this can leave. encodeTurn :: Text -> Turn -> [Value] encodeTurn _ (UserMessage parts) = - [ object - [ "role" .= ("user" :: Text), - "content" .= encodeUserContent parts + case encodeUserContent parts of + Array a | null a -> [] + content -> + [ object + [ "role" .= ("user" :: Text), + "content" .= content + ] ] - ] encodeTurn currentModel (AssistantMessage parts) = - [ object - [ "role" .= ("assistant" :: Text), - "content" .= mapMaybe (encodeAssistantPart currentModel . stripForeignOpaque claudeProviderName) parts + case mapMaybe (encodeAssistantPart currentModel . stripForeignOpaque claudeProviderName) parts of + [] -> [] + content -> + [ object + [ "role" .= ("assistant" :: Text), + "content" .= content + ] ] - ] encodeTurn _ (ToolTurn results) = [ object [ "role" .= ("user" :: Text), @@ -369,13 +379,15 @@ go :: Maybe Text -> Value -> Parser [ContentPart] go mv = withObject "ClaudeResponse" $ \o -> do content <- o .: "content" :: Parser [Value] - mapM (parseBlock mv) content + catMaybes <$> mapM (parseBlock mv) content -parseBlock :: Maybe Text -> Value -> Parser ContentPart +-- | Parse one content block. Unknown block types (e.g. @server_tool_use@, +-- @web_search_tool_result@) give 'Nothing' and are skipped, as in streaming. +parseBlock :: Maybe Text -> Value -> Parser (Maybe ContentPart) parseBlock mv = withObject "content_block" $ \o -> do typ <- o .: "type" :: Parser Text case typ of - "text" -> textPart <$> o .: "text" + "text" -> Just . textPart <$> o .: "text" "thinking" -> do thinkingTxt <- o .:? "thinking" :: Parser (Maybe Text) signature <- o .:? "signature" :: Parser (Maybe Text) @@ -393,7 +405,7 @@ mText = case thinkingTxt of Just t | not (T.null t) -> Just t _ -> Nothing - pure $ thinkingPart (ThinkingContent mText (Just opaque)) + pure $ Just $ thinkingPart (ThinkingContent mText (Just opaque)) "redacted_thinking" -> do data_ <- o .: "data" :: Parser Text let payload = @@ -407,13 +419,13 @@ poModel = mv, poPayload = payload } - pure $ thinkingPart (ThinkingContent Nothing (Just opaque)) + pure $ Just $ 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 + pure $ Just $ toolCallPart (mkToolCall cid name args) + _ -> pure Nothing parseClaudeUsage :: Value -> Maybe Usage parseClaudeUsage = parseMaybe $ withObject "ClaudeResponse" $ \o -> do 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 @@ -54,7 +54,7 @@ ReqBodyJson (ReqBodyJson), Url, https, - jsonResponse, + bsResponse, req, reqBr, responseBody, @@ -114,7 +114,7 @@ runReq lenientConfig $ do let url = baseUrl /: "v1" /: "chat" /: "completions" opts = baseOpts <> authHeader apiKey - resp <- req POST url (ReqBodyJson body) jsonResponse opts + resp <- req POST url (ReqBodyJson body) bsResponse opts pure (responseStatusCode resp, responseBody resp) parseObject :: Value -> Parser Text parseObject = withObject "OpenAIObjectResponse" $ \o -> do 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 @@ -12,6 +12,7 @@ where import Control.Applicative ((<|>)) +import Control.Monad (forM) import Data.Aeson ( KeyValue ((.=)), Object, @@ -26,11 +27,14 @@ 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, mapMaybe) +import Data.Maybe (catMaybes, fromMaybe, mapMaybe) +import Data.Char (isDigit) +import Data.Foldable (for_) import Data.Text (Text) import Data.Text qualified as T import Data.Text.Encoding (encodeUtf8) -import Data.Unique (hashUnique, newUnique) +import Data.UUID.Types (UUID) +import Data.UUID.Types qualified as UUID import LLM.Core.LLMProvider (LLMProvider (..), toGateway) import LLM.Core.ProviderUtils (handleStreamResponse, lenientConfig, stripBoundsAndComments, stripJsonFences) import LLM.Core.SSE (SSEEvent (sseData), readSSEEvents) @@ -46,7 +50,7 @@ reqTools ), ContentPart (..), - LLMError (EmptyResponse), + LLMError (EmptyResponse, HttpError, ParseError, ResponseBlocked), LLMGateway, LLMObjectResult, LLMTextResult, @@ -76,7 +80,7 @@ Url, header, https, - jsonResponse, + bsResponse, req, reqBr, responseBody, @@ -85,6 +89,7 @@ (/:), (=:), ) +import System.Random (randomIO) geminiProviderName :: Text geminiProviderName = "gemini" @@ -136,7 +141,9 @@ "responseSchema" .= stripBoundsAndComments schema ] ++ ["temperature" .= t | Just t <- [r.reqTemperature]] - ++ thinkingConfigPairs r + -- The object API returns no reasoning, so don't ask + -- for thought summaries. + ++ thinkingConfigPairs False r ) ] ++ [ "system_instruction" .= object ["parts" .= [object ["text" .= sys]]] @@ -158,7 +165,7 @@ /: "models" /: (model <> ":generateContent") opts = baseOpts <> geminiAuthOpts apiKey - resp <- req POST url (ReqBodyJson (stripModel body)) jsonResponse opts + resp <- req POST url (ReqBodyJson (stripModel body)) bsResponse opts pure (responseStatusCode resp, responseBody resp) geminiAuthOpts :: Text -> Option scheme @@ -175,10 +182,17 @@ parseGeminiStream reader callback = do partsRef <- newIORef ([] :: [ContentPart]) usageRef <- newIORef Nothing + errorRef <- newIORef (Nothing :: Maybe LLMError) + stopRef <- newIORef (Nothing :: Maybe GeminiStop) readSSEEvents (HC.brRead reader) $ \sse -> do case decodeStrict' (encodeUtf8 sse.sseData) of Nothing -> pure () Just v -> do + -- The stream stays HTTP 200 when it fails midway; the error arrives + -- as a data event. Keep the first one. + for_ (parseMaybe parseErrorCode v) $ \code -> + modifyIORef' errorRef (<|> Just (HttpError code sse.sseData)) + for_ (parseGeminiStop v) (writeIORef stopRef . Just) let modelVer = parseMaybe parseModelVersion v case parseMaybe (parseChunkParts modelVer) v of Just parts -> do @@ -190,9 +204,13 @@ Nothing -> pure () parts <- readIORef partsRef usage <- readIORef usageRef - if null parts - then pure $ Left EmptyResponse - else pure $ Right (mkChatResponse parts usage) + mError <- readIORef errorRef + stop <- readIORef stopRef + pure $ case mError of + Just err -> Left err + Nothing + | null parts -> Left (noContentError stop) + | otherwise -> Right (mkChatResponse parts usage) where assignToolId :: (StreamEvent -> IO ()) -> ContentPart -> IO ContentPart assignToolId cb (ContentPart (TextPart t) hint) = do @@ -244,24 +262,19 @@ -- | Encode a turn for Gemini. Foreign opaque thinking / tool metadata is omitted -- at encode time so fallbacks never replay Claude (or other) state. +-- A user or model turn whose parts encode to nothing (e.g. a Claude turn +-- holding only redacted thinking) is dropped, since Gemini rejects +-- @"parts": []@. encodeTurn :: Text -> Turn -> [Value] encodeTurn _ (UserMessage parts) = - [ object - [ "role" .= ("user" :: Text), - "parts" .= mapMaybe encodeUserPart parts - ] - ] + nonEmptyTurn "user" (mapMaybe encodeUserPart parts) 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" .= mapMaybe (encodeAssistantPart currentModel) cleaned - ] - ] + nonEmptyTurn "model" (mapMaybe (encodeAssistantPart currentModel) cleaned) where cleaned = map (stripForeignOpaque geminiProviderName) parts encodeTurn _ (ToolTurn results) = @@ -271,6 +284,10 @@ ] ] +nonEmptyTurn :: Text -> [Value] -> [Value] +nonEmptyTurn _ [] = [] +nonEmptyTurn role ps = [object ["role" .= role, "parts" .= ps]] + encodeImagePart :: ImageSource -> Value encodeImagePart (ImageUrl url) = object @@ -364,11 +381,13 @@ ] ] +-- | Give a Gemini function call an id. A random UUID, so ids never repeat +-- across processes: a reloaded history must not get new calls with ids it +-- already holds, or Claude / OpenAI reject it after a fallback. normalizeToolCallId :: ToolCall -> IO ToolCall normalizeToolCallId tc = do - u <- newUnique - let callId = "call_" <> T.pack (show (hashUnique u)) - pure tc {tcId = callId} + u <- randomIO :: IO UUID + pure tc {tcId = "call_" <> UUID.toText u} normalizePart :: ContentPart -> IO ContentPart normalizePart (ContentPart (ToolCallPart tc) hint) = @@ -380,15 +399,16 @@ object $ ("maxOutputTokens" .= r.reqMaxTokens) : ["temperature" .= t | Just t <- [r.reqTemperature]] - ++ thinkingConfigPairs r + ++ thinkingConfigPairs True 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 = +-- without effort requests dynamic budget (@-1@). With @includeThoughts@, +-- the response carries thought summaries as parts marked @"thought": true@. +thinkingConfigPairs :: Bool -> ChatRequest -> [Pair] +thinkingConfigPairs includeThoughts r = case r.reqThinking of Nothing -> [] Just tm @@ -396,27 +416,15 @@ ["thinkingConfig" .= object ["thinkingBudget" .= (0 :: Int)]] | Just e <- tm.tmEffort, e `elem` ["minimal", "low", "medium", "high"] -> - [ "thinkingConfig" - .= object - [ "thinkingLevel" .= e, - "includeThoughts" .= True - ] - ] + thinkingConfig ["thinkingLevel" .= e] | Just e <- tm.tmEffort, Just n <- readMaybeInt e -> - [ "thinkingConfig" - .= object - [ "thinkingBudget" .= n, - "includeThoughts" .= True - ] - ] + thinkingConfig ["thinkingBudget" .= n] | otherwise -> - [ "thinkingConfig" - .= object - [ "thinkingBudget" .= (-1 :: Int), - "includeThoughts" .= True - ] - ] + thinkingConfig ["thinkingBudget" .= (-1 :: Int)] + where + thinkingConfig pairs = + ["thinkingConfig" .= object (pairs ++ ["includeThoughts" .= True | includeThoughts])] readMaybeInt :: Text -> Maybe Int readMaybeInt t = @@ -424,13 +432,63 @@ [(n, "")] -> Just n _ -> Nothing +-- | Why Gemini stopped: the prompt was blocked (@promptFeedback.blockReason@) +-- or the first candidate finished (@finishReason@). +data GeminiStop = GeminiPromptBlocked Text | GeminiFinished Text + +parseGeminiStop :: Value -> Maybe GeminiStop +parseGeminiStop v = parseMaybe blocked v <|> parseMaybe finished v + where + blocked = withObject "GeminiResponse" $ \o -> do + fb <- o .: "promptFeedback" + GeminiPromptBlocked <$> withObject "promptFeedback" (.: "blockReason") fb + finished = withObject "GeminiResponse" $ \o -> do + (cand : _) <- o .: "candidates" :: Parser [Value] + GeminiFinished <$> withObject "candidate" (.: "finishReason") cand + +-- | The error for a response without usable content, named by why Gemini +-- stopped. Responses that do have content are returned as they are, even +-- when cut off by the token limit. +noContentError :: Maybe GeminiStop -> LLMError +noContentError = \case + Just (GeminiPromptBlocked reason) -> + ResponseBlocked ("Gemini blocked the prompt (blockReason: " <> reason <> ")") + Just (GeminiFinished reason) + | reason `elem` blockedFinishReasons -> + ResponseBlocked ("Gemini blocked the response (finishReason: " <> reason <> ")") + | reason == "MAX_TOKENS" -> + ParseError "Gemini returned no content (finishReason: MAX_TOKENS); increase max tokens, which thinking also uses" + | reason /= "STOP" -> + ParseError ("Gemini returned no content (finishReason: " <> reason <> ")") + _ -> EmptyResponse + +-- | Finish reasons that mean a filter stopped the response. +blockedFinishReasons :: [Text] +blockedFinishReasons = + [ "SAFETY", + "RECITATION", + "LANGUAGE", + "BLOCKLIST", + "PROHIBITED_CONTENT", + "SPII", + "IMAGE_SAFETY", + "IMAGE_PROHIBITED_CONTENT", + "IMAGE_RECITATION" + ] + +-- | The status code of an error event: @{"error": {"code": 503, ...}}@. +parseErrorCode :: Value -> Parser Int +parseErrorCode = withObject "GeminiError" $ \o -> do + e <- o .: "error" + withObject "error" (.: "code") e + parseGeminiResponse :: Value -> IO LLMTextResult parseGeminiResponse v = case parseMaybe (go modelVer) v of - Nothing -> pure $ Left EmptyResponse + Nothing -> pure $ Left (noContentError (parseGeminiStop v)) Just parts -> do parts' <- mapM normalizePart parts case parts' of - [] -> pure $ Left EmptyResponse + [] -> pure $ Left (noContentError (parseGeminiStop v)) _ -> pure $ Right (mkChatResponse parts' (parseGeminiUsage v)) where modelVer = parseMaybe parseModelVersion v @@ -517,8 +575,21 @@ _ -> Nothing signatureForModel _ _ = Nothing +-- | Whether a recorded model (the response's @modelVersion@) is the model +-- being called. Ignores only a @models/@ prefix and a stable version +-- suffix (@gemini-2.5-pro-001@ = @gemini-2.5-pro@); anything else must be +-- equal, so @gemini-2.5-flash@ and @gemini-2.5-flash-lite@ differ. A false +-- mismatch only costs a placeholder signature; a false match is a 400. modelsMatch :: Text -> Text -> Bool -modelsMatch a b = a == b || T.isPrefixOf a b || T.isPrefixOf b a +modelsMatch a b = canonical a == canonical b + where + canonical = dropStableVersion . dropModelsPrefix + dropModelsPrefix t = fromMaybe t (T.stripPrefix "models/" t) + dropStableVersion t = + let (stem, suffix) = T.breakOnEnd "-" t + in if T.length suffix == 3 && T.all isDigit suffix && T.length stem > 1 + then T.dropEnd 1 stem + else t parseGeminiUsage :: Value -> Maybe Usage parseGeminiUsage = parseMaybe $ withObject "GeminiResponse" $ \o -> do @@ -542,12 +613,25 @@ parseGeminiObjectResponse :: Value -> IO LLMObjectResult parseGeminiObjectResponse v = case parseMaybe go v of - Nothing -> pure $ Left EmptyResponse + Nothing -> pure $ Left (noContentError stop) Just text -> case decodeStrict' (encodeUtf8 (stripJsonFences text)) of - Nothing -> pure $ Left EmptyResponse + Nothing + | Just (GeminiFinished "MAX_TOKENS") <- stop -> + pure $ Left $ ParseError "Gemini structured output was cut off (finishReason: MAX_TOKENS); increase max tokens" + | otherwise -> pure $ Left EmptyResponse Just obj -> pure $ Right (obj, parseGeminiUsage v) where + stop = parseGeminiStop v + + -- The JSON is the text of all parts that are not thought summaries + -- (@"thought": true@), which come first when thinking is enabled. go :: Value -> Parser Text go = withObject "GeminiObjectResponse" $ \o -> do (cand : _) <- o .: "candidates" :: Parser [Value] - withObject "candidate" (\co -> co .: "content" >>= withObject "content" (\cco -> cco .: "parts" >>= \case (p : _) -> withObject "part" (.: "text") p; _ -> fail "No parts")) cand + parts <- withObject "candidate" (\co -> co .: "content" >>= withObject "content" (.: "parts")) cand + texts <- forM parts $ withObject "part" $ \p -> do + thought <- p .:? "thought" .!= False + if thought then pure Nothing else p .:? "text" + case catMaybes texts of + [] -> fail "No text parts" + ts -> pure (T.concat ts) diff --git a/src/LLM/Providers/Ollama.hs b/src/LLM/Providers/Ollama.hs --- a/src/LLM/Providers/Ollama.hs +++ b/src/LLM/Providers/Ollama.hs @@ -14,7 +14,7 @@ withObject, (.:), ) -import Data.Aeson.Types (Parser, parseMaybe) +import Data.Aeson.Types (Pair, Parser, parseMaybe) import Data.Text (Text) import Data.Text.Encoding (encodeUtf8) import LLM.Core.LLMProvider (LLMProvider (..), toGateway) @@ -31,14 +31,14 @@ LLMGateway, defaultMessageEncodeOptions, ) -import LLM.Providers.OpenAI (buildMessages, encodeToolDef, openAIBuildBodyPairs, parseOpenAIResponse, parseOpenAIStream, parseOpenAIUsage) +import LLM.Providers.OpenAI (buildMessages, encodeToolDef, parseOpenAIResponse, parseOpenAIStream, parseOpenAIUsage) import Network.HTTP.Req ( Option, POST (POST), ReqBodyJson (ReqBodyJson), Url, http, - jsonResponse, + bsResponse, port, req, reqBr, @@ -68,9 +68,7 @@ -- Ollama rejects image URLs ("not currently supported, please use -- base64 encoded data instead"); only data URLs work. providerImageUrlSupport = False, - -- Ollama uses the same request format as OpenAI, but without stream_options - -- (Ollama doesn't support include_usage in streaming). - buildBody = ollamaBuildBody, + buildBody = \stream r -> object (ollamaBuildBodyPairs stream r), sendRequest = sendRequest, openStream = \body k -> runReq lenientConfig $ do @@ -81,7 +79,7 @@ parseResponse = pure . parseOpenAIResponse, buildObjectBody = \r schema -> object $ - openAIBuildBodyPairs False r + ollamaBuildBodyPairs False r <> [ "response_format" .= object [ "type" .= ("json_schema" :: Text), @@ -104,22 +102,23 @@ sendRequest body = runReq lenientConfig $ do let url = baseUrl /: "v1" /: "chat" /: "completions" - resp <- req POST url (ReqBodyJson body) jsonResponse baseOpts + resp <- req POST url (ReqBodyJson body) bsResponse baseOpts pure (responseStatusCode resp, responseBody resp) parseObject :: Value -> Parser Text parseObject = withObject "OpenAIObjectResponse" $ \o -> do (choice : _) <- o .: "choices" :: Parser [Value] withObject "choice" (\co -> co .: "message" >>= withObject "message" (.: "content")) choice --- | Build request body — same as OpenAI but without stream_options --- since Ollama doesn't support include_usage. -ollamaBuildBody :: Bool -> ChatRequest -> Value -ollamaBuildBody stream r = - object $ - [ "model" .= r.reqModel, - "messages" .= buildMessages defaultMessageEncodeOptions r - ] - ++ ["num_predict" .= r.reqMaxTokens] - ++ ["temperature" .= t | Just t <- [r.reqTemperature]] - ++ ["tools" .= map encodeToolDef r.reqTools | not (null r.reqTools)] - ++ ["stream" .= True | stream] +-- | Request body for Ollama's OpenAI-compatible endpoint. Like OpenAI's, but +-- the token limit is @max_tokens@: Ollama ignores @max_completion_tokens@ +-- (and @num_predict@ belongs to its native API). +ollamaBuildBodyPairs :: Bool -> ChatRequest -> [Pair] +ollamaBuildBodyPairs stream r = + [ "model" .= r.reqModel, + "max_tokens" .= r.reqMaxTokens, + "messages" .= buildMessages defaultMessageEncodeOptions r + ] + ++ ["temperature" .= t | Just t <- [r.reqTemperature]] + ++ ["tools" .= map encodeToolDef r.reqTools | not (null r.reqTools)] + ++ ["stream" .= True | stream] + ++ ["stream_options" .= object ["include_usage" .= True] | stream] 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 @@ -37,7 +37,7 @@ import Data.Foldable (foldl', for_, forM_) import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef) import Data.Map.Strict qualified as Map -import Data.Maybe (catMaybes, fromMaybe, listToMaybe, mapMaybe) +import Data.Maybe (catMaybes, fromMaybe, mapMaybe) import Control.Applicative ((<|>)) import Data.Text (Text) import Data.Text qualified as T @@ -47,6 +47,7 @@ import LLM.Core.SSE (SSEEvent (sseData), readSSEEvents) import LLM.Core.Types ( commonImageMediaTypes, + projectReasoning, ChatRequest ( reqConversation, reqMaxTokens, @@ -87,7 +88,7 @@ Url, header, https, - jsonResponse, + bsResponse, req, reqBr, responseBody, @@ -165,7 +166,7 @@ runReq lenientConfig $ do let url = baseUrl /: "v1" /: "chat" /: "completions" opts = baseOpts <> authHeader apiKey - resp <- req POST url (ReqBodyJson body) jsonResponse opts + resp <- req POST url (ReqBodyJson body) bsResponse opts pure (responseStatusCode resp, responseBody resp) authHeader :: Text -> Option scheme @@ -216,13 +217,7 @@ 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) - ] + mReasoning = projectReasoning cleaned calls = [tc | ContentPart (ToolCallPart tc) _ <- cleaned] requireReasoning | opts.meoRequireReasoning = Just . fromMaybe "" diff --git a/src/LLM/Tools/CopyFile.hs b/src/LLM/Tools/CopyFile.hs --- a/src/LLM/Tools/CopyFile.hs +++ b/src/LLM/Tools/CopyFile.hs @@ -2,11 +2,12 @@ import Autodocodec qualified as AC import Data.Aeson (FromJSON) +import Data.Maybe (isJust) import Data.Text (Text) import Data.Text qualified as T import GHC.Generics (Generic) import LLM.Core.Types (TypedTool (..)) -import LLM.Tools.FsConfig (FsConfig, isSymlink, sandboxPath, sandboxWritePath) +import LLM.Tools.FsConfig (FsConfig, sandboxPath, sandboxSymlink, sandboxWritePath) import System.Directory (copyFile, doesDirectoryExist, doesFileExist, doesPathExist) data CopyFileToolArgs = CopyFileToolArgs @@ -39,14 +40,15 @@ copyFileExecTyped cfg args = do let src = args._cfSrc dst = args._cfDst - srcResolved <- sandboxPath cfg (T.unpack src) -- Refuse symbolic-link sources: copying would either dereference the -- link (potentially pulling content from outside the sandbox) or be -- ambiguous about whether the link itself should be reproduced. - srcIsLink <- isSymlink srcResolved - if srcIsLink + -- Checked before resolving, since `sandboxPath` follows the link. + srcLink <- sandboxSymlink cfg (T.unpack src) + if isJust srcLink then pure $ "Error: refusing to copy symbolic link source: " <> src else do + srcResolved <- sandboxPath cfg (T.unpack src) srcExists <- doesPathExist srcResolved if not srcExists then pure $ "Error: source does not exist: " <> src diff --git a/src/LLM/Tools/FileInfo.hs b/src/LLM/Tools/FileInfo.hs --- a/src/LLM/Tools/FileInfo.hs +++ b/src/LLM/Tools/FileInfo.hs @@ -10,11 +10,11 @@ import Data.ByteString qualified as BS import Data.Text (Text) import Data.Text qualified as T -import Data.Text.IO qualified as TIO import Data.Time.Clock (UTCTime) import GHC.Generics (Generic) import LLM.Core.Types (TypedTool (..)) import LLM.Tools.FsConfig (FsConfig, sandboxPath) +import LLM.Tools.FsLimits (readTextFileLenient) import System.Directory ( Permissions (..), doesDirectoryExist, @@ -219,7 +219,7 @@ safeLineCount :: FilePath -> IO (Maybe Int) safeLineCount p = do - r <- try (TIO.readFile p) :: IO (Either IOException Text) + r <- try (readTextFileLenient p) :: IO (Either IOException Text) pure $ case r of Right content -> Just (length (T.lines content)) Left _ -> Nothing diff --git a/src/LLM/Tools/FindFiles.hs b/src/LLM/Tools/FindFiles.hs --- a/src/LLM/Tools/FindFiles.hs +++ b/src/LLM/Tools/FindFiles.hs @@ -131,24 +131,26 @@ [Text] -> Int -> IO ([FilePath], Int, Bool) -walk root includeHidden kind pat cap = go root [] 0 0 [] False +walk root includeHidden kind pat cap = do + WalkState acc _ scanned truncated <- go root [] 0 (WalkState [] 0 0 False) + pure (reverse acc, scanned, truncated) where -- relSegs is the path from root to current dir, in reverse order -- depth is current directory depth (0 = root) - go dir relSegs depth scanned acc truncated - | length acc >= cap = pure (reverse acc, scanned, True) - | depth > maxDepth = pure (reverse acc, scanned, truncated) + go dir relSegs depth st@(WalkState acc n scanned _) + | n >= cap = pure (WalkState acc n scanned True) + | depth > maxDepth = pure st | otherwise = do entries <- sort <$> safeListDirectory dir let visible = if includeHidden then entries else filter (not . isFileHidden) entries - step visible dir relSegs depth scanned acc truncated + step visible dir relSegs depth st - step [] _ _ _ scanned acc truncated = pure (reverse acc, scanned, truncated) - step (name : rest) dir relSegs depth scanned acc truncated - | length acc >= cap = pure (reverse acc, scanned + 1, True) + step [] _ _ _ st = pure st + step (name : rest) dir relSegs depth (WalkState acc n scanned truncated) + | n >= cap = pure (WalkState acc n (scanned + 1) True) | otherwise = do let full = dir name entryRel = reverse (T.pack name : relSegs) @@ -156,7 +158,7 @@ -- Skip symlinks: they could point outside the sandbox. isLink <- isSymlink full if isLink - then step rest dir relSegs depth (scanned + 1) acc truncated + then step rest dir relSegs depth (WalkState acc n (scanned + 1) truncated) else do isDir <- doesDirectoryExist full let entryKind = if isDir then DirectoryOnly else FileOnly @@ -165,13 +167,19 @@ FileOnly -> entryKind == FileOnly DirectoryOnly -> entryKind == DirectoryOnly matched = kindOk && matchGlob pat entryRel - acc' = if matched then relPath : acc else acc - (descendAcc, descendScanned, descendTrunc) <- + st' + | matched = WalkState (relPath : acc) (n + 1) (scanned + 1) truncated + | otherwise = WalkState acc n (scanned + 1) truncated + st'' <- if isDir - then go full (T.pack name : relSegs) (depth + 1) (scanned + 1) acc' truncated - else pure (reverse acc', scanned + 1, truncated) - step rest dir relSegs depth descendScanned (reverse descendAcc) descendTrunc + then go full (T.pack name : relSegs) (depth + 1) st' + else pure st' + step rest dir relSegs depth st'' +-- | Walk state: matches so far (newest first), their count, entries +-- inspected, and whether the cap truncated the output. +data WalkState = WalkState [FilePath] Int Int Bool + safeListDirectory :: FilePath -> IO [FilePath] safeListDirectory path = do r <- try (listDirectory path) :: IO (Either IOException [FilePath]) @@ -207,17 +215,20 @@ -- | Match a single glob segment (no '/' allowed in either side). -- Supports '*' (any run of chars, possibly empty) and '?' (one char). +-- On a mismatch only the most recent '*' takes one more char, which is +-- enough for this glob syntax and keeps matching O(pattern * name) instead +-- of exponential in the number of stars. matchSegment :: String -> String -> Bool -matchSegment [] [] = True -matchSegment ('*' : ps) ss = - -- '*' matches zero or more chars; try each split - matchSegment ps ss - || case ss of - [] -> False - (_ : rest) -> matchSegment ('*' : ps) rest -matchSegment ('?' : ps) (_ : ss) = matchSegment ps ss -matchSegment (p : ps) (s : ss) = p == s && matchSegment ps ss -matchSegment _ _ = False +matchSegment pat0 str0 = go pat0 str0 Nothing + where + -- lastStar: the pattern after the most recent '*' and the input it + -- was last retried from. + go ('*' : ps) ss _ = go ps ss (Just (ps, ss)) + go (p : ps) (s : ss) lastStar + | p == '?' || p == s = go ps ss lastStar + go [] [] _ = True + go _ _ (Just (ps, _ : ss)) = go ps ss (Just (ps, ss)) + go _ _ _ = False -- --------------------------------------------------------------------------- -- Rendering diff --git a/src/LLM/Tools/FsConfig.hs b/src/LLM/Tools/FsConfig.hs --- a/src/LLM/Tools/FsConfig.hs +++ b/src/LLM/Tools/FsConfig.hs @@ -5,6 +5,7 @@ mkFsConfig, sandboxPath, sandboxWritePath, + sandboxSymlink, isSandboxRoot, isFileHidden, isSymlink, @@ -83,6 +84,27 @@ -- 'directory' package. verifyNoSymlinksInChain base (joinPath baseComps) tailComps pure canonical + +-- | If @relPath@ names a symbolic link, return the link's own path +-- (not its target); otherwise 'Nothing'. +-- +-- 'sandboxPath' resolves links, so a tool that checks 'isSymlink' on its +-- result never sees one. Destructive tools call this first instead. Only +-- the directory containing the link is resolved and checked against the +-- sandbox (throwing 'SandboxViolation' on escape); the link's target may +-- lie anywhere, since acting on the link itself never touches it. +sandboxSymlink :: FsConfig -> FilePath -> IO (Maybe FilePath) +sandboxSymlink cfg relPath = do + let candidate = collapseDots (normalise (cfg.fsBasePath relPath)) + (parentRaw, name) = splitFileName candidate + isLink <- if null name then pure False else isSymlink candidate + if not isLink + then pure Nothing + else do + parent <- sandboxPath cfg (dropTrailingPathSeparator parentRaw) + let link = parent name + stillLink <- isSymlink link + pure (if stillLink then Just link else Nothing) -- | True when a path already resolved by 'sandboxPath' is the workspace -- root itself (e.g. the model passed @\".\"@, @\"\"@ or @\"sub\/..\"@). diff --git a/src/LLM/Tools/FsLimits.hs b/src/LLM/Tools/FsLimits.hs --- a/src/LLM/Tools/FsLimits.hs +++ b/src/LLM/Tools/FsLimits.hs @@ -24,6 +24,9 @@ binarySniffBytes, detectBinary, readBoundedTextFile, + readBoundedUtf8File, + readTextFileLenient, + writeUtf8File, ) where @@ -31,7 +34,7 @@ import Data.ByteString qualified as BS import Data.Text (Text) import Data.Text qualified as T -import Data.Text.Encoding (decodeUtf8With) +import Data.Text.Encoding (decodeUtf8', decodeUtf8With, encodeUtf8) import Data.Text.Encoding.Error (lenientDecode) import System.IO (IOMode (ReadMode), hFileSize, withBinaryFile) @@ -81,10 +84,36 @@ Right bs -> pure (BS.elem 0 bs) Left _ -> pure True --- | Read a text file up to @cap@ bytes. Refuses binary files, checks size --- before reading, and performs a bounded read to mitigate TOCTOU growth. +-- | Read a text file up to @cap@ bytes for display. Invalid UTF-8 bytes +-- become U+FFFD, so never write the result back; edit tools use +-- 'readBoundedUtf8File'. readBoundedTextFile :: FilePath -> Text -> Int -> IO (Either Text Text) -readBoundedTextFile path display cap = do +readBoundedTextFile path display cap = + fmap (decodeUtf8With lenientDecode) <$> readBoundedBytes path display cap + +-- | Read a text file up to @cap@ bytes for editing. Refuses files that are +-- not valid UTF-8, since writing them back would destroy their bytes. +readBoundedUtf8File :: FilePath -> Text -> Int -> IO (Either Text Text) +readBoundedUtf8File path display cap = do + r <- readBoundedBytes path display cap + pure $ + r >>= \bs -> case decodeUtf8' bs of + Right t -> Right t + Left _ -> Left $ display <> " is not valid UTF-8; refusing to edit it." + +-- | Read a whole file as UTF-8, replacing invalid bytes with U+FFFD. +-- Independent of the locale. For read-only use. +readTextFileLenient :: FilePath -> IO Text +readTextFileLenient path = decodeUtf8With lenientDecode <$> BS.readFile path + +-- | Write text as UTF-8, independent of the locale. +writeUtf8File :: FilePath -> Text -> IO () +writeUtf8File path = BS.writeFile path . encodeUtf8 + +-- | Read a file up to @cap@ bytes. Refuses binary files, checks size +-- before reading, and performs a bounded read to mitigate TOCTOU growth. +readBoundedBytes :: FilePath -> Text -> Int -> IO (Either Text BS.ByteString) +readBoundedBytes path display cap = do isBinary <- detectBinary path if isBinary then @@ -135,4 +164,4 @@ <> T.pack (show cap) <> "-byte cap during the read." | otherwise -> - pure (Right (decodeUtf8With lenientDecode bs)) + pure (Right bs) diff --git a/src/LLM/Tools/Grep.hs b/src/LLM/Tools/Grep.hs --- a/src/LLM/Tools/Grep.hs +++ b/src/LLM/Tools/Grep.hs @@ -11,11 +11,10 @@ import Data.List (sort) import Data.Text (Text) import Data.Text qualified as T -import Data.Text.IO qualified as TIO import GHC.Generics (Generic) import LLM.Core.Types (TypedTool (..)) import LLM.Tools.FsConfig (FsConfig, isFileHidden, isSymlink, sandboxPath) -import LLM.Tools.FsLimits (binarySniffBytes, expensiveWalkDepth) +import LLM.Tools.FsLimits (binarySniffBytes, expensiveWalkDepth, readTextFileLenient) import System.Directory ( doesDirectoryExist, doesFileExist, @@ -117,15 +116,26 @@ resolved <- sandboxPath cfg rel isDir <- doesDirectoryExist resolved isFile <- doesFileExist resolved + -- An explicit file gets the same size cap as files found by the walk. + sizeOk <- if isFile then isSizeOk resolved else pure True if not (isDir || isFile) then pure $ "Error: path does not exist: " <> args._grepPath - else do - files <- - if isFile - then pure [resolved] - else collectFiles resolved 0 exts - (matches, scanned, truncated) <- searchFiles resolved files pat ci cap - pure $ renderResults args._grepPath pat matches scanned truncated + else + if not sizeOk + then + pure $ + "Error: " + <> args._grepPath + <> " exceeds the " + <> T.pack (show maxFileSizeBytes) + <> "-byte cap for grep." + else do + files <- + if isFile + then pure [resolved] + else collectFiles resolved 0 exts + (matches, scanned, truncated) <- searchFiles resolved files pat ci cap + pure $ renderResults args._grepPath pat matches scanned truncated -- | Lowercase and strip leading dots so callers can pass "hs" or ".hs". normalizeExtensions :: [Text] -> [Text] @@ -207,7 +217,7 @@ then go rest acc (scanned + 1) truncated else do r <- - try (TIO.readFile f) :: + try (readTextFileLenient f) :: IO (Either IOException Text) case r of Left _ -> go rest acc (scanned + 1) truncated diff --git a/src/LLM/Tools/MoveFile.hs b/src/LLM/Tools/MoveFile.hs --- a/src/LLM/Tools/MoveFile.hs +++ b/src/LLM/Tools/MoveFile.hs @@ -2,11 +2,12 @@ import Autodocodec qualified as AC import Data.Aeson (FromJSON) +import Data.Maybe (isJust) import Data.Text (Text) import Data.Text qualified as T import GHC.Generics (Generic) import LLM.Core.Types (TypedTool (..)) -import LLM.Tools.FsConfig (FsConfig, isSandboxRoot, isSymlink, sandboxPath, sandboxWritePath) +import LLM.Tools.FsConfig (FsConfig, isSandboxRoot, sandboxPath, sandboxSymlink, sandboxWritePath) import System.Directory (doesDirectoryExist, doesFileExist, doesPathExist, renameFile) data MoveFileToolArgs = MoveFileToolArgs @@ -38,15 +39,16 @@ moveFileExecTyped cfg args = do let src = args._mfSrc dst = args._mfDst - srcResolved <- sandboxPath cfg (T.unpack src) -- Refuse to operate on symbolic-link sources: this would either move -- the link itself (surprising) or, after `sandboxPath` resolution, - -- potentially target a file outside the sandbox. - srcIsLink <- isSymlink srcResolved + -- move the link's target. Checked before resolving, since + -- `sandboxPath` follows the link. + srcLink <- sandboxSymlink cfg (T.unpack src) + srcResolved <- sandboxPath cfg (T.unpack src) if isSandboxRoot cfg srcResolved then pure $ "Error: refusing to move the workspace root: " <> src else - if srcIsLink + if isJust srcLink then pure $ "Error: refusing to move symbolic link source: " <> src else do srcExists <- doesPathExist srcResolved diff --git a/src/LLM/Tools/MultiReplaceInFile.hs b/src/LLM/Tools/MultiReplaceInFile.hs --- a/src/LLM/Tools/MultiReplaceInFile.hs +++ b/src/LLM/Tools/MultiReplaceInFile.hs @@ -4,11 +4,10 @@ import Data.Aeson (FromJSON) import Data.Text (Text) import Data.Text qualified as T -import Data.Text.IO qualified as TIO import GHC.Generics (Generic) import LLM.Core.Types (TypedTool (..)) import LLM.Tools.FsConfig (FsConfig, sandboxPath) -import LLM.Tools.FsLimits (maxReadBytes, readBoundedTextFile) +import LLM.Tools.FsLimits (maxReadBytes, readBoundedUtf8File, writeUtf8File) data Replacement = Replacement { _repOld :: Text, @@ -59,13 +58,13 @@ replaceExecTyped cfg args = do let MultiReplaceInFileToolArgs {_mrifPath, _mrifReplacements} = args resolved <- sandboxPath cfg (T.unpack _mrifPath) - contentE <- readBoundedTextFile resolved _mrifPath maxReadBytes + contentE <- readBoundedUtf8File resolved _mrifPath maxReadBytes case contentE of Left err -> pure $ "Error: " <> err Right content -> case applyReplacements _mrifReplacements content of Left err -> pure err Right result -> do - TIO.writeFile resolved result + writeUtf8File resolved result pure $ "Successfully applied " <> T.pack (show (length _mrifReplacements)) diff --git a/src/LLM/Tools/ReadFilePaginated.hs b/src/LLM/Tools/ReadFilePaginated.hs --- a/src/LLM/Tools/ReadFilePaginated.hs +++ b/src/LLM/Tools/ReadFilePaginated.hs @@ -6,22 +6,19 @@ import Autodocodec qualified as AC import Data.Aeson (FromJSON) +import Data.ByteString qualified as BS import Data.Text (Text) import Data.Text qualified as T -import Data.Text.IO qualified as TIO +import Data.Text.Encoding (decodeUtf8With) +import Data.Text.Encoding.Error (lenientDecode) import GHC.Generics (Generic) import LLM.Core.Types (TypedTool (..)) import LLM.Tools.FsConfig (FsConfig, sandboxPath) import LLM.Tools.FsLimits (detectBinary, maxPaginatedFileBytes) import System.IO - ( Handle, - IOMode (ReadMode), + ( IOMode (ReadMode), hFileSize, - hIsEOF, - hSetEncoding, - utf8, withBinaryFile, - withFile, ) -- | Default number of lines returned per page when the caller omits 'limit'. @@ -104,38 +101,15 @@ (lns, hasMore) <- readPage resolved offset limit pure $ renderPage args._rfpPath size offset lns hasMore --- | Skip @offset - 1@ lines, take up to @limit@ lines, then probe one --- past the last returned line to determine whether more content remains. --- Reads strictly within 'withFile' so the handle is released eagerly. +-- | Take up to @limit@ lines starting at line @offset@, and report whether +-- more lines follow. Reads at most 'maxPaginatedFileBytes' (the file may +-- have grown since the size check) and decodes leniently, so invalid UTF-8 +-- bytes show as U+FFFD instead of failing the read. readPage :: FilePath -> Int -> Int -> IO ([Text], Bool) -readPage path offset limit = withFile path ReadMode $ \h -> do - hSetEncoding h utf8 - skipLines h (offset - 1) - collected <- takeLines h limit - more <- not <$> hIsEOF h - pure (collected, more) - -skipLines :: Handle -> Int -> IO () -skipLines _ 0 = pure () -skipLines h n = do - eof <- hIsEOF h - if eof - then pure () - else do - _ <- TIO.hGetLine h - skipLines h (n - 1) - -takeLines :: Handle -> Int -> IO [Text] -takeLines h n0 = go n0 id - where - go 0 acc = pure (acc []) - go k acc = do - eof <- hIsEOF h - if eof - then pure (acc []) - else do - line <- TIO.hGetLine h - go (k - 1) (acc . (line :)) +readPage path offset limit = do + bytes <- withBinaryFile path ReadMode (`BS.hGet` fromIntegral maxPaginatedFileBytes) + let rest = drop (offset - 1) (T.lines (decodeUtf8With lenientDecode bytes)) + pure (take limit rest, not (null (drop limit rest))) renderPage :: Text -> Integer -> Int -> [Text] -> Bool -> Text renderPage path size offset lns hasMore = diff --git a/src/LLM/Tools/RemoveDirectory.hs b/src/LLM/Tools/RemoveDirectory.hs --- a/src/LLM/Tools/RemoveDirectory.hs +++ b/src/LLM/Tools/RemoveDirectory.hs @@ -7,7 +7,7 @@ import Data.Text qualified as T import GHC.Generics (Generic) import LLM.Core.Types (TypedTool (..)) -import LLM.Tools.FsConfig (FsConfig, isSandboxRoot, isSymlink, sandboxPath) +import LLM.Tools.FsConfig (FsConfig, isSandboxRoot, isSymlink, sandboxPath, sandboxSymlink) import System.Directory ( doesDirectoryExist, listDirectory, @@ -34,7 +34,8 @@ { ttoolName = "remove_directory", ttoolDescription = "Recursively delete a directory and all its contents at the given path (relative to the workspace). " - <> "Use with caution — this operation is irreversible.", + <> "Use with caution — this operation is irreversible. " + <> "A symbolic link is removed itself; its target is left unchanged.", ttoolReadonly = False, ttoolExecute = const (removeDirectoryExecTyped cfg) } @@ -42,12 +43,18 @@ removeDirectoryExecTyped :: FsConfig -> RemoveDirectoryToolArgs -> IO Text removeDirectoryExecTyped cfg args = do let p = args._rdPath - resolved <- sandboxPath cfg (T.unpack p) - if isSandboxRoot cfg resolved - then pure $ "Error: refusing to remove the workspace root: " <> p - else do - safeRemoveDirectoryRecursive resolved - pure $ "Successfully removed directory " <> p + link <- sandboxSymlink cfg (T.unpack p) + case link of + Just l -> do + removeFile l + pure $ "Successfully removed symbolic link " <> p <> " (its target was left unchanged)" + Nothing -> do + resolved <- sandboxPath cfg (T.unpack p) + if isSandboxRoot cfg resolved + then pure $ "Error: refusing to remove the workspace root: " <> p + else do + safeRemoveDirectoryRecursive resolved + pure $ "Successfully removed directory " <> p -- | Like 'System.Directory.removeDirectoryRecursive', but refuses to -- follow directory symlinks. A symbolic link encountered during the @@ -55,10 +62,10 @@ -- the sandbox pointing at an external directory cannot be used to -- delete files outside the sandbox. -- --- Note: the top-level argument has already been resolved through --- 'sandboxPath', so it is guaranteed to be a real directory inside --- the sandbox at the moment of the call. We still re-check it --- defensively in case of an unusual TOCTOU race. +-- Note: a top-level link is handled by the caller via 'sandboxSymlink', +-- and the argument has been resolved through 'sandboxPath', so it is a +-- real directory inside the sandbox at the moment of the call. We still +-- re-check it defensively in case of an unusual TOCTOU race. safeRemoveDirectoryRecursive :: FilePath -> IO () safeRemoveDirectoryRecursive dir = do topIsLink <- isSymlink dir diff --git a/src/LLM/Tools/RemoveFile.hs b/src/LLM/Tools/RemoveFile.hs --- a/src/LLM/Tools/RemoveFile.hs +++ b/src/LLM/Tools/RemoveFile.hs @@ -6,7 +6,7 @@ import Data.Text qualified as T import GHC.Generics (Generic) import LLM.Core.Types (TypedTool (..)) -import LLM.Tools.FsConfig (FsConfig, isSandboxRoot, sandboxPath) +import LLM.Tools.FsConfig (FsConfig, isSandboxRoot, sandboxPath, sandboxSymlink) import System.Directory (doesDirectoryExist, removeFile) newtype RemoveFileToolArgs = RemoveFileToolArgs @@ -26,7 +26,8 @@ { ttoolName = "remove_file", ttoolDescription = "Delete a file at the given path (relative to the workspace). " - <> "Fails if the path does not exist or is a directory.", + <> "Fails if the path does not exist or is a directory. " + <> "A symbolic link is removed itself; its target is left unchanged.", ttoolReadonly = False, ttoolExecute = const (removeFileExecTyped cfg) } @@ -34,13 +35,19 @@ removeFileExecTyped :: FsConfig -> RemoveFileToolArgs -> IO Text removeFileExecTyped cfg args = do let p = args._rftPath - resolved <- sandboxPath cfg (T.unpack p) - isDir <- doesDirectoryExist resolved - if isSandboxRoot cfg resolved - then pure $ "Error: refusing to remove the workspace root: " <> p - else - if isDir - then pure $ "Error: path is a directory, not a file (use remove_directory): " <> p - else do - removeFile resolved - pure $ "Successfully removed " <> p + link <- sandboxSymlink cfg (T.unpack p) + case link of + Just l -> do + removeFile l + pure $ "Successfully removed symbolic link " <> p <> " (its target was left unchanged)" + Nothing -> do + resolved <- sandboxPath cfg (T.unpack p) + isDir <- doesDirectoryExist resolved + if isSandboxRoot cfg resolved + then pure $ "Error: refusing to remove the workspace root: " <> p + else + if isDir + then pure $ "Error: path is a directory, not a file (use remove_directory): " <> p + else do + removeFile resolved + pure $ "Successfully removed " <> p diff --git a/src/LLM/Tools/ReplaceInFile.hs b/src/LLM/Tools/ReplaceInFile.hs --- a/src/LLM/Tools/ReplaceInFile.hs +++ b/src/LLM/Tools/ReplaceInFile.hs @@ -4,11 +4,10 @@ import Data.Aeson (FromJSON) import Data.Text (Text) import Data.Text qualified as T -import Data.Text.IO qualified as TIO import GHC.Generics (Generic) import LLM.Core.Types (TypedTool (..)) import LLM.Tools.FsConfig (FsConfig, sandboxPath) -import LLM.Tools.FsLimits (maxReadBytes, readBoundedTextFile) +import LLM.Tools.FsLimits (maxReadBytes, readBoundedUtf8File, writeUtf8File) data ReplaceInFileToolArgs = ReplaceInFileToolArgs { _rifPath :: Text, @@ -47,14 +46,14 @@ let ReplaceInFileToolArgs {_rifPath, _rifOld, _rifNew} = args p = _rifPath resolved <- sandboxPath cfg (T.unpack p) - contentE <- readBoundedTextFile resolved p maxReadBytes + contentE <- readBoundedUtf8File resolved p maxReadBytes case contentE of Left err -> pure $ "Error: " <> err Right content -> case countOccurrences _rifOld content of 0 -> pure "Error: the 'old' string was not found in the file" 1 -> do let replaced = replaceFirst _rifOld _rifNew content - TIO.writeFile resolved replaced + writeUtf8File resolved replaced pure $ "Successfully replaced text in " <> p n -> pure $ diff --git a/src/LLM/Tools/Writefile.hs b/src/LLM/Tools/Writefile.hs --- a/src/LLM/Tools/Writefile.hs +++ b/src/LLM/Tools/Writefile.hs @@ -4,10 +4,10 @@ import Data.Aeson (FromJSON) import Data.Text (Text) import Data.Text qualified as T -import Data.Text.IO qualified as TIO import GHC.Generics (Generic) import LLM.Core.Types (TypedTool (..)) import LLM.Tools.FsConfig (FsConfig, sandboxWritePath) +import LLM.Tools.FsLimits (writeUtf8File) data WritefileToolArgs = WritefileToolArgs { _wfPath :: Text, @@ -41,5 +41,5 @@ let p = args._wfPath c = args._wfContent resolved <- sandboxWritePath cfg (T.unpack p) - TIO.writeFile resolved c + writeUtf8File resolved c pure $ "Successfully wrote " <> T.pack (show (T.length c)) <> " characters to " <> p diff --git a/test/LLM/ChatSpec.hs b/test/LLM/ChatSpec.hs --- a/test/LLM/ChatSpec.hs +++ b/test/LLM/ChatSpec.hs @@ -33,6 +33,7 @@ LLMGateway (..), LLMHooks (..), ToolDef (ToolDef, toolDescription, toolName, toolParameters, toolReadonly), + ToolResult (..), mkToolCall, ) import LLM.Core.Usage (PricingInfo (..), defaultPricingInfo, mkUsage) @@ -375,7 +376,11 @@ models = ModelWithFallbacks (mockModel infiniteToolGateway) [] result <- runGenerate agent models toolMap Nothing [UserTurn "test"] case result of - Left GenerateErrorResult {gerError = GErrToolExceeded} -> pure () + Left GenerateErrorResult {gerError = GErrToolExceeded, gerPartialNewMessages = msgs} -> + -- The last round's results are returned, though the model never saw them. + case reverse msgs of + ToolTurn [_] : AssistantMessage _ : _ -> length msgs `shouldBe` 4 + _ -> expectationFailure $ "Expected the last round's ToolTurn last, got: " <> show msgs other -> expectationFailure $ "Expected GErrToolExceeded, got: " <> show other it "falls back to next model on retryable error" $ do @@ -541,7 +546,14 @@ models = ModelWithFallbacks (mockModel twoCallGw) [] result <- runGenerate agent models tm (Just sig) [UserTurn "go"] case result of - Left GenerateErrorResult {gerError = GErrAborted} -> pure () + Left GenerateErrorResult {gerError = GErrAborted, gerPartialNewMessages} -> + -- Every tool call in the history must have a result, or the next + -- request is rejected. c1 ran; c2 never did. + case reverse gerPartialNewMessages of + ToolTurn results : _ -> + [(r.trCallId, r.trContent) | r <- results] + `shouldBe` [("c1", "done"), ("c2", "Aborted before execution")] + _ -> expectationFailure $ "Expected history to end with a ToolTurn, got: " <> show gerPartialNewMessages other -> expectationFailure $ "Expected GErrAborted during tools, got: " <> show other describe "sticky fallback" $ stickyFallbackSpec runGenerate toolMap diff --git a/test/LLM/ClaudeSpec.hs b/test/LLM/ClaudeSpec.hs --- a/test/LLM/ClaudeSpec.hs +++ b/test/LLM/ClaudeSpec.hs @@ -74,6 +74,50 @@ hasToolCalls resp `shouldBe` False Left err -> expectationFailure $ "Parse failed: " <> show err + it "skips unknown content block types" $ do + let val = + object + [ "model" .= ("claude-haiku-4-5-20251001" :: Text), + "stop_reason" .= ("end_turn" :: Text), + "content" + .= [ object + [ "type" .= ("server_tool_use" :: Text), + "id" .= ("srvtoolu_1" :: Text), + "name" .= ("web_search" :: Text), + "input" .= object ["query" .= ("weather" :: Text)] + ], + object + [ "type" .= ("web_search_tool_result" :: Text), + "tool_use_id" .= ("srvtoolu_1" :: Text), + "content" .= ([] :: [Value]) + ], + object ["type" .= ("text" :: Text), "text" .= ("Sunny." :: Text)] + ] + ] + case parseClaudeResponse val of + Right resp -> do + resp.respText `shouldBe` "Sunny." + hasToolCalls resp `shouldBe` False + Left err -> expectationFailure $ "Parse failed: " <> show err + + it "keeps signed thinking blocks separate and projects all their reasoning" $ do + let block t s = object ["type" .= ("thinking" :: Text), "thinking" .= (t :: Text), "signature" .= (s :: Text)] + val = + object + [ "model" .= ("claude-haiku-4-5-20251001" :: Text), + "stop_reason" .= ("end_turn" :: Text), + "content" + .= [ block "First thought. " "sig1", + block "Second thought." "sig2", + object ["type" .= ("text" :: Text), "text" .= ("Done." :: Text)] + ] + ] + case parseClaudeResponse val of + Right resp -> do + resp.respReasoning `shouldBe` Just "First thought. Second thought." + length [() | ContentPart (ThinkingPart _) _ <- resp.respContent] `shouldBe` 2 + Left err -> expectationFailure $ "Parse failed: " <> show err + it "parses a tool_use response" $ do Right val <- eitherDecodeFileStrict' "test/fixtures/claude-tool-use.json" case parseClaudeResponse val of @@ -250,6 +294,20 @@ lookupText "name" (Object toolO) `shouldBe` Just "get_weather" KM.lookup "thoughtSignature" toolO `shouldBe` Nothing _ -> expectationFailure "expected text + tool_use only" + + it "drops an assistant turn whose content encodes to nothing" $ do + -- A DeepSeek turn cut off during reasoning: foreign thinking without a + -- Claude signature, and empty text. "content": [] is a 400. + let turn = + AssistantMessage + [ thinkingPart (ThinkingContent (Just "deepseek reasoning") Nothing), + textPart "" + ] + encodeTurn "claude-haiku-4-5-20251001" turn `shouldBe` [] + + it "drops a user turn whose content encodes to nothing" $ do + let turn = UserMessage [thinkingPart (ThinkingContent (Just "stray") Nothing)] + encodeTurn "claude-haiku-4-5-20251001" turn `shouldBe` [] describe "thinking request mapping" $ do it "maps effort to budget_tokens and omits temperature when thinking is on" $ do diff --git a/test/LLM/DeepSeekSpec.hs b/test/LLM/DeepSeekSpec.hs --- a/test/LLM/DeepSeekSpec.hs +++ b/test/LLM/DeepSeekSpec.hs @@ -39,6 +39,16 @@ msg = head (encodeTurn deepSeekMessageEncodeOptions turn) lookupText "reasoning_content" msg `shouldBe` Just "I should call the weather tool." + it "replays the reasoning of every thinking part, not just the first" $ do + let turn = + AssistantMessage + [ ContentPart (ThinkingPart (ThinkingContent (Just "first, ") Nothing)) Nothing, + ContentPart (ThinkingPart (ThinkingContent (Just "second") Nothing)) Nothing, + textPart "Done." + ] + msg = head (encodeTurn deepSeekMessageEncodeOptions turn) + lookupText "reasoning_content" msg `shouldBe` Just "first, second" + it "omits reasoning_content for OpenAI-compatible default encoding" $ do let turn = assistantTurn "Hello" (Just "thinking") [] msg = head (encodeTurn defaultMessageEncodeOptions turn) diff --git a/test/LLM/FsToolsSpec.hs b/test/LLM/FsToolsSpec.hs --- a/test/LLM/FsToolsSpec.hs +++ b/test/LLM/FsToolsSpec.hs @@ -1,9 +1,13 @@ module LLM.FsToolsSpec (spec) where +import Control.Exception (bracket_, evaluate) import Data.Aeson (Value, object, (.=)) +import Data.ByteString qualified as BS import Data.Map qualified as Map import Data.Text (Text) import Data.Text qualified as T +import Data.Text.Encoding (encodeUtf8) +import GHC.IO.Encoding (getLocaleEncoding, setLocaleEncoding) import Heptapod (generate) import LLM.Agent.Types (RuntimeArgs (..), Tool (..), ToolContext (..)) import LLM.Generate.GenerateUtils (llmHooks) @@ -11,10 +15,16 @@ import LLM.Load.FsTools (fsTools) import System.Directory ( createDirectory, + createDirectoryLink, + createFileLink, doesDirectoryExist, doesFileExist, + doesPathExist, + pathIsSymbolicLink, ) -import System.FilePath (()) +import System.FilePath (takeDirectory, ()) +import System.IO (TextEncoding, char8) +import GHC.Clock (getMonotonicTime) import System.IO.Temp (withSystemTempDirectory) import Test.Hspec ( Spec, @@ -24,6 +34,7 @@ shouldContain, shouldNotContain, shouldReturn, + shouldSatisfy, ) spec :: Spec @@ -164,6 +175,132 @@ runTool_ tm "remove_directory" (object ["path" .= ("tmp" :: Text)]) doesDirectoryExist (sb "tmp") `shouldReturn` False + describe "symbolic links" $ do + it "remove_file removes the link and leaves its target intact" $ withSandbox $ \sb -> do + tm <- fsTools sb + writeFile (sb "real.txt") "keep" + createFileLink "real.txt" (sb "link") + result <- runTool tm "remove_file" (object ["path" .= ("link" :: Text)]) + T.unpack result `shouldContain` "symbolic link" + doesPathExist (sb "link") `shouldReturn` False + readFile (sb "real.txt") `shouldReturn` "keep" + + it "remove_file removes a link that points outside the workspace" $ withSandbox $ \sb -> do + tm <- fsTools sb + let outside = takeDirectory sb "outside.txt" + writeFile outside "keep" + createFileLink outside (sb "link") + runTool_ tm "remove_file" (object ["path" .= ("link" :: Text)]) + doesPathExist (sb "link") `shouldReturn` False + readFile outside `shouldReturn` "keep" + + it "remove_directory removes the link and leaves the linked directory intact" $ withSandbox $ \sb -> do + tm <- fsTools sb + createDirectory (sb "real") + writeFile (sb "real" "x.txt") "keep" + createDirectoryLink "real" (sb "link") + result <- runTool tm "remove_directory" (object ["path" .= ("link" :: Text)]) + T.unpack result `shouldContain` "symbolic link" + doesPathExist (sb "link") `shouldReturn` False + readFile (sb "real" "x.txt") `shouldReturn` "keep" + + it "move_file refuses a link source" $ withSandbox $ \sb -> do + tm <- fsTools sb + writeFile (sb "real.txt") "keep" + createFileLink "real.txt" (sb "link") + result <- runTool tm "move_file" (object ["source" .= ("link" :: Text), "destination" .= ("moved.txt" :: Text)]) + T.unpack result `shouldContain` "symbolic link" + pathIsSymbolicLink (sb "link") `shouldReturn` True + readFile (sb "real.txt") `shouldReturn` "keep" + doesPathExist (sb "moved.txt") `shouldReturn` False + + it "copy_file refuses a link source" $ withSandbox $ \sb -> do + tm <- fsTools sb + writeFile (sb "real.txt") "keep" + createFileLink "real.txt" (sb "link") + result <- runTool tm "copy_file" (object ["source" .= ("link" :: Text), "destination" .= ("copy.txt" :: Text)]) + T.unpack result `shouldContain` "symbolic link" + doesPathExist (sb "copy.txt") `shouldReturn` False + + describe "text encoding" $ do + -- "café foo\n" in Latin-1: 0xE9 is not valid UTF-8. + let latin1 = BS.pack [0x63, 0x61, 0x66, 0xE9, 0x20, 0x66, 0x6F, 0x6F, 0x0A] + it "replace_in_file refuses a non-UTF-8 file and leaves its bytes unchanged" $ withSandbox $ \sb -> do + tm <- fsTools sb + BS.writeFile (sb "latin1.txt") latin1 + result <- runTool tm "replace_in_file" (object ["path" .= ("latin1.txt" :: Text), "old" .= ("foo" :: Text), "new" .= ("bar" :: Text)]) + T.unpack result `shouldContain` "UTF-8" + BS.readFile (sb "latin1.txt") `shouldReturn` latin1 + + it "multi_replace_in_file refuses a non-UTF-8 file and leaves its bytes unchanged" $ withSandbox $ \sb -> do + tm <- fsTools sb + BS.writeFile (sb "latin1.txt") latin1 + result <- + runTool + tm + "multi_replace_in_file" + ( object + [ "path" .= ("latin1.txt" :: Text), + "replacements" .= [object ["old" .= ("foo" :: Text), "new" .= ("bar" :: Text)]] + ] + ) + T.unpack result `shouldContain` "UTF-8" + BS.readFile (sb "latin1.txt") `shouldReturn` latin1 + + it "writefile writes UTF-8 regardless of the locale" $ withSandbox $ \sb -> do + tm <- fsTools sb + withLocaleEncoding char8 $ + runTool_ tm "writefile" (object ["path" .= ("u.txt" :: Text), "content" .= ("äöü" :: Text)]) + BS.readFile (sb "u.txt") `shouldReturn` encodeUtf8 "äöü" + + it "replace_in_file writes UTF-8 regardless of the locale" $ withSandbox $ \sb -> do + tm <- fsTools sb + BS.writeFile (sb "u.txt") (encodeUtf8 "x = 1\n") + withLocaleEncoding char8 $ + runTool_ tm "replace_in_file" (object ["path" .= ("u.txt" :: Text), "old" .= ("1" :: Text), "new" .= ("äöü" :: Text)]) + BS.readFile (sb "u.txt") `shouldReturn` encodeUtf8 "x = äöü\n" + + it "grep finds matches in a non-UTF-8 file" $ withSandbox $ \sb -> do + tm <- fsTools sb + BS.writeFile (sb "latin1.txt") latin1 + result <- runTool tm "grep" (object ["pattern" .= ("foo" :: Text), "path" .= ("." :: Text)]) + T.unpack result `shouldContain` "latin1.txt" + + it "read_file_paginated reads a non-UTF-8 file leniently" $ withSandbox $ \sb -> do + tm <- fsTools sb + BS.writeFile (sb "latin1.txt") latin1 + result <- runTool tm "read_file_paginated" (object ["path" .= ("latin1.txt" :: Text)]) + T.unpack result `shouldContain` "caf\xFFFD foo" + + describe "file tool limits" $ do + it "grep refuses an explicit file over the size cap" $ withSandbox $ \sb -> do + tm <- fsTools sb + BS.writeFile (sb "big.txt") (BS.replicate (1024 * 1024) 0x61 <> "needle\n") + result <- runTool tm "grep" (object ["pattern" .= ("needle" :: Text), "path" .= ("big.txt" :: Text)]) + T.unpack result `shouldContain` "Error" + T.unpack result `shouldContain` "exceeds" + + it "find_files matches '*' and '?' within a segment" $ withSandbox $ \sb -> do + tm <- fsTools sb + mapM_ (\f -> BS.writeFile (sb f) "x") ["abc.hs", "a.hs", "abxc.hs", "ab.hs", "b.hs"] + result <- runTool tm "find_files" (object ["pattern" .= ("a*b?c.hs" :: Text)]) + T.lines result `shouldContain` ["abxc.hs"] + T.unpack result `shouldNotContain` "abc.hs" + result2 <- runTool tm "find_files" (object ["pattern" .= ("*a*.hs" :: Text)]) + filter (not . T.isPrefixOf "===") (T.lines result2) `shouldBe` ["a.hs", "ab.hs", "abc.hs", "abxc.hs"] + + it "find_files does not backtrack exponentially" $ withSandbox $ \sb -> do + tm <- fsTools sb + -- The old backtracking matcher took ~5 s on this; its loop does not + -- allocate, so 'timeout' could not interrupt it. Measure instead. + BS.writeFile (sb replicate 40 'a') "x" + t0 <- getMonotonicTime + result <- runTool tm "find_files" (object ["pattern" .= ("*a*a*a*a*a*a*a*a*b" :: Text)]) + _ <- evaluate (T.length result) + t1 <- getMonotonicTime + T.unpack result `shouldContain` "0 match(es)" + (t1 - t0) `shouldSatisfy` (< 1) + describe "workspace root guard" $ do let rootPaths = [".", "", "./", "sub/.."] :: [Text] it "remove_directory refuses the workspace root" $ withSandbox $ \sb -> do @@ -208,6 +345,11 @@ result <- runTool tm "remove_file" (object ["path" .= ("sub" :: Text)]) T.unpack result `shouldContain` "is a directory" doesDirectoryExist (sb "sub") `shouldReturn` True + +withLocaleEncoding :: TextEncoding -> IO a -> IO a +withLocaleEncoding enc act = do + old <- getLocaleEncoding + bracket_ (setLocaleEncoding enc) (setLocaleEncoding old) act withSandbox :: (FilePath -> IO a) -> IO a withSandbox act = withSystemTempDirectory "llm-simple-fstools" $ \root -> do diff --git a/test/LLM/GeminiSpec.hs b/test/LLM/GeminiSpec.hs --- a/test/LLM/GeminiSpec.hs +++ b/test/LLM/GeminiSpec.hs @@ -2,20 +2,28 @@ module LLM.GeminiSpec (spec) where -import Data.Aeson (Value (Array, Object, String), eitherDecodeFileStrict', object, (.=)) +import Data.Aeson (Value (Array, Object, String), decodeStrict', eitherDecodeFileStrict', object, (.=)) import Data.Aeson.Key qualified as K import Data.Aeson.KeyMap qualified as KM +import Data.ByteString qualified as BS +import Data.Char (isAsciiLower, isAsciiUpper, isDigit) +import Data.IORef (newIORef, readIORef, writeIORef) +import Data.Maybe (isJust) import Data.Text (Text) import Data.Text qualified as T +import Data.Text.Encoding qualified as TE +import Data.UUID.Types qualified as UUID import Data.Vector qualified as V import LLM.Core.LLMProvider (LLMProvider (..)) import LLM.Core.ProviderUtils (normalizeSchemaOpenAI, stripBoundsAndComments) import LLM.Core.Types ( ChatRequest (..), - ChatResponse (respText), + ChatResponse (respContent, respReasoning, respText), + LLMError (..), ToolDef (..), ProviderOpaque (..), ThinkingContent (..), + ThinkingMode (..), ToolCall (..), Turn (..), mkToolCall, @@ -35,6 +43,7 @@ expectationFailure, it, shouldBe, + shouldSatisfy, ) spec :: Spec @@ -144,6 +153,42 @@ lookupPath ["generationConfig", "responseSchema"] body `shouldBe` Just (object ["type" .= ("integer" :: Text)]) + describe "structured output with thinking" $ do + it "decodes the object when a thought summary part comes first" $ do + let resp = + object + [ "candidates" + .= [ object + [ "content" + .= object + [ "role" .= ("model" :: Text), + "parts" + .= [ object ["text" .= ("Let me think about the answer." :: Text), "thought" .= True], + object ["text" .= ("{\"answer\": " :: Text)], + object ["text" .= ("42}" :: Text)] + ] + ] + ] + ] + ] + result <- (geminiProvider "").parseObjectResponse resp + fmap fst result `shouldBe` Right (object ["answer" .= (42 :: Int)]) + + it "does not request thought summaries for object requests" $ do + let req = + ChatRequest + { reqModel = "gemini-2.5-flash", + reqConversation = [], + reqSystem = Nothing, + reqMaxTokens = 100, + reqTemperature = Nothing, + reqTools = [], + reqThinking = Just ThinkingMode {tmEnabled = True, tmEffort = Just "1024"} + } + body = (geminiProvider "").buildObjectBody req (object ["type" .= ("object" :: Text)]) + lookupPath ["generationConfig", "thinkingConfig"] body + `shouldBe` Just (object ["thinkingBudget" .= (1024 :: Int)]) + describe "parseGeminiResponse" $ do it "parses a text response" $ do Right val <- eitherDecodeFileStrict' "test/fixtures/gemini-text.json" @@ -164,6 +209,83 @@ tc.tcName `shouldBe` "get_weather" Left err -> expectationFailure $ "Parse failed: " <> show err + -- Data.Unique ids restarted at call_1 in every process, so a reloaded + -- history could get duplicate ids. A random UUID cannot repeat. + it "gives each function call a random UUID id that Claude accepts" $ do + Right val <- eitherDecodeFileStrict' "test/fixtures/gemini-tool-use.json" + Right r <- parseGeminiResponse val + let [tc] = getToolCalls r + (UUID.fromText =<< T.stripPrefix "call_" tc.tcId) `shouldSatisfy` isJust + tc.tcId `shouldSatisfy` T.all (\c -> isAsciiUpper c || isAsciiLower c || isDigit c || c == '_' || c == '-') + + describe "streamed reasoning" $ do + it "keeps every thought chunk, not just the first" $ do + let chunk t thought = + "data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"" + <> t + <> "\"" + <> (if thought then ",\"thought\":true" else "") + <> "}]}}],\"modelVersion\":\"gemini-2.5-flash\"}\n\n" + sse = chunk "Let me " True <> chunk "think about " True <> chunk "this." True <> chunk "Answer." False + reader <- mkBodyReader sse + result <- (geminiProvider "key").parseStream (object []) reader (\_ -> pure ()) + case result of + Right resp -> do + resp.respReasoning `shouldBe` Just "Let me think about this." + resp.respText `shouldBe` "Answer." + length resp.respContent `shouldBe` 2 + Left err -> expectationFailure ("unexpected error: " <> show err) + + describe "errors and finish reasons" $ do + let stream sse = do + reader <- mkBodyReader sse + (geminiProvider "key").parseStream (object []) reader (\_ -> pure ()) + textChunk = "data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"Hel\"}]}}]}\n\n" + errorData = "{\"error\":{\"code\":503,\"message\":\"The model is overloaded.\",\"status\":\"UNAVAILABLE\"}}" + parse json = case decodeStrict' json of + Just v -> parseGeminiResponse v + Nothing -> error "bad fixture" + isParseErrorWith s = \case + Left (ParseError msg) -> s `T.isInfixOf` msg + _ -> False + + -- Shape as the Gemini API sends it with alt=sse: HTTP 200, then an + -- error object as a data event after some content. + it "turns a mid-stream error chunk into HttpError with its code" $ do + result <- stream (textChunk <> "data: " <> errorData <> "\n\n") + result `shouldBe` Left (HttpError 503 (TE.decodeUtf8 errorData)) + + it "reports a blocked prompt as ResponseBlocked" $ do + result <- stream "data: {\"promptFeedback\":{\"blockReason\":\"SAFETY\"}}\n\n" + result `shouldSatisfy` \case + Left (ResponseBlocked msg) -> "SAFETY" `T.isInfixOf` msg + _ -> False + parsed <- parse "{\"promptFeedback\":{\"blockReason\":\"PROHIBITED_CONTENT\"}}" + parsed `shouldSatisfy` \case + Left (ResponseBlocked msg) -> "PROHIBITED_CONTENT" `T.isInfixOf` msg + _ -> False + + it "reports an empty response stopped by a safety filter as ResponseBlocked" $ do + parsed <- parse "{\"candidates\":[{\"finishReason\":\"RECITATION\"}]}" + parsed `shouldSatisfy` \case + Left (ResponseBlocked msg) -> "RECITATION" `T.isInfixOf` msg + _ -> False + + it "explains an empty response cut off by the token limit" $ do + result <- stream "data: {\"candidates\":[{\"content\":{\"role\":\"model\"},\"finishReason\":\"MAX_TOKENS\"}]}\n\n" + result `shouldSatisfy` isParseErrorWith "MAX_TOKENS" + parsed <- parse "{\"candidates\":[{\"content\":{\"role\":\"model\"},\"finishReason\":\"MAX_TOKENS\"}]}" + parsed `shouldSatisfy` isParseErrorWith "MAX_TOKENS" + + it "keeps partial content cut off by the token limit" $ do + result <- stream (textChunk <> "data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"lo\"}]},\"finishReason\":\"MAX_TOKENS\"}]}\n\n") + fmap (.respText) result `shouldBe` Right "Hello" + + it "explains structured output cut off by the token limit" $ do + let Just resp = decodeStrict' "{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"{\\\"a\\\": \"}]},\"finishReason\":\"MAX_TOKENS\"}]}" + result <- (geminiProvider "").parseObjectResponse resp + result `shouldSatisfy` isParseErrorWith "MAX_TOKENS" + describe "thought signature round-trip" $ do it "replays signature only for a matching model, else sends the placeholder" $ do let tc = @@ -184,6 +306,20 @@ signatureForModel "gemini-2.5-pro" tc.tcProviderMeta `shouldBe` Just "sig-xyz" signatureForModel "gemini-2.5-flash" tc.tcProviderMeta `shouldBe` Nothing + it "treats models that share a name prefix as different models" $ do + let meta m = + Just + ProviderOpaque + { poProvider = "gemini", + poModel = Just m, + poPayload = object ["thoughtSignature" .= ("sig" :: Text)] + } + signatureForModel "gemini-2.5-flash-lite" (meta "gemini-2.5-flash") `shouldBe` Nothing + signatureForModel "gemini-2.5-flash" (meta "gemini-2.5-flash-lite") `shouldBe` Nothing + signatureForModel "gemini-2.5-flash" (meta "gemini-2.5-flash-preview-09-2025") `shouldBe` Nothing + signatureForModel "models/gemini-2.5-flash" (meta "gemini-2.5-flash") `shouldBe` Just "sig" + signatureForModel "gemini-2.5-flash" (meta "gemini-2.5-flash") `shouldBe` Just "sig" + it "omits foreign Claude thinking opaque while keeping text and tool calls" $ do let turn = AssistantMessage @@ -223,6 +359,27 @@ encoded = head (encodeTurn "gemini-3.1-flash-lite" turn) thoughtSignatures encoded `shouldBe` [skipThoughtSignatureValidator, skipThoughtSignatureValidator] + it "drops an assistant turn whose content encodes to nothing" $ do + -- A Claude turn holding only redacted thinking: the foreign opaque is + -- stripped and there is no text, so "parts": [] would be a 400. + let redacted = + ThinkingContent + { thinkingText = Nothing, + thinkingOpaque = + Just + ProviderOpaque + { poProvider = "claude", + poModel = Just "claude-haiku-4-5", + poPayload = object ["type" .= ("redacted_thinking" :: Text), "data" .= ("xyz" :: Text)] + } + } + turn = AssistantMessage [thinkingPart redacted, textPart ""] + encodeTurn "gemini-3.1-flash-lite" turn `shouldBe` [] + + it "drops a user turn whose content encodes to nothing" $ do + let turn = UserMessage [thinkingPart (ThinkingContent (Just "stray") Nothing)] + encodeTurn "gemini-3.1-flash-lite" turn `shouldBe` [] + describe "image request encoding" $ do it "encodes URL fileData and base64 inlineData parts" $ do Right b64 <- pure $ imageBase64Part "image/jpeg" "aGVsbG8=" @@ -320,6 +477,18 @@ Just (Array arr) -> V.toList arr _ -> [] assistantParts _ = [] + +-- | Yield the whole payload once, then 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 -- | Follow object keys / array indices (as decimal strings). lookupPath :: [Text] -> Value -> Maybe Value diff --git a/test/LLM/HttpConfigSpec.hs b/test/LLM/HttpConfigSpec.hs new file mode 100644 --- /dev/null +++ b/test/LLM/HttpConfigSpec.hs @@ -0,0 +1,54 @@ +module LLM.HttpConfigSpec (spec) where + +import Control.Concurrent (forkIO, killThread, threadDelay) +import Control.Exception (SomeException, bracket, finally, try) +import Control.Monad (forever, void) +import Data.Aeson (Value, object) +import Data.ByteString qualified as BS +import Data.IORef (atomicModifyIORef', newIORef, readIORef) +import LLM.Core.ProviderUtils (mkLenientConfig) +import Network.HTTP.Client qualified as HC +import Network.HTTP.Req (POST (..), ReqBodyJson (..), http, ignoreResponse, port, req, runReq, (/:)) +import Network.Socket qualified as N +import Network.Socket.ByteString qualified as NB +import Test.Hspec (Spec, describe, it, shouldBe) + +spec :: Spec +spec = describe "HTTP config" $ do + it "sends a POST exactly once when the response times out" $ do + posts <- newIORef (0 :: Int) + -- Answer each request after 2 s; the client gives up after 100 ms. + withSlowServer 2_000_000 (atomicModifyIORef' posts (\n -> (n + 1, ()))) $ \p -> do + mgr <- HC.newManager HC.defaultManagerSettings {HC.managerResponseTimeout = HC.responseTimeoutMicro 100_000} + r <- + try $ + runReq (mkLenientConfig mgr) $ + req POST (http "127.0.0.1" /: "v1") (ReqBodyJson (object [] :: Value)) ignoreResponse (port p) + case r of + Left (_ :: SomeException) -> pure () + Right _ -> fail "expected the request to time out" + readIORef posts >>= (`shouldBe` 1) + +-- | A TCP server on a free local port that counts each request it receives +-- and answers it only after the given delay. +withSlowServer :: Int -> IO () -> (Int -> IO a) -> IO a +withSlowServer delay onRequest action = + bracket open N.close $ \sock -> do + p <- fromIntegral <$> N.socketPort sock + tid <- forkIO $ forever $ do + (conn, _) <- N.accept sock + void $ forkIO $ flip finally (N.close conn) $ do + chunk <- NB.recv conn 4096 + if "POST" `BS.isPrefixOf` chunk + then do + onRequest + threadDelay delay + NB.sendAll conn "HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n" + else pure () + action p `finally` killThread tid + where + open = do + sock <- N.socket N.AF_INET N.Stream N.defaultProtocol + N.bind sock (N.SockAddrInet 0 (N.tupleToHostAddress (127, 0, 0, 1))) + N.listen sock 16 + pure sock diff --git a/test/LLM/LLMProviderSpec.hs b/test/LLM/LLMProviderSpec.hs new file mode 100644 --- /dev/null +++ b/test/LLM/LLMProviderSpec.hs @@ -0,0 +1,47 @@ +module LLM.LLMProviderSpec (spec) where + +import Data.ByteString qualified as BS +import LLM.Core.LLMProvider (LLMProvider (..), genericGenerateText) +import LLM.Core.Types (ChatRequest (..), LLMError (..), pattern UserTurn) +import LLM.Generate.GenerateUtils (llmHooks) +import LLM.Generate.Logger (noHooks) +import LLM.Providers.OpenAI (openAIProvider) +import Test.Hspec (Spec, describe, expectationFailure, it, shouldBe) + +spec :: Spec +spec = describe "LLMProvider" $ do + describe "non-streaming error responses" $ do + it "keeps a JSON error body as the raw JSON text" $ do + let raw = "{\"error\":{\"message\":\"bad key\"}}" + result <- generateWith 401 raw + result `shouldBe` Just (HttpError 401 "{\"error\":{\"message\":\"bad key\"}}") + + it "keeps the status of a non-JSON error body" $ do + result <- generateWith 502 "Bad Gateway" + result `shouldBe` Just (HttpError 502 "Bad Gateway") + + it "reports a non-JSON 200 body as a ParseError" $ do + result <- generateWith 200 "oops" + case result of + Just (ParseError _) -> pure () + other -> expectationFailure ("expected ParseError, got " <> show other) + +-- | Run a non-streaming chat against a provider that answers with the given +-- status and raw body; return the error, if any. +generateWith :: Int -> BS.ByteString -> IO (Maybe LLMError) +generateWith status raw = do + let p = (openAIProvider "key") {sendRequest = \_ -> pure (status, raw)} + result <- genericGenerateText p (llmHooks noHooks) request + pure (either Just (const Nothing) result) + +request :: ChatRequest +request = + ChatRequest + { reqModel = "gpt-test", + reqConversation = [UserTurn "hi"], + reqSystem = Nothing, + reqTools = [], + reqMaxTokens = 10, + reqTemperature = Nothing, + reqThinking = Nothing + } diff --git a/test/LLM/LoadSpec.hs b/test/LLM/LoadSpec.hs --- a/test/LLM/LoadSpec.hs +++ b/test/LLM/LoadSpec.hs @@ -3,6 +3,7 @@ import Control.Exception (IOException, SomeException, catch, displayException, fromException, try) import Control.Monad.Except (runExceptT) import Data.Aeson (Value (Array), eitherDecodeFileStrict') +import Data.ByteString qualified as BS import Data.List (find) import Data.Map qualified as Map import Data.Maybe (isJust) @@ -21,6 +22,7 @@ loadProviderCatalogForModelCatalog, ) import LLM.Load.Types (LoadConfigError (..)) +import LLM.Load.Utils (getSystemPrompt) import System.FilePath (()) import System.IO.Temp (withSystemTempDirectory) import Test.Hspec @@ -65,14 +67,14 @@ 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 + fmap (.capabilities) (find (\i -> i.modelConfigName == "haiku_4_5") items) `shouldBe` Just - ( ModelCapabilities - { capThinking = True, - capVision = True, - capPromptCaching = True - } + ( Just + ModelCapabilities + { capThinking = True, + capVision = True, + capPromptCaching = True + } ) Left err -> expectationFailure err @@ -92,6 +94,38 @@ Map.member "llama_3_2" catalog `shouldBe` True Map.member "mistral" catalog `shouldBe` True + describe "providers that are unavailable" $ do + it "loads a model when another model's provider has no API key" $ + withCatalogs mixedCatalog mixedProviders $ \path -> do + cfg <- loadModelOrThrow path "llama_3_2" + cfg.mcModel `shouldBe` "llama3.2:latest" + + it "names the missing env var when the requested model's provider has no API key" $ + withCatalogs mixedCatalog mixedProviders $ \path -> do + result <- try @LoadConfigError (loadModelOrThrow path "haiku") + case result of + Left (LoadModelConfigError msg) -> do + T.unpack msg `shouldContain` "haiku" + T.unpack msg `shouldContain` "claude" + T.unpack msg `shouldContain` "LLM_SIMPLE_TEST_UNSET_KEY" + Left other -> expectationFailure $ "expected LoadModelConfigError, got: " <> show other + Right _ -> expectationFailure "expected failure" + + it "ignores an invalid baseUrl of a provider no requested model uses" $ + withCatalogs mixedCatalog mixedProviders $ \path -> do + (_, catalog, _) <- loadModelOrThrow_ path "llama_3_2" + Map.keys catalog `shouldBe` ["llama_3_2"] + + it "reports the invalid baseUrl when the requested model's provider has one" $ + withCatalogs mixedCatalog mixedProviders $ \path -> do + result <- try @LoadConfigError (loadModelOrThrow path "broken_model") + case result of + Left (LoadModelConfigError msg) -> do + T.unpack msg `shouldContain` "broken" + T.unpack msg `shouldContain` "baseUrl" + Left other -> expectationFailure $ "expected LoadModelConfigError, got: " <> show other + Right _ -> expectationFailure "expected failure" + describe "loadGateways" $ do it "always includes ollama" $ do gateways <- loadGateways @@ -193,12 +227,54 @@ Left other -> expectationFailure $ "expected LoadModelCatalogError, got: " <> show other Right _ -> expectationFailure "expected failure" + describe "getSystemPrompt" $ do + it "reports a file that is not valid UTF-8" $ + withSystemTempDirectory "llm-simple-prompt" $ \root -> do + let path = root "prompt.txt" + BS.writeFile path (BS.pack [0x63, 0x61, 0x66, 0xE9]) + result <- runExceptT $ getSystemPrompt (Just (T.pack ("file:" <> path))) + case result of + Left (LoadSystemPromptError msg) -> do + T.unpack msg `shouldContain` path + T.unpack msg `shouldContain` "UTF-8" + Left other -> expectationFailure $ "expected LoadSystemPromptError, got: " <> show other + Right p -> expectationFailure $ "expected failure, got: " <> show p + withTempCatalog :: String -> (FilePath -> IO a) -> IO a withTempCatalog contents act = withSystemTempDirectory "llm-simple-load" $ \root -> do let path = root "catalog.json" writeFile path contents act path `catch` \(_ :: IOException) -> act path + +-- | Write a model catalog and a @providers.json@ next to it. +withCatalogs :: String -> String -> (FilePath -> IO a) -> IO a +withCatalogs models providers act = + withSystemTempDirectory "llm-simple-catalogs" $ \root -> do + let catalogPath = root "model-catalog.json" + writeFile catalogPath models + writeFile (root "providers.json") providers + act catalogPath + +-- | One model each for Ollama (no key), Claude (key env var unset) and a +-- provider with an invalid base URL. +mixedCatalog :: String +mixedCatalog = + "[" <> catalogEntry "llama_3_2" "ollama" "llama3.2:latest" + <> "," <> catalogEntry "haiku" "claude" "claude-haiku-4-5" + <> "," <> catalogEntry "broken_model" "broken" "some-model" + <> "]" + where + catalogEntry name provider model = + "{\"modelConfigName\":\"" <> name <> "\",\"providerName\":\"" <> provider + <> "\",\"modelName\":\"" <> model + <> "\",\"pricing\":{\"pricePerMillionInput\":0,\"pricePerMillionOutput\":0}" + <> ",\"maxTokens\":1024,\"retryCount\":0,\"jitterBackoff\":0}" + +mixedProviders :: String +mixedProviders = + "[{\"providerName\":\"claude\",\"protocol\":\"claude\",\"baseUrl\":\"https://api.anthropic.com\",\"apiKeyEnv\":\"LLM_SIMPLE_TEST_UNSET_KEY\"}" + <> ",{\"providerName\":\"broken\",\"protocol\":\"openai\",\"baseUrl\":\"not a url\"}]" withProviderCatalogDir :: FilePath -> (FilePath -> IO a) -> IO a withProviderCatalogDir providersFixture act = diff --git a/test/LLM/OllamaSpec.hs b/test/LLM/OllamaSpec.hs new file mode 100644 --- /dev/null +++ b/test/LLM/OllamaSpec.hs @@ -0,0 +1,44 @@ +module LLM.OllamaSpec (spec) where + +import Data.Aeson (Value (Object), object, toJSON, (.=)) +import Data.Aeson.KeyMap qualified as KM +import Data.Text (Text) +import LLM.Core.LLMProvider (LLMProvider (..)) +import LLM.Core.Types (ChatRequest (..), pattern UserTurn) +import LLM.Providers.Ollama (ollamaProvider) +import Test.Hspec (Spec, describe, it, shouldBe) + +spec :: Spec +spec = describe "Ollama" $ do + describe "request body" $ do + -- Ollama's OpenAI-compatible endpoint honours max_tokens and ignores + -- num_predict and max_completion_tokens. + it "sends max_tokens for chat requests" $ do + let body = ollamaProvider.buildBody False req + field "max_tokens" body `shouldBe` Just (toJSON (77 :: Int)) + field "num_predict" body `shouldBe` Nothing + field "max_completion_tokens" body `shouldBe` Nothing + + it "sends max_tokens for object requests" $ do + let body = ollamaProvider.buildObjectBody req (object ["type" .= ("object" :: Text)]) + field "max_tokens" body `shouldBe` Just (toJSON (77 :: Int)) + field "num_predict" body `shouldBe` Nothing + field "max_completion_tokens" body `shouldBe` Nothing + + it "requests usage in streams" $ do + let body = ollamaProvider.buildBody True req + field "stream_options" body `shouldBe` Just (object ["include_usage" .= True]) + where + req = + ChatRequest + { reqModel = "llama3.2", + reqConversation = [UserTurn "hi"], + reqSystem = Nothing, + reqMaxTokens = 77, + reqTemperature = Nothing, + reqTools = [], + reqThinking = Nothing + } + field k = \case + Object o -> KM.lookup k o + _ -> Nothing diff --git a/test/LLM/StreamingSpec.hs b/test/LLM/StreamingSpec.hs --- a/test/LLM/StreamingSpec.hs +++ b/test/LLM/StreamingSpec.hs @@ -4,13 +4,16 @@ import Control.Concurrent (threadDelay) import Control.Monad (when) +import Control.Retry (RetryPolicyM (getRetryPolicyM), RetryStatus (rsIterNumber), defaultRetryStatus) import Data.Aeson (object, (.=)) -import Data.IORef (atomicModifyIORef', modifyIORef', newIORef, readIORef) +import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef, readIORef) import Data.Text (Text) +import LLM.Core.Abort (abort, newAbortSignal) import LLM.Core.Types ( ChatResponse (..), LLMError (..), LLMGateway (..), + LLMTextResult, supportedImageMediaTypes, StreamEvent (..), Turn (..), @@ -21,16 +24,17 @@ toolCallPart, ) import LLM.Core.Usage (PricingInfo (..), defaultPricingInfo, mkUsage) -import LLM.Generate.Generate (generateTextLLM, streamTextLLM, streamTextWithFallbacks) +import LLM.Generate.Generate (generateTextLLM, generateTextWithFallbacks, streamTextLLM, streamTextWithFallbacks) import LLM.Generate.GenerateUtils (llmHooks) import LLM.Generate.Logger (noHooks) -import LLM.Generate.ModelConfig (ModelConfig (..), ModelWithFallbacks (..), defaultModelCapabilities) +import LLM.Generate.ModelConfig (ModelConfig (..), ModelWithFallbacks (..), defaultModelCapabilities, modelRetryPolicy) import LLM.Generate.Types ( GenRequest (..), + GenerateError (..), RoundTextRole (..), StreamChunk (..), ) -import Test.Hspec (Spec, describe, it, shouldBe) +import Test.Hspec (Spec, describe, it, shouldBe, shouldReturn) spec :: Spec spec = describe "Streaming" $ do @@ -124,6 +128,64 @@ r <- streamTextLLM (\_ -> pure ()) (testGenRequest []) mc fmap (.respText) r `shouldBe` Left TimeoutError + describe "abort signal" $ do + it "does not call the model when the signal fired before the call" $ do + (gw, calls) <- countingGateway (Right (ChatResponse "hi" [textPart "hi"] Nothing Nothing)) + sig <- newAbortSignal + abort sig + r <- generateTextWithFallbacks (testGenRequest []) {grAbortSignal = Just sig} (ModelWithFallbacks (testModelConfig gw) []) + fmap (.respText) r `shouldBe` Left GErrAborted + readIORef calls `shouldReturn` 0 + + it "stops retrying once the signal fired" $ do + sig <- newAbortSignal + (gw, calls) <- countingGateway (Left (HttpError 503 "overloaded")) + let gw' = gw {gwGenerateText = \h req -> abort sig >> gw.gwGenerateText h req} + mc = (testModelConfig gw') {mcRetryCount = 3} + r <- generateTextWithFallbacks (testGenRequest []) {grAbortSignal = Just sig} (ModelWithFallbacks mc []) + fmap (.respText) r `shouldBe` Left GErrAborted + readIORef calls `shouldReturn` 1 + + it "does not fall back to the next model once the signal fired" $ do + sig <- newAbortSignal + (failing, _) <- countingGateway (Left (HttpError 400 "bad")) + (ok, okCalls) <- countingGateway (Right (ChatResponse "hi" [textPart "hi"] Nothing Nothing)) + let failing' = failing {gwGenerateText = \h req -> abort sig >> failing.gwGenerateText h req} + models = ModelWithFallbacks (testModelConfig failing') [testModelConfig ok] + r <- generateTextWithFallbacks (testGenRequest []) {grAbortSignal = Just sig} models + fmap (.respText) r `shouldBe` Left GErrAborted + readIORef okCalls `shouldReturn` 0 + + it "cuts off a stream once the signal fired" $ do + sig <- newAbortSignal + let resp = ChatResponse "abc" [textPart "abc"] Nothing Nothing + gw = + (streamGateway [] resp) + { gwStreamText = \_ _ onEvent -> do + onEvent (StreamDelta "a") + abort sig + onEvent (StreamDelta "b") + onEvent (StreamDelta "c") + pure (Right resp) + } + ref <- newIORef [] + r <- streamTextWithFallbacks (\c -> modifyIORef' ref (c :)) (testGenRequest []) {grAbortSignal = Just sig} (ModelWithFallbacks (testModelConfig gw) []) + fmap (.respText) r `shouldBe` Left GErrAborted + reverse <$> readIORef ref `shouldReturn` [TextDelta "a"] + + describe "modelRetryPolicy" $ do + it "caps the backoff delay at 30 s" $ do + let mc = + (testModelConfig (streamGateway [] (ChatResponse "" [] Nothing Nothing))) + { mcRetryCount = 20, + mcJitterBackoff = 1000 + } + delays <- + mapM + (\n -> getRetryPolicyM (modelRetryPolicy mc) defaultRetryStatus {rsIterNumber = n}) + [10 .. 19] + all (maybe False (<= 30_000_000)) delays `shouldBe` True + describe "stream reset between attempts" $ do it "emits StreamAttemptReset before a retry when the failed attempt streamed" $ do gw <- flakyGateway "abc" "xyz" @@ -185,6 +247,17 @@ onEvent (StreamDelta second) pure (Right resp) } + +-- | A gateway whose non-streaming call always returns @result@, with a +-- counter of how often it was called. +countingGateway :: LLMTextResult -> IO (LLMGateway, IORef Int) +countingGateway result = do + calls <- newIORef 0 + let gw = + (streamGateway [] (ChatResponse "" [] Nothing Nothing)) + { gwGenerateText = \_ _ -> modifyIORef' calls (+ 1) >> pure result + } + pure (gw, calls) streamGateway :: [StreamEvent] -> ChatResponse -> LLMGateway streamGateway events resp = diff --git a/test/LLM/TestKit.hs b/test/LLM/TestKit.hs --- a/test/LLM/TestKit.hs +++ b/test/LLM/TestKit.hs @@ -16,7 +16,7 @@ ) where -import Data.Aeson (FromJSON, ToJSON, Value (Array, Object, String), decodeFileStrict) +import Data.Aeson (FromJSON, ToJSON, Value (Array, Object, String), decodeFileStrict, encode) import Data.Aeson.KeyMap qualified as KM import Data.Aeson.Encode.Pretty (encodePretty) import Data.ByteString qualified as BS @@ -109,7 +109,7 @@ case M.lookup val mp of Nothing -> error (show val <> "\n" <> show (fst $ head $ M.toList mp)) - Just rr -> pure (200, rr.response), + Just rr -> pure (200, BSL.toStrict (encode rr.response)), openStream = \body k -> do case M.lookup body mp of Nothing -> error ("value not found for:" <> show body) diff --git a/test/LLM/TypesSpec.hs b/test/LLM/TypesSpec.hs --- a/test/LLM/TypesSpec.hs +++ b/test/LLM/TypesSpec.hs @@ -216,6 +216,8 @@ isRetryable (HttpError 429 "rate limited") `shouldBe` True it "retries on 503" $ do isRetryable (HttpError 503 "overloaded") `shouldBe` True + it "retries on transient server errors (500, 502, 504)" $ do + map (\s -> isRetryable (HttpError s "")) [500, 502, 504] `shouldBe` [True, True, True] it "retries on network errors" $ do isRetryable (NetworkError "connection refused") `shouldBe` True it "retries on timeouts" $ do