{-# LANGUAGE OverloadedStrings #-}

{- |
Module      : Data.Ollama.Load
Copyright   : (c) 2025 Tushar Adhatrao
License     : MIT
Maintainer  : Tushar Adhatrao <tusharadhatrao@gmail.com>
Stability   : experimental
Description : High-level functions for loading and unloading models in the Ollama client.

This module provides functions to load and unload generative models in the Ollama server.
It includes both IO-based functions ('loadGenModel', 'unloadGenModel') and monadic versions
('loadGenModelM', 'unloadGenModelM') for use in 'MonadIO' contexts. The operations are
performed via POST requests to the "/api//generate" endpoint, leveraging the 'GenerateOps'
configuration from the 'Data.Ollama.Generate' module.

Loading a model keeps it in memory for faster subsequent requests, while unloading frees
up memory by setting the keep-alive duration to zero.

Example:

>>> loadGenModel "gemma3"
Right ()
>>> unloadGenModel "gemma3"
Right ()
-}
module Data.Ollama.Load
  ( -- * Load and Unload Model APIs
    loadGenModel
  , unloadGenModel
  , loadGenModelM
  , unloadGenModelM
  ) where

import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.Ollama.Common.Error
import Data.Ollama.Common.Utils (commonNonStreamingHandler, withOllamaRequest)
import Data.Ollama.Generate qualified as Gen
import Data.Text (Text)

{- | Loads a generative model into memory.

Sends a POST request to the "/api//generate" endpoint to load the specified model into
memory, ensuring faster response times for subsequent requests. Returns 'Right ()' on
success or 'Left' with an 'OllamaError' on failure.
--
-- @since 0.2.0.0
-}
loadGenModel ::
  -- |  Model name (e.g., "gemma3")
  Text ->
  IO (Either OllamaError ())
loadGenModel :: Text -> IO (Either OllamaError ())
loadGenModel Text
m = do
  let ops :: GenerateOps
ops = GenerateOps
Gen.defaultGenerateOps {Gen.modelName = m}
  Text
-> ByteString
-> Maybe GenerateOps
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError ()))
-> IO (Either OllamaError ())
forall payload response.
ToJSON payload =>
Text
-> ByteString
-> Maybe payload
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError response))
-> IO (Either OllamaError response)
withOllamaRequest Text
"/api//generate" ByteString
"POST" (GenerateOps -> Maybe GenerateOps
forall a. a -> Maybe a
Just GenerateOps
ops) Maybe OllamaConfig
forall a. Maybe a
Nothing Response BodyReader -> IO (Either OllamaError ())
forall a.
FromJSON a =>
Response BodyReader -> IO (Either OllamaError a)
commonNonStreamingHandler

{- | Unloads a generative model from memory.

Sends a POST request to the "/api//generate" endpoint with a keep-alive duration of zero
to unload the specified model from memory, freeing up resources. Returns 'Right ()' on
success or 'Left' with an 'OllamaError' on failure.
--
-- @since 0.2.0.0
-}
unloadGenModel ::
  -- | Model name (e.g., "gemma3")
  Text ->
  IO (Either OllamaError ())
unloadGenModel :: Text -> IO (Either OllamaError ())
unloadGenModel Text
m = do
  let ops :: GenerateOps
ops = GenerateOps
Gen.defaultGenerateOps {Gen.modelName = m, Gen.keepAlive = Just 0}
  Text
-> ByteString
-> Maybe GenerateOps
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError ()))
-> IO (Either OllamaError ())
forall payload response.
ToJSON payload =>
Text
-> ByteString
-> Maybe payload
-> Maybe OllamaConfig
-> (Response BodyReader -> IO (Either OllamaError response))
-> IO (Either OllamaError response)
withOllamaRequest Text
"/api//generate" ByteString
"POST" (GenerateOps -> Maybe GenerateOps
forall a. a -> Maybe a
Just GenerateOps
ops) Maybe OllamaConfig
forall a. Maybe a
Nothing Response BodyReader -> IO (Either OllamaError ())
forall a.
FromJSON a =>
Response BodyReader -> IO (Either OllamaError a)
commonNonStreamingHandler

{- | MonadIO version of 'loadGenModel' for use in monadic contexts.

Lifts the 'loadGenModel' function into a 'MonadIO' context, allowing it to be used in
monadic computations.

Example:

>>> import Control.Monad.IO.Class
>>> runReaderT (loadGenModelM "gemma3") someContext
Right ()
--
-- @since 0.2.0.0
-}
loadGenModelM :: MonadIO m => Text -> m (Either OllamaError ())
loadGenModelM :: forall (m :: * -> *).
MonadIO m =>
Text -> m (Either OllamaError ())
loadGenModelM Text
t = IO (Either OllamaError ()) -> m (Either OllamaError ())
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either OllamaError ()) -> m (Either OllamaError ()))
-> IO (Either OllamaError ()) -> m (Either OllamaError ())
forall a b. (a -> b) -> a -> b
$ Text -> IO (Either OllamaError ())
loadGenModel Text
t

{- | MonadIO version of 'unloadGenModel' for use in monadic contexts.

Lifts the 'unloadGenModel' function into a 'MonadIO' context, allowing it to be used in
monadic computations.

Example:

>>> import Control.Monad.IO.Class
>>> runReaderT (unloadGenModelM "gemma3") someContext
Right ()
--
-- @since 0.2.0.0
-}
unloadGenModelM :: MonadIO m => Text -> m (Either OllamaError ())
unloadGenModelM :: forall (m :: * -> *).
MonadIO m =>
Text -> m (Either OllamaError ())
unloadGenModelM Text
t = IO (Either OllamaError ()) -> m (Either OllamaError ())
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either OllamaError ()) -> m (Either OllamaError ()))
-> IO (Either OllamaError ()) -> m (Either OllamaError ())
forall a b. (a -> b) -> a -> b
$ Text -> IO (Either OllamaError ())
unloadGenModel Text
t