Haskell ollama call showing how to respond to a tool usage and parsing

ollamaTool2.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE QuasiQuotes #-}

module App where

import           Verset
import qualified Data.Aeson as Ae
import           Data.Aeson ((.:))
import qualified Data.ByteString.Lazy as BSL
import qualified Data.Map.Strict as Map
import qualified Data.Text as Txt
import qualified Data.Text.Encoding as TxtE
import qualified Data.List.NonEmpty as NE
import qualified Ollama as O
import           Text.Pretty.Simple (pPrint)
import           Text.RawString.QQ (r)
import           Control.Exception.Safe (throwString)


app :: IO ()
app = do
  -- Model to use. Install with `ollama pull <model_name>`
  -- Make sure it has the tool capability (use `ollama show <model_name>` to check)
  -- And look at the Berkeley leaderboard to pick a good model for function calls. https://gorilla.cs.berkeley.edu/leaderboard.html
  let modelName = "qwen2.5:7b-instruct"
  let myCopt = mkCopt modelName

  toolsJson <-
       case Ae.eitherDecode toolsJson' :: Either [Char] [Ae.Value] of
         Right v -> pure v
         Left e -> throwString e


  -- Works a bit like a system message
  let msg0 =
        O.Message
          { role = O.User
          , content = "When using a function just give a basic answer like the result from <function_name> is <result>. Don't give any other text."
          , images = Nothing
          , tool_calls = Nothing
          }

  -- First message e.g. "from the user"
  let msg1 =
        O.Message
          { role = O.User
          , content = "Please squash 2 and 3"
          , images = Nothing
          , tool_calls = Nothing
          }

  -- Options for the first call. Send the 1st message and set the tools
  let chatOpts1 =
        (myCopt (NE.fromList [msg0, msg1]))
          { O.tools = Just toolsJson
          }

  -- Chat call
  chatRes1' <- O.chat chatOpts1

  -- Check for errors
  chatRes1 <-
    case chatRes1' of
      Left e -> throwString e
      Right r' -> pure r'

  putText ""
  putText "=================="
  pPrint chatRes1


  -- Parse any tool calls from the response
  toolCalls <- parseResponseToolCalls chatRes1

  case toolCalls of
    Nothing -> do
      pass

    Just tcs -> do
      -- Get the response tool message as an array
      let chatRes1Message = maybe [] (:[]) chatRes1.message

      -- Get the tool response as JSON
      toolAnswers1  <- traverse runTool tcs
      toolAnswers2 <- case sequenceA toolAnswers1 of
        Left e -> throwString $ Txt.unpack e
        Right r' -> pure r'
      let toolAnswers3 = TxtE.decodeUtf8 . BSL.toStrict . Ae.encode $ Map.fromList toolAnswers2


      -- Response with the tool result (fake in this example)
      let msg2 =
            O.Message
              { role = O.Tool
              , content = toolAnswers3
              , images = Nothing
              , tool_calls = Nothing
              }

      -- Call the LLM again with the message history and the tool result
      chatRes2 <- O.chat $ myCopt (NE.fromList $ [msg0, msg1] <> chatRes1Message <> [msg2])

      putText ""
      putText "=================="
      pPrint chatRes2






mkCopt :: Text -> NE.NonEmpty O.Message -> O.ChatOps
mkCopt modelName msg =
  -- Chat options, you could use O.defaultChatOptions
  O.ChatOps
    { chatModelName = modelName
    , messages = msg
    , tools = Nothing
    , format = Nothing
    , keepAlive = Nothing
    , hostUrl = Nothing
    , responseTimeOut = Nothing
    , options = Nothing
    , stream = Nothing
    }
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------



--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
-- Tools handler
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
runTool :: ToolCall -> IO (Either Text (Text, Text))
runTool (ToolCall (FunctionCall name args)) = do
  case name of
    "squash" -> do
      let a = Map.lookup "a" args
      let b = Map.lookup "b" args
      case (a, b) of
        (Just (Ae.Number a'), Just (Ae.Number b')) -> pure . Right $ (name, show (a' * b'))
        _ -> pure . Left $ "Invalid arguments for squash function. Expected two numbers."

    _ -> pure . Left $ "Unknown function: " <> name
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------




--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
-- FunctionCall and ToolCall types and helpers
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
data FunctionCall =
  FunctionCall
    { fcName :: !Text
    , fcArguments :: !(Map Text Ae.Value)
    } deriving (Show, Eq)

data ToolCall =
  ToolCall
    { tcFunction :: !FunctionCall
    } deriving (Show, Eq)


instance Ae.FromJSON FunctionCall where
  parseJSON = Ae.withObject "FunctionCall" $ \o ->
    FunctionCall
      <$> o .: "name"
      <*> o .: "arguments"

instance Ae.FromJSON ToolCall where
  parseJSON = Ae.withObject "ToolCall" $ \o ->
    ToolCall
      <$> o .: "function"



parseToolCalls :: O.Message -> Either [Char] (Maybe [ToolCall])
parseToolCalls msg =
  case msg.tool_calls of
    Nothing -> Right Nothing
    Just tc -> case traverse Ae.fromJSON tc of
      Ae.Error e -> Left e
      Ae.Success tcs -> pure . Just $ tcs


parseResponseToolCalls :: O.ChatResponse -> IO (Maybe [ToolCall])
parseResponseToolCalls resp =
  case resp.message of
    Nothing -> pure Nothing
    Just msg ->
      case parseToolCalls msg of
        Left e -> throwString e
        Right tcs ->
          case tcs of
            Nothing -> pure Nothing
            Just tcs' -> pure (Just tcs')
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------



--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
-- Example tool JSON. Using "squash" as a made up function name
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
toolsJson' :: BSL.ByteString
toolsJson' = [r|
      [
        {
          "type": "function",
          "function": {
            "name": "squash",
            "description": "Squash two numbers",
            "parameters": {
              "type": "object",
              "properties": {
                "a": {
                  "type": "number",
                  "description": "First number to squash"
                },
                "b": {
                  "type": "number",
                  "description": "Second number to squash"
                }
              },
              "required": ["a", "b"]
            }
          }
        }
      ]
    |]
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------