Haskell ollama example REPL with basic tools support

ollamaRepl.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)
import qualified System.Console.Haskeline as HL

data ReplState = ReplState
  { rsVerbose :: !Bool
  , rsHistory :: ![O.Message]
  , rsTools :: ![Ae.Value]
  , rsSysMessage :: !O.Message
  }


app :: IO ()
app = do
  toolsJson <-
       case Ae.eitherDecode toolsJson' :: Either [Char] [Ae.Value] of
         Right v -> pure v
         Left e -> throwString e


  let st0 =
       ReplState
         { rsVerbose = False
         , rsTools = toolsJson
         , rsSysMessage = O.Message O.User "When using a function just give a basic answer like the result from <function_name> is <result>. Don't give any other text." Nothing Nothing
         , rsHistory = []
         }


  putText ""
  putText ""
  putText "Welcome to the REPL. Type :quit to exit."
  putText "Type :verbose to see the full response."
  putText "Type :quiet to hide the full response."
  putText ""

  HL.runInputT
    (HL.Settings
      { HL.historyFile = Just "llm.repl.txt"
      , HL.complete = HL.noCompletion
      , HL.autoAddHistory = True
      }
    )
    (loop st0)

  where
     loop :: ReplState -> HL.InputT IO ()
     loop st1 = do
        minput <- HL.getInputLine "▷ "
        case minput of
          Nothing -> pure ()
          Just ":quit" -> pure ()
          Just ":quiet" -> loop $ st1 { rsVerbose = False }
          Just ":verbose" -> loop $ st1 { rsVerbose = True }
          Just input -> do
            st2 <- liftIO . handleInput st1 $ Txt.pack input
            loop st2



handleInput :: ReplState -> Text -> IO ReplState
handleInput st1 input = do
  handleMessage
    st1
    O.Message
      { role = O.User
      , content = input
      , images = Nothing
      , tool_calls = Nothing
      }



handleMessage :: ReplState -> O.Message -> IO ReplState
handleMessage st1 inputMsg = do
  let msgs =
       [st1.rsSysMessage] -- Keep the system message
       <> (reverse . take 10 . reverse . excludeSysMsg $ st1.rsHistory)  -- Keep the last 10 messages
       <> [inputMsg] -- The current user message

  chatRes <- chat msgs

  parseResponseToolCalls chatRes >>= \case
   -- No tool calls, just print the response
    Nothing -> do
      if st1.rsVerbose
        then do
          putText ""
          putText "=================="
          pPrint chatRes
        else
          putText $ maybe "" O.content chatRes.message

      pure $ st1 { rsHistory = msgs <> maybe [] (:[]) chatRes.message }


    -- Tool calls, run the tools and call the LLM again
    Just tcs -> do
      -- We need to send the response back to the LLM (i.e. the response that it sent us with the tool calls)
      let prevResp = maybe [] (:[]) chatRes.message

      -- Run tools
      toolAnswers1 <- traverse runTool tcs

      -- Check for errors
      toolAnswers2 <- case sequenceA toolAnswers1 of
        Left e -> throwString $ Txt.unpack e
        Right r' -> pure r'

      -- Get the tool response as JSON
      let toolAnswers3 = TxtE.decodeUtf8 . BSL.toStrict . Ae.encode $ Map.fromList toolAnswers2

      -- Response with the tool result
      let toolMsg =
            O.Message
              { role = O.Tool
              , content = toolAnswers3
              , images = Nothing
              , tool_calls = Nothing
              }

      -- The tool response message must be added to the message history
      let st2 = st1 { rsHistory = msgs <> prevResp <> [toolMsg] }
      putText "ƒ loop"
      handleMessage st2 toolMsg


  where
    -- 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
    modelName = "qwen2.5:7b-instruct"
    myCopt = mkCopt modelName (Just st1.rsTools)

    chat msgs = do
      let chatOpts =
            (myCopt (NE.fromList msgs))
              { O.tools = Just st1.rsTools
              }

      chatRes' <- O.chat chatOpts

      case chatRes' of
        Left e -> throwString e
        Right r' -> pure r'

    excludeSysMsg msgs =
      filter (\m -> m.content /= st1.rsSysMessage.content) msgs





mkCopt :: Text -> Maybe [Ae.Value] -> NE.NonEmpty O.Message -> O.ChatOps
mkCopt modelName tools msg =
  -- Chat options, you could use O.defaultChatOptions
  O.ChatOps
    { chatModelName = modelName
    , messages = msg
    , tools = tools
    , 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"]
            }
          }
        }
      ]
    |]
--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------