path: root/Omni/Agent/Provider.hs

{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NoImplicitPrelude #-}

-- | LLM Provider abstraction for multi-backend support.
--
-- Supports multiple LLM backends:
-- - OpenRouter (cloud API, multiple models)
-- - Ollama (local models)
-- - Amp CLI (subprocess)
--
-- : out omni-agent-provider
-- : dep aeson
-- : dep http-conduit
-- : dep http-client-tls
-- : dep http-types
-- : dep case-insensitive
module Omni.Agent.Provider
  ( Provider (..),
    ProviderConfig (..),
    ChatResult (..),
    Message (..),
    Role (..),
    ToolCall (..),
    FunctionCall (..),
    Usage (..),
    ToolApi (..),
    StreamChunk (..),
    defaultOpenRouter,
    defaultOllama,
    chat,
    chatWithUsage,
    chatStream,
    main,
    test,
  )
where

import Alpha
import Data.Aeson ((.=))
import qualified Data.Aeson as Aeson
import qualified Data.Aeson.KeyMap as KeyMap
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import qualified Data.CaseInsensitive as CI
import Data.IORef (modifyIORef, newIORef, readIORef, writeIORef)
import qualified Data.IntMap.Strict as IntMap
import qualified Data.Text as Text
import qualified Data.Text.Encoding as TE
import qualified Network.HTTP.Client as HTTPClient
import qualified Network.HTTP.Client.TLS as HTTPClientTLS
import qualified Network.HTTP.Simple as HTTP
import Network.HTTP.Types.Status (statusCode)
import qualified Omni.Test as Test
import qualified System.Timeout as Timeout

main :: IO ()
main = Test.run test

test :: Test.Tree
test =
  Test.group
    "Omni.Agent.Provider"
    [ Test.unit "defaultOpenRouter has correct endpoint" <| do
        case defaultOpenRouter "" "test-model" of
          OpenRouter cfg -> providerBaseUrl cfg Test.@=? "https://openrouter.ai/api/v1"
          _ -> Test.assertFailure "Expected OpenRouter",
      Test.unit "defaultOllama has correct endpoint" <| do
        case defaultOllama "test-model" of
          Ollama cfg -> providerBaseUrl cfg Test.@=? "http://localhost:11434"
          _ -> Test.assertFailure "Expected Ollama",
      Test.unit "ChatResult preserves message" <| do
        let msg = Message User "test" Nothing Nothing
            result = ChatResult msg Nothing
        chatMessage result Test.@=? msg
    ]

-- | HTTP request timeout in microseconds (60 seconds)
httpTimeoutMicros :: Int
httpTimeoutMicros = 60 * 1000000

-- | Maximum number of retries for transient failures
maxRetries :: Int
maxRetries = 3

-- | Initial backoff delay in microseconds (1 second)
initialBackoffMicros :: Int
initialBackoffMicros = 1000000

-- | Retry an IO action with exponential backoff
-- Retries on timeout, connection errors, and 5xx status codes
retryWithBackoff :: Int -> Int -> IO (Either Text a) -> IO (Either Text a)
retryWithBackoff retriesLeft backoff action
  | retriesLeft <= 0 = action
  | otherwise = do
      result <- Timeout.timeout httpTimeoutMicros action
      case result of
        Nothing -> do
          threadDelay backoff
          retryWithBackoff (retriesLeft - 1) (backoff * 2) action
        Just (Left err)
          | isRetryable err -> do
              threadDelay backoff
              retryWithBackoff (retriesLeft - 1) (backoff * 2) action
        Just r -> pure r
  where
    isRetryable err =
      "HTTP error: 5"
        `Text.isPrefixOf` err
        || "connection"
        `Text.isInfixOf` Text.toLower err
        || "timeout"
        `Text.isInfixOf` Text.toLower err

data Provider
  = OpenRouter ProviderConfig
  | Ollama ProviderConfig
  | AmpCLI FilePath
  deriving (Show, Eq, Generic)

data ProviderConfig = ProviderConfig
  { providerBaseUrl :: Text,
    providerApiKey :: Text,
    providerModel :: Text,
    providerExtraHeaders :: [(ByteString, ByteString)]
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON ProviderConfig where
  toJSON c =
    Aeson.object
      [ "baseUrl" .= providerBaseUrl c,
        "apiKey" .= providerApiKey c,
        "model" .= providerModel c
      ]

instance Aeson.FromJSON ProviderConfig where
  parseJSON =
    Aeson.withObject "ProviderConfig" <| \v ->
      (ProviderConfig </ (v Aeson..: "baseUrl"))
        <*> (v Aeson..: "apiKey")
        <*> (v Aeson..: "model")
        <*> pure []

defaultOpenRouter :: Text -> Text -> Provider
defaultOpenRouter apiKey model =
  OpenRouter
    ProviderConfig
      { providerBaseUrl = "https://openrouter.ai/api/v1",
        providerApiKey = apiKey,
        providerModel = model,
        providerExtraHeaders =
          [ ("HTTP-Referer", "https://omni.dev"),
            ("X-Title", "Omni Agent")
          ]
      }

defaultOllama :: Text -> Provider
defaultOllama model =
  Ollama
    ProviderConfig
      { providerBaseUrl = "http://localhost:11434",
        providerApiKey = "",
        providerModel = model,
        providerExtraHeaders = []
      }

data Role = System | User | Assistant | ToolRole
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON Role where
  toJSON System = Aeson.String "system"
  toJSON User = Aeson.String "user"
  toJSON Assistant = Aeson.String "assistant"
  toJSON ToolRole = Aeson.String "tool"

instance Aeson.FromJSON Role where
  parseJSON = Aeson.withText "Role" parseRole
    where
      parseRole "system" = pure System
      parseRole "user" = pure User
      parseRole "assistant" = pure Assistant
      parseRole "tool" = pure ToolRole
      parseRole _ = empty

data Message = Message
  { msgRole :: Role,
    msgContent :: Text,
    msgToolCalls :: Maybe [ToolCall],
    msgToolCallId :: Maybe Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON Message where
  toJSON m =
    Aeson.object
      <| catMaybes
        [ Just ("role" .= msgRole m),
          Just ("content" .= msgContent m),
          ("tool_calls" .=) </ msgToolCalls m,
          ("tool_call_id" .=) </ msgToolCallId m
        ]

instance Aeson.FromJSON Message where
  parseJSON =
    Aeson.withObject "Message" <| \v ->
      (Message </ (v Aeson..: "role"))
        <*> (v Aeson..:? "content" Aeson..!= "")
        <*> (v Aeson..:? "tool_calls")
        <*> (v Aeson..:? "tool_call_id")

data ToolCall = ToolCall
  { tcId :: Text,
    tcType :: Text,
    tcFunction :: FunctionCall
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON ToolCall where
  toJSON tc =
    Aeson.object
      [ "id" .= tcId tc,
        "type" .= tcType tc,
        "function" .= tcFunction tc
      ]

instance Aeson.FromJSON ToolCall where
  parseJSON =
    Aeson.withObject "ToolCall" <| \v ->
      (ToolCall </ (v Aeson..: "id"))
        <*> (v Aeson..:? "type" Aeson..!= "function")
        <*> (v Aeson..: "function")

data FunctionCall = FunctionCall
  { fcName :: Text,
    fcArguments :: Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.ToJSON FunctionCall where
  toJSON fc =
    Aeson.object
      [ "name" .= fcName fc,
        "arguments" .= fcArguments fc
      ]

instance Aeson.FromJSON FunctionCall where
  parseJSON =
    Aeson.withObject "FunctionCall" <| \v ->
      (FunctionCall </ (v Aeson..: "name"))
        <*> (v Aeson..:? "arguments" Aeson..!= "{}")

data Usage = Usage
  { usagePromptTokens :: Int,
    usageCompletionTokens :: Int,
    usageTotalTokens :: Int,
    usageCost :: Maybe Double
  }
  deriving (Show, Eq, Generic)

instance Aeson.FromJSON Usage where
  parseJSON =
    Aeson.withObject "Usage" <| \v ->
      (Usage </ (v Aeson..: "prompt_tokens"))
        <*> (v Aeson..: "completion_tokens")
        <*> (v Aeson..: "total_tokens")
        <*> (v Aeson..:? "cost")

data ChatResult = ChatResult
  { chatMessage :: Message,
    chatUsage :: Maybe Usage
  }
  deriving (Show, Eq)

data ToolApi = ToolApi
  { toolApiName :: Text,
    toolApiDescription :: Text,
    toolApiParameters :: Aeson.Value
  }
  deriving (Generic)

instance Aeson.ToJSON ToolApi where
  toJSON t =
    Aeson.object
      [ "type" .= ("function" :: Text),
        "function"
          .= Aeson.object
            [ "name" .= toolApiName t,
              "description" .= toolApiDescription t,
              "parameters" .= toolApiParameters t
            ]
      ]

data ChatCompletionRequest = ChatCompletionRequest
  { reqModel :: Text,
    reqMessages :: [Message],
    reqTools :: Maybe [ToolApi]
  }
  deriving (Generic)

instance Aeson.ToJSON ChatCompletionRequest where
  toJSON r =
    Aeson.object
      <| catMaybes
        [ Just ("model" .= reqModel r),
          Just ("messages" .= reqMessages r),
          ("tools" .=) </ reqTools r,
          Just ("usage" .= Aeson.object ["include" .= True])
        ]

data Choice = Choice
  { choiceIndex :: Int,
    choiceMessage :: Message,
    choiceFinishReason :: Maybe Text
  }
  deriving (Show, Eq, Generic)

instance Aeson.FromJSON Choice where
  parseJSON =
    Aeson.withObject "Choice" <| \v ->
      (Choice </ (v Aeson..: "index"))
        <*> (v Aeson..: "message")
        <*> (v Aeson..:? "finish_reason")

data ChatCompletionResponse = ChatCompletionResponse
  { respId :: Text,
    respChoices :: [Choice],
    respModel :: Text,
    respUsage :: Maybe Usage
  }
  deriving (Show, Eq, Generic)

instance Aeson.FromJSON ChatCompletionResponse where
  parseJSON =
    Aeson.withObject "ChatCompletionResponse" <| \v ->
      (ChatCompletionResponse </ (v Aeson..: "id"))
        <*> (v Aeson..: "choices")
        <*> (v Aeson..: "model")
        <*> (v Aeson..:? "usage")

chat :: Provider -> [ToolApi] -> [Message] -> IO (Either Text Message)
chat provider tools messages = do
  result <- chatWithUsage provider tools messages
  pure (chatMessage </ result)

chatWithUsage :: Provider -> [ToolApi] -> [Message] -> IO (Either Text ChatResult)
chatWithUsage (OpenRouter cfg) tools messages = chatOpenAI cfg tools messages
chatWithUsage (Ollama cfg) tools messages = chatOllama cfg tools messages
chatWithUsage (AmpCLI _promptFile) _tools _messages = do
  pure (Left "Amp CLI provider not yet implemented")

chatOpenAI :: ProviderConfig -> [ToolApi] -> [Message] -> IO (Either Text ChatResult)
chatOpenAI cfg tools messages = do
  let url = Text.unpack (providerBaseUrl cfg) <> "/chat/completions"
  req0 <- HTTP.parseRequest url
  let body =
        ChatCompletionRequest
          { reqModel = providerModel cfg,
            reqMessages = messages,
            reqTools = if null tools then Nothing else Just tools
          }
      baseReq =
        HTTP.setRequestMethod "POST"
          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
          <| HTTP.setRequestHeader "Authorization" ["Bearer " <> TE.encodeUtf8 (providerApiKey cfg)]
          <| HTTP.setRequestBodyLBS (Aeson.encode body)
          <| req0
      req = foldr addHeader baseReq (providerExtraHeaders cfg)
      addHeader (name, value) = HTTP.addRequestHeader (CI.mk name) value

  retryWithBackoff maxRetries initialBackoffMicros <| do
    response <- HTTP.httpLBS req
    let status = HTTP.getResponseStatusCode response
        respBody = HTTP.getResponseBody response
        cleanedBody = BL.dropWhile (\b -> b `elem` [0x0a, 0x0d, 0x20]) respBody
    if status >= 200 && status < 300
      then case Aeson.eitherDecode cleanedBody of
        Right resp ->
          case respChoices resp of
            (c : _) -> pure (Right (ChatResult (choiceMessage c) (respUsage resp)))
            [] -> pure (Left "No choices in response")
        Left err -> do
          let bodyPreview = TE.decodeUtf8 (BL.toStrict (BL.take 500 cleanedBody))
          pure (Left ("Failed to parse response: " <> Text.pack err <> " | Body: " <> bodyPreview))
      else pure (Left ("HTTP error: " <> tshow status <> " - " <> TE.decodeUtf8 (BL.toStrict respBody)))

chatOllama :: ProviderConfig -> [ToolApi] -> [Message] -> IO (Either Text ChatResult)
chatOllama cfg tools messages = do
  let url = Text.unpack (providerBaseUrl cfg) <> "/api/chat"
  req0 <- HTTP.parseRequest url
  let body =
        Aeson.object
          [ "model" .= providerModel cfg,
            "messages" .= messages,
            "tools" .= if null tools then Aeson.Null else Aeson.toJSON tools,
            "stream" .= False
          ]
      req =
        HTTP.setRequestMethod "POST"
          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
          <| HTTP.setRequestBodyLBS (Aeson.encode body)
          <| req0

  retryWithBackoff maxRetries initialBackoffMicros <| do
    response <- HTTP.httpLBS req
    let status = HTTP.getResponseStatusCode response
    if status >= 200 && status < 300
      then case Aeson.decode (HTTP.getResponseBody response) of
        Just resp -> parseOllamaResponse resp
        Nothing -> pure (Left ("Failed to parse Ollama response: " <> TE.decodeUtf8 (BL.toStrict (HTTP.getResponseBody response))))
      else pure (Left ("HTTP error: " <> tshow status <> " - " <> TE.decodeUtf8 (BL.toStrict (HTTP.getResponseBody response))))

parseOllamaResponse :: Aeson.Value -> IO (Either Text ChatResult)
parseOllamaResponse val =
  case val of
    Aeson.Object obj -> do
      let msgResult = do
            msgObj <- case KeyMap.lookup "message" obj of
              Just m -> Right m
              Nothing -> Left "No message in response"
            case Aeson.fromJSON msgObj of
              Aeson.Success msg -> Right msg
              Aeson.Error e -> Left (Text.pack e)
          usageResult = case KeyMap.lookup "prompt_eval_count" obj of
            Just (Aeson.Number promptTokens) ->
              case KeyMap.lookup "eval_count" obj of
                Just (Aeson.Number evalTokens) ->
                  Just
                    Usage
                      { usagePromptTokens = round promptTokens,
                        usageCompletionTokens = round evalTokens,
                        usageTotalTokens = round promptTokens + round evalTokens,
                        usageCost = Nothing
                      }
                _ -> Nothing
            _ -> Nothing
      case msgResult of
        Right msg -> pure (Right (ChatResult msg usageResult))
        Left e -> pure (Left e)
    _ -> pure (Left "Expected object response from Ollama")

data StreamChunk
  = StreamContent Text
  | StreamToolCall ToolCall
  | StreamToolCallDelta ToolCallDelta
  | StreamDone ChatResult
  | StreamError Text
  deriving (Show, Eq)

data ToolCallDelta = ToolCallDelta
  { tcdIndex :: Int,
    tcdId :: Maybe Text,
    tcdFunctionName :: Maybe Text,
    tcdFunctionArgs :: Maybe Text
  }
  deriving (Show, Eq)

chatStream :: Provider -> [ToolApi] -> [Message] -> (StreamChunk -> IO ()) -> IO (Either Text ChatResult)
chatStream (OpenRouter cfg) tools messages onChunk = chatStreamOpenAI cfg tools messages onChunk
chatStream (Ollama _cfg) _tools _messages _onChunk = pure (Left "Streaming not implemented for Ollama")
chatStream (AmpCLI _) _tools _messages _onChunk = pure (Left "Streaming not implemented for AmpCLI")

chatStreamOpenAI :: ProviderConfig -> [ToolApi] -> [Message] -> (StreamChunk -> IO ()) -> IO (Either Text ChatResult)
chatStreamOpenAI cfg tools messages onChunk = do
  let url = Text.unpack (providerBaseUrl cfg) <> "/chat/completions"
      managerSettings =
        HTTPClientTLS.tlsManagerSettings
          { HTTPClient.managerResponseTimeout = HTTPClient.responseTimeoutMicro httpTimeoutMicros
          }
  manager <- HTTPClient.newManager managerSettings
  req0 <- HTTP.parseRequest url
  let body =
        Aeson.object
          <| catMaybes
            [ Just ("model" .= providerModel cfg),
              Just ("messages" .= messages),
              if null tools then Nothing else Just ("tools" .= tools),
              Just ("stream" .= True),
              Just ("usage" .= Aeson.object ["include" .= True])
            ]
      baseReq =
        HTTP.setRequestMethod "POST"
          <| HTTP.setRequestHeader "Content-Type" ["application/json"]
          <| HTTP.setRequestHeader "Authorization" ["Bearer " <> TE.encodeUtf8 (providerApiKey cfg)]
          <| HTTP.setRequestBodyLBS (Aeson.encode body)
          <| req0
      req = foldr addHeader baseReq (providerExtraHeaders cfg)
      addHeader (name, value) = HTTP.addRequestHeader (CI.mk name) value

  result <-
    try <| HTTPClient.withResponse req manager <| \response -> do
      let status = HTTPClient.responseStatus response
          code = statusCode status
      if code >= 200 && code < 300
        then processSSEStream (HTTPClient.responseBody response) onChunk
        else do
          bodyChunks <- readAllBody (HTTPClient.responseBody response)
          let errBody = TE.decodeUtf8 (BS.concat bodyChunks)
          pure (Left ("HTTP error: " <> tshow code <> " - " <> errBody))
  case result of
    Left (e :: SomeException) -> pure (Left ("Stream request failed: " <> tshow e))
    Right r -> pure r

readAllBody :: IO BS.ByteString -> IO [BS.ByteString]
readAllBody readBody = go []
  where
    go acc = do
      chunk <- readBody
      if BS.null chunk
        then pure (reverse acc)
        else go (chunk : acc)

data ToolCallAccum = ToolCallAccum
  { tcaId :: Text,
    tcaName :: Text,
    tcaArgs :: Text
  }

processSSEStream :: IO BS.ByteString -> (StreamChunk -> IO ()) -> IO (Either Text ChatResult)
processSSEStream readBody onChunk = do
  accumulatedContent <- newIORef ("" :: Text)
  toolCallAccum <- newIORef (IntMap.empty :: IntMap.IntMap ToolCallAccum)
  lastUsage <- newIORef (Nothing :: Maybe Usage)
  buffer <- newIORef ("" :: Text)

  let loop = do
        chunk <- readBody
        if BS.null chunk
          then do
            content <- readIORef accumulatedContent
            accum <- readIORef toolCallAccum
            usage <- readIORef lastUsage
            let toolCalls = map accumToToolCall (IntMap.elems accum)
                finalMsg =
                  Message
                    { msgRole = Assistant,
                      msgContent = content,
                      msgToolCalls = if null toolCalls then Nothing else Just toolCalls,
                      msgToolCallId = Nothing
                    }
            pure (Right (ChatResult finalMsg usage))
          else do
            modifyIORef buffer (<> TE.decodeUtf8 chunk)
            buf <- readIORef buffer
            let (events, remaining) = parseSSEEvents buf
            writeIORef buffer remaining
            forM_ events <| \event -> do
              case parseStreamEvent event of
                Just (StreamContent txt) -> do
                  modifyIORef accumulatedContent (<> txt)
                  onChunk (StreamContent txt)
                Just (StreamToolCallDelta delta) -> do
                  modifyIORef toolCallAccum (mergeToolCallDelta delta)
                Just (StreamToolCall tc) -> do
                  modifyIORef toolCallAccum (mergeCompleteToolCall tc)
                  onChunk (StreamToolCall tc)
                Just (StreamDone result) -> do
                  writeIORef lastUsage (chatUsage result)
                Just (StreamError err) -> do
                  onChunk (StreamError err)
                Nothing -> pure ()
            loop

  loop

accumToToolCall :: ToolCallAccum -> ToolCall
accumToToolCall acc =
  ToolCall
    { tcId = tcaId acc,
      tcType = "function",
      tcFunction = FunctionCall (tcaName acc) (tcaArgs acc)
    }

mergeToolCallDelta :: ToolCallDelta -> IntMap.IntMap ToolCallAccum -> IntMap.IntMap ToolCallAccum
mergeToolCallDelta delta accum =
  let idx = tcdIndex delta
      existing = IntMap.lookup idx accum
      updated = case existing of
        Nothing ->
          ToolCallAccum
            { tcaId = fromMaybe "" (tcdId delta),
              tcaName = fromMaybe "" (tcdFunctionName delta),
              tcaArgs = fromMaybe "" (tcdFunctionArgs delta)
            }
        Just a ->
          a
            { tcaId = fromMaybe (tcaId a) (tcdId delta),
              tcaName = fromMaybe (tcaName a) (tcdFunctionName delta),
              tcaArgs = tcaArgs a <> fromMaybe "" (tcdFunctionArgs delta)
            }
   in IntMap.insert idx updated accum

mergeCompleteToolCall :: ToolCall -> IntMap.IntMap ToolCallAccum -> IntMap.IntMap ToolCallAccum
mergeCompleteToolCall tc accum =
  let nextIdx = if IntMap.null accum then 0 else fst (IntMap.findMax accum) + 1
      newAccum =
        ToolCallAccum
          { tcaId = tcId tc,
            tcaName = fcName (tcFunction tc),
            tcaArgs = fcArguments (tcFunction tc)
          }
   in IntMap.insert nextIdx newAccum accum

parseSSEEvents :: Text -> ([Text], Text)
parseSSEEvents input =
  let lines' = Text.splitOn "\n" input
      (events, remaining) = go [] [] lines'
   in (events, remaining)
  where
    go events current [] = (reverse events, Text.intercalate "\n" (reverse current))
    go events current (line : rest)
      | Text.null line && not (null current) =
          go (Text.intercalate "\n" (reverse current) : events) [] rest
      | otherwise =
          go events (line : current) rest

parseStreamEvent :: Text -> Maybe StreamChunk
parseStreamEvent eventText = do
  let dataLines = filter ("data:" `Text.isPrefixOf`) (Text.lines eventText)
  case dataLines of
    [] -> Nothing
    (dataLine : _) -> do
      let jsonStr = Text.strip (Text.drop 5 dataLine)
      if jsonStr == "[DONE]"
        then Nothing
        else case Aeson.decode (BL.fromStrict (TE.encodeUtf8 jsonStr)) of
          Nothing -> Nothing
          Just (Aeson.Object obj) -> parseStreamChunk obj
          _ -> Nothing

parseStreamChunk :: Aeson.Object -> Maybe StreamChunk
parseStreamChunk obj = do
  case KeyMap.lookup "error" obj of
    Just (Aeson.Object errObj) -> do
      let errMsg = case KeyMap.lookup "message" errObj of
            Just (Aeson.String m) -> m
            _ -> "Unknown error"
      Just (StreamError errMsg)
    _ -> do
      let usageChunk = case KeyMap.lookup "usage" obj of
            Just usageVal -> case Aeson.fromJSON usageVal of
              Aeson.Success usage -> Just (StreamDone (ChatResult (Message Assistant "" Nothing Nothing) (Just usage)))
              _ -> Nothing
            _ -> Nothing
      case KeyMap.lookup "choices" obj of
        Just (Aeson.Array choices) | not (null choices) -> do
          case toList choices of
            (Aeson.Object choice : _) -> do
              case KeyMap.lookup "delta" choice of
                Just (Aeson.Object delta) -> do
                  let contentChunk = case KeyMap.lookup "content" delta of
                        Just (Aeson.String c) | not (Text.null c) -> Just (StreamContent c)
                        _ -> Nothing
                      toolCallChunk = case KeyMap.lookup "tool_calls" delta of
                        Just (Aeson.Array tcs)
                          | not (null tcs) ->
                              parseToolCallDelta (toList tcs)
                        _ -> Nothing
                  contentChunk <|> toolCallChunk <|> usageChunk
                _ -> usageChunk
            _ -> usageChunk
        _ -> usageChunk

parseToolCallDelta :: [Aeson.Value] -> Maybe StreamChunk
parseToolCallDelta [] = Nothing
parseToolCallDelta (Aeson.Object tcObj : _) = do
  idx <- case KeyMap.lookup "index" tcObj of
    Just (Aeson.Number n) -> Just (round n)
    _ -> Nothing
  let tcId' = case KeyMap.lookup "id" tcObj of
        Just (Aeson.String s) -> Just s
        _ -> Nothing
      funcObj = case KeyMap.lookup "function" tcObj of
        Just (Aeson.Object f) -> Just f
        _ -> Nothing
      funcName = case funcObj of
        Just f -> case KeyMap.lookup "name" f of
          Just (Aeson.String s) -> Just s
          _ -> Nothing
        Nothing -> Nothing
      funcArgs = case funcObj of
        Just f -> case KeyMap.lookup "arguments" f of
          Just (Aeson.String s) -> Just s
          _ -> Nothing
        Nothing -> Nothing
  Just
    ( StreamToolCallDelta
        ToolCallDelta
          { tcdIndex = idx,
            tcdId = tcId',
            tcdFunctionName = funcName,
            tcdFunctionArgs = funcArgs
          }
    )
parseToolCallDelta _ = Nothing